#!/usr/bin/env python3
"""Plot smooth capacity-transition prediction against FA trajectory ensembles."""

from __future__ import annotations

import argparse
import csv
import json
import math
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np


@dataclass(frozen=True)
class SummaryRow:
    width: float
    parameter_count: float
    task_dimension: float
    fa_constraint_rank: float
    bp_capacity_margin: float
    fa_capacity_margin: float
    train_gap_mean: float


@dataclass(frozen=True)
class RunRow:
    width: float
    feedback_seed: int
    run_type: str
    train_gap_to_bp: float
    fa_capacity_margin: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--run-dir",
        type=Path,
        default=Path("outputs/downstream_capacity_random_main_fast"),
    )
    parser.add_argument("--samples", type=int, default=2048)
    parser.add_argument(
        "--transition-width",
        type=float,
        default=70.0,
        help="Soft capacity-transition width measured in scalar output constraints.",
    )
    parser.add_argument("--xmin", type=float, default=-430.0)
    parser.add_argument("--xmax", type=float, default=600.0)
    parser.add_argument("--outname", default="capacity_transition_theory_vs_trajectories.png")
    return parser.parse_args()


def read_summary(path: Path) -> list[SummaryRow]:
    rows: list[SummaryRow] = []
    with path.open() as handle:
        for row in csv.DictReader(handle):
            rows.append(
                SummaryRow(
                    width=float(row["width"]),
                    parameter_count=float(row["parameter_count"]),
                    task_dimension=float(row["task_dimension"]),
                    fa_constraint_rank=float(row["fa_constraint_rank"]),
                    bp_capacity_margin=float(row["bp_capacity_margin"]),
                    fa_capacity_margin=float(row["fa_capacity_margin"]),
                    train_gap_mean=float(row["train_gap_mean"]),
                )
            )
    return rows


def read_runs(path: Path) -> list[RunRow]:
    rows: list[RunRow] = []
    with path.open() as handle:
        for row in csv.DictReader(handle):
            rows.append(
                RunRow(
                    width=float(row["width"]),
                    feedback_seed=int(row["feedback_seed"]),
                    run_type=row["run_type"],
                    train_gap_to_bp=float(row["train_gap_to_bp"]),
                    fa_capacity_margin=float(row["fa_capacity_margin"]),
                )
            )
    return rows


def soft_overlap_gap_counts(
    width: np.ndarray,
    input_dim: float,
    output_dim: float,
    task_dim: float,
    samples: int,
    transition_width: float,
    rng: np.random.Generator,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    parameter_count = input_dim * width + width * width + width * output_dim
    constraint_rank = np.maximum(width * width - 1.0, 0.0) + np.maximum(
        width * output_dim - 1.0, 0.0
    )
    hard_fa_margin = parameter_count - constraint_rank - task_dim

    raw_gaps = np.zeros((samples, width.size), dtype=np.float64)

    def softplus_hinge(value: np.ndarray | float) -> np.ndarray:
        scaled = np.asarray(value, dtype=np.float64) / transition_width
        return transition_width * np.logaddexp(0.0, scaled)

    for index, (p_count, k_rank) in enumerate(zip(parameter_count, constraint_rank)):
        bp_missing = softplus_hinge(task_dim - p_count)
        task_subspace_dim = min(task_dim, p_count)

        if task_subspace_dim >= p_count or k_rank <= 0:
            overlap = np.full(samples, min(k_rank, p_count), dtype=np.float64)
        else:
            mean = k_rank * task_subspace_dim / p_count
            var = (
                2.0
                * k_rank
                * task_subspace_dim
                * (p_count - k_rank)
                * (p_count - task_subspace_dim)
                / (p_count * p_count * (p_count - 1.0) * (p_count + 2.0))
            )
            overlap = rng.normal(mean, math.sqrt(max(var, 0.0)), size=samples)
            overlap = np.clip(overlap, 0.0, min(k_rank, task_subspace_dim))

        finite_width_noise = rng.normal(0.0, transition_width, size=samples)
        fa_missing = softplus_hinge(
            task_dim - (parameter_count[index] - overlap) + finite_width_noise
        )
        raw_gaps[:, index] = np.maximum(fa_missing - bp_missing, 0.0)

    mean = raw_gaps.mean(axis=0)
    lo = np.quantile(raw_gaps, 0.1, axis=0)
    hi = np.quantile(raw_gaps, 0.9, axis=0)
    return hard_fa_margin, mean, lo, hi


def infer_dims(run_dir: Path) -> tuple[float, float, float]:
    payload = json.loads((run_dir / "summary.json").read_text())
    config = payload["config"]
    return (
        float(config["input_dim"]),
        float(config["output_dim"]),
        float(config["train_samples"] * config["output_dim"]),
    )


def calibrate_scale(summary: list[SummaryRow], raw_mean_at_widths: np.ndarray) -> float:
    observed = np.array([row.train_gap_mean for row in summary], dtype=np.float64)
    denom = float(np.dot(raw_mean_at_widths, raw_mean_at_widths))
    if denom <= 0:
        return 1.0
    return float(np.dot(raw_mean_at_widths, observed) / denom)


def plot_transition(
    run_dir: Path,
    samples: int,
    transition_width: float,
    outname: str,
    xmin: float,
    xmax: float,
) -> Path:
    summary = read_summary(run_dir / "width_summary.csv")
    runs = read_runs(run_dir / "runs.csv")
    input_dim, output_dim, task_dim = infer_dims(run_dir)
    rng = np.random.default_rng(0)

    widths = np.array([row.width for row in summary], dtype=np.float64)
    dense_width = np.linspace(widths.min(), widths.max(), 240)

    _, raw_mean_at_widths, _, _ = soft_overlap_gap_counts(
        widths, input_dim, output_dim, task_dim, samples, transition_width, rng
    )
    scale = calibrate_scale(summary, raw_mean_at_widths)

    dense_margin, dense_mean_raw, dense_lo_raw, dense_hi_raw = soft_overlap_gap_counts(
        dense_width,
        input_dim,
        output_dim,
        task_dim,
        samples,
        transition_width,
        np.random.default_rng(1),
    )
    dense_mean = scale * dense_mean_raw
    dense_lo = scale * dense_lo_raw
    dense_hi = scale * dense_hi_raw

    summary_margin = np.array([row.fa_capacity_margin for row in summary], dtype=np.float64)
    summary_gap = np.array([row.train_gap_mean for row in summary], dtype=np.float64)

    by_seed: dict[int, list[RunRow]] = defaultdict(list)
    for row in runs:
        if row.run_type == "fa":
            by_seed[row.feedback_seed].append(row)

    fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2), sharex=True, sharey=True)

    for axis, show_trajectories in zip(axes, [False, True]):
        axis.fill_between(
            dense_margin,
            dense_lo,
            dense_hi,
            color="tab:orange",
            alpha=0.22,
            linewidth=0,
            label="theory 10-90% band",
        )
        axis.plot(
            dense_margin,
            dense_mean,
            color="black",
            linewidth=2.4,
            label="theory mean",
        )
        axis.axvline(0.0, color="black", linestyle="--", linewidth=1.1, alpha=0.8)
        axis.axhline(0.0, color="black", linewidth=0.9, alpha=0.75)
        axis.set_xlim(xmin, xmax)
        axis.set_xlabel("hard FA capacity margin  P - K_FA - N*out")
        axis.grid(True, color="#dddddd", linewidth=0.8, alpha=0.7)

        if show_trajectories:
            for seed, rows in sorted(by_seed.items()):
                ordered = sorted(rows, key=lambda item: item.fa_capacity_margin)
                axis.plot(
                    [row.fa_capacity_margin for row in ordered],
                    [row.train_gap_to_bp for row in ordered],
                    color="tab:blue",
                    alpha=0.12,
                    linewidth=0.9,
                )
            axis.plot(
                summary_margin,
                summary_gap,
                color="tab:blue",
                linewidth=2.2,
                alpha=0.95,
                label="empirical mean",
            )
            axis.scatter(
                summary_margin,
                summary_gap,
                color="black",
                s=34,
                zorder=4,
                label="empirical points",
            )

    axes[0].set_title("P1: smooth capacity-transition prediction")
    axes[1].set_title("P2: transparent trajectory ensemble overlay")
    axes[0].set_ylabel("FA train MSE - BP train MSE")
    axes[0].legend(loc="upper right")
    axes[1].legend(loc="upper right")
    fig.suptitle(
        f"Random-label memorization transition, scale={scale:.4g}, transition width={transition_width:g}",
        y=1.02,
    )
    fig.tight_layout()

    outpath = run_dir / outname
    fig.savefig(outpath, dpi=220, bbox_inches="tight")
    plt.close(fig)
    return outpath


def main() -> None:
    args = parse_args()
    outpath = plot_transition(
        args.run_dir,
        args.samples,
        args.transition_width,
        args.outname,
        args.xmin,
        args.xmax,
    )
    print(outpath)


if __name__ == "__main__":
    main()