"""Plot pure-component calculations against experimental JSON data."""

from collections import defaultdict
import re

import matplotlib.pyplot as plt
import numpy as np

from equi_gel.save_load_json import (
    BASE_FOLDER_CALC,
    BASE_FOLDER_EXP,
    load_lle_record,
)
from equi_gel.plot_ternary import STYLE_PATH

DEFAULT_COMPONENTS = ["water", "n_butanol", "acetic_acid", "butyl_acetate"]
T_LIMITS = (290, 370)


def discover_pure_component_files(component_list=None):
    """Find calculation/experimental JSON files for pure components."""
    calc_paths = defaultdict(list)
    exp_paths = defaultdict(list)

    component_pattern = "|".join(
        map(re.escape, component_list or DEFAULT_COMPONENTS)
    )

    calc_pattern = re.compile(
        rf"^CAL_PUR_({component_pattern})_.*_2026_01_27_.*\.json$",
        re.IGNORECASE,
    )
    exp_pattern = re.compile(
        rf"^EXP_PUR_({component_pattern})_.*\.json$",
        re.IGNORECASE,
    )

    for path in sorted(BASE_FOLDER_CALC.glob("CAL_PUR_*.json")):
        match = calc_pattern.match(path.name)
        if match:
            calc_paths[match.group(1)].append(path)

    for path in sorted(BASE_FOLDER_EXP.glob("EXP_PUR_*.json")):
        match = exp_pattern.match(path.name)
        if match:
            exp_paths[match.group(1)].append(path)

    if component_list is None:
        component_list = sorted(set(calc_paths) | set(exp_paths))

    return component_list, calc_paths, exp_paths


def load_records(path_list):
    """Load a list of LLE records from paths."""
    return [load_lle_record(path) for path in path_list]


def component_title(component_label):
    """Human-friendly title for component names."""
    return component_label.replace("_", " ").title()


def short_component_name(component_label):
    """Compact label for LaTeX table rows."""
    mapping = {
        "water": "W",
        "n_butanol": "B",
        "acetic_acid": "AA",
        "butyl_acetate": "BA",
    }
    return mapping.get(component_label, component_label)


def calc_name_from_comment(calc_record, default_idx):
    """Compact model/source name from calc comment."""
    comment = calc_record.get("comment", f"calc {default_idx + 1}")

    name = re.split(r"[ ,.-]", comment)[0]

    if "bahrabadi" in name:
        name = "jovein"

    return name.capitalize()


def format_error(err):
    """Format error for LaTeX table."""
    if err is None:
        return ""
    return f"${err:.1f}\\%$"


def first_token(value):
    """Return the leading token of a source string for compact legends."""
    return re.split(r"[ ,.-]", value)[0]


def build_error_table(component_list, calc_paths, exp_paths):
    """Build nested error table data."""
    table_data = {}
    _calc_names = []

    for _component_id in component_list:
        calc_records = load_records(calc_paths.get(_component_id, []))
        exp_records = load_records(exp_paths.get(_component_id, []))

        if not calc_records or not exp_records:
            continue

        component_errors = {
            "rho": {},
            "p_lv": {},
        }

        for calc_idx, calc_record in enumerate(calc_records):
            calc_name = calc_name_from_comment(calc_record, calc_idx)
            if calc_name not in _calc_names:
                _calc_names.append(calc_name)

            rho_errors = []
            p_errors = []

            for exp_record in exp_records:
                rho_err = mean_relative_error(
                    exp_record["temperature_K"],
                    exp_record["liquid_density_kmol_per_cbm"],
                    calc_record,
                    "liquid_density_kmol_per_cbm",
                )
                if rho_err is not None:
                    rho_errors.append(rho_err)

                p_err = mean_relative_error(
                    exp_record["temperature_K"],
                    exp_record["vapor_pressure_pa"],
                    calc_record,
                    "vapor_pressure_pa",
                    scale=1 / 1000.0,
                )
                if p_err is not None:
                    p_errors.append(p_err)

            component_errors["rho"][calc_name] = (
                float(np.mean(rho_errors)) if rho_errors else None
            )
            component_errors["p_lv"][calc_name] = (
                float(np.mean(p_errors)) if p_errors else None
            )

        table_data[_component_id] = component_errors

    return _calc_names, table_data


def write_latex_error_table(filepath, component_list, _calc_names, table_data):
    """Write LaTeX tabular with mean relative errors to a text file."""
    colspec = "ll" + "c" * len(_calc_names)
    lines = [
        rf"\begin{{tabular}}{{{colspec}}}",
        r"\hline",
        "Component & Property "
        + "".join(f"\n& {name} " for name in _calc_names)
        + r"\\",
        r"\hline",
    ]

    for _component_id in component_list:
        if _component_id not in table_data:
            continue

        comp_short = short_component_name(_component_id)
        rho_row = table_data[_component_id]["rho"]
        p_row = table_data[_component_id]["p_lv"]

        rho_cells = " & ".join(
            format_error(rho_row.get(name)) for name in _calc_names
        )
        p_cells = " & ".join(
            format_error(p_row.get(name)) for name in _calc_names
        )

        lines.append(
            rf"{comp_short}  & $\rho_{{\text{{liq}}}}$ & {rho_cells} \\"
        )
        lines.append(rf"   & $P_{{\text{{LV}}}}$     & {p_cells} \\")
        lines.append(r"\hline")

    lines.append(r"\end{tabular}")

    with open(filepath, "w", encoding="utf-8") as handle:
        handle.write("\n".join(lines))


def mean_relative_error(exp_t, exp_y, calc_record, key, *, scale=1.0):
    """Return mean relative error (%) for calc data vs experiments."""
    exp_t = np.asarray(exp_t, dtype=float)
    exp_y = np.asarray(exp_y, dtype=float) * scale

    calc_t = np.asarray(calc_record["temperature_K"], dtype=float)
    calc_y = np.asarray(calc_record[key], dtype=float) * scale

    min_t = np.min(calc_t)
    max_t = np.max(calc_t)
    mask = (exp_t >= min_t) & (exp_t <= max_t)
    if not np.any(mask):
        return None

    interp_vals = np.interp(exp_t[mask], calc_t, calc_y)
    exp_vals = exp_y[mask]
    denom = np.where(exp_vals != 0, np.abs(exp_vals), np.nan)
    rel = np.abs(interp_vals - exp_vals) / denom
    rel = rel[np.isfinite(rel)]
    if rel.size == 0:
        return None
    return 100.0 * float(np.mean(rel))


# pylint: disable=too-many-locals


def plot_pure_component(
    component_label, calc_records, exp_records, *, t_limits
):
    """Plot density and vapor pressure for a pure component."""
    color_cycler = plt.rcParams["axes.prop_cycle"]
    colors = color_cycler.by_key()["color"]
    exp_markers = ["o", "s", "^", "D", "v"]
    component_title_str = component_title(component_label)

    # Liquid density
    fig, ax = plt.subplots()
    for calc_idx, calc_record in enumerate(calc_records):
        label = calc_record.get("comment", f"calc {calc_idx + 1}")
        errors = []
        for exp_record in exp_records:
            err = mean_relative_error(
                exp_record["temperature_K"],
                exp_record["liquid_density_kmol_per_cbm"],
                calc_record,
                "liquid_density_kmol_per_cbm",
            )
            if err is None:
                continue
            errors.append(err)
        if errors:
            label = f"{label} err={float(np.mean(errors)):.1f}%"
        ax.plot(
            calc_record["temperature_K"],
            calc_record["liquid_density_kmol_per_cbm"],
            color=colors[calc_idx % len(colors)],
            linestyle="-.",
            linewidth=1.5,
            label=label,
        )
    for exp_idx, exp_record in enumerate(exp_records):
        source = first_token(exp_record.get("source", "exp"))
        ax.plot(
            exp_record["temperature_K"],
            exp_record["liquid_density_kmol_per_cbm"],
            color="black",
            linestyle=":",
            marker=exp_markers[exp_idx % len(exp_markers)],
            markersize=4,
            label=source,
        )

    ax.set_title(f"{component_title_str} - Liquid density")
    ax.set_xlim(t_limits)
    ax.set_xlabel(r"Temperature [K]")
    ax.set_ylabel(r"$\rho_\mathrm{liq} [\mathrm{kmol/cbm}]$")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="lower left", frameon=False)
    fig.tight_layout()

    # Vapor pressure
    fig, ax = plt.subplots()
    for calc_idx, calc_record in enumerate(calc_records):
        label = calc_record.get("comment", f"calc {calc_idx + 1}")
        errors = []
        for exp_record in exp_records:
            err = mean_relative_error(
                exp_record["temperature_K"],
                exp_record["vapor_pressure_pa"],
                calc_record,
                "vapor_pressure_pa",
                scale=1 / 1000.0,
            )
            if err is None:
                continue
            errors.append(err)
        if errors:
            label = f"{label} err={float(np.mean(errors)):.1f}%"
        ax.plot(
            calc_record["temperature_K"],
            [v / 1000 for v in calc_record["vapor_pressure_pa"]],
            color=colors[calc_idx % len(colors)],
            linestyle="-",
            linewidth=1.5,
            label=label,
        )
    for exp_idx, exp_record in enumerate(exp_records):
        source = first_token(exp_record.get("source", "exp"))
        ax.plot(
            exp_record["temperature_K"],
            [v / 1000 for v in exp_record["vapor_pressure_pa"]],
            color="black",
            linestyle=":",
            marker=exp_markers[exp_idx % len(exp_markers)],
            markersize=4,
            label=source,
        )

    ax.set_yscale("log")
    ax.set_title(f"{component_title_str} - Vapor pressure")
    ax.set_xlim(t_limits)
    ax.set_xlabel(r"Temperature [K]")
    ax.set_ylabel(r"$p_\mathrm{LV} [\mathrm{kPa}]$")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="upper left", frameon=False)
    fig.tight_layout()


plt.style.use(STYLE_PATH)

pure_component_list, pure_calc_paths, pure_exp_paths = (
    discover_pure_component_files(DEFAULT_COMPONENTS)
)

for component_id in pure_component_list:
    pure_calc_records = load_records(pure_calc_paths.get(component_id, []))
    pure_exp_records = load_records(pure_exp_paths.get(component_id, []))

    if not pure_calc_records or not pure_exp_records:
        print(
            f"Skipping {component_id}: "
            f"calc={len(pure_calc_records)} exp={len(pure_exp_records)}"
        )
        continue

    plot_pure_component(
        component_id,
        pure_calc_records,
        pure_exp_records,
        t_limits=T_LIMITS,
    )

calc_names, error_table_data = build_error_table(
    pure_component_list,
    pure_calc_paths,
    pure_exp_paths,
)

write_latex_error_table(
    "tables/pure_component_errors_table.txt",
    pure_component_list,
    calc_names,
    error_table_data,
)
