from __future__ import annotations

import csv
import os
from pathlib import Path

import numpy as np


ROOT = Path(r"D:\LDA")
PIC = ROOT / "PIC"
MPLCONFIGDIR = ROOT / ".mplconfig_pic_fig30f_early_late"
MPLCONFIGDIR.mkdir(exist_ok=True)
for lock_file in MPLCONFIGDIR.glob("*.matplotlib-lock"):
    try:
        lock_file.unlink()
    except OSError:
        pass
os.environ["MPLCONFIGDIR"] = str(MPLCONFIGDIR)

RESULT_DIR = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "results"
    / "phyloc_inflation_shared"
)
TEST_CASE_PATH = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "data"
    / "test_case.npz"
)

OUT_BASE = PIC / "Fig07_cycles_early_late_combined_cumulative"
OUT_SUMMARY = OUT_BASE.with_name(OUT_BASE.name + "_summary.csv")

DEALIAS_FRAC = 2.0 / 3.0
AXIS_LABEL_FONT_SIZE = 12.2
PANEL_TAG_FONT_SIZE = 13.0
SPECTRUM_Y_MIN = 1.0e-18
SPECTRUM_Y_SPLIT = 1.0e-7
SPECTRUM_Y_MAX = 1.0e-1
SPECTRUM_LOW_FRACTION = 0.35
SPECTRUM_TICKS = np.array([1.0e-17, 1.0e-15, 1.0e-13, 1.0e-11, 1.0e-9, 1.0e-7, 1.0e-5, 1.0e-3, 1.0e-1])

METHODS = [
    ("PHY", "N20_test_PHY-EnSRF-loc.npz", "#1F4E79"),
    ("GLA_256", "N20_test_GDLR-EnSRF-loc_global_Z256.npz", "#7B3294"),
    ("SLA_256", "N20_test_SLAT-EnSRF-loc_spatial_C2_Kz128.npz", "#2E7D32"),
    ("SLA_128", "N20_test_SLAT-EnSRF-loc_spatial_C1_Kz128.npz", "#E69F00"),
]

WINDOWS = [
    {
        "label": "Cycles 1-50",
        "analysis_cycles": "1-50",
        "analysis_sl": slice(0, 50),
        "growth_posterior_cycles": "1-49",
        "growth_next_prior_cycles": "2-50",
        "growth_post_sl": slice(0, 49),
        "growth_next_sl": slice(1, 50),
    },
    {
        "label": "Cycles 51-200",
        "analysis_cycles": "51-200",
        "analysis_sl": slice(50, 200),
        "growth_posterior_cycles": "51-199",
        "growth_next_prior_cycles": "52-200",
        "growth_post_sl": slice(50, 199),
        "growth_next_sl": slice(51, 200),
    },
]


def one_sided_power(fields: np.ndarray) -> np.ndarray:
    n = fields.shape[-1]
    fft = np.fft.rfft(fields, axis=-1)
    power = (np.abs(fft) ** 2) / float(n * n)
    if n > 1:
        power[..., 1:-1] *= 2.0
    return power


def band_sum(spec: np.ndarray, lo: int, hi: int) -> float:
    return float(np.nansum(spec[lo : hi + 1]))


def piecewise_log_power_axis(values: np.ndarray | float) -> np.ndarray:
    vals = np.asarray(values, dtype=np.float64)
    clipped = np.clip(vals, SPECTRUM_Y_MIN, SPECTRUM_Y_MAX)
    log_vals = np.log10(clipped)
    log_min = np.log10(SPECTRUM_Y_MIN)
    log_split = np.log10(SPECTRUM_Y_SPLIT)
    log_max = np.log10(SPECTRUM_Y_MAX)

    transformed = np.empty_like(log_vals)
    low_tail = log_vals <= log_split
    transformed[low_tail] = SPECTRUM_LOW_FRACTION * (log_vals[low_tail] - log_min) / (log_split - log_min)
    transformed[~low_tail] = SPECTRUM_LOW_FRACTION + (1.0 - SPECTRUM_LOW_FRACTION) * (
        log_vals[~low_tail] - log_split
    ) / (log_max - log_split)
    return transformed


def spectrum_tick_labels(ticks: np.ndarray) -> list[str]:
    return [rf"$10^{{{int(np.round(np.log10(tick)))}}}$" for tick in ticks]


def compute_window(
    truth_at_cycles: np.ndarray,
    result_arrays: dict[str, tuple[np.ndarray, np.ndarray]],
    window: dict[str, object],
    max_mode: int,
) -> tuple[dict[str, dict[str, np.ndarray]], list[dict[str, object]]]:
    analysis_sl = window["analysis_sl"]
    growth_post_sl = window["growth_post_sl"]
    growth_next_sl = window["growth_next_sl"]
    assert isinstance(analysis_sl, slice)
    assert isinstance(growth_post_sl, slice)
    assert isinstance(growth_next_sl, slice)

    analysis_truth = truth_at_cycles[analysis_sl]
    growth_post_truth = truth_at_cycles[growth_post_sl]
    growth_next_truth = truth_at_cycles[growth_next_sl]

    curves: dict[str, dict[str, np.ndarray]] = {}
    rows: list[dict[str, object]] = []

    for label, _filename, _color in METHODS:
        prior_mean_all, analysis_mean_all = result_arrays[label]

        prior_error = prior_mean_all[analysis_sl] - analysis_truth
        posterior_error = analysis_mean_all[analysis_sl] - analysis_truth
        prior_power = one_sided_power(prior_error)
        posterior_power = one_sided_power(posterior_error)
        reduction_spec = (prior_power - posterior_power).mean(axis=0)

        growth_post_error = analysis_mean_all[growth_post_sl] - growth_post_truth
        growth_next_error = prior_mean_all[growth_next_sl] - growth_next_truth
        growth_spec = (one_sided_power(growth_next_error) - one_sided_power(growth_post_error)).mean(axis=0)

        prior_spec = prior_power.mean(axis=0)
        posterior_spec = posterior_power.mean(axis=0)
        cumulative_reduction = np.cumsum(reduction_spec)
        cumulative_growth = np.cumsum(growth_spec)

        curves[label] = {
            "prior": prior_spec,
            "posterior": posterior_spec,
            "cumulative_reduction": cumulative_reduction,
            "cumulative_growth": cumulative_growth,
        }

        total_prior = band_sum(prior_spec, 1, max_mode)
        total_post = band_sum(posterior_spec, 1, max_mode)
        total_reduction = band_sum(reduction_spec, 1, max_mode)
        total_growth = band_sum(growth_spec, 1, max_mode)
        rows.append(
            {
                "window": window["label"],
                "method": label,
                "analysis_cycles": window["analysis_cycles"],
                "growth_posterior_cycles": window["growth_posterior_cycles"],
                "growth_next_prior_cycles": window["growth_next_prior_cycles"],
                "total_prior_error_power_k1_682": f"{total_prior:.10e}",
                "total_posterior_error_power_k1_682": f"{total_post:.10e}",
                "total_prior_to_posterior_reduction_k1_682": f"{total_reduction:.10e}",
                "total_posterior_to_next_prior_growth_k1_682": f"{total_growth:.10e}",
                "net_reduction_minus_growth_k1_682": f"{(total_reduction - total_growth):.10e}",
                "cumulative_reduction_k64": f"{cumulative_reduction[64]:.10e}",
                "cumulative_growth_k64": f"{cumulative_growth[64]:.10e}",
                "cumulative_reduction_k256": f"{cumulative_reduction[256]:.10e}",
                "cumulative_growth_k256": f"{cumulative_growth[256]:.10e}",
            }
        )

    return curves, rows


def main() -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D

    case = np.load(TEST_CASE_PATH)
    truth_all = np.asarray(case["truth"], dtype=np.float64)
    obs_times = np.asarray(case["obs_times"], dtype=np.int64)
    truth_at_cycles = truth_all[obs_times]

    n = truth_at_cycles.shape[-1]
    max_mode = int(np.floor((n // 2) * DEALIAS_FRAC))
    modes = np.arange(n // 2 + 1)
    plot_modes = modes[1 : max_mode + 1]

    result_arrays: dict[str, tuple[np.ndarray, np.ndarray]] = {}
    for label, filename, _color in METHODS:
        with np.load(RESULT_DIR / filename) as data:
            result_arrays[label] = (
                np.asarray(data["prior_mean"], dtype=np.float64),
                np.asarray(data["analysis_mean"], dtype=np.float64),
            )

    all_curves: list[dict[str, dict[str, np.ndarray]]] = []
    rows: list[dict[str, object]] = []
    for window in WINDOWS:
        curves, window_rows = compute_window(truth_at_cycles, result_arrays, window, max_mode)
        all_curves.append(curves)
        rows.extend(window_rows)

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

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

    fig, axes = plt.subplots(2, 3, figsize=(14.5, 8.1), facecolor="white", sharex="col")

    for row_idx, (window, curves) in enumerate(zip(WINDOWS, all_curves)):
        row_axes = axes[row_idx]
        for label, _filename, color in METHODS:
            row_axes[0].semilogx(
                plot_modes,
                piecewise_log_power_axis(curves[label]["prior"][1 : max_mode + 1]),
                color=color,
                lw=1.55,
            )
            row_axes[1].semilogx(
                plot_modes,
                piecewise_log_power_axis(curves[label]["posterior"][1 : max_mode + 1]),
                color=color,
                lw=1.55,
            )
            row_axes[2].semilogx(
                plot_modes,
                curves[label]["cumulative_reduction"][1 : max_mode + 1],
                color=color,
                lw=1.65,
            )
            row_axes[2].semilogx(
                plot_modes,
                curves[label]["cumulative_growth"][1 : max_mode + 1],
                color=color,
                lw=1.65,
                ls="--",
        )

        row_axes[0].set_ylabel(f"{window['label']}\nError power")
        row_axes[1].set_ylabel("Error power")
        row_axes[2].set_ylabel("Cumulative error power")

    titles = ["Prior error power", "Posterior error power", "Cumulative error reduction and growth"]
    for col, title in enumerate(titles):
        axes[0, col].set_title(title, pad=7)

    tags = ["(a)", "(b)", "(c)", "(d)", "(e)", "(f)"]
    for ax, tag in zip(axes.ravel(), tags):
        ax.set_xlim(1, max_mode)
        ax.set_xlabel("")
        ax.text(
            0.025,
            0.975,
            tag,
            transform=ax.transAxes,
            ha="left",
            va="top",
            fontsize=PANEL_TAG_FONT_SIZE,
            bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.82, "pad": 1.5},
        )
        ax.grid(True, which="major", color="0.88", lw=0.8)
        ax.grid(True, which="minor", color="0.93", lw=0.5)
        ax.set_axisbelow(True)

    for ax in axes.ravel():
        ax.xaxis.label.set_size(AXIS_LABEL_FONT_SIZE)
        ax.yaxis.label.set_size(AXIS_LABEL_FONT_SIZE)

    for ax in axes[1, :]:
        ax.set_xlabel("Wavenumber")

    for ax in axes[0, :]:
        ax.tick_params(axis="x", labelbottom=False)

    spectrum_ticks = SPECTRUM_TICKS[(SPECTRUM_TICKS >= SPECTRUM_Y_MIN) & (SPECTRUM_TICKS <= SPECTRUM_Y_MAX)]
    spectrum_tick_positions = piecewise_log_power_axis(spectrum_ticks)
    for ax in list(axes[:, 0]) + list(axes[:, 1]):
        ax.set_ylim(float(piecewise_log_power_axis(SPECTRUM_Y_MIN)), float(piecewise_log_power_axis(SPECTRUM_Y_MAX)))
        ax.set_yticks(spectrum_tick_positions)
        ax.set_yticklabels(spectrum_tick_labels(spectrum_ticks))
        ax.axhline(float(piecewise_log_power_axis(SPECTRUM_Y_SPLIT)), color="0.78", lw=0.65, ls=":")

    for ax in axes[:, 2]:
        ax.axhline(0.0, color="0.2", lw=0.85)
        ax.axvline(64, color="0.6", lw=0.85, ls=":")
        ax.axvline(256, color="0.6", lw=0.85, ls=":")

    method_handles = [Line2D([0], [0], color=color, lw=2.5, label=label) for label, _filename, color in METHODS]
    style_handles = [
        Line2D([0], [0], color="0.2", lw=1.8, label="analysis error reduction"),
        Line2D([0], [0], color="0.2", lw=1.8, ls="--", label="forecast error growth"),
    ]
    fig.legend(
        handles=method_handles + style_handles,
        loc="upper center",
        ncol=6,
        frameon=False,
        bbox_to_anchor=(0.54, 0.975),
        columnspacing=1.25,
        handlelength=2.0,
    )
    fig.tight_layout(rect=(0.02, 0.02, 0.995, 0.905))
    fig.savefig(OUT_BASE.with_suffix(".png"), dpi=300)
    fig.savefig(OUT_BASE.with_suffix(".pdf"))
    plt.close(fig)

    print(OUT_BASE.with_suffix(".png"))
    print(OUT_BASE.with_suffix(".pdf"))
    print(OUT_SUMMARY)


if __name__ == "__main__":
    main()
