"""v41: BitConv — depthwise causal ±1 convolution replaces attention.

Findings so far:
  - v38 (global static T×T mix, shared-across-channels) failed catastrophically (3.12 BPC).
    Too few degrees of freedom: one pattern for 256 channels can't capture diverse
    per-channel routing.
  - v39/v40 (attention + content-addressable memory) failed to beat v17. Adding
    more routing options didn't help; the underlying local routing was already OK.
  - Analysis said 102/120 heads are local-recent (mean argmax-distance < 3 tokens).

v41 takes the "attention is local" finding literally and replaces attention with a
CHANNEL-WISE causal ±1 convolution. Each of d_model channels gets its own learnable
K-length ±1 kernel over the past K tokens. Per-channel (not shared) is the key
distinction from v38. Parameter cost is trivial (D·K), so most of the 5M budget
goes to FFN.

  y[b, t, d] = sign( sum_{k=0}^{K-1}  sign(kernel[d, k]) * x[b, t-k, d] / sqrt(K) )

Everything stays ±1. No Q/K/V, no softmax, no ALiBi, no multi-head.

Config (5M): d_model=256, n_layers=8, K=16, d_ff=800.
"""
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


class BitCausalConv(nn.Module):
    """Depth-wise causal ±1 1D convolution. One K-length kernel per channel."""
    def __init__(self, d_model, kernel_size=16):
        super().__init__()
        self.d_model = d_model
        self.K = kernel_size
        # Latent float kernel; forward uses sign(kernel).
        self.kernel = nn.Parameter(torch.randn(d_model, kernel_size) * 0.02)
        # Scale so sum of K ±1 values has unit variance on average.
        self.scale = 1.0 / math.sqrt(kernel_size)

    def forward(self, x):
        # x: (B, T, D), ±1
        B, T, D = x.shape
        W = sign_ste(self.kernel)          # (D, K) ±1
        # F.conv1d expects (B, C, T_in) and weight (C_out, C_in/groups, K).
        # Depthwise: C_out = C_in = D, groups = D, weight shape (D, 1, K).
        W_conv = W.unsqueeze(1)            # (D, 1, K)
        x_t = x.transpose(1, 2)            # (B, D, T)
        # Causal: pad K-1 on the left. Output length = T.
        x_pad = F.pad(x_t, (self.K - 1, 0))
        y = F.conv1d(x_pad, W_conv, groups=D)   # (B, D, T)
        y = y.transpose(1, 2) * self.scale      # (B, T, D), ~unit scale
        return sign_ste_clipped(y)


class BitBlockV41(nn.Module):
    def __init__(self, d_model, d_ff, kernel_size):
        super().__init__()
        self.conv = BitCausalConv(d_model, kernel_size)
        self.ffn = BitFFN(d_model, d_ff)

    def forward(self, x):
        c = self.conv(x)
        f = self.ffn(x)
        return sign_ste(x + c + f)


class BitLMv41(nn.Module):
    def __init__(self, vocab_size=128, d_model=256, n_layers=8, d_ff=800,
                 kernel_size=16, 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.kernel_size = kernel_size
        self.embed = BinaryEmbedding(vocab_size, d_model)
        self.blocks = nn.ModuleList([
            BitBlockV41(d_model, d_ff, kernel_size) 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__':
    for d_ff in (768, 800, 816):
        for K in (8, 16, 32):
            m = BitLMv41(d_model=256, n_layers=8, d_ff=d_ff, kernel_size=K)
            n = sum(p.numel() for p in m.parameters())
            print(f'd_ff={d_ff}, K={K}: {n:,} ({n/1e6:.3f}M)')
    m = BitLMv41()
    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')