from __future__ import annotations

import csv
import os
from pathlib import Path

import numpy as np


ROOT = Path(r"D:\LDA")
PIC = ROOT / "PIC"
OUT_BASE = PIC / "Fig05_obsdensity_scaledR_32_128_rmse_diff"
DISCARD = 50

SCALED_RESULT_DIR = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "results"
    / "N20_obsdensity_scaledR_fixed_phy_loc"
)
TEST_CASE_PATH = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "data"
    / "test_case.npz"
)
RMSE_SCALE = float(np.std(np.load(TEST_CASE_PATH)["truth"]))

CASES = [
    {
        "label": "32 obs\n$\\sigma_o=0.25$",
        "case_label": "obs64_half_Rhalf",
        "obs_count": 32,
        "obs_error_std": 0.25,
        "keys": {
            "PHY": "obs64_half_Rhalf_test_PHY.npz",
            "GLA_256": "obs64_half_Rhalf_test_GLA_256.npz",
            "SLA_256": "obs64_half_Rhalf_test_SLA_256.npz",
            "SLA_128": "obs64_half_Rhalf_test_SLA_128.npz",
        },
    },
    {
        "label": "128 obs\n$\\sigma_o=1.0$",
        "case_label": "obs16_double_Rdouble",
        "obs_count": 128,
        "obs_error_std": 1.0,
        "keys": {
            "PHY": "obs16_double_Rdouble_test_PHY.npz",
            "GLA_256": "obs16_double_Rdouble_test_GLA_256.npz",
            "SLA_256": "obs16_double_Rdouble_test_SLA_256.npz",
            "SLA_128": "obs16_double_Rdouble_test_SLA_128.npz",
        },
    },
]

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


def load_analysis(filename: str) -> np.ndarray:
    path = SCALED_RESULT_DIR / filename
    if not path.exists():
        raise FileNotFoundError(path)
    return np.load(path)["analysis_nrmse"].astype(float) * RMSE_SCALE


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

    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.rcParams.update(
        {
            "font.family": "DejaVu Serif",
            "font.size": 13,
            "mathtext.fontset": "dejavuserif",
            "axes.linewidth": 1.15,
            "xtick.major.width": 1.0,
            "ytick.major.width": 1.0,
        }
    )

    fig, axes = plt.subplots(
        nrows=3,
        ncols=2,
        figsize=(10.7, 8.6),
        sharex=True,
        sharey=True,
        constrained_layout=False,
    )

    summary_rows: list[dict[str, object]] = []
    panel_letters = iter("abcdef")
    ylim = (-1.15, 0.65)

    for col, case in enumerate(CASES):
        phy = load_analysis(case["keys"]["PHY"])
        x = np.arange(1, phy.size + 1)
        for row, (title, key, color) in enumerate(METHODS):
            ax = axes[row, col]
            target = load_analysis(case["keys"][key])
            diff = target - phy
            mean_after = float(np.nanmean(diff[DISCARD:]))
            frac_better = float(np.nanmean(diff[DISCARD:] < 0.0))

            ax.axhspan(0, ylim[1], color="#F7E6E6", alpha=0.75, zorder=0)
            ax.axhspan(ylim[0], 0, color="#E7F3E8", alpha=0.85, zorder=0)
            ax.axvspan(1, DISCARD, color="0.88", alpha=0.8, zorder=0)
            ax.plot(x, diff, color=color, lw=1.75)
            ax.axhline(0, color="0.25", lw=1.0)
            ax.grid(axis="y", color="0.86", lw=0.8)
            ax.set_ylim(*ylim)
            ax.set_xlim(1, phy.size)
            ax.text(
                0.03,
                0.90,
                f"({next(panel_letters)})",
                transform=ax.transAxes,
                ha="left",
                va="center",
                fontsize=14,
            )
            ax.text(
                0.04,
                0.08,
                title,
                transform=ax.transAxes,
                ha="left",
                va="bottom",
                fontsize=12.5,
                bbox=dict(facecolor="white", edgecolor="none", alpha=0.88, pad=1.5),
            )
            ax.text(
                0.96,
                0.08,
                f"mean = {mean_after:+.3f}",
                transform=ax.transAxes,
                ha="right",
                va="bottom",
                fontsize=12.5,
                color=color,
                bbox=dict(facecolor="white", edgecolor="none", alpha=0.88, pad=1.5),
            )
            if col == 0:
                ax.set_ylabel("Analysis RMSE difference", fontsize=13)
            if row == 2:
                ax.set_xlabel("DA cycle", fontsize=13)

            summary_rows.append(
                {
                    "obs_label": str(case["label"]).replace("\n", " "),
                    "case_label": case["case_label"],
                    "obs_count": case["obs_count"],
                    "obs_error_std": case["obs_error_std"],
                    "comparison": title,
                    "mean_rmse_difference_after_discard": mean_after,
                    "fraction_better_after_discard": frac_better,
                    "rmse_scale": RMSE_SCALE,
                }
            )

    for ax in axes.ravel():
        ax.tick_params(labelsize=12, length=4)
    fig.subplots_adjust(left=0.095, right=0.985, top=0.925, bottom=0.085, wspace=0.10, hspace=0.14)

    png = OUT_BASE.with_suffix(".png")
    pdf = OUT_BASE.with_suffix(".pdf")
    fig.savefig(png, dpi=300)
    fig.savefig(pdf)
    plt.close(fig)

    csv_path = OUT_BASE.with_name(OUT_BASE.name + "_sources.csv")
    with csv_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(
            f,
            fieldnames=[
                "obs_label",
                "case_label",
                "obs_count",
                "obs_error_std",
                "comparison",
                "mean_rmse_difference_after_discard",
                "fraction_better_after_discard",
                "rmse_scale",
            ],
        )
        writer.writeheader()
        writer.writerows(summary_rows)

    print(png)
    print(pdf)
    print(csv_path)


if __name__ == "__main__":
    main()
