#!/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:
    source_id: int
    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("--extra-run-dir", type=Path, action="append", default=[])
    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(
        "--effective-task-dim",
        type=float,
        default=None,
        help="Use this task dimension in the smooth theory while keeping the hard-margin x-axis unchanged.",
    )
    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, source_id: int) -> list[RunRow]:
    rows: list[RunRow] = []
    with path.open() as handle:
        for row in csv.DictReader(handle):
            rows.append(
                RunRow(
                    source_id=source_id,
                    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,
    hard_task_dim: float,
    effective_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 - hard_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(effective_task_dim - p_count)
        task_subspace_dim = min(effective_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(
            effective_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,
    extra_run_dirs: list[Path],
    samples: int,
    transition_width: float,
    effective_task_dim_arg: float | None,
    outname: str,
    xmin: float,
    xmax: float,
) -> Path:
    summary = read_summary(run_dir / "width_summary.csv")
    runs = read_runs(run_dir / "runs.csv", 0)
    for source_id, extra_run_dir in enumerate(extra_run_dirs, start=1):
        runs.extend(read_runs(extra_run_dir / "runs.csv", source_id))
    input_dim, output_dim, hard_task_dim = infer_dims(run_dir)
    effective_task_dim = (
        hard_task_dim if effective_task_dim_arg is None else effective_task_dim_arg
    )
    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,
        hard_task_dim,
        effective_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,
        hard_task_dim,
        effective_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

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

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

    axes[0].fill_between(
        dense_margin,
        dense_lo,
        dense_hi,
        color="tab:orange",
        alpha=0.22,
        linewidth=0,
        label="theory 10-90% band",
    )
    axes[0].plot(
        dense_margin,
        dense_mean,
        color="black",
        linewidth=2.4,
        label="theory mean",
    )
    axes[0].axvline(0.0, color="black", linestyle="--", linewidth=1.1, alpha=0.8)

    for axis in axes:
        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)

    for _seed, rows in sorted(by_seed.items()):
        ordered = sorted(rows, key=lambda item: item.fa_capacity_margin)
        axes[1].plot(
            [row.fa_capacity_margin for row in ordered],
            [row.train_gap_to_bp for row in ordered],
            color="tab:blue",
            alpha=0.07,
            linewidth=0.85,
        )

    axes[1].text(
        0.98,
        0.96,
        f"{len(by_seed)} empirical FA trajectories",
        transform=axes[1].transAxes,
        ha="right",
        va="top",
        fontsize=10,
        bbox={
            "boxstyle": "round,pad=0.25",
            "facecolor": "white",
            "edgecolor": "#cccccc",
            "alpha": 0.85,
        },
    )

    axes[0].set_title("P1: smooth capacity-transition prediction")
    axes[1].set_title("P2: empirical trajectory overlay")
    axes[0].set_ylabel("FA train MSE - BP train MSE")
    axes[0].legend(loc="upper right")
    fig.suptitle(
        f"Random-label memorization transition, scale={scale:.4g}, transition width={transition_width:g}, d_eff={effective_task_dim: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.extra_run_dir,
        args.samples,
        args.transition_width,
        args.effective_task_dim,
        args.outname,
        args.xmin,
        args.xmax,
    )
    print(outpath)


if __name__ == "__main__":
    main()