"""Exact 2,990,000-parameter Sol Milkshake model implemented in MLX."""
from __future__ import annotations
import math
from collections import OrderedDict
import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_flatten
from sol_config import DEPLOYED_PARAMS, SolConfig

def unit_norm(x):
    xf=x.astype(mx.float32)
    return (xf*mx.rsqrt(mx.maximum(mx.sum(xf*xf,axis=-1,keepdims=True),mx.array(1e-12)))).astype(x.dtype)

def _count(module): return sum(int(x.size) for _,x in tree_flatten(module.trainable_parameters()))

def _rope(q,k,theta):
    n,d=q.shape[-2:]; inv=mx.exp(mx.arange(0,d,2,dtype=mx.float32)*(-math.log(theta)/d))
    ang=mx.arange(n,dtype=mx.float32)[:,None]*inv[None,:]; cs,sn=mx.cos(ang)[None,None],mx.sin(ang)[None,None]
    def rot(x):
        xf=x.astype(mx.float32); a,b=xf[...,::2],xf[...,1::2]
        return mx.stack((a*cs-b*sn,a*sn+b*cs),axis=-1).reshape(x.shape).astype(x.dtype)
    return rot(q),rot(k)

class TiedTokens(nn.Module):
    def __init__(self,c):
        super().__init__(); self.weight=unit_norm(mx.random.normal((c.vocab_size,c.d_model)))
    def encode(self,ids): return self.weight[ids]
    def decode(self,x,scale): return x.astype(mx.float32)@unit_norm(self.weight).astype(mx.float32).T*mx.exp(scale)

class TNGram(nn.Module):
    """Causal hashed CP retrieval for orders 2--5 (exactly 573,460 parameters)."""
    def __init__(self,c):
        super().__init__(); r=c.tn_rank; self.token_factor=mx.random.normal((c.vocab_size,r))*.02
        self.hash_tables=mx.random.normal((4,4950,r))*.02; self.order_factors=mx.ones((4,r))
        self.rank_to_hidden=mx.random.normal((r,c.d_model))*.02; self.output_bias=mx.zeros((c.d_model,))
        self.gates=mx.full((4,),-2.0); self.order_mix=mx.eye(4,r); self.context_controls=mx.zeros((16,))
    def __call__(self,ids):
        b,t=ids.shape; out=mx.broadcast_to(self.output_bias,(b,t,self.output_bias.shape[0]))
        for oi,order in enumerate(range(2,6)):
            h=mx.zeros_like(ids); prod=mx.ones((b,t,self.token_factor.shape[1]),dtype=mx.float32); valid=mx.arange(t)>=order-1
            for lag in range(1,order):
                shifted=mx.zeros((b,t),dtype=ids.dtype) if lag>=t else mx.concatenate((mx.zeros((b,lag),dtype=ids.dtype),ids[:,:-lag]),axis=1)
                h=(h*(131+oi*6)+shifted)%4950; prod=prod*self.token_factor[shifted].astype(mx.float32)
            rank=prod*self.hash_tables[oi,h]*self.order_factors[oi]
            rank=rank+self.order_mix[oi][None,None,:]*mx.mean(mx.tanh(self.context_controls[oi*4:(oi+1)*4]))
            out=out+(rank@self.rank_to_hidden)*mx.sigmoid(self.gates[oi])*valid[None,:,None]
        return out

class ShakeMemory(nn.Module):
    """Completed-chunk rolling memory (exactly 51,520 parameters)."""
    def __init__(self,c):
        super().__init__(); w=c.memory_width; self.chunk_size=c.chunk_size; self.slots=c.memory_slots
        self.summary_router=nn.Linear(c.d_model,1,bias=False); self.compress=nn.Linear(c.d_model,w,bias=False)
        self.query=nn.Linear(c.d_model,w,bias=False); self.key=nn.Linear(w,w,bias=False)
        self.value=nn.Linear(w,w,bias=False); self.output=nn.Linear(w,c.d_model,bias=False)
        self.slot_positions=mx.random.normal((c.memory_slots,c.d_model))*.01
        self.decay=mx.zeros((w,)); self.update_gate=mx.zeros((w,))
    def __call__(self,x):
        b,t,d=x.shape; chunks=[]; reads=[]
        for start in range(0,t,self.chunk_size):
            end=min(t,start+self.chunk_size); length=end-start
            if chunks:
                mem=mx.stack(chunks[-self.slots:],axis=1); count=mem.shape[1]
                mem=mem+self.compress(self.slot_positions[:count])[None]*mx.sigmoid(self.decay)[None,None]
                a=mx.softmax((self.query(x[:,start:end]).astype(mx.float32)@self.key(mem).astype(mx.float32).transpose(0,2,1))/math.sqrt(mem.shape[-1]),axis=-1)
                reads.append(self.output((a@self.value(mem).astype(mx.float32)).astype(x.dtype)))
            else: reads.append(mx.zeros((b,length,d),dtype=x.dtype))
            if length==self.chunk_size:
                weights=mx.softmax(self.summary_router(x[:,start:end]).astype(mx.float32).squeeze(-1),axis=-1)
                summary=mx.sum(weights[:,:,None]*x[:,start:end].astype(mx.float32),axis=1)
                chunks.append(self.compress(summary.astype(x.dtype))*mx.sigmoid(self.update_gate))
        return mx.concatenate(reads,axis=1)

class Attention(nn.Module):
    def __init__(self,c):
        super().__init__(); d=c.d_model; self.q=nn.Linear(d,d,bias=False); self.k=nn.Linear(d,c.n_kv_heads*c.head_dim,bias=False)
        self.v=nn.Linear(d,c.n_kv_heads*c.head_dim,bias=False); self.o=nn.Linear(d,d,bias=False)
        self.hq=c.n_q_heads; self.hk=c.n_kv_heads; self.hd=c.head_dim; self.theta=c.rope_theta
        self.value_current=mx.zeros((1,)); self.value_prior=mx.zeros((1,))
    def __call__(self,x,prior=None):
        b,t,_=x.shape; q=unit_norm(self.q(x).reshape(b,t,self.hq,self.hd)).transpose(0,2,1,3)
        k=unit_norm(self.k(x).reshape(b,t,self.hk,self.hd)).transpose(0,2,1,3); v=self.v(x).reshape(b,t,self.hk,self.hd).transpose(0,2,1,3)
        if prior is not None: v=mx.sigmoid(self.value_current)*v+mx.sigmoid(self.value_prior)*prior
        q,k=_rope(q,k,self.theta); k=mx.repeat(k,self.hq//self.hk,axis=1); vr=mx.repeat(v,self.hq//self.hk,axis=1)
        scores=(q.astype(mx.float32)@k.astype(mx.float32).transpose(0,1,3,2))/math.sqrt(self.hd)
        pos=mx.arange(t); mask=mx.where(pos[:,None]>=pos[None,:],0.0,-1e9)
        attended=(mx.softmax(scores+mask[None,None],axis=-1)@vr.astype(mx.float32)).astype(x.dtype)
        vh=unit_norm(vr); attended=attended-mx.sum(attended.astype(mx.float32)*vh.astype(mx.float32),axis=-1,keepdims=True).astype(x.dtype)*vh
        return self.o(attended.transpose(0,2,1,3).reshape(b,t,-1)),v

class Block(nn.Module):
    def __init__(self,c):
        super().__init__(); self.attention=Attention(c); self.gate=nn.Linear(c.d_model,c.ffn_hidden,bias=False)
        self.up=nn.Linear(c.d_model,c.ffn_hidden,bias=False); self.down=nn.Linear(c.ffn_hidden,c.d_model,bias=False)
        self.attn_step=mx.zeros((c.d_model,)); self.ffn_step=mx.zeros((c.d_model,))
    def __call__(self,x,prior=None):
        a,v=self.attention(unit_norm(x),prior); x=unit_norm(x+mx.sigmoid(self.attn_step)*(unit_norm(a)-x))
        f=self.down(nn.silu(self.gate(x))*self.up(x)); return unit_norm(x+mx.sigmoid(self.ffn_step)*(unit_norm(f)-x)),v

class SolMilkshake(nn.Module):
    def __init__(self,c=None):
        super().__init__(); self.config=c or SolConfig(); self.tokens=TiedTokens(self.config); self.tn_gram=TNGram(self.config)
        self.shake_memory=ShakeMemory(self.config); self.blocks=[Block(self.config) for _ in range(5)]
        self.pass_embeddings=mx.zeros((3,self.config.d_model)); self.mod_routers=[nn.Linear(self.config.d_model,1) for _ in range(6)]
        self.logit_scale=mx.zeros((self.config.vocab_size,)); self.misc_controls=mx.zeros((12,)); self._verify_ledger()
    def _select(self,x,old,router,capacity,lam):
        delta=1-mx.sum(unit_norm(x).astype(mx.float32)*unit_norm(old).astype(mx.float32),axis=-1)
        score=router(x).squeeze(-1)+lam*delta
        positions=mx.arange(x.shape[1])
        mask=(positions%4!=3) if capacity==.75 else (positions%2==0)
        return mx.broadcast_to(mask[None,:],score.shape),mx.sigmoid(score)
    def __call__(self,ids):
        x=self.tokens.encode(ids); x=unit_norm(x+mx.sigmoid(self.misc_controls[0])*self.tn_gram(ids))
        x=unit_norm(x+mx.sigmoid(self.misc_controls[1])*self.shake_memory(x)); x,v=self.blocks[0](x); ri=0
        for p in range(3):
            old=x; x=unit_norm(x+self.pass_embeddings[p]); prior=None
            for bi in range(1,4):
                proposal,prior=self.blocks[bi](x,prior)
                if p==0: x=proposal
                else:
                    mask,route_gate=self._select(x,old,self.mod_routers[ri],.75 if p==1 else .5,self.misc_controls[2+ri]); ri+=1
                    routed=unit_norm(x+route_gate[:,:,None]*(proposal-x))
                    x=mx.where(mask[:,:,None],routed,x)
        x,_=self.blocks[4](x,v); return self.tokens.decode(x,self.logit_scale)
    def parameter_ledger(self):
        ledger=OrderedDict(token_embedding_tied_head=int(self.tokens.weight.size),attention_total=sum(_count(b.attention)-2 for b in self.blocks),
          swiglu_total=sum(_count(b.gate)+_count(b.up)+_count(b.down) for b in self.blocks),tn_gram=_count(self.tn_gram),shake_memory=_count(self.shake_memory),
          ngpt_steps=sum(int(b.attn_step.size+b.ffn_step.size) for b in self.blocks),pass_embeddings=int(self.pass_embeddings.size),
          mod_routers=sum(_count(r) for r in self.mod_routers),logit_scale=int(self.logit_scale.size),value_residual=10,
          misc_controls=int(self.misc_controls.size),xsa=0); ledger['deployed_total']=sum(ledger.values()); return dict(ledger)
    def deployed_parameter_count(self): return _count(self)
    def _verify_ledger(self):
        expected={"token_embedding_tied_head":393216,"attention_total":491520,"swiglu_total":1474560,"tn_gram":573460,"shake_memory":51520,"ngpt_steps":1920,"pass_embeddings":576,"mod_routers":1158,"logit_scale":2048,"value_residual":10,"misc_controls":12,"xsa":0,"deployed_total":DEPLOYED_PARAMS}
        if self.parameter_ledger()!=expected or self.deployed_parameter_count()!=DEPLOYED_PARAMS: raise RuntimeError(f"parameter ledger mismatch: {self.parameter_ledger()}, actual={self.deployed_parameter_count()}")

def future_token_causality_error(model,sequence_length=35):
    ids=mx.arange(sequence_length,dtype=mx.int32)[None]%model.config.vocab_size; boundary=sequence_length//2
    changed=mx.concatenate((ids[:,:boundary],(ids[:,boundary:]+17)%model.config.vocab_size),axis=1)
    a,b=model(ids),model(changed); mx.eval(a,b); return float(mx.max(mx.abs(a[:,:boundary]-b[:,:boundary])))

def load_model(model_dir):
    """Load the standalone MLX model and frozen tokenizer from a snapshot."""
    from pathlib import Path
    from tokenizers import Tokenizer

    model_dir = Path(model_dir)
    model = SolMilkshake(SolConfig())
    model.load_weights(str(model_dir / "model.npz"))
    mx.eval(model.parameters())
    tokenizer = Tokenizer.from_file(str(model_dir / "tokenizer.json"))
    return model, tokenizer

def generate(model, tokenizer, prompt, max_new_tokens=64, temperature=0.0, seed=0):
    """Generate a single continuation with greedy or temperature sampling."""
    mx.random.seed(seed)
    tokens = tokenizer.encode(prompt).ids
    if not tokens:
        tokens = [0]
    for _ in range(max_new_tokens):
        context = tokens[-model.config.max_context:]
        logits = model(mx.array([context], dtype=mx.int32))[0, -1]
        if temperature and temperature > 0:
            next_token = int(mx.random.categorical(logits / temperature).item())
        else:
            next_token = int(mx.argmax(logits).item())
        tokens.append(next_token)
        if next_token == 1:
            break
    return tokenizer.decode(tokens)