from __future__ import annotations

import csv
import os
import sys
from pathlib import Path
from typing import Any

import numpy as np
import yaml


ROOT = Path(r"D:\LDA")
PIC = ROOT / "PIC"
MPLCONFIGDIR = ROOT / ".mplconfig_pic_fig08_projection_increment"
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)

sys.path.insert(0, str(ROOT))

from lda.ks_ae import (  # noqa: E402
    checkpoint_path,
    decode_global_latent,
    decode_latent,
    encode_global_physical,
    encode_physical,
    global_checkpoint_path,
    load_autoencoder,
    load_global_autoencoder,
)


RESULT_DIR = (
    ROOT
    / "outputs"
    / "PATH_NAME"
    / "results"
    / "phyloc_inflation_shared"
)
SPATIAL_C2_CONFIG = (
    ROOT
    / "configs"
    / "ks_200cyc.yaml"
)
SPATIAL_C1_CONFIG = (
    ROOT
    / "configs"
    / "ks_200cyc.yaml"
)
GLOBAL_Z256_CONFIG = (
    ROOT
    / "configs"
    / "ks_200cyc.yaml"
)

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

DEALIAS_FRAC = 2.0 / 3.0
DEVICE = "cpu"
SPECTRUM_Y_MIN = 1.0e-17
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])

ACTUAL_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"),
]

PROJECTION_METHODS = [
    ("PHY", "#1F4E79"),
    ("PHY -> GLA_256", "#7B3294"),
    ("PHY -> SLA_256", "#2E7D32"),
    ("PHY -> SLA_128", "#E69F00"),
]

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


def load_yaml(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as f:
        return yaml.safe_load(f)


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


def mean_increment_spectrum(increments: np.ndarray, window: slice) -> np.ndarray:
    return one_sided_power(increments[window]).mean(axis=0)


def project_spatial(states: np.ndarray, cfg_path: Path) -> np.ndarray:
    cfg = load_yaml(cfg_path)
    model, mean, std = load_autoencoder(checkpoint_path(cfg), device=DEVICE)
    latent = encode_physical(model, states, mean, std, DEVICE)
    return decode_latent(model, latent, mean, std, DEVICE)


def project_global(states: np.ndarray, cfg_path: Path) -> np.ndarray:
    cfg = load_yaml(cfg_path)
    model, mean, std = load_global_autoencoder(global_checkpoint_path(cfg), device=DEVICE)
    latent = encode_global_physical(model, states, mean, std, DEVICE)
    return decode_global_latent(model, latent, mean, std, DEVICE)


def load_actual_increments() -> tuple[dict[str, np.ndarray], np.ndarray, np.ndarray]:
    actual: dict[str, np.ndarray] = {}
    phy_prior: np.ndarray | None = None
    phy_analysis: np.ndarray | None = None
    for label, filename, _color in ACTUAL_METHODS:
        with np.load(RESULT_DIR / filename) as data:
            prior = np.asarray(data["prior_mean"], dtype=np.float64)
            analysis = np.asarray(data["analysis_mean"], dtype=np.float64)
        actual[label] = analysis - prior
        if label == "PHY":
            phy_prior = prior
            phy_analysis = analysis
    if phy_prior is None or phy_analysis is None:
        raise RuntimeError("Missing PHY prior/analysis arrays.")
    return actual, phy_prior, phy_analysis


def compute_projection_only_increments(phy_prior: np.ndarray, phy_analysis: np.ndarray) -> dict[str, np.ndarray]:
    projected: dict[str, np.ndarray] = {
        "PHY": phy_analysis - phy_prior,
    }
    projected["PHY -> GLA_256"] = project_global(phy_analysis, GLOBAL_Z256_CONFIG) - phy_prior
    projected["PHY -> SLA_256"] = project_spatial(phy_analysis, SPATIAL_C2_CONFIG) - phy_prior
    projected["PHY -> SLA_128"] = project_spatial(phy_analysis, SPATIAL_C1_CONFIG) - phy_prior
    return projected


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 main() -> None:
    import matplotlib

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

    actual_increments, phy_prior, phy_analysis = load_actual_increments()
    projection_increments = compute_projection_only_increments(phy_prior, phy_analysis)

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

    plt.rcParams.update(
        {
            "font.family": "DejaVu Serif",
            "font.size": 11.0,
            "mathtext.fontset": "dejavuserif",
            "axes.linewidth": 0.9,
        }
    )
    fig, axes = plt.subplots(2, 2, figsize=(11.2, 7.5), sharex=True, sharey=True, facecolor="white")

    rows: list[dict[str, object]] = []
    panel_specs = [
        (axes[0, 0], "actual", "Cycles 1-50", WINDOWS[0][1], actual_increments, ACTUAL_METHODS),
        (axes[0, 1], "projection_only", "Cycles 1-50", WINDOWS[0][1], projection_increments, PROJECTION_METHODS),
        (axes[1, 0], "actual", "Cycles 51-200", WINDOWS[1][1], actual_increments, ACTUAL_METHODS),
        (axes[1, 1], "projection_only", "Cycles 51-200", WINDOWS[1][1], projection_increments, PROJECTION_METHODS),
    ]

    for ax, family, window_label, window_slice, increments_by_label, method_specs in panel_specs:
        for label, *rest in method_specs:
            color = rest[-1]
            spec = mean_increment_spectrum(increments_by_label[label], window_slice)
            ax.semilogx(plot_modes, piecewise_log_power_axis(spec[1 : max_mode + 1]), color=color, lw=1.8)
            rows.append(
                {
                    "panel_family": family,
                    "window": window_label,
                    "method": label,
                    "total_power_k1_682": f"{band_sum(spec, 1, max_mode):.10e}",
                    "low_power_k1_64": f"{band_sum(spec, 1, 64):.10e}",
                    "active_power_k65_256": f"{band_sum(spec, 65, 256):.10e}",
                    "tail_power_k257_682": f"{band_sum(spec, 257, max_mode):.10e}",
                }
            )
        ax.grid(True, which="major", color="0.88", lw=0.8)
        ax.grid(True, which="minor", color="0.93", lw=0.5)
        ax.set_xlim(1, max_mode)
        ax.set_ylim(float(piecewise_log_power_axis(SPECTRUM_Y_MIN)), float(piecewise_log_power_axis(SPECTRUM_Y_MAX)))
        ax.set_axisbelow(True)

    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 axes.ravel():
        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=":")

    axes[0, 0].set_title("Actual DA increment")
    axes[0, 1].set_title("PHY analysis projected through AE")
    axes[1, 0].set_xlabel("Wavenumber")
    axes[1, 1].set_xlabel("Wavenumber")
    axes[0, 0].set_ylabel("Cycles 1-50\nincrement power")
    axes[1, 0].set_ylabel("Cycles 51-200\nincrement power")

    tags = ["(a)", "(b)", "(c)", "(d)"]
    for ax, tag in zip(axes.ravel(), tags):
        ax.text(
            0.03,
            0.94,
            tag,
            transform=ax.transAxes,
            ha="left",
            va="top",
            fontsize=13,
            bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.82, "pad": 1.4},
        )

    actual_handles = [Line2D([0], [0], color=color, lw=2.4, label=label) for label, _filename, color in ACTUAL_METHODS]
    projection_handles = [
        Line2D([0], [0], color=color, lw=2.4, label=label) for label, color in PROJECTION_METHODS
    ]
    leg1 = axes[0, 0].legend(handles=actual_handles, loc="lower left", frameon=False, fontsize=9.2)
    axes[0, 0].add_artist(leg1)
    axes[0, 1].legend(handles=projection_handles, loc="lower left", frameon=False, fontsize=9.2)

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

    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)

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


if __name__ == "__main__":
    main()
