"""
LaminarNet v0.3.2 — Selective State Space Architecture (Rock-Solid Stability)
Numerically stabilized selective recurrence with Log-Sum-Exp trick.
"""

import math
from dataclasses import dataclass
from typing import Optional, List

import torch
import torch.nn as nn
import torch.nn.functional as F


@dataclass
class LaminarNetConfig:
    vocab_size: int = 32000
    d_model: int = 256
    n_heads: int = 8
    n_layers: int = 6
    d_ff: int = 512
    n_strata: int = 2
    strata_ratios: tuple = (1, 4)
    seq_len: int = 2048
    dropout: float = 0.1
    conv_kernel: int = 5


class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(d))
        self.eps = eps
    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return self.scale * x * rms


# ─────────────────────────────────────────────────────────────
# 1. Selective Geometric Drift (SGD) — High-Throughput v3.3
# ─────────────────────────────────────────────────────────────

class GeometricDriftField(nn.Module):
    """
    Geometric Drift Field v3.3 — High-Throughput Selective Scan.
    
    Optimized for VRAM and Speed by replacing the (T, T) kernel with a 
    more efficient chunked recurrence and using smaller block sizes.
    """

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        
        self.conv1d = nn.Conv1d(d_model, d_model, kernel_size=4, padding=3, groups=d_model)
        
        self.dt_proj = nn.Linear(d_model, d_model, bias=True)
        self.to_v = nn.Linear(d_model, d_model, bias=False)
        self.to_gate = nn.Linear(d_model, d_model, bias=False)
        self.to_theta = nn.Linear(d_model, d_model)
        
        self.to_output = nn.Linear(d_model, d_model)
        self.norm = RMSNorm(d_model)
        self.dropout = nn.Dropout(dropout)

        # Initialize dt bias to ensure stable starting decay
        nn.init.constant_(self.dt_proj.bias, -3.0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.norm(x)
        B, N, D = x.shape
        
        # 1. Local Context
        x_conv = self.conv1d(x.transpose(1, 2))[..., :N].transpose(1, 2)
        x_conv = F.silu(x_conv)
        
        # 2. Selective Parameters
        # dt: (B, N, D)
        dt = F.softplus(self.dt_proj(x_conv)).clamp(min=0.001, max=5.0) 
        v = self.to_v(x_conv)
        gate = torch.sigmoid(self.to_gate(x))
        
        # 3. Positional Rotation
        theta = torch.tanh(self.to_theta(x)) * 0.1
        theta = theta.view(B, N, self.n_heads, self.d_head)
        cum_theta = torch.cumsum(theta, dim=1)
        
        half = self.d_head // 2
        v_view = v.view(B, N, self.n_heads, self.d_head)
        v_real, v_imag = v_view[..., :half], v_view[..., half:]
        c, s = torch.cos(cum_theta[..., :half]), torch.sin(cum_theta[..., :half])
        v_rotated = torch.cat([v_real * c - v_imag * s, v_real * s + v_imag * c], dim=-1).reshape(B, N, D)
        
        # 4. Optimized Selective Scan (Chunked)
        # h_t = exp(-dt_t) * h_{t-1} + dt_t * v_t
        # We process in chunks to balance parallelism and memory.
        chunk_size = 64 # Smaller chunk = much less VRAM (T^2 dependency)
        v_in = v_rotated * dt
        log_alpha = -dt
        
        memory_state = torch.zeros_like(v_in)
        h_prev = torch.zeros(B, D, device=x.device, dtype=x.dtype)
        
        for i in range(0, N, chunk_size):
            end = min(i + chunk_size, N)
            v_c = v_in[:, i:end, :]
            la_c = log_alpha[:, i:end, :]
            
            # Local L (relative cumsum)
            L_c = torch.cumsum(la_c, dim=1)
            
            # Stable Intra-chunk computation:
            # (B, T, T, D) matrix but T is now only 64.
            # Memory per layer: 8 * 64 * 64 * 256 * 4 bytes = 32 MB (4x less than v0.3.2)
            rel_L = L_c.unsqueeze(2) - L_c.unsqueeze(1)
            # Masking for causality
            mask = torch.tril(torch.ones(end-i, end-i, device=x.device, dtype=torch.bool)).unsqueeze(0).unsqueeze(-1)
            kernel = torch.exp(rel_L).masked_fill(~mask, 0.0)
            
            chunk_out = torch.einsum('btjd,bjd->btd', kernel, v_c)
            
            if i > 0:
                chunk_out = chunk_out + h_prev.unsqueeze(1) * torch.exp(L_c)
            
            memory_state[:, i:end, :] = chunk_out
            h_prev = chunk_out[:, -1, :]

        # 5. Final Assembly
        out = (memory_state * gate)
        out = self.to_output(out)
        out = self.dropout(out)
        
        return residual + out


# ─────────────────────────────────────────────────────────────
# 2. Cross-Stratum Routing (CSR)
# ─────────────────────────────────────────────────────────────

class CrossStratumRouting(nn.Module):
    def __init__(self, d_model: int, stride: int):
        super().__init__()
        self.stride = stride
        kernel = stride * 2
        self.pad_up = nn.ConstantPad1d((kernel - 1, 0), 0.0)
        self.conv_up = nn.Conv1d(d_model, d_model, kernel_size=kernel, stride=stride, padding=0)
        self.conv_down = nn.ConvTranspose1d(d_model, d_model, kernel_size=kernel, stride=stride, padding=0)
        self.gate = nn.Linear(d_model, d_model)

    def forward(self, h_fine: torch.Tensor, h_coarse: torch.Tensor):
        Lf, Lc = h_fine.shape[1], h_coarse.shape[1]
        x_f = h_fine.transpose(1, 2)
        up = self.conv_up(self.pad_up(x_f)).transpose(1, 2)[:, :Lc, :]
        h_coarse = h_coarse + up * torch.sigmoid(self.gate(up))
        x_c = h_coarse.transpose(1, 2)
        down = self.conv_down(x_c).transpose(1, 2)
        down = F.pad(down, (0, 0, self.stride, 0))[:, :Lf, :] 
        h_fine = h_fine + down * torch.sigmoid(self.gate(down))
        return h_fine, h_coarse


class LocalMixing(nn.Module):
    def __init__(self, d_model: int, kernel_size: int = 5):
        super().__init__()
        self.norm = RMSNorm(d_model)
        self.pad = nn.ConstantPad1d((kernel_size - 1, 0), 0.0)
        self.conv = nn.Conv1d(d_model, d_model, kernel_size=kernel_size, groups=d_model)
        self.pt_conv = nn.Conv1d(d_model, d_model, kernel_size=1)

    def forward(self, x):
        res = x
        x = self.norm(x).transpose(1, 2)
        x = self.pt_conv(F.silu(self.conv(self.pad(x)))).transpose(1, 2)
        return res + x

class SwiGLUFFN(nn.Module):
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.norm = RMSNorm(d_model)
        self.w1 = nn.Linear(d_model, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d_model, bias=False)
        self.w3 = nn.Linear(d_model, d_ff, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        res = x
        x = self.norm(x)
        x = self.w2(F.silu(self.w1(x)) * self.w3(x))
        return res + self.dropout(x)

class LaminarNet(nn.Module):
    def __init__(self, config: LaminarNetConfig):
        super().__init__()
        self.config = config
        d = config.d_model
        self.tok_emb = nn.Embedding(config.vocab_size, d)
        self.pos_emb = nn.Embedding(config.seq_len, d)
        self.dropout = nn.Dropout(config.dropout)
        self.strata_init = nn.ModuleList([nn.Sequential(nn.ConstantPad1d((r*2-1, 0), 0.0), nn.Conv1d(d, d, kernel_size=r*2, stride=r)) for r in config.strata_ratios[1:]])
        self.blocks = nn.ModuleList([LaminarBlock(config) for _ in range(config.n_layers)])
        self.norm_out = RMSNorm(d)
        self.head = nn.Linear(d, config.vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight
        self.apply(self._init_weights)

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

    def forward(self, input_ids):
        B, N = input_ids.shape
        pos = torch.arange(N, device=input_ids.device)
        x = self.dropout(self.tok_emb(input_ids) + self.pos_emb(pos))
        strata = [x] + [layer(x.transpose(1, 2)).transpose(1, 2) for layer in self.strata_init]
        for block in self.blocks: strata = block(strata)
        return self.head(self.norm_out(strata[0]))

    def count_parameters(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

class LaminarBlock(nn.Module):
    def __init__(self, config: LaminarNetConfig):
        super().__init__()
        self.S = config.n_strata
        self.gdfs = nn.ModuleList([GeometricDriftField(config.d_model, config.n_heads, config.dropout) for _ in range(self.S)])
        self.mixers = nn.ModuleList([LocalMixing(config.d_model, config.conv_kernel) for _ in range(self.S)])
        self.csrs = nn.ModuleList([CrossStratumRouting(config.d_model, config.strata_ratios[s+1]//config.strata_ratios[s]) for s in range(self.S-1)])
        self.ffns = nn.ModuleList([SwiGLUFFN(config.d_model, config.d_ff, config.dropout) for _ in range(self.S)])

    def forward(self, strata: List[torch.Tensor]):
        for s in range(self.S): strata[s] = self.gdfs[s](strata[s])
        for s in range(self.S): strata[s] = self.mixers[s](strata[s])
        for s in range(self.S - 1): strata[s], strata[s+1] = self.csrs[s](strata[s], strata[s+1])
        for s in range(self.S): strata[s] = self.ffns[s](strata[s])
        return strata

if __name__ == "__main__":
    conf = LaminarNetConfig()
    model = LaminarNet(conf)
    print(f"LaminarNet v0.3.2 | Params: {model.count_parameters():,}")
