"""
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) — Parallel FFT Convolution
# ─────────────────────────────────────────────────────────────

class GeometricDriftField(nn.Module):
    """
    Geometric Drift Field v2.7 — Fast Global Convolution (FFT).
    
    Replaces the JIT sequential recurrence with an O(N log N) Global Convolution 
    computed via Fast Fourier Transform (FFT). This restores the massive speed 
    advantage of the architecture while maintaining mathematical equivalence 
    to the linear recurrence y_t = decay * y_{t-1} + x_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
        self.decay_log = 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 & Content ──
        gate = torch.sigmoid(self.to_gate(torch.cat([x, local_context], dim=-1)))
        
        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)
        v_in = v * select_gate  # (B, N, H, DH)
        
        # ── 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. Global Convolution via FFT (Memory) ──
        # 1. Prepare Decay Kernel
        # Decay rate lambda = -exp(decay_log) ensuring stability (< 0)
        decay_lambda = -torch.exp(self.decay_log) # (H, DH)
        
        # Time steps [0, 1, ..., N-1]
        t = torch.arange(N, device=x.device, dtype=torch.float32)
        
        # Kernel: k[t] = exp(lambda * t)
        # Shape: (N, H, DH)
        kernel = torch.exp(torch.einsum('n,hd->nhd', t, decay_lambda))
        
        # 2. Prepare Input
        # (B, N, H, DH) -> (B, H, DH, N) for FFT
        # Perform FFT in float32 for stability
        v_fft_in = v_in.permute(0, 2, 3, 1).to(torch.float32) 
        k_fft_in = kernel.permute(1, 2, 0).to(torch.float32)  # (H, DH, N)
        
        # 3. FFT Convolution
        # Pad to 2N to avoid circular aliasing (causal conv)
        n_fft = 2 * N
        
        v_f = torch.fft.rfft(v_fft_in, n=n_fft)
        k_f = torch.fft.rfft(k_fft_in, n=n_fft)
        
        # Element-wise multiplication in frequency domain
        # Broadcast kernel over batch dimension
        out_f = v_f * k_f.unsqueeze(0)
        
        # Inverse FFT
        out_time = torch.fft.irfft(out_f, n=n_fft)
        
        # Slice to original length and restore shape
        # (B, H, DH, 2N) -> (B, H, DH, N) -> (B, N, H, DH)
        memory_state = out_time[..., :N].permute(0, 3, 1, 2).to(x.dtype)
        
        # ── 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.")