from __future__ import annotations

from pathlib import Path
import csv
import math
import os

import numpy as np


SCRIPT_DIR = Path(__file__).resolve().parent
ROOT = SCRIPT_DIR.parent
MPLCONFIGDIR = ROOT / ".mplconfig_pic_fig04_sawtooth_n20"
MPLCONFIGDIR.mkdir(exist_ok=True)
for lock_file in MPLCONFIGDIR.glob("*.matplotlib-lock"):
    try:
        lock_file.unlink()
    except OSError:
        pass
os.environ.setdefault("MPLCONFIGDIR", str(MPLCONFIGDIR))

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D


NOINFL_DIR = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "results"
    / "four_convae_testoracle_unlocinf100_physharedloc"
)
PHYSHARED_DIR = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "results"
    / "phyloc_inflation_shared"
)

OUT_PNG = SCRIPT_DIR / "Fig04_physhared_N20_sawtooth_ladder.png"
OUT_PDF = SCRIPT_DIR / "Fig04_physhared_N20_sawtooth_ladder.pdf"
OUT_CSV = SCRIPT_DIR / "Fig04_physhared_N20_sawtooth_ladder_sources.csv"
TEST_CASE_PATH = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "data"
    / "test_case.npz"
)
RMSE_SCALE = float(np.std(np.load(TEST_CASE_PATH)["truth"]))

N_CYCLES = 200
DISCARD_CYCLES = 50

METHODS = [
    ("PHY", "#1F4E79"),
    ("GLA_256", "#7B3294"),
    ("GLA_128", "#D81B60"),
    ("SLA_256", "#2E7D32"),
    ("SLA_128", "#E69F00"),
]
NO_DA_COLOR = "0.0"
LEFT_LINEWIDTH = 1.15

LEFT_ROWS = [
    (
        "No localization, no inflation",
        NOINFL_DIR,
        {
            "No DA": "test_NODA.npz",
            "PHY": "test_PHY-EnSRF-unloc.npz",
            "GLA_256": "test_GLAT-EnSRF-unloc_global_Z256.npz",
            "GLA_128": "test_GLAT-EnSRF-unloc_global_Z128.npz",
            "SLA_256": "test_LAT-EnSRF-unloc_spatial_C2_Kz128.npz",
            "SLA_128": "test_LAT-EnSRF-unloc_spatial_C1_Kz128.npz",
        },
    ),
    (
        "Inflation only",
        PHYSHARED_DIR,
        {
            "No DA": "N20_test_NODA.npz",
            "PHY": "N20_test_PHY-EnSRF-unloc.npz",
            "GLA_256": "N20_test_GLAT-EnSRF-unloc_global_Z256.npz",
            "GLA_128": "N20_test_GLAT-EnSRF-unloc_global_Z128.npz",
            "SLA_256": "N20_test_LAT-EnSRF-unloc_spatial_C2_Kz128.npz",
            "SLA_128": "N20_test_LAT-EnSRF-unloc_spatial_C1_Kz128.npz",
        },
    ),
    (
        "Localization + inflation",
        PHYSHARED_DIR,
        {
            "No DA": "N20_test_NODA.npz",
            "PHY": "N20_test_PHY-EnSRF-loc.npz",
            "GLA_256": "N20_test_GDLR-EnSRF-loc_global_Z256.npz",
            "GLA_128": "N20_test_GDLR-EnSRF-loc_global_Z128.npz",
            "SLA_256": "N20_test_SLAT-EnSRF-loc_spatial_C2_Kz128.npz",
            "SLA_128": "N20_test_SLAT-EnSRF-loc_spatial_C1_Kz128.npz",
        },
    ),
]

RIGHT_DIFFS = [
    ("GLA_256 - PHY", "GLA_256", "#7B3294"),
    ("SLA_256 - PHY", "SLA_256", "#2E7D32"),
    ("SLA_128 - PHY", "SLA_128", "#E69F00"),
]


def load_curve(base_dir: Path, filename: str) -> dict[str, np.ndarray | str | bool | int | float]:
    path = base_dir / filename
    if not path.exists():
        raise FileNotFoundError(path)
    data = np.load(path)
    prior = data["prior_nrmse"].astype(float) * RMSE_SCALE
    analysis = data["analysis_nrmse"].astype(float) * RMSE_SCALE
    finite = np.isfinite(prior) & np.isfinite(analysis)
    first_bad = int(np.argmax(~finite)) if not np.all(finite) else -1
    return {
        "path": str(path),
        "prior": prior,
        "analysis": analysis,
        "stable": bool(np.all(finite)),
        "first_bad": first_bad,
        "mean_prior": float(np.nanmean(prior[DISCARD_CYCLES:])),
        "mean_analysis": float(np.nanmean(analysis[DISCARD_CYCLES:])),
    }


def sawtooth_xy(prior: np.ndarray, analysis: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    x = np.empty(prior.size * 2, dtype=float)
    y = np.empty(prior.size * 2, dtype=float)
    cycles = np.arange(1, prior.size + 1, dtype=float)
    x[0::2] = cycles - 0.35
    x[1::2] = cycles + 0.35
    y[0::2] = prior
    y[1::2] = analysis
    return x, y


def write_source_table(loaded: dict[str, dict[str, dict[str, np.ndarray | str | bool | int | float]]]) -> None:
    fields = [
        "setting",
        "method_label",
        "source_npz",
        "normalization_truth_std",
        "stable",
        "first_nonfinite_cycle_1based",
        "mean_prior_rmse_after_discard",
        "mean_analysis_rmse_after_discard",
    ]
    with OUT_CSV.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()
        for setting, curves in loaded.items():
            for label, curve in curves.items():
                first_bad = int(curve["first_bad"])
                mean_prior = float(curve["mean_prior"])
                mean_analysis = float(curve["mean_analysis"])
                writer.writerow(
                    {
                        "setting": setting,
                        "method_label": label,
                        "source_npz": curve["path"],
                        "normalization_truth_std": f"{RMSE_SCALE:.8f}",
                        "stable": curve["stable"],
                        "first_nonfinite_cycle_1based": "" if first_bad < 0 else first_bad + 1,
                        "mean_prior_rmse_after_discard": f"{mean_prior:.8f}"
                        if math.isfinite(mean_prior)
                        else "nan",
                        "mean_analysis_rmse_after_discard": f"{mean_analysis:.8f}"
                        if math.isfinite(mean_analysis)
                        else "nan",
                    }
                )


def main() -> None:
    loaded: dict[str, dict[str, dict[str, np.ndarray | str | bool | int | float]]] = {}
    for setting, base_dir, files in LEFT_ROWS:
        loaded[setting] = {label: load_curve(base_dir, filename) for label, filename in files.items()}
    write_source_table(loaded)

    plt.rcParams.update(
        {
            "font.family": "DejaVu Serif",
            "font.size": 10.5,
            "mathtext.fontset": "dejavuserif",
            "axes.linewidth": 0.9,
        }
    )

    fig, axes = plt.subplots(
        3,
        2,
        figsize=(12.3, 8.8),
        sharex="col",
        gridspec_kw={"width_ratios": [1.45, 1.0], "wspace": 0.18, "hspace": 0.20},
        facecolor="white",
    )

    left_panel_labels = ["(a)", "(b)", "(c)"]
    right_panel_labels = ["(d)", "(e)", "(f)"]

    left_ylims = [
        (1.00, 2.30),
        (1.00, 2.30),
        (0.10, 1.70),
    ]
    right_ylim = (-0.55 * RMSE_SCALE, 0.58 * RMSE_SCALE)

    for row_index, (setting, _, _) in enumerate(LEFT_ROWS):
        ax = axes[row_index, 0]
        ax.axvspan(0.5, DISCARD_CYCLES + 0.5, color="0.92", lw=0, zorder=0)
        for label, color in METHODS:
            curve = loaded[setting][label]
            x, y = sawtooth_xy(curve["prior"], curve["analysis"])
            ax.plot(x, y, color=color, lw=LEFT_LINEWIDTH, alpha=0.88, label=label)
            if not curve["stable"]:
                first_bad = int(curve["first_bad"])
                if first_bad > 0:
                    ax.plot(
                        first_bad + 0.5,
                        float(curve["analysis"][first_bad - 1]),
                        marker="x",
                        color=color,
                        ms=7,
                        mew=1.6,
                    )
                    ax.text(
                        first_bad + 4.0,
                        float(curve["analysis"][first_bad - 1]) + 0.04,
                        "unstable",
                        color=color,
                        fontsize=8.5,
                    )
        ax.text(0.012, 0.93, left_panel_labels[row_index], transform=ax.transAxes, ha="left", va="top")
        if row_index < 3:
            noda_curve = loaded[setting]["No DA"]
            ax.plot(
                np.arange(1, N_CYCLES + 1, dtype=float),
                noda_curve["analysis"],
                color=NO_DA_COLOR,
                lw=LEFT_LINEWIDTH,
                ls=(0, (4.0, 2.2)),
                alpha=0.85,
                label="No DA",
            )
        ax.set_ylabel("RMSE")
        ax.set_xlim(0.5, N_CYCLES + 0.5)
        ax.set_ylim(*left_ylims[row_index])
        ax.grid(True, axis="y", color="0.88", lw=0.8)
        ax.set_axisbelow(True)
        if row_index == 0:
            legend_handles = [
                Line2D([0], [0], color=NO_DA_COLOR, lw=LEFT_LINEWIDTH, ls=(0, (4.0, 2.2)), label="NoDA"),
                Line2D([0], [0], color="#1F4E79", lw=LEFT_LINEWIDTH, label="PHY"),
                Line2D([0], [0], color="#7B3294", lw=LEFT_LINEWIDTH, label="GLA_256"),
                Line2D([0], [0], color="#D81B60", lw=LEFT_LINEWIDTH, label="GLA_128"),
                Line2D([0], [0], color="#2E7D32", lw=LEFT_LINEWIDTH, label="SLA_256"),
                Line2D([0], [0], color="#E69F00", lw=LEFT_LINEWIDTH, label="SLA_128"),
            ]
            ax.legend(
                handles=legend_handles,
                loc="upper center",
                bbox_to_anchor=(0.58, 1.28),
                ncol=6,
                frameon=False,
                handlelength=2.0,
                columnspacing=0.85,
            )

    loc_setting = "Localization + inflation"
    phy_analysis = loaded[loc_setting]["PHY"]["analysis"]
    cycles = np.arange(1, N_CYCLES + 1, dtype=float)
    right_text_layouts = [
        {
            "title_xy": (0.03, 0.08),
            "title_ha": "left",
            "title_va": "bottom",
            "mean_xy": (0.98, 0.08),
            "mean_ha": "right",
            "mean_va": "bottom",
        },
        {
            "title_xy": (0.03, 0.08),
            "title_ha": "left",
            "title_va": "bottom",
            "mean_xy": (0.98, 0.08),
            "mean_ha": "right",
            "mean_va": "bottom",
        },
        {
            "title_xy": (0.03, 0.08),
            "title_ha": "left",
            "title_va": "bottom",
            "mean_xy": (0.98, 0.08),
            "mean_ha": "right",
            "mean_va": "bottom",
        },
    ]
    for row_index, (title, label, color) in enumerate(RIGHT_DIFFS):
        ax = axes[row_index, 1]
        layout = right_text_layouts[row_index]
        ax.axhspan(right_ylim[0], 0.0, color="#EAF4EA", lw=0, zorder=-2)
        ax.axhspan(0.0, right_ylim[1], color="#F8ECEC", lw=0, zorder=-2)
        ax.axvspan(0.5, DISCARD_CYCLES + 0.5, color="0.92", lw=0, zorder=0)
        curve = loaded[loc_setting][label]
        diff = curve["analysis"] - phy_analysis
        ax.axhline(0.0, color="0.25", lw=0.9)
        ax.plot(cycles, diff, color=color, lw=1.35)
        ax.text(0.02, 0.93, right_panel_labels[row_index], transform=ax.transAxes, ha="left", va="top")
        ax.text(
            layout["title_xy"][0],
            layout["title_xy"][1],
            title,
            transform=ax.transAxes,
            ha=layout["title_ha"],
            va=layout["title_va"],
            fontsize=9.0,
            bbox=dict(facecolor="white", edgecolor="none", alpha=0.92, pad=1.8),
        )
        mean_diff = float(np.nanmean(diff[DISCARD_CYCLES:]))
        ax.text(
            layout["mean_xy"][0],
            layout["mean_xy"][1],
            f"mean = {mean_diff:+.3f}",
            transform=ax.transAxes,
            ha=layout["mean_ha"],
            va=layout["mean_va"],
            fontsize=9.0,
            color=color,
            bbox=dict(facecolor="white", edgecolor="none", alpha=0.88, pad=1.6),
        )
        ax.set_ylabel("Analysis RMSE difference")
        ax.set_xlim(0.5, N_CYCLES + 0.5)
        ax.set_ylim(*right_ylim)
        ax.grid(True, axis="y", color="0.88", lw=0.8)
        ax.set_axisbelow(True)

    for ax in axes[-1, :]:
        ax.set_xlabel("DA cycle")

    for ax in axes.ravel():
        ax.tick_params(direction="out", length=4, width=0.8)

    fig.savefig(OUT_PNG, dpi=300, bbox_inches="tight")
    fig.savefig(OUT_PDF, bbox_inches="tight")
    print(f"saved {OUT_PNG}")
    print(f"saved {OUT_PDF}")
    print(f"saved {OUT_CSV}")


if __name__ == "__main__":
    main()
