"""
Self-contained HuggingFace wrapper for the BabyLM entry (LoopLM) and monolith (LM), so the
models load as a stock AutoModelForCausalLM (trust_remote_code) for babylm-eval / leaderboard.
Model code is INLINED (no import of train_*.py) so this file is portable on the HF hub.
The ACTIVE class defs (LoopLMv2/Bind2 for arch "loop2", LM for the monolith) are byte-for-byte
the current training defs (train_loop.py / train_stage1.py) so state_dicts load exactly; the
legacy v1 defs (LoopLM/Bind) are retained ONLY to load the already-published v1 bypass
checkpoint (paper §4b diagnostic) and no longer exist in train_loop.py. forward() runs the whole loop inside a standard causal pass and
returns CausalLMOutput(logits, loss); empty-context, stateless across examples.

BabyLMModel (AutoModel entry) exists for the GLUE finetuning pipeline, which pools
last_hidden_state through its own classifier head. attention_mask is honored only on that
path (padded batches); the causal-LM path is unchanged — attn_mask=None reproduces the
exact zero-shot behavior the published eval numbers came from.
"""
import math, torch, torch.nn as nn, torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import CausalLMOutput, BaseModelOutput

def build_rope(T, D, device, base=10000.0):
    inv = 1.0/(base**(torch.arange(0,D,2,device=device).float()/D)); t = torch.arange(T,device=device).float()
    f = torch.outer(t, inv); emb = torch.cat([f, f], dim=-1); return emb.cos(), emb.sin()
def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1); return torch.cat((-x2, x1), dim=-1)
def apply_rope(x, cos, sin):
    return x*cos[None,None] + rotate_half(x)*sin[None,None]

class Attn(nn.Module):
    def __init__(self, d, nh):
        super().__init__(); self.nh=nh; self.hd=d//nh
        self.qkv=nn.Linear(d,3*d,bias=False); self.o=nn.Linear(d,d,bias=False)
    def forward(self, x, cos, sin, attn_mask=None):
        B,T,D=x.shape; qkv=self.qkv(x).view(B,T,3,self.nh,self.hd).permute(2,0,3,1,4)
        q,k,v=qkv[0],qkv[1],qkv[2]; q=apply_rope(q,cos,sin); k=apply_rope(k,cos,sin)
        if attn_mask is None: o=F.scaled_dot_product_attention(q,k,v,is_causal=True)
        else: o=F.scaled_dot_product_attention(q,k,v,attn_mask=attn_mask)
        return self.o(o.transpose(1,2).reshape(B,T,D))
class SwiGLU(nn.Module):
    def __init__(self, d, h):
        super().__init__(); self.w1=nn.Linear(d,h,bias=False); self.w3=nn.Linear(d,h,bias=False); self.w2=nn.Linear(h,d,bias=False)
    def forward(self, x): return self.w2(F.silu(self.w1(x))*self.w3(x))
class Block(nn.Module):
    def __init__(self, d, nh, h):
        super().__init__(); self.n1=nn.RMSNorm(d); self.attn=Attn(d,nh); self.n2=nn.RMSNorm(d); self.mlp=SwiGLU(d,h)
    def forward(self, x, cos, sin, attn_mask=None):
        x=x+self.attn(self.n1(x),cos,sin,attn_mask); return x+self.mlp(self.n2(x))

class LM(nn.Module):                       # monolith (train_stage1.LM)
    def __init__(self, vocab, d=384, nl=12, nh=6):
        super().__init__(); h=((int(8/3*d)+63)//64)*64
        self.emb=nn.Embedding(vocab,d); self.blocks=nn.ModuleList([Block(d,nh,h) for _ in range(nl)])
        self.nf=nn.RMSNorm(d); self.head=nn.Linear(d,vocab,bias=False); self.head.weight=self.emb.weight
        self.d=d; self.nh=nh
    def hidden(self, ids, attn_mask=None):
        cos,sin=build_rope(ids.shape[1], self.d//self.nh, ids.device); h=self.emb(ids)
        for b in self.blocks: h=b(h,cos,sin,attn_mask)
        return self.nf(h)
    def forward(self, ids): return self.head(self.hidden(ids))

class Bind(nn.Module):                     # label + trust (train_loop.Bind)
    def __init__(self, d, K=16, dr=64):
        super().__init__(); self.role=nn.Linear(d,K,bias=False); self.R=nn.Parameter(torch.randn(K,dr)*0.02)
        self.up=nn.Linear(dr,d,bias=False); self.trust=nn.Linear(d,1)
    def forward(self, h):
        a=torch.softmax(self.role(h),dim=-1); lab=a@self.R; tau=torch.sigmoid(self.trust(h)); return h+tau*self.up(lab)
class LoopLM(nn.Module):                    # entry (train_loop.LoopLM)
    def __init__(self, vocab, d=384, in_n=3, core_n=4, out_n=3, nh=6, T=3, K=16):
        super().__init__(); hdim=((int(8/3*d)+63)//64)*64
        self.emb=nn.Embedding(vocab,d)
        self.inb=nn.ModuleList([Block(d,nh,hdim) for _ in range(in_n)])
        self.core=nn.ModuleList([Block(d,nh,hdim) for _ in range(core_n)])
        self.outb=nn.ModuleList([Block(d,nh,hdim) for _ in range(out_n)])
        self.bind=Bind(d,K); self.vhead=nn.Linear(d,1)
        self.nf=nn.RMSNorm(d); self.head=nn.Linear(d,vocab,bias=False); self.head.weight=self.emb.weight
        self.d=d; self.nh=nh; self.T=T
    def hidden(self, ids, attn_mask=None):
        cos,sin=build_rope(ids.shape[1], self.d//self.nh, ids.device); h=self.emb(ids)
        for b in self.inb: h=b(h,cos,sin,attn_mask)
        for _ in range(self.T):
            z=self.bind(h); h2=z
            for b in self.core: h2=b(h2,cos,sin,attn_mask)
            v=torch.sigmoid(self.vhead(h2)); h=h+(1.0-v)*(h2-h)
        for b in self.outb: h=b(h,cos,sin,attn_mask)
        return self.nf(h)
    def forward(self, ids): return self.head(self.hidden(ids))

class Bind2(nn.Module):                    # v2 label+trust (train_loop.Bind, arch "loop2"): verdict-driven trust + experience prior + role-slice re-stamp
    def __init__(self, d, K=16, dr=64):
        super().__init__()
        self.dr = dr
        self.role = nn.Linear(d, K, bias=False)
        self.role_scale = nn.Parameter(torch.ones(1))
        self.R = nn.Parameter(torch.randn(K, dr) * 0.02)
        self.trust = nn.Linear(d, 1)
        self.v_gain = nn.Parameter(torch.zeros(1))
        self.vasana = nn.Parameter(torch.zeros(K))
    def forward(self, h, v_prev):
        a = torch.softmax(self.role_scale * self.role(h), dim=-1)
        lab = a @ self.R
        tau = torch.sigmoid(self.trust(h) + (a @ self.vasana)[..., None] + self.v_gain * (0.5 - v_prev))
        s = h[..., -self.dr:]
        return torch.cat([h[..., :-self.dr], (1.0 - tau) * s + tau * lab], dim=-1), a, tau

class LoopLMv2(nn.Module):                  # entry v2 (train_loop.LoopLM, arch "loop2")
    def __init__(self, vocab, d=384, in_n=3, core_n=4, out_n=3, nh=6, T=3, K=16):
        super().__init__(); hdim=((int(8/3*d)+63)//64)*64
        self.emb=nn.Embedding(vocab,d)
        self.inb=nn.ModuleList([Block(d,nh,hdim) for _ in range(in_n)])
        self.core=nn.ModuleList([Block(d,nh,hdim) for _ in range(core_n)])
        self.outb=nn.ModuleList([Block(d,nh,hdim) for _ in range(out_n)])
        self.bind=Bind2(d,K); self.vhead=nn.Linear(d,1)
        self.nf=nn.RMSNorm(d); self.head=nn.Linear(d,vocab,bias=False); self.head.weight=self.emb.weight
        self.d=d; self.nh=nh; self.T=T
    def hidden(self, ids, attn_mask=None):
        cos,sin=build_rope(ids.shape[1], self.d//self.nh, ids.device); h=self.emb(ids)
        for b in self.inb: h=b(h,cos,sin,attn_mask)
        v=torch.full_like(h[..., :1], 0.5)
        for _ in range(self.T):
            z,a,tau=self.bind(h,v); h2=z
            for b in self.core: h2=b(h2,cos,sin,attn_mask)
            v=torch.sigmoid(self.vhead(h2)); h=h2          # state flows through the loop (no bypass)
        for b in self.outb: h=b(h,cos,sin,attn_mask)
        return self.nf(h)
    def forward(self, ids): return self.head(self.hidden(ids))

# --- delta-rule + forced-bottleneck (arch "bind2_0"); class defs byte-for-byte from modeling_bind2_0.py
#     (train_bind2_0_babylm.py) so state_dicts load exactly. fla is imported lazily inside GDNBlock so
#     this module still imports without fla for the mono/loop2 paths. ---
class ChunkedAttn(nn.Module):
    """Forced bottleneck: causal attention restricted to within non-overlapping chunks of size C."""
    def __init__(self, d, nh, chunk):
        super().__init__()
        self.nh=nh; self.hd=d//nh; self.chunk=chunk
        self.qkv=nn.Linear(d,3*d,bias=False); self.o=nn.Linear(d,d,bias=False)
    def forward(self, x, cos, sin):
        B,T,D=x.shape
        qkv=self.qkv(x).view(B,T,3,self.nh,self.hd).permute(2,0,3,1,4)
        q,k,v=qkv[0],qkv[1],qkv[2]
        q=apply_rope(q,cos,sin); k=apply_rope(k,cos,sin)
        idx=torch.arange(T,device=x.device)
        same=(idx[:,None]//self.chunk)==(idx[None,:]//self.chunk)
        causal=idx[:,None]>=idx[None,:]
        keep=same&causal
        mask=torch.zeros(T,T,device=x.device,dtype=q.dtype).masked_fill(~keep,float("-inf"))
        o=F.scaled_dot_product_attention(q,k,v,attn_mask=mask)
        return self.o(o.transpose(1,2).reshape(B,T,D))
class GDNBlock(nn.Module):
    def __init__(self, d, idx, mlp_hidden, gdn_heads=4, gdn_hd=72):
        super().__init__()
        from fla.layers import GatedDeltaNet    # lazy: only bind2_0 needs fla
        self.n1=nn.RMSNorm(d)
        self.gdn=GatedDeltaNet(hidden_size=d, num_heads=gdn_heads, head_dim=gdn_hd, layer_idx=idx)
        self.n2=nn.RMSNorm(d); self.mlp=SwiGLU(d, mlp_hidden)
    def forward(self, x):
        m=self.gdn(self.n1(x))[0]   # fla returns (output, attn, cache)
        x=x+m
        return x+self.mlp(self.n2(x))
class AttnBlock(nn.Module):
    def __init__(self, d, nh, chunk, mlp_hidden):
        super().__init__()
        self.n1=nn.RMSNorm(d); self.attn=ChunkedAttn(d,nh,chunk)
        self.n2=nn.RMSNorm(d); self.mlp=SwiGLU(d,mlp_hidden)
    def forward(self, x, cos, sin):
        x=x+self.attn(self.n1(x),cos,sin)
        return x+self.mlp(self.n2(x))
class Bind2_0LM(nn.Module):                 # delta-rule + forced-bottleneck (modeling_bind2_0.Bind2_0LM)
    def __init__(self, vocab, d=384, depth=12, nh=6, chunk=32, mlp_hidden=576, gdn_heads=4, gdn_hd=72):
        super().__init__()
        self.emb=nn.Embedding(vocab,d)
        self.kinds=["attn" if (i+1)%4==0 else "gdn" for i in range(depth)]   # 3:1 GDN:attn
        self.blocks=nn.ModuleList([
            GDNBlock(d,i,mlp_hidden,gdn_heads,gdn_hd) if k=="gdn" else AttnBlock(d,nh,chunk,mlp_hidden)
            for i,k in enumerate(self.kinds)])
        self.nf=nn.RMSNorm(d); self.head=nn.Linear(d,vocab,bias=False); self.head.weight=self.emb.weight
        self.d=d; self.nh=nh; self.chunk=chunk
    def hidden(self, ids, attn_mask=None):   # attn_mask unused: chunked attn carries its own intra-chunk
        cos,sin=build_rope(ids.shape[1], self.d//self.nh, ids.device); h=self.emb(ids)   # mask (pad-mask
        for blk,k in zip(self.blocks,self.kinds):                                          # for GLUE is TODO,
            h=blk(h) if k=="gdn" else blk(h,cos,sin)                                        # zero-shot unaffected)
        return self.nf(h)
    def forward(self, ids): return self.head(self.hidden(ids))

def _build_backbone(config):
    if config.arch == "bind2_0":
        return Bind2_0LM(config.vocab_size, config.dim, config.depth, config.nhead,
                         chunk=config.chunk, mlp_hidden=config.mlp_hidden,
                         gdn_heads=config.gdn_heads, gdn_hd=config.gdn_hd)
    if config.arch == "loop2":
        return LoopLMv2(config.vocab_size, config.dim, config.in_n, config.core_n,
                        config.out_n, config.nhead, config.T, config.K)
    if config.arch == "loop":
        return LoopLM(config.vocab_size, config.dim, config.in_n, config.core_n,
                      config.out_n, config.nhead, config.T, config.K)
    return LM(config.vocab_size, config.dim, config.n_layer, config.nhead)

class BabyLMConfig(PretrainedConfig):
    model_type = "babylm"
    # the GLUE finetuning classifier reads config.hidden_size
    attribute_map = {"hidden_size": "dim", "num_attention_heads": "nhead", "num_hidden_layers": "n_layer"}
    def __init__(self, arch="loop", vocab_size=16000, dim=384, in_n=3, core_n=4, out_n=3,
                 T=3, K=16, nhead=6, n_layer=12,
                 depth=12, chunk=32, mlp_hidden=576, gdn_heads=4, gdn_hd=72, **kw):
        self.arch=arch; self.vocab_size=vocab_size; self.dim=dim; self.in_n=in_n; self.core_n=core_n
        self.out_n=out_n; self.T=T; self.K=K; self.nhead=nhead; self.n_layer=n_layer
        self.depth=depth; self.chunk=chunk; self.mlp_hidden=mlp_hidden; self.gdn_heads=gdn_heads; self.gdn_hd=gdn_hd
        super().__init__(**kw)

class BabyLMForCausalLM(PreTrainedModel):
    config_class = BabyLMConfig
    def __init__(self, config):
        super().__init__(config)
        self.backbone = _build_backbone(config)
        # Untie the LM head for a clean HF save (no shared tensors). Inference-equivalent: the head
        # weight is loaded from the checkpoint, which equals the tied embedding used at train time.
        self.backbone.head = nn.Linear(config.dim, config.vocab_size, bias=False)
        self.config.tie_word_embeddings = False
        self.post_init()
    def tie_weights(self, *args, **kwargs):
        pass  # head intentionally untied for export
    def get_input_embeddings(self): return self.backbone.emb
    def set_input_embeddings(self, v): self.backbone.emb = v
    def get_output_embeddings(self): return self.backbone.head
    def forward(self, input_ids=None, labels=None, attention_mask=None, **kw):
        logits = self.backbone(input_ids)
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits[:, :-1].reshape(-1, logits.size(-1)).float(), labels[:, 1:].reshape(-1))
        return CausalLMOutput(loss=loss, logits=logits)

def padding_causal_mask(attention_mask):
    # bool SDPA mask (B,1,T,T): attend where causal AND the key is a real (non-pad) token.
    # Pad-query rows would be fully masked (softmax NaN) with left padding, so the diagonal
    # stays open; their outputs are finite and get zero weight from every real query.
    B, T = attention_mask.shape; dev = attention_mask.device
    causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=dev))
    m = causal[None, None] & attention_mask.to(torch.bool)[:, None, None, :]
    return m | torch.eye(T, dtype=torch.bool, device=dev)[None, None]

class BabyLMModel(PreTrainedModel):
    """AutoModel entry (base model, no LM head applied) for the GLUE finetuning pipeline.
    Same backbone module tree as BabyLMForCausalLM so the exported checkpoint loads key-for-key."""
    config_class = BabyLMConfig
    def __init__(self, config):
        super().__init__(config)
        self.backbone = _build_backbone(config)
        self.backbone.head = nn.Linear(config.dim, config.vocab_size, bias=False)
        self.config.tie_word_embeddings = False
        self.post_init()
    def tie_weights(self, *args, **kwargs):
        pass  # head intentionally untied for export
    def get_input_embeddings(self): return self.backbone.emb
    def set_input_embeddings(self, v): self.backbone.emb = v
    def forward(self, input_ids=None, attention_mask=None, **kw):
        attn_mask = None
        if attention_mask is not None and not bool(attention_mask.all()):
            attn_mask = padding_causal_mask(attention_mask)
        return BaseModelOutput(last_hidden_state=self.backbone.hidden(input_ids, attn_mask))