#!/usr/bin/env python3
"""Validate random-subspace functional capacity overlap formulas.

Let S be a d-dimensional task-sensitive subspace in R^P and E be a k-dimensional
alignment constraint subspace. For generic subspaces:

  hard rank loss = max(0, k - (P - d))

and the soft overlap T = tr(P_E P_S) satisfies:

  E[T] = kd/P.
"""

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


@dataclass(frozen=True)
class RunConfig:
    parameters: int
    task_rank: int
    constraint_ranks: list[int]
    trials: int
    seed: int
    rank_tol: float
    outdir: str
    plot: bool


@dataclass(frozen=True)
class CapacityRow:
    constraint_rank: int
    hard_loss_empirical_mean: float
    hard_loss_empirical_std: float
    hard_loss_theory: float
    soft_overlap_empirical_mean: float
    soft_overlap_empirical_var: float
    soft_overlap_theory_mean: float
    soft_overlap_theory_var: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Validate redundancy-exhaustion capacity formulas."
    )
    parser.add_argument(
        "--parameters",
        type=int,
        default=128,
        help="Ambient parameter dimension P.",
    )
    parser.add_argument(
        "--task-rank",
        type=int,
        default=32,
        help="Task-sensitive subspace dimension d.",
    )
    parser.add_argument(
        "--constraint-ranks",
        type=int,
        nargs="+",
        default=[0, 16, 32, 64, 96, 112, 128],
        help="Alignment constraint ranks k to test.",
    )
    parser.add_argument("--trials", type=int, default=200, help="Monte Carlo trials.")
    parser.add_argument("--seed", type=int, default=0, help="Random seed.")
    parser.add_argument(
        "--rank-tol",
        type=float,
        default=1e-9,
        help="Numerical tolerance for rank estimation.",
    )
    parser.add_argument(
        "--outdir",
        type=Path,
        default=Path("outputs/functional_capacity_overlap"),
        help="Output directory.",
    )
    parser.add_argument("--plot", action="store_true", help="Save plots.")
    return parser.parse_args()


def random_orthonormal(
    rng: np.random.Generator, ambient_dim: int, subspace_dim: int
) -> np.ndarray:
    if subspace_dim == 0:
        return np.empty((ambient_dim, 0), dtype=np.float64)
    values = rng.standard_normal((ambient_dim, subspace_dim))
    q, _ = np.linalg.qr(values, mode="reduced")
    return q[:, :subspace_dim]


def hard_rank_loss(
    task_basis: np.ndarray, constraint_basis: np.ndarray, rank_tol: float
) -> int:
    task_rank = task_basis.shape[1]
    if constraint_basis.shape[1] == 0:
        return 0
    projected = task_basis - constraint_basis @ (constraint_basis.T @ task_basis)
    remaining_rank = np.linalg.matrix_rank(projected, tol=rank_tol)
    return int(task_rank - remaining_rank)


def soft_overlap(task_basis: np.ndarray, constraint_basis: np.ndarray) -> float:
    if task_basis.shape[1] == 0 or constraint_basis.shape[1] == 0:
        return 0.0
    cross = constraint_basis.T @ task_basis
    return float(np.sum(cross * cross))


def theory_hard_loss(parameters: int, task_rank: int, constraint_rank: int) -> int:
    return max(0, constraint_rank - (parameters - task_rank))


def theory_soft_variance(parameters: int, task_rank: int, constraint_rank: int) -> float:
    p = parameters
    d = task_rank
    k = constraint_rank
    if p <= 2:
        return 0.0
    numerator = 2 * k * d * (p - k) * (p - d)
    denominator = p * p * (p - 1) * (p + 2)
    return numerator / denominator


def run_trials(config: RunConfig) -> list[CapacityRow]:
    rng = np.random.default_rng(config.seed)
    rows: list[CapacityRow] = []

    for constraint_rank in config.constraint_ranks:
        hard_losses = np.empty(config.trials, dtype=np.float64)
        overlaps = np.empty(config.trials, dtype=np.float64)

        for trial in range(config.trials):
            task_basis = random_orthonormal(
                rng, config.parameters, config.task_rank
            )
            constraint_basis = random_orthonormal(
                rng, config.parameters, constraint_rank
            )
            hard_losses[trial] = hard_rank_loss(
                task_basis, constraint_basis, config.rank_tol
            )
            overlaps[trial] = soft_overlap(task_basis, constraint_basis)

        rows.append(
            CapacityRow(
                constraint_rank=constraint_rank,
                hard_loss_empirical_mean=float(np.mean(hard_losses)),
                hard_loss_empirical_std=float(np.std(hard_losses, ddof=1))
                if config.trials > 1
                else 0.0,
                hard_loss_theory=float(
                    theory_hard_loss(
                        config.parameters, config.task_rank, constraint_rank
                    )
                ),
                soft_overlap_empirical_mean=float(np.mean(overlaps)),
                soft_overlap_empirical_var=float(np.var(overlaps, ddof=1))
                if config.trials > 1
                else 0.0,
                soft_overlap_theory_mean=constraint_rank
                * config.task_rank
                / config.parameters,
                soft_overlap_theory_var=theory_soft_variance(
                    config.parameters, config.task_rank, constraint_rank
                ),
            )
        )

    return rows


def write_outputs(config: RunConfig, rows: list[CapacityRow], outdir: Path) -> None:
    outdir.mkdir(parents=True, exist_ok=True)
    payload = {
        "config": asdict(config),
        "rows": [asdict(row) for row in rows],
    }
    (outdir / "summary.json").write_text(
        json.dumps(payload, indent=2, sort_keys=True) + "\n"
    )

    with (outdir / "summary.csv").open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(asdict(rows[0]).keys()))
        writer.writeheader()
        for row in rows:
            writer.writerow(asdict(row))


def save_plots(rows: list[CapacityRow], outdir: Path) -> list[Path]:
    outdir.mkdir(parents=True, exist_ok=True)
    paths: list[Path] = []
    k = np.array([row.constraint_rank for row in rows])

    hard_path = outdir / "hard_rank_loss.png"
    plt.figure(figsize=(7, 4.5))
    plt.plot(
        k,
        [row.hard_loss_empirical_mean for row in rows],
        marker="o",
        label="empirical",
    )
    plt.plot(
        k,
        [row.hard_loss_theory for row in rows],
        linestyle="--",
        color="black",
        label="theory",
    )
    plt.xlabel("constraint rank k")
    plt.ylabel("hard functional rank loss")
    plt.title("Redundancy-exhaustion transition")
    plt.legend()
    plt.tight_layout()
    plt.savefig(hard_path, dpi=180)
    plt.close()
    paths.append(hard_path)

    soft_path = outdir / "soft_overlap.png"
    plt.figure(figsize=(7, 4.5))
    plt.plot(
        k,
        [row.soft_overlap_empirical_mean for row in rows],
        marker="o",
        label="empirical mean",
    )
    plt.plot(
        k,
        [row.soft_overlap_theory_mean for row in rows],
        linestyle="--",
        color="black",
        label="theory mean",
    )
    plt.xlabel("constraint rank k")
    plt.ylabel("soft overlap tr(P_E P_S)")
    plt.title("Soft functional capacity overlap")
    plt.legend()
    plt.tight_layout()
    plt.savefig(soft_path, dpi=180)
    plt.close()
    paths.append(soft_path)

    return paths


def main() -> None:
    args = parse_args()
    if args.parameters < 1:
        raise ValueError("--parameters must be positive.")
    if not 0 <= args.task_rank <= args.parameters:
        raise ValueError("--task-rank must be in [0, parameters].")
    if any(k < 0 or k > args.parameters for k in args.constraint_ranks):
        raise ValueError("All constraint ranks must be in [0, parameters].")
    if args.trials < 1:
        raise ValueError("--trials must be positive.")

    config = RunConfig(
        parameters=args.parameters,
        task_rank=args.task_rank,
        constraint_ranks=args.constraint_ranks,
        trials=args.trials,
        seed=args.seed,
        rank_tol=args.rank_tol,
        outdir=str(args.outdir),
        plot=args.plot,
    )

    rows = run_trials(config)
    write_outputs(config, rows, args.outdir)
    plot_paths = save_plots(rows, args.outdir) if args.plot else []

    print(f"parameters P: {args.parameters}")
    print(f"task rank d: {args.task_rank}")
    print(f"trials: {args.trials}")
    print(f"table: {args.outdir / 'summary.csv'}")
    for row in rows:
        print(
            f"k={row.constraint_rank}: "
            f"hard={row.hard_loss_empirical_mean:.6g} "
            f"(theory {row.hard_loss_theory:.6g}), "
            f"soft={row.soft_overlap_empirical_mean:.6g} "
            f"(theory {row.soft_overlap_theory_mean:.6g})"
        )
    for path in plot_paths:
        print(f"plot: {path}")


if __name__ == "__main__":
    main()