import os
import sys
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as cp
from dataclasses import dataclass

@dataclass
class ModelConfig:
    vocab_size: int = 50272       # Resized from 50257 to 50272 for CUDA warp-coalesced alignment
    d_model: int = 768            # Hidden dimension size
    n_iterations: int = 16        # Number of recursive passes (effective layers)
    n_heads: int = 12             # Query attention heads
    n_kv_heads: int = 4           # KV attention heads for GQA (3:1 ratio)
    d_ff: int = 2048              # SwiGLU intermediate dimension
    max_seq_len: int = 512        # Sequence window context limit
    bias: bool = False            # True LLaMA/Gemma style is bias-free for stability

class RMSNorm(nn.Module):
    """Llama-style Root Mean Square Normalization with optional step-conditioned adaptive scale (AdaRMSNorm).
    Includes FP32 upcasting to prevent FP16 numerical underflow/overflow NaN corruption.
    """
    def __init__(self, dim: int, n_iterations: int = None, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
        
        # Step-specific adaptive scales (RingFormer style) to prevent KV-cache semantic shift
        if n_iterations is not None:
            self.step_scales = nn.Parameter(torch.zeros(n_iterations, dim))
        else:
            self.step_scales = None

    def forward(self, x, r_idx: int = None):
        x_fp32 = x.to(torch.float32)
        variance = x_fp32.pow(2).mean(-1, keepdim=True)
        normed = x_fp32 * torch.rsqrt(variance + self.eps)
        normed = normed.to(x.dtype)
        
        if self.step_scales is not None and r_idx is not None:
            idx = r_idx % self.step_scales.shape[0]
            scale = self.weight + self.step_scales[idx]
            return normed * scale
        else:
            return normed * self.weight

class RoPE(nn.Module):
    """Rotary Positional Embeddings (RoPE) applied to Query and Key states.
    Includes a Dynamic Frequency Extension safeguard to support sequence lengths beyond max_seq_len.
    """
    def __init__(self, dim: int, max_seq_len: int = 512, theta: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.theta = theta
        
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        
        t = torch.arange(max_seq_len, dtype=torch.float32)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)

    def _rotate_half(self, x):
        half_dim = self.dim // 2
        x1 = x[..., :half_dim]
        x2 = x[..., half_dim:]
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, x, seq_len: int, start_pos: int = 0):
        end_pos = start_pos + seq_len
        
        if end_pos > self.max_seq_len:
            t = torch.arange(end_pos, dtype=torch.float32, device=x.device)
            freqs = torch.outer(t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos()[start_pos:end_pos, :].unsqueeze(0).unsqueeze(2)
            sin = emb.sin()[start_pos:end_pos, :].unsqueeze(0).unsqueeze(2)
        else:
            cos = self.cos_cached[start_pos:end_pos, :].unsqueeze(0).unsqueeze(2)
            sin = self.sin_cached[start_pos:end_pos, :].unsqueeze(0).unsqueeze(2)
            
        cos = cos.to(device=x.device, dtype=x.dtype)
        sin = sin.to(device=x.device, dtype=x.dtype)
        
        return (x * cos) + (self._rotate_half(x) * sin)

class ModulatedLinear(nn.Module):
    """A linear layer with frozen weights augmented with SVD-initialized low-rank bases 
    which are dynamically scaled by an input-conditioned modulation vector (Ouroboros Weight Modulation).
    """
    def __init__(self, in_features: int, out_features: int, rank_ctrl: int = 64, bias: bool = False):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.rank_ctrl = rank_ctrl
        
        # Base projection (frozen during core training)
        self.base_layer = nn.Linear(in_features, out_features, bias=bias)
        self.base_layer.weight.requires_grad = False
        if self.base_layer.bias is not None:
            self.base_layer.bias.requires_grad = False
            
        # Low-rank modulation bases
        self.A = nn.Parameter(torch.zeros(out_features, rank_ctrl))
        self.B = nn.Parameter(torch.zeros(in_features, rank_ctrl))
        
        self.svd_initialized = False

    def initialize_svd(self):
        if self.svd_initialized:
            return
        with torch.no_grad():
            W_base = self.base_layer.weight.float()
            U, S, Vh = torch.linalg.svd(W_base, full_matrices=False)
            rank = min(self.rank_ctrl, S.numel())
            self.A.copy_((U[:, :rank] * torch.sqrt(S[:rank])).to(self.A.dtype))
            self.B.copy_((Vh[:rank, :].t() * torch.sqrt(S[:rank])).to(self.B.dtype))
        self.svd_initialized = True

    def forward(self, x, mod_vector=None):
        out_base = self.base_layer(x)
        if mod_vector is None:
            return out_base
            
        # Optimized execution of (W_base + A * Diag(m) * B^T) * x^T:
        # hb = x * B  [B, T, r_ctrl]
        # hb_scaled = hb * m  [B, T, r_ctrl]
        # output = out_base + hb_scaled * A^T
        hb = torch.matmul(x, self.B)
        if mod_vector.dim() == hb.dim() - 1:
            hb_scaled = hb * mod_vector.unsqueeze(-2)
        elif mod_vector.dim() == hb.dim():
            hb_scaled = hb * mod_vector
        else:
            hb_scaled = hb * mod_vector
        out_extra = torch.matmul(hb_scaled, self.A.t())
        
        return out_base + out_extra

class ControllerHypernetwork(nn.Module):
    """Generates dynamic step-dependent weight modulation diagonal scalars from the mean-pooled state."""
    def __init__(self, d_model: int, num_modulated_projs: int = 6, rank_ctrl: int = 64):
        super().__init__()
        self.d_model = d_model
        self.num_modulated_projs = num_modulated_projs
        self.rank_ctrl = rank_ctrl
        
        self.fc1 = nn.Linear(d_model, 256, bias=False)
        self.fc2 = nn.Linear(256, num_modulated_projs * rank_ctrl, bias=False)
        
        nn.init.normal_(self.fc1.weight, std=0.01)
        nn.init.normal_(self.fc2.weight, std=0.01)

    def forward(self, h, step_emb):
        if h.dim() == 2:
            inp = h + step_emb.unsqueeze(0)
            out = self.fc2(F.silu(self.fc1(inp)))
            mod = 1.0 + out.view(-1, self.num_modulated_projs, self.rank_ctrl)
        else:
            inp = h + step_emb.view(1, 1, -1)
            out = self.fc2(F.silu(self.fc1(inp)))
            B, T, _ = h.shape
            mod = 1.0 + out.view(B, T, self.num_modulated_projs, self.rank_ctrl)
        return mod

class LoRAExit(nn.Module):
    """Decoupled low-rank intermediate early exit adapter (LoRAExit) to prevent gradient conflict."""
    def __init__(self, d_model: int, r: int = 32):
        super().__init__()
        self.up_proj = nn.Linear(d_model, r, bias=False)
        self.down_proj = nn.Linear(r, d_model, bias=False)
        
        nn.init.zeros_(self.down_proj.weight)
        nn.init.normal_(self.up_proj.weight, std=0.02)

    def initialize_svd(self, W_i, W_j):
        with torch.no_grad():
            diff = W_i.float() - W_j.float()
            U, S, Vh = torch.linalg.svd(diff, full_matrices=False)
            r = min(self.up_proj.out_features, S.numel())
            
            A = U[:, :r] * S[:r]
            B = Vh[:r, :]
            
            self.down_proj.weight.copy_(A.to(self.down_proj.weight.dtype))
            self.up_proj.weight.copy_(B.to(self.up_proj.weight.dtype))

    def forward(self, x):
        return self.down_proj(self.up_proj(x))

class ContextAnchoredMLA(nn.Module):
    """Context-Anchored Recurrent Attention (CART) using Multi-Head Latent Attention (MLA)."""
    def __init__(self, config: ModelConfig, r_latent: int = 128, d_r: int = 64):
        super().__init__()
        self.n_heads = config.n_heads
        self.head_dim = config.d_model // config.n_heads
        self.d_r = d_r
        self.r_latent = r_latent
        
        # Modulated query projections
        self.q_proj = ModulatedLinear(config.d_model, config.n_heads * self.head_dim, bias=config.bias)
        self.q_rope_proj = ModulatedLinear(config.d_model, config.n_heads * d_r, bias=config.bias)
        
        # Key & Value projections from the cached anchor C_KV
        self.k_proj = nn.Linear(r_latent, config.n_heads * self.head_dim, bias=config.bias)
        self.k_rope_proj = nn.Linear(r_latent, config.n_heads * d_r, bias=config.bias)
        self.v_proj = nn.Linear(r_latent, config.n_heads * self.head_dim, bias=config.bias)
        
        # Modulated output projection
        self.out_proj = ModulatedLinear(config.n_heads * self.head_dim, config.d_model, bias=config.bias)
        
        self.rope = RoPE(dim=d_r, max_seq_len=config.max_seq_len)
        
        # Fallback dummy projection if called without anchor (e.g., in diagnostics)
        self.dkv_weight_dummy = nn.Parameter(torch.empty(r_latent, config.d_model))
        nn.init.normal_(self.dkv_weight_dummy, std=0.02)

    def forward(self, x, r_idx: int = 0, anchor_ckv=None, mod_vector=None, kv_cache=None):
        B, T, _ = x.shape
        
        if anchor_ckv is None:
            anchor_ckv = F.linear(x, self.dkv_weight_dummy)
            
        T_anchor = anchor_ckv.shape[1]
        
        q_mod = mod_vector[..., 0, :] if mod_vector is not None else None
        q_rope_mod = mod_vector[..., 1, :] if mod_vector is not None else None
        out_mod = mod_vector[..., 2, :] if mod_vector is not None else None
        
        q_c = self.q_proj(x, q_mod).view(B, T, self.n_heads, self.head_dim)
        q_r = self.q_rope_proj(x, q_rope_mod).view(B, T, self.n_heads, self.d_r)
        
        k_c = self.k_proj(anchor_ckv).view(B, T_anchor, self.n_heads, self.head_dim)
        k_r = self.k_rope_proj(anchor_ckv).view(B, T_anchor, self.n_heads, self.d_r)
        v = self.v_proj(anchor_ckv).view(B, T_anchor, self.n_heads, self.head_dim)
        
        q_r = self.rope(q_r, T, start_pos=T_anchor - T)
        k_r = self.rope(k_r, T_anchor, start_pos=0)
        
        q = torch.cat([q_c, q_r], dim=-1).transpose(1, 2)
        k = torch.cat([k_c, k_r], dim=-1).transpose(1, 2)
        v = v.transpose(1, 2)
        
        is_causal = (T > 1) and (T == T_anchor)
        
        try:
            context = F.scaled_dot_product_attention(
                q, k, v, is_causal=is_causal
            )
        except Exception:
            scores = torch.matmul(q, k.transpose(-2, -1)) / ((self.head_dim + self.d_r) ** 0.5)
            if T > 1 and T == T_anchor:
                mask = torch.triu(torch.full((T, T_anchor), float('-inf'), device=x.device), diagonal=1)
                scores = scores + mask.unsqueeze(0).unsqueeze(1)
            scores_fp32 = scores.to(torch.float32)
            attn = F.softmax(scores_fp32, dim=-1).to(x.dtype)
            context = torch.matmul(attn, v)
            
        context = context.transpose(1, 2).contiguous().view(B, T, -1)
        return self.out_proj(context, out_mod)

class GQAAttention(nn.Module):
    """Grouped-Query Attention (GQA) with multi-step recursive KV Caching."""
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.n_heads = config.n_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim = config.d_model // config.n_heads
        self.num_queries_per_kv = config.n_heads // config.n_kv_heads
        
        self.q_proj = nn.Linear(config.d_model, config.n_heads * self.head_dim, bias=config.bias)
        self.k_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=config.bias)
        self.v_proj = nn.Linear(config.d_model, config.n_kv_heads * self.head_dim, bias=config.bias)
        self.out_proj = nn.Linear(config.n_heads * self.head_dim, config.d_model, bias=config.bias)
        
        self.rope = RoPE(dim=self.head_dim, max_seq_len=config.max_seq_len)

    def forward(self, x, r_idx: int = 0, kv_cache_info=None):
        B, T, _ = x.shape
        
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
        
        if kv_cache_info is not None:
            cache_obj, layer_idx, zone_name = kv_cache_info
            if zone_name == "prelude":
                start_pos = cache_obj.prelude_lengths[layer_idx]
            else: # "coda"
                start_pos = cache_obj.coda_lengths[layer_idx]
        else:
            start_pos = 0
            
        q = self.rope(q, T, start_pos=start_pos)
        k = self.rope(k, T, start_pos=start_pos)
        
        if kv_cache_info is not None:
            cache_obj, layer_idx, zone_name = kv_cache_info
            if zone_name == "prelude":
                k, v = cache_obj.update_prelude(k, v, layer_idx)
            else: # "coda"
                k, v = cache_obj.update_coda(k, v, layer_idx)
            T_total = k.shape[1]
        else:
            T_total = T
            
        q = q.transpose(1, 2)  # [B, H_q, T, head_dim]
        
        if self.num_queries_per_kv > 1:
            k_expanded = k.repeat_interleave(self.num_queries_per_kv, dim=2)
            v_expanded = v.repeat_interleave(self.num_queries_per_kv, dim=2)
        else:
            k_expanded = k
            v_expanded = v
            
        k_expanded = k_expanded.transpose(1, 2)  # [B, H_q, T_total, head_dim]
        v_expanded = v_expanded.transpose(1, 2)  # [B, H_q, T_total, head_dim]
        
        is_causal = (T > 1)
        
        try:
            context = F.scaled_dot_product_attention(
                q, 
                k_expanded, 
                v_expanded, 
                is_causal=is_causal
            )
        except Exception:
            scores = torch.matmul(q, k_expanded.transpose(-2, -1)) / (self.head_dim ** 0.5)
            if T > 1:
                mask = torch.triu(torch.full((T, T_total), float('-inf'), device=x.device), diagonal=T_total - T + 1)
                scores = scores + mask.unsqueeze(0).unsqueeze(1)
            scores_fp32 = scores.to(torch.float32)
            attn = F.softmax(scores_fp32, dim=-1).to(x.dtype)
            context = torch.matmul(attn, v_expanded)
            
        context = context.transpose(1, 2).contiguous().view(B, T, -1)
        return self.out_proj(context)

class SwiGLUFFN(nn.Module):
    """Standard Gated Linear Unit with Swish (SiLU) activation for LLMs."""
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.d_model, config.d_ff, bias=config.bias)
        self.up_proj = nn.Linear(config.d_model, config.d_ff, bias=config.bias)
        self.down_proj = nn.Linear(config.d_ff, config.d_model, bias=config.bias)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))

class ModulatedSwiGLUFFN(nn.Module):
    """Modulated Gated Linear Unit with Swish (SiLU) activation."""
    def __init__(self, config: ModelConfig, rank_ctrl: int = 64):
        super().__init__()
        self.gate_proj = ModulatedLinear(config.d_model, config.d_ff, rank_ctrl=rank_ctrl, bias=config.bias)
        self.up_proj = ModulatedLinear(config.d_model, config.d_ff, rank_ctrl=rank_ctrl, bias=config.bias)
        self.down_proj = ModulatedLinear(config.d_ff, config.d_model, rank_ctrl=rank_ctrl, bias=config.bias)

    def forward(self, x, mod_vector=None):
        gate_mod = mod_vector[..., 3, :] if mod_vector is not None else None
        up_mod = mod_vector[..., 4, :] if mod_vector is not None else None
        down_mod = mod_vector[..., 5, :] if mod_vector is not None else None
        
        return self.down_proj(F.silu(self.gate_proj(x, gate_mod)) * self.up_proj(x, up_mod), down_mod)

class StepAdapter(nn.Module):
    """A lightweight, low-rank step adapter to allow depth-specific specialization."""
    def __init__(self, d_model: int, r: int = 64, alpha: int = 128):
        super().__init__()
        self.up_proj = nn.Linear(d_model, r, bias=False)
        self.down_proj = nn.Linear(r, d_model, bias=False)
        self.scale = alpha / r
        
        nn.init.zeros_(self.down_proj.weight)
        nn.init.normal_(self.up_proj.weight, std=0.02)

    def forward(self, x):
        return self.down_proj(F.silu(self.up_proj(x))) * self.scale

class TransformerBlock(nn.Module):
    """A unified block class that acts as either a standard GQA layer or a recurrent core layer."""
    def __init__(self, config: ModelConfig, layer_type: str = "standard"):
        super().__init__()
        self.layer_type = layer_type
        
        self.attn_norm = RMSNorm(config.d_model, n_iterations=config.n_iterations)
        self.ffn_norm = RMSNorm(config.d_model, n_iterations=config.n_iterations)
        
        if layer_type == "recurrent_core":
            self.attn = ContextAnchoredMLA(config)
            self.ffn = ModulatedSwiGLUFFN(config)
            self.adapters = nn.ModuleList([
                StepAdapter(config.d_model, r=64) for _ in range(config.n_iterations)
            ])
        else:
            self.attn = GQAAttention(config)
            self.ffn = SwiGLUFFN(config)
            # Dummy adapters to keep diagnostic code happy if it accesses block.adapters directly
            self.adapters = nn.ModuleList([
                StepAdapter(config.d_model, r=64) for _ in range(config.n_iterations)
            ])

    def forward(self, x, r_idx: int = 0, kv_cache=None, anchor_ckv=None, mod_vector=None):
        if self.layer_type == "recurrent_core":
            attn_norm_out = self.attn_norm(x, r_idx)
            attn_out = self.attn(attn_norm_out, r_idx=r_idx, anchor_ckv=anchor_ckv, mod_vector=mod_vector)
            adapter_out = self.adapters[r_idx](attn_out)
            
            h = x + attn_out + adapter_out
            ffn_norm_out = self.ffn_norm(h, r_idx)
            ffn_out = self.ffn(ffn_norm_out, mod_vector=mod_vector)
            
            return attn_out + adapter_out + ffn_out
        else:
            # GQA self-attention layers (Prelude and Coda)
            attn_norm_out = self.attn_norm(x, r_idx)
            attn_out = self.attn(attn_norm_out, r_idx=r_idx, kv_cache_info=kv_cache)
            
            h = x + attn_out
            ffn_norm_out = self.ffn_norm(h, r_idx)
            ffn_out = self.ffn(ffn_norm_out)
            
            return h + ffn_out

class RecursiveCausalLM(nn.Module):
    """The main Unified Recurrent Language Model (Prelude-Core-Coda Layout)."""
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.embeddings = nn.Embedding(config.vocab_size, config.d_model)
        
        # Positional step-embeddings
        self.step_embeddings = nn.Parameter(torch.zeros(config.n_iterations, config.d_model))
        nn.init.normal_(self.step_embeddings, std=0.02)
        
        # Ouroboros Depth Gating parameter with Smooth Sigmoid Scaling
        self.depth_gate = nn.Parameter(torch.ones(config.n_iterations, config.d_model) * 1.73) # ~0.85 Sigmoid value
        
        # Decoupled LoRAExit early-exit adapters
        self.lora_exits = nn.ModuleList([
            LoRAExit(config.d_model, r=32) for _ in range(config.n_iterations)
        ])
        
        # AdaExit 1-parameter binary halt-classifier
        self.halt_head = nn.Linear(config.d_model, 1, bias=False)
        nn.init.zeros_(self.halt_head.weight)
        
        # Partition 1: Prelude (4 unshared standard layers)
        self.prelude = nn.ModuleList([TransformerBlock(config, layer_type="standard") for _ in range(4)])
        
        # Partition 2: Recurrent Core (2 stage-tied core layers)
        self.core_blocks = nn.ModuleList([TransformerBlock(config, layer_type="recurrent_core") for _ in range(2)])
        
        # Partition 3: Coda (4 unshared standard layers)
        self.coda = nn.ModuleList([TransformerBlock(config, layer_type="standard") for _ in range(4)])
        
        # Controller Hypernetwork (outputs diagonal modulation scales for 6 projections)
        self.controller = ControllerHypernetwork(config.d_model, num_modulated_projs=6, rank_ctrl=64)
        
        # MLA anchor projection down-projection weights
        self.dkv_weight = nn.Parameter(torch.empty(128, config.d_model))
        nn.init.normal_(self.dkv_weight, std=0.02)
        
        self.final_norm = RMSNorm(config.d_model)
        self.lm_head_bias = nn.Parameter(torch.zeros(config.vocab_size)) if config.bias else None
        
        # Speculative parallel decoding heads
        self.speculative_projs = nn.ModuleList([
            nn.Linear(config.d_model, config.d_model, bias=config.bias) for _ in range(4)
        ])
        self.speculative_biases = nn.ParameterList([
            nn.Parameter(torch.zeros(config.vocab_size)) for _ in range(4)
        ])
        
        # First-pass initialization
        self.apply(self._init_weights)
        
        # Explicit zero-init for speculative biases
        for bias in self.speculative_biases:
            nn.init.zeros_(bias)
            
        # Re-zero step adapters to preserve identities at step 0
        for core_block in self.core_blocks:
            for adapter in core_block.adapters:
                nn.init.zeros_(adapter.down_proj.weight)
            
        # Scaled initialization for residual projection layers
        std_scale = 1.0 / (2 * config.n_iterations) ** 0.5
        with torch.no_grad():
            for core_block in self.core_blocks:
                core_block.attn.out_proj.base_layer.weight.mul_(std_scale)
                core_block.ffn.down_proj.base_layer.weight.mul_(std_scale)

    @property
    def core_block(self):
        """Property to support legacy code that accesses core_block directly."""
        return self.core_blocks[0]

    @property
    def block(self):
        """Property to keep diagnostic and audit scripts perfectly backward compatible."""
        return self.core_block

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, std=0.02)

    def load_state_dict(self, state_dict, strict=False):
        """Saves representational accuracy by automatically copying unshared block weights 
        and initializing SVD low-rank bases.
        """
        new_state_dict = {}
        for k, v in state_dict.items():
            if k in self.state_dict():
                new_state_dict[k] = v
                continue
                
            if k.startswith("block."):
                # 1. Map to prelude layers
                for idx in range(4):
                    pre_key = k.replace("block.", f"prelude.{idx}.")
                    if pre_key in self.state_dict():
                        new_state_dict[pre_key] = v
                        
                # 2. Map to recurrent core blocks (including bases)
                for b_idx in range(2):
                    core_key = k.replace("block.", f"core_blocks.{b_idx}.")
                    if "attn.k_proj" in core_key or "attn.v_proj" in core_key:
                        pass
                    elif core_key in self.state_dict():
                        new_state_dict[core_key] = v
                    else:
                        if ".q_proj." in k:
                            new_state_dict[core_key.replace(".q_proj.", ".q_proj.base_layer.")] = v
                        elif ".out_proj." in k:
                            new_state_dict[core_key.replace(".out_proj.", ".out_proj.base_layer.")] = v
                        elif ".gate_proj." in k:
                            new_state_dict[core_key.replace(".gate_proj.", ".gate_proj.base_layer.")] = v
                        elif ".up_proj." in k:
                            new_state_dict[core_key.replace(".up_proj.", ".up_proj.base_layer.")] = v
                        elif ".down_proj." in k:
                            new_state_dict[core_key.replace(".down_proj.", ".down_proj.base_layer.")] = v
                        
                # 3. Map to coda layers
                for idx in range(4):
                    cod_key = k.replace("block.", f"coda.{idx}.")
                    if cod_key in self.state_dict():
                        new_state_dict[cod_key] = v
                        
        res = super().load_state_dict(new_state_dict, strict=False)
        
        # Trigger SVD initialization on ModulatedLinear modules
        with torch.no_grad():
            for name, module in self.named_modules():
                if isinstance(module, ModulatedLinear):
                    module.initialize_svd()
                    
        return res

    def get_num_params(self, unique_only: bool = True):
        """Returns the physical parameters saved on disk vs unrolled virtual parameter capacity."""
        if unique_only:
            return sum(p.numel() for p in self.parameters())
        else:
            embeds = self.embeddings.weight.numel()
            steps = self.step_embeddings.numel()
            gate = self.depth_gate.numel()
            halt = self.halt_head.weight.numel()
            norm = sum(p.numel() for p in self.final_norm.parameters())
            bias_count = self.lm_head_bias.numel() if self.lm_head_bias is not None else 0
            
            # Virtual unrolled layers capacity
            pre_params = sum(p.numel() for p in self.prelude.parameters())
            core_unrolled = sum(p.numel() for p in self.core_blocks.parameters()) * 4 # 2 core blocks unrolled 4 times each (total 8)
            coda_params = sum(p.numel() for p in self.coda.parameters())
            
            return embeds + steps + gate + halt + norm + bias_count + pre_params + core_unrolled + coda_params

    def forward(self, input_ids, targets=None, kv_cache=None, return_spec: bool = False, spec_coef: float = 0.0):
        B, T = input_ids.shape
        x = self.embeddings(input_ids)
        
        # 1. Prelude Phase (Layers 0 to 3)
        for idx in range(4):
            if self.training and targets is not None:
                def make_prelude_fn(layer_idx):
                    def custom_forward(tensor_in):
                        return self.prelude[layer_idx](tensor_in, r_idx=layer_idx)
                    return custom_forward
                x = cp.checkpoint(make_prelude_fn(idx), x, use_reentrant=False)
            else:
                cache_info = (kv_cache, idx, "prelude") if kv_cache is not None else None
                x = self.prelude[idx](x, r_idx=idx, kv_cache=cache_info)
                
        # 2. Extract & Cache compressed MLA Anchor
        c_kv = F.linear(x, self.dkv_weight) # [B, T, 128]
        if kv_cache is not None:
            anchor_ckv = kv_cache.update_anchor(c_kv)
        else:
            anchor_ckv = c_kv
            
        # 3. Recurrent Core Phase (8 iterations: steps 4 to 11)
        xs = []
        for r in range(4, 12): # Recurrent core runs exactly R=8 times
            # Add step-dependent embedding to indicate virtual depth level
            x = x + self.step_embeddings[r].view(1, 1, -1)
            
            # Ouroboros Depth Gating with Smooth Sigmoid Scaling
            gate = (0.02 + 0.96 * torch.sigmoid(self.depth_gate[r])).view(1, 1, -1)
            
            # Controller generates input-conditioned modulation scales directly from hidden states
            mod_vector = self.controller(x, self.step_embeddings[r])
            
            block_idx = 0 if r < 8 else 1
            curr_core_block = self.core_blocks[block_idx]
            
            if self.training and targets is not None:
                def make_core_fn(r_val, b_idx):
                    def custom_forward(tensor_in, anchor, mod):
                        return self.core_blocks[b_idx](tensor_in, r_idx=r_val, anchor_ckv=anchor, mod_vector=mod)
                    return custom_forward
                block_out = cp.checkpoint(make_core_fn(r, block_idx), x, anchor_ckv, mod_vector, use_reentrant=False)
            else:
                block_out = curr_core_block(x, r_idx=r, anchor_ckv=anchor_ckv, mod_vector=mod_vector)
                
            x = gate * x + (1.0 - gate) * block_out
            
            if targets is not None:
                xs.append(x)
                
            # Early exit checking (Inference only, non-cached mode)
            is_training_mode = self.training or (targets is not None)
            is_cached_inference = (kv_cache is not None)
            
            if not is_training_mode and not is_cached_inference and T == 1:
                # Run through early-exit LoRAExit adapter to decouple intermediate gradients
                x_exit = x[:, -1, :] + self.lora_exits[r](x[:, -1, :])
                halt_logit = self.halt_head(x_exit)
                halt_prob = torch.sigmoid(halt_logit).min().item()
                if halt_prob > 0.95:
                    break
                    
        # 4. Coda Phase (Layers 12 to 15)
        for idx in range(4):
            virtual_idx = 12 + idx
            if self.training and targets is not None:
                def make_coda_fn(layer_idx):
                    def custom_forward(tensor_in):
                        return self.coda[layer_idx](tensor_in, r_idx=virtual_idx)
                    return custom_forward
                x = cp.checkpoint(make_coda_fn(idx), x, use_reentrant=False)
            else:
                cache_info = (kv_cache, idx, "coda") if kv_cache is not None else None
                x = self.coda[idx](x, r_idx=virtual_idx, kv_cache=cache_info)
                
        x = self.final_norm(x)
        logits = F.linear(x, self.embeddings.weight, self.lm_head_bias)
        
        # Speculative Parallel Decoding projection
        spec_logits = []
        for k in range(4):
            h_k = F.silu(self.speculative_projs[k](x))
            logits_k = F.linear(h_k, self.embeddings.weight, self.speculative_biases[k])
            spec_logits.append(logits_k)
            
        loss = None
        if targets is not None:
            loss_lm = F.cross_entropy(logits.to(torch.float32).view(-1, logits.size(-1)), targets.view(-1))
            
            # Early-Exit (Halt Head) Distillation Loss using Decoupled LoRAExit
            loss_halt = 0.0
            num_halt_steps = len(xs) - 1
            if num_halt_steps > 0:
                final_x = xs[-1].detach()
                for i in range(num_halt_steps):
                    r = 4 + i
                    sim = F.cosine_similarity(xs[i], final_x, dim=-1)
                    target_halt = (sim >= 0.985).to(dtype=xs[i].dtype)
                    
                    # Decoupled LoRAExit adapter path to evaluate exit classification
                    x_exit = xs[i] + self.lora_exits[r](xs[i])
                    halt_logits = self.halt_head(x_exit).squeeze(-1)
                    loss_halt += F.binary_cross_entropy_with_logits(
                        halt_logits.to(torch.float32), 
                        target_halt.to(torch.float32)
                    )
                loss_halt = loss_halt / num_halt_steps
                
            # Speculative heads training loss
            if spec_coef > 0.0:
                loss_spec = 0.0
                for k in range(4):
                    shift_len = k + 1
                    if T > shift_len:
                        logits_slice = spec_logits[k][:, :-shift_len, :].contiguous()
                        targets_slice = targets[:, shift_len:].contiguous()
                        loss_spec_k = F.cross_entropy(
                            logits_slice.to(torch.float32).view(-1, logits_slice.size(-1)), 
                            targets_slice.view(-1)
                        )
                        loss_spec += loss_spec_k
                loss_spec = loss_spec / 4.0
                loss = loss_lm + 0.1 * loss_halt + spec_coef * loss_spec
            else:
                loss = loss_lm + 0.1 * loss_halt
                
        if return_spec:
            stacked_spec_logits = torch.stack(spec_logits, dim=2)
            return logits, loss, stacked_spec_logits
        else:
            return logits, loss

class KVCache:
    """A highly-optimized key-value cache that stores states recursively across prelude, anchor, and coda phases."""
    def __init__(self, config: ModelConfig, max_batch_size: int, device: str, dtype: torch.dtype = torch.float16, max_seq_len: int = None):
        super().__init__()
        self.max_batch_size = max_batch_size
        self.head_dim = config.d_model // config.n_heads
        self.dtype = dtype
        
        cache_seq_len = max_seq_len if max_seq_len is not None else config.max_seq_len
        self.cache_seq_len = cache_seq_len
        
        # Prelude Cache (4 layers)
        self.prelude_k = torch.zeros(4, max_batch_size, config.n_kv_heads, cache_seq_len, self.head_dim, device=device, dtype=dtype)
        self.prelude_v = torch.zeros(4, max_batch_size, config.n_kv_heads, cache_seq_len, self.head_dim, device=device, dtype=dtype)
        self.prelude_lengths = [0] * 4
        
        # Anchor C_KV cache
        self.anchor_ckv = torch.zeros(max_batch_size, cache_seq_len, 128, device=device, dtype=dtype)
        self.anchor_length = 0
        
        # Coda Cache (4 layers)
        self.coda_k = torch.zeros(4, max_batch_size, config.n_kv_heads, cache_seq_len, self.head_dim, device=device, dtype=dtype)
        self.coda_v = torch.zeros(4, max_batch_size, config.n_kv_heads, cache_seq_len, self.head_dim, device=device, dtype=dtype)
        self.coda_lengths = [0] * 4
        
        # Compatibility current lengths
        self.current_lengths = [0] * config.n_iterations

    def update_prelude(self, k_new, v_new, layer_idx: int):
        B, T_new, H_kv, d_k = k_new.shape
        start_pos = self.prelude_lengths[layer_idx]
        end_pos = start_pos + T_new
        
        k_new = k_new.to(self.prelude_k.dtype)
        v_new = v_new.to(self.prelude_v.dtype)
        
        self.prelude_k[layer_idx, :B, :, start_pos:end_pos, :] = k_new.transpose(1, 2)
        self.prelude_v[layer_idx, :B, :, start_pos:end_pos, :] = v_new.transpose(1, 2)
        
        self.prelude_lengths[layer_idx] = end_pos
        self.current_lengths[layer_idx] = end_pos
        
        k_out = self.prelude_k[layer_idx, :B, :, :end_pos, :].transpose(1, 2)
        v_out = self.prelude_v[layer_idx, :B, :, :end_pos, :].transpose(1, 2)
        return k_out, v_out

    def update_coda(self, k_new, v_new, layer_idx: int):
        B, T_new, H_kv, d_k = k_new.shape
        start_pos = self.coda_lengths[layer_idx]
        end_pos = start_pos + T_new
        
        k_new = k_new.to(self.coda_k.dtype)
        v_new = v_new.to(self.coda_v.dtype)
        
        self.coda_k[layer_idx, :B, :, start_pos:end_pos, :] = k_new.transpose(1, 2)
        self.coda_v[layer_idx, :B, :, start_pos:end_pos, :] = v_new.transpose(1, 2)
        
        self.coda_lengths[layer_idx] = end_pos
        self.current_lengths[12 + layer_idx] = end_pos
        
        k_out = self.coda_k[layer_idx, :B, :, :end_pos, :].transpose(1, 2)
        v_out = self.coda_v[layer_idx, :B, :, :end_pos, :].transpose(1, 2)
        return k_out, v_out

    def update_anchor(self, c_kv_new):
        B, T_new, r_lat = c_kv_new.shape
        start_pos = self.anchor_length
        end_pos = start_pos + T_new
        
        self.anchor_ckv[:B, start_pos:end_pos, :] = c_kv_new.to(self.anchor_ckv.dtype)
        self.anchor_length = end_pos
        
        return self.anchor_ckv[:B, :end_pos, :]