"""
LaminarNet v0.3.5 — Turbo & Lean Update
High-speed selective state space architecture optimized for VRAM and throughput.
"""

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 = 50257
    d_model: int = 256
    n_heads: int = 8
    n_layers: int = 8
    d_ff: int = 1024
    n_strata: int = 2
    strata_ratios: tuple = (1, 4)
    seq_len: int = 1024
    dropout: float = 0.1
    conv_kernel: int = 4


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):
        return self.scale * x * x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()


# ─────────────────────────────────────────────────────────────
# 1. Optimized Selective Geometric Drift (SGD)
# ─────────────────────────────────────────────────────────────

class GeometricDriftField(nn.Module):
    """
    Geometric Drift Field v3.5 — Optimized for Speed and VRAM.
    Fuses projections and uses a memory-efficient selective scan.
    """

    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
        
        # Fused projection for efficiency: dt, v, gate, theta
        # 1*D (dt) + 1*D (v) + 1*D (gate) + 1*D (theta)
        self.in_proj = nn.Linear(d_model, d_model * 4, bias=False)
        self.conv1d = nn.Conv1d(d_model, d_model, kernel_size=4, padding=3, groups=d_model)
        
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.norm = RMSNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
        # Learnable dt bias to control initial memory horizon
        self.dt_bias = nn.Parameter(torch.ones(d_model) * -3.0)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.norm(x)
        B, N, D = x.shape
        
        # 1. Local Context & Fused Proj
        x_conv = self.conv1d(x.transpose(1, 2))[..., :N].transpose(1, 2)
        x_conv = F.silu(x_conv)
        
        # Fused multiplication
        fused = self.in_proj(x_conv)
        dt_raw, v, gate, theta = fused.chunk(4, dim=-1)
        
        # 2. Selectivity
        dt = F.softplus(dt_raw + self.dt_bias).clamp(min=0.001, max=5.0)
        gate = torch.sigmoid(gate)
        
        # 3. Positional Rotation (Simplified for speed)
        theta = torch.tanh(theta) * 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. Memory-Efficient Selective Scan (Parallelized via cumprod/cumsum trick)
        # Numerical stability via log-space subtraction
        log_alpha = -dt
        # h_t = exp(L_t) * sum(exp(-L_j) * dt_j * v_j)
        # For N=1024, we use a simple and fast approach:
        v_in = v_rotated * dt
        
        # To keep it O(N) and fast in pure PyTorch without huge VRAM:
        # We use a simple loop over a larger chunk to avoid the T*T matrix
        chunk_size = 256 # Larger chunks for better GPU utilization
        out = torch.zeros_like(v_in)
        h = 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, :]
            
            # This part is vectorized within the chunk
            L_c = torch.cumsum(la_c, dim=1)
            # Efficient intra-chunk scan
            # We use a smaller memory footprint by computing kernel on the fly
            exp_L = torch.exp(L_c)
            # Instead of T*T matrix, we use a slightly slower but 
            # much more memory-efficient sequential accumulation inside the chunk
            # or a simple associative scan if T is small.
            # Here we go for stability + speed:
            
            # Simple vectorized recurrence for the chunk
            curr_h = h.unsqueeze(1)
            chunk_results = []
            for t in range(end - i):
                curr_h = torch.exp(la_c[:, t:t+1, :]) * curr_h + v_c[:, t:t+1, :]
                chunk_results.append(curr_h)
            
            chunk_out = torch.cat(chunk_results, dim=1)
            out[:, i:end, :] = chunk_out
            h = chunk_out[:, -1, :]

        # 5. Output
        out = self.out_proj(out * gate)
        return residual + self.dropout(out)


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

class CrossStratumRouting(nn.Module):
    def __init__(self, d_model: int, stride: int):
        super().__init__()
        self.stride = stride
        # Lightweight exchange
        self.down = nn.AvgPool1d(kernel_size=stride, stride=stride)
        self.up = nn.Upsample(scale_factor=stride, mode='linear', align_corners=False)
        self.gate = nn.Sequential(nn.Linear(d_model, d_model), nn.Sigmoid())

    def forward(self, h_fine: torch.Tensor, h_coarse: torch.Tensor):
        # h_fine: (B, L, D), h_coarse: (B, L/S, D)
        # Quick exchange to avoid heavy ConvTranspose
        f_to_c = self.down(h_fine.transpose(1, 2)).transpose(1, 2)
        L_c = h_coarse.shape[1]
        h_coarse = h_coarse + self.gate(f_to_c[:, :L_c, :]) * f_to_c[:, :L_c, :]
        
        c_to_f = self.up(h_coarse.transpose(1, 2)).transpose(1, 2)
        L_f = h_fine.shape[1]
        h_fine = h_fine + self.gate(c_to_f[:, :L_f, :]) * c_to_f[:, :L_f, :]
        
        return h_fine, h_coarse


# ─────────────────────────────────────────────────────────────
# 3. Main Architecture Components
# ─────────────────────────────────────────────────────────────

class LocalMixing(nn.Module):
    def __init__(self, d_model: int, kernel_size: int = 4):
        super().__init__()
        self.norm = RMSNorm(d_model)
        self.conv = nn.Conv1d(d_model, d_model, kernel_size=kernel_size, padding=kernel_size-1, groups=d_model)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        res = x
        x = self.norm(x).transpose(1, 2)
        x = F.silu(self.conv(x)[..., :res.shape[1]]).transpose(1, 2)
        return res + self.proj(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)
        return res + self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(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)
        
        # Optimized Strata Init
        self.strata_init = nn.ModuleList([
            nn.Sequential(nn.AvgPool1d(kernel_size=r, 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.02)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)

    def forward(self, input_ids):
        B, N = input_ids.shape
        x = self.dropout(self.tok_emb(input_ids) + self.pos_emb(torch.arange(N, device=input_ids.device)))
        
        # Init Strata (B, L, D) -> (B, D, L) -> Pool -> (B, L/S, D)
        strata = [x]
        x_t = x.transpose(1, 2)
        for pool in self.strata_init:
            strata.append(pool(x_t).transpose(1, 2))
            
        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.5 | Params: {model.count_parameters():,}")