#!/usr/bin/env python3
"""Frozen epoch-zero control for the million-epoch Deep-UFM run."""

from __future__ import annotations

import argparse
import io
import json
import zipfile
from pathlib import Path

import numpy as np
import torch

from verify_native_relu_ufm import D, K, N, N_PER_CLASS, analyse

INITIAL_STD = 0.1


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--seed", type=int, default=71)
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)

    torch.manual_seed(args.seed)
    arrays = {
        "Y": torch.eye(K, dtype=torch.float32)
        .repeat_interleave(N_PER_CLASS, dim=1)
        .numpy(),
        "H1": (torch.randn(D, N, dtype=torch.float32) * INITIAL_STD).numpy(),
    }
    for index in range(1, 5):
        arrays[f"W{index}"] = (
            torch.randn(D, D, dtype=torch.float32) * INITIAL_STD
        ).numpy()
    arrays["W5"] = (
        torch.randn(K, D, dtype=torch.float32) * INITIAL_STD
    ).numpy()
    state = args.output / "epoch_zero_state.npz"
    with zipfile.ZipFile(state, "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())
    oracle = analyse(state)
    control = {
        "control": (
            "same source-scale architecture, target, seed and initialization "
            "before any of the registered gradient-descent epochs"
        ),
        "oracle": oracle,
        "training_accuracy_below_one": oracle["fit"]["accuracy"] < 1.0,
        "nine_outlier_claim_absent_at_initialization": not oracle["hessian"][
            "nine_outlier_gate"
        ],
    }
    control["destructive_control_pass"] = bool(
        control["training_accuracy_below_one"]
        and control["nine_outlier_claim_absent_at_initialization"]
    )
    (args.output / "epoch_zero_control.json").write_text(
        json.dumps(control, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(control, indent=2, sort_keys=True))
    if not control["destructive_control_pass"]:
        raise SystemExit(2)


if __name__ == "__main__":
    main()