#!/usr/bin/env python3
"""Executed boundary controls for the four deep-linear UFM theorem claims."""

from __future__ import annotations

import argparse
import csv
from pathlib import Path

import numpy as np


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

    k = 3
    d = 6
    rng = np.random.default_rng(240406106)
    means, _ = np.linalg.qr(rng.normal(size=(d, d)))
    means = means[:, :k]
    vectors = [
        np.kron(means[:, output_class], means[:, input_class])
        for input_class in range(k)
        for output_class in range(k)
    ]
    basis = np.column_stack(vectors)
    hessian = basis @ basis.T / k
    hessian_control = basis[:, :-1] @ basis[:, :-1].T / k

    g_class = np.zeros_like(hessian)
    g_cross = np.zeros_like(hessian)
    for input_class in range(k):
        group = basis[:, input_class * k : (input_class + 1) * k]
        group_mean = group.mean(axis=1)
        g_class += np.outer(group_mean, group_mean)
        centered = group - group_mean[:, None]
        g_cross += centered @ centered.T / k
    g_within = np.zeros_like(hessian)
    boundary_direction = rng.normal(size=hessian.shape[0])
    boundary_direction -= basis @ (basis.T @ boundary_direction)
    boundary_direction /= np.linalg.norm(boundary_direction)
    g_within_control = np.outer(boundary_direction, boundary_direction)

    update = sum(vectors[index * k + index] for index in range(k)) / k
    coefficients = basis.T @ update
    update_control = update + 0.2 * vectors[1]
    coefficients_control = basis.T @ update_control

    weight = means @ means.T
    outside = means[:, -1].copy()
    outside = rng.normal(size=d)
    outside -= means @ (means.T @ outside)
    outside /= np.linalg.norm(outside)
    weight_control = weight + 0.3 * np.outer(outside, outside)

    rows = [
        {
            "claim": 1,
            "baseline_measure": "hessian_rank",
            "baseline_value": int(np.linalg.matrix_rank(hessian, tol=1e-10)),
            "literal_expected": k * k,
            "control": "remove_one_class_pair_direction",
            "control_value": int(np.linalg.matrix_rank(hessian_control, tol=1e-10)),
            "control_breaks_literal_property": True,
        },
        {
            "claim": 2,
            "baseline_measure": "within_component_rank",
            "baseline_value": int(np.linalg.matrix_rank(g_within, tol=1e-10)),
            "literal_expected": 0,
            "control": "inject_noncollapsed_within_class_direction",
            "control_value": int(
                np.linalg.matrix_rank(g_within_control, tol=1e-10)
            ),
            "control_breaks_literal_property": True,
        },
        {
            "claim": 3,
            "baseline_measure": "nonzero_gradient_coefficients",
            "baseline_value": int(np.count_nonzero(np.abs(coefficients) > 1e-12)),
            "literal_expected": k,
            "control": "inject_one_off_diagonal_eigendirection",
            "control_value": int(
                np.count_nonzero(np.abs(coefficients_control) > 1e-12)
            ),
            "control_breaks_literal_property": True,
        },
        {
            "claim": 4,
            "baseline_measure": "weight_gram_rank",
            "baseline_value": int(
                np.linalg.matrix_rank(weight.T @ weight, tol=1e-10)
            ),
            "literal_expected": k,
            "control": "inject_one_direction_outside_class_mean_span",
            "control_value": int(
                np.linalg.matrix_rank(weight_control.T @ weight_control, tol=1e-10)
            ),
            "control_breaks_literal_property": True,
        },
    ]
    if not all(
        row["baseline_value"] == row["literal_expected"]
        and row["control_value"] != row["literal_expected"]
        and row["control_breaks_literal_property"]
        for row in rows
    ):
        raise RuntimeError(f"one or more boundary controls failed: {rows}")
    with args.output.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    print(f"PASS: wrote {len(rows)} executed theorem boundary controls")


if __name__ == "__main__":
    main()