from __future__ import annotations

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

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

from pymoo.algorithms.soo.nonconvex.de import DE
from pymoo.optimize import minimize
from pymoo.parallelization.starmap import StarmapParallelization
from pymoo.operators.sampling.lhs import LHS


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,
    "sampling_iterations": 50,
    "de_variant": "DE/rand/1/bin",
    "CR": 0.3,
    "dither": "vector",
    "jitter": False,
    "return_least_infeasible": True,
}


def _build_case(
    mode: str,
    *,
    boundaries: dict,
    same_fluid: Optional[bool],
    heat_losses: Optional[float],
    COP: Optional[float],
    q_dot: Optional[float],
) -> Tuple[str, Dict[str, Any], Dict[str, Any]]:
    """
    Prepare all cycle-specific inputs for the optimization.

    Depending on the optimization mode, this function returns the function
    type to optimize, the boundaries for the decision variables, and additional
    parameters required for the computation.

    Args:
        mode (str): Optimization mode. Must be one of 'hp', 'orc', or 'cb'.
        boundaries (dict): Dictionary containing lower and upper bounds for 
            each optimization variable.
        same_fluid (Optional[bool]): Whether the same working fluid should be used 
            (required for 'cb' mode).
        heat_losses (Optional[float]): Heat losses fraction (required for 'cb' mode).
        COP (Optional[float]): Coefficient of performance (required for 'orc' mode).
        q_dot (Optional[float]): Heat input/output rate (required for 'orc' mode).

    Raises:
        ValueError: If `mode` is invalid or required parameters for a mode are missing.

    Returns:
        Tuple[str, Dict[str, Any], Dict[str, Any]]: 
            - Function type for optimization ('hp', 'orc', or 'cb').
            - Boundaries dictionary for optimization variables.
            - Additional keyword arguments required by the optimization function.
    """

    if mode == "hp":
        return "hp", boundaries, {}

    if mode == "orc":
        if COP is None or q_dot is None:
            raise ValueError("For mode='orc' you must provide COP and q_dot.")
        return "orc", boundaries, {"COP": COP, "q_dot": q_dot}

    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,
    COP: Optional[float] = None,
    q_dot: Optional[float] = None,
    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"],
    sampling_iterations: int = DEFAULT_OPT["sampling_iterations"],
    de_variant: str = DEFAULT_OPT["de_variant"],
    CR: float = DEFAULT_OPT["CR"],
    dither: str = DEFAULT_OPT["dither"],
    jitter: bool = DEFAULT_OPT["jitter"],
    n_gen: int = DEFAULT_OPT["n_gen"],
    period: int = DEFAULT_OPT["period"],
    ftol: float = DEFAULT_OPT["ftol"],
    return_least_infeasible: bool = DEFAULT_OPT["return_least_infeasible"],
):
    """
    Perform the global optimization for the selected cycle.

    This function sets up the optimization problem, initializes the
    differential evolution algorithm, handles parallelization, and
    executes the optimization using Pymoo.

    Args:
        mode (str): Optimization mode ('hp', 'orc', or '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 for all cycles ('cb' mode only).
        heat_losses (Optional[float]): Heat losses fraction ('cb' mode only).
        COP (Optional[float]): Coefficient of performance ('orc' mode only).
        q_dot (Optional[float]): Heat input/output rate ('orc' mode only).
        verbose (bool): Whether to print detailed optimization progress.
        n_processes (Optional[int]): Number of processes for parallelization. Defaults to auto-detection.
        maxtasksperchild (int): Maximum tasks per worker process in the pool.
        n_ieq_constr (int): Number of inequality constraints in the optimization problem.
        pop_size (int): Population size for the DE algorithm.
        sampling_iterations (int): Number of iterations for the Latin Hypercube sampling.
        de_variant (str): Variant of the differential evolution algorithm.
        CR (float): Crossover probability for DE.
        dither (str): Dithering strategy for DE.
        jitter (bool): Whether to apply jitter in DE.
        n_gen (int): Maximum number of generations.
        period (int): Period for robust termination checking.
        ftol (float): Function tolerance for termination.
        return_least_infeasible (bool): Whether to return the least infeasible solution if optimization fails.

    Raises:
        ValueError: If `boundaries` is None or required parameters for the mode are missing.

    Returns:
        pymoo.model.result.Result: The optimization result object containing the best solution and additional 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,
        COP=COP,
        q_dot=q_dot,
    )

    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=1,
            n_ieq_constr=n_ieq_constr,
            xl=xl,
            xu=xu,
            elementwise_runner=runner,
            **problem_kwargs,
        )

        sampling = LHS(iterations=sampling_iterations)

        algorithm = DE(
            pop_size=pop_size,
            sampling=sampling,
            variant=de_variant,
            CR=CR,
            dither=dither,
            jitter=jitter,
        )

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

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

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