#!/usr/bin/env python3
"""Match feedback-initialization coverage distributions from M_mu spectra.

For a feedback initialization distribution mu, define

    M_mu = E[b_hat b_hat^T].

For a uniformly random target direction a, the expected squared alignment over
feedback draws is A(a) = a^T M_mu a. If lambda_i are eigenvalues of M_mu, then

    A(a) = sum_i lambda_i G_i / sum_i G_i,  G_i ~ chi^2_1.

This script compares the population-predicted distribution from known spectra
with the empirical distribution obtained by estimating M_mu from sampled
feedback directions.
"""

from __future__ import annotations

import argparse
import csv
import json
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:
    dimension: int
    target_samples: int
    feedback_samples: int
    batch_size: int
    seed: int
    schemes: list[str]
    subspace_dim: int
    anisotropy: float
    outdir: str
    plot: bool


@dataclass(frozen=True)
class MatchRow:
    scheme: str
    dimension: int
    target_samples: int
    feedback_samples: int
    population_lambda_min: float
    empirical_lambda_min: float
    population_lambda_max: float
    empirical_lambda_max: float
    predicted_mean: float
    empirical_mean: float
    predicted_std: float
    empirical_std: float
    predicted_q01: float
    empirical_q01: float
    predicted_q50: float
    empirical_q50: float
    predicted_q99: float
    empirical_q99: float
    ks_2sample_statistic: float
    ks_2sample_pvalue: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Distribution matching for feedback initialization coverage."
    )
    parser.add_argument("--dimension", type=int, default=128)
    parser.add_argument("--target-samples", type=int, default=100_000)
    parser.add_argument("--feedback-samples", type=int, default=100_000)
    parser.add_argument("--batch-size", type=int, default=8192)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--schemes",
        nargs="+",
        default=["isotropic", "rademacher", "subspace", "axis", "geometric_axis"],
        choices=["isotropic", "rademacher", "subspace", "axis", "geometric_axis"],
    )
    parser.add_argument("--subspace-dim", type=int, default=8)
    parser.add_argument("--anisotropy", type=float, default=64.0)
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/initialization_distribution_matching"),
    )
    parser.add_argument("--plot", action="store_true")
    return parser.parse_args()


def validate_config(config: RunConfig) -> None:
    if config.dimension < 2:
        raise ValueError("--dimension must be at least 2.")
    if config.target_samples < 1:
        raise ValueError("--target-samples must be positive.")
    if config.feedback_samples < 1:
        raise ValueError("--feedback-samples must be positive.")
    if config.batch_size < 1:
        raise ValueError("--batch-size must be positive.")
    if config.subspace_dim < 1:
        raise ValueError("--subspace-dim must be positive.")
    if config.anisotropy < 1:
        raise ValueError("--anisotropy must be at least 1.")


def population_eigenvalues(
    scheme: str, dimension: int, subspace_dim: int, anisotropy: float
) -> np.ndarray:
    if scheme in {"isotropic", "rademacher"}:
        return np.full(dimension, 1.0 / dimension)

    if scheme == "axis":
        values = np.zeros(dimension, dtype=np.float64)
        values[-1] = 1.0
        return values

    if scheme == "subspace":
        active = min(subspace_dim, dimension)
        values = np.zeros(dimension, dtype=np.float64)
        values[-active:] = 1.0 / active
        return values

    if scheme == "geometric_axis":
        weights = np.geomspace(1.0, anisotropy, num=dimension)
        return weights / weights.sum()

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


def sample_feedback(
    rng: np.random.Generator,
    scheme: str,
    count: int,
    dimension: int,
    eigenvalues: np.ndarray,
    subspace_dim: int,
) -> np.ndarray:
    if scheme == "isotropic":
        values = rng.standard_normal((count, dimension))
        norms = np.linalg.norm(values, axis=1, keepdims=True)
        return values / np.maximum(norms, np.finfo(float).tiny)

    if scheme == "rademacher":
        values = rng.choice(np.array([-1.0, 1.0]), size=(count, dimension))
        return values / np.sqrt(dimension)

    if scheme == "axis":
        values = np.zeros((count, dimension), dtype=np.float64)
        values[:, -1] = rng.choice(np.array([-1.0, 1.0]), size=count)
        return values

    if scheme == "subspace":
        active = min(subspace_dim, dimension)
        values = np.zeros((count, dimension), dtype=np.float64)
        block = rng.standard_normal((count, active))
        norms = np.linalg.norm(block, axis=1, keepdims=True)
        values[:, -active:] = block / np.maximum(norms, np.finfo(float).tiny)
        return values

    if scheme == "geometric_axis":
        indices = rng.choice(dimension, size=count, p=eigenvalues)
        values = np.zeros((count, dimension), dtype=np.float64)
        signs = rng.choice(np.array([-1.0, 1.0]), size=count)
        values[np.arange(count), indices] = signs
        return values

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


def estimate_moment(
    rng: np.random.Generator,
    scheme: str,
    dimension: int,
    eigenvalues: np.ndarray,
    subspace_dim: int,
    feedback_samples: int,
    batch_size: int,
) -> np.ndarray:
    moment = np.zeros((dimension, dimension), dtype=np.float64)
    offset = 0
    while offset < feedback_samples:
        count = min(batch_size, feedback_samples - offset)
        feedback = sample_feedback(
            rng, scheme, count, dimension, eigenvalues, subspace_dim
        )
        moment += feedback.T @ feedback
        offset += count
    return moment / feedback_samples


def target_coverage_from_eigenvalues(
    rng: np.random.Generator, eigenvalues: np.ndarray, samples: int, batch_size: int
) -> np.ndarray:
    values = np.empty(samples, dtype=np.float64)
    dimension = len(eigenvalues)
    offset = 0
    while offset < samples:
        count = min(batch_size, samples - offset)
        g = rng.chisquare(df=1.0, size=(count, dimension))
        values[offset : offset + count] = (g @ eigenvalues) / np.sum(g, axis=1)
        offset += count
    return values


def summarize(
    scheme: str,
    dimension: int,
    target_samples: int,
    feedback_samples: int,
    population_eigs: np.ndarray,
    empirical_eigs: np.ndarray,
    predicted: np.ndarray,
    empirical: np.ndarray,
) -> MatchRow:
    if np.std(predicted) < 1e-14:
        ks_statistic = float("nan")
        ks_pvalue = float("nan")
    else:
        ks = stats.ks_2samp(predicted, empirical)
        ks_statistic = float(ks.statistic)
        ks_pvalue = float(ks.pvalue)
    return MatchRow(
        scheme=scheme,
        dimension=dimension,
        target_samples=target_samples,
        feedback_samples=feedback_samples,
        population_lambda_min=float(np.min(population_eigs)),
        empirical_lambda_min=float(np.min(empirical_eigs)),
        population_lambda_max=float(np.max(population_eigs)),
        empirical_lambda_max=float(np.max(empirical_eigs)),
        predicted_mean=float(np.mean(predicted)),
        empirical_mean=float(np.mean(empirical)),
        predicted_std=float(np.std(predicted, ddof=1)),
        empirical_std=float(np.std(empirical, ddof=1)),
        predicted_q01=float(np.quantile(predicted, 0.01)),
        empirical_q01=float(np.quantile(empirical, 0.01)),
        predicted_q50=float(np.quantile(predicted, 0.50)),
        empirical_q50=float(np.quantile(empirical, 0.50)),
        predicted_q99=float(np.quantile(predicted, 0.99)),
        empirical_q99=float(np.quantile(empirical, 0.99)),
        ks_2sample_statistic=ks_statistic,
        ks_2sample_pvalue=ks_pvalue,
    )


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,
    rows: list[MatchRow],
    spectra: dict[str, dict[str, list[float]]],
    outdir: Path,
) -> None:
    outdir.mkdir(parents=True, exist_ok=True)
    write_csv(outdir / "distribution_match.csv", rows)
    payload = {
        "config": asdict(config),
        "distribution_match": [asdict(row) for row in rows],
        "spectra": spectra,
    }
    (outdir / "summary.json").write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n"
    )


def save_plots(
    rows: list[MatchRow],
    distributions: dict[str, tuple[np.ndarray, np.ndarray]],
    spectra: dict[str, dict[str, list[float]]],
    outdir: Path,
) -> list[Path]:
    outdir.mkdir(parents=True, exist_ok=True)
    paths: list[Path] = []

    spectrum_path = outdir / "population_vs_empirical_spectra.png"
    plt.figure(figsize=(7, 4.5))
    for scheme, values in sorted(spectra.items()):
        empirical = np.array(values["empirical"])
        plt.plot(np.sort(empirical), label=scheme)
    plt.xlabel("eigenvalue index")
    plt.ylabel("empirical eigenvalue of M_mu")
    plt.title("Estimated feedback second-moment spectra")
    plt.legend()
    plt.tight_layout()
    plt.savefig(spectrum_path, dpi=180)
    plt.close()
    paths.append(spectrum_path)

    hist_path = outdir / "coverage_histograms.png"
    plt.figure(figsize=(8, 5))
    for scheme, (predicted, empirical) in sorted(distributions.items()):
        plt.hist(
            empirical,
            bins=80,
            density=True,
            histtype="step",
            linewidth=1.5,
            label=f"{scheme} empirical",
        )
    plt.xlabel("target coverage a^T M_mu a")
    plt.ylabel("density")
    plt.title("Coverage distributions from estimated M_mu")
    plt.legend(fontsize=8)
    plt.tight_layout()
    plt.savefig(hist_path, dpi=180)
    plt.close()
    paths.append(hist_path)

    for scheme, (predicted, empirical) in sorted(distributions.items()):
        if np.std(predicted) < 1e-14:
            continue
        qq_path = outdir / f"qq_{scheme}.png"
        probs = (np.arange(1, len(predicted) + 1) - 0.5) / len(predicted)
        pred_q = np.quantile(predicted, probs)
        emp_q = np.quantile(empirical, probs)
        max_value = float(max(pred_q[-1], emp_q[-1]))
        min_value = float(min(pred_q[0], emp_q[0]))
        plt.figure(figsize=(5, 5))
        plt.scatter(pred_q, emp_q, s=4, alpha=0.25)
        plt.plot([min_value, max_value], [min_value, max_value], color="black", linewidth=1)
        plt.xlabel("population-predicted quantile")
        plt.ylabel("empirical-estimated quantile")
        plt.title(f"Q-Q coverage: {scheme}")
        plt.tight_layout()
        plt.savefig(qq_path, dpi=180)
        plt.close()
        paths.append(qq_path)

    worst_path = outdir / "worst_best_coverage.png"
    names = [row.scheme for row in rows]
    x = np.arange(len(names))
    plt.figure(figsize=(7, 4.5))
    plt.bar(
        x - 0.18,
        [row.empirical_lambda_min for row in rows],
        width=0.36,
        label="worst target",
    )
    plt.bar(
        x + 0.18,
        [row.empirical_lambda_max for row in rows],
        width=0.36,
        label="best target",
    )
    plt.axhline(1.0 / rows[0].dimension, color="black", linestyle="--", linewidth=1, label="1/D")
    plt.xticks(x, names, rotation=25, ha="right")
    plt.ylabel("expected squared alignment")
    plt.title("Worst and best target coverage")
    plt.legend()
    plt.tight_layout()
    plt.savefig(worst_path, dpi=180)
    plt.close()
    paths.append(worst_path)

    return paths


def parse_config(args: argparse.Namespace) -> RunConfig:
    return RunConfig(
        dimension=args.dimension,
        target_samples=args.target_samples,
        feedback_samples=args.feedback_samples,
        batch_size=args.batch_size,
        seed=args.seed,
        schemes=args.schemes,
        subspace_dim=args.subspace_dim,
        anisotropy=args.anisotropy,
        outdir=str(args.outdir),
        plot=args.plot,
    )


def main() -> None:
    args = parse_args()
    config = parse_config(args)
    validate_config(config)
    rng = np.random.default_rng(config.seed)

    rows: list[MatchRow] = []
    spectra: dict[str, dict[str, list[float]]] = {}
    distributions: dict[str, tuple[np.ndarray, np.ndarray]] = {}

    for scheme in config.schemes:
        population_eigs = population_eigenvalues(
            scheme, config.dimension, config.subspace_dim, config.anisotropy
        )
        moment = estimate_moment(
            rng,
            scheme,
            config.dimension,
            population_eigs,
            config.subspace_dim,
            config.feedback_samples,
            config.batch_size,
        )
        empirical_eigs = np.linalg.eigvalsh(moment)
        predicted = target_coverage_from_eigenvalues(
            rng, population_eigs, config.target_samples, config.batch_size
        )
        empirical = target_coverage_from_eigenvalues(
            rng, empirical_eigs, config.target_samples, config.batch_size
        )
        row = summarize(
            scheme,
            config.dimension,
            config.target_samples,
            config.feedback_samples,
            population_eigs,
            empirical_eigs,
            predicted,
            empirical,
        )
        rows.append(row)
        spectra[scheme] = {
            "population": population_eigs.tolist(),
            "empirical": empirical_eigs.tolist(),
        }
        distributions[scheme] = (predicted, empirical)
        print(
            f"{scheme}: "
            f"lambda_min_pop={row.population_lambda_min:.6g}, "
            f"lambda_min_emp={row.empirical_lambda_min:.6g}, "
            f"KS={row.ks_2sample_statistic:.6g}, "
            f"mean_pred={row.predicted_mean:.6g}, "
            f"mean_emp={row.empirical_mean:.6g}"
        )

    outdir = Path(config.outdir)
    write_outputs(config, rows, spectra, outdir)
    plot_paths = save_plots(rows, distributions, spectra, outdir) if config.plot else []

    finite_ks = [row.ks_2sample_statistic for row in rows if np.isfinite(row.ks_2sample_statistic)]
    max_ks = max(finite_ks) if finite_ks else float("nan")
    max_mean_error = max(abs(row.predicted_mean - row.empirical_mean) for row in rows)
    print(f"max_ks_2sample: {max_ks:.8g}")
    print(f"max_abs_mean_error: {max_mean_error:.8g}")
    print(f"distribution_match: {outdir / 'distribution_match.csv'}")
    for path in plot_paths:
        print(f"plot: {path}")


if __name__ == "__main__":
    main()