#!/usr/bin/env python3
"""Source-scale reproduction of the nonlinear Deep-UFM experiment.

The paper fixes K=3, d=65, n=40, L=5, the fourth hidden layer, normal
initialisation, full-batch gradient descent, and 10^6 epochs.  It does not
publish the random seed, learning rate, regularisation coefficient, or
initialisation variance.  Those otherwise-unregistered choices are frozen
below and emitted in the result artifact.

The update is an explicit back-propagation implementation of the exact
registered MSE + L2 objective.  ``autograd_equivalence`` compares every
explicit gradient with PyTorch autograd before the million-epoch run.
"""

from __future__ import annotations

import argparse
import hashlib
import io
import json
import platform
import time
import zipfile
from pathlib import Path

import numpy as np
import torch

K = 3
D = 65
N_PER_CLASS = 40
N = K * N_PER_CLASS
DEPTH = 5
LAYER = 4
EPOCHS = 1_000_000
LEARNING_RATE = 0.01
WEIGHT_DECAY = 5e-4
INITIAL_STD = 0.1
CHECKPOINTS = {0, 1_000, 10_000, 100_000, 1_000_000}


def sha256(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            h.update(block)
    return h.hexdigest()


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


def gradients_explicit(
    h: torch.Tensor, weights: list[torch.Tensor], target: torch.Tensor
) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]:
    output, activations, preactivations = forward_explicit(h, weights)
    residual_over_n = (output - target) / N
    gradients: list[torch.Tensor] = [torch.empty_like(w) for w in weights]
    gradients[-1] = (
        residual_over_n @ activations[-1].T + WEIGHT_DECAY * weights[-1]
    )
    delta = (weights[-1].T @ residual_over_n) * (preactivations[-1] > 0)
    for index in range(len(weights) - 2, -1, -1):
        gradients[index] = (
            delta @ activations[index].T + WEIGHT_DECAY * weights[index]
        )
        if index:
            delta = (weights[index].T @ delta) * (
                preactivations[index - 1] > 0
            )
    gradient_h = weights[0].T @ delta + WEIGHT_DECAY * h
    return gradient_h, gradients, output


def objective(
    h: torch.Tensor, weights: list[torch.Tensor], target: torch.Tensor
) -> torch.Tensor:
    output, _, _ = forward_explicit(h, weights)
    value = 0.5 * (output - target).square().sum() / N
    value = value + 0.5 * WEIGHT_DECAY * h.square().sum()
    for weight in weights:
        value = value + 0.5 * WEIGHT_DECAY * weight.square().sum()
    return value


def autograd_equivalence(
    h: torch.Tensor, weights: list[torch.Tensor], target: torch.Tensor
) -> dict:
    h_ref = h.detach().clone().requires_grad_(True)
    weights_ref = [
        weight.detach().clone().requires_grad_(True) for weight in weights
    ]
    loss = objective(h_ref, weights_ref, target)
    loss.backward()
    explicit_h, explicit_weights, _ = gradients_explicit(h, weights, target)
    errors = [
        float((explicit_h - h_ref.grad).abs().max().detach().cpu())
    ]
    errors.extend(
        float((actual - reference.grad).abs().max().detach().cpu())
        for actual, reference in zip(explicit_weights, weights_ref)
    )
    return {
        "objective": float(loss.detach().cpu()),
        "max_abs_gradient_error": max(errors),
        "per_parameter_max_abs_error": errors,
        "tolerance": 2e-6,
        "pass": max(errors) <= 2e-6,
    }


def checkpoint_metrics(
    epoch: int,
    h: torch.Tensor,
    weights: list[torch.Tensor],
    target: torch.Tensor,
) -> dict:
    output, activations, preactivations = forward_explicit(h, weights)
    mse = float((0.5 * (output - target).square().sum() / N).detach().cpu())
    objective_value = float(objective(h, weights, target).detach().cpu())
    prediction = output.argmax(dim=0)
    truth = target.argmax(dim=0)
    accuracy = float((prediction == truth).float().mean().detach().cpu())
    active = [
        float((preactivation > 0).float().mean().detach().cpu())
        for preactivation in preactivations
    ]
    means = activations[LAYER].reshape(D, K, N_PER_CLASS).mean(dim=2)
    within = activations[LAYER].reshape(D, K, N_PER_CLASS) - means[:, :, None]
    within_norm = float(within.square().mean().sqrt().detach().cpu())
    mean_norm = float(means.square().mean().sqrt().detach().cpu())
    return {
        "epoch": epoch,
        "objective": objective_value,
        "unregularized_mse": mse,
        "training_accuracy": accuracy,
        "relu_active_fractions": active,
        "layer4_within_class_rms": within_norm,
        "layer4_class_mean_rms": mean_norm,
        "layer4_within_to_mean_ratio": within_norm / max(mean_norm, 1e-30),
    }


def save_state(
    path: Path, h: torch.Tensor, weights: list[torch.Tensor], target: torch.Tensor
) -> None:
    arrays = {
        "H1": h.detach().cpu().numpy(),
        "Y": target.detach().cpu().numpy(),
    }
    arrays.update(
        {f"W{index + 1}": weight.detach().cpu().numpy()
         for index, weight in enumerate(weights)}
    )
    with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_STORED) as archive:
        for name, array in arrays.items():
            payload = io.BytesIO()
            np.lib.format.write_array(
                payload, np.asanyarray(array), allow_pickle=False
            )
            info = zipfile.ZipInfo(f"{name}.npy", (1980, 1, 1, 0, 0, 0))
            info.compress_type = zipfile.ZIP_STORED
            info.external_attr = 0o600 << 16
            archive.writestr(info, payload.getvalue())


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--seed", type=int, default=71)
    parser.add_argument("--epochs", type=int, default=EPOCHS)
    parser.add_argument(
        "--device",
        choices=("auto", "mps", "cpu"),
        default="auto",
    )
    args = parser.parse_args()
    if args.epochs != EPOCHS:
        raise RuntimeError("release run must execute the registered 1,000,000 epochs")
    args.output.mkdir(parents=True, exist_ok=True)
    device = (
        "mps"
        if args.device == "auto" and torch.backends.mps.is_available()
        else "cpu"
        if args.device == "auto"
        else args.device
    )
    if device == "mps" and not torch.backends.mps.is_available():
        raise RuntimeError("MPS requested but unavailable")

    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)
    )

    equivalence = autograd_equivalence(h, weights, target)
    if not equivalence["pass"]:
        raise RuntimeError(f"explicit gradient failed autograd check: {equivalence}")

    started = time.time()
    checkpoints = [checkpoint_metrics(0, h, weights, target)]
    with torch.no_grad():
        for epoch in range(1, args.epochs + 1):
            gradient_h, gradients, _ = gradients_explicit(h, weights, target)
            h -= LEARNING_RATE * gradient_h
            for weight, gradient in zip(weights, gradients):
                weight -= LEARNING_RATE * gradient
            if epoch in CHECKPOINTS:
                if device == "mps":
                    torch.mps.synchronize()
                checkpoints.append(
                    checkpoint_metrics(epoch, h, weights, target)
                )
                print(
                    json.dumps(
                        {
                            "epoch": epoch,
                            "objective": checkpoints[-1]["objective"],
                            "accuracy": checkpoints[-1]["training_accuracy"],
                            "elapsed_seconds": time.time() - started,
                        }
                    ),
                    flush=True,
                )

    if device == "mps":
        torch.mps.synchronize()
    state_path = args.output / "final_state.npz"
    save_state(state_path, h, weights, target)
    result = {
        "paper": {
            "openreview_id": "RwiGcN2feP",
            "title": "Unifying Low Dimensional Spectra in Deep Learning",
            "source_revision": "arXiv:2404.06106v1",
            "literal_claim": (
                "In the non-linear (ReLU) Deep UFM, K^2=9 Hessian "
                "outliers separate but do not fully converge to equal values, "
                "and the gradient has K non-zero coefficients that remain "
                "unequal, unlike the linear case (Figure 9, Table 2)."
            ),
        },
        "registered_configuration": {
            "K": K,
            "d": D,
            "n_per_class": N_PER_CLASS,
            "training_examples": N,
            "L": DEPTH,
            "audited_layer_l": LAYER,
            "activation": "ReLU on W1 through W4; W5 linear",
            "optimizer": "full-batch gradient descent",
            "epochs": args.epochs,
            "normal_initialization": True,
        },
        "source_omissions_frozen_by_reproduction": {
            "seed": args.seed,
            "learning_rate": LEARNING_RATE,
            "l2_coefficient_all_weights_and_H1": WEIGHT_DECAY,
            "normal_initialization_standard_deviation": INITIAL_STD,
        },
        "implementation": {
            "device": device,
            "dtype": "float32",
            "explicit_update_equivalence_to_autograd": equivalence,
            "python": platform.python_version(),
            "torch": torch.__version__,
            "platform": platform.platform(),
        },
        "checkpoints": checkpoints,
        "final_state_sha256": sha256(state_path),
    }
    (args.output / "training_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 __name__ == "__main__":
    main()