#!/usr/bin/env python3
"""Real-data (MNIST) validation of the two core no-fit claims.

Part A — initial-erosion theorem. For any MLP trained with squared loss and
independent zero-mean feedback, E_B[e_0 | W, r] = hidden BP speed share
(notes 29/35). The theorem is architecture- and data-agnostic; this script
checks it on MNIST-subset MLPs by running the *real FA backward pass* for
every feedback draw.

Part B — early operator-velocity estimator. The no-fit finite-time predictor
K_hat_t = K_0 + (t/s)(K_s - K_0), rolled through
r_{t+1} = (I - eta K_hat_t / N) r_t, predicts the T-step FA/BP train gap
(notes 15-17). This script repeats the synthetic protocol on MNIST subsets:
empirical full-batch SGD trajectories vs fixed-K(0) and linear-velocity
predictions, across subset sizes and feedback seeds.

Tangent operators are computed with a layerwise gram identity
    K[(n,c),(m,c')] = sum_l <delta_l(n,c), delta'_l(m,c')> <h_{l-1}(n), h_{l-1}(m)>
(delta' = BP deltas for K_BP, FA deltas for K_FA = J J_tilde^T), which avoids
materializing the (N C) x P Jacobian; `--self-test` verifies it against an
explicit autograd Jacobian on a tiny network.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import sys
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch

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

Tensor = torch.Tensor


# ---------------------------------------------------------------------------
# data and model utilities
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="MNIST validation of e0 theorem and estimator.")
    p.add_argument("--part", choices=["a", "b", "both"], default="both")
    p.add_argument("--self-test", action="store_true", help="verify kernel code and exit")
    p.add_argument("--data-root", type=Path, default=Path("data"))
    # part A
    p.add_argument("--e0-widths", type=int, nargs="+", default=[64, 256])
    p.add_argument("--e0-depths", type=int, nargs="+", default=[1, 2])
    p.add_argument("--e0-init-seeds", type=int, default=4)
    p.add_argument("--e0-train-samples", type=int, default=256)
    p.add_argument("--feedback-draws", type=int, default=512)
    # part B
    p.add_argument("--est-widths", type=int, nargs="+", default=[128, 256])
    p.add_argument("--est-depth", type=int, default=2)
    p.add_argument("--est-train-sizes", type=int, nargs="+", default=[128, 256, 512])
    p.add_argument("--est-init-seeds", type=int, default=2)
    p.add_argument("--est-feedback-seeds", type=int, default=8)
    p.add_argument("--horizon", type=int, default=50)
    p.add_argument("--early-step", type=int, default=20)
    p.add_argument("--lr-safety", type=float, default=0.3,
                   help="lr = safety * N / lam_max(K_BP_0)")
    p.add_argument("--feedback-scale", choices=["relu", "fan-in", "unit"], default="relu")
    p.add_argument("--torch-threads", type=int, default=16)
    p.add_argument("--outdir", type=Path, default=Path("outputs/real_data_validation_mnist"))
    return p.parse_args()


def load_mnist_subset(root: Path, samples: int, seed: int = 0) -> tuple[Tensor, Tensor]:
    """Flattened, standardized MNIST images with one-hot regression targets."""
    from torchvision import datasets

    train = datasets.MNIST(root=str(root), train=True, download=True)
    images = train.data.to(torch.float64) / 255.0
    labels = train.targets
    generator = torch.Generator().manual_seed(seed)
    index = torch.randperm(images.shape[0], generator=generator)[:samples]
    x = images[index].reshape(samples, -1)
    x = (x - 0.1307) / 0.3081
    x = x / math.sqrt(x.shape[1])  # unit-scale rows so kernel scales match synthetic runs
    y = torch.zeros(samples, 10, dtype=torch.float64)
    y[torch.arange(samples), labels[index]] = 1.0
    return x, y


def init_weights_dims(dims: list[int], seed: int) -> list[Tensor]:
    """He init for hidden layers, 1/sqrt(fan_in) for output (project convention)."""
    generator = torch.Generator(device="cpu")
    generator.manual_seed(seed)
    weights: list[Tensor] = []
    for layer, (fan_in, fan_out) in enumerate(zip(dims[:-1], dims[1:])):
        scale = math.sqrt(2.0 / fan_in) if layer < len(dims) - 2 else 1.0 / math.sqrt(fan_in)
        weights.append(torch.randn(fan_out, fan_in, generator=generator,
                                   dtype=torch.float64) * scale)
    return weights


def init_feedback_dims(dims: list[int], seed: int, mode: str) -> list[Tensor]:
    """feedback[l] has shape (dims[l+1], dims[l+2]) and replaces W_{l+1}^T."""
    generator = torch.Generator(device="cpu")
    generator.manual_seed(seed)
    feedback: list[Tensor] = []
    for l in range(len(dims) - 2):
        rows, cols = dims[l + 1], dims[l + 2]
        feedback.append(torch.randn(rows, cols, generator=generator, dtype=torch.float64)
                        * dcs.feedback_scale(rows, mode))
    return feedback


def speeds(weights: list[Tensor], x: Tensor, y: Tensor,
           feedback: list[Tensor] | None) -> tuple[float, float, float]:
    """(speed_BP, speed_FA, output_speed); speed_FA = sum_l <g_l^BP, g_l^FA>."""
    bp = dcs.gradients(weights, x, y, None)
    speed_bp = float(sum((g * g).sum() for g in bp))
    output_speed = float((bp[-1] * bp[-1]).sum())
    if feedback is None:
        return speed_bp, speed_bp, output_speed
    fa = dcs.gradients(weights, x, y, feedback)
    speed_fa = float(sum((gb * gf).sum() for gb, gf in zip(bp, fa)))
    return speed_bp, speed_fa, output_speed


# ---------------------------------------------------------------------------
# tangent operators via the layerwise gram identity
# ---------------------------------------------------------------------------

def collect_states(weights: list[Tensor], x: Tensor) -> tuple[list[Tensor], list[Tensor]]:
    activations, preacts = dcs.forward(weights, x)
    gates = [(preacts[l] > 0).to(torch.float64) for l in range(len(weights) - 1)]
    return activations, gates


def output_deltas(weights: list[Tensor], gates: list[Tensor],
                  feedback: list[Tensor] | None) -> list[Tensor]:
    """D[l]: (N, C, dims[l+1]) per-output-coordinate deltas at layer l."""
    n = gates[0].shape[0]
    c = weights[-1].shape[0]
    layers = len(weights)
    deltas: list[Tensor] = [torch.empty(0)] * layers
    deltas[layers - 1] = torch.eye(c, dtype=torch.float64).unsqueeze(0).expand(n, c, c).clone()
    for l in range(layers - 2, -1, -1):
        if feedback is None:
            back = deltas[l + 1] @ weights[l + 1]
        else:
            back = deltas[l + 1] @ feedback[l].T
        deltas[l] = back * gates[l].unsqueeze(1)
    return deltas


def tangent_kernel(weights: list[Tensor], x: Tensor,
                   feedback: list[Tensor] | None) -> Tensor:
    """K_BP = J J^T (feedback=None) or K_FA = J J_tilde^T, shape (NC, NC)."""
    activations, gates = collect_states(weights, x)
    d_bp = output_deltas(weights, gates, None)
    d_rule = d_bp if feedback is None else output_deltas(weights, gates, feedback)
    n, c = x.shape[0], weights[-1].shape[0]
    kernel = torch.zeros(n * c, n * c, dtype=torch.float64)
    for l in range(len(weights)):
        left = d_bp[l].reshape(n * c, -1)
        right = d_rule[l].reshape(n * c, -1)
        gram = activations[l] @ activations[l].T
        gram_x = gram.repeat_interleave(c, dim=0).repeat_interleave(c, dim=1)
        kernel += (left @ right.T) * gram_x
    return kernel


def explicit_jacobians(weights: list[Tensor], x: Tensor,
                       feedback: list[Tensor] | None) -> Tensor:
    """Explicit (NC, P) pseudo-Jacobian via per-coordinate deltas (reference)."""
    activations, gates = collect_states(weights, x)
    deltas = output_deltas(weights, gates, feedback)
    rows = []
    n, c = x.shape[0], weights[-1].shape[0]
    for sample in range(n):
        for coord in range(c):
            parts = [torch.outer(deltas[l][sample, coord], activations[l][sample]).reshape(-1)
                     for l in range(len(weights))]
            rows.append(torch.cat(parts))
    return torch.stack(rows, dim=0)


def autograd_jacobian(weights: list[Tensor], x: Tensor) -> Tensor:
    leaves = [w.clone().requires_grad_(True) for w in weights]
    pred = dcs.predict(leaves, x).reshape(-1)
    rows = []
    for i in range(pred.numel()):
        grads = torch.autograd.grad(pred[i], leaves, retain_graph=True)
        rows.append(torch.cat([g.reshape(-1) for g in grads]))
    return torch.stack(rows, dim=0)


def self_test() -> None:
    torch.manual_seed(0)
    dims = [5, 7, 6, 3]
    weights = init_weights_dims(dims, seed=11)
    feedback = init_feedback_dims(dims, seed=21, mode="relu")
    x = torch.randn(4, dims[0], dtype=torch.float64)

    jac_auto = autograd_jacobian(weights, x)
    jac_manual = explicit_jacobians(weights, x, None)
    err_j = float((jac_auto - jac_manual).abs().max())

    k_bp_gram = tangent_kernel(weights, x, None)
    k_bp_ref = jac_auto @ jac_auto.T
    err_bp = float((k_bp_gram - k_bp_ref).abs().max())

    jac_fa = explicit_jacobians(weights, x, feedback)
    k_fa_gram = tangent_kernel(weights, x, feedback)
    k_fa_ref = jac_auto @ jac_fa.T
    err_fa = float((k_fa_gram - k_fa_ref).abs().max())

    # gradient consistency: J^T r and J_tilde^T r reproduce dcs.gradients
    y = torch.randn(4, dims[-1], dtype=torch.float64)
    r = (dcs.predict(weights, x) - y).reshape(-1)
    batch = x.shape[0]
    g_bp_vec = jac_auto.T @ (r / batch)
    g_bp = torch.cat([g.reshape(-1) for g in dcs.gradients(weights, x, y, None)])
    err_g_bp = float((g_bp_vec - g_bp).abs().max())
    g_fa_vec = jac_fa.T @ (r / batch)
    g_fa = torch.cat([g.reshape(-1) for g in dcs.gradients(weights, x, y, feedback)])
    err_g_fa = float((g_fa_vec - g_fa).abs().max())

    print(f"self-test: |J_auto - J_manual|_max = {err_j:.3e}")
    print(f"self-test: |K_BP_gram - J J^T|_max = {err_bp:.3e}")
    print(f"self-test: |K_FA_gram - J J_tilde^T|_max = {err_fa:.3e}")
    print(f"self-test: |J^T r - g_BP|_max = {err_g_bp:.3e}")
    print(f"self-test: |J_tilde^T r - g_FA|_max = {err_g_fa:.3e}")
    tol = 1e-9
    assert max(err_j, err_bp, err_fa, err_g_bp, err_g_fa) < tol, "self-test FAILED"
    print("self-test PASSED")


# ---------------------------------------------------------------------------
# part A: e0 theorem on MNIST
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class E0Row:
    width: int
    depth: int
    init_seed: int
    feedback_draws: int
    hidden_share: float
    empirical_mean: float
    empirical_stderr: float
    empirical_std: float
    mean_abs_error: float


def run_part_a(args: argparse.Namespace, outdir: Path) -> list[E0Row]:
    x, y = load_mnist_subset(args.data_root, args.e0_train_samples)
    rows: list[E0Row] = []
    for depth in args.e0_depths:
        for width in args.e0_widths:
            dims = [x.shape[1], *([width] * depth), y.shape[1]]
            for init_seed in range(args.e0_init_seeds):
                weights = init_weights_dims(dims, seed=10_000 + init_seed)
                speed_bp, _, output_speed = speeds(weights, x, y, None)
                hidden_share = 1.0 - output_speed / speed_bp
                erosions = []
                for draw in range(args.feedback_draws):
                    feedback = init_feedback_dims(
                        dims, seed=500_000 + 1_000 * init_seed + draw,
                        mode=args.feedback_scale)
                    _, speed_fa, _ = speeds(weights, x, y, feedback)
                    erosions.append(1.0 - speed_fa / speed_bp)
                erosions_arr = np.array(erosions)
                row = E0Row(
                    width=width, depth=depth, init_seed=init_seed,
                    feedback_draws=args.feedback_draws,
                    hidden_share=hidden_share,
                    empirical_mean=float(erosions_arr.mean()),
                    empirical_stderr=float(erosions_arr.std(ddof=1) / math.sqrt(len(erosions))),
                    empirical_std=float(erosions_arr.std(ddof=1)),
                    mean_abs_error=float(abs(erosions_arr.mean() - hidden_share)),
                )
                rows.append(row)
                print(f"[A] depth={depth} width={width} init={init_seed}: "
                      f"hidden_share={row.hidden_share:.4f} "
                      f"E_B[e0]={row.empirical_mean:.4f} +/- {row.empirical_stderr:.4f} "
                      f"|diff|={row.mean_abs_error:.4f}", flush=True)

    with (outdir / "e0_mnist_rows.csv").open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=list(E0Row.__annotations__.keys()))
        writer.writeheader()
        for row in rows:
            writer.writerow(asdict(row))

    fig, ax = plt.subplots(figsize=(5.6, 5.2), dpi=180)
    markers = {1: "o", 2: "s", 3: "^"}
    for depth in sorted({row.depth for row in rows}):
        sub = [row for row in rows if row.depth == depth]
        ax.errorbar([row.hidden_share for row in sub],
                    [row.empirical_mean for row in sub],
                    yerr=[2 * row.empirical_stderr for row in sub],
                    fmt=markers.get(depth, "o"), ms=5.5, capsize=2.0, lw=0, elinewidth=1.0,
                    label=f"{depth} hidden layer{'s' if depth > 1 else ''}")
    lo = min(min(row.hidden_share for row in rows), min(row.empirical_mean for row in rows))
    hi = max(max(row.hidden_share for row in rows), max(row.empirical_mean for row in rows))
    pad = 0.04 * (hi - lo + 1e-9)
    ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], color="black", lw=1.0)
    ax.set_xlabel("predicted: hidden BP speed share (one BP pass, no fit)")
    ax.set_ylabel("measured: E_B[e0] over real FA backward draws")
    ax.set_title("MNIST: initial FA erosion = hidden speed share")
    ax.legend(fontsize=8)
    ax.grid(alpha=0.16)
    fig.tight_layout()
    fig.savefig(outdir / "e0_mnist_calibration.png", bbox_inches="tight")
    plt.close(fig)
    return rows


# ---------------------------------------------------------------------------
# part B: early operator-velocity estimator on MNIST
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class EstimatorRow:
    width: int
    depth: int
    train_samples: int
    init_seed: int
    feedback_seed: int
    lr: float
    lam_max: float
    horizon: int
    early_step: int
    hidden_share: float
    loss_bp_emp: float
    loss_fa_emp: float
    gap_emp: float
    gap_fixed: float
    gap_velocity: float


def power_lam_max(kernel: Tensor, iters: int = 60) -> float:
    vec = torch.randn(kernel.shape[0], dtype=torch.float64)
    vec /= vec.norm()
    lam = 0.0
    for _ in range(iters):
        nxt = kernel @ vec
        lam = float(nxt.norm())
        vec = nxt / (lam + 1e-30)
    return lam


def train_with_snapshots(weights0: list[Tensor], x: Tensor, y: Tensor, lr: float,
                         steps: int, feedback: list[Tensor] | None,
                         snap_steps: set[int]) -> tuple[list[Tensor], dict[int, list[Tensor]]]:
    weights = dcs.clone_weights(weights0)
    snaps: dict[int, list[Tensor]] = {}
    for t in range(steps + 1):
        if t in snap_steps:
            snaps[t] = dcs.clone_weights(weights)
        if t == steps:
            break
        grads = dcs.gradients(weights, x, y, feedback)
        weights = [w - lr * g for w, g in zip(weights, grads)]
    return weights, snaps


def rollout_loss(k_zero: Tensor, k_slope: Tensor | None, r0: Tensor, lr: float,
                 n: int, steps: int, early_step: int) -> float:
    """Roll r_{t+1} = r_t - (lr/N) K_hat_t r_t with K_hat_t = K0 + (t/s) dK."""
    r = r0.clone()
    for t in range(steps):
        kr = k_zero @ r
        if k_slope is not None:
            kr = kr + (t / early_step) * (k_slope @ r)
        r = r - (lr / n) * kr
    return float(r.pow(2).sum() / (2 * n))


def run_part_b(args: argparse.Namespace, outdir: Path) -> list[EstimatorRow]:
    rows: list[EstimatorRow] = []
    t_total, s_step = args.horizon, args.early_step
    for est_width, n_train in [(w, n) for w in args.est_widths
                               for n in args.est_train_sizes]:
        x, y = load_mnist_subset(args.data_root, n_train)
        dims = [x.shape[1], *([est_width] * args.est_depth), y.shape[1]]
        n = x.shape[0]
        for init_seed in range(args.est_init_seeds):
            weights0 = init_weights_dims(dims, seed=10_000 + init_seed)
            r0 = (dcs.predict(weights0, x) - y).reshape(-1)
            k_bp_0 = tangent_kernel(weights0, x, None)
            lam_max = power_lam_max(k_bp_0)
            lr = args.lr_safety * n / lam_max
            speed_bp, _, output_speed = speeds(weights0, x, y, None)
            hidden_share = 1.0 - output_speed / speed_bp

            t0 = time.time()
            bp_final, bp_snaps = train_with_snapshots(
                weights0, x, y, lr, t_total, None, {0, s_step})
            loss_bp_emp = dcs.mse(bp_final, x, y)
            k_bp_s = tangent_kernel(bp_snaps[s_step], x, None)
            bp_loss_fixed = rollout_loss(k_bp_0, None, r0, lr, n, t_total, s_step)
            bp_loss_vel = rollout_loss(k_bp_0, k_bp_s - k_bp_0, r0, lr, n, t_total, s_step)

            for feedback_index in range(args.est_feedback_seeds):
                feedback_seed = 100_000 + init_seed * 1_000 + feedback_index
                feedback = init_feedback_dims(dims, feedback_seed, args.feedback_scale)
                fa_final, fa_snaps = train_with_snapshots(
                    weights0, x, y, lr, t_total, feedback, {0, s_step})
                loss_fa_emp = dcs.mse(fa_final, x, y)
                k_fa_0 = tangent_kernel(weights0, x, feedback)
                k_fa_s = tangent_kernel(fa_snaps[s_step], x, feedback)
                fa_loss_fixed = rollout_loss(k_fa_0, None, r0, lr, n, t_total, s_step)
                fa_loss_vel = rollout_loss(k_fa_0, k_fa_s - k_fa_0, r0, lr, n, t_total, s_step)
                row = EstimatorRow(
                    width=est_width, depth=args.est_depth, train_samples=n_train,
                    init_seed=init_seed, feedback_seed=feedback_seed, lr=lr,
                    lam_max=lam_max, horizon=t_total, early_step=s_step,
                    hidden_share=hidden_share,
                    loss_bp_emp=loss_bp_emp, loss_fa_emp=loss_fa_emp,
                    gap_emp=loss_fa_emp - loss_bp_emp,
                    gap_fixed=fa_loss_fixed - bp_loss_fixed,
                    gap_velocity=fa_loss_vel - bp_loss_vel,
                )
                rows.append(row)
                print(f"[B] w={est_width} N={n_train} init={init_seed} fb={feedback_index}: "
                      f"gap_emp={row.gap_emp:.5f} fixed={row.gap_fixed:.5f} "
                      f"velocity={row.gap_velocity:.5f}", flush=True)
            print(f"[B] w={est_width} N={n_train} init={init_seed} block done "
                  f"in {time.time() - t0:.1f}s", flush=True)

    with (outdir / "estimator_mnist_rows.csv").open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=list(EstimatorRow.__annotations__.keys()))
        writer.writeheader()
        for row in rows:
            writer.writerow(asdict(row))

    emp = np.array([row.gap_emp for row in rows])
    metrics = []
    for name, pred in [("fixed K(0)", np.array([row.gap_fixed for row in rows])),
                       ("linear velocity", np.array([row.gap_velocity for row in rows]))]:
        metrics.append({
            "predictor": name,
            "rows": len(rows),
            "mae": float(np.mean(np.abs(pred - emp))),
            "bias": float(np.mean(pred - emp)),
            "corr": float(np.corrcoef(pred, emp)[0, 1]),
        })
    with (outdir / "estimator_mnist_metrics.csv").open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=list(metrics[0].keys()))
        writer.writeheader()
        writer.writerows(metrics)
    for metric in metrics:
        print(f"[B] {metric['predictor']}: MAE={metric['mae']:.5f} "
              f"bias={metric['bias']:+.5f} corr={metric['corr']:.5f}", flush=True)

    fig, ax = plt.subplots(figsize=(5.8, 5.4), dpi=180)
    lims = [min(emp.min(), 0.0), max(emp.max(), 0.0)]
    pad = 0.06 * (lims[1] - lims[0] + 1e-9)
    ax.plot([lims[0] - pad, lims[1] + pad], [lims[0] - pad, lims[1] + pad],
            color="black", lw=1.0)
    ax.scatter([row.gap_fixed for row in rows], emp, s=22, alpha=0.65,
               color="#9aa5b1", marker="^",
               label=f"fixed K(0)  (MAE {metrics[0]['mae']:.4f})")
    ax.scatter([row.gap_velocity for row in rows], emp, s=22, alpha=0.8,
               color="#c65f16", marker="o",
               label=f"early velocity  (MAE {metrics[1]['mae']:.4f}, corr {metrics[1]['corr']:.3f})")
    ax.set_xlabel("predicted FA/BP train gap at T (no fit)")
    ax.set_ylabel("empirical FA/BP train gap at T")
    ax.set_title(f"MNIST: finite-time gap prediction\n"
                 f"widths {args.est_widths}, depth {args.est_depth}, "
                 f"T={t_total}, s={s_step}, N in {args.est_train_sizes}")
    ax.legend(fontsize=8)
    ax.grid(alpha=0.16)
    fig.tight_layout()
    fig.savefig(outdir / "estimator_mnist_scatter.png", bbox_inches="tight")
    plt.close(fig)
    return rows


def main() -> None:
    args = parse_args()
    torch.set_num_threads(args.torch_threads)
    if args.self_test:
        self_test()
        return
    args.outdir.mkdir(parents=True, exist_ok=True)
    payload: dict = {"config": {k: (str(v) if isinstance(v, Path) else v)
                                for k, v in vars(args).items()}}
    if args.part in ("a", "both"):
        rows_a = run_part_a(args, args.outdir)
        payload["e0_max_abs_error"] = max(row.mean_abs_error for row in rows_a)
        payload["e0_rows"] = len(rows_a)
    if args.part in ("b", "both"):
        rows_b = run_part_b(args, args.outdir)
        emp = np.array([row.gap_emp for row in rows_b])
        vel = np.array([row.gap_velocity for row in rows_b])
        fixed = np.array([row.gap_fixed for row in rows_b])
        payload["estimator_velocity_corr"] = float(np.corrcoef(vel, emp)[0, 1])
        payload["estimator_velocity_mae"] = float(np.mean(np.abs(vel - emp)))
        payload["estimator_fixed_mae"] = float(np.mean(np.abs(fixed - emp)))
        payload["estimator_rows"] = len(rows_b)
    (args.outdir / "summary.json").write_text(json.dumps(payload, indent=2) + "\n")
    print(f"summary: {args.outdir / 'summary.json'}")


if __name__ == "__main__":
    main()