from __future__ import annotations

import csv
import os
import shutil
from pathlib import Path

import numpy as np


ROOT = Path(r"D:\LDA")
PIC = ROOT / "PIC"
PRIOR_DIR = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "results"
    / "prior_ensembles"
)

ORG_FIG_PATH = PIC / "Organized" / "Fig09_abs_N10_N40_corr_difference_early_late_full_PHY_CBR_wet_vmax055.png"
ORG_FIG_PDF = PIC / "Organized" / "Fig09_abs_N10_N40_corr_difference_early_late_full_PHY_CBR_wet_vmax055.pdf"
SUMMARY_CSV = PIC / "Fig09_abs_N10_N40_corr_difference_early_late_full_PHY_summary.csv"
CBR_WET = Path(r"CBR_wet.rgb")

PHYSICAL_LOC_RADIUS = 192
LATENT_GRID_SIZE = 128
LATENT_GRID_SPACING = 2048 // LATENT_GRID_SIZE
LATENT_LOC_RADIUS = PHYSICAL_LOC_RADIUS // LATENT_GRID_SPACING
COLOR_LIMIT = 0.55

N10_CASE = "N10_seed"
N40_CASE = "N40_seed"

WINDOWS = [
    ("Cycles 1-50", slice(0, 50)),
    ("Cycles 51-200", slice(50, 200)),
]

SPACES = [
    ("PHY", "physical", "PHY"),
    ("GLA_256", "global_latent", "GLA_256"),
    ("SLA_256", "stored_latent", "SLA_256"),
    ("SLA_128", "stored_latent", "SLA_128"),
]


def load_prior_npz(case_key: str, method: str) -> dict[str, np.ndarray]:
    path = PRIOR_DIR / f"{case_key}_{method}_physhared_prior_ensembles.npz"
    if not path.exists():
        raise FileNotFoundError(path)
    with np.load(path) as data:
        return {key: data[key] for key in data.files}


def all_cycle_variables(data: dict[str, np.ndarray], space_kind: str) -> np.ndarray:
    if space_kind == "physical":
        return data["prior_ensemble"]
    if space_kind in {"stored_latent", "global_latent"}:
        latent = data["prior_latent_ensemble"]
        return latent.reshape(latent.shape[0], latent.shape[1], -1)
    raise ValueError(space_kind)


def corr_matrix(variables: np.ndarray) -> np.ndarray:
    x = np.asarray(variables, dtype=np.float64)
    x = x - x.mean(axis=0, keepdims=True)
    std = x.std(axis=0, ddof=1)
    cov = x.T @ x / float(x.shape[0] - 1)
    denom = std[:, None] * std[None, :]
    corr = np.divide(cov, denom, out=np.zeros_like(cov), where=denom > 1.0e-14)
    return np.clip(corr, -1.0, 1.0)


def mean_abs_difference(
    n10_vars: np.ndarray,
    n40_vars: np.ndarray,
    window: slice,
) -> tuple[np.ndarray, int]:
    n_cycles = min(n10_vars.shape[0], n40_vars.shape[0])
    start = int(window.start or 0)
    stop = int(window.stop if window.stop is not None else n_cycles)
    stop = min(stop, n_cycles)
    abs_sum: np.ndarray | None = None
    count = 0
    for cycle in range(start, stop):
        diff_abs = np.abs(corr_matrix(n10_vars[cycle]) - corr_matrix(n40_vars[cycle]))
        if abs_sum is None:
            abs_sum = np.zeros_like(diff_abs)
        abs_sum += diff_abs
        count += 1
    if abs_sum is None or count == 0:
        raise RuntimeError("No paired cycles were available.")
    return abs_sum / count, count


def offdiag_mask(n: int) -> np.ndarray:
    return ~np.eye(n, dtype=bool)


def periodic_band_mask(n: int, radius: int) -> np.ndarray:
    idx = np.arange(n)
    d = np.abs(idx[:, None] - idx[None, :])
    d = np.minimum(d, n - d)
    return (d <= radius) & offdiag_mask(n)


def sla256_localization_mask() -> np.ndarray:
    n = 2 * LATENT_GRID_SIZE
    cell = np.arange(n) % LATENT_GRID_SIZE
    d = np.abs(cell[:, None] - cell[None, :])
    d = np.minimum(d, LATENT_GRID_SIZE - d)
    return (d <= LATENT_LOC_RADIUS) & offdiag_mask(n)


def localization_mask(label: str, n: int) -> np.ndarray | None:
    if label == "PHY":
        return periodic_band_mask(n, PHYSICAL_LOC_RADIUS)
    if label == "SLA_128":
        return periodic_band_mask(n, LATENT_LOC_RADIUS)
    if label == "SLA_256":
        return sla256_localization_mask()
    return None


def draw_periodic_band_block(ax, x_start: int, y_start: int, size: int, radius: int) -> None:
    x_end = x_start + size - 1
    y_end = y_start + size - 1
    r = int(radius)
    if r <= 0 or r >= size:
        return
    style = {
        "color": "0.0",
        "lw": 1.75,
        "ls": (0, (4, 3)),
        "alpha": 0.78,
        "solid_capstyle": "butt",
        "zorder": 4,
    }
    ax.plot([x_start, x_end - r], [y_start + r, y_end], **style)
    ax.plot([x_start + r, x_end], [y_start, y_end - r], **style)
    ax.plot([x_start, x_start + r], [y_end - r, y_end], **style)
    ax.plot([x_end - r, x_end], [y_start, y_start + r], **style)


def add_localization_guides(ax, label: str) -> None:
    if label == "PHY":
        draw_periodic_band_block(ax, 0, 0, 2048, PHYSICAL_LOC_RADIUS)
    elif label == "SLA_128":
        draw_periodic_band_block(ax, 0, 0, LATENT_GRID_SIZE, LATENT_LOC_RADIUS)
    elif label == "SLA_256":
        for x0 in (0, LATENT_GRID_SIZE):
            for y0 in (0, LATENT_GRID_SIZE):
                draw_periodic_band_block(ax, x0, y0, LATENT_GRID_SIZE, LATENT_LOC_RADIUS)


def add_sla256_block_guides(ax) -> None:
    boundary = LATENT_GRID_SIZE - 0.5
    style = {
        "color": "0.28",
        "lw": 2.05,
        "ls": (0, (3, 3)),
        "alpha": 0.9,
        "zorder": 5,
    }
    ax.axvline(boundary, **style)
    ax.axhline(boundary, **style)


def load_cbr_wet_cmap():
    import matplotlib.colors as mcolors

    if not CBR_WET.exists():
        return "magma_r"
    colors = []
    with CBR_WET.open("r", encoding="utf-8") as f:
        for line in f:
            parts = line.split()
            if len(parts) < 3:
                continue
            try:
                rgb = [float(parts[0]), float(parts[1]), float(parts[2])]
            except ValueError:
                continue
            colors.append(rgb)
    arr = np.asarray(colors, dtype=float)
    if arr.size == 0:
        return "magma_r"
    if np.nanmax(arr) > 1.0:
        arr = arr / 255.0
    return mcolors.LinearSegmentedColormap.from_list("CBR_wet", arr)


def set_axis_labels(ax, label: str, row: int) -> None:
    if label == "PHY":
        ticks = np.array([0, 512, 1024, 1536, 2047])
        ax.set_xticks(ticks)
        ax.set_yticks(ticks)
        ax.set_yticklabels([str(int(t)) for t in ticks])
        if row == 1:
            ax.set_xticklabels([str(int(t)) for t in ticks])
            ax.set_xlabel("Physical grid index")
        else:
            ax.set_xticklabels([])
        ax.set_ylabel("Physical grid index")
    else:
        n = 128 if label == "SLA_128" else 256
        step = 25 if n == 128 else 50
        ticks = np.arange(0, n, step)
        ax.set_xticks(ticks)
        ax.set_yticks(ticks)
        ax.set_yticklabels([str(int(t)) for t in ticks])
        if row == 1:
            ax.set_xticklabels([str(int(t)) for t in ticks])
            ax.set_xlabel("Latent variable index")
        else:
            ax.set_xticklabels([])
        ax.set_ylabel("Latent variable index")


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

    import matplotlib

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

    loaded: dict[tuple[str, str], np.ndarray] = {}
    for label, space_kind, method in SPACES:
        for case_key in (N10_CASE, N40_CASE):
            data = load_prior_npz(case_key, method)
            loaded[(case_key, label)] = all_cycle_variables(data, space_kind)

    maps: dict[tuple[str, str], np.ndarray] = {}
    rows: list[dict[str, object]] = []
    for window_label, window_slice in WINDOWS:
        for label, _space_kind, _method in SPACES:
            abs_diff, count = mean_abs_difference(
                loaded[(N10_CASE, label)],
                loaded[(N40_CASE, label)],
                window_slice,
            )
            maps[(window_label, label)] = abs_diff

            offdiag = offdiag_mask(abs_diff.shape[0])
            loc = localization_mask(label, abs_diff.shape[0])
            row = {
                "window": window_label,
                "label": label,
                "n_variables": int(abs_diff.shape[0]),
                "paired_cycles": int(count),
                "mean_abs_offdiag_cycle_abs_N10_minus_N40": f"{float(np.mean(abs_diff[offdiag])):.10e}",
                "median_abs_offdiag_cycle_abs_N10_minus_N40": f"{float(np.median(abs_diff[offdiag])):.10e}",
                "p90_abs_offdiag_cycle_abs_N10_minus_N40": f"{float(np.percentile(abs_diff[offdiag], 90)):.10e}",
                "mean_abs_localized_cycle_abs_N10_minus_N40": "",
                "median_abs_localized_cycle_abs_N10_minus_N40": "",
                "p90_abs_localized_cycle_abs_N10_minus_N40": "",
            }
            if loc is not None:
                row["mean_abs_localized_cycle_abs_N10_minus_N40"] = f"{float(np.mean(abs_diff[loc])):.10e}"
                row["median_abs_localized_cycle_abs_N10_minus_N40"] = f"{float(np.median(abs_diff[loc])):.10e}"
                row["p90_abs_localized_cycle_abs_N10_minus_N40"] = f"{float(np.percentile(abs_diff[loc], 90)):.10e}"
            rows.append(row)

    with SUMMARY_CSV.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)

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

    fig, axes = plt.subplots(2, 4, figsize=(16.8, 6.4), constrained_layout=False)
    panel_letters = list("abcdefgh")
    cmap = load_cbr_wet_cmap()
    for row_idx, (window_label, _window_slice) in enumerate(WINDOWS):
        for col_idx, (label, _kind, _method) in enumerate(SPACES):
            ax = axes[row_idx, col_idx]
            data = maps[(window_label, label)]
            im = ax.imshow(
                data,
                origin="lower",
                cmap=cmap,
                vmin=0.0,
                vmax=COLOR_LIMIT,
                interpolation="nearest",
                aspect="auto",
            )
            add_localization_guides(ax, label)
            if label == "SLA_256":
                add_sla256_block_guides(ax)
            set_axis_labels(ax, label, row_idx)
            letter = panel_letters[row_idx * 4 + col_idx]
            ax.text(
                0.02,
                0.96,
                f"({letter}) {label}",
                transform=ax.transAxes,
                ha="left",
                va="top",
                fontsize=13.0,
                bbox={"facecolor": "white", "alpha": 0.78, "edgecolor": "none", "pad": 1.5},
            )
    fig.subplots_adjust(left=0.105, right=0.91, bottom=0.105, top=0.96, wspace=0.36, hspace=0.20)
    for row_idx, (window_label, _window_slice) in enumerate(WINDOWS):
        pos = axes[row_idx, 0].get_position()
        fig.text(
            0.052,
            0.5 * (pos.y0 + pos.y1),
            window_label,
            rotation=90,
            ha="center",
            va="center",
            fontsize=13.0,
            family="serif",
        )
    cax = fig.add_axes([0.925, 0.19, 0.018, 0.62])
    cb = fig.colorbar(im, cax=cax, extend="max")
    cb.set_label(r"Mean $|r_{10}-r_{40}|$", rotation=90, labelpad=17)

    fig.savefig(FIG_PATH, dpi=300)
    fig.savefig(FIG_PDF)
    ORG_FIG_PATH.parent.mkdir(exist_ok=True)
    shutil.copy2(FIG_PATH, ORG_FIG_PATH)
    shutil.copy2(FIG_PDF, ORG_FIG_PDF)
    print(f"Saved {FIG_PATH}")
    print(f"Saved {ORG_FIG_PATH}")
    print(f"Saved {SUMMARY_CSV}")


if __name__ == "__main__":
    main()
