"""v39: BitHop — binary attention + content-addressable memory bank.

v38 (pure BitMixer, no attention) tanked at 3.12 BPC vs v17's 1.68, proving
that binary attention — even though the analysis showed heads look locally
biased — is still doing data-dependent work a static mix cannot replicate.

But the analysis was not wrong about the *missing* capability: binary attention
with ±1 QK and ALiBi cannot form long-range content-addressable routing.
v39 adds it explicitly via a Hopfield-style prototype memory bank per layer:

  BitHopHead: M learnable ±1 prototype (key, value) pairs.
              Each token: q = BitLinear(x); scores = q @ keys.T (integer);
              Gumbel-argmax over M picks one prototype; out = BitLinear(V_sel).

All signal paths ±1. Prototypes are latent-float weights whose sign() is used
at forward — identical discipline to BitLinear weights. At eval, argmax is a
pure integer compare over M options.

Each v39 block: attention(x) + hop(x) + ffn(x), all summed into residual.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from model import sign_ste, sign_ste_clipped, BitLinear, BitFFN, BinaryEmbedding
from model_v16 import _get_tau
from model_v18 import IntBinaryAttention


class BitHopHead(nn.Module):
    """Content-addressable memory: M ±1 prototype key/value pairs, top-1 routing."""
    def __init__(self, d_model, n_proto=32):
        super().__init__()
        self.n_proto = n_proto
        self.q_proj = BitLinear(d_model, d_model, binarize_input=True)
        self.o_proj = BitLinear(d_model, d_model, binarize_input=True)
        # Latent-float prototypes; forward uses sign(proto).
        self.key_proto = nn.Parameter(torch.randn(n_proto, d_model) * 0.02)
        self.val_proto = nn.Parameter(torch.randn(n_proto, d_model) * 0.02)

    def forward(self, x):
        B, T, D = x.shape
        q = self.q_proj(x)            # (B, T, D) ±1
        K = sign_ste(self.key_proto)  # (M, D) ±1
        V = sign_ste(self.val_proto)  # (M, D) ±1
        scores = q @ K.t()            # (B, T, M) integer popcount
        tau = _get_tau(scores.device)
        if scores.requires_grad:
            g = -torch.log(-torch.log(torch.rand_like(scores).clamp(min=1e-9)) + 1e-9)
            y_soft = F.softmax((scores + g) / tau, dim=-1)
            y_hard = torch.zeros_like(y_soft).scatter_(-1, y_soft.argmax(-1, keepdim=True), 1.0)
            A = y_soft + (y_hard - y_soft).detach()
        else:
            A = torch.zeros_like(scores).scatter_(-1, scores.argmax(-1, keepdim=True), 1.0)
        out = A @ V                   # (B, T, D)
        return self.o_proj(out)


class BitBlockV39(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, n_proto):
        super().__init__()
        self.attn = IntBinaryAttention(d_model, n_heads)
        self.hop = BitHopHead(d_model, n_proto)
        self.ffn = BitFFN(d_model, d_ff)

    def forward(self, x):
        a = self.attn(x)
        h = self.hop(x)
        f = self.ffn(x)
        return sign_ste(x + a + h + f)


class BitLMv39(nn.Module):
    def __init__(self, vocab_size=128, d_model=256, n_layers=8, n_heads=8,
                 d_ff=288, n_proto=32, max_seq_len=256):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.n_layers = n_layers
        self.max_seq_len = max_seq_len
        self.n_proto = n_proto
        self.embed = BinaryEmbedding(vocab_size, d_model)
        self.blocks = nn.ModuleList([
            BitBlockV39(d_model, n_heads, d_ff, n_proto) for _ in range(n_layers)
        ])
        self.out_codebook = nn.Parameter(torch.randn(vocab_size, d_model) * 0.02)
        self.logit_scale = nn.Parameter(torch.tensor(1.0 / math.sqrt(d_model)))
        self.out_bias = nn.Parameter(torch.zeros(vocab_size))

    def forward(self, idx, targets=None):
        x = self.embed(idx)
        for blk in self.blocks:
            x = blk(x)
        W_out = sign_ste(self.out_codebook)
        scores = torch.matmul(x, W_out.t())
        logits = scores * self.logit_scale + self.out_bias
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, self.vocab_size), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens=200, temperature=1.0, top_k=None):
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.max_seq_len:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-5)
            if top_k is not None:
                v, _ = torch.topk(logits, top_k)
                logits[logits < v[:, [-1]]] = -float('inf')
            probs = F.softmax(logits, dim=-1)
            nxt = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, nxt], dim=1)
        return idx


if __name__ == '__main__':
    from model_v16 import set_gumbel_tau
    set_gumbel_tau(0.5)
    for d_ff in (256, 272, 288, 304):
        m = BitLMv39(d_ff=d_ff)
        n = sum(p.numel() for p in m.parameters())
        print(f'd_ff={d_ff}: {n:,} ({n/1e6:.3f}M)')
    m = BitLMv39()
    x = torch.randint(0, 128, (2, 64))
    y = torch.randint(0, 128, (2, 64))
    logits, loss = m(x, y)
    loss.backward()
    print(f'loss={loss.item():.3f}, backward OK')