from __future__ import annotations

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

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

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


def _default_n_processes() -> int:
    physical = psutil.cpu_count(logical=False)
    return max(1, int(physical) - 2)


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, dict, dict]:

    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"],
):
    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()
