#!/usr/bin/env python3
"""Compare soft-erosion theory distributions with empirical FA trajectories."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
from scipy import stats

SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

import downstream_capacity_sweep as dcs  # noqa: E402
from fa_tangent_kernel_capacity import flatten_grads, pseudo_jacobian  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--run-dir",
        type=Path,
        default=Path("outputs/phase_transition_dense_T30000_352traj"),
    )
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/soft_erosion_theory_vs_empirical"),
    )
    parser.add_argument("--theory-samples", type=int, default=20_000)
    parser.add_argument("--local-widths", type=int, nargs="+", default=[])
    parser.add_argument("--seed", type=int, default=0)
    return parser.parse_args()


def load_config(run_dir: Path) -> dcs.RunConfig:
    payload = json.loads((run_dir / "summary.json").read_text())
    raw = payload["config"]
    raw["outdir"] = str(run_dir)
    return dcs.RunConfig(**raw)


def loss_from_spectrum(
    eigenvalues: np.ndarray,
    coefficients: np.ndarray,
    rho: np.ndarray,
    lr: float,
    steps: int,
    samples: int,
) -> np.ndarray:
    tau = lr * steps / samples
    positive = np.clip(eigenvalues, 0.0, None)
    # Shape: theory_samples x modes.
    decay = np.exp(-2.0 * tau * rho[:, None] * positive[None, :])
    return 0.5 * (decay @ (coefficients * coefficients)) / samples


def bp_spectrum_for_width(
    config: dcs.RunConfig,
    x_train: torch.Tensor,
    y_train: torch.Tensor,
    width: int,
    init_seed: int,
) -> tuple[np.ndarray, np.ndarray]:
    weights = dcs.initialize_weights(config, width, init_seed)
    residual = (dcs.predict(weights, x_train) - y_train).reshape(-1).cpu().numpy()
    jacobian = pseudo_jacobian(weights, x_train, feedback=None).cpu().numpy()
    kernel = jacobian @ jacobian.T
    eigenvalues, eigenvectors = np.linalg.eigh(0.5 * (kernel + kernel.T))
    coefficients = eigenvectors.T @ residual
    return eigenvalues, coefficients


def theory_gap_distribution(
    config: dcs.RunConfig,
    summary: pd.Series,
    eigenvalues: np.ndarray,
    coefficients: np.ndarray,
    rng: np.random.Generator,
    samples: int,
) -> np.ndarray:
    p_count = int(summary["parameter_count"])
    k_count = int(summary["fa_constraint_rank"])
    if not 0 < k_count < p_count:
        raise ValueError(f"Invalid random-subspace dimensions: P={p_count}, k={k_count}")
    rho = rng.beta((p_count - k_count) / 2.0, k_count / 2.0, size=samples)
    fa_loss = loss_from_spectrum(
        eigenvalues,
        coefficients,
        rho,
        config.lr,
        config.steps,
        config.train_samples,
    )
    bp_loss = loss_from_spectrum(
        eigenvalues,
        coefficients,
        np.ones(1, dtype=np.float64),
        config.lr,
        config.steps,
        config.train_samples,
    )[0]
    return fa_loss - bp_loss


def summarize(values: np.ndarray) -> dict[str, float]:
    return {
        "mean": float(np.mean(values)),
        "q05": float(np.quantile(values, 0.05)),
        "q25": float(np.quantile(values, 0.25)),
        "q50": float(np.quantile(values, 0.50)),
        "q75": float(np.quantile(values, 0.75)),
        "q95": float(np.quantile(values, 0.95)),
    }


def flatten_gradient_list(grads: list[torch.Tensor]) -> np.ndarray:
    return flatten_grads(grads).detach().cpu().numpy()


def local_erosion_rows(
    config: dcs.RunConfig,
    x_train: torch.Tensor,
    y_train: torch.Tensor,
    runs: pd.DataFrame,
    widths: list[int],
) -> pd.DataFrame:
    records: list[dict[str, float | int]] = []
    for width in widths:
        width_runs = runs[(runs["width"] == width) & (runs["run_type"] == "fa")]
        if width_runs.empty:
            continue
        for init_seed, init_group in width_runs.groupby("init_seed"):
            weights = dcs.initialize_weights(config, width, int(init_seed))
            bp_grad = flatten_gradient_list(dcs.gradients(weights, x_train, y_train, None))
            denom = float(np.dot(bp_grad, bp_grad))
            for _, row in init_group.iterrows():
                feedback = dcs.init_feedback(config, width, int(row["feedback_seed"]))
                fa_grad = flatten_gradient_list(
                    dcs.gradients(weights, x_train, y_train, feedback)
                )
                ratio = float(np.dot(bp_grad, fa_grad) / denom)
                records.append(
                    {
                        "width": int(width),
                        "init_seed": int(init_seed),
                        "feedback_seed": int(row["feedback_seed"]),
                        "fa_capacity_margin": int(row["fa_capacity_margin"]),
                        "speed_ratio": ratio,
                        "erosion": 1.0 - ratio,
                    }
                )
    return pd.DataFrame(records)


def plot_gap_distributions(
    empirical: pd.DataFrame,
    theory: pd.DataFrame,
    outdir: Path,
) -> Path:
    fig, ax = plt.subplots(figsize=(9.2, 5.6), dpi=180)
    theory = theory.sort_values("fa_capacity_margin", ascending=False)
    empirical = empirical.sort_values("fa_capacity_margin", ascending=False)
    x_theory = theory["fa_capacity_margin"].to_numpy()
    x_emp = empirical["fa_capacity_margin"].to_numpy()

    ax.fill_between(
        x_theory,
        np.clip(theory["q05"], 1e-9, None),
        np.clip(theory["q95"], 1e-9, None),
        color="#174f91",
        alpha=0.16,
        linewidth=0,
        label="theory 5-95%",
    )
    ax.fill_between(
        x_theory,
        np.clip(theory["q25"], 1e-9, None),
        np.clip(theory["q75"], 1e-9, None),
        color="#174f91",
        alpha=0.28,
        linewidth=0,
        label="theory 25-75%",
    )
    ax.plot(
        x_theory,
        np.clip(theory["mean"], 1e-9, None),
        color="#174f91",
        marker="o",
        linewidth=2.2,
        label="theory mean",
    )
    ax.errorbar(
        x_emp,
        np.clip(empirical["mean"], 1e-9, None),
        yerr=[
            np.clip(empirical["mean"] - empirical["q05"], 0.0, None),
            np.clip(empirical["q95"] - empirical["mean"], 0.0, None),
        ],
        color="#c65f16",
        marker="o",
        linewidth=2.0,
        capsize=3,
        label="empirical 5-95%",
    )
    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
    ax.set_yscale("log")
    ax.set_xlim(max(x_emp.max(), x_theory.max()) + 20, min(x_emp.min(), x_theory.min()) - 20)
    ax.set_xlabel("hard FA capacity margin P - K_FA - N*out")
    ax.set_ylabel("FA train MSE - BP train MSE")
    ax.set_title("Soft-erosion random-subspace theory vs empirical FA gap")
    ax.grid(alpha=0.18, which="both")
    ax.legend(frameon=True, ncol=2)
    fig.tight_layout()
    path = outdir / "gap_distribution_theory_vs_empirical.png"
    fig.savefig(path)
    plt.close(fig)
    return path


def plot_local_erosion(
    local: pd.DataFrame,
    summary: pd.DataFrame,
    rng: np.random.Generator,
    theory_samples: int,
    outdir: Path,
) -> Path:
    widths = sorted(local["width"].unique())
    fig, axes = plt.subplots(1, len(widths), figsize=(5.2 * len(widths), 4.6), dpi=180)
    if len(widths) == 1:
        axes = [axes]
    for ax, width in zip(axes, widths, strict=True):
        values = local.loc[local["width"] == width, "erosion"].to_numpy(dtype=np.float64)
        row = summary.loc[summary["width"] == width].iloc[0]
        p_count = int(row["parameter_count"])
        k_count = int(row["fa_constraint_rank"])
        theory = rng.beta(k_count / 2.0, (p_count - k_count) / 2.0, size=theory_samples)
        bins = np.linspace(
            min(values.min(), theory.min(), 0.0),
            max(values.max(), theory.max(), 1.0),
            45,
        )
        ax.hist(
            values,
            bins=bins,
            density=True,
            alpha=0.46,
            color="#c65f16",
            label="empirical FA",
        )
        ax.hist(
            theory,
            bins=bins,
            density=True,
            histtype="step",
            linewidth=2.0,
            color="#174f91",
            label="Beta subspace",
        )
        ax.axvline(float(np.mean(values)), color="#c65f16", linestyle="--", linewidth=1.2)
        ax.axvline(float(np.mean(theory)), color="#174f91", linestyle="--", linewidth=1.2)
        ax.set_title(f"width={width}, margin={int(row['fa_capacity_margin'])}")
        ax.set_xlabel("one-step erosion 1 - ΔL_FA/ΔL_BP")
        ax.grid(alpha=0.16)
        ax.legend(frameon=True, fontsize=8)
    axes[0].set_ylabel("density")
    fig.suptitle("Local operator erosion: real FA vs random-subspace null")
    fig.tight_layout()
    path = outdir / "local_erosion_distribution_theory_vs_empirical.png"
    fig.savefig(path)
    plt.close(fig)
    return path


def plot_local_scalar_gap(
    empirical: pd.DataFrame,
    local_theory: pd.DataFrame,
    outdir: Path,
) -> Path:
    fig, ax = plt.subplots(figsize=(9.2, 5.6), dpi=180)
    theory = local_theory.sort_values("fa_capacity_margin", ascending=False)
    empirical = empirical.sort_values("fa_capacity_margin", ascending=False)
    x_theory = theory["fa_capacity_margin"].to_numpy()
    x_emp = empirical["fa_capacity_margin"].to_numpy()
    ax.fill_between(
        x_theory,
        np.clip(theory["q05"], 1e-9, None),
        np.clip(theory["q95"], 1e-9, None),
        color="#2b8a3e",
        alpha=0.16,
        linewidth=0,
        label="operator theory 5-95%",
    )
    ax.fill_between(
        x_theory,
        np.clip(theory["q25"], 1e-9, None),
        np.clip(theory["q75"], 1e-9, None),
        color="#2b8a3e",
        alpha=0.28,
        linewidth=0,
        label="operator theory 25-75%",
    )
    ax.plot(
        x_theory,
        np.clip(theory["mean"], 1e-9, None),
        color="#2b8a3e",
        marker="o",
        linewidth=2.2,
        label="operator theory mean",
    )
    ax.errorbar(
        x_emp,
        np.clip(empirical["mean"], 1e-9, None),
        yerr=[
            np.clip(empirical["mean"] - empirical["q05"], 0.0, None),
            np.clip(empirical["q95"] - empirical["mean"], 0.0, None),
        ],
        color="#c65f16",
        marker="o",
        linewidth=2.0,
        capsize=3,
        label="empirical 5-95%",
    )
    ax.axvline(0, color="black", linestyle="--", linewidth=1.1)
    ax.set_yscale("log")
    ax.set_xlim(max(x_emp.max(), x_theory.max()) + 20, min(x_emp.min(), x_theory.min()) - 20)
    ax.set_xlabel("hard FA capacity margin P - K_FA - N*out")
    ax.set_ylabel("FA train MSE - BP train MSE")
    ax.set_title("Local-operator scalar theory vs empirical FA gap")
    ax.grid(alpha=0.18, which="both")
    ax.legend(frameon=True, ncol=2)
    fig.tight_layout()
    path = outdir / "local_operator_scalar_gap_theory_vs_empirical.png"
    fig.savefig(path)
    plt.close(fig)
    return path


def main() -> None:
    args = parse_args()
    if args.theory_samples < 1:
        raise ValueError("--theory-samples must be positive")
    rng = np.random.default_rng(args.seed)
    outdir = args.outdir
    outdir.mkdir(parents=True, exist_ok=True)

    config = load_config(args.run_dir)
    x_train, y_train, _x_test, _y_test, _x_probe, _teacher = dcs.make_data(config)
    runs = pd.read_csv(args.run_dir / "runs.csv")
    summary = pd.read_csv(args.run_dir / "width_summary.csv")

    theory_records = []
    empirical_records = []
    spectra: dict[int, tuple[np.ndarray, np.ndarray]] = {}
    for _, width_row in summary.sort_values("width").iterrows():
        width = int(width_row["width"])
        init_seed = int(
            runs.loc[(runs["width"] == width) & (runs["run_type"] == "bp"), "init_seed"].iloc[0]
        )
        print(f"width={width}: computing BP spectrum", flush=True)
        eigenvalues, coefficients = bp_spectrum_for_width(
            config, x_train, y_train, width, init_seed
        )
        spectra[width] = (eigenvalues, coefficients)
        theory_gaps = theory_gap_distribution(
            config,
            width_row,
            eigenvalues,
            coefficients,
            rng,
            args.theory_samples,
        )
        t_summary = summarize(np.clip(theory_gaps, 1e-12, None))
        t_summary.update(
            {
                "width": width,
                "fa_capacity_margin": int(width_row["fa_capacity_margin"]),
                "parameter_count": int(width_row["parameter_count"]),
                "fa_constraint_rank": int(width_row["fa_constraint_rank"]),
            }
        )
        theory_records.append(t_summary)

        emp_values = runs.loc[
            (runs["width"] == width) & (runs["run_type"] == "fa"),
            "train_gap_to_bp",
        ].to_numpy(dtype=np.float64)
        e_summary = summarize(np.clip(emp_values, 1e-12, None))
        e_summary.update(
            {
                "width": width,
                "fa_capacity_margin": int(width_row["fa_capacity_margin"]),
                "parameter_count": int(width_row["parameter_count"]),
                "fa_constraint_rank": int(width_row["fa_constraint_rank"]),
            }
        )
        empirical_records.append(e_summary)

    theory_df = pd.DataFrame(theory_records)
    empirical_df = pd.DataFrame(empirical_records)
    theory_csv = outdir / "theory_gap_summary.csv"
    empirical_csv = outdir / "empirical_gap_summary.csv"
    theory_df.to_csv(theory_csv, index=False)
    empirical_df.to_csv(empirical_csv, index=False)

    local_widths = args.local_widths or sorted(summary["width"].astype(int).tolist())
    local = local_erosion_rows(config, x_train, y_train, runs, local_widths)
    local_csv = outdir / "local_erosion_rows.csv"
    local.to_csv(local_csv, index=False)

    local_gap_records = []
    for width, group in local.groupby("width"):
        eigenvalues, coefficients = spectra[int(width)]
        bp_loss = loss_from_spectrum(
            eigenvalues,
            coefficients,
            np.ones(1, dtype=np.float64),
            config.lr,
            config.steps,
            config.train_samples,
        )[0]
        rho = group["speed_ratio"].to_numpy(dtype=np.float64)
        # Negative local speed ratios are valid predictions of local instability,
        # but the scalar frozen-kernel rollout can explode. Keep the raw value in
        # local_erosion_rows.csv and clip only for this conservative distribution
        # plot.
        rho_for_rollout = np.clip(rho, 0.0, 2.0)
        predicted_loss = loss_from_spectrum(
            eigenvalues,
            coefficients,
            rho_for_rollout,
            config.lr,
            config.steps,
            config.train_samples,
        )
        predicted_gap = predicted_loss - bp_loss
        row = summary.loc[summary["width"] == int(width)].iloc[0]
        local_summary = summarize(np.clip(predicted_gap, 1e-12, None))
        local_summary.update(
            {
                "width": int(width),
                "fa_capacity_margin": int(row["fa_capacity_margin"]),
                "parameter_count": int(row["parameter_count"]),
                "fa_constraint_rank": int(row["fa_constraint_rank"]),
            }
        )
        local_gap_records.append(local_summary)
    local_gap_df = pd.DataFrame(local_gap_records)
    local_gap_csv = outdir / "local_operator_scalar_gap_summary.csv"
    local_gap_df.to_csv(local_gap_csv, index=False)

    paths = [
        plot_gap_distributions(empirical_df, theory_df, outdir),
        plot_local_erosion(local, summary, rng, args.theory_samples, outdir),
        plot_local_scalar_gap(empirical_df, local_gap_df, outdir),
    ]
    print(f"theory summary: {theory_csv}")
    print(f"empirical summary: {empirical_csv}")
    print(f"local erosion rows: {local_csv}")
    print(f"local scalar gap summary: {local_gap_csv}")
    for path in paths:
        print(f"plot: {path}")


if __name__ == "__main__":
    main()