"""Load, merge, and scale pure-component and binary parameter sets."""

import re
import copy
import itertools
from datetime import datetime
from pathlib import Path
import numpy as np
import pandas as pd


BASE_DIR = Path(__file__).resolve().parents[3]
PATH_LIT_PURE = BASE_DIR / "01_Parameters" / "LIT_PURE.csv"
PATH_LIT_BINARY = BASE_DIR / "01_Parameters" / "LIT_BINARY.csv"
PATH_CAL_FULLSETS = BASE_DIR / "01_Parameters" / "CAL_FULLSETS.csv"


DF_LIT_PURE = pd.read_csv(PATH_LIT_PURE)
DF_LIT_BINARY = pd.read_csv(PATH_LIT_BINARY)


def _nan_to_none(value):
    return None if pd.isna(value) else value


def _none_to_nan(value):
    return pd.NA if value is None else value


def get_component_params(component, source, idx):
    """Return the component parameter row as a dict.

    Raises KeyError if the component/source combination is missing.
    """
    row = DF_LIT_PURE[
        (DF_LIT_PURE["comp$_name"] == component)
        & (DF_LIT_PURE["source"] == source)
    ]

    if row.empty:
        raise KeyError(f"Component/source not found: {component}, {source}")

    return {
        k.replace("$", str(idx)): v
        for k, v in row.iloc[0].to_dict().items()
        if "$" in k
    }


def get_binary_params(comp_i, comp_j, source, idx_i, idx_j):
    """Return the binary parameter row as a dict.

    Raises KeyError if the binary/source combination is missing.
    """

    if source == "None":
        return {}

    comp1, comp2 = sorted((comp_i, comp_j))

    row = DF_LIT_BINARY[
        (DF_LIT_BINARY["comp1"] == comp1)
        & (DF_LIT_BINARY["comp2"] == comp2)
        & (DF_LIT_BINARY["source"] == source)
    ]

    if row.empty:
        raise KeyError(
            f"Binary/source not found: {comp_i}, {comp_j}, {source}"
        )

    return {
        k.replace("$", str(idx_i)).replace("§", str(idx_j)): v
        for k, v in row.iloc[0].to_dict().items()
        if "$" in k or "§" in k
    }


def add_mixing_association(param_dict: dict):
    """
    For all components in param_dict, set missing binary cross
    association parameters using the standard mixing rule (k_ij = 0).

    Notation of mixing based on:
    Rehner et al. 2023, https://doi.org/10.1007/s10765-023-03290-3
    """

    n = param_dict["number_components"]

    if n <= 1:
        return param_dict

    for i, j in itertools.combinations(range(1, n + 1), 2):
        eps_key = f"comp{i}{j}_epsilon_k_ab"
        kap_key = f"comp{i}{j}_kappa_ab"

        eps_val = param_dict.get(eps_key)
        if eps_val is None or (
            isinstance(eps_val, float) and np.isnan(eps_val)
        ):
            param_dict[eps_key] = 0.5 * (
                param_dict[f"comp{i}_epsilon_k_ab"]
                + param_dict[f"comp{j}_epsilon_k_ab"]
            )

        kap_val = param_dict.get(kap_key)
        if kap_val is None or (
            isinstance(kap_val, float) and np.isnan(kap_val)
        ):
            param_dict[kap_key] = np.sqrt(
                param_dict[f"comp{i}_kappa_ab"]
                * param_dict[f"comp{j}_kappa_ab"]
            )

    return param_dict


# used for the reindexing
_COMP_RE = re.compile(r"^comp(?P<digits>[0-9]+)_")


def get_subset_of_parameters(param_dict: dict, idx_list):
    """
    creating a subset of the given parameter set with reindexing
    """
    new_param_dict = copy.deepcopy(param_dict)

    keep = set(str(i) for i in idx_list)
    # map old digit -> new digit (1-based)
    old_to_new = {str(old): str(new + 1) for new, old in enumerate(idx_list)}

    keys_to_delete = []
    rename_map = {}

    for key in list(new_param_dict.keys()):
        m = _COMP_RE.match(key)
        if not m:
            continue

        digits = m.group("digits")

        # delete if any digit not kept
        if any(d not in keep for d in digits):
            keys_to_delete.append(key)
            continue

        # remap each digit
        new_digits = "".join(old_to_new[d] for d in digits)
        new_key = _COMP_RE.sub(f"comp{new_digits}_", key, 1)

        if new_key != key:
            rename_map[key] = new_key

    for k in keys_to_delete:
        new_param_dict.pop(k, None)

    for old_key, new_key in rename_map.items():
        new_param_dict[new_key] = new_param_dict.pop(old_key)

    new_param_dict["number_components"] = len(idx_list)
    return new_param_dict


def get_parameter_set_from_tables(
    name_components,
    source,
    binary_sources=None,
):
    """
    Build parameter dict from literature tables.

    Supports either:

    • Single source for all components + binaries
    • Multiple sources per component
      → then `binary_sources` must be provided per pair

    Parameters
    ----------
    name_components : list[str]
        Component names.

    source : str | list[str]
        Literature source(s) for pure components.
        If list, length must match components.

    binary_sources : dict, optional
        Required if multiple component sources are used.
        Keys can be (i, j) or (name_i, name_j).

    Returns
    -------
    dict
        Parameter dictionary.

    Examples
    --------
    Single source:
        get_parameter_set_from_tables(
            ["A", "B"], "Smith2020"
        )

    Multiple sources:
        get_parameter_set_from_tables(
            ["A", "B"],
            ["Smith2020", "Lee2019"],
            binary_sources={
                (1, 2): "MixPaper"
            },
        )
    """

    n_components = len(name_components)

    params = {
        "number_components": n_components,
        "gel": False,
        "elasticity": "",
        "C": 0,
        "ref_volume_frac": 0.0,
        "max_volume_poly_frac": 0.0,
    }

    # normalize sources
    if isinstance(source, (list, tuple)):
        if len(source) != n_components:
            raise ValueError("Length of 'source' must match components.")
        comp_sources = list(source)
        multi_source = True
    else:
        comp_sources = [source] * n_components
        multi_source = False

    # pure components
    for i, (comp, src) in enumerate(
        zip(name_components, comp_sources), start=1
    ):
        params |= get_component_params(comp, src, i)

    # binaries
    if n_components > 1:
        if multi_source and binary_sources is None:
            raise ValueError(
                "Provide 'binary_sources' when using multiple sources."
            )

        for (i, comp_i), (j, comp_j) in itertools.combinations(
            enumerate(name_components, start=1), 2
        ):
            if multi_source:
                bsrc = binary_sources.get((i, j)) or binary_sources.get(
                    (comp_i, comp_j)
                )
                if bsrc is None:
                    params |= {
                        "comp$§_kij_intercept".replace("$", str(i)).replace(
                            "§", str(j)
                        ): 0.0,
                        "comp$§_kij_slope".replace("$", str(i)).replace(
                            "§", str(j)
                        ): 0.0,
                    }
                    continue
            else:
                bsrc = source

            params |= get_binary_params(comp_i, comp_j, bsrc, i, j)

    return {k: _nan_to_none(v) for k, v in params.items()}


def save_parameter_set(
    param_dict: dict,
    key: str,
    comment: str | None = None,
):
    """Save a parameter set into the calibrated fullsets table."""
    if not key or not isinstance(key, str):
        raise ValueError("Key must be a non-empty string")

    row = {
        "key": key,
        "comment": comment,
        "timestamp": datetime.utcnow().isoformat(),
    }

    row |= {k: _none_to_nan(v) for k, v in param_dict.items()}
    df_new = pd.DataFrame([row])

    if PATH_CAL_FULLSETS.exists():
        df = pd.read_csv(PATH_CAL_FULLSETS)

        if "key" in df.columns and key in df["key"].values:
            raise ValueError(f"Key '{key}' already exists")

        df = pd.concat([df, df_new], ignore_index=True)
    else:
        df = df_new

    df.to_csv(PATH_CAL_FULLSETS, index=False)


def load_parameter_set(
    key: str,
) -> dict:
    """Load a parameter set by key from the calibrated fullsets table."""
    if not PATH_CAL_FULLSETS.exists():
        raise FileNotFoundError(PATH_CAL_FULLSETS)

    df = pd.read_csv(PATH_CAL_FULLSETS)

    row = df[df["key"] == key]
    if row.empty:
        raise KeyError(f"No parameter set with key '{key}'")

    record = row.iloc[0].to_dict()

    # strip metadata
    record.pop("key", None)
    record.pop("comment", None)
    record.pop("timestamp", None)

    # NaN -> None
    return {k: _nan_to_none(v) for k, v in record.items()}


def scale_component_molarweight(
    params: dict,
    comp_idx: int,
    new_molarweight: float,
    *,
    in_place: bool = True,
) -> dict:
    """Scale component m proportional to the new molar weight."""
    mw_key = f"comp{comp_idx}_molarweight"
    m_key = f"comp{comp_idx}_m"
    na_key = f"comp{comp_idx}_na"
    nb_key = f"comp{comp_idx}_nb"

    if mw_key not in params or m_key not in params:
        raise KeyError(f"Missing keys: {mw_key} and/or {m_key}")

    current_mw = float(params[mw_key])
    current_m = float(params[m_key])
    if current_mw <= 0:
        raise ValueError(f"Invalid molar weight {current_mw} for {mw_key}")
    if new_molarweight <= 0:
        raise ValueError(
            f"Invalid new molar weight {new_molarweight} for {mw_key}"
        )

    scale = new_molarweight / current_mw
    updated = params if in_place else dict(params)
    updated[mw_key] = new_molarweight
    updated[m_key] = current_m * scale
    updated[na_key] = float(round(updated[m_key]))
    updated[nb_key] = float(round(updated[m_key]))
    return updated
