import copy
import logging
import os
import time
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Iterator,
    List,
    Literal,
    Optional,
    Tuple,
    Type,
    Union,
    cast,
)

import pandas as pd

from MelodieInfra import Config, MelodieExceptions
from MelodieInfra.parallel.parallel_manager import ParallelManager, ThreadParallelManager
from MelodieInfra.parallel.utils import resolve_parallel_mode
from MelodieInfra.utils import underline_to_camel

from .algorithms import AlgorithmParameters
from .algorithms.ga import MelodieGA
from .data_loader import DataLoader
from .model import Model
from .scenario_manager import Scenario
from .simulator import BaseModellingManager

logger = logging.getLogger(__name__)


class GACalibratorParams(AlgorithmParameters):
    def __init__(
        self,
        id: int,
        path_num: int,
        generation_num: int,
        strategy_population: int,
        mutation_prob: int,
        strategy_param_code_length: int,
        **kw,
    ):
        super().__init__(id, path_num)

        self.generation_num = generation_num
        self.strategy_population = strategy_population
        self.mutation_prob = mutation_prob
        self.strategy_param_code_length = strategy_param_code_length
        self.parse_params(kw)

    @staticmethod
    def from_dataframe_record(
        record: Dict[str, Union[int, float]]
    ) -> "GACalibratorParams":
        s = GACalibratorParams(
            record["id"],
            record["path_num"],
            record["generation_num"],
            record["strategy_population"],
            record["mutation_prob"],
            record["strategy_param_code_length"],
        )
        s.parse_params(record)
        return s

    def __hash__(self):
        return hash(
            (
                self.id,
                self.path_num,
                self.generation_num,
                self.strategy_population,
                self.mutation_prob,
                self.strategy_param_code_length,
            )
        )


class CalibratorAlgorithmMeta:
    """
    CalibratorAlgorithmMeta Records the status, including current scenario, params scenario, path and generation, of
     the calibrator.
    """

    def __init__(self):
        # If new attributes can be added by attribute assignment.
        self._freeze = False
        self.id_calibrator_scenario = 0
        self.id_calibrator_params_scenario = 1
        self.id_path = 0
        self.id_generation = 0

    def to_dict(self, public_only=False):
        if public_only:
            return {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
        return copy.copy(self.__dict__)

    def __repr__(self):
        return f"<{self.to_dict()}>"

    def __setattr__(self, key, value):
        if (not hasattr(self, "_freeze")) or (not self._freeze):
            super().__setattr__(key, value)
        else:
            if key in self.__dict__:
                super().__setattr__(key, value)
            else:
                raise MelodieExceptions.General.NoAttributeError(self, key)


class GACalibratorAlgorithmMeta(CalibratorAlgorithmMeta):
    def __init__(self):
        super(GACalibratorAlgorithmMeta, self).__init__()
        self.id_chromosome = 0
        self._freeze = True


class GACalibratorAlgorithm:
    """
    (Internal) The genetic algorithm implementation for the Calibrator.

    This class orchestrates the GA process, including managing the population of
    chromosomes (parameter sets), distributing simulation tasks to parallel

    workers, caching results, and evolving the population across generations.
    """

    def __init__(
        self,
        env_param_names: List[str],
        recorded_env_properties: List[str],
        recorded_agent_properties: Dict[str, List[str]],
        params: GACalibratorParams,
        target_func: "Callable[[Model], Union[float, int]]",
        manager: "Calibrator" = None,
        processors=1,
        parallel_mode: Optional[Literal["process", "thread"]] = None,
    ):
        self.manager = manager
        self.params = params
        self.chromosomes = 20
        self.target_func = target_func
        self.env_param_names = env_param_names
        self.recorded_env_properties = recorded_env_properties
        self.recorded_agent_properties = recorded_agent_properties
        lb, ub = self.params.bounds(self.env_param_names)
        self.algorithm: Optional[MelodieGA] = MelodieGA(
            cast("function", self.generate_target_function()),
            len(self.env_param_names),
            self.params.strategy_population,
            self.params.generation_num,
            self.params.mutation_prob,
            lb,
            ub,
            precision=1e-5,
        )
        self.cache: Dict[Tuple[int, int], float] = {}
        # self.agent_container_getters: Dict[str, Callable[[Model], AgentList]] = {}
        # self.agent_ids: Dict[str, List[int]] = {}  # {category : [agent_id]}
        # self.agent_params_defined: Dict[str, List[str]] = {}  # category: [param_name1, param_name2...]
        # self.recorded_agent_properties: Dict[str, List[str]] = {}  # category: [prop1, prop2...]
        # self.recorded_env_properties: List[str] = []  # category: [prop1, prop2...]

        self._chromosome_counter = 0
        self._current_generation = 0
        self.processors = processors
        self.parallel_mode = resolve_parallel_mode(parallel_mode)

        if self.parallel_mode == "process":
            d = {
                "model": (
                    self.manager.model_cls.__name__,
                    self.manager.model_cls.__module__,
                ),
                "scenario": (
                    self.manager.scenario_cls.__name__,
                    self.manager.scenario_cls.__module__,
                ),
                "trainer": (
                    self.manager.__class__.__name__,
                    self.manager.__class__.__module__,
                ),
                "data_loader": (
                    self.manager.df_loader_cls.__name__,
                    self.manager.df_loader_cls.__module__,
                ),
            }
            self.parallel_manager = ParallelManager(
                self.processors,
                configs=(d, self.manager.config.to_dict()),
                port=self.manager.config.parallel_port,
            )
            self.parallel_manager.run("calibrator")
        elif self.parallel_mode == "thread":
            self.parallel_manager = ThreadParallelManager(
                cores=self.processors,
                worker_func=self._thread_worker_func,
                worker_init_args=(self.manager,),
            )
            self.parallel_manager.run("calibrator")
    def _thread_worker_func(
        self, core_id: int, task: Tuple, calibrator: "Calibrator"
    ) -> Tuple:
        """
        Worker function for thread-based parallelism.
        Runs a single simulation for a given chromosome.
        """
        import logging
        import time
        from Melodie import Environment
        from MelodieInfra.config.global_configs import MelodieGlobalConfig

        # Set up logging for this thread worker (similar to process workers)
        thread_logger = logging.getLogger(f"Calibrator-processor-{core_id}")
        
        t0 = time.time()
        chrom, scenario_json, env_params = task
        thread_logger.debug(f"processor {core_id} got chrom {chrom}")
        
        scenario = calibrator.scenario_cls()
        scenario.manager = calibrator
        scenario._setup(scenario_json)
        scenario.set_params(env_params, asserts_key_exist=False)

        model = calibrator.model_cls(calibrator.config, scenario)
        model.create()
        model._setup()
        model.run()

        agent_data = {}
        for container_name, props in calibrator.recorded_agent_properties.items():
            agent_container = getattr(model, container_name)
            df = agent_container.to_list(props)
            agent_data[container_name] = df
            for row in df:
                row["agent_id"] = row.pop("id")

        env: Environment = model.environment
        env_data = env.to_dict(calibrator.watched_env_properties)
        env_data.update({prop: scenario.to_dict()[prop] for prop in calibrator.properties})
        env_data["target_function_value"] = env_data["distance"] = calibrator.distance(
            model
        )

        t1 = time.time()
        thread_logger.info(
            f"Processor {core_id}, chromosome {chrom}, time: {MelodieGlobalConfig.Logger.round_elapsed_time(t1 - t0)}s"
        )

        return (chrom, agent_data, env_data)

    def stop(self):
        self.parallel_manager.close()

    def get_params(self, id_chromosome: int) -> Dict[str, Any]:
        """
        Decode a chromosome into a dictionary of scenario parameters.

        :param id_chromosome: The index of the chromosome in the current
            population.
        :return: A dictionary mapping parameter names to their float values.
        """
        chromosome_value = self.algorithm.chrom2x(self.algorithm.Chrom)[id_chromosome]
        env_parameters_dict = {}
        for i, param_name in enumerate(self.env_param_names):
            # Convert numpy float to native Python float for serialization
            env_parameters_dict[param_name] = float(chromosome_value[i])
        return env_parameters_dict

    def target_function_to_cache(
        self,
        env_data,
        generation: int,
        id_chromosome: int,
    ):
        """
        (Internal) Store the output of the target function (distance) in a cache.
        """
        self.cache[(generation, id_chromosome)] = env_data["target_function_value"]

    def generate_target_function(self) -> Callable[[], float]:
        """
        (Internal) Create the fitness function for the genetic algorithm.

        This function retrieves pre-computed distance values from the cache for
        each chromosome, allowing the GA to evaluate fitness without re-running
        the simulation.
        """

        def f(*args):
            self._chromosome_counter += 1
            value = self.cache[(self._current_generation, self._chromosome_counter)]
            return value

        return f

    def record_agent_properties(
        self,
        agent_data: Dict[str, List[Dict[str, Any]]],
        env_data: Dict[str, Any],
        meta: GACalibratorAlgorithmMeta,
    ):
        """
        Record the property of each agent in the current chromosome.

        :param agent_data: {<agent_container_name>: [{id: 0, prop1: xxx, prop2: xxx, ...}]}
        :param env_data: {env_prop1: xxx, env_prop2: yyy, ...}
        :param meta:
        :return:
        """

        agent_records = {}
        environment_record = {}
        meta_dict = meta.to_dict(public_only=True)

        for container_name, _ in self.recorded_agent_properties.items():
            agent_records[container_name] = []
            data = agent_data[container_name]
            for agent_container_data in data:
                d = {}
                d.update(meta_dict)
                d.update(agent_container_data)

                agent_records[container_name].append(d)
        environment_record.update(meta_dict)
        environment_record.update(env_data)
        environment_record.pop("target_function_value")

        self.manager._write_to_table(
            "csv", "Result_Calibrator_Environment", pd.DataFrame([environment_record])
        )
        return agent_records, environment_record

    def calc_cov_df(
        self,
        agent_container_df_dict: "Dict[str, pd.DataFrame]",
        env_df: "pd.DataFrame",
        meta,
    ):
        """
        Calculate the coefficient of variation

        :param agent_container_df_dict:
        :param env_df:
        :param meta:
        :return:
        """

        pd.set_option("max_colwidth", 500)
        pd.set_option("display.max_columns", None)
        pd.set_option("display.max_rows", None)
        meta_dict = meta.to_dict(public_only=True)
        meta_dict.pop("id_chromosome")
        for container_name in self.recorded_agent_properties.keys():
            df = agent_container_df_dict[container_name]
            container_agent_record_list = []
            for agent_id in df["agent_id"].unique():
                agent_data = df.loc[df["agent_id"] == agent_id]
                cov_records = {}
                cov_records.update(meta_dict)
                cov_records["agent_id"] = agent_id
                for prop_name in self.recorded_agent_properties[container_name] + [
                    "distance"
                ]:
                    p: pd.Series = agent_data[prop_name]
                    mean = p.mean()
                    cov = float("nan") if mean == 0 else p.std() / mean
                    cov_records.update(
                        {prop_name + "_mean": mean, prop_name + "_cov": cov}
                    )
                container_agent_record_list.append(cov_records)

            self.manager._write_to_table(
                "csv",
                f"Result_Calibrator_{underline_to_camel(container_name)}_Cov",
                pd.DataFrame(container_agent_record_list),
            )
        env_record = {}
        env_record.update(meta_dict)
        for prop_name in (
            self.env_param_names + self.recorded_env_properties + ["distance"]
        ):
            mean = env_df[prop_name].mean()
            cov = float("nan") if mean == 0 else env_df[prop_name].std() / mean
            env_record.update({prop_name + "_mean": mean, prop_name + "_cov": cov})

        self.manager._write_to_table(
            "csv", "Result_Calibrator_Environment_Cov", pd.DataFrame([env_record])
        )

    def pre_check(self, meta):
        """
        Check at the beginning of `run()`

        :return:
        """
        logger.info(
            f"""Algorithm will run with:
        Meta value: {meta}
        Recording environment parameters: {self.recorded_env_properties}
        Recording Agent agent_lists: {self.recorded_agent_properties}\n"""
        )

    def run(self, scenario: Scenario, meta: Union[GACalibratorAlgorithmMeta]):
        self.pre_check(meta)

        for i in range(self.params.generation_num):
            t0 = time.time()
            self._current_generation = i
            meta.id_generation = i
            logger.info(
                f"======================="
                f"Path {meta.id_path} Generation {i + 1}/{self.params.generation_num}"
                f"======================="
            )

            for id_chromosome in range(self.params.strategy_population):
                params = self.get_params(id_chromosome)
                # params_queue.put(
                #     json.dumps((id_chromosome, scenario.to_json(), params))
                # )
                self.parallel_manager.put_task(
                    (id_chromosome, scenario.to_json(), params)
                )

            agent_records_collector: Dict[str, List[Dict[str, Any]]] = {
                container_name: []
                for container_name in self.recorded_agent_properties.keys()
            }
            env_records_list: List[Dict[str, Any]] = []

            for _id_chromosome in range(self.params.strategy_population):
                chrom, agents_data, env_data = self.parallel_manager.get_result()
                if chrom is None or agents_data is None or env_data is None:
                    raise RuntimeError(
                        "Parallel calibrator worker failed; see worker traceback above."
                    )
                meta.id_chromosome = chrom
                agent_records, env_record = self.record_agent_properties(
                    agents_data, env_data, meta
                )
                for container_name, records in agent_records.items():
                    agent_records_collector[container_name] += records
                env_records_list.append(env_record)
                self.target_function_to_cache(env_data, i, chrom)

            self.calc_cov_df(
                {k: pd.DataFrame(v) for k, v in agent_records_collector.items()},
                pd.DataFrame(env_records_list),
                meta,
            )

            self._chromosome_counter = -1
            self.algorithm.run(1)
            t1 = time.time()
            logger.info("=" * 20 + f"Time Elapsed: {t1 - t0}s" + "=" * 20)


class Calibrator(BaseModellingManager):
    """
    The ``Calibrator`` uses a genetic algorithm to tune scenario parameters.

    It aims to find the parameter set that minimizes the "distance" between a
    model's output and a predefined target (e.g., empirical data).
    """

    def __init__(
        self,
        config: "Config",
        scenario_cls: "Optional[Type[Scenario]]",
        model_cls: "Optional[Type[Model]]",
        data_loader_cls: Type["DataLoader"] = None,
        processors: int = 1,
        parallel_mode: Optional[Literal["process", "thread"]] = None,
    ):
        """
        :param config: The project :class:`~Melodie.Config` object.
        :param scenario_cls: The :class:`~Melodie.Scenario` subclass for the model.
        :param model_cls: The :class:`~Melodie.Model` subclass for the model.
        :param data_loader_cls: The :class:`~Melodie.DataLoader` subclass for the
            model.
        :param processors: The number of processor cores to use for parallel
            computation of the genetic algorithm.
        :param parallel_mode: The parallelization mode. If explicitly set to
            ``"process"`` or ``"thread"``, that value is used directly. If left
            as ``None``, Melodie auto-selects ``"thread"`` on Python 3.13+ and
            ``"process"`` on older Python versions.
        """
        super().__init__(
            config=config,
            scenario_cls=scenario_cls,
            model_cls=model_cls,
            data_loader_cls=data_loader_cls,
        )
        self.processes = processors
        self.parallel_mode = resolve_parallel_mode(parallel_mode)
        self.training_strategy: "Optional[Type[SearchingAlgorithm]]" = None
        self.container_name: str = ""

        self.properties: List[str] = []
        self.watched_env_properties: List[str] = []
        self.recorded_agent_properties: Dict[str, List[str]] = {}
        self.algorithm: Optional[GACalibratorAlgorithm] = None
        self.algorithm_instance: Dict[str, Any] = {}

        self.model: Optional[Model] = None

        self.current_algorithm_meta = GACalibratorAlgorithmMeta()

    def setup(self):
        """
        A hook for setting up the Calibrator.

        This method should be overridden in a subclass to define which
        scenario properties to calibrate using
        :meth:`add_scenario_calibrating_property`.
        """
        pass

    def collect_data(self):
        """
        (Optional) A hook to define which agent and environment properties to record.

        This is not required for calibration itself but is useful for saving
        detailed simulation data during the calibration process. Use
        :meth:`add_environment_property` to register properties.
        """
        pass

    def generate_scenarios(self) -> List["Scenario"]:
        """
        Generate scenarios from the ``CalibratorScenarios`` table.

        :return: A list of ``Scenario`` objects.
        """
        return self.data_loader.generate_scenarios("Calibrator")

    def get_params_scenarios(self) -> List:
        """
        Load the genetic algorithm parameters from the
        ``CalibratorParamsScenarios`` table.

        :return: A list of dictionaries, where each dictionary contains the GA
            parameters for one calibration run.
        """

        calibrator_scenarios_table = self.get_dataframe("CalibratorParamsScenarios")
        assert isinstance(
            calibrator_scenarios_table, pd.DataFrame
        ), "No learning scenarios table specified!"
        return calibrator_scenarios_table.to_dict(orient="records")

    def run(self):
        """
        The main entry point for starting the calibration process.
        """
        self.setup()
        self.pre_run()

        scenario_cls = GACalibratorParams
        for scenario in self.scenarios:
            self.current_algorithm_meta.id_calibrator_scenario = scenario.id
            calibration_scenarios = self.get_params_scenarios()
            for calibrator_scenario in calibration_scenarios:
                calibrator_scenario = scenario_cls.from_dataframe_record(
                    calibrator_scenario
                )
                self.current_algorithm_meta.id_calibrator_params_scenario = (
                    calibrator_scenario.id
                )
                for id_trainer_path in range(calibrator_scenario.path_num):
                    self.current_algorithm_meta.id_path = id_trainer_path

                    self.run_once_new(scenario, calibrator_scenario)

    def run_once_new(self, scenario: Scenario, calibration_params: GACalibratorParams):
        """
        (Internal) Run a single calibration path.
        """
        self.algorithm = GACalibratorAlgorithm(
            self.properties,
            self.watched_env_properties,
            {},
            calibration_params,
            self.target_function,
            manager=self,
            processors=self.processes,
            parallel_mode=self.parallel_mode,
        )
        self.algorithm.run(scenario, self.current_algorithm_meta)
        self.algorithm.stop()

    def target_function(self, model: "Model") -> Union[float, int]:
        """
        (Internal) A wrapper for the user-defined distance function.
        """
        return self.distance(model)

    def distance(self, model: "Model") -> float:
        """
        The distance function to be minimized by the calibrator.

        This method **must be overridden** in a subclass. It should take a
        ``Model`` object (representing the final state of a simulation run) and
        return a single float value representing the "distance" or error between
        the model's output and the desired target.

        :param model: The ``Model`` object after a simulation run.
        :return: A float representing the distance.
        """
        raise NotImplementedError(
            "Calibrator.distance(model) must be overridden in sub-class!"
        )

    def add_scenario_calibrating_property(self, prop: str):
        """
        Register a scenario property to be calibrated.

        The specified property will be tuned by the genetic algorithm within the
        bounds defined in the ``CalibratorParamsScenarios`` table.

        :param prop: The name of the scenario property to calibrate.
        """
        assert (
            prop not in self.properties
        ), f'Property "{prop}" is already in the calibrating training_properties!'
        self.properties.append(prop)

    def add_environment_property(self, prop: str):
        """
        Register an environment property to be recorded during calibration.

        This allows for saving the values of specific environment properties
        for each simulation run within the calibration process.

        :param prop: The name of the environment property to record.
        """
        assert prop not in self.watched_env_properties
        self.watched_env_properties.append(prop)
