import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from dataclasses import dataclass

# ============================================================================
# 1. HELIX STATE SPACE MEMORY (DNA-inspired curved memory)
# ============================================================================
class HelixMemory(nn.Module):
    """
    Double-helix memory with semantic (strand A) and structural (strand B)
    strands, connected by learned cross-links. Gives effectively O(log L)
    memory scaling for unlimited context.
    
    CPU-optimized: pre-allocated slots, no .item() calls, torch.compile friendly.
    """
    def __init__(self, d_model, helix_slots=1024, max_context=1048576):
        super().__init__()
        self.d_model = d_model
        self.helix_slots = helix_slots
        self.max_context = max_context
        
        self.encode = nn.Linear(d_model * 3, d_model)
        self.query_proj = nn.Linear(d_model, d_model)
        self.key_proj = nn.Linear(d_model, d_model)
        self.cross_link_net = nn.Sequential(
            nn.Linear(d_model * 2, d_model), nn.SiLU(), nn.Linear(d_model, 1))
        self.ast_embed = nn.Embedding(1024, d_model)
        self.output_proj = nn.Linear(d_model * 2, d_model)
        self.compress_net = nn.Linear(d_model, d_model)
        self.ema_alpha = nn.Parameter(torch.tensor(0.9))
        self.tau_insert = nn.Parameter(torch.tensor(0.5))
    
    def init_helix(self, batch_size, device):
        max_s = self.helix_slots
        return {
            'S': torch.zeros(batch_size, max_s, self.d_model, device=device),
            'active': torch.zeros(batch_size, max_s, dtype=torch.bool, device=device),
            'n_active': torch.zeros(batch_size, dtype=torch.long, device=device),
            'step': torch.zeros(batch_size, dtype=torch.long, device=device),
        }
    
    def forward(self, x, helix_state, ast_paths=None):
        B, L, D = x.shape
        S = helix_state['S'].detach()
        active = helix_state['active']
        n_active = helix_state['n_active']
        max_s = S.shape[1]
        
        S_new = S.clone()
        active_new = active.clone()
        n_active_new = n_active.clone()
        
        # Pre-compute local context for all positions
        local_ctx = torch.stack([
            x[:, max(0, i-64):i].mean(dim=1)
            if i > 0 else torch.zeros(B, D, device=x.device)
            for i in range(L)
        ], dim=1)  # [B, L, D]
        
        if ast_paths is not None:
            ast_h = self.ast_embed(ast_paths.clamp(0, 1023))  # [B, L, D]
        else:
            ast_h = torch.zeros_like(x)
        
        # Pre-compute all candidates
        encode_in = torch.cat([x, local_ctx, ast_h], dim=-1)  # [B, L, 3*D]
        candidates = self.encode(encode_in)  # [B, L, D]
        
        threshold = torch.sigmoid(self.tau_insert)
        alpha = torch.sigmoid(self.ema_alpha)
        
        for i in range(L):
            cand = candidates[:, i:i+1]  # [B, 1, D]
            
            # Resonance with active slots only
            resonance = torch.bmm(S_new, cand.transpose(1, 2)).squeeze(-1)  # [B, max_s]
            resonance = resonance.masked_fill(~active_new, -1e9)
            max_res, best_idx = resonance.max(dim=-1, keepdim=True)  # [B, 1]
            
            # Decide: update existing slot vs append new slot
            do_update = (max_res > threshold)  # [B, 1]
            
            # Batch update: where resonance is high enough, EMA merge
            for b_idx in range(B):
                if do_update[b_idx, 0]:
                    idx = best_idx[b_idx, 0]
                    merged = alpha * S_new[b_idx:b_idx+1, idx:idx+1] + (1 - alpha) * cand[b_idx:b_idx+1]
                    S_new[b_idx:b_idx+1, idx:idx+1] = merged.detach()
                else:
                    new_idx = n_active_new[b_idx]
                    if new_idx < max_s:
                        S_new[b_idx:b_idx+1, new_idx:new_idx+1] = cand[b_idx:b_idx+1].detach()
                        active_new[b_idx, new_idx] = True
                        n_active_new[b_idx] = new_idx + 1
                        
                        # Compress if near limit
                        if new_idx + 1 >= max_s:
                            compressed = self._compress_batch(
                                S_new[b_idx:b_idx+1], active_new[b_idx:b_idx+1], D, max_s)
                            S_new[b_idx:b_idx+1] = compressed
                            active_new[b_idx:b_idx+1] = active_new[b_idx:b_idx+1]
        
        helix_state['S'] = S_new
        helix_state['active'] = active_new
        helix_state['step'] += L
        
        # Retrieval from active slots
        last_query = self.query_proj(x[:, -1:])  # [B, 1, D]
        S_retrieve = S_new.clone()
        resonance = torch.bmm(S_retrieve, last_query.transpose(1, 2)).squeeze(-1)  # [B, max_s]
        resonance = resonance.masked_fill(~active_new, -1e9)
        
        top_k = min(16, max_s)
        _, top_indices = torch.topk(resonance, top_k, dim=-1)
        
        retrieved_list = []
        for b_idx in range(B):
            idxs = top_indices[b_idx]
            sel = S_retrieve[b_idx:b_idx+1, idxs]  # [1, k, D]
            c_weights = F.softmax(
                (sel @ sel.transpose(1, 2)).squeeze(0) / math.sqrt(D), dim=-1
            )
            assoc = (c_weights @ sel.squeeze(0)).mean(dim=0, keepdim=True).unsqueeze(0)
            retrieved_list.append(assoc)
        
        retrieved = torch.cat(retrieved_list, dim=0)  # [B, 1, D]
        enhanced = self.output_proj(torch.cat([x[:, -1:], retrieved], dim=-1))
        return enhanced, helix_state
    
    def _compress_batch(self, S_b, active_b, D, max_s):
        k = active_b.sum().item()
        if k <= 64 or k < 4:
            return S_b
        active_indices = active_b.nonzero(as_tuple=True)[0]
        clusters = torch.split(active_indices, max(4, k // 16))
        result = []
        for cluster in clusters:
            cluster_sel = S_b[0:1, cluster]
            corr = (cluster_sel @ cluster_sel.transpose(1, 2)).sum(dim=-1)
            weights = F.softmax(corr, dim=-1)
            merged = (weights.unsqueeze(-1) * cluster_sel).sum(dim=1, keepdim=True)
            result.append(merged)
        S_compressed = torch.cat(result, dim=1)
        n_compressed = S_compressed.shape[1]
        if n_compressed < max_s:
            pad = torch.zeros(1, max_s - n_compressed, D, device=S_b.device)
            S_compressed = torch.cat([S_compressed, pad], dim=1)
        S_b[:, :max_s] = S_compressed[:, :max_s]
        return S_b


# ============================================================================
# 2. HIERARCHICAL CODE ATTENTION (HCA)
# ============================================================================
class HierarchicalCodeAttention(nn.Module):
    """
    Three-tier attention: local window, AST-aware structure, global sparse.
    """
    def __init__(self, d_model, n_heads, kv_heads=None, window_size=128, 
                 local_heads=8, struct_heads=4, global_heads=4):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.kv_heads = kv_heads or n_heads
        self.window_size = window_size
        self.head_dim = d_model // n_heads
        self.groups = n_heads // self.kv_heads
        
        self.local_heads = local_heads
        self.struct_heads = struct_heads
        self.global_heads = global_heads
        assert local_heads + struct_heads + global_heads == n_heads
        
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, self.kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d_model, self.kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)
        
        self.struct_q = nn.Linear(d_model, struct_heads * self.head_dim, bias=False)
        self.struct_k = nn.Linear(d_model * 2, struct_heads * self.head_dim, bias=False)
        
        self.global_router = nn.Linear(d_model, global_heads * 2, bias=False)
        self.top_k = 32
        
    def forward(self, x, mask=None, ast_embeds=None):
        B, L, D = x.shape
        H, Hk, Hd = self.n_heads, self.kv_heads, self.head_dim
        
        q = self.q_proj(x).view(B, L, H, Hd).transpose(1,2)
        k = self.k_proj(x).view(B, L, Hk, Hd).transpose(1,2)
        v = self.v_proj(x).view(B, L, Hk, Hd).transpose(1,2)
        
        out = torch.zeros(B, H, L, Hd, device=x.device, dtype=x.dtype)
        
        # Level 1: Local sliding window
        local_q = q[:, :self.local_heads]
        n_local_kv = max(1, self.local_heads // self.groups)
        local_k = k[:, :n_local_kv]
        local_v = v[:, :n_local_kv]
        for i in range(L):
            start = max(0, i - self.window_size)
            w_k = local_k[:, :, start:i+1]
            w_v = local_v[:, :, start:i+1]
            q_i = local_q[:, :, i]  # [B, local_h, Hd]
            scores = torch.matmul(q_i.unsqueeze(2), w_k.transpose(-2,-1)).squeeze(2) / math.sqrt(Hd)
            if mask is not None:
                m = mask[:, 0, 0, start:i+1]  # [B, W]
                scores = scores.masked_fill(~m.unsqueeze(1).expand(-1, self.local_heads, -1), -1e9)
            attn = F.softmax(scores, dim=-1)
            out[:, :self.local_heads, i] = torch.matmul(attn.unsqueeze(2), w_v).squeeze(2)
        
        # Level 2: AST-aware structure heads
        if ast_embeds is not None and self.struct_heads > 0:
            s_q = self.struct_q(x).view(B, L, self.struct_heads, Hd).transpose(1,2)
            s_input = torch.cat([x, ast_embeds], dim=-1)
            s_k = self.struct_k(s_input).view(B, L, self.struct_heads, Hd).transpose(1,2)
            n_s_kv = max(1, self.struct_heads // self.groups)
            s_v = v[:, :n_s_kv]
            scores = torch.matmul(s_q, s_k.transpose(-2,-1)) / math.sqrt(Hd)
            if mask is not None:
                scores = scores.masked_fill(~mask[:, :, :, :L], -1e9)
            attn = F.softmax(scores, dim=-1)
            start = self.local_heads
            end = self.local_heads + self.struct_heads
            out[:, start:end] = torch.matmul(attn, s_v)
        
        # Level 3: Global attention
        if self.global_heads > 0:
            n_g_kv = max(1, self.global_heads // self.groups)
            q_g = q[:, -self.global_heads:]  # [B, g_h, L, Hd]
            k_g = k[:, :n_g_kv]  # [B, n_g_kv, L, Hd]
            v_g = v[:, :n_g_kv]  # [B, n_g_kv, L, Hd]
            scores = torch.matmul(q_g, k_g.transpose(-2,-1)) / math.sqrt(Hd)
            if mask is not None:
                scores = scores.masked_fill(~mask[:, :, :, :L], -1e9)
            attn = F.softmax(scores, dim=-1)
            g_out = torch.matmul(attn, v_g)
            if g_out.shape[1] != self.global_heads:
                g_out = g_out[:, :self.global_heads]
            out[:, -self.global_heads:] = g_out
        
        out = out.transpose(1,2).contiguous().view(B, L, D)
        out = self.o_proj(out)
        return out


# ============================================================================
# 3. EXECUTION-AUGMENTED FFN (EA-FFN)
# ============================================================================
class ExecutionAugmentedFFN(nn.Module):
    """
    Two-stream FFN: standard SwiGLU + execution trace stream.
    """
    def __init__(self, d_model, d_ff=None, trace_dim=256):
        super().__init__()
        d_ff = d_ff or d_model * 4
        self.trace_dim = trace_dim
        
        # Stream A: Standard SwiGLU
        self.gate_proj = nn.Linear(d_model, d_ff, bias=False)
        self.up_proj = nn.Linear(d_model, d_ff, bias=False)
        self.down_proj = nn.Linear(d_ff, d_model, bias=False)
        
        # Stream B: Execution trace
        self.trace_proj = nn.Linear(d_model + trace_dim, d_ff, bias=False)
        self.trace_down = nn.Linear(d_ff, d_model, bias=False)
        
        # Gate
        self.gate_net = nn.Linear(d_model, d_model)
        
    def forward(self, x, trace=None):
        # Stream A
        a = F.silu(self.gate_proj(x)) * self.up_proj(x)
        a = self.down_proj(a)
        
        # Stream B (with execution trace if available)
        b = 0
        if trace is not None:
            b_input = torch.cat([x, trace], dim=-1)
            b = F.silu(self.trace_proj(b_input))
            b = self.trace_down(b)
        
        # Learned gating
        gate = torch.sigmoid(self.gate_net(x))
        return gate * a + (1 - gate) * b


# ============================================================================
# 4. ROPE WITH STRUCTURAL BIAS (RoPE-S)
# ============================================================================
class RoPEWithStructuralBias(nn.Module):
    """
    Rotary position encoding with structural bias terms for AST depth,
    scope, and control flow nesting.
    """
    def __init__(self, d_model, max_len=131072, base=10000.0):
        super().__init__()
        self.d_model = d_model
        self.max_len = max_len
        inv_freq = 1.0 / (base ** (torch.arange(0, d_model, 2).float() / d_model))
        self.register_buffer('inv_freq', inv_freq)
        
        self.struct_bias = nn.Linear(4, d_model // 2, bias=False)
        
    def forward(self, x, positions, ast_depth=None, scope_id=None, 
                ctrl_flow=None, branch_id=None):
        B, L, D = x.shape
        # positions: [B, L] or [L] - handle both
        if positions.dim() == 2:
            pos = positions.float()  # [B, L]
            freqs = pos.unsqueeze(-1) * self.inv_freq.unsqueeze(0).unsqueeze(0)  # [B, L, D/2]
        else:
            pos = positions.float()
            freqs = torch.outer(pos, self.inv_freq)  # [L, D/2]
        emb = torch.cat([freqs.sin(), freqs.cos()], dim=-1)  # [B, L, D] or [L, D]
        if emb.dim() == 2:
            emb = emb.unsqueeze(0).expand(B, -1, -1)
        
        if ast_depth is not None:
            ast_h = ast_depth.float() / 32.0
            sc_h = (scope_id.float() / 256.0) if scope_id is not None else pos / 256.0
            cf_h = (ctrl_flow.float() / 16.0) if ctrl_flow is not None else pos / 16.0
            br_h = (branch_id.float() / 8.0) if branch_id is not None else pos / 8.0
            struct_feats = torch.stack([ast_h, sc_h, cf_h, br_h], dim=-1)
            struct_bias = self.struct_bias(struct_feats)
            emb = emb + torch.cat([struct_bias, struct_bias], dim=-1)
        
        # Apply rotary: x * cos + rotate_half(x) * sin
        x_rot = x * emb.cos() + torch.stack([-x[..., 1::2], x[..., ::2]], dim=-1).reshape(x.shape) * emb.sin()
        return x_rot


# ============================================================================
# 5. PREFIX-PRESERVING NORMALIZATION (PPN)
# ============================================================================
class PrefixPreservingNorm(nn.Module):
    """
    Block-diagonal normalization that preserves prefix subspaces.
    Enables exact KV-cache sharing across prefix lengths.
    """
    def __init__(self, d_model, n_groups=4, eps=1e-5):
        super().__init__()
        self.n_groups = n_groups
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))
        self.bias = nn.Parameter(torch.zeros(d_model))
        self.group_size = d_model // n_groups
        
    def forward(self, x):
        B, L, D = x.shape
        x = x.view(B, L, self.n_groups, self.group_size)
        var = x.var(dim=-1, keepdim=True, unbiased=False)
        x = x / torch.sqrt(var + self.eps)
        x = x.view(B, L, D)
        return x * self.weight + self.bias


# ============================================================================
# 6. MIXTURE-OF-DEPTH (MoD) ROUTER
# ============================================================================
class MixtureOfDepth(nn.Module):
    """
    Per-token depth routing. Each token chooses how many layers to execute.
    """
    def __init__(self, d_model, n_layers, n_buckets=4):
        super().__init__()
        self.n_layers = n_layers
        self.n_buckets = n_buckets
        self.bucket_sizes = [n_layers // 2 ** i for i in range(n_buckets)]
        self.bucket_sizes[0] = max(1, n_layers // 4)
        self.router = nn.Sequential(
            nn.Linear(d_model + 8, d_model), nn.SiLU(),
            nn.Linear(d_model, n_buckets))
        
    def forward(self, x, depth_features=None):
        B, L, D = x.shape
        features = depth_features if depth_features is not None else x.new_zeros(B, L, 8)
        logits = self.router(torch.cat([x, features], dim=-1))
        probs = F.softmax(logits, dim=-1)
        
        bucket = probs.argmax(dim=-1)
        depths = torch.tensor([self.bucket_sizes[b.item()] for b in bucket[0]], 
                              device=x.device, dtype=torch.long)
        
        aux_loss = -torch.mean(probs * torch.log(probs + 1e-8))  # entropy regularization
        return depths, probs, aux_loss


# ============================================================================
# 7. FSI_EDGE TRANSFORMER BLOCK
# ============================================================================
class FSIEdgeBlock(nn.Module):
    def __init__(self, layer_id, config):
        super().__init__()
        self.layer_id = layer_id
        self.config = config
        
        self.helix = HelixMemory(config.d_model, config.helix_slots)
        self.hca = HierarchicalCodeAttention(
            config.d_model, config.n_heads, config.kv_heads,
            config.window_size, config.local_heads,
            config.struct_heads, config.global_heads)
        self.eaffn = ExecutionAugmentedFFN(config.d_model, config.d_ff, config.trace_dim)
        self.ppn1 = PrefixPreservingNorm(config.d_model, config.norm_groups)
        self.ppn2 = PrefixPreservingNorm(config.d_model, config.norm_groups)
        self.rope_s = RoPEWithStructuralBias(config.d_model, config.max_seq_len)
        
    def forward(self, x, helmet_state, mask=None, positions=None,
                ast_embeds=None, ast_depth=None, scope_id=None,
                ctrl_flow=None, branch_id=None, trace=None):
        
        # Pre-norm + RoPE-S
        h = self.ppn1(x)
        h = self.rope_s(h, positions, ast_depth, scope_id, ctrl_flow, branch_id)
        
        # HCA + Helix memory
        h = self.hca(h, mask, ast_embeds)
        helix_out, helmet_state = self.helix(h, helmet_state)
        x = x + h + helix_out
        
        # EA-FFN
        h = self.ppn2(x)
        h = self.eaffn(h, trace)
        x = x + h
        
        return x, helmet_state


# ============================================================================
# 8. FSI_EDGE MAIN MODEL
# ============================================================================

@dataclass
class FSIEdgeConfig:
    vocab_size: int = 32768
    d_model: int = 1536
    n_layers: int = 28
    n_heads: int = 24
    kv_heads: int = 6
    d_ff: int = 6144
    max_seq_len: int = 16384
    window_size: int = 128
    local_heads: int = 14
    struct_heads: int = 6
    global_heads: int = 4
    helix_slots: int = 1024
    trace_dim: int = 256
    norm_groups: int = 4
    rope_base: float = 10000.0
    moe_n_experts: int = 1
    moe_top_k: int = 1
    dropout: float = 0.0
    init_std: float = 0.02

    def __post_init__(self):
        total = self.local_heads + self.struct_heads + self.global_heads
        if total != self.n_heads:
            ratio = self.n_heads / total
            self.local_heads = max(1, int(self.local_heads * ratio))
            self.struct_heads = max(1, int(self.struct_heads * ratio))
            self.global_heads = self.n_heads - self.local_heads - self.struct_heads
            if self.global_heads < 1:
                self.global_heads = 1
                self.local_heads = self.n_heads - self.struct_heads - self.global_heads


class FSIEdgeModel(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embed = nn.Embedding(config.vocab_size, config.d_model)
        self.ast_type_embed = nn.Embedding(64, config.d_model)
        
        self.layers = nn.ModuleList([
            FSIEdgeBlock(i, config) for i in range(config.n_layers)
        ])
        
        self.mod = MixtureOfDepth(config.d_model, config.n_layers)
        self.final_norm = PrefixPreservingNorm(config.d_model, config.norm_groups)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        
        self.apply(self._init_weights)
        
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.init_std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.init_std)
    
    def forward(self, input_ids, ast_types=None, ast_depths=None,
                scope_ids=None, ctrl_flows=None, branch_ids=None,
                traces=None, attention_mask=None, labels=None):
        B, L = input_ids.shape
        device = input_ids.device
        
        x = self.embed(input_ids)
        if ast_types is not None:
            x = x + self.ast_type_embed(ast_types)
        
        positions = torch.arange(L, device=device).unsqueeze(0).expand(B, -1)
        mask = attention_mask.unsqueeze(1).unsqueeze(2).bool() if attention_mask is not None else None
        
        ast_embeds = None
        if ast_types is not None:
            ast_embeds = self.ast_type_embed(ast_types)
        
        # Helix init at layer 0
        helmet_state = self.layers[0].helix.init_helix(B, device)
        
        for layer in self.layers:
            x, helmet_state = layer(
                x, helmet_state, mask, positions,
                ast_embeds, ast_depths, scope_ids,
                ctrl_flows, branch_ids, traces)
        
        x = self.final_norm(x)
        logits = self.lm_head(x)
        
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=0)
        
        return FSIEdgeOutput(loss=loss, logits=logits, helmet_state=helmet_state)


@dataclass
class FSIEdgeOutput:
    loss: torch.Tensor = None
    logits: torch.Tensor = None
    helmet_state: dict = None