"""
LaminarNet v0.5.2 — Numerical Stability Fix, Fully Causal, O(N)
Faster than Transformer: larger chunks, streamlined architecture, vectorized carry.
All temporal operations are strictly causal — no future information leakage.
"""

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


# ─────────────────────────────────────────────────────────────
# 1. O(N) Selective Geometric Drift Field — Speed Optimized
# ─────────────────────────────────────────────────────────────

class GeometricDriftField(nn.Module):
    """
    Geometric Drift Field v5.0 — O(N) Vectorized Parallel Scan.
    chunk_size=256 with vectorized carry for maximum GPU throughput.
    """

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1, conv_kernel: int = 4):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        
        # Fused Projections
        self.in_proj = nn.Linear(d_model, d_model * 4, bias=False)
        self.conv1d = nn.Conv1d(d_model, d_model, kernel_size=conv_kernel,
                                padding=conv_kernel-1, groups=d_model)
        
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.norm = RMSNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
        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. Context & Proj (causal conv)
        x_conv = self.conv1d(x.transpose(1, 2))[..., :N].transpose(1, 2)
        x_conv = F.silu(x_conv)
        
        fused = self.in_proj(x_conv)
        dt_raw, v, gate, theta = fused.chunk(4, dim=-1)
        
        # 2. Selective Parameters
        dt = F.softplus(dt_raw + self.dt_bias).clamp(min=0.001, max=2.0)
        gate = torch.sigmoid(gate)
        log_alpha = -dt
        
        # 3. Positional Rotation
        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. O(N) Vectorized Parallel Scan — chunk_size=256 for speed
        chunk_size = 256
        v_in = v_rotated * dt
        
        # Pad to nearest chunk multiple
        orig_N = N
        remainder = N % chunk_size
        if remainder != 0:
            pad_len = chunk_size - remainder
            v_in = F.pad(v_in, (0, 0, 0, pad_len))
            log_alpha = F.pad(log_alpha, (0, 0, 0, pad_len))
        
        N_padded = v_in.shape[1]
        num_chunks = N_padded // chunk_size
        
        v_chunks = v_in.view(B, num_chunks, chunk_size, D)
        la_chunks = log_alpha.view(B, num_chunks, chunk_size, D)
        
        # Cumulative log-decay within each chunk
        L_chunks = torch.cumsum(la_chunks, dim=2)  # (B, C, T, D)
        
        # O(N) intra-chunk scan via cumsum (float32 for precision + stability)
        L_f32 = L_chunks.float().clamp(min=-20.0, max=0.0)
        # Stabilized: subtract per-chunk max to prevent exp overflow
        L_max = L_f32.max(dim=2, keepdim=True).values
        L_stable = L_f32 - L_max
        exp_neg_L_stable = torch.exp(-L_stable)
        scaled_v = exp_neg_L_stable * v_chunks.float()
        cum_scaled = torch.cumsum(scaled_v, dim=2)
        exp_L_stable = torch.exp(L_stable)
        chunk_out = (exp_L_stable * cum_scaled).to(x.dtype)  # (B, C, T, D)
        
        # Vectorized inter-chunk carry (compute carries, apply ALL at once)
        if num_chunks > 1:
            # Compute carry states: h_carry[c] = h_carry[c-1]*decay[c-1] + out[c-1,-1]
            carry_list = [torch.zeros(B, D, device=x.device, dtype=x.dtype)]
            for c in range(1, num_chunks):
                # out-of-place accumulation for autograd stability
                next_carry = carry_list[-1] * torch.exp(L_chunks[:, c-1, -1, :]) + chunk_out[:, c-1, -1, :]
                carry_list.append(next_carry)
            carries = torch.stack(carry_list, dim=1)  # (B, C, D)
            # Apply all carries in parallel (single tensor op for ALL positions)
            final_out = chunk_out + carries.unsqueeze(2) * torch.exp(L_chunks)
            final_out = final_out.view(B, -1, D)
        else:
            final_out = chunk_out.view(B, -1, D)
        
        final_out = final_out[:, :orig_N, :]
        
        # 5. Output
        out = self.out_proj(final_out * gate)
        return residual + self.dropout(out)


# ─────────────────────────────────────────────────────────────
# 2. Standard Infrastructure
# ─────────────────────────────────────────────────────────────

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

    def forward(self, h_fine: torch.Tensor, h_coarse: torch.Tensor):
        fine_t = h_fine.transpose(1, 2)
        k = self.down.kernel_size if isinstance(self.down.kernel_size, int) else self.down.kernel_size[0]
        fine_t = F.pad(fine_t, (k - 1, 0))
        f_to_c = self.down(fine_t).transpose(1, 2)
        Lc = h_coarse.shape[1]
        if f_to_c.shape[1] < Lc:
            f_to_c = F.pad(f_to_c, (0, 0, 0, Lc - f_to_c.shape[1]))
        h_coarse = h_coarse + self.gate(f_to_c[:, :Lc, :]) * f_to_c[:, :Lc, :]
        c_to_f = self.up(h_coarse.transpose(1, 2)).transpose(1, 2)
        Lf = h_fine.shape[1]
        if c_to_f.shape[1] < Lf:
            c_to_f = F.pad(c_to_f, (0, 0, 0, Lf - c_to_f.shape[1]))
        h_fine = h_fine + self.gate(c_to_f[:, :Lf, :]) * c_to_f[:, :Lf, :]
        return h_fine, h_coarse

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, self.w2, self.w3 = nn.Linear(d_model, d_ff, bias=False), nn.Linear(d_ff, d_model, bias=False), 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, d = config, 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.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(lambda m: nn.init.normal_(m.weight, std=0.02) if isinstance(m, (nn.Linear, nn.Embedding)) else None)

    def forward(self, ids):
        B, N = ids.shape
        pos = torch.arange(N, device=ids.device).clamp(max=self.config.seq_len - 1)
        x = self.dropout(self.tok_emb(ids) + self.pos_emb(pos))
        # CAUSAL strata init: left-pad so each coarse position only sees past/current
        coarse_strata = []
        for pool in self.strata_init:
            x_t = x.transpose(1, 2)
            k = pool.kernel_size[0]
            x_t = F.pad(x_t, (k - 1, 0))
            coarse_strata.append(pool(x_t).transpose(1, 2))
        strata = [x] + coarse_strata
        for b in self.blocks: strata = b(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):
        super().__init__()
        self.S = config.n_strata
        self.gdfs = nn.ModuleList([GeometricDriftField(config.d_model, config.n_heads, config.dropout, 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):
        for s in range(self.S): strata[s] = self.gdfs[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)
    x = torch.randint(0, conf.vocab_size, (2, 128))
    print(f"LaminarNet v0.5.1 | Out: {model(x).shape}")
