from __future__ import annotations

import multiprocessing
from typing import Any, Dict, Optional, Tuple, List

from carbatpy.optimizations.helpers_optimization import (
    opti_Problem,
    CombinedTermination,
    extract_boundary_with_path,
    _default_n_processes,
    )

from pymoo.algorithms.moo.nsga2 import NSGA2
from pymoo.optimize import minimize
from pymoo.parallelization.starmap import StarmapParallelization
from pymoo.operators.sampling.lhs import LHS
from pymoo.operators.survival.rank_and_crowding import RankAndCrowding


DEFAULT_OPT: Dict[str, Any] = {
    "verbose": True,
    "n_processes": None,
    "maxtasksperchild": 20,
    "n_ieq_constr": 3,
    "n_gen": 100,
    "period": 20,
    "ftol": 1e-4,
    "pop_size": 50,
    "n_offspring": None,
    "sampling_iterations": 50,
    "crowding_func": "pcd",
    "eliminate_duplicates": True,
    "save_history": False,
    "return_least_infeasible": False,
    "to_optimize": ["RTE", "CAPEX"],
}


def _build_case(
    mode: str,
    boundaries: dict,
    same_fluid: Optional[bool],
    heat_losses: Optional[float],
) -> Tuple[str, dict, dict]:
    """
    Prepare the optimization case for NSGA2.

    Currently, only the 'cb' (combined) mode is supported for NSGA2. 
    Other modes will raise a ValueError.

    Args:
        mode (str): Optimization mode. Must be 'cb'.
        boundaries (dict): Boundaries for all optimization variables.
        same_fluid (Optional[bool]): Whether the working fluid is the same across all cycles.
        heat_losses (Optional[float]): Fraction of heat losses in the system.

    Raises:
        ValueError: If the mode is not 'cb'.

    Returns:
        Tuple[str, dict, dict]:
            - fun_to_optimize (str): Function identifier used in opti_Problem.
            - bounds (dict): Boundaries dictionary.
            - problem_kwargs (dict): Additional keyword arguments required by the problem.
    """
    if mode == "hp":
        raise ValueError("Mode HP not yet implemented for NSGA2")

    if mode == "orc":
        raise ValueError("Mode ORC not yet implemented for NSGA2")

    if mode == "cb":
        return "cb", boundaries, {"same_fluid": same_fluid, "heat_losses": heat_losses}

    raise ValueError("mode must be one of: 'hp' | 'orc' | 'cb'")


def optimize(
    mode: str,
    config: dict | str,
    boundaries: dict,
    same_fluid: Optional[bool] = True,
    heat_losses: Optional[float] = 0.0,
    to_optimize: List[str] = DEFAULT_OPT["to_optimize"],
    verbose: bool = DEFAULT_OPT["verbose"],
    n_processes: Optional[int] = DEFAULT_OPT["n_processes"],
    maxtasksperchild: int = DEFAULT_OPT["maxtasksperchild"],
    n_ieq_constr: int = DEFAULT_OPT["n_ieq_constr"],
    pop_size: int = DEFAULT_OPT["pop_size"],
    n_offspring: Optional[int] = DEFAULT_OPT["n_offspring"],
    sampling_iterations: int = DEFAULT_OPT["sampling_iterations"],
    crowding_func: str = DEFAULT_OPT["crowding_func"],
    eliminate_duplicates: bool = DEFAULT_OPT["eliminate_duplicates"],
    n_gen: int = DEFAULT_OPT["n_gen"],
    period: int = DEFAULT_OPT["period"],
    ftol: float = DEFAULT_OPT["ftol"],
    save_history: bool = DEFAULT_OPT["save_history"],
    return_least_infeasible: bool = DEFAULT_OPT["return_least_infeasible"],
):
    """
    Perform a multi-objective optimization using NSGA2.

    Sets up the problem, initializes the NSGA2 algorithm, handles parallelization,
    and executes the optimization using Pymoo.

    Args:
        mode (str): Optimization mode. Must currently be 'cb'.
        config (dict | str): Configuration dictionary or path to configuration file.
        boundaries (dict): Lower and upper bounds for all optimization variables.
        same_fluid (Optional[bool]): Whether the working fluid is the same ('cb' mode only).
        heat_losses (Optional[float]): Fraction of heat losses ('cb' mode only).
        to_optimize (List[str]): List of objective names to optimize.
        verbose (bool): Whether to print progress information.
        n_processes (Optional[int]): Number of parallel processes. Auto-detected if None.
        maxtasksperchild (int): Maximum number of tasks per worker process.
        n_ieq_constr (int): Number of inequality constraints in the problem.
        pop_size (int): Population size for NSGA2.
        n_offspring (Optional[int]): Number of offspring per generation. Defaults to NSGA2 rules if None.
        sampling_iterations (int): Iterations for Latin Hypercube sampling.
        crowding_func (str): Crowding distance function for survival selection.
        eliminate_duplicates (bool): Whether to remove duplicate solutions in population.
        n_gen (int): Maximum number of generations.
        period (int): Period for robust termination checking.
        ftol (float): Function tolerance for termination criteria.
        save_history (bool): Whether to save the optimization history.
        return_least_infeasible (bool): Whether to return the least infeasible solution if no feasible solution is found.

    Raises:
        ValueError: If `boundaries` is None or the mode is not supported.

    Returns:
        pymoo.model.result.Result: Result object containing the optimization outcomes, including
        the Pareto front, best solutions, and other metadata.
    """
    multiprocessing.freeze_support()

    if boundaries is None:
        raise ValueError("boundaries must be provided.")

    if n_processes is None:
        n_processes = _default_n_processes()

    fun_to_optimize, bounds, problem_kwargs = _build_case(
        mode,
        boundaries=boundaries,
        same_fluid=same_fluid,
        heat_losses=heat_losses,
    )

    paths, xl, xu = extract_boundary_with_path(bounds)

    pool = multiprocessing.Pool(n_processes, maxtasksperchild=maxtasksperchild)
    runner = StarmapParallelization(pool.starmap)
    try:
        problem = opti_Problem(
            dir_config=config,
            paths=paths,
            opti_fun=fun_to_optimize,
            n_var=len(paths),
            n_obj=len(to_optimize),
            n_ieq_constr=n_ieq_constr,
            xl=xl,
            xu=xu,
            elementwise_runner=runner,
            **problem_kwargs,
        )

        sampling = LHS(iterations=sampling_iterations)
        survival = RankAndCrowding(crowding_func=crowding_func)

        algorithm = NSGA2(
            pop_size=pop_size,
            n_offsprings=n_offspring,
            sampling=sampling,
            survival=survival,
            eliminate_duplicates=eliminate_duplicates,
        )

        termination = CombinedTermination(max_gen=n_gen, period=period, ftol=ftol)

        result = minimize(
            problem,
            algorithm,
            termination=termination,
            save_history=save_history,
            return_least_infeasible=return_least_infeasible,
            verbose=verbose,
        )
        return result

    finally:
        pool.close()
        pool.join()
