# model_dna_tern.py - ternary-QAT wrapper for the DNA-DiskChat-2B controller.
# Swaps every dense nn.Linear in the controller for a BitNet b1.58 ternary-STE
# linear (weights -> {-1,0,+1}*scale, 8-bit activations). Embedding, tied head,
# the codon table and its codebook, and ca/cb code tables stay full precision.
import sys; sys.path.insert(0, '/root/dna')
import torch, torch.nn as nn, torch.nn.functional as F
from model_dna import DnaChat, counts

def weight_quant_ternary(w):
    scale = w.abs().mean().clamp_(min=1e-5)
    w_tern = torch.round(w / scale).clamp_(-1, 1)
    return w_tern, scale

def act_quant_int8(x):
    scale = x.abs().amax(dim=-1, keepdim=True).clamp_(min=1e-5) / 127.0
    return torch.round(x / scale).clamp_(-128, 127), scale

class TernLinear(nn.Module):
    """Drop-in for a bias-free nn.Linear: same .weight [out,in], ternary-STE forward."""
    def __init__(self, lin: nn.Linear):
        super().__init__()
        self.in_features = lin.in_features; self.out_features = lin.out_features
        self.weight = lin.weight  # share the SAME parameter -> warm-start is a no-op
    def forward(self, x):
        x_q, x_scale = act_quant_int8(x)
        x_ste = x + (x_q * x_scale - x).detach()
        w_tern, w_scale = weight_quant_ternary(self.weight)
        w_ste = self.weight + (w_tern * w_scale - self.weight).detach()
        return F.linear(x_ste, w_ste)

def ternarize_(model, skip=('head2',)):
    """In-place: replace controller nn.Linear modules with TernLinear."""
    n = 0
    for name, mod in list(model.named_modules()):
        for child_name, child in list(mod.named_children()):
            if isinstance(child, nn.Linear) and child.bias is None:
                full = f'{name}.{child_name}' if name else child_name
                if any(s in full for s in skip):
                    continue
                setattr(mod, child_name, TernLinear(child)); n += 1
    return n

def controller_byte_budget(model):
    """Bytes read per token at batch-1 decode for the controller's matrix weights."""
    tern_params = 0; fp_params = 0
    for name, mod in model.named_modules():
        if isinstance(mod, TernLinear):
            tern_params += mod.weight.numel()
        elif isinstance(mod, nn.Linear) and mod.bias is None:
            fp_params += mod.weight.numel()
    # ternary packs at 1.6 bits/param = 0.2 B; fp linears at bf16 = 2 B
    return {
        'tern_params': tern_params, 'fp_linear_params': fp_params,
        'bytes_per_token_tern': tern_params * 0.2 + fp_params * 2,
        'bytes_per_token_bf16': (tern_params + fp_params) * 2,
    }