import json, math, random
from pathlib import Path
import numpy as np

SEED = 2327
np.random.seed(SEED)
random.seed(SEED)


def frob(x):
    return float(np.linalg.norm(x, 'fro'))


def comm(a, b):
    return a @ b - b @ a


def run_math():
    n = 6
    sigma = np.arange(n)[::-1]
    kappa = np.array([1.0, 1.7, 0.8, 1.3, 2.0, 0.6])
    G = np.zeros((n, n))
    G[np.arange(n), sigma] = kappa

    # A matrix polynomial in G commutes exactly with weighted reflection.
    P0 = np.eye(n) + 0.055 * G + 0.009 * (G @ G)
    rng = np.random.default_rng(SEED + 1)
    E = rng.normal(size=(n, n))
    E /= frob(E)
    C0 = frob(comm(G, E))

    # Prediction 1: ||[G,P0+eps E]|| is exactly linear in eps.
    epses = np.array([0.0, .001, .003, .01, .03, .1, .3])
    cvals = np.array([frob(comm(G, P0 + e * E)) for e in epses])
    slope = float(np.dot(epses, cvals) / max(np.dot(epses, epses), 1e-30))
    linear_rel_err = float(np.max(np.abs(cvals - epses * C0)) / max(cvals[-1], 1e-30))

    # Prediction 2: exact commutation propagates to powers and resolvents.
    power_errors = []
    Ppow = np.eye(n)
    for t in range(0, 9):
        power_errors.append(frob(comm(G, Ppow)))
        Ppow = Ppow @ P0
    gamma0 = 0.37
    R0 = np.linalg.inv(np.eye(n) - (1 - gamma0) * P0)
    resolvent_zero = frob(comm(G, R0))

    # Prediction 3: [G,R] = (1-gamma) R [G,P] R; it vanishes at gamma=1
    # and is approximately linear in alpha=1-gamma near alpha=0.
    P1 = P0 + 0.08 * E
    gammas = np.array([0.0, .1, .25, .5, .7, .85, .95, 1.0])
    rvals, identity_errors = [], []
    for gamma in gammas:
        alpha = 1.0 - gamma
        R = np.linalg.inv(np.eye(n) - alpha * P1)
        lhs = comm(G, R)
        rhs = alpha * R @ comm(G, P1) @ R
        rvals.append(frob(lhs))
        identity_errors.append(frob(lhs-rhs) / max(frob(lhs), 1e-30))
    rvals = np.array(rvals)
    alpha = 1-gammas
    near = alpha <= .3
    through_origin = float(np.dot(alpha[near], rvals[near]) / max(np.dot(alpha[near], alpha[near]), 1e-30))
    # report normalized values so the damping trend is easy to inspect
    normalized = (rvals / max(rvals[0], 1e-30)).tolist()

    return {
        'weighted_G': G.tolist(),
        'prediction_1_linear_perturbation': {
            'eps': epses.tolist(), 'observed_commutators': cvals.tolist(),
            'predicted_slope_norm_[G,E]': C0, 'fit_slope': slope,
            'max_relative_deviation': linear_rel_err
        },
        'prediction_2_exact_multi_step': {
            'power_commutator_norms_t0_to_8': power_errors,
            'resolvent_commutator_norm_gamma_0.37': resolvent_zero
        },
        'prediction_3_damping_resolvent': {
            'gamma': gammas.tolist(), 'observed_norms': rvals.tolist(),
            'normalized_to_gamma_0': normalized,
            'resolvent_identity_relative_errors': identity_errors,
            'near_gamma_1_fit_slope_vs_(1-gamma)': through_origin,
            'gamma_1_norm': float(rvals[-1])
        }
    }


def attention_trial(comm_lambda, seed, steps=180):
    import torch
    torch.manual_seed(seed); np.random.seed(seed)
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        dev = torch.device(device)
        N, L, D = 768, 8, 8
        x = torch.randn(N, L, D, device=dev)
        # Reversal-invariant signal, with deliberately unequal noise scales by side.
        latent = x[:, :, 0].sum(1) + .35 * x[:, :, 1].sum(1)
        x[:, :L//2] += .45 * torch.randn(N, L//2, D, device=dev)
        x[:, L//2:] += 1.25 * torch.randn(N, L//2, D, device=dev)
        y = (latent > 0).long()
        tr, va = torch.arange(0, 600, device=dev), torch.arange(600, N, device=dev)
        q = torch.nn.Linear(D, D, bias=False).to(dev)
        k = torch.nn.Linear(D, D, bias=False).to(dev)
        v = torch.nn.Linear(D, D, bias=False).to(dev)
        out = torch.nn.Linear(D, 2).to(dev)
        opt = torch.optim.Adam(list(q.parameters())+list(k.parameters())+list(v.parameters())+list(out.parameters()), lr=.012)
        rev = torch.arange(L-1, -1, -1, device=dev)
        kap = torch.tensor([1., 1.35, .85, 1.15, 1.15, .85, 1.35, 1.], device=dev)
        G = torch.zeros(L,L,device=dev); G[torch.arange(L),rev] = kap
        last_loss = None
        for _ in range(steps):
            xb, yb = x[tr], y[tr]
            P = torch.softmax(q(xb) @ k(xb).transpose(1,2) / math.sqrt(D), dim=-1)
            z = (P @ v(xb)).mean(1)
            loss = torch.nn.functional.cross_entropy(out(z), yb)
            GP = G.unsqueeze(0) @ P; PG = P @ G.unsqueeze(0)
            penalty = ((GP-PG)**2).sum() / (G.square().sum()*P.square().sum().mean()+1e-8)
            total = loss + comm_lambda * penalty
            opt.zero_grad(); total.backward(); opt.step(); last_loss=float(loss.detach())
        with torch.no_grad():
            P = torch.softmax(q(x[va]) @ k(x[va]).transpose(1,2) / math.sqrt(D), dim=-1)
            pred = out((P @ v(x[va])).mean(1)).argmax(1)
            acc = float((pred == y[va]).float().mean())
            ce = float(torch.nn.functional.cross_entropy(out((P @ v(x[va])).mean(1)), y[va]))
            cm = torch.linalg.vector_norm(G@P-P@G).item() / math.sqrt(len(va))
        return {'accuracy':acc, 'val_loss':ce, 'commutator_rms':cm, 'device':str(dev), 'train_loss':last_loss}
    except Exception as exc:
        # CUDA can be shared; retrying on CPU is part of the experiment contract.
        if device == 'cuda':
            torch.cuda.empty_cache()
            old = torch.cuda.is_available
            # A clean CPU subprocess is simpler than mutating torch's device state.
            return attention_trial_cpu(comm_lambda, seed, steps)
        raise


def attention_trial_cpu(comm_lambda, seed, steps):
    import torch
    # Force CPU by temporarily using a CPU-only equivalent implementation.
    # The compact fallback duplicates the computation with tensors on CPU.
    torch.manual_seed(seed); np.random.seed(seed)
    N,L,D=768,8,8; dev=torch.device('cpu')
    x=torch.randn(N,L,D); latent=x[:,:,0].sum(1)+.35*x[:,:,1].sum(1)
    x[:,:4]+=.45*torch.randn(N,4,D); x[:,4:]+=1.25*torch.randn(N,4,D); y=(latent>0).long()
    q=torch.nn.Linear(D,D,bias=False); k=torch.nn.Linear(D,D,bias=False); v=torch.nn.Linear(D,D,bias=False); out=torch.nn.Linear(D,2)
    opt=torch.optim.Adam(list(q.parameters())+list(k.parameters())+list(v.parameters())+list(out.parameters()),lr=.012)
    rev=torch.arange(7,-1,-1); kap=torch.tensor([1.,1.35,.85,1.15,1.15,.85,1.35,1.]); G=torch.zeros(L,L); G[torch.arange(L),rev]=kap
    for _ in range(steps):
        P=torch.softmax(q(x[:600])@k(x[:600]).transpose(1,2)/math.sqrt(D),-1); z=(P@v(x[:600])).mean(1); loss=torch.nn.functional.cross_entropy(out(z),y[:600]); pen=((G@P-P@G)**2).sum()/(G.square().sum()*P.square().sum().mean()+1e-8); opt.zero_grad(); (loss+comm_lambda*pen).backward(); opt.step()
    with torch.no_grad():
        P=torch.softmax(q(x[600:])@k(x[600:]).transpose(1,2)/math.sqrt(D),-1); logits=out((P@v(x[600:])).mean(1)); return {'accuracy':float((logits.argmax(1)==y[600:]).float().mean()),'val_loss':float(torch.nn.functional.cross_entropy(logits,y[600:])),'commutator_rms':float(torch.linalg.vector_norm(G@P-P@G).item()/math.sqrt(168)),'device':'cpu','train_loss':float(loss)}


def main():
    math_result=run_math()
    attention={}
    for lam in (0.0, 1.0):
        vals=[attention_trial(lam,s) for s in (11,22)]
        attention[str(lam)]={'seeds':vals,'mean_accuracy':float(np.mean([v['accuracy'] for v in vals])),'mean_val_loss':float(np.mean([v['val_loss'] for v in vals])),'mean_commutator_rms':float(np.mean([v['commutator_rms'] for v in vals]))}
    result={'seed':SEED,'math':math_result,'attention_mini_experiment':attention}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()