"""
lilm1-200m model definition — lilm hybrid local-convolution/attention model
═════════════════════════════════════════════════════════════════════════════

Architecture:
  - 10 local causal-convolution layers + 6 attention layers (16 total)
  - GQA with n_head=16, kv_heads=4  (4:1 ratio)
  - SwiGLU FFN (intermediate_size=2752)
  - RMSNorm (learnable, eps=1e-6)
  - RoPE positional encoding
  - Attention via PyTorch SDPA (FlashAttention/math backend auto-dispatch)
  - Residual dropout on all residual paths
  - Untied embedding / lm_head weights
  - Causal depthwise 1D convolution for local token mixing

Paste this entire cell into the "PASTE YOUR MODEL DEFINITION BELOW" cell.
═════════════════════════════════════════════════════════════════════════════
"""

from dataclasses import dataclass
from typing import Optional
import math

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


# ── Config ────────────────────────────────────────────────────────────────────

@dataclass
class LilmConfig:
    vocab_size:        int   = 49152
    n_layer:           int   = 16          # 10 Conv + 6 Attention
    n_embd:            int   = 1024
    n_head:            int   = 16
    kv_heads:          int   = 4           # GQA 4:1
    intermediate_size: int   = 2752        # SwiGLU hidden dim
    kernel_size:       int   = 4           # Conv window
    attn_dropout_p:    float = 0.1         # Passed to PyTorch SDPA
    resid_dropout_p:   float = 0.05        # nn.Dropout on residual paths
    n_conv_layers:     int   = 10          # First 10 layers are Conv
    max_seq_len:       int   = 2048        # For RoPE precomputation
    rope_theta:        float = 10_000.0    # RoPE base frequency
    use_native_gqa:     bool  = True        # Use SDPA's native GQA when available
    activation_checkpointing: bool = False  # Recompute layer activations to save memory

    @property
    def head_dim(self) -> int:
        return self.n_embd // self.n_head

    @property
    def n_attn_layers(self) -> int:
        return self.n_layer - self.n_conv_layers


# ── RMSNorm ───────────────────────────────────────────────────────────────────

class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization (Zhang & Sennrich, 2019)."""

    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * norm).type_as(x) * self.weight


# ── RoPE ──────────────────────────────────────────────────────────────────────

def precompute_rope_freqs(
    head_dim: int,
    max_seq_len: int,
    theta: float = 10_000.0,
    device: Optional[torch.device] = None,
) -> torch.Tensor:
    """Precompute complex exponentials for RoPE: shape (max_seq_len, head_dim//2)."""
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(max_seq_len, device=device).float()
    freqs = torch.outer(t, freqs)                       # (T, head_dim//2)
    return torch.polar(torch.ones_like(freqs), freqs)    # complex64


def apply_rope(
    x: torch.Tensor,                   # (B, T, n_head, head_dim)
    freqs_cis: torch.Tensor,           # (T, head_dim//2)
) -> torch.Tensor:
    """Apply rotary positional embeddings to q or k."""
    B, T, H, D = x.shape
    # View as pairs of (real, imag)
    x_complex = torch.view_as_complex(x.float().reshape(B, T, H, D // 2, 2))
    freqs = freqs_cis[:T].unsqueeze(0).unsqueeze(2)     # (1, T, 1, head_dim//2)
    x_rotated = x_complex * freqs
    return torch.view_as_real(x_rotated).reshape(B, T, H, D).type_as(x)


# ── SwiGLU FFN ────────────────────────────────────────────────────────────────

class SwiGLUFFN(nn.Module):
    """
    SwiGLU Feed-Forward Network (Shazeer 2020, Touvron et al. 2023).
    gate = Swish(x W_gate)
    out  = (gate ⊙ (x W_up)) W_down
    """

    def __init__(self, config: LilmConfig):
        super().__init__()
        self.w_gate = nn.Linear(config.n_embd, config.intermediate_size, bias=False)
        self.w_up   = nn.Linear(config.n_embd, config.intermediate_size, bias=False)
        self.w_down = nn.Linear(config.intermediate_size, config.n_embd, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


# ── GQA Attention ────────────────────────────────────────────────────────────

class LilmAttention(nn.Module):
    """
    Grouped Query Attention using PyTorch scaled_dot_product_attention.

    n_head=16 query heads, kv_heads=4 key/value heads → 4:1 GQA ratio.
    KV heads are expanded explicitly for broad PyTorch compatibility; SDPA
    still dispatches to the best available CUDA backend for the resulting
    dense-head causal attention.
    """

    def __init__(self, config: LilmConfig):
        super().__init__()
        self.n_head   = config.n_head
        self.kv_heads = config.kv_heads
        self.head_dim = config.head_dim
        self.n_rep    = config.n_head // config.kv_heads   # 4 — repeat factor for GQA
        self.attn_dropout_p = config.attn_dropout_p
        self.use_native_gqa = (
            config.use_native_gqa
            and "enable_gqa" in (F.scaled_dot_product_attention.__doc__ or "")
        )

        # Projections
        self.q_proj = nn.Linear(config.n_embd, config.n_head   * config.head_dim, bias=False)
        self.k_proj = nn.Linear(config.n_embd, config.kv_heads * config.head_dim, bias=False)
        self.v_proj = nn.Linear(config.n_embd, config.kv_heads * config.head_dim, bias=False)
        self.o_proj = nn.Linear(config.n_head * config.head_dim, config.n_embd,   bias=False)

    def _expand_kv(self, x: torch.Tensor) -> torch.Tensor:
        """Repeat KV heads to match query head count: (B,T,kv_heads,D) → (B,T,n_head,D)."""
        if self.n_rep == 1:
            return x
        B, T, H, D = x.shape
        return x[:, :, :, None, :].expand(B, T, H, self.n_rep, D).reshape(B, T, H * self.n_rep, D)

    def forward(
        self,
        x: torch.Tensor,                       # (B, T, n_embd)
        freqs_cis: torch.Tensor,                # (max_seq_len, head_dim//2)
    ) -> torch.Tensor:
        B, T, _ = x.shape

        # Project → (B, T, n_heads, head_dim)
        q = self.q_proj(x).reshape(B, T, self.n_head,   self.head_dim)
        k = self.k_proj(x).reshape(B, T, self.kv_heads, self.head_dim)
        v = self.v_proj(x).reshape(B, T, self.kv_heads, self.head_dim)

        # Apply RoPE to queries and keys
        q = apply_rope(q, freqs_cis)
        k = apply_rope(k, freqs_cis)

        dropout_p = self.attn_dropout_p if self.training else 0.0

        # Use PyTorch SDPA, which auto-dispatches to the best available
        # attention backend on supported hardware.
        # SDPA expects (B, H, T, D) layout
        q = q.transpose(1, 2)
        if self.use_native_gqa:
            k = k.transpose(1, 2)
            v = v.transpose(1, 2)
            out = F.scaled_dot_product_attention(
                q, k, v,
                dropout_p=dropout_p,
                is_causal=True,
                enable_gqa=True,
            )
        else:
            # Compatibility fallback for PyTorch builds without native GQA.
            k = self._expand_kv(k).transpose(1, 2)
            v = self._expand_kv(v).transpose(1, 2)
            out = F.scaled_dot_product_attention(
                q, k, v,
                dropout_p=dropout_p,
                is_causal=True,
            )
        out = out.transpose(1, 2)    # back to (B, T, H, D)

        # Merge heads → project out
        out = out.reshape(B, T, self.n_head * self.head_dim)
        return self.o_proj(out)


# ── Causal Depthwise Conv ────────────────────────────────────────────────────

class CausalDepthwiseConv1d(nn.Module):
    """
    Causal depthwise 1D convolution (Mamba/Hyena style).

    Each channel is convolved independently with its own kernel.
    Left-padded to preserve causality — output at position t depends
    only on inputs at positions [t - kernel_size + 1, ..., t].
    """

    def __init__(self, channels: int, kernel_size: int):
        super().__init__()
        self.kernel_size = kernel_size
        # groups=channels → depthwise (each channel has its own filter)
        self.conv = nn.Conv1d(
            in_channels  = channels,
            out_channels = channels,
            kernel_size  = kernel_size,
            padding      = 0,          # We handle causal padding manually
            groups       = channels,
            bias         = True,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """x: (B, T, C) → (B, T, C)"""
        # Transpose to (B, C, T) for Conv1d
        x = x.transpose(1, 2)
        # Causal left-padding: pad (kernel_size - 1) zeros on the left
        x = F.pad(x, (self.kernel_size - 1, 0))
        x = self.conv(x)
        return x.transpose(1, 2)       # Back to (B, T, C)


# ── Conv Mixer Block (for Conv layers) ───────────────────────────────────────

class ConvMixerBlock(nn.Module):
    """
    Conv-based block (used in layers 0–9):
      x → RMSNorm → CausalDepthwiseConv → SiLU → Linear(project) → residual dropout → +x
      x → RMSNorm → SwiGLU FFN → residual dropout → +x

    The conv branch replaces attention with a local causal mixing mechanism.
    SiLU activation after conv + a linear projection keeps dimensionality matched.
    """

    def __init__(self, config: LilmConfig):
        super().__init__()
        self.norm1 = RMSNorm(config.n_embd)
        self.norm2 = RMSNorm(config.n_embd)

        # Conv branch: depthwise conv → activation → pointwise projection
        self.conv1d    = CausalDepthwiseConv1d(config.n_embd, config.kernel_size)
        self.conv_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)

        # FFN
        self.mlp = SwiGLUFFN(config)

        # Residual dropout
        self.resid_drop = nn.Dropout(p=config.resid_dropout_p)

    def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor:
        # Conv mixing path
        h = self.norm1(x)
        h = self.conv1d(h)
        h = F.silu(h)
        h = self.conv_proj(h)
        x = x + self.resid_drop(h)

        # FFN path
        x = x + self.resid_drop(self.mlp(self.norm2(x)))
        return x


# ── Attention Block (for Attention layers) ───────────────────────────────────

class AttentionBlock(nn.Module):
    """
    Standard transformer block (used in layers 10–15):
      x → RMSNorm → GQA Attention (SDPA) → residual dropout → +x
      x → RMSNorm → SwiGLU FFN → residual dropout → +x
    """

    def __init__(self, config: LilmConfig):
        super().__init__()
        self.norm1 = RMSNorm(config.n_embd)
        self.norm2 = RMSNorm(config.n_embd)

        self.attn = LilmAttention(config)
        self.mlp  = SwiGLUFFN(config)

        self.resid_drop = nn.Dropout(p=config.resid_dropout_p)

    def forward(
        self,
        x: torch.Tensor,
        freqs_cis: torch.Tensor,
    ) -> torch.Tensor:
        x = x + self.resid_drop(self.attn(self.norm1(x), freqs_cis))
        x = x + self.resid_drop(self.mlp(self.norm2(x)))
        return x


# ── Full Model ────────────────────────────────────────────────────────────────

class LilmModel(nn.Module):
    """
    lilm hybrid local-convolution/attention model.

    - Untied embedding + lm_head
    - RoPE applied only in attention layers
    - Pre-norm (RMSNorm) architecture
    - Final RMSNorm before lm_head

    Forward returns logits: (B, T, vocab_size)
    """

    def __init__(self, config: LilmConfig):
        super().__init__()
        self.config = config

        # Token embedding (untied — separate from lm_head)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.n_embd)

        # Hybrid layers: first n_conv_layers are Conv, rest are Attention
        self.layers = nn.ModuleList()
        for i in range(config.n_layer):
            if i < config.n_conv_layers:
                self.layers.append(ConvMixerBlock(config))
            else:
                self.layers.append(AttentionBlock(config))

        # Final norm + output head (untied)
        self.final_norm = RMSNorm(config.n_embd)
        self.lm_head    = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        # Precompute RoPE frequencies (registered as buffer — not a parameter)
        freqs_cis = precompute_rope_freqs(
            config.head_dim,
            config.max_seq_len,
            config.rope_theta,
        )
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)

        # Initialise weights
        self._init_weights()

        # Report parameter count
        n_params = sum(p.numel() for p in self.parameters())
        print(f"LilmModel — {n_params / 1e6:.2f}M parameters")
        print(f"  Conv layers : {config.n_conv_layers} (layers 0–{config.n_conv_layers - 1})")
        print(f"  Attn layers : {config.n_attn_layers} (layers {config.n_conv_layers}–{config.n_layer - 1})")
        print(f"  GQA ratio   : {config.n_head}:{config.kv_heads}")
        print(f"  Head dim    : {config.head_dim}")

    def _init_weights(self):
        """
        Weight initialisation following GPT-NeoX / LLaMA conventions:
        - Embeddings:  N(0, 0.02)
        - Linear:      N(0, 0.02) with residual scaling on output projections
        - Conv1d:      N(0, 0.02)
        - RMSNorm:     ones (already default)
        - Biases:      zeros
        """
        residual_scale = 1.0 / math.sqrt(2.0 * self.config.n_layer)

        for name, module in self.named_modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                # Scale down residual-path output projections to stabilise deep nets
                if any(tag in name for tag in ("o_proj", "w_down", "conv_proj")):
                    module.weight.data.mul_(residual_scale)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
            elif isinstance(module, nn.Conv1d):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """
        Args:
            input_ids: (B, T) token indices
        Returns:
            logits: (B, T, vocab_size)
        """
        B, T = input_ids.shape
        assert T <= self.config.max_seq_len, (
            f"Sequence length {T} exceeds max_seq_len {self.config.max_seq_len}"
        )

        # Embed tokens
        x = self.embed_tokens(input_ids)        # (B, T, n_embd)

        # Pass through all layers
        if self.config.activation_checkpointing and self.training:
            from torch.utils.checkpoint import checkpoint

            for layer in self.layers:
                if isinstance(layer, AttentionBlock):
                    x = checkpoint(
                        lambda hidden, layer=layer: layer(hidden, freqs_cis=self.freqs_cis),
                        x,
                        use_reentrant=False,
                    )
                else:
                    x = checkpoint(layer, x, use_reentrant=False)
        else:
            for layer in self.layers:
                if isinstance(layer, AttentionBlock):
                    x = layer(x, freqs_cis=self.freqs_cis)
                else:
                    x = layer(x)

        # Final norm + logits
        x = self.final_norm(x)
        logits = self.lm_head(x)                # (B, T, vocab_size)

        return logits

    def gradient_checkpointing_enable(self):
        """Enable gradient checkpointing for memory-efficient training."""
        self.config.activation_checkpointing = True


# Deprecated compatibility aliases for older notebooks/checkpoints that imported
# previous class names. New code should use LilmConfig and LilmModel.
LFMUpgradedConfig = LilmConfig
LiquidAttention = LilmAttention
LiquidUpgradedModel = LilmModel