#!/usr/bin/env python3
"""Plot exact initial FA erosion distributions for one-hidden-layer MLPs.

For a fixed forward initialization and residual, one-hidden-layer FA has

    speed_FA = output_speed + <C, B>

with Gaussian feedback B. Therefore the one-step erosion

    e0 = 1 - speed_FA / speed_BP

is exactly Gaussian conditional on the forward weights and data.
"""

from __future__ import annotations

import argparse
import csv
import math
from dataclasses import asdict, dataclass
from pathlib import Path

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


Tensor = torch.Tensor


@dataclass(frozen=True)
class DistributionRow:
    width: int
    init_seed: int
    feedback_samples: int
    bp_speed: float
    output_speed: float
    hidden_bp_speed: float
    theory_mean: float
    theory_std: float
    empirical_mean: float
    empirical_std: float
    ks_stat: float
    ks_pvalue: float
    mean_error: float
    std_error: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Exact e0 distribution validation.")
    parser.add_argument("--input-dim", type=int, default=16)
    parser.add_argument("--output-dim", type=int, default=4)
    parser.add_argument("--widths", type=int, nargs="+", default=[16, 32, 64, 128])
    parser.add_argument("--train-samples", type=int, default=128)
    parser.add_argument("--init-seeds", type=int, default=3)
    parser.add_argument("--feedback-samples", type=int, default=4096)
    parser.add_argument("--data-seed", type=int, default=123)
    parser.add_argument(
        "--feedback-scale",
        choices=["relu", "fan-in", "unit"],
        default="relu",
    )
    parser.add_argument("--torch-threads", type=int, default=8)
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/actual_fa_initial_erosion_distribution"),
    )
    return parser.parse_args()


def feedback_scale(rows: int, mode: str) -> float:
    if mode == "relu":
        return math.sqrt(2.0 / rows)
    if mode == "fan-in":
        return math.sqrt(1.0 / rows)
    if mode == "unit":
        return 1.0
    raise ValueError(f"Unknown feedback scale: {mode}")


def make_data(samples: int, input_dim: int, output_dim: int, seed: int) -> tuple[Tensor, Tensor]:
    generator = torch.Generator(device="cpu")
    generator.manual_seed(seed)
    x = torch.randn(samples, input_dim, generator=generator, dtype=torch.float64)
    y = torch.randn(samples, output_dim, generator=generator, dtype=torch.float64)
    return x, y


def init_weights(input_dim: int, output_dim: int, width: int, seed: int) -> tuple[Tensor, Tensor]:
    generator = torch.Generator(device="cpu")
    generator.manual_seed(seed)
    w1 = torch.randn(width, input_dim, generator=generator, dtype=torch.float64) * math.sqrt(2.0 / input_dim)
    w2 = torch.randn(output_dim, width, generator=generator, dtype=torch.float64) / math.sqrt(width)
    return w1, w2


def forward(w1: Tensor, w2: Tensor, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
    z = x @ w1.T
    h = torch.relu(z)
    out = h @ w2.T
    return z, h, out


def theory_terms(
    w1: Tensor,
    w2: Tensor,
    x: Tensor,
    y: Tensor,
    scale: float,
) -> tuple[float, float, float, float, float, Tensor]:
    z, h, out = forward(w1, w2, x)
    batch = x.shape[0]
    delta_out = (out - y) / batch
    gate = (z > 0).to(torch.float64)

    grad_out = delta_out.T @ h
    delta_hidden_bp = (delta_out @ w2) * gate
    grad_hidden_bp = delta_hidden_bp.T @ x

    output_speed = float(torch.sum(grad_out * grad_out))
    hidden_bp_speed = float(torch.sum(grad_hidden_bp * grad_hidden_bp))
    bp_speed = output_speed + hidden_bp_speed

    # g_hidden_FA[a, d] = sum_c B[a, c] S[a, c, d]
    # C[a, c] = sum_d g_hidden_BP[a, d] S[a, c, d]
    # Use einsum over samples n and input dimension d:
    # S[a,c,d] = sum_n gate[n,a] * x[n,d] * delta_out[n,c]
    s_tensor = torch.einsum("na,nd,nc->acd", gate, x, delta_out)
    coeff = torch.einsum("ad,acd->ac", grad_hidden_bp, s_tensor)
    z_std = scale * float(torch.linalg.norm(coeff))
    theory_mean = 1.0 - output_speed / bp_speed
    theory_std = z_std / bp_speed
    return bp_speed, output_speed, hidden_bp_speed, theory_mean, theory_std, coeff


def sample_erosions(
    bp_speed: float,
    output_speed: float,
    coeff: Tensor,
    scale: float,
    samples: int,
    seed: int,
) -> np.ndarray:
    generator = torch.Generator(device="cpu")
    generator.manual_seed(seed)
    values = []
    for _ in range(samples):
        feedback = torch.randn(
            coeff.shape,
            generator=generator,
            dtype=torch.float64,
        ) * scale
        hidden_mixed = float(torch.sum(feedback * coeff))
        speed_fa = output_speed + hidden_mixed
        values.append(1.0 - speed_fa / bp_speed)
    return np.array(values, dtype=np.float64)


def run_one(args: argparse.Namespace, x: Tensor, y: Tensor, width: int, init_seed: int) -> tuple[DistributionRow, np.ndarray]:
    w1, w2 = init_weights(args.input_dim, args.output_dim, width, init_seed)
    scale = feedback_scale(width, args.feedback_scale)
    bp_speed, output_speed, hidden_bp_speed, theory_mean, theory_std, coeff = theory_terms(
        w1,
        w2,
        x,
        y,
        scale,
    )
    erosions = sample_erosions(
        bp_speed,
        output_speed,
        coeff,
        scale,
        args.feedback_samples,
        seed=1_000_000 + 10_000 * init_seed + width,
    )
    standardized = (erosions - theory_mean) / theory_std
    ks = stats.kstest(standardized, "norm")
    row = DistributionRow(
        width=width,
        init_seed=init_seed,
        feedback_samples=args.feedback_samples,
        bp_speed=bp_speed,
        output_speed=output_speed,
        hidden_bp_speed=hidden_bp_speed,
        theory_mean=theory_mean,
        theory_std=theory_std,
        empirical_mean=float(np.mean(erosions)),
        empirical_std=float(np.std(erosions, ddof=1)),
        ks_stat=float(ks.statistic),
        ks_pvalue=float(ks.pvalue),
        mean_error=float(np.mean(erosions) - theory_mean),
        std_error=float(np.std(erosions, ddof=1) - theory_std),
    )
    return row, erosions


def write_rows(path: Path, rows: list[DistributionRow]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(DistributionRow.__annotations__.keys()))
        writer.writeheader()
        for row in rows:
            writer.writerow(asdict(row))


def plot_overlay(
    rows: list[DistributionRow],
    samples: dict[tuple[int, int], np.ndarray],
    outdir: Path,
) -> Path:
    selected: list[DistributionRow] = []
    for width in sorted({row.width for row in rows}):
        width_rows = [row for row in rows if row.width == width]
        selected.append(width_rows[0])

    cols = 2
    panel_rows = int(math.ceil(len(selected) / cols))
    fig, axes = plt.subplots(panel_rows, cols, figsize=(11.0, 4.2 * panel_rows), dpi=180, squeeze=False)
    for ax in axes.ravel():
        ax.axis("off")

    for ax, row in zip(axes.ravel(), selected):
        ax.axis("on")
        values = samples[(row.width, row.init_seed)]
        lo = min(float(np.quantile(values, 0.001)), row.theory_mean - 4.0 * row.theory_std)
        hi = max(float(np.quantile(values, 0.999)), row.theory_mean + 4.0 * row.theory_std)
        grid = np.linspace(lo, hi, 500)
        density = stats.norm.pdf(grid, loc=row.theory_mean, scale=row.theory_std)
        ax.hist(values, bins=70, density=True, color="#c65f16", alpha=0.42, label="empirical samples")
        ax.plot(grid, density, color="#174f91", linewidth=2.4, label="theory Gaussian")
        ax.axvline(row.theory_mean, color="#174f91", linewidth=1.4, linestyle="--")
        ax.axvline(row.empirical_mean, color="#a84708", linewidth=1.4, linestyle=":")
        ax.set_title(
            f"width={row.width}, init={row.init_seed} | "
            f"KS={row.ks_stat:.3f}, p={row.ks_pvalue:.2f}"
        )
        ax.set_xlabel("initial FA operator erosion e0")
        ax.set_ylabel("density")
        ax.legend(fontsize=8)
        ax.grid(alpha=0.16)

    fig.suptitle("Exact e0 distribution: theory density vs random-feedback samples", y=1.01)
    fig.tight_layout()
    path = outdir / "e0_distribution_theory_vs_empirical.png"
    fig.savefig(path, bbox_inches="tight")
    plt.close(fig)
    return path


def plot_moments(rows: list[DistributionRow], outdir: Path) -> Path:
    fig, axes = plt.subplots(1, 2, figsize=(11.2, 4.8), dpi=180)
    widths = sorted({row.width for row in rows})
    for width in widths:
        sub = [row for row in rows if row.width == width]
        axes[0].scatter(
            [row.theory_mean for row in sub],
            [row.empirical_mean for row in sub],
            s=34,
            alpha=0.72,
            label=f"w={width}",
        )
        axes[1].scatter(
            [row.theory_std for row in sub],
            [row.empirical_std for row in sub],
            s=34,
            alpha=0.72,
            label=f"w={width}",
        )
    for ax, title in zip(axes, ["mean", "standard deviation"]):
        values = []
        if title == "mean":
            values = [v for row in rows for v in (row.theory_mean, row.empirical_mean)]
        else:
            values = [v for row in rows for v in (row.theory_std, row.empirical_std)]
        lo, hi = min(values), max(values)
        pad = 0.04 * (hi - lo + 1e-12)
        ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], color="black", linewidth=1.0)
        ax.set_title(f"theory vs empirical {title}")
        ax.set_xlabel("theory")
        ax.set_ylabel("empirical")
        ax.grid(alpha=0.16)
    axes[0].legend(fontsize=8, ncols=2)
    fig.tight_layout()
    path = outdir / "e0_distribution_moment_calibration.png"
    fig.savefig(path, bbox_inches="tight")
    plt.close(fig)
    return path


def main() -> None:
    args = parse_args()
    torch.set_num_threads(args.torch_threads)
    x, y = make_data(args.train_samples, args.input_dim, args.output_dim, args.data_seed)
    rows: list[DistributionRow] = []
    sample_map: dict[tuple[int, int], np.ndarray] = {}
    for width in args.widths:
        for init_seed in range(args.init_seeds):
            row, erosions = run_one(args, x, y, width, init_seed)
            rows.append(row)
            sample_map[(width, init_seed)] = erosions
            print(
                f"width={width} init={init_seed}: "
                f"mean theory={row.theory_mean:.6f} empirical={row.empirical_mean:.6f}; "
                f"std theory={row.theory_std:.6f} empirical={row.empirical_std:.6f}; "
                f"KS={row.ks_stat:.4f} p={row.ks_pvalue:.3g}",
                flush=True,
            )

    args.outdir.mkdir(parents=True, exist_ok=True)
    csv_path = args.outdir / "e0_distribution_rows.csv"
    write_rows(csv_path, rows)
    paths = [
        plot_overlay(rows, sample_map, args.outdir),
        plot_moments(rows, args.outdir),
    ]
    print(f"rows: {csv_path}")
    for path in paths:
        print(f"plot: {path}")


if __name__ == "__main__":
    main()