"""
LaminarNet — Structured Orthogonal State Space Sequence Model
A novel neural architecture for sequence modeling merging SSM, RNN, and Hierarchical processing.

Core mechanisms:
  1. Geometric Drift Fields (GDF) — Parallelizable rotation-based state evolution (via Prefix Sum)
  2. Cross-Stratum Routing (CSR) — Bidirectional multi-resolution exchange
  3. Phase Mesh Encoding (PME) — Stabilized coupled oscillator position encoding
"""

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          # number of resolution strata
    strata_ratios: tuple = (1, 4)  # compression ratios
    n_oscillators: int = 16    # phase mesh oscillators
    seq_len: int = 2048        # Extended context length capability
    dropout: float = 0.1
    conv_kernel: int = 5       # base local mixing kernel


# ─────────────────────────────────────────────────────────────
# 1. Geometric Drift Field (GDF) — JIT Recurrence (Stable)
# ─────────────────────────────────────────────────────────────

@torch.jit.script
def fused_recurrence(x: torch.Tensor, decay: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
    """
    Stable linear recurrence: h_t = decay * h_{t-1} + gate * x_t
    x: (B, N, H, D)
    decay: (H, D)
    gate: (B, N, H, D)
    """
    B, N, H, D = x.shape
    h = torch.zeros(B, H, D, device=x.device, dtype=x.dtype)
    out = torch.zeros_like(x)
    
    # Iterate over sequence length (dim 1)
    # In JIT this is much faster than Python loop
    for t in range(N):
        x_t = x[:, t, :, :]
        g_t = gate[:, t, :, :]
        h = h * decay + x_t * g_t
        out[:, t, :, :] = h
        
    return out

class GeometricDriftField(nn.Module):
    """
    Geometric Drift Field v2.6 — JIT-Stable Recurrence.
    
    Replaces the unstable parallel scan (v2.5) with a JIT-compiled sequential
    recurrence. This ensures numerical stability (no NaNs) while maintaining
    reasonable speed on GPUs.
    
    Mechanism:
      h_t = decay * h_{t-1} + gate * x_t
      y_t = Rot(x_t) + h_t
    """

    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.to_theta = nn.Linear(d_model, d_model)
        self.to_input = nn.Linear(d_model, d_model)
        self.context_conv = nn.Conv1d(d_model, d_model, kernel_size=5, 
                                      padding=4, groups=d_model)
        
        self.to_gate = nn.Linear(d_model * 2, d_model)
        self.to_select = nn.Linear(d_model, d_model)
        
        # Log-space decay parameter for stability
        # initialized to approx 0.9 range (sigmoid(-2) approx 0.1 -> decay 0.9)
        self.decay_param = nn.Parameter(torch.randn(n_heads, self.d_head))
        
        self.to_output = nn.Linear(d_model, d_model)
        self.norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        residual = x
        x = self.norm(x)
        
        B, N, D = x.shape
        H = self.n_heads
        D_H = self.d_head
        
        # ── A. Stable Local Context ──
        x_feat = x.transpose(1, 2)
        local_context = self.context_conv(x_feat)[..., :N].transpose(1, 2)
        
        # ── B. Gates & Values ──
        gate = torch.sigmoid(self.to_gate(torch.cat([x, local_context], dim=-1)))
        
        # Content to store
        v = self.to_input(x).view(B, N, H, D_H)
        select_gate = torch.sigmoid(self.to_select(x)).view(B, N, H, D_H)
        
        # ── C. Rotation (Positional) ──
        theta = torch.tanh(self.to_theta(x)) * 0.1
        theta = theta.view(B, N, H, D_H)
        cum_theta = torch.cumsum(theta, dim=1)
        
        half = D_H // 2
        v_real, v_imag = v[..., :half], v[..., half:]
        cos_rot = torch.cos(cum_theta[..., :half])
        sin_rot = torch.sin(cum_theta[..., :half])
        
        out_real = v_real * cos_rot - v_imag * sin_rot
        out_imag = v_real * sin_rot + v_imag * cos_rot
        v_rotated = torch.cat([out_real, out_imag], dim=-1)
        
        # ── D. JIT Linear Recurrence (Memory) ──
        # Decay: Sigmoid to force (0, 1) range for stability
        decay = torch.sigmoid(self.decay_param) # (H, D_H)
        
        # Run fused scan
        memory_state = fused_recurrence(v, decay, select_gate)
        
        # ── E. Output ──
        combined = v_rotated + memory_state
        v_gated = combined * gate.view(B, N, H, D_H)
        
        out = v_gated.reshape(B, N, D)
        out = self.to_output(out)
        out = self.dropout(out)
        
        return residual + out


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

class CrossStratumRouting(nn.Module):
    """
    Bidirectional information exchange between adjacent strata with robust padding.
    """

    def __init__(self, d_model: int, stride: int):
        super().__init__()
        self.stride = stride
        kernel = stride * 2
        
        # Promotion (Fine -> Coarse)
        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.gate_up = nn.Linear(d_model, d_model)
        
        # Demotion (Coarse -> Fine)
        self.conv_down = nn.ConvTranspose1d(d_model, d_model, kernel_size=kernel, 
                                            stride=stride, padding=0)
        self.gate_down = nn.Linear(d_model, d_model)

    def forward(self, h_fine: torch.Tensor, h_coarse: torch.Tensor):
        L_fine = h_fine.shape[1]
        L_coarse = h_coarse.shape[1]
        
        # ── PROMOTE ──
        x_fine = (h_fine * torch.sigmoid(self.gate_up(h_fine))).transpose(1, 2)
        x_padded = self.pad_up(x_fine)
        promote = self.conv_up(x_padded).transpose(1, 2)
        
        if promote.shape[1] > L_coarse:
            promote = promote[:, :L_coarse, :]
        elif promote.shape[1] < L_coarse:
            pad_len = L_coarse - promote.shape[1]
            promote = F.pad(promote, (0, 0, 0, pad_len))
            
        h_coarse = h_coarse + promote
        
        # ── DEMOTE ──
        x_coarse = (h_coarse * torch.sigmoid(self.gate_down(h_coarse))).transpose(1, 2)
        demote = self.conv_down(x_coarse).transpose(1, 2)
        demote = F.pad(demote, (0, 0, self.stride, 0)) # Causal shift
        
        if demote.shape[1] > L_fine:
            demote = demote[:, :L_fine, :]
        elif demote.shape[1] < L_fine:
            pad_len = L_fine - demote.shape[1]
            demote = F.pad(demote, (0, 0, 0, pad_len))
            
        h_fine = h_fine + demote
        
        return h_fine, h_coarse


# ─────────────────────────────────────────────────────────────
# 3. Phase Mesh Encoding (PME) — Stabilized
# ─────────────────────────────────────────────────────────────

class PhaseMeshEncoding(nn.Module):
    """
    Stabilized Positional Encoding.
    """

    def __init__(self, d_model: int, n_oscillators: int = 16, seq_len: int = 2048):
        super().__init__()
        self.d_model = d_model
        self.emb = nn.Embedding(seq_len, d_model)

    def forward(self, input_len: int, device: torch.device):
        positions = torch.arange(input_len, device=device)
        return self.emb(positions)


# ─────────────────────────────────────────────────────────────
# 4. Local Mixing & FFN
# ─────────────────────────────────────────────────────────────

class LocalMixing(nn.Module):
    """Depthwise Separable Causal Conv1d"""
    def __init__(self, d_model: int, kernel_size: int = 5):
        super().__init__()
        self.norm = nn.LayerNorm(d_model)
        self.conv = nn.Conv1d(d_model, d_model, kernel_size=kernel_size, 
                              groups=d_model, padding=0)
        self.pad = nn.ConstantPad1d((kernel_size - 1, 0), 0.0)
        self.pt_conv = nn.Conv1d(d_model, d_model, kernel_size=1)
        self.act = nn.SiLU()

    def forward(self, x):
        residual = x
        x = self.norm(x).transpose(1, 2)
        x = self.pad(x)
        x = self.conv(x)
        x = self.act(x)
        x = self.pt_conv(x).transpose(1, 2)
        return residual + x


class SwiGLUFFN(nn.Module):
    """Gated FFN"""
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.norm = nn.LayerNorm(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):
        residual = x
        x = self.norm(x)
        x = F.silu(self.w1(x)) * self.w3(x)
        x = self.w2(x)
        x = self.dropout(x)
        return residual + x


# ─────────────────────────────────────────────────────────────
# 5. Laminar Block & Main Model
# ─────────────────────────────────────────────────────────────

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, kernel_size=config.conv_kernel * config.strata_ratios[s])
            for s in range(self.S)
        ])
        
        self.csrs = nn.ModuleList()
        for s in range(self.S - 1):
            stride = config.strata_ratios[s+1] // config.strata_ratios[s]
            self.csrs.append(CrossStratumRouting(config.d_model, stride))
            
        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


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.phase_mesh = PhaseMeshEncoding(d, config.n_oscillators, config.seq_len)
        self.dropout = nn.Dropout(config.dropout)
        
        self.strata_init = nn.ModuleList()
        for s in range(1, config.n_strata):
            r = config.strata_ratios[s]
            self.strata_init.append(nn.Sequential(
                nn.ConstantPad1d((r*2 - 1, 0), 0.0),
                nn.Conv1d(d, d, kernel_size=r*2, stride=r, padding=0)
            ))
            
        self.blocks = nn.ModuleList([
            LaminarBlock(config) for _ in range(config.n_layers)
        ])
        
        self.norm_out = nn.LayerNorm(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, module):
        if isinstance(module, nn.Linear):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
        elif isinstance(module, nn.Conv1d):
            torch.nn.init.kaiming_normal_(module.weight)

    def forward(self, input_ids):
        B, N = input_ids.shape
        device = input_ids.device
        
        x = self.tok_emb(input_ids)
        pos = self.phase_mesh(N, device)
        x = self.dropout(x + pos)
        
        strata = [x]
        for init_layer in self.strata_init:
            feat = x.transpose(1, 2)
            down = init_layer(feat).transpose(1, 2)
            strata.append(down)
            
        for block in self.blocks:
            strata = block(strata)
            
        out = self.norm_out(strata[0])
        logits = self.head(out)
        return logits

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


if __name__ == "__main__":
    conf = LaminarNetConfig()
    model = LaminarNet(conf)
    print(f"Model Params: {model.count_parameters():,}")
    
    x = torch.randint(0, conf.vocab_size, (2, 256))
    y = model(x)
    print(f"Input: {x.shape} -> Output: {y.shape}")
    assert y.shape == (2, 256, conf.vocab_size)
    print("Test Passed: Parallel LaminarNet is operational.")