#!/usr/bin/env python3
"""Validate the beta law for random feedback alignment.

For independent isotropic matrix directions A, B in R^{rows x cols}, the
squared Frobenius cosine

    Q = <A, B>_F^2 / (||A||_F^2 ||B||_F^2)

has distribution Beta(1/2, (D - 1)/2), where D = rows * cols.
"""

from __future__ import annotations

import argparse
import json
from dataclasses import asdict, dataclass
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from scipy import special, stats


@dataclass(frozen=True)
class RunConfig:
    rows: int
    cols: int
    samples: int
    batch_size: int
    seed: int
    dist_a: str
    dist_b: str
    thresholds: list[float]
    outdir: str
    plot: bool


@dataclass(frozen=True)
class Summary:
    dimension: int
    beta_alpha: float
    beta_beta: float
    empirical_mean: float
    theoretical_mean: float
    empirical_var: float
    theoretical_var: float
    ks_statistic: float
    ks_pvalue: float
    quantiles: dict[str, float]
    tail_probabilities: dict[str, dict[str, float]]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Monte Carlo validation of the FA static alignment beta law."
    )
    parser.add_argument("--rows", type=int, default=16, help="Matrix row count.")
    parser.add_argument("--cols", type=int, default=16, help="Matrix column count.")
    parser.add_argument("--samples", type=int, default=50_000, help="Number of pairs.")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=10_000,
        help="Number of pairs to sample per batch.",
    )
    parser.add_argument("--seed", type=int, default=0, help="Random seed.")
    parser.add_argument(
        "--dist-a",
        choices=["gaussian", "sphere", "rademacher"],
        default="gaussian",
        help="Distribution for A.",
    )
    parser.add_argument(
        "--dist-b",
        choices=["gaussian", "sphere", "rademacher"],
        default="gaussian",
        help="Distribution for B.",
    )
    parser.add_argument(
        "--thresholds",
        type=float,
        nargs="*",
        default=None,
        help="Tail thresholds q. Defaults to {1, 2, 5, 10}/D.",
    )
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/static_alignment_beta"),
        help="Directory for summary and plots.",
    )
    parser.add_argument("--plot", action="store_true", help="Save diagnostic plots.")
    return parser.parse_args()


def sample_vectors(
    rng: np.random.Generator, count: int, dim: int, distribution: str
) -> np.ndarray:
    if distribution == "gaussian":
        return rng.standard_normal((count, dim))

    if distribution == "sphere":
        values = rng.standard_normal((count, dim))
        norms = np.linalg.norm(values, axis=1, keepdims=True)
        return values / np.maximum(norms, np.finfo(values.dtype).tiny)

    if distribution == "rademacher":
        return rng.choice(np.array([-1.0, 1.0]), size=(count, dim))

    raise ValueError(f"Unknown distribution: {distribution}")


def squared_cosines(config: RunConfig) -> np.ndarray:
    dim = config.rows * config.cols
    rng = np.random.default_rng(config.seed)
    values = np.empty(config.samples, dtype=np.float64)

    offset = 0
    while offset < config.samples:
        count = min(config.batch_size, config.samples - offset)
        a = sample_vectors(rng, count, dim, config.dist_a)
        b = sample_vectors(rng, count, dim, config.dist_b)

        dot = np.einsum("ij,ij->i", a, b)
        a_norm_sq = np.einsum("ij,ij->i", a, a)
        b_norm_sq = np.einsum("ij,ij->i", b, b)
        values[offset : offset + count] = (dot * dot) / (a_norm_sq * b_norm_sq)
        offset += count

    return values


def summarize(q_values: np.ndarray, dimension: int, thresholds: list[float]) -> Summary:
    alpha = 0.5
    beta_param = (dimension - 1) / 2
    beta_dist = stats.beta(alpha, beta_param)

    quantile_levels = [0.01, 0.05, 0.10, 0.50, 0.90, 0.95, 0.99]
    quantiles = {
        f"{level:.2f}": float(np.quantile(q_values, level))
        for level in quantile_levels
    }

    tail_probabilities: dict[str, dict[str, float]] = {}
    for threshold in thresholds:
        empirical_tail = float(np.mean(q_values >= threshold))
        theoretical_tail = float(1.0 - special.betainc(alpha, beta_param, threshold))
        tail_probabilities[f"{threshold:.12g}"] = {
            "empirical": empirical_tail,
            "theoretical": theoretical_tail,
            "absolute_error": abs(empirical_tail - theoretical_tail),
        }

    ks = stats.kstest(q_values, beta_dist.cdf)

    return Summary(
        dimension=dimension,
        beta_alpha=alpha,
        beta_beta=beta_param,
        empirical_mean=float(np.mean(q_values)),
        theoretical_mean=float(beta_dist.mean()),
        empirical_var=float(np.var(q_values, ddof=1)),
        theoretical_var=float(beta_dist.var()),
        ks_statistic=float(ks.statistic),
        ks_pvalue=float(ks.pvalue),
        quantiles=quantiles,
        tail_probabilities=tail_probabilities,
    )


def write_summary(config: RunConfig, summary: Summary, outdir: Path) -> Path:
    outdir.mkdir(parents=True, exist_ok=True)
    path = outdir / "summary.json"
    payload = {
        "config": asdict(config),
        "summary": asdict(summary),
    }
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    return path


def save_plots(q_values: np.ndarray, dimension: int, outdir: Path) -> list[Path]:
    outdir.mkdir(parents=True, exist_ok=True)
    alpha = 0.5
    beta_param = (dimension - 1) / 2
    beta_dist = stats.beta(alpha, beta_param)

    paths: list[Path] = []

    hist_path = outdir / "histogram_beta_overlay.png"
    x_min = float(max(beta_dist.ppf(1e-5), np.finfo(float).tiny))
    x_max = float(max(np.quantile(q_values, 0.999), beta_dist.ppf(0.999)))
    xs = np.linspace(x_min, x_max, 500)
    beta_pdf = np.exp(np.clip(beta_dist.logpdf(xs), -745, 80))
    plt.figure(figsize=(7, 4.5))
    plt.hist(q_values, bins=80, density=True, alpha=0.45, label="empirical")
    plt.plot(xs, beta_pdf, color="black", linewidth=2, label="beta law")
    plt.xlabel("Q = squared Frobenius cosine")
    plt.ylabel("density")
    plt.title(f"Static alignment distribution, D={dimension}")
    plt.legend()
    plt.tight_layout()
    plt.savefig(hist_path, dpi=180)
    plt.close()
    paths.append(hist_path)

    qq_path = outdir / "qq_plot.png"
    probs = (np.arange(1, len(q_values) + 1) - 0.5) / len(q_values)
    empirical = np.sort(q_values)
    theoretical = beta_dist.ppf(probs)
    max_value = float(max(empirical[-1], theoretical[-1]))
    plt.figure(figsize=(5, 5))
    plt.scatter(theoretical, empirical, s=4, alpha=0.35)
    plt.plot([0, max_value], [0, max_value], color="black", linewidth=1)
    plt.xlabel("theoretical beta quantile")
    plt.ylabel("empirical quantile")
    plt.title("Q-Q plot")
    plt.tight_layout()
    plt.savefig(qq_path, dpi=180)
    plt.close()
    paths.append(qq_path)

    return paths


def main() -> None:
    args = parse_args()
    dimension = args.rows * args.cols
    if dimension < 2:
        raise ValueError("rows * cols must be at least 2 for the beta law.")
    if args.samples < 1:
        raise ValueError("--samples must be positive.")
    if args.batch_size < 1:
        raise ValueError("--batch-size must be positive.")

    thresholds = args.thresholds
    if thresholds is None:
        thresholds = [1 / dimension, 2 / dimension, 5 / dimension, 10 / dimension]
    thresholds = [q for q in thresholds if 0 <= q <= 1]

    config = RunConfig(
        rows=args.rows,
        cols=args.cols,
        samples=args.samples,
        batch_size=args.batch_size,
        seed=args.seed,
        dist_a=args.dist_a,
        dist_b=args.dist_b,
        thresholds=thresholds,
        outdir=str(args.outdir),
        plot=args.plot,
    )

    q_values = squared_cosines(config)
    summary = summarize(q_values, dimension, thresholds)
    summary_path = write_summary(config, summary, args.outdir)

    plot_paths: list[Path] = []
    if args.plot:
        plot_paths = save_plots(q_values, dimension, args.outdir)

    print(f"dimension: {dimension}")
    print(f"samples: {args.samples}")
    print(
        "mean: "
        f"empirical={summary.empirical_mean:.8g}, "
        f"theoretical={summary.theoretical_mean:.8g}"
    )
    print(
        "variance: "
        f"empirical={summary.empirical_var:.8g}, "
        f"theoretical={summary.theoretical_var:.8g}"
    )
    print(
        "KS: "
        f"statistic={summary.ks_statistic:.6g}, "
        f"pvalue={summary.ks_pvalue:.6g}"
    )
    print(f"summary: {summary_path}")
    for path in plot_paths:
        print(f"plot: {path}")


if __name__ == "__main__":
    main()