"""Read, write, and resample experimental or calculated JSON data."""

from __future__ import annotations

import json
import re
from datetime import datetime
from pathlib import Path
from scipy.interpolate import interp1d
import numpy as np


PROJECT_ROOT = Path(__file__).resolve().parents[3]
BASE_FOLDER_EXP = PROJECT_ROOT / "03_ExperimentalData"
BASE_FOLDER_CALC = PROJECT_ROOT / "02_CalculationData"
BASE_FOLDER_EXP.mkdir(parents=True, exist_ok=True)
BASE_FOLDER_CALC.mkdir(parents=True, exist_ok=True)

TYPE_FLAGS = {
    "tie_line": "EXP_LLE",
    "solubilities": "EXP_SOL",
    "calculation": "CAL_LLE",
    "calculation_pure": "CAL_PUR",
}


def _normalize_array(obj):
    if isinstance(obj, np.ndarray):
        return obj.tolist()
    if isinstance(obj, np.generic):
        return obj.item()
    if isinstance(obj, list):
        return [_normalize_array(v) for v in obj]
    if isinstance(obj, dict):
        return {k: _normalize_array(v) for k, v in obj.items()}
    return obj


def _clean_name(value: str) -> str:
    return value.replace(" ", "")


def _first_token(value: str) -> str:
    return re.split(r"[ ,.-]", value)[0]


# pylint: disable=too-many-locals


def build_lle_filename(data: dict, ext: str) -> str:
    """Build a descriptive filename for an LLE record."""
    comp1 = "_" + _clean_name(data["comp1_name"])

    try:
        type_flag = TYPE_FLAGS[data["type"]]
    except KeyError as exc:
        raise ValueError(f"Unknown LLE record type: {data['type']}") from exc

    comp2 = ""
    if name := data.get("comp2_name"):
        comp2 = "_" + _clean_name(name)

    ternary_flag = ""
    quaternary_flag = ""
    temperature_flag = ""

    if comp3 := data.get("comp3_name"):
        ternary_flag = "_" + _clean_name(comp3)

        temperature = data["temperature_K"]
        if isinstance(temperature, list):
            temperature = temperature[0]
        temperature_flag = f"_T{temperature:.2f}K"

    if comp4 := data.get("comp4_name"):
        quaternary_flag = "_" + _clean_name(comp4)

    if "EXP" in type_flag:
        terminal = _first_token(data["source"])
    else:
        timestamp = datetime.now().strftime("%Y_%m_%d_%H_%M")
        terminal = f"{data['comment']}_{timestamp}"

    flags = (
        type_flag,
        comp1,
        comp2,
        ternary_flag,
        quaternary_flag,
        temperature_flag,
    )

    return f"{''.join(flags)}_{terminal}.{ext}"


def save_lle_records(records):
    """Save a list of LLE records to JSON files."""
    for record in records:
        save_lle_record(record)


def save_lle_record(record):
    """Save a single LLE record to JSON."""
    if record["type"] in {"solubilities", "tie_line"}:
        filename = build_lle_filename(record, "json")
        filepath = BASE_FOLDER_EXP / filename
    elif "calculation" in record["type"]:
        filename = build_lle_filename(record, "json")
        filepath = BASE_FOLDER_CALC / filename
    else:
        raise ValueError(f"Unknown LLE record type: {record['type']}")

    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(_normalize_array(record), f, indent=2)


def load_lle_record(path):
    """Load a LLE record from JSON.

    Resolves relative paths against the project root and data folders.
    """
    path_obj = Path(path)
    if not path_obj.is_absolute():
        candidates = [
            path_obj,
            PROJECT_ROOT / path_obj,
            BASE_FOLDER_CALC / path_obj.name,
            BASE_FOLDER_EXP / path_obj.name,
        ]
        for candidate in candidates:
            if candidate.exists():
                path_obj = candidate
                break
        else:
            tried = ", ".join(str(c) for c in candidates)
            raise FileNotFoundError(f"File not found: {path}. Tried: {tried}")

    with open(path_obj, "r", encoding="utf-8") as f:
        return json.load(f)


# pylint: disable=too-many-locals


def find_lle_paths_by_components(components: list[str], is_exp: bool = True):
    """Return paths for LLE records matching the given components."""
    components_sorted = sorted(components)
    folder = BASE_FOLDER_EXP if is_exp else BASE_FOLDER_CALC
    paths = []

    for path in sorted(folder.glob("*.json")):
        record = load_lle_record(path)
        record_type = record.get("type")
        if is_exp:
            if record_type not in {"tie_line", "solubilities"}:
                continue
        else:
            if record_type != "calculation":
                continue

        record_components = []
        for idx in range(1, len(components) + 1):
            name = record.get(f"comp{idx}_name")
            if name:
                record_components.append(name)
        if sorted(record_components) != components_sorted:
            continue
        paths.append(str(path))

    return paths


# pylint: disable=too-many-locals


def load_record_groups(
    path_calc_lst=None, path_exp_lst=None, component_list=None
):
    """Load calculation/experimental records and verify component order."""
    dict_calc_lst = (
        [load_lle_record(p) for p in path_calc_lst] if path_calc_lst else []
    )
    dict_exp_lst = (
        [load_lle_record(p) for p in path_exp_lst] if path_exp_lst else []
    )

    # check for any dict
    source = next(
        (lst for lst in (dict_calc_lst, dict_exp_lst) if lst),
        None,
    )
    if source is None:
        raise ValueError("No calculation or experimental records provided")

    # given order of the components
    if component_list:
        number_of_components = len(component_list)

    # order of components based on the first calc dict
    else:
        number_of_components = dict_calc_lst[0]["params"]["number_components"]
        component_list = [
            source[0][f"comp{i+1}_name"] for i in range(number_of_components)
        ]
    permutation_list = [[], []]

    for dict_lst_idx, dict_lst in enumerate([dict_calc_lst, dict_exp_lst]):
        for dict_idx, current_dict in enumerate(dict_lst):
            current_comp_list = [
                current_dict[f"comp{i+1}_name"]
                for i in range(number_of_components)
            ]

            if set(current_comp_list) != set(component_list):
                raise RuntimeError(
                    "File "
                    f"{path_calc_lst[dict_idx]} using different components "
                    f"{current_comp_list} instead of {component_list}"
                )

            permutation_list[dict_lst_idx].append(
                [current_comp_list.index(comp) for comp in component_list]
            )

    return dict_calc_lst, dict_exp_lst, permutation_list


def load_exp_pnipaam_water_all(glob_pattern):
    """Load all PNIPAAm-water experimental solubilities."""
    paths = sorted(BASE_FOLDER_EXP.glob(glob_pattern))
    if not paths:
        raise FileNotFoundError(
            f"Missing experimental files for pattern: {glob_pattern}"
        )
    return [load_lle_record(path) for path in paths]


def load_exp_pnipaam_density_kany(path):
    """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 resample_exp_data(exp_dict, temp_weight=20, points=25):
    """Resample experimental data to constant arc length"""

    if exp_dict.get("type") != "solubilities":
        raise RuntimeError(
            "Expected exp_dict['type'] == 'solubilities'"
            + f", got {exp_dict.get('type')!r}"
        )

    # composition (scaled)
    w = np.array([wi[0] for wi in exp_dict["w_I_lst"]]) * 100
    t = np.array(exp_dict["temperature_K"])

    # sort by composition
    idx = np.argsort(w)
    w = w[idx]
    t = t[idx]

    # arc length

    dw = np.diff(w)
    dt = np.diff(t) * temp_weight  # scaling: 1K equals temp_weight mass %
    ds = np.sqrt(dw**2 + dt**2)

    s = np.concatenate([[0], np.cumsum(ds)])
    s /= s[-1]

    # resample
    s_new = np.linspace(0, 1, points)

    fw = interp1d(s, w, kind="linear")
    ft = interp1d(s, t, kind="linear")

    w_new = fw(s_new)
    t_new = ft(s_new)

    return exp_dict | {
        "type": "solubilities",
        "x_I_lst": [],
        "w_I_lst": [[wi / 100, 1.0 - wi / 100] for wi in w_new],
        "temperature_K": t_new,
    }


def load_dos_record(path):
    """Load a DOS record from JSON and add w_polymer = 1/DOS.
    If the system is >binary, also extract solvent compositions.
    """
    path_obj = Path(path)
    if not path_obj.is_absolute():
        candidates = [
            path_obj,
            PROJECT_ROOT / path_obj,
            BASE_FOLDER_CALC / path_obj.name,
            BASE_FOLDER_EXP / path_obj.name,
        ]
        for candidate in candidates:
            if candidate.exists():
                path_obj = candidate
                break
        else:
            tried = ", ".join(str(c) for c in candidates)
            raise FileNotFoundError(f"File not found: {path}. Tried: {tried}")

    with open(path_obj, "r", encoding="utf-8") as f:
        rec = json.load(f)

    dos_vals = rec.get("DOS_g_per_g") or rec.get("DOS") or rec.get("dos")
    if dos_vals is None:
        raise KeyError(
            "No DOS array found. Expected one of: 'DOS_g_per_g', 'DOS', 'dos'."
        )

    w_polymer = []
    for v in dos_vals:
        if v is None:
            w_polymer.append(None)
            continue
        dv = float(v)
        w_polymer.append(None if dv == 0.0 else 1.0 / dv)
    rec["w_polymer"] = w_polymer

    # Determine number of components
    n_components = 0
    while f"component_{n_components + 1}" in rec:
        n_components += 1
    rec["n_components"] = n_components

    # If more than binary, collect compositions (best-effort across schemas)
    if n_components > 2:
        compositions = {}

        # Common patterns used:
        for i in range(2, n_components + 1):
            keys_to_try = [
                f"w_{i}_mol_per_mol",
                f"w_{i}_g_per_g",
                f"x{i}_mol_per_mol",
                f"x{i}",
            ]
            val = None
            for k in keys_to_try:
                if k in rec:
                    val = rec[k]
                    compositions[k] = val
                    break

        # Also accept generic "w_3" etc if someone simplified keys
        for i in range(2, n_components + 1):
            k = f"w_{i}"
            if k in rec and k not in compositions:
                compositions[k] = rec[k]

        if not compositions:
            # No compositions found
            rec["compositions"] = {}
            rec["compositions_found"] = False
        else:
            rec["compositions"] = compositions
            rec["compositions_found"] = True

    return rec
