import json
import os
import unittest.mock
from functools import partial

import allure
import pytest
from pathlib import Path

from idmtools import IdmConfigParser
from idmtools.core.platform_factory import Platform
from idmtools.entities.simulation import Simulation
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_dummy_suite
from idmtools_test import COMMON_INPUT_PATH
from idmtools_test.utils.test_execute_platform import clear_execute_platform
from idmtools_test.utils.test_task import TestTask
from idmtools.entities.experiment import Experiment
from idmtools_test.utils.utils import get_performance_scale, clear_id_cache


@pytest.mark.serial
@pytest.mark.assets
@pytest.mark.smoke
@allure.story("Entities")
@allure.story("Plugins")
@allure.suite("idmtools_core")
class TestItemSequence(unittest.TestCase):

    def setUp(self):
        clear_execute_platform()

    @classmethod
    def tearDownClass(cls) -> None:
        clear_id_cache()
        clear_execute_platform()

    @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

    @pytest.mark.serial
    @pytest.mark.performance
    def test_1(self):
        clear_id_cache()
        parser = IdmConfigParser()
        parser._load_config_file(file_name='idmtools_1.ini')
        sequence_file = self.get_sequence_file()
        platform = Platform('MTest', missing_ok=True, job_directory="DEST")
        task = JSONConfiguredPythonTask(script_path=os.path.join(COMMON_INPUT_PATH, "python", "model1.py"),
                                        envelope="parameters", parameters=(dict(c=0)))
        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(10 * get_performance_scale()))
        builder.add_sweep_definition(partial(param_update, param="b"), range(50))
        e.simulations.add_builder(builder)
        add_dummy_suite(e)
        e.run(platform=platform, wait_until_done=False)