"""Claim 5: for a deep linear unconstrained-features model (UFM) with K = 3
classes, the loss Hessian has K^2 = 9 outlier eigenvalues separating from the
bulk.

The logbook states the figures remain source-reported with no independent
training run. A UFM needs no external data -- the features are free variables --
so it is trained here and the exact Hessian spectrum is computed.

Deep linear UFM: logits = W_L ... W_1 H, with H the free feature matrix
(d x N) and W_l the layer weights. Loss = MSE to one-hot targets + weight decay.
The Hessian is formed exactly by autograd over ALL parameters and diagonalised.
"""
import json, numpy as np, torch

RES = {}
torch.set_default_dtype(torch.float64)


def build(K=3, n=8, d=6, depth=2, seed=0):
    g = torch.Generator().manual_seed(seed)
    N = K*n
    Y = torch.zeros(K, N)
    for c in range(K): Y[c, c*n:(c+1)*n] = 1.0
    H = (torch.randn(d, N, generator=g)*0.3).requires_grad_(True)
    Ws = []
    dims = [K]+[d]*(depth-1)+[d]
    for i in range(depth):
        Ws.append((torch.randn(dims[i], dims[i+1], generator=g)*0.3).requires_grad_(True))
    return H, Ws, Y


def loss_fn(H, Ws, Y, wd=5e-4):
    Z = H
    for W in reversed(Ws): Z = W @ Z
    l = ((Z-Y)**2).mean()
    for W in Ws: l = l+wd*(W**2).sum()
    return l+wd*(H**2).sum()


def train(H, Ws, Y, steps=4000, lr=0.05):
    params = [H]+Ws
    opt = torch.optim.Adam(params, lr=lr)
    for _ in range(steps):
        opt.zero_grad(); l = loss_fn(H, Ws, Y); l.backward(); opt.step()
    return float(l.item())


def hessian_spectrum(H, Ws, Y):
    params = [H]+Ws
    flat = torch.cat([p.reshape(-1) for p in params]).detach()
    n = flat.numel()
    def f(v):
        i = 0; ps = []
        for p in params:
            k = p.numel(); ps.append(v[i:i+k].view_as(p)); i += k
        Z = ps[0]
        for W in reversed(ps[1:]): Z = W @ Z
        l = ((Z-Y)**2).mean()
        for W in ps[1:]: l = l+5e-4*(W**2).sum()
        return l+5e-4*(ps[0]**2).sum()
    Hs = torch.autograd.functional.hessian(f, flat)
    ev = torch.linalg.eigvalsh((Hs+Hs.T)/2).numpy()
    return np.sort(ev)[::-1]


def count_outliers(ev, K):
    """Outliers = eigenvalues above the largest gap in the top of the spectrum."""
    top = ev[:4*K*K]
    gaps = top[:-1]-top[1:]
    j = int(np.argmax(gaps))
    return j+1, float(gaps[j]), float(top[j]), float(top[j+1])


def main():
    rows = []
    for K in (2, 3, 4):
        for depth in (2, 3):
            H, Ws, Y = build(K=K, depth=depth, seed=K*10+depth)
            fin = train(H, Ws, Y)
            ev = hessian_spectrum(H, Ws, Y)
            n_out, gap, above, below = count_outliers(ev, K)
            rows.append({"K": K, "depth": depth, "final_loss": round(fin, 8),
                         "params": int(sum(p.numel() for p in [H]+Ws)),
                         "K_squared": K*K, "outliers_found": n_out,
                         "matches_K2": bool(n_out == K*K),
                         "gap_size": round(gap, 6),
                         "last_outlier": round(above, 6), "first_bulk": round(below, 6),
                         "separation_ratio": round(above/max(below, 1e-12), 2)})
            print("  K=%d depth=%d  loss=%.2e  params=%d  K^2=%d  outliers found=%d  match=%s  gap=%.4f (%.4f -> %.4f, ratio %.1fx)"
                  % (K, depth, fin, rows[-1]["params"], K*K, n_out, n_out == K*K,
                     gap, above, below, rows[-1]["separation_ratio"]), flush=True)
    RES["claim5_ufm_hessian"] = {"rows": rows,
        "matches_in": sum(r["matches_K2"] for r in rows), "cells": len(rows),
        "min_separation_ratio": min(r["separation_ratio"] for r in rows)}
    R = RES["claim5_ufm_hessian"]
    print("  outlier count equals K^2 in %d/%d cells; min separation ratio %.1fx"
          % (R["matches_in"], R["cells"], R["min_separation_ratio"]), flush=True)
    json.dump(RES, open("ufm_results.json", "w"), indent=1)


if __name__ == "__main__":
    main(); print("DONE")