import json, math, time
import numpy as np
import torch
from gated_attention import local_attention, global_linear_attention, dense_attention

torch.manual_seed(7)
np.random.seed(7)
torch.set_num_threads(4)
device = "cuda" if torch.cuda.is_available() else "cpu"

def graph(n, degree):
    src, dst = [], []
    for i in range(n):
        for t in range(degree):
            src.append(i); dst.append((i + t) % n)
    return torch.tensor(src, device=device), torch.tensor(dst, device=device)

def kernel_explicit(x, wq, wk, wv):
    q, k, v = x @ wq, x @ wk, x @ wv
    pq, pk = torch.nn.functional.elu(q) + 1, torch.nn.functional.elu(k) + 1
    weights = pq @ pk.T
    return (weights @ v) / (weights.sum(1, keepdim=True) + 1e-6)

def timed(fn, repeats=8):
    if device == "cuda": torch.cuda.synchronize()
    for _ in range(2): fn()
    if device == "cuda": torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(repeats): fn()
    if device == "cuda": torch.cuda.synchronize()
    return (time.perf_counter() - t0) / repeats * 1000

# Prediction 1: local attention is exactly dense attention with non-neighbors removed.
n, d, r = 31, 12, 7
x = torch.randn(n, d, device=device)
wq, wk, wv = [torch.randn(d, r, device=device) / math.sqrt(d) for _ in range(3)]
src, dst = graph(n, 4)
local = local_attention(x, src, dst, wq, wk, wv)
q, k, v = x @ wq, x @ wk, x @ wv
logits = torch.full((n, n), -torch.inf, device=device)
logits[src, dst] = (q[src] * k[dst]).sum(-1) / math.sqrt(r)
masked = torch.softmax(logits, dim=-1) @ v
mask_error = (local - masked).abs().max().item()

# Prediction 2: feature-map linear attention is the exact reassociation of its
# pairwise kernel, up to the stated epsilon in the denominator.
explicit = kernel_explicit(x, wq, wk, wv)
linear = global_linear_attention(x, wq, wk, wv)
linear_error = (explicit - linear).abs().max().item()

# Prediction 3: operation proxy has dense N^2*r versus E*r + N*r^2;
# for fixed degree, dense/hybrid proxy ratio grows linearly with N.
width = 16
sizes = [64, 128, 256, 512]
degree = 8
ratios, slopes = [], []
for nn in sizes:
    e = nn * degree
    dense_ops = nn * nn * width
    hybrid_ops = e * width + nn * width * width
    ratios.append(dense_ops / hybrid_ops)
# log-log slope of ratio versus N; predicted slope approaches 1 when N >> r.
slope = float(np.polyfit(np.log(sizes), np.log(ratios), 1)[0])

# Small learning task: target combines one-hop neighbor mean and global mean.
# Dense baseline uses full attention; hybrid uses local+global branches with a
# learned gate. Both have the same projections and output head dimensions.
def make_data(nn=48, dd=10, deg=5, batches=24):
    xx, yy = [], []
    ss, tt = graph(nn, deg)
    for _ in range(batches):
        z = torch.randn(nn, dd, device=device)
        neigh = torch.zeros_like(z)
        neigh.index_add_(0, ss, z[tt])
        counts = torch.zeros(nn, device=device)
        counts.index_add_(0, ss, torch.ones_like(ss, dtype=torch.float))
        neigh = neigh / counts[:, None]
        target = 0.65 * neigh[:, :1] + 0.35 * z.mean(0, keepdim=True)[:, :1]
        xx.append(z); yy.append(target)
    return xx, yy, ss, tt

class DenseModel(torch.nn.Module):
    def __init__(self, dd=10, rr=8):
        super().__init__(); self.q=torch.nn.Linear(dd,rr,bias=False); self.k=torch.nn.Linear(dd,rr,bias=False); self.v=torch.nn.Linear(dd,rr,bias=False); self.head=torch.nn.Linear(rr,1)
    def forward(self,x): return self.head(dense_attention(x,self.q.weight.T,self.k.weight.T,self.v.weight.T))

class HybridModel(torch.nn.Module):
    def __init__(self, dd=10, rr=8):
        super().__init__(); self.q=torch.nn.Linear(dd,rr,bias=False); self.k=torch.nn.Linear(dd,rr,bias=False); self.v=torch.nn.Linear(dd,rr,bias=False); self.g=torch.nn.Linear(dd,1); self.head=torch.nn.Linear(rr,1)
    def forward(self,x,ss,tt):
        l=local_attention(x,ss,tt,self.q.weight.T,self.k.weight.T,self.v.weight.T); g=global_linear_attention(x,self.q.weight.T,self.k.weight.T,self.v.weight.T); gate=torch.sigmoid(self.g(x)); return self.head(gate*l+(1-gate)*g)

def train(model, hybrid, xs, ys, ss, tt):
    opt=torch.optim.Adam(model.parameters(),lr=0.02)
    for _ in range(80):
        for x,y in zip(xs,ys):
            pred=model(x,ss,tt) if hybrid else model(x)
            loss=((pred-y)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        vals=[]
        for x,y in zip(xs,ys): vals.append(float(((model(x,ss,tt) if hybrid else model(x))-y).pow(2).mean()))
    return float(np.mean(vals))

xs, ys, ss, tt = make_data()
dense_model=DenseModel().to(device); hybrid_model=HybridModel().to(device)
dense_loss=train(dense_model,False,xs,ys,ss,tt); hybrid_loss=train(hybrid_model,True,xs,ys,ss,tt)

# Timing uses one representative graph and identical projection width.
timing=[]
for nn in sizes:
    z=torch.randn(nn,d,device=device); a,b=graph(nn,degree)
    W=[torch.randn(d,width,device=device) for _ in range(3)]
    td=timed(lambda: dense_attention(z,*W), repeats=5)
    th=timed(lambda: (local_attention(z,a,b,*W), global_linear_attention(z,*W)), repeats=5)
    timing.append({'N':nn,'dense_ms':td,'hybrid_ms':th,'speedup':td/th})

result={'device':device,'predictions':{'local_mask_max_error':mask_error,'linear_reassociation_max_error':linear_error,'proxy_ratio_N64_to_N512':[ratios[0],ratios[-1]],'proxy_loglog_slope':slope},'learning_mse':{'dense':dense_loss,'hybrid':hybrid_loss},'timing':timing}
print(json.dumps(result,indent=2))
with open('results.json','w') as f: json.dump(result,f,indent=2)