from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Dict, Iterable, List, Tuple

import os
import yaml
import pickle
import psutil
import pandas as pd
import carbatpy as cb

from pymoo.core.problem import ElementwiseProblem
from pymoo.core.termination import Termination
from pymoo.termination.ftol import MultiObjectiveSpaceTermination
from pymoo.termination.max_gen import MaximumGenerationTermination
from pymoo.termination.robust import RobustTermination

def _default_n_processes() -> int:
    """
    Determine the default number of processes to use for parallelization.

    By default, this is the number of physical CPU cores minus 2, with a minimum of 1.
    If the number of physical cores cannot be determined, it falls back to the
    total logical cores, or 1 as a last resort.

    Returns:
        int: Recommended number of worker processes.
    """
    physical: int | None = psutil.cpu_count(logical=False)
    if physical is None:
        physical = psutil.cpu_count(logical=True) or 1

    return max(1, physical - 2)

def save_data(path: str, data: Any) -> None:
    """
    Saves the provided data to a file. The format is inferred from the file extension.

    Supported formats:
        - YAML (.yaml)
        - Pickle (.pkl)
        - CSV (.csv)

    Args:
        path (str): File path where the data will be saved.
        data (Any): Data to save. For CSV, must be a pandas DataFrame.

    Raises:
        ValueError: If the file extension is unknown.
    """
    file_type = os.path.splitext(path)[1].lower()
    match file_type:
        case ".yaml":
            with open(path, "w", encoding="utf-8") as f:
                yaml.safe_dump(data, f, sort_keys=False, allow_unicode=True, indent=2)
        case ".pkl":
            with open(path, "wb") as f:
                pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
        case ".csv":
            data.to_csv(path, index=False, sep=";", decimal=",")
        case _:
            raise ValueError("Unknown format {file_type}")


def get_data_from_file(path: str) -> Any | None:
    """
    Loads data from a file based on its extension.

    Supported formats:
        - YAML (.yaml)
        - Pickle (.pkl)
        - CSV (.csv)

    Args:
        path (str): Path to the file.

    Returns:
        Any | None: Loaded data, or None if loading fails.

    Raises:
        ValueError: If the file extension is unsupported.
    """
    file_type = os.path.splitext(path)[1].lower()
    try:
        match file_type:
            case ".yaml":
                with open(path, "r") as f:
                    return yaml.safe_load(f)
            case ".pkl":
                with open(path, "rb") as f:
                    return pickle.load(f)
            case ".csv":
                return pd.read_csv(path, sep=";", decimal=",")
            case _:
                raise ValueError(
                    f"Unsupported file format '{file_type}' for file: {path}"
                )
    except Exception as e:
        print(f"Error: {e}")
    return None


def extract_boundary_with_path(
    bounds: Dict[str, Any],
) -> Tuple[List[List[str]], Tuple[Any, ...], Tuple[Any, ...]]:
    """
    Extracts the variable paths and their lower/upper boundaries from a nested dictionary.

    Args:
        bounds (Dict[str, Any]): Nested dictionary with variable boundaries, e.g. {'param': [l, u]}.

    Returns:
        Tuple[List[List[str]], Tuple[Any, ...], Tuple[Any, ...]]:
            - List of variable paths (each path is a list of keys)
            - Tuple of lower bounds
            - Tuple of upper bounds
    """
    entries = traverse_dict(bounds)
    path = [p for (p, _, _) in entries]
    lower = tuple(l for (_, l, _) in entries)
    upper = tuple(u for (_, _, u) in entries)
    return path, lower, upper


def traverse_dict(
    d: Dict[str, Any],
    current_path: List[str] | None = None,
    entries: List[Tuple[List[str], Any, Any]] | None = None,
) -> List[Tuple[List[str], Any, Any]]:
    """
    Recursively traverses a nested dictionary to extract variable boundaries.

    Args:
        d (Dict[str, Any]): Dictionary to traverse.
        current_path (List[str] | None, optional): Current key path. Defaults to None.
        entries (List[Tuple[List[str], Any, Any]] | None, optional): Collected entries. Defaults to None.

    Returns:
        List[Tuple[List[str], Any, Any]]: List of tuples (path, lower, upper).
    """
    if current_path is None:
        current_path = []
    if entries is None:
        entries = []

    for key, val in d.items():
        new_path = current_path + [key]

        if isinstance(val, dict):
            traverse_dict(val, new_path, entries)
        elif key == "fractions":
            for i, (l, u) in enumerate(val):
                entries.append((new_path + [str(i)], l, u))
        else:
            l, u = val
            entries.append((new_path, l, u))

    return entries


def create_config(values: Iterable[float], paths: List[List[str]]) -> Dict[str, Any]:
    """
    Builds a nested configuration dictionary from flattened optimizer variables and paths.

    Fractions are handled specially: The last fraction is calculated as 1-sum(others).

    Args:
        values (Iterable[float]): Variable values from the optimizer.
        paths (List[List[str]]): Corresponding paths in the config dictionary.

    Returns:
        Dict[str, Any]: Nested dictionary usable by carbatpy.
    """
    config = {}
    normal_paths = []
    fraction_paths = []

    # Separate fractions from other parameters
    for path, value in zip(paths, values):
        if "fractions" in path:
            fraction_paths.append((path, float(value)))
        else:
            normal_paths.append((path, float(value)))

    # Create config with other parameters
    for path, value in normal_paths:
        current = config
        for key in path[:-1]:
            if key not in current:
                current[key] = {}
            current = current[key]
        current[path[-1]] = value

    # Create list with fractions (last one is 1-sum(fractions))
    fractions_by_base_path = {}
    for path, value in fraction_paths:
        base_path = path[:-1]
        index = path[-1]
        base_key = tuple(base_path)
        if base_key not in fractions_by_base_path:
            fractions_by_base_path[base_key] = {}
        fractions_by_base_path[base_key][index] = value

    for base_key, fractions_dict in fractions_by_base_path.items():
        fractions_list = [val for _, val in fractions_dict.items()]
        sum_existing = sum(fractions_list)
        missing_value = 1 - sum_existing
        fractions_list.append(missing_value)

        # Add fractions to config
        current = config
        for key in base_key[:-1]:
            if key not in current:
                current[key] = {}
            current = current[key]
        current[base_key[-1]] = fractions_list

    return config


def warning_score(res: Dict[str, Any], system_key: None | str = None) -> float:
    """
    Computes a total warning score from carbatpy simulation results.

    Args:
        res (Dict[str, Any]): Simulation results from carbatpy.
        system_key (str | None, optional): Specific system ("hp" or "orc") to compute warnings. Defaults to None for all.

    Returns:
        float: Absolute sum of warnings (0 if none).
    """
    w = res["warnings"] if system_key is None else res[system_key]["warnings"]
    return abs(sum(item.value for item in w.values()))


def apply_dt_min_start(
    conf_act: Dict[str, Any], opti_fun: str, paths: List[List[str]]
) -> Dict[str, Any]:
    """
    Sets starting delta-T (dT_min) values in the config based on heat exchanger values.

    Args:
        conf_act (Dict[str, Any]): Current configuration.
        opti_fun (str): Cycle type ("hp", "orc", "cb").
        paths (List[List[str]]): Paths to check if dT_min is optimized.

    Returns:
        Dict[str, Any]: Updated configuration with starting dT_min.
    """
    if not any("dt_min" in p for p in paths):
        return conf_act

    if opti_fun == "cb":
        hp = conf_act.setdefault("hp", {})
        orc = conf_act.setdefault("orc", {})

        hp_start = hp.setdefault("start", {})
        orc_start = orc.setdefault("start", {})

        hp_start["dt_min"] = hp["evaporator"]["dt_min"] + 0.001
        orc_start["dt_min"] = orc["condenser"]["dt_min"]

    elif opti_fun == "hp":
        start = conf_act.setdefault("start", {})
        start["dt_min"] = conf_act["evaporator"]["dt_min"]

    elif opti_fun == "orc":
        start = conf_act.setdefault("start", {})
        start["dt_min"] = conf_act["condenser"]["dt_min"]

    return conf_act


def calc_costs(
    results: Dict[str, Any], system: str, verbose: bool = False
) -> Tuple[float, Dict[str, Any], Dict[str, Any]]:
    """
    Calculate total system costs and distribution for HP, ORC, or CB.

    For CB, distinguishes between HP and ORC components to avoid double counting.

    Args:
        results (Dict[str, Any]): Simulation results with costs.
        system (str): System type ("hp", "orc", or "cb").
        verbose (bool, optional): Print details about ORC components. Defaults to False.

    Returns:
        Tuple[float, Dict[str, Any], Dict[str, Any]]: Total cost, HP costs, ORC costs.
    """
    select = ("total_costs", "cost_distribution")

    if system in ("hp", "orc"):
        return results["costs"]["total_costs"], {}, {}

    if system != "cb":
        raise ValueError(system)

    costs_hp = {k: results["hp"]["costs"][k] for k in select}
    costs_orc = {k: results["orc"]["costs"][k] for k in select}

    hp_total = sum(costs_hp["cost_distribution"].values())

    new_orc_components = {
        "expander",
        "pump",
        "condenser",
        "evaporator",
    }

    orc_unique = sum(
        c
        for comp, c in costs_orc["cost_distribution"].items()
        if comp in new_orc_components
    )

    if verbose:
        print(
            f"CB cost assumption: ORC unique components counted: {sorted(new_orc_components)}"
        )

    return hp_total + orc_unique, costs_hp, costs_orc


@dataclass(frozen=True)
class EvalResult:
    """
    Container for the results of a carbatpy optimization evaluation.

    Attributes:
        performance (float): Performance metric of the system (e.g., COP or thermal efficiency).
        costs (float): Total system cost calculated for the optimization case.
        p_low (float): Minimum pressure value in the system, used for constraints.
        warn (float): Warning score representing thermodynamic or design violations.
    """
    performance: float
    costs: float
    p_low: float
    warn: float


def opti_func(
    x: Iterable[float],
    *,
    dir_config: Any,
    paths: List[List[str]],
    opti_fun: str,
    same_fluid: bool = True,
    heat_losses: float = 0.0,
    COP: float | None = None,
    q_dot: float | None = None,
) -> EvalResult:
    """
    Performs a carbatpy simulation for the specified cycle (HP, ORC, or CB) 
    and calculates the evaluation metrics for optimization.

    The function first constructs a complete configuration from the optimizer
    variables `x` and the corresponding `paths`. It then simulates the chosen
    cycle using carbatpy. Any warnings are assessed via `warning_score()`.  
    In case of violations or infeasible results, heavily penalized values 
    are returned to guide the optimizer away from invalid solutions.

    Args:
        x (Iterable[float]): Variable values generated by the optimizer.
        dir_config (Any): General configuration for carbatpy.
        paths (List[List[str]]): Paths indicating where each value in `x` should 
                                 be placed in the configuration dictionary.
        opti_fun (str): Cycle to optimize ("hp", "orc", or "cb").
        same_fluid (bool, optional): Whether to reuse the HP working fluid in the ORC. Defaults to True.
        heat_losses (float, optional): Thermal losses in the storage components. Defaults to 0.0.
        COP (float | None, optional): COP of the HP for ORC optimization. Required if `opti_fun="orc"`.
        q_dot (float | None, optional): Heat flow of the HP for ORC optimization. Required if `opti_fun="orc"`.

    Raises:
        ValueError: If an unknown cycle is specified or required parameters are missing.

    Returns:
        EvalResult: Contains performance metrics for the optimizer:
            - `performance`: Cycle performance (e.g., COP or efficiency)
            - `costs`: Total system costs
            - `p_low`: Low-side pressure at start
            - `warn`: Warning score (0 = no violations; high values indicate infeasible solutions)

    Notes:
        - In the presence of warnings, costs are heavily penalized (1e12) 
          to discourage infeasible solutions in the optimization process.
        - This function serves as the main interface between the optimizer and carbatpy.
    """

    try:
        conf_act = create_config(x, paths)
        conf_act = apply_dt_min_start(conf_act, opti_fun, paths)

        if opti_fun == "cb":
            performance, res = cb.cb_comp.cb_calc(
                dir_config,
                same_fluid=same_fluid,
                heat_losses=heat_losses,
                config=conf_act,
                plotting=False,
                verbose=False,
            )

            p_low = res["hp"]["output"]["start"]["p_low"]

            warn_hp = warning_score(res, system_key="hp")
            if warn_hp > 0:
                return EvalResult(0.0, 1e12, p_low, 100.0 + float(warn_hp))
            warn_orc = warning_score(res, system_key="orc")
            if warn_orc > 0:
                return EvalResult(0.0, 1e12, p_low, float(warn_orc))

            costs, _, _ = calc_costs(res, opti_fun)
            return EvalResult(float(performance), float(costs), p_low, 0.0)

        elif opti_fun == "hp":
            res = cb.hp_comp.heat_pump(
                dir_config, config=conf_act, verbose=False, plotting=False
            )

            p_low = res["output"]["start"]["p_low"]
            warn = float(warning_score(res))
            if warn > 0:
                return EvalResult(0.0, 1e12, p_low, warn)

            performance = res["COP"]
            costs, _, _ = calc_costs(res, opti_fun)

            return EvalResult(performance, float(costs), p_low, 0.0)

        elif opti_fun == "orc":
            res = cb.orc_comp.orc(
                dir_config, COP, q_dot, config=conf_act, verbose=False, plotting=False
            )

            p_low = res["output"]["start"]["p_low"]
            warn = float(warning_score(res))

            if warn > 0:
                return EvalResult(0.0, 1e12, 1e4, warn)

            performance = res["eta_th"]
            costs, _, _ = calc_costs(res, opti_fun)

            return EvalResult(performance, float(costs), p_low, 0.0)

        raise ValueError(f"Unknown opti_fun={opti_fun}")

    except:
        return EvalResult(0.0, 1e12, 1e4, 1e6)


# Pymoo initialization
class CombinedTermination(Termination):
    """
    Termination criteria combining multiple conditions for Pymoo optimizations.

    Combines:
      - Maximum number of generations
      - Multi-objective function tolerance (ftol) with robust checking over periods

    Attributes:
        max_gen_termination (MaximumGenerationTermination): Termination based on max generations.
        ftol_termination (RobustTermination): Termination based on objective function tolerance.
        criteria (List[Termination]): List of criteria combined.
    """
    def __init__(
        self,
        max_gen: int,
        ftol: float,
        period: int,
        n_skip: int = 5,
    ) -> None:
        """
        Initializes the combined termination.

        Args:
            max_gen (int): Maximum number of generations.
            ftol (float): Tolerance for objective function changes.
            period (int): Number of generations to check tolerance over.
            n_skip (int, optional): Number of initial generations to skip for robust checking. Defaults to 5.
        """

        super().__init__()

        self.max_gen_termination = MaximumGenerationTermination(max_gen)

        self.ftol_termination = RobustTermination(
            MultiObjectiveSpaceTermination(ftol, only_feas=True, n_skip=n_skip),
            period=period,
        )

        # Combine all criteria
        self.criteria = [
            self.ftol_termination,
            self.max_gen_termination,
        ]

    def _update(self, algorithm):
        """
        Update termination criteria based on current algorithm state.

        Args:
            algorithm: Pymoo algorithm instance.

        Returns:
            int: Maximum value of all termination criteria (1 if any reached, else 0).
        """
        p = [criterion.update(algorithm) for criterion in self.criteria]
        criteria_names = [
            "Objective values tolerance",
            "Max number of generations",
        ]
        for i in range(len(p)):
            if p[i] == 1:
                print(f"{criteria_names[i]} reached!")

        return max(p)


# Define class for pymoo
class opti_Problem(ElementwiseProblem):
    """
    Pymoo ElementwiseProblem wrapper for carbatpy optimizations.

    Handles different optimization modes (HP, ORC, CB) and sets up variables, 
    bounds, constraints, and the evaluation function to interface with Pymoo.

    Attributes:
        dir_config (Any): General configuration passed to carbatpy.
        paths_var (List[List[str]]): Paths to variables in the config.
        opti_fun (str): Optimized cycle ("hp", "orc", or "cb").
        same_fluid (bool): Whether fluids are reused across cycles.
        heat_losses (float): Heat loss fraction in CB mode.
        COP (float | None): COP used for ORC optimizations.
        q_dot (float | None): Heat flow used for ORC optimizations.
    """
    def __init__(self, **kwargs):
        self.dir_config: Any = kwargs.get("dir_config", None)
        self.paths_var: List[List[str]] = kwargs.get("paths", [])
        self.opti_fun: str = kwargs.get("opti_fun", "")
        self.same_fluid: bool = kwargs.get("same_fluid", True)
        self.heat_losses: float | int = kwargs.get("heat_losses", 0.0)

        # Needed for ORC optimization
        self.COP = kwargs.get("COP", None)
        self.q_dot = kwargs.get("q_dot", None)

        super().__init__(**kwargs)

    def _evaluate(self, x, out) -> None:
        """
        Evaluate a single solution.

        Args:
            x: List or array of decision variable values.
            out: Dictionary to store objective ('F') and constraint ('G') results.

        Sets:
            out["F"]: Objective value(s).
            out["G"]: Constraint violations (e.g., minimum performance, pressure limits, warnings).
        """
        res = opti_func(
            x,
            dir_config=self.dir_config,
            paths=self.paths_var,
            opti_fun=self.opti_fun,
            same_fluid=self.same_fluid,
            heat_losses=self.heat_losses,
            COP=self.COP,
            q_dot=self.q_dot,
        )

        PERF_MIN = 0.05
        HP_P_LOW_MIN = 1e5

        g1 = PERF_MIN - res.performance
        g2 = (HP_P_LOW_MIN - res.p_low) / HP_P_LOW_MIN
        g3 = res.warn
        out["G"] = [g1, g2, g3]

        f1 = -res.performance

        if self.n_obj == 1:
            out["F"] = f1

        elif self.n_obj == 2:
            f2 = res.costs / 1e6  # Scale to 1e6 for better handling

            out["F"] = [f1, f2]
