from __future__ import annotations

import csv
import json
import math
import os
import shutil
from pathlib import Path

import numpy as np


ROOT = Path(r"D:\LDA")
PIC = ROOT / "PIC"
ORGANIZED = PIC / "Organized"

LATENT_TUNING_CSV = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "tuning"
    / "method_specific_tuning_results.csv"
)
PHY_TUNING_CSV = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "figures"
    / "test_oracle_phy_tuning_results.csv"
)
NORM_SOURCE_JSON = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "results"
    / "N20_current_ae_loc_method_tuning_tuning"
    / "test_SLA_256_inf1p10_rad192p00_summary.json"
)

FIG_PATH = PIC / "Fig10_N20_method_specific_tuning_heatmaps.png"
FIG_PDF = PIC / "Fig10_N20_method_specific_tuning_heatmaps.pdf"
SUMMARY_CSV = PIC / "Fig10_N20_method_specific_tuning_heatmaps_values.csv"

METHODS = [
    ("GLA_256", "GLA_256 - PHY"),
    ("SLA_256", "SLA_256 - PHY"),
    ("SLA_128", "SLA_128 - PHY"),
]


def read_rows(path: Path) -> list[dict[str, str]]:
    with path.open("r", newline="", encoding="utf-8") as f:
        return list(csv.DictReader(f))


def as_float(text: str) -> float:
    try:
        return float(text)
    except (TypeError, ValueError):
        return math.nan


def rmse_std() -> float:
    with NORM_SOURCE_JSON.open("r", encoding="utf-8") as f:
        data = json.load(f)
    return float(data["normalization_std"])


def build_surface(
    rows: list[dict[str, str]],
    method_label: str | None,
    inflations: list[float],
    radii: list[float],
) -> tuple[np.ndarray, np.ndarray]:
    values = np.full((len(inflations), len(radii)), np.nan, dtype=float)
    unstable = np.zeros((len(inflations), len(radii)), dtype=bool)
    infl_index = {round(v, 8): i for i, v in enumerate(inflations)}
    rad_index = {round(v, 8): i for i, v in enumerate(radii)}

    for row in rows:
        if method_label is not None and row.get("method_label") != method_label:
            continue
        inf = as_float(row.get("inflation", ""))
        rad = as_float(row.get("loc_radius", ""))
        if not math.isfinite(inf) or not math.isfinite(rad):
            continue
        i = infl_index.get(round(inf, 8))
        j = rad_index.get(round(rad, 8))
        if i is None or j is None:
            continue
        values[i, j] = as_float(row.get("mean_analysis_nrmse", ""))
        unstable[i, j] = row.get("stable", "").strip().lower() != "true"
    return values, unstable


def main() -> None:
    mpl_config = ROOT / ".mplconfig_fig34_tuning_heatmaps"
    mpl_config.mkdir(exist_ok=True)
    os.environ["MPLCONFIGDIR"] = str(mpl_config)

    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.colors import TwoSlopeNorm

    latent_rows = read_rows(LATENT_TUNING_CSV)
    phy_rows = read_rows(PHY_TUNING_CSV)

    inflations = sorted({as_float(r["inflation"]) for r in latent_rows if r.get("inflation")})
    radii = sorted({as_float(r["loc_radius"]) for r in latent_rows if r.get("loc_radius")})
    norm_std = rmse_std()

    phy_nrmse, phy_unstable = build_surface(phy_rows, None, inflations, radii)
    if np.any(phy_unstable):
        raise RuntimeError("The PHY tuning surface contains unstable rows; this plot expects a stable PHY reference.")

    method_data: dict[str, dict[str, np.ndarray]] = {}
    finite_abs_max = 0.0
    summary_rows: list[dict[str, object]] = []
    for method_label, _title in METHODS:
        method_nrmse, method_unstable = build_surface(latent_rows, method_label, inflations, radii)
        diff_rmse = (method_nrmse - phy_nrmse) * norm_std
        diff_rmse[method_unstable | ~np.isfinite(method_nrmse)] = np.nan
        finite = diff_rmse[np.isfinite(diff_rmse)]
        if finite.size:
            finite_abs_max = max(finite_abs_max, float(np.max(np.abs(finite))))
        method_data[method_label] = {
            "diff_rmse": diff_rmse,
            "unstable": method_unstable | ~np.isfinite(method_nrmse),
            "method_nrmse": method_nrmse,
            "method_rmse": method_nrmse * norm_std,
        }

        for i, inf in enumerate(inflations):
            for j, rad in enumerate(radii):
                summary_rows.append(
                    {
                        "method_label": method_label,
                        "inflation": inf,
                        "loc_radius": rad,
                        "latent_analysis_nrmse": method_nrmse[i, j],
                        "phy_analysis_nrmse": phy_nrmse[i, j],
                        "analysis_rmse_difference_latent_minus_phy": diff_rmse[i, j],
                        "unstable_or_missing": bool(method_data[method_label]["unstable"][i, j]),
                    }
                )

    with SUMMARY_CSV.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "method_label",
                "inflation",
                "loc_radius",
                "latent_analysis_nrmse",
                "phy_analysis_nrmse",
                "analysis_rmse_difference_latent_minus_phy",
                "unstable_or_missing",
            ],
        )
        writer.writeheader()
        writer.writerows(summary_rows)

    matplotlib.rcParams.update(
        {
            "font.family": "serif",
            "font.size": 11.0,
            "axes.linewidth": 1.0,
            "xtick.direction": "out",
            "ytick.direction": "out",
        }
    )

    color_limit = max(0.05, math.ceil(finite_abs_max * 100.0) / 100.0)
    norm = TwoSlopeNorm(vmin=-color_limit, vcenter=0.0, vmax=color_limit)
    cmap = plt.get_cmap("RdBu_r").copy()
    cmap.set_bad(color="0.72")

    fig, axes = plt.subplots(1, 3, figsize=(12.8, 3.9), constrained_layout=False)
    fig.subplots_adjust(left=0.075, right=0.895, bottom=0.18, top=0.88, wspace=0.10)

    im = None
    letters = ["(a)", "(b)", "(c)"]
    x = np.arange(len(radii) + 1)
    y = np.arange(len(inflations) + 1)
    for ax, (method_label, title), letter in zip(axes, METHODS, letters):
        diff = method_data[method_label]["diff_rmse"]
        unstable = method_data[method_label]["unstable"]
        im = ax.pcolormesh(x, y, diff, cmap=cmap, norm=norm, shading="flat", edgecolors="white", linewidth=0.55)
        ax.set_xticks(np.arange(len(radii)) + 0.5)
        ax.set_xticklabels([str(int(r)) for r in radii], rotation=35, ha="right")
        ax.set_yticks(np.arange(len(inflations)) + 0.5)
        ax.set_yticklabels([f"{v:.2f}".rstrip("0").rstrip(".") for v in inflations])
        ax.set_xlabel("Localization radius")
        ax.set_title(f"{letter} {title}", fontsize=13.0)
        if ax is axes[0]:
            ax.set_ylabel("Inflation")
        else:
            ax.set_ylabel("")
            ax.tick_params(axis="y", labelleft=False)

        ys, xs = np.where(unstable)
        for yy, xx in zip(ys, xs):
            ax.text(
                xx + 0.5,
                yy + 0.5,
                "×",
                ha="center",
                va="center",
                color="black",
                fontsize=11.0,
                fontweight="bold",
            )

        method_rmse = method_data[method_label]["method_rmse"]
        stable_rmse = np.where(unstable, np.nan, method_rmse)
        best_flat = int(np.nanargmin(stable_rmse)) if np.any(np.isfinite(stable_rmse)) else -1
        best_yy, best_xx = np.unravel_index(best_flat, stable_rmse.shape) if best_flat >= 0 else (-1, -1)
        for yy in range(len(inflations)):
            for xx in range(len(radii)):
                if unstable[yy, xx] or not np.isfinite(method_rmse[yy, xx]):
                    continue
                text = f"{method_rmse[yy, xx]:.2f}"
                if yy == best_yy and xx == best_xx:
                    text += "*"
                ax.text(
                    xx + 0.5,
                    yy + 0.5,
                    text,
                    ha="center",
                    va="center",
                    fontsize=7.4,
                    color="black",
                    fontweight="bold" if yy == best_yy and xx == best_xx else "normal",
                )

    if im is not None:
        cax = fig.add_axes([0.915, 0.23, 0.018, 0.62])
        cbar = fig.colorbar(im, cax=cax, extend="both")
        cbar.set_label("Analysis RMSE difference\n(latent - PHY)")

    fig.savefig(FIG_PATH, dpi=300)
    fig.savefig(FIG_PDF)
    plt.close(fig)

    if ORGANIZED.exists():
        shutil.copy2(FIG_PATH, ORGANIZED / FIG_PATH.name)
        shutil.copy2(FIG_PDF, ORGANIZED / FIG_PDF.name)
        shutil.copy2(FIG_PATH, ORGANIZED / "Fig10_N20_method_specific_tuning_heatmaps.png")
        shutil.copy2(FIG_PDF, ORGANIZED / "Fig10_N20_method_specific_tuning_heatmaps.pdf")
        shutil.copy2(SUMMARY_CSV, ORGANIZED / "Fig10_N20_method_specific_tuning_heatmaps_values.csv")

    print(FIG_PATH)
    print(FIG_PDF)
    print(SUMMARY_CSV)


if __name__ == "__main__":
    main()
