#!/usr/bin/env python3
"""Paper-scale Deep Linear UFM trajectory for registered Claim 5."""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path

import numpy as np
import torch

K = 3
D = 60
N_PER_CLASS = 40
N = K * N_PER_CLASS
DEPTH = 5
LAYER = 3
EPOCHS = 1_000_000
LEARNING_RATE = 0.01
WEIGHT_DECAY = 5e-4
INITIAL_STD = 0.1
CHECKPOINTS = {
    0,
    50,
    100,
    250,
    500,
    1_000,
    2_500,
    5_000,
    10_000,
    20_000,
    40_000,
    100_000,
    250_000,
    500_000,
    1_000_000,
}


def forward(h: torch.Tensor, weights: list[torch.Tensor]) -> tuple[torch.Tensor, list[torch.Tensor]]:
    activations = [h]
    x = h
    for weight in weights:
        x = weight @ x
        activations.append(x)
    return x, activations


def gradients(
    h: torch.Tensor, weights: list[torch.Tensor], target: torch.Tensor
) -> tuple[torch.Tensor, list[torch.Tensor]]:
    output, activations = forward(h, weights)
    delta = (output - target) / N
    gradients: list[torch.Tensor] = [torch.empty_like(weight) for weight in weights]
    for index in range(len(weights) - 1, -1, -1):
        gradients[index] = delta @ activations[index].T + WEIGHT_DECAY * weights[index]
        delta = weights[index].T @ delta
    return delta + WEIGHT_DECAY * h, gradients


def metrics(
    epoch: int, h: torch.Tensor, weights: list[torch.Tensor], target: torch.Tensor
) -> dict:
    output, activations = forward(h, weights)
    input_features = activations[LAYER - 1]
    output_features = activations[LAYER]
    tail = weights[-1]
    for index in range(len(weights) - 2, LAYER - 1, -1):
        tail = tail @ weights[index]
    left = tail.T @ tail
    right = input_features @ input_features.T / N
    with torch.no_grad():
        left_values = torch.linalg.eigvalsh(left).detach().cpu().numpy()
        right_values = torch.linalg.eigvalsh(right).detach().cpu().numpy()
    hessian_values = np.sort(np.outer(left_values, right_values).reshape(-1))[::-1]
    top9 = hessian_values[: K * K]

    means_in = input_features.reshape(D, K, N_PER_CLASS).mean(dim=2)
    means_out = output_features.reshape(D, K, N_PER_CLASS).mean(dim=2)
    alignments: list[float] = []
    for output_class in range(K):
        u = means_out[:, output_class]
        lu = left @ u
        for input_class in range(K):
            v = means_in[:, input_class]
            rv = right @ v
            numerator = (u @ lu).square() * (v @ rv).square()
            denominator = (
                u.square().sum()
                * lu.square().sum()
                * v.square().sum()
                * rv.square().sum()
            )
            alignments.append(float((numerator / denominator.clamp_min(1e-30)).cpu()))
    residual = output - target
    return {
        "epoch": epoch,
        "objective": float(
            (
                0.5 * residual.square().sum() / N
                + 0.5 * WEIGHT_DECAY * h.square().sum()
                + sum(0.5 * WEIGHT_DECAY * weight.square().sum() for weight in weights)
            ).cpu()
        ),
        "training_accuracy": float(
            (output.argmax(dim=0) == target.argmax(dim=0)).float().mean().cpu()
        ),
        "top9_max_to_min_ratio": float(top9[0] / max(top9[-1], 1e-30)),
        "ninth_to_tenth_ratio": float(top9[-1] / max(hessian_values[9], 1e-30)),
        "mean_alignment": float(np.mean(alignments)),
        "minimum_alignment": float(np.min(alignments)),
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--seed", type=int, default=53)
    parser.add_argument("--device", choices=("mps", "cpu"), default="mps")
    args = parser.parse_args()
    if args.device == "mps" and not torch.backends.mps.is_available():
        raise RuntimeError("MPS is unavailable")
    args.output.mkdir(parents=True, exist_ok=True)
    device = args.device
    torch.manual_seed(args.seed)
    target = torch.eye(K, dtype=torch.float32).repeat_interleave(N_PER_CLASS, dim=1).to(device)
    h = (torch.randn(D, N, dtype=torch.float32) * INITIAL_STD).to(device)
    weights = [
        (torch.randn(D, D, dtype=torch.float32) * INITIAL_STD).to(device)
        for _ in range(DEPTH - 1)
    ]
    weights.append((torch.randn(K, D, dtype=torch.float32) * INITIAL_STD).to(device))

    rows = [metrics(0, h, weights, target)]
    with torch.no_grad():
        for epoch in range(1, EPOCHS + 1):
            gradient_h, gradient_weights = gradients(h, weights, target)
            h -= LEARNING_RATE * gradient_h
            for weight, gradient in zip(weights, gradient_weights):
                weight -= LEARNING_RATE * gradient
            if epoch in CHECKPOINTS:
                rows.append(metrics(epoch, h, weights, target))
    with (args.output / "linear_native_trajectory.csv").open(
        "w", encoding="utf-8", newline=""
    ) as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    result = {
        "paper": "RwiGcN2feP",
        "registered_configuration": {
            "K": K,
            "d": D,
            "n_per_class": N_PER_CLASS,
            "L": DEPTH,
            "audited_layer_l": LAYER,
            "normal_initialization": True,
            "optimizer": "full-batch gradient descent",
        },
        "frozen_source_omissions": {
            "seed": args.seed,
            "epochs": EPOCHS,
            "learning_rate": LEARNING_RATE,
            "weight_decay": WEIGHT_DECAY,
            "initialization_standard_deviation": INITIAL_STD,
        },
        "checkpoints": rows,
        "literal_gates": {
            "nine_outliers_separate": rows[-1]["ninth_to_tenth_ratio"] >= 3.0,
            "top9_converge_near_equality": rows[-1]["top9_max_to_min_ratio"] <= 1.1,
            "alignment_converges_to_one": rows[-1]["minimum_alignment"] >= 0.99,
            "initial_alignment_is_not_about_point_two": rows[0]["mean_alignment"] < 0.1,
        },
    }
    result["all_literal_gates_pass"] = all(result["literal_gates"].values())
    (args.output / "linear_native_results.json").write_text(
        json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(json.dumps(result, indent=2, sort_keys=True))
    if not result["all_literal_gates_pass"]:
        raise SystemExit(2)


if __name__ == "__main__":
    main()