#!/usr/bin/env python3
"""Large empirical validation of FA alignment capacity costs.

By rotational invariance, for an isotropic target A and feedback B in R^D,
cos^2(A, B) has the same law as the squared first coordinate of a random
Gaussian direction:

    Q = z_1^2 / ||z||^2,  z ~ N(0, I_D).

This script uses that fixed-target sampler to validate the beta law, empirical
tail probabilities, log-volume capacity costs, and multilayer product scaling.
"""

from __future__ import annotations

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

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


@dataclass(frozen=True)
class RunConfig:
    dimensions: list[int]
    samples: int
    batch_size: int
    seed: int
    chance_cs: list[float]
    fixed_qs: list[float]
    multilayer_cs: list[float]
    multilayer_counts: list[int]
    outdir: str
    plot: bool


@dataclass(frozen=True)
class DistributionRow:
    dimension: int
    samples: int
    empirical_mean: float
    theoretical_mean: float
    empirical_var: float
    theoretical_var: float
    ks_statistic: float
    ks_pvalue: float
    q50_empirical: float
    q50_theoretical: float
    q90_empirical: float
    q90_theoretical: float
    q99_empirical: float
    q99_theoretical: float


@dataclass(frozen=True)
class CapacityTailRow:
    dimension: int
    samples: int
    threshold_mode: str
    threshold_label: str
    q_threshold: float
    hits: int
    empirical_tail: float
    empirical_tail_smoothed: float
    theoretical_tail: float
    expected_hits: float
    empirical_cost_nats: float
    empirical_cost_smoothed_nats: float
    theoretical_cost_nats: float
    cost_error_smoothed_nats: float
    relative_tail_error: float


@dataclass(frozen=True)
class MultilayerRow:
    dimension: int
    samples: int
    grouped_samples: int
    layers: int
    c: float
    q_threshold: float
    hits: int
    empirical_all_tail: float
    empirical_all_tail_smoothed: float
    theoretical_all_tail: float
    expected_hits: float
    empirical_total_cost_smoothed_nats: float
    theoretical_total_cost_nats: float
    cost_error_smoothed_nats: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Large empirical validation for FA capacity costs."
    )
    parser.add_argument(
        "--dimensions",
        type=int,
        nargs="+",
        default=[64, 128, 256, 512, 1024, 2048, 4096],
        help="Flattened matrix dimensions D to validate.",
    )
    parser.add_argument(
        "--samples",
        type=int,
        default=100_000,
        help="Monte Carlo samples per dimension.",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=2048,
        help="Gaussian direction samples per batch.",
    )
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--chance-cs",
        type=float,
        nargs="+",
        default=[0.5, 1.0, 2.0, 5.0, 10.0],
        help="Chance-level thresholds q=c/D.",
    )
    parser.add_argument(
        "--fixed-qs",
        type=float,
        nargs="*",
        default=[0.001, 0.002, 0.005, 0.01],
        help="Fixed thresholds q.",
    )
    parser.add_argument(
        "--multilayer-cs",
        type=float,
        nargs="+",
        default=[0.5, 1.0],
        help="Chance-level c values for multilayer all-event validation.",
    )
    parser.add_argument(
        "--multilayer-counts",
        type=int,
        nargs="+",
        default=[1, 2, 4, 8],
        help="Layer counts for multilayer all-event validation.",
    )
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/capacity_empirical_validation"),
    )
    parser.add_argument("--plot", action="store_true")
    return parser.parse_args()


def sample_q_fixed_target(
    rng: np.random.Generator, dimension: int, samples: int, batch_size: int
) -> np.ndarray:
    values = np.empty(samples, dtype=np.float64)
    offset = 0
    while offset < samples:
        count = min(batch_size, samples - offset)
        z = rng.standard_normal((count, dimension))
        numerator = z[:, 0] * z[:, 0]
        denominator = np.einsum("ij,ij->i", z, z)
        values[offset : offset + count] = numerator / denominator
        offset += count
    return values


def beta_dist(dimension: int) -> stats.rv_continuous:
    return stats.beta(0.5, (dimension - 1) / 2)


def distribution_summary(q_values: np.ndarray, dimension: int) -> DistributionRow:
    dist = beta_dist(dimension)
    ks = stats.kstest(q_values, dist.cdf)
    return DistributionRow(
        dimension=dimension,
        samples=len(q_values),
        empirical_mean=float(np.mean(q_values)),
        theoretical_mean=float(dist.mean()),
        empirical_var=float(np.var(q_values, ddof=1)),
        theoretical_var=float(dist.var()),
        ks_statistic=float(ks.statistic),
        ks_pvalue=float(ks.pvalue),
        q50_empirical=float(np.quantile(q_values, 0.50)),
        q50_theoretical=float(dist.ppf(0.50)),
        q90_empirical=float(np.quantile(q_values, 0.90)),
        q90_theoretical=float(dist.ppf(0.90)),
        q99_empirical=float(np.quantile(q_values, 0.99)),
        q99_theoretical=float(dist.ppf(0.99)),
    )


def finite_cost(tail: float) -> float:
    if tail <= 0:
        return math.inf
    return -math.log(tail)


def capacity_tail_row(
    q_values: np.ndarray,
    dimension: int,
    threshold_mode: str,
    threshold_label: str,
    q_threshold: float,
) -> CapacityTailRow:
    samples = len(q_values)
    dist = beta_dist(dimension)
    hits = int(np.count_nonzero(q_values >= q_threshold))
    empirical_tail = hits / samples
    empirical_tail_smoothed = (hits + 0.5) / (samples + 1.0)
    theoretical_tail = float(dist.sf(q_threshold))
    expected_hits = samples * theoretical_tail
    theoretical_cost = finite_cost(theoretical_tail)
    smoothed_cost = finite_cost(empirical_tail_smoothed)
    relative_error = (
        abs(empirical_tail - theoretical_tail) / theoretical_tail
        if theoretical_tail > 0
        else math.nan
    )
    return CapacityTailRow(
        dimension=dimension,
        samples=samples,
        threshold_mode=threshold_mode,
        threshold_label=threshold_label,
        q_threshold=q_threshold,
        hits=hits,
        empirical_tail=empirical_tail,
        empirical_tail_smoothed=empirical_tail_smoothed,
        theoretical_tail=theoretical_tail,
        expected_hits=expected_hits,
        empirical_cost_nats=finite_cost(empirical_tail),
        empirical_cost_smoothed_nats=smoothed_cost,
        theoretical_cost_nats=theoretical_cost,
        cost_error_smoothed_nats=smoothed_cost - theoretical_cost,
        relative_tail_error=relative_error,
    )


def multilayer_row(
    q_values: np.ndarray, dimension: int, layers: int, c: float
) -> MultilayerRow:
    samples = len(q_values)
    grouped_samples = samples // layers
    usable = grouped_samples * layers
    q_threshold = c / dimension
    grouped = q_values[:usable].reshape(grouped_samples, layers)
    hits = int(np.count_nonzero(np.all(grouped >= q_threshold, axis=1)))
    empirical_tail = hits / grouped_samples
    empirical_tail_smoothed = (hits + 0.5) / (grouped_samples + 1.0)
    per_layer_tail = float(beta_dist(dimension).sf(q_threshold))
    theoretical_tail = per_layer_tail**layers
    expected_hits = grouped_samples * theoretical_tail
    smoothed_cost = finite_cost(empirical_tail_smoothed)
    theoretical_cost = finite_cost(theoretical_tail)
    return MultilayerRow(
        dimension=dimension,
        samples=samples,
        grouped_samples=grouped_samples,
        layers=layers,
        c=c,
        q_threshold=q_threshold,
        hits=hits,
        empirical_all_tail=empirical_tail,
        empirical_all_tail_smoothed=empirical_tail_smoothed,
        theoretical_all_tail=theoretical_tail,
        expected_hits=expected_hits,
        empirical_total_cost_smoothed_nats=smoothed_cost,
        theoretical_total_cost_nats=theoretical_cost,
        cost_error_smoothed_nats=smoothed_cost - theoretical_cost,
    )


def validate_args(config: RunConfig) -> None:
    if any(d < 2 for d in config.dimensions):
        raise ValueError("All dimensions must be at least 2.")
    if config.samples < 1:
        raise ValueError("--samples must be positive.")
    if config.batch_size < 1:
        raise ValueError("--batch-size must be positive.")
    if any(c < 0 for c in config.chance_cs):
        raise ValueError("--chance-cs must be non-negative.")
    if any(c < 0 for c in config.multilayer_cs):
        raise ValueError("--multilayer-cs must be non-negative.")
    if any(q < 0 or q > 1 for q in config.fixed_qs):
        raise ValueError("--fixed-qs must lie in [0, 1].")
    if any(layers < 1 for layers in config.multilayer_counts):
        raise ValueError("--multilayer-counts must be positive.")


def write_csv(path: Path, rows: list[object]) -> None:
    if not rows:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as handle:
        first = asdict(rows[0])  # type: ignore[arg-type]
        writer = csv.DictWriter(handle, fieldnames=list(first.keys()))
        writer.writeheader()
        for row in rows:
            writer.writerow(asdict(row))  # type: ignore[arg-type]


def write_outputs(
    config: RunConfig,
    distribution_rows: list[DistributionRow],
    capacity_rows: list[CapacityTailRow],
    multilayer_rows: list[MultilayerRow],
    outdir: Path,
) -> None:
    outdir.mkdir(parents=True, exist_ok=True)
    write_csv(outdir / "distribution_summary.csv", distribution_rows)
    write_csv(outdir / "capacity_tails.csv", capacity_rows)
    write_csv(outdir / "multilayer_capacity.csv", multilayer_rows)
    payload = {
        "config": asdict(config),
        "distribution_summary": [asdict(row) for row in distribution_rows],
        "capacity_tails": [asdict(row) for row in capacity_rows],
        "multilayer_capacity": [asdict(row) for row in multilayer_rows],
    }
    (outdir / "summary.json").write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n"
    )


def save_plots(
    distribution_rows: list[DistributionRow],
    capacity_rows: list[CapacityTailRow],
    multilayer_rows: list[MultilayerRow],
    outdir: Path,
) -> list[Path]:
    outdir.mkdir(parents=True, exist_ok=True)
    paths: list[Path] = []

    ks_path = outdir / "ks_by_dimension.png"
    plt.figure(figsize=(7, 4.5))
    dims = [row.dimension for row in distribution_rows]
    ks_values = [row.ks_statistic for row in distribution_rows]
    plt.plot(dims, ks_values, marker="o")
    plt.xscale("log", base=2)
    plt.xlabel("dimension D")
    plt.ylabel("KS statistic")
    plt.title("Beta-law distribution calibration")
    plt.tight_layout()
    plt.savefig(ks_path, dpi=180)
    plt.close()
    paths.append(ks_path)

    chance_path = outdir / "chance_capacity_empirical_vs_theory.png"
    chance_rows = [row for row in capacity_rows if row.threshold_mode == "chance"]
    plt.figure(figsize=(7, 4.5))
    for label in sorted({row.threshold_label for row in chance_rows}):
        rows = sorted(
            [row for row in chance_rows if row.threshold_label == label],
            key=lambda row: row.dimension,
        )
        plt.plot(
            [row.theoretical_cost_nats for row in rows],
            [row.empirical_cost_smoothed_nats for row in rows],
            marker="o",
            linestyle="",
            label=label,
        )
    max_cost = max(
        max(row.theoretical_cost_nats, row.empirical_cost_smoothed_nats)
        for row in chance_rows
        if np.isfinite(row.empirical_cost_smoothed_nats)
    )
    plt.plot([0, max_cost], [0, max_cost], color="black", linewidth=1)
    plt.xlabel("theoretical capacity cost (nats)")
    plt.ylabel("empirical smoothed cost (nats)")
    plt.title("Chance-threshold capacity calibration")
    plt.legend()
    plt.tight_layout()
    plt.savefig(chance_path, dpi=180)
    plt.close()
    paths.append(chance_path)

    multilayer_path = outdir / "multilayer_capacity_empirical_vs_theory.png"
    plt.figure(figsize=(7, 4.5))
    for c in sorted({row.c for row in multilayer_rows}):
        rows = [row for row in multilayer_rows if row.c == c]
        plt.scatter(
            [row.theoretical_total_cost_nats for row in rows],
            [row.empirical_total_cost_smoothed_nats for row in rows],
            label=f"c={c:g}",
            alpha=0.8,
        )
    max_cost = max(
        max(row.theoretical_total_cost_nats, row.empirical_total_cost_smoothed_nats)
        for row in multilayer_rows
        if np.isfinite(row.empirical_total_cost_smoothed_nats)
    )
    plt.plot([0, max_cost], [0, max_cost], color="black", linewidth=1)
    plt.xlabel("theoretical total cost (nats)")
    plt.ylabel("empirical smoothed total cost (nats)")
    plt.title("Multilayer product-capacity calibration")
    plt.legend()
    plt.tight_layout()
    plt.savefig(multilayer_path, dpi=180)
    plt.close()
    paths.append(multilayer_path)

    return paths


def parse_config(args: argparse.Namespace) -> RunConfig:
    return RunConfig(
        dimensions=args.dimensions,
        samples=args.samples,
        batch_size=args.batch_size,
        seed=args.seed,
        chance_cs=args.chance_cs,
        fixed_qs=args.fixed_qs,
        multilayer_cs=args.multilayer_cs,
        multilayer_counts=args.multilayer_counts,
        outdir=str(args.outdir),
        plot=args.plot,
    )


def main() -> None:
    args = parse_args()
    config = parse_config(args)
    validate_args(config)

    rng = np.random.default_rng(config.seed)
    distribution_rows: list[DistributionRow] = []
    capacity_rows: list[CapacityTailRow] = []
    multilayer_rows: list[MultilayerRow] = []

    for dimension in config.dimensions:
        q_values = sample_q_fixed_target(
            rng, dimension, config.samples, config.batch_size
        )
        dist_row = distribution_summary(q_values, dimension)
        distribution_rows.append(dist_row)

        for c in config.chance_cs:
            q_threshold = min(c / dimension, 1.0)
            capacity_rows.append(
                capacity_tail_row(
                    q_values,
                    dimension,
                    "chance",
                    f"c={c:g}/D",
                    q_threshold,
                )
            )

        for q in config.fixed_qs:
            capacity_rows.append(
                capacity_tail_row(q_values, dimension, "fixed", f"q={q:g}", q)
            )

        for c in config.multilayer_cs:
            for layers in config.multilayer_counts:
                if layers <= len(q_values):
                    multilayer_rows.append(
                        multilayer_row(q_values, dimension, layers, c)
                    )

        print(
            f"D={dimension}: "
            f"mean_emp={dist_row.empirical_mean:.6g}, "
            f"mean_theory={dist_row.theoretical_mean:.6g}, "
            f"KS={dist_row.ks_statistic:.6g}"
        )

    outdir = Path(config.outdir)
    write_outputs(config, distribution_rows, capacity_rows, multilayer_rows, outdir)
    plot_paths = (
        save_plots(distribution_rows, capacity_rows, multilayer_rows, outdir)
        if config.plot
        else []
    )

    reliable_rows = [
        row
        for row in capacity_rows
        if row.expected_hits >= 20 and row.threshold_mode == "chance"
    ]
    max_abs_cost_error = max(
        (abs(row.cost_error_smoothed_nats) for row in reliable_rows), default=math.nan
    )
    max_ks = max(row.ks_statistic for row in distribution_rows)
    print(f"dimensions: {config.dimensions}")
    print(f"samples_per_dimension: {config.samples}")
    print(f"max_ks_statistic: {max_ks:.8g}")
    print(
        "max_abs_chance_cost_error_nats_expected_hits_ge_20: "
        f"{max_abs_cost_error:.8g}"
    )
    print(f"distribution_summary: {outdir / 'distribution_summary.csv'}")
    print(f"capacity_tails: {outdir / 'capacity_tails.csv'}")
    print(f"multilayer_capacity: {outdir / 'multilayer_capacity.csv'}")
    for path in plot_paths:
        print(f"plot: {path}")


if __name__ == "__main__":
    main()