import json, math, random, time
import numpy as np


def empty(k):
    return np.zeros(k, dtype=np.float64), np.zeros((k, k), dtype=np.float64), np.zeros(k, dtype=np.float64)


def update(state, v):
    u, A, Q = state
    v = np.asarray(v, dtype=np.float64)
    return (u + v,
            A + 0.5 * (np.outer(u, v) - np.outer(v, u)),
            Q + v * v)


def compose(left, right):
    u, A, Q = left
    v, B, R = right
    return (u + v,
            A + B + 0.5 * (np.outer(u, v) - np.outer(v, u)),
            Q + R)


def sequence(xs):
    s = empty(len(xs[0]))
    for x in xs:
        s = update(s, x)
    return s


def maxdiff(a, b):
    return max(float(np.max(np.abs(x - y))) for x, y in zip(a, b))


def math_checks():
    # Prediction 1: for orthogonal increments a e1 then b e2, A12=ab/2;
    # reversing order gives -ab/2 and therefore cancels exactly.
    reversal = []
    for a in [0.25, 0.5, 1.0, 2.0, 4.0]:
        forward = sequence([np.array([a, 0.0]), np.array([0.0, 1.0])])[1][0, 1]
        reverse = sequence([np.array([0.0, 1.0]), np.array([a, 0.0])])[1][0, 1]
        reversal.append({"a": a, "predicted_abs": a / 2.0,
                         "observed_abs": abs(forward),
                         "reverse_sum": forward + reverse})

    # Prediction 2: n repeated ordered pairs accumulate linearly: A12=n/2.
    accumulation = []
    for n in range(1, 9):
        xs = []
        for _ in range(n):
            xs.extend([np.array([1.0, 0.0]), np.array([0.0, 1.0])])
        observed = sequence(xs)[1][0, 1]
        accumulation.append({"n": n, "predicted": n / 2.0, "observed": observed})

    # Prediction 3: under x -> lambda*x, u has degree one while A,Q have degree two.
    rng = np.random.default_rng(4)
    xs = rng.normal(size=(7, 3))
    base = sequence(xs)
    dilation = []
    for lam in [0.25, 0.5, 1.0, 2.0, 4.0]:
        got = sequence(lam * xs)
        dilation.append({
            "lambda": lam,
            "u_ratio_to_lambda": float(np.linalg.norm(got[0]) / (np.linalg.norm(base[0]) * lam)),
            "A_ratio_to_lambda2": float(np.linalg.norm(got[1]) / (np.linalg.norm(base[1]) * lam * lam)),
            "Q_ratio_to_lambda2": float(np.linalg.norm(got[2]) / (np.linalg.norm(base[2]) * lam * lam)),
        })

    # Associativity/chunking prediction: recursively summarized chunks equal token scan.
    rng = np.random.default_rng(11)
    xs = rng.normal(size=(13, 5))
    token = sequence(xs)
    chunks = [sequence(xs[:4]), sequence(xs[4:9]), sequence(xs[9:])]
    merged = compose(compose(chunks[0], chunks[1]), chunks[2])
    return {"reversal_sweep": reversal, "linear_accumulation": accumulation,
            "dilation_sweep": dilation, "chunk_max_abs_error": maxdiff(token, merged)}


def benchmark_impl(seed=123, epochs=80):
    # Task deliberately requires order: label is whether the first symbol precedes
    # the second symbol, while the multiset of symbols is identical.
    import torch
    import torch.nn as nn
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        torch.zeros(1, device=device)
    except Exception:
        device = "cpu"
    ntrain, ntest, length, vocab, k = 768, 256, 12, 8, 8
    rng = np.random.default_rng(seed)
    def make(n):
        x = np.zeros((n, length), dtype=np.int64); y=np.zeros(n, dtype=np.int64)
        for i in range(n):
            a,b=rng.choice(vocab, 2, replace=False)
            x[i,:]=rng.integers(0,vocab,length)
            pos=rng.choice(length,2,replace=False); x[i,pos[0]]=a; x[i,pos[1]]=b
            y[i]=int(pos[0] < pos[1])
        return torch.tensor(x), torch.tensor(y)
    tr_x,tr_y=make(ntrain); te_x,te_y=make(ntest)

    class Additive(nn.Module):
        def __init__(self):
            super().__init__(); self.emb=nn.Embedding(vocab,k); self.head=nn.Linear(k,2)
        def forward(self,x): return self.head(self.emb(x).sum(1))

    class Levy(nn.Module):
        def __init__(self):
            super().__init__(); self.emb=nn.Embedding(vocab,k); self.head=nn.Linear(k+k*(k-1)//2+k,2)
        def forward(self,x):
            z=torch.tanh(self.emb(x)); u=torch.zeros(x.size(0),k,device=x.device); A=torch.zeros(x.size(0),k,k,device=x.device); q=torch.zeros_like(u)
            for t in range(length):
                v=z[:,t]; A=A+0.5*(u[:,:,None]*v[:,None,:]-v[:,:,None]*u[:,None,:]); u=u+v; q=q+v*v
            iu=torch.triu_indices(k,k,1,device=x.device); feat=torch.cat((u/math.sqrt(length),A[:,iu[0],iu[1]]/length,q/length),1)
            return self.head(feat)

    class GRUModel(nn.Module):
        def __init__(self):
            super().__init__(); self.emb=nn.Embedding(vocab,k); self.gru=nn.GRU(k,k,batch_first=True); self.head=nn.Linear(k,2)
        def forward(self,x): return self.head(self.gru(self.emb(x))[0][:,-1])

    out={}
    for name, cls in [("additive",Additive),("levy",Levy),("gru",GRUModel)]:
        torch.manual_seed(seed)
        model=cls().to(device); opt=torch.optim.Adam(model.parameters(),lr=0.02); lossfn=nn.CrossEntropyLoss()
        tx,ty=tr_x.to(device),tr_y.to(device); vx,vy=te_x.to(device),te_y.to(device)
        t0=time.time()
        for _ in range(epochs):
            opt.zero_grad(); loss=lossfn(model(tx),ty); loss.backward(); opt.step()
        with torch.no_grad():
            pred=model(vx).argmax(1); acc=float((pred==vy).float().mean().cpu())
        out[name]={"test_accuracy":acc,"final_train_loss":float(loss.detach().cpu()),"seconds":time.time()-t0,"parameters":sum(p.numel() for p in model.parameters())}
    out["device"]=device
    return out


def benchmark(seed=123, epochs=80):
    import torch
    try:
        return benchmark_impl(seed, epochs)
    except Exception as exc:
        # cuDNN/GPU allocation can fail on the shared device; rerun identically on CPU.
        print("CUDA benchmark failed; falling back to CPU:", repr(exc))
        old = torch.cuda.is_available
        torch.cuda.is_available = lambda: False
        try:
            return benchmark_impl(seed, epochs)
        finally:
            torch.cuda.is_available = old


if __name__ == "__main__":
    result={"math":math_checks(),"benchmark":benchmark()}
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))