"""Fit PNIPAAm-water parameters to binodal, LCST, and density data
across molar masses."""

# pylint: disable=duplicate-code

from __future__ import annotations

import json
from pathlib import Path
import numpy as np
from equi_gel.optimize_parameters import (
    ObjectiveWrapperBase,
    ParameterOptimization,
)
from equi_gel.cost_func_binodal import get_minimal_distance
from equi_gel.alternating_tangents import find_lcst_temperature
from equi_gel.interface_feos import (
    get_feos_npt_state,
    is_density_iteration_sane,
)

from equi_gel.newton_raphson_tracing import (
    build_initial_guess_binary,
    interpret_y_vector,
    trace_parameter,
)

from equi_gel.parameters import (
    get_parameter_set_from_tables,
    scale_component_molarweight,
)
from equi_gel.save_load_json import (
    BASE_FOLDER_EXP,
    load_lle_record,
    resample_exp_data,
)
from scipy.interpolate import interp1d
from si_units import GRAM, METER

components = ["pnipaam", "water"]
SOURCE = "arndt"

BASE_SYSTEM = {
    "number_components": 2,
    "gel": False,
    "elasticity": "",
    "C": 0,
    "ref_volume_frac": 0.0,
    "max_volume_poly_frac": 0.0,
}
PRESSURE = 1.13
T_HIGH = 350
T_LOW = 290
T_LCST = 300
MN_LIST = [3000.0, 6000.0, 9000.0]
EXP_PNIPAAM_WATER_AFROZE_ALL = "EXP_SOL_PNIPAAM_water_Afroze_Mn36*.json"
EXP_PNIPAAM_DENS_KANY = BASE_FOLDER_EXP / "EXP_POLY_DENS_Kany.json"
MASS_DENSITY_CONVERSION = GRAM / (METER * 1e-2) ** 3

T_SCALE = 100  # K per percent for binodal fitting


def interpolate_dict(dict1: dict, dict2: dict, lam: float) -> dict:
    """Return interpolated dictionary (1 - lam)*dict1 + lam*dict2."""
    inter_dict = {}
    for key in dict1.keys():
        if isinstance(dict1[key], (int, float)):
            inter_dict[key] = (1 - lam) * dict1[key] + lam * dict2[key]
        else:
            inter_dict[key] = dict1[key]

    return inter_dict


def load_exp_pnipaam_water_afroze_all(
    pattern: str = EXP_PNIPAAM_WATER_AFROZE_ALL,
) -> list[dict]:
    """Load all Afroze PNIPAAm-water experimental solubilities."""
    paths = sorted(BASE_FOLDER_EXP.glob(pattern))
    if not paths:
        raise FileNotFoundError(
            f"Missing experimental files for pattern: {pattern}"
        )
    return [load_lle_record(path) for path in paths]


def load_exp_pnipaam_density_kany(_path=EXP_PNIPAAM_DENS_KANY):
    """Load PNIPAAm-water density records from Kany dataset."""
    if not _path.exists():
        raise FileNotFoundError(f"Missing experimental file: {_path}")
    data = json.loads(_path.read_text())
    return next(iter(data.values()))


def _extract_constant_value(constant_rows, key):
    """Extract a float value from the Constant Value table."""
    for row in constant_rows:
        if not row:
            continue
        if row[0] == key and len(row) > 1:
            try:
                return float(row[1])
            except (TypeError, ValueError):
                return None
    return None


def _interpolate_pure_water_density(w_vals, rho_vals):
    """Interpolate/extrapolate pure-water density at w1=0.0"""
    w_arr = np.asarray(w_vals, dtype=float)
    rho_arr = np.asarray(rho_vals, dtype=float)
    order = np.argsort(w_arr)
    w_sorted = w_arr[order]
    rho_sorted = rho_arr[order]
    interp = interp1d(
        w_sorted,
        rho_sorted,
        kind="linear",
        fill_value="extrapolate",
    )
    return float(interp(0.0))


def _build_exp_pure_water_table(exp_density_records):
    """Build rho_water(T) table from experimental density records."""
    temp_to_rho = {}
    for record in exp_density_records:
        constant_vals = record.get("Constant Value", [])
        temperature = _extract_constant_value(constant_vals, "Temperature")
        data_rows = record.get("Data Table", [])
        if temperature is None or len(data_rows) < 2:
            continue
        data = data_rows[1:]
        w1 = [row[0] for row in data]
        rho = [row[1] for row in data]
        rho_pure = _interpolate_pure_water_density(w1, rho)
        if rho_pure is None:
            continue
        temp_to_rho.setdefault(float(temperature), []).append(rho_pure)

    if not temp_to_rho:
        return np.array([], dtype=float), np.array([], dtype=float)

    temps = np.array(sorted(temp_to_rho.keys()), dtype=float)
    rho_vals = np.array(
        [np.mean(temp_to_rho[temp]) for temp in temps], dtype=float
    )
    return temps, rho_vals


def _interp_extrap_1d(x_vals, y_vals, x_query):
    """Linear interpolate/extrapolate y(x) at x_query."""
    x_arr = np.asarray(x_vals, dtype=float)
    y_arr = np.asarray(y_vals, dtype=float)
    if x_arr.size == 0:
        return np.nan
    if x_arr.size == 1:
        return float(y_arr[0])
    order = np.argsort(x_arr)
    x_arr = x_arr[order]
    y_arr = y_arr[order]
    if x_query <= x_arr[0]:
        x0, x1 = x_arr[0], x_arr[1]
        y0, y1 = y_arr[0], y_arr[1]
    elif x_query >= x_arr[-1]:
        x0, x1 = x_arr[-2], x_arr[-1]
        y0, y1 = y_arr[-2], y_arr[-1]
    else:
        return float(np.interp(x_query, x_arr, y_arr))
    if x1 == x0:
        return float(y0)
    return float(y0 + (y1 - y0) * (x_query - x0) / (x1 - x0))


def _mass_density_g_cm3(params, temperature, w1):
    """Return mass density in g/cm^3 at given w1 and temperature."""
    try:
        wi_lst = [float(w1), float(1.0 - w1)]
        state = get_feos_npt_state(
            temperature=temperature,
            pressure=PRESSURE,
            params=params,
            wi_lst=wi_lst,
        )
        return float(state.mass_density() / MASS_DENSITY_CONVERSION)
    except Exception:  # pylint: disable=broad-exception-caught
        return np.nan


def _pure_water_density_g_cm3(params, temperature):
    """Return pure-water mass density in g/cm^3 at given temperature."""
    try:
        state = get_feos_npt_state(
            temperature=temperature,
            pressure=PRESSURE,
            params=params,
            xi_lst=[0.0, 1.0],
        )
        return float(state.mass_density() / MASS_DENSITY_CONVERSION)
    except Exception:  # pylint: disable=broad-exception-caught
        return np.nan


dict_exp_lst = load_exp_pnipaam_water_afroze_all()
dict_exp_lst = [resample_exp_data(exp_dict) for exp_dict in dict_exp_lst]


class PNIPAAmObjectiveWrapper(ObjectiveWrapperBase):
    """Objective wrapper for PNIPAAm-water parameter optimization."""

    # pylint: disable=too-many-locals,too-many-branches,too-many-statements
    # pylint: disable=broad-exception-caught
    def __init__(self, base_parameters, _keys_to_optimize):
        super().__init__(base_parameters, _keys_to_optimize)
        self.exp_density_records = load_exp_pnipaam_density_kany()
        self.exp_pure_water_temps, self.exp_pure_water_rhos = (
            _build_exp_pure_water_table(self.exp_density_records)
        )

    def _density_score(self, current_args: dict) -> float:
        if not self.exp_density_records or self.exp_pure_water_temps.size == 0:
            return 0.0

        temps_in_records = set()
        for record in self.exp_density_records:
            constant_vals = record.get("Constant Value", [])
            temperature = _extract_constant_value(constant_vals, "Temperature")
            if temperature is None:
                continue
            temps_in_records.add(float(temperature))

        model_pure_by_temp = {
            temp: _pure_water_density_g_cm3(current_args, temp)
            for temp in sorted(temps_in_records)
        }

        precomputed = []
        for record in self.exp_density_records:
            constant_vals = record.get("Constant Value", [])
            temperature = _extract_constant_value(constant_vals, "Temperature")
            mn = _extract_constant_value(constant_vals, "Mn")
            if temperature is None:
                continue

            data_rows = record.get("Data Table", [])
            if len(data_rows) < 2:
                continue

            w1_vals = []
            rho_exp_vals = []
            for row in data_rows[1:]:
                try:
                    w1_vals.append(float(row[0]))
                    rho_exp_vals.append(float(row[1]))
                except (TypeError, ValueError):
                    continue

            if not w1_vals:
                continue

            params_mn = (
                scale_component_molarweight(
                    dict(current_args), 1, float(mn), in_place=True
                )
                if mn is not None
                else dict(current_args)
            )

            rho_model_vals = [
                _mass_density_g_cm3(params_mn, float(temperature), w1)
                for w1 in w1_vals
            ]

            precomputed.append(
                (
                    float(temperature),
                    np.array(rho_exp_vals, dtype=float),
                    np.array(rho_model_vals, dtype=float),
                )
            )

        rel_errors = []
        for temperature, rho_exp_vals, rho_model_vals in precomputed:
            rho_model_pure = model_pure_by_temp.get(float(temperature), np.nan)
            rho_exp_pure = _interp_extrap_1d(
                self.exp_pure_water_temps,
                self.exp_pure_water_rhos,
                float(temperature),
            )
            if (
                not np.isfinite(rho_model_pure)
                or not np.isfinite(rho_exp_pure)
                or rho_exp_pure <= 0
            ):
                continue

            delta_rho = rho_model_pure - rho_exp_pure

            for rho_exp, rho_model in zip(rho_exp_vals, rho_model_vals):
                rho_shifted = rho_exp + delta_rho
                if not np.isfinite(rho_model) or rho_shifted <= 0:
                    rel_errors.append(1.0)
                    continue
                rel_error = abs(rho_model - rho_shifted) / rho_shifted
                rel_errors.append(rel_error)

        if not rel_errors:
            return 0.0
        return 100.0 * float(np.mean(rel_errors))

    def objective(self, current_args: dict, verbosity: bool = False) -> float:
        """Evaluate performance of PNIPAAm-water parameter set."""
        try:
            verbosity = False
            p_flt = PRESSURE
            upper_t_limit = T_HIGH
            lower_t_limit = T_LOW
            stage_1_penalty = 1e4

            density_score = self._density_score(current_args)

            if verbosity:
                print("Stage 1: Check Density Sanity Behaviour")

            t_lst = np.linspace(lower_t_limit, upper_t_limit, 8)
            params_by_mn = []
            lcst_mismatches = []
            for mn in MN_LIST:
                params_mn = scale_component_molarweight(
                    dict(current_args), 1, mn, in_place=True
                )
                for t_flt in t_lst:
                    if not is_density_iteration_sane(t_flt, p_flt, params_mn):
                        if verbosity:
                            print("Failed!")
                        return stage_1_penalty
                params_by_mn.append(params_mn)
                estimated_lcst = find_lcst_temperature(
                    T_LOW,
                    T_HIGH,
                    PRESSURE,
                    params_mn,
                    xtol=0.1,
                )
                lcst_mismatches.append(abs(estimated_lcst - T_LCST))

            if verbosity:
                print("Passed!")

            max_lcst_mismatch = max(lcst_mismatches)
            lcst_mismatch = float(np.mean(lcst_mismatches))

            if max_lcst_mismatch > 10:
                if verbosity:
                    print("LCST Failed!")
                    print(f"Max mismatch: {max_lcst_mismatch:.2f} K")
                    print(f"Aim.: {T_LCST:.1f}")
                return 1000 * max_lcst_mismatch

            score_binodal = 0.0
            for params_mn in params_by_mn:
                y_init_binary = build_initial_guess_binary(
                    T_HIGH, PRESSURE, params_mn
                )

                y_lst, _, _, _, _, _ = trace_parameter(
                    trace_index=-1,
                    y_init=y_init_binary,
                    trace_start=T_HIGH,
                    trace_target=T_LOW,
                    pressure=PRESSURE,
                    params=params_mn,
                    step_size_init=0.001,
                    newton_tol=1e-5,
                    critical_point_tol=1e-3,
                    max_step_fraction=1,
                )
                binodal_dict = interpret_y_vector(y_lst, PRESSURE, params_mn)

                # close critical point for the evaluation
                binodal_dict["w_I_lst"].append(binodal_dict["w_II_lst"][-1])
                binodal_dict["x_I_lst"].append(binodal_dict["x_II_lst"][-1])
                binodal_dict["w_II_lst"].append(binodal_dict["w_II_lst"][-1])
                binodal_dict["x_II_lst"].append(binodal_dict["x_II_lst"][-1])
                binodal_dict["temperature_K"].append(
                    binodal_dict["temperature_K"][-1]
                )

                score_binodal += get_minimal_distance(
                    binodal_dict,
                    dict_exp_lst,
                    permutation_list=[
                        [[0, 1]],
                        [[0, 1] for _ in range(len(dict_exp_lst))],
                    ],
                    binodal_sampling=8,
                    which_idx_binary=0,
                    t_scale=T_SCALE,
                )[0]

            score_binodal /= len(MN_LIST)

            if verbosity:
                print("Binodal Score")
                print(score_binodal)

            if verbosity:
                print("Density Score")
                print(density_score)

            return score_binodal + lcst_mismatch * 100 + density_score * 50

        except Exception as exc:
            print("Error:", exc)
            return float("inf")


keys_to_optimize = [
    "comp12_kij_intercept",
    "comp12_kij_slope",
    "comp1_epsilon_k_ab",
    # "comp12_epsilon_k_ab",
    "comp1_kappa_ab",
    "comp1_epsilon_k",
    #    "comp1_m",
    "comp1_sigma",
]

max_change_of_parameters = [
    0.4,  # comp12_kij_intercept
    0.1,  # comp12_kij_slope
    690,  # comp1_epsilon_k_ab
    # 2000,  # comp12_epsilon_k_ab
    0.08,  # comp1_kappa_ab
    600,  # comp1_epsilon
    #   30,  # comp1_m
    4,  # comp1_sigma
]


WATER_MODELS = ["arndt", "grob", "esper", "bahrabadi"]

if __name__ == "__main__":

    for idx, water_model in enumerate(WATER_MODELS):

        print("##################")
        print(f"Starting transition No. {idx}, {water_model}")
        print("##################")
        current_parameter_set = get_parameter_set_from_tables(
            ["pnipaam", "water"],
            ["arndt", water_model],
            binary_sources={
                (1, 2): "arndt",
            },
        )

        current_parameter_set["comp1_epsilon_k_ab"] = 700

        initial_guess = np.array(
            [current_parameter_set[key] for key in keys_to_optimize]
        )

        objective_wrapper = PNIPAAmObjectiveWrapper(
            current_parameter_set, keys_to_optimize
        )
        optimizer = ParameterOptimization(objective_wrapper)

        path = (
            f"optimization_water_models/"
            f"pnipaam_water_{water_model}_const_m_T_scale_{T_SCALE}"
        )

        res = optimizer.run_optimization(
            initial_guess=initial_guess,
            max_change_of_parameters=max_change_of_parameters,
            verbosity=True,
            de_kwargs={
                "popsize": 8,
                "maxiter": 2000,
                "prefix_file": path,
            },
            population_kwargs={
                "popsize": 8,
                "sigma": 0.05,
                "sigma_random": 0.02,
            },
        )
        # Objective file 1:
        print(objective_wrapper.objective(current_parameter_set))

        best_params = current_parameter_set | {
            k: v for k, v in zip(keys_to_optimize, res["best_unscaled"][0])
        }
        print(objective_wrapper.objective(best_params))

        def _jsonable(x):
            if isinstance(x, (np.floating, np.integer)):
                return x.item()
            if isinstance(x, np.ndarray):
                return x.tolist()
            return x

        keys_tag = "_".join(keys_to_optimize)

        outpath = Path(f"{path}.json")
        outpath.write_text(
            json.dumps(best_params, indent=2, default=_jsonable),
            encoding="utf-8",
        )
        print(f"Saved: {outpath.resolve()}")
