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

import os
import yaml
import pickle
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


CSV_KW = dict(sep=";", decimal=",")


def save_data(path: str, data: Any) -> None:
    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, **CSV_KW)
        case _:
            raise ValueError("Unknown format {file_type}")


def get_data_from_file(path):
    file_type = os.path.splitext(path)[1].lower()
    try:
        match file_type:
            case ".yaml":
                with open(path, "r") as f:
                    data = yaml.safe_load(f)
            case ".pkl":
                with open(path, "rb") as f:
                    data = pickle.load(f)
            case ".csv":
                data = pd.read_csv(path, **CSV_KW)
            case _:
                raise ValueError(
                    f"Unsupported file format '{file_type}' for file: {path}"
                )
    except:
        data = None
    return data


def extract_boundary_with_path(bounds: Dict[str, Any]):
    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=None, entries=None):
    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]:
    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, system_key=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, opti_fun, paths):
    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, system, verbose=False):
    select = ("total_costs", "cost_distribution")

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

    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:
    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:

    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):
    def __init__(
        self,
        max_gen: int,
        ftol: float,
        period: int,
        n_skip: int = 5,
    ) -> None:

        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):
        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):
    def __init__(self, **kwargs):
        self.dir_config = kwargs.get("dir_config", None)
        self.paths_var = kwargs.get("paths", None)
        self.opti_fun = kwargs.get("opti_fun", None)
        self.same_fluid = kwargs.get("same_fluid", True)
        self.heat_losses = 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):

        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]
