"""v53: Bit-plane binary weights. Every stored bit ±1; effective weight is
a K-bit signed integer (K ±1 matrices combined as Σ 2^k · sign(W_k)).

No softmax, no RMSNorm, no float scales, no float residual. Pure ±1 everywhere:
  - Attention: Gumbel hard-argmax (v16 style) — one-hot per query
  - Residual: sign_ste(x + a + f)
  - Activations: ±1 via sign_ste
  - Weights: K ±1 planes summed with powers-of-2 weights

Effective weight values: signed 2^K-level integer. K=4 gives 16 distinct values
per weight while still storing every bit strictly as ±1.

Config: d_model=224, n_layers=4, n_heads=8, d_ff=160, K=4 → 5.0M ±1 weights.
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

from model import sign_ste, sign_ste_clipped, BinaryEmbedding
from model_v16 import gumbel_hard_attention


class BitPlaneLinear(nn.Module):
    """K stacked ±1 weight matrices; effective weight is signed K-bit integer."""
    def __init__(self, in_features, out_features, K=4, binarize_input=True):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.K = K
        self.binarize_input = binarize_input
        # K latent-float weight matrices; each sign()'d at forward.
        self.weights = nn.ParameterList([
            nn.Parameter(torch.randn(out_features, in_features) * 0.02)
            for _ in range(K)
        ])
        self.threshold = nn.Parameter(torch.zeros(out_features))
        # Accumulator max magnitude = sum_k 2^k * sqrt(in) ≈ (2^K - 1) sqrt(in)
        self.scale = 1.0 / ((2 ** K - 1) * math.sqrt(in_features))

    def forward(self, x):
        if self.binarize_input:
            x = sign_ste_clipped(x)
        acc = 0
        for k, w in enumerate(self.weights):
            W = sign_ste(w)
            acc = acc + (2 ** k) * F.linear(x, W)
        s = acc * self.scale - self.threshold
        return sign_ste_clipped(s)


class BitPlaneFFN(nn.Module):
    def __init__(self, d_model, d_ff, K=4):
        super().__init__()
        self.gate = BitPlaneLinear(d_model, d_ff, K=K, binarize_input=True)
        self.up = BitPlaneLinear(d_model, d_ff, K=K, binarize_input=True)
        self.down = BitPlaneLinear(d_ff, d_model, K=K, binarize_input=True)

    def forward(self, x):
        return self.down(self.gate(x) * self.up(x))


class BitPlaneAttention(nn.Module):
    """Gumbel hard-argmax attention with bit-plane Q/K/V/O projections.
    Attention matrix is one-hot per query (strict ±1 / {0,1})."""
    def __init__(self, d_model, n_heads, K=4):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.q_proj = BitPlaneLinear(d_model, d_model, K=K)
        self.k_proj = BitPlaneLinear(d_model, d_model, K=K)
        self.v_proj = BitPlaneLinear(d_model, d_model, K=K)
        self.o_proj = BitPlaneLinear(d_model, d_model, K=K)
        slopes = torch.tensor([1 << i for i in range(n_heads)], dtype=torch.long)
        self.register_buffer('alibi_slopes_int', slopes)

    def forward(self, x):
        B, T, D = x.shape
        H, Dh = self.n_heads, self.head_dim
        Q = self.q_proj(x).view(B, T, H, Dh).transpose(1, 2)
        K = self.k_proj(x).view(B, T, H, Dh).transpose(1, 2)
        V = self.v_proj(x).view(B, T, H, Dh).transpose(1, 2)

        scores = torch.matmul(Q, K.transpose(-2, -1))
        pos = torch.arange(T, device=x.device)
        dist = (pos.unsqueeze(0) - pos.unsqueeze(1)).abs()
        alibi = self.alibi_slopes_int.view(1, H, 1, 1).to(scores.dtype) \
                * dist.view(1, 1, T, T).to(scores.dtype)
        scores = scores - alibi

        mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1)
        A = gumbel_hard_attention(scores, mask=mask)    # one-hot per query
        O = torch.matmul(A, V)
        O = O.transpose(1, 2).contiguous().view(B, T, D)
        return self.o_proj(O)


class BitBlockV53(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, K=4):
        super().__init__()
        self.attn = BitPlaneAttention(d_model, n_heads, K=K)
        self.ffn = BitPlaneFFN(d_model, d_ff, K=K)

    def forward(self, x):
        a = self.attn(x)
        f = self.ffn(x)
        return sign_ste(x + a + f)          # strict ±1 residual


class BitLMv53(nn.Module):
    def __init__(self, vocab_size=128, d_model=224, n_layers=4, n_heads=8,
                 d_ff=160, K=4, 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.K = K
        self.embed = BinaryEmbedding(vocab_size, d_model)
        self.blocks = nn.ModuleList([
            BitBlockV53(d_model, n_heads, d_ff, K=K) 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


if __name__ == '__main__':
    from model_v16 import set_gumbel_tau
    set_gumbel_tau(0.5)
    for (D, d_ff, K) in ((224, 160, 4), (200, 160, 4), (256, 128, 4), (224, 192, 4)):
        m = BitLMv53(d_model=D, d_ff=d_ff, K=K)
        n = sum(p.numel() for p in m.parameters())
        print(f'D={D} d_ff={d_ff} K={K}: {n:,} ({n/1e6:.3f}M)')
    m = BitLMv53()
    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')