import os
import unittest
from functools import partial
from pathlib import Path

from idmtools import IdmConfigParser
from idmtools.core.platform_factory import Platform
from idmtools.entities.experiment import Experiment
from idmtools.entities.templated_simulation import TemplatedSimulations
from idmtools_models.python.json_python_task import JSONConfiguredPythonTask
from idmtools_platform_slurm.platform_operations.utils import add_dammy_suite
from idmtools_test import COMMON_INPUT_PATH

class TestItemSequence(unittest.TestCase):
    def tearDown(self) -> None:
        IdmConfigParser.clear_instance()
        super().tearDown()

    @staticmethod
    def get_sequence_file():
        sequence_file = Path(IdmConfigParser.get_option("item_sequence", "sequence_file", 'item_sequences.json')).expanduser()
        if sequence_file.exists():
            sequence_file.unlink()
        return sequence_file

    def test_id(self):
        parser = IdmConfigParser()
        parser._load_config_file(file_name='idmtools_slurm.ini')
        parser.ensure_init(file_name='idmtools_slurm.ini', force=True)
        sequence_file = self.get_sequence_file()
        task = JSONConfiguredPythonTask(script_path=os.path.join(COMMON_INPUT_PATH, "python", "model3.py"),
                                        envelope="parameters", parameters=(dict(c=0)))
        #job_directory = os.path.join(os.path.expanduser('~'), "DESC")
        platform = Platform('SLURM_TEST')  #, job_directory=job_directory)
        ts = TemplatedSimulations(base_task=task)
        e = Experiment.from_template(ts)
        from idmtools.builders import SimulationBuilder
        builder = SimulationBuilder()

        def param_update(simulation, param, value):
            return simulation.task.set_parameter(param, value)

        builder.add_sweep_definition(partial(param_update, param="a"), range(3))
        builder.add_sweep_definition(partial(param_update, param="b"), range(3))
        e.simulations.add_builder(builder)
        simulations = e.simulations.items
        print("Before run")
        [print(sim.id) for sim in simulations]
        suite = add_dammy_suite(e)
        suite.run(platform=platform, dry_run=True, wait_until_done=False, wait_on_done=False)
        print("After run")
        simulations = e.simulations.items
        [print(sim.id) for sim in simulations]