import json, math, random
import numpy as np
import torch
from torch import nn

SEED = 447
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(min(12, torch.get_num_threads()))


def free_moments(tau, K=3):
    """m_p = 1/p [w^(p-1)] (1+w)^p exp(-p*tau/(1+w))."""
    out = []
    for p in range(1, K + 1):
        n, a = p - 1, p * tau
        # exp(-a/(1+w)) = exp(-a) exp(sum_j>=1 a*(-1)^(j+1) w^j)
        q = np.zeros(n + 1); q[0] = 1.0
        b = np.zeros(n + 1)
        for j in range(1, n + 1):
            b[j] = a * ((-1) ** (j + 1))
        for k in range(1, n + 1):
            q[k] = sum(j * b[j] * q[k-j] for j in range(1, k + 1)) / k
        coeff = 0.0
        for j in range(n + 1):
            coeff += math.comb(p, j) * q[n-j]
        out.append(math.exp(-a) * coeff / p)
    return np.array(out)


def haar_unitary(n):
    z = (np.random.randn(n, n) + 1j*np.random.randn(n, n)) / np.sqrt(2)
    q, r = np.linalg.qr(z)
    d = np.diag(r)
    return q * (d / np.abs(d)).conj()


def product_check(n=48, taus=(.25, .5, 1., 2.), reps=25):
    rows = []
    for tau in taus:
        L = max(1, round(tau * n))
        vals, atom = [], []
        P = np.diag([1.0] * (n-1) + [0.0])
        for _ in range(reps):
            B = np.eye(n, dtype=complex)
            for _ in range(L):
                B = P @ haar_unitary(n) @ B
            eig = np.linalg.eigvalsh(B.conj().T @ B).real
            vals.append([np.mean(eig), np.mean(eig**2), np.mean(eig**3)])
            atom.append(np.mean(eig > 1 - 1e-9))
        emp = np.mean(vals, axis=0); target = free_moments(tau, 3)
        rows.append(dict(tau=tau, L=L, target=target.tolist(), empirical=emp.tolist(),
                         abs_err=np.abs(emp-target).tolist(), atom_emp=float(np.mean(atom)),
                         atom_target=max(0., 1.-tau)))
    return rows


class MLP(nn.Module):
    def __init__(self, d=16, width=32, depth=4, out=4):
        super().__init__()
        layers = []
        for i in range(depth):
            layers += [nn.Linear(d if i == 0 else width, width), nn.Tanh()]
        layers.append(nn.Linear(width, out))
        self.net = nn.Sequential(*layers)
    def forward(self, x): return self.net(x)


def jacobian_moments(model, x, K=2):
    """Exact per-example J^T J moments, with scalar output probes replaced by full J."""
    x = x.detach().requires_grad_(True)
    y = model(x)
    estimates = []
    for b in range(x.shape[0]):
        Jrows = []
        for o in range(y.shape[1]):
            g = torch.autograd.grad(y[b, o], x, retain_graph=True, create_graph=True)[0][b]
            Jrows.append(g)
        J = torch.stack(Jrows)
        G = J.T @ J
        pows = G
        estimates.append([torch.trace(pows) / x.shape[1]])
        for p in range(2, K + 1):
            pows = pows @ G
            estimates[-1].append(torch.trace(pows) / x.shape[1])
    return torch.stack([torch.stack(a) for a in estimates]).mean(0)


def train_variant(use_spec, steps=120):
    torch.manual_seed(SEED + (1 if use_spec else 0))
    d, out, width, depth = 16, 4, 32, 4
    model = MLP(d, width, depth, out)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3)
    x = torch.randn(64, d)
    true_w = torch.randn(d, out)
    labels = x @ true_w + .15 * torch.randn(64, out)
    tau = depth / width
    target = torch.tensor(free_moments(tau, 2), dtype=x.dtype)
    history = []
    for step in range(steps):
        pred = model(x)
        task = ((pred-labels)**2).mean()
        spec = torch.tensor(0.)
        if use_spec and step < steps // 3 and step % 8 == 0:
            moms = jacobian_moments(model, x[:8], 2)
            spec = ((torch.log(moms + 1e-5) - torch.log(target + 1e-5))**2).sum()
        loss = task + (0.03 * spec.clamp(max=10.0) if use_spec else 0.)
        opt.zero_grad(); loss.backward(); opt.step()
        if step in (0, 19, 39, 79, steps-1):
            hist_m = jacobian_moments(model, x[:8], 2).detach().numpy()
            history.append(dict(step=step, task=float(task), spec=float(spec), moments=hist_m.tolist()))
    final_m = jacobian_moments(model, x[:8], 2).detach().numpy()
    return dict(final_task=float(((model(x)-labels)**2).mean()), target=target.numpy().tolist(),
                final_moments=final_m.tolist(), history=history)


def main():
    result = dict(seed=SEED, product_check=product_check(),
                  baseline=train_variant(False), spectral=train_variant(True))
    with open('results.json', 'w') as f: json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))

if __name__ == '__main__':
    main()