import math, random, time
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F

SEED = 153

def seed_all(seed=SEED):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)


def math_check():
    seed_all(7)
    beta, eta, steps, n, dim = 0.07, 0.013, 80, 12000, 3
    z = torch.zeros(n, dim)
    amp = math.sqrt(2 * beta * eta)
    for _ in range(steps):
        z += amp * torch.randn_like(z)
    empirical = z.var(0, unbiased=True).mean().item()
    predicted = 2 * beta * eta * steps
    x = torch.ones(n)
    for _ in range(steps):
        x *= (1 - eta)
    expected = (1 - eta) ** steps
    return {
        'diffusion_empirical_variance': empirical,
        'diffusion_predicted_variance': predicted,
        'diffusion_relative_error': abs(empirical - predicted) / predicted,
        'quadratic_contraction_empirical': x.mean().item(),
        'quadratic_contraction_expected': expected,
    }


class ParticleAttention(nn.Module):
    def __init__(self, d_model=24, heads=8, particle=False, beta=0.0, noise_eta=0.03):
        super().__init__()
        assert d_model % heads == 0
        self.d, self.h, self.dk = d_model, heads, d_model // heads
        self.particle, self.beta, self.noise_eta = particle, beta, noise_eta
        self.q = nn.Parameter(torch.randn(heads, d_model, self.dk) * .12)
        self.k = nn.Parameter(torch.randn(heads, d_model, self.dk) * .12)
        self.v = nn.Parameter(torch.randn(heads, d_model, self.dk) * .12)
        self.o = nn.Parameter(torch.randn(heads, self.dk, d_model) * .12)
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x, noise=True):
        # x: batch, sequence, model dimension. Each particle is one attention head.
        q = torch.einsum('bnd,hdk->bhnk', x, self.q)
        k = torch.einsum('bnd,hdk->bhnk', x, self.k)
        v = torch.einsum('bnd,hdk->bhnk', x, self.v)
        a = torch.softmax(torch.einsum('bhnk,bhmk->bhnm', q, k) / math.sqrt(self.dk), dim=-1)
        hv = torch.einsum('bhnm,bhmk->bhnk', a, v)
        outs = torch.einsum('bhnk,hkd->bhnd', hv, self.o)
        y = outs.mean(1)
        return self.norm(x + y), a, outs

    @torch.no_grad()
    def langevin_noise(self):
        if self.particle and self.beta > 0:
            amp = math.sqrt(2 * self.beta * self.noise_eta)
            for p in (self.q, self.k, self.v, self.o):
                p.add_(amp * torch.randn_like(p))

    @torch.no_grad()
    def diversity(self):
        # Mean pairwise cosine similarity of flattened particle parameters.
        w = torch.cat([z.flatten(1) for z in (self.q, self.k, self.v, self.o)], dim=1)
        c = F.normalize(w, dim=1) @ F.normalize(w, dim=1).T
        h = self.h
        return ((c.sum() - h) / (h * (h - 1))) if h > 1 else 1.0

    def dropped_forward(self, x, keep):
        q = torch.einsum('bnd,hdk->bhnk', x, self.q[:keep])
        k = torch.einsum('bnd,hdk->bhnk', x, self.k[:keep])
        v = torch.einsum('bnd,hdk->bhnk', x, self.v[:keep])
        a = torch.softmax(torch.einsum('bhnk,bhmk->bhnm', q, k) / math.sqrt(self.dk), dim=-1)
        hv = torch.einsum('bhnm,bhmk->bhnk', a, v)
        y = torch.einsum('bhnk,hkd->bhnd', hv, self.o[:keep]).mean(1)
        return self.norm(x + y)


class TinyClassifier(nn.Module):
    def __init__(self, heads, beta):
        super().__init__()
        self.attn = ParticleAttention(heads=heads, particle=(beta > 0), beta=beta)
        self.fc = nn.Sequential(nn.Linear(24, 32), nn.Tanh(), nn.Linear(32, 2))
    def forward(self, x):
        y, a, outs = self.attn(x)
        return self.fc(y[:, 0]), a, outs


def data(seed=SEED):
    g = torch.Generator().manual_seed(seed)
    ntr, nte, seq, d = 768, 256, 8, 24
    # Label is encoded by a shared direction at token zero; other tokens are distractors.
    direction = F.normalize(torch.randn(d, generator=g), dim=0)
    def make(n):
        y = torch.randint(0, 2, (n,), generator=g)
        x = .8 * torch.randn(n, seq, d, generator=g)
        x[:, 0] += (2*y.float()-1)[:, None] * direction[None, :] * 1.4
        return x, y
    return make(ntr), make(nte)


def run_one(heads, beta, epochs=18):
    seed_all(SEED + heads + int(beta * 1e6))
    (xt, yt), (xv, yv) = data()
    model = TinyClassifier(heads, beta)
    opt = torch.optim.Adam(model.parameters(), lr=0.012)
    bs = 64; losses=[]; t0=time.time()
    for ep in range(epochs):
        perm = torch.randperm(len(xt))
        model.train()
        for ix in perm.split(bs):
            logits, _, _ = model(xt[ix])
            loss = F.cross_entropy(logits, yt[ix])
            opt.zero_grad(); loss.backward(); opt.step()
            model.attn.langevin_noise()
            losses.append(loss.item())
    model.eval()
    with torch.no_grad():
        logits, maps, outs = model(xv)
        val_loss = F.cross_entropy(logits, yv).item()
        acc = (logits.argmax(1) == yv).float().mean().item()
        keep = max(1, heads // 2)
        drop_acc = (model.attn.dropped_forward(xv, keep)[:,0] if False else
                    model.fc(model.attn.dropped_forward(xv, keep)[:,0]).argmax(1) == yv).float().mean().item()
        # Attention map similarity across particles, averaged over examples/queries.
        flat = maps.permute(1,0,2,3).reshape(heads, -1)
        sim = (F.normalize(flat, dim=1) @ F.normalize(flat, dim=1).T)
        map_sim = ((sim.sum()-heads)/(heads*(heads-1))).item() if heads > 1 else 1.0
    return {'heads':heads, 'beta':beta, 'val_loss':val_loss, 'val_accuracy':acc,
            'half_particle_accuracy':drop_acc, 'parameter_cosine':float(model.attn.diversity()),
            'attention_map_cosine':map_sim, 'last_train_loss':float(np.mean(losses[-20:])),
            'seconds':time.time()-t0}


def main():
    print({'math_check': math_check()})
    results=[]
    # Standard fixed 8-head attention versus the same particle population with Langevin noise.
    for beta in (0.0, 0.001, 0.005):
        results.append(run_one(8, beta))
    print({'results':results})

if __name__ == '__main__':
    main()