"""
LatentRecurrentFlow (LRF) - Core Architecture Modules

Architecture Overview:
=====================
The LRF architecture consists of 4 main components:

1. CompactEncoder/Decoder (VAE): f=32 spatial compression with tiny decoder
2. TextConditioner: Lightweight text encoding (TinyCLIP or small LM)
3. RecursiveLatentCore: The novel HRM-inspired denoising backbone
4. FlowScheduler: Rectified flow for training and sampling

The RecursiveLatentCore is the key innovation:
- It contains N_blocks GLD (Gated Linear Diffusion) blocks
- These blocks are applied recursively T_outer * T_inner times
- The same parameters are reused across recursions (weight sharing)
- Training uses IFT (Implicit Function Theorem) for O(1) memory backprop
- This gives effective depth of T_outer * T_inner * N_blocks layers
  from only N_blocks parameter sets

Memory budget at inference (1024x1024, INT8):
- Text encoder: ~150MB (TinyCLIP-ViT-B/16)
- VAE encoder: ~100MB (f32 encoder, only needed for editing)
- VAE decoder: ~6MB (SnapGen-style tiny decoder)
- LRF core: ~200-400MB (depending on config)
- Activations: ~500MB peak
- Total: ~1-1.5GB model + ~500MB activations = 1.5-2GB
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
from typing import Optional, Tuple, Dict, Any


# ============================================================================
# Utility Modules
# ============================================================================

class RMSNorm(nn.Module):
    """RMSNorm - more stable than LayerNorm for small models."""
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    
    def forward(self, x):
        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * norm).type_as(x) * self.weight


class SwiGLU(nn.Module):
    """SwiGLU FFN - better than GELU for small models, mobile-friendly (SiLU not GELU)."""
    def __init__(self, dim: int, hidden_dim: Optional[int] = None, dropout: float = 0.0):
        super().__init__()
        hidden_dim = hidden_dim or int(dim * 8 / 3)
        # Round to nearest multiple of 8 for efficiency
        hidden_dim = ((hidden_dim + 7) // 8) * 8
        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x):
        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))


class DepthwiseSeparableConv2d(nn.Module):
    """Mobile-optimized convolution."""
    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3):
        super().__init__()
        padding = kernel_size // 2
        self.dw = nn.Conv2d(in_channels, in_channels, kernel_size, padding=padding, groups=in_channels, bias=False)
        self.pw = nn.Conv2d(in_channels, out_channels, 1, bias=False)
    
    def forward(self, x):
        return self.pw(self.dw(x))


# ============================================================================
# 2D Positional Encoding
# ============================================================================

class RotaryPositionEncoding2D(nn.Module):
    """2D RoPE for spatial tokens - resolution-independent."""
    def __init__(self, dim: int, max_res: int = 64):
        super().__init__()
        self.dim = dim
        half_dim = dim // 4  # Split into 4 parts: sin_h, cos_h, sin_w, cos_w
        freqs = torch.exp(torch.arange(half_dim) * -(math.log(10000.0) / half_dim))
        self.register_buffer('freqs', freqs)
    
    def forward(self, h: int, w: int, device=None):
        device = device or self.freqs.device
        pos_h = torch.arange(h, device=device).float()
        pos_w = torch.arange(w, device=device).float()
        
        freqs_h = torch.outer(pos_h, self.freqs.to(device))  # [H, D/4]
        freqs_w = torch.outer(pos_w, self.freqs.to(device))  # [W, D/4]
        
        # Expand to [H, W, D/4] each
        freqs_h = freqs_h.unsqueeze(1).expand(-1, w, -1)
        freqs_w = freqs_w.unsqueeze(0).expand(h, -1, -1)
        
        # Concatenate: [H, W, D/2] for sin, [H, W, D/2] for cos
        freqs = torch.cat([freqs_h, freqs_w], dim=-1)  # [H, W, D/2]
        
        sin_enc = freqs.sin()
        cos_enc = freqs.cos()
        
        return sin_enc.reshape(h * w, -1), cos_enc.reshape(h * w, -1)


def apply_rope_2d(x, sin_enc, cos_enc):
    """Apply 2D RoPE to queries/keys."""
    d = x.shape[-1]
    half_d = d // 2
    x1, x2 = x[..., :half_d], x[..., half_d:]
    # Expand sin/cos to match batch dims
    while sin_enc.dim() < x1.dim():
        sin_enc = sin_enc.unsqueeze(0)
        cos_enc = cos_enc.unsqueeze(0)
    return torch.cat([x1 * cos_enc - x2 * sin_enc, x2 * cos_enc + x1 * sin_enc], dim=-1)


# ============================================================================
# Gated Linear Diffusion (GLD) Block - The Core Spatial Mixer
# ============================================================================

class GatedLinearAttention(nn.Module):
    """
    Gated Linear Attention for 2D spatial mixing.
    O(N) complexity instead of O(N²) softmax attention.
    
    Based on ViG/GLA research but adapted for diffusion:
    - Bidirectional scan (forward + backward)
    - 2D locality injection via depthwise conv gating
    - Token-differential operator to prevent oversmoothing (from DyDiLA)
    
    Math:
        Q, K, V = linear(x), linear(x), linear(x)
        Q = phi(Q), K = phi(K)  where phi = 1 + elu (non-negative feature map)
        
        Forward scan:  S_i = decay * S_{i-1} + K_i^T V_i;  O_i = Q_i S_i
        Backward scan: same in reverse
        
        Output = gate * (O_fwd + O_bwd) * local_gate
        
    Complexity: O(N * d²) where d is head dimension, N is sequence length
    """
    def __init__(self, dim: int, num_heads: int = 8, head_dim: int = 32, dropout: float = 0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        inner_dim = num_heads * head_dim
        
        self.qkv = nn.Linear(dim, 3 * inner_dim, bias=False)
        self.out_proj = nn.Linear(inner_dim, dim, bias=False)
        
        # Learnable decay for recurrence (per-head)
        self.log_decay = nn.Parameter(torch.zeros(num_heads))
        
        # Gate for output
        self.gate = nn.Linear(dim, inner_dim, bias=False)
        
        # 2D locality injection (depthwise conv) - critical for spatial structure
        self.local_conv = nn.Conv2d(inner_dim, inner_dim, 3, padding=1, groups=inner_dim, bias=False)
        self.local_gate = nn.Linear(dim, inner_dim, bias=False)
        
        # Token differential parameter (from DyDiLA - prevents oversmoothing)
        self.diff_lambda = nn.Parameter(torch.tensor(0.1))
        
        self.dropout = nn.Dropout(dropout)
        self.norm = RMSNorm(inner_dim)
    
    def _feature_map(self, x):
        """Non-negative feature map: 1 + elu(x)"""
        return 1.0 + F.elu(x)
    
    def _scan(self, Q, K, V, reverse=False):
        """Linear recurrent scan - O(N * d²) per direction."""
        B, H, N, D = Q.shape
        
        decay = torch.sigmoid(self.log_decay).view(1, H, 1, 1)  # [1, H, 1, 1]
        
        if reverse:
            Q = Q.flip(2)
            K = K.flip(2)
            V = V.flip(2)
        
        # Chunk-wise computation for memory efficiency
        chunk_size = min(64, N)
        outputs = []
        S = torch.zeros(B, H, D, D, device=Q.device, dtype=Q.dtype)
        
        for i in range(0, N, chunk_size):
            q_chunk = Q[:, :, i:i+chunk_size]  # [B, H, C, D]
            k_chunk = K[:, :, i:i+chunk_size]
            v_chunk = V[:, :, i:i+chunk_size]
            
            chunk_len = q_chunk.shape[2]
            
            # Update state: S = decay * S + K^T V
            kv = torch.einsum('bhcd,bhce->bhde', k_chunk, v_chunk)
            S = decay * S + kv
            
            # Query state: O = Q S
            o_chunk = torch.einsum('bhcd,bhde->bhce', q_chunk, S)
            outputs.append(o_chunk)
        
        output = torch.cat(outputs, dim=2)
        
        if reverse:
            output = output.flip(2)
        
        return output
    
    def forward(self, x, h: int, w: int):
        """
        Args:
            x: [B, N, D] where N = H*W
            h, w: spatial dimensions
        Returns:
            [B, N, D]
        """
        B, N, D = x.shape
        
        # Project to Q, K, V
        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)
        
        # Reshape to heads
        q = rearrange(q, 'b n (h d) -> b h n d', h=self.num_heads)
        k = rearrange(k, 'b n (h d) -> b h n d', h=self.num_heads)
        v = rearrange(v, 'b n (h d) -> b h n d', h=self.num_heads)
        
        # Token differential (prevents oversmoothing)
        # Q_diff = Q_i - lambda * Q_{i-1}, K_diff = K_i - lambda * K_{i-1}
        lam = torch.sigmoid(self.diff_lambda)
        q_shifted = F.pad(q[:, :, :-1], (0, 0, 1, 0))
        k_shifted = F.pad(k[:, :, :-1], (0, 0, 1, 0))
        q = q - lam * q_shifted
        k = k - lam * k_shifted
        
        # Apply feature map (non-negative)
        q = self._feature_map(q)
        k = self._feature_map(k)
        
        # Bidirectional scan
        o_fwd = self._scan(q, k, v, reverse=False)
        o_bwd = self._scan(q, k, v, reverse=True)
        output = o_fwd + o_bwd
        
        # Normalize
        output = rearrange(output, 'b h n d -> b n (h d)')
        output = self.norm(output)
        
        # 2D locality injection (GaLI from ViG)
        x_2d = rearrange(x, 'b (h w) d -> b d h w', h=h, w=w)
        gate_input = rearrange(x, 'b n d -> b n d')
        local_feat = self.local_conv(rearrange(self.local_gate(gate_input), 'b (h w) d -> b d h w', h=h, w=w))
        local_feat = rearrange(local_feat, 'b d h w -> b (h w) d')
        
        # Gated output
        g = torch.sigmoid(self.gate(x))
        output = g * output * torch.sigmoid(local_feat)
        
        return self.dropout(self.out_proj(output))


class GLDBlock(nn.Module):
    """
    Gated Linear Diffusion Block.
    
    Components:
    1. GatedLinearAttention for spatial mixing (O(N) complexity)
    2. SwiGLU FFN for channel mixing
    3. Timestep + condition modulation (adaptive layer norm)
    4. 2D RoPE for position encoding
    
    This replaces the standard transformer block in diffusion models.
    """
    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        head_dim: int = 32,
        ffn_mult: float = 2.67,
        dropout: float = 0.0,
        cond_dim: int = 256,
    ):
        super().__init__()
        self.norm1 = RMSNorm(dim)
        self.norm2 = RMSNorm(dim)
        
        self.attn = GatedLinearAttention(dim, num_heads, head_dim, dropout)
        self.ffn = SwiGLU(dim, int(dim * ffn_mult), dropout)
        
        # Adaptive modulation (scale, shift, gate for each sub-layer)
        # Conditioned on timestep + text embedding
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(cond_dim, 6 * dim, bias=False),
        )
        
        # Cross-attention to text (lightweight - only when text is available)
        self.cross_norm = RMSNorm(dim)
        self.cross_q = nn.Linear(dim, dim, bias=False)
        self.cross_kv = nn.Linear(cond_dim, 2 * dim, bias=False)
        self.cross_out = nn.Linear(dim, dim, bias=False)
        self.cross_gate = nn.Parameter(torch.zeros(1))  # Zero-init for residual
    
    def forward(
        self,
        x: torch.Tensor,        # [B, N, D]
        cond: torch.Tensor,      # [B, cond_dim] - timestep + global condition
        text_ctx: Optional[torch.Tensor] = None,  # [B, T, cond_dim] - text tokens
        h: int = 32,
        w: int = 32,
    ) -> torch.Tensor:
        B, N, D = x.shape
        
        # Compute modulation parameters
        mod = self.adaLN_modulation(cond)  # [B, 6*D]
        shift1, scale1, gate1, shift2, scale2, gate2 = mod.chunk(6, dim=-1)
        
        # Pre-norm + modulate + GLA
        x_norm = self.norm1(x)
        x_norm = x_norm * (1 + scale1.unsqueeze(1)) + shift1.unsqueeze(1)
        x = x + gate1.unsqueeze(1) * self.attn(x_norm, h, w)
        
        # Cross-attention to text (if available)
        if text_ctx is not None:
            x_cross = self.cross_norm(x)
            q = self.cross_q(x_cross)
            kv = self.cross_kv(text_ctx)
            k, v = kv.chunk(2, dim=-1)
            
            # Simple dot-product attention (text sequence is short, so O(N*T) is fine)
            scale = q.shape[-1] ** -0.5
            attn_weights = torch.bmm(q, k.transpose(-2, -1)) * scale
            attn_weights = F.softmax(attn_weights, dim=-1)
            cross_out = torch.bmm(attn_weights, v)
            x = x + torch.tanh(self.cross_gate) * self.cross_out(cross_out)
        
        # Pre-norm + modulate + FFN
        x_norm = self.norm2(x)
        x_norm = x_norm * (1 + scale2.unsqueeze(1)) + shift2.unsqueeze(1)
        x = x + gate2.unsqueeze(1) * self.ffn(x_norm)
        
        return x


# ============================================================================
# Recursive Latent Refinement (RLR) Core - THE KEY INNOVATION
# ============================================================================

class RecursiveLatentCore(nn.Module):
    """
    The Recursive Latent Refinement (RLR) Core.
    
    This is the key architectural innovation of LRF. Instead of stacking
    many unique transformer layers (like DiT with 28 layers), we use a 
    small set of GLD blocks applied RECURSIVELY through an HRM-inspired
    iterative refinement loop.
    
    Architecture:
    - N_blocks GLD blocks (typically 4-6, shared across recursions)
    - T_inner recursive applications per outer step (typically 4-6)
    - T_outer outer steps with slow abstract state update (typically 2-3)
    
    Effective depth: T_outer * T_inner * N_blocks = 2*4*4 = 32 effective layers
    Actual parameters: only N_blocks sets = 4 unique block parameter sets
    
    Training uses IFT (Implicit Function Theorem):
    - Forward: run full recursion with torch.no_grad() for warmup
    - Backward: only backprop through the LAST recursion step
    - This gives O(1) memory cost regardless of recursion depth!
    
    Mathematical formulation:
    
    Let z be the noisy latent, c be the condition embedding.
    
    Outer loop (j = 1..T_outer):
        z_abstract = f_slow(z, c)                    # Abstract planning update
        Inner loop (i = 1..T_inner):
            z = f_blocks(z, z_abstract, c)            # Apply N shared GLD blocks
        
    Where f_blocks applies the same N GLD blocks in sequence.
    
    The model learns a FIXED POINT: z* = f(z*, c)
    At convergence, the output is the denoised prediction v(z_t, t, c).
    """
    
    def __init__(
        self,
        dim: int = 384,
        cond_dim: int = 256,
        num_blocks: int = 4,
        num_heads: int = 6,
        head_dim: int = 64,
        T_inner: int = 4,
        T_outer: int = 2,
        ffn_mult: float = 2.67,
        dropout: float = 0.0,
        use_ift_training: bool = True,
    ):
        super().__init__()
        self.dim = dim
        self.cond_dim = cond_dim
        self.num_blocks = num_blocks
        self.T_inner = T_inner
        self.T_outer = T_outer
        self.use_ift_training = use_ift_training
        
        # The shared GLD blocks (applied recursively)
        self.blocks = nn.ModuleList([
            GLDBlock(
                dim=dim,
                num_heads=num_heads,
                head_dim=head_dim,
                ffn_mult=ffn_mult,
                dropout=dropout,
                cond_dim=cond_dim,
            )
            for _ in range(num_blocks)
        ])
        
        # Abstract state updater (the "slow" H-module from HRM)
        # This updates a global abstract representation every T_inner steps
        self.abstract_norm = RMSNorm(dim)
        self.abstract_update = nn.Sequential(
            nn.Linear(dim * 2, dim, bias=False),
            nn.SiLU(),
            nn.Linear(dim, dim, bias=False),
        )
        self.abstract_gate = nn.Parameter(torch.zeros(1))  # Zero-init
        
        # Input projection
        self.input_proj = nn.Linear(dim, dim, bias=False)
        
        # Timestep embedding
        self.time_embed = nn.Sequential(
            nn.Linear(256, cond_dim),
            nn.SiLU(),
            nn.Linear(cond_dim, cond_dim),
        )
        
        # Output projection (predicts velocity v for rectified flow)
        self.out_norm = RMSNorm(dim)
        self.out_proj = nn.Sequential(
            nn.Linear(dim, dim, bias=False),
            nn.SiLU(),
            nn.Linear(dim, dim, bias=False),
        )
        
        # Recursion depth embedding (tells the model which recursion step it's on)
        self.recursion_embed = nn.Embedding(T_outer * T_inner + 1, cond_dim)
        
        # 2D positional encoding
        self.rope = RotaryPositionEncoding2D(head_dim)
    
    def _sinusoidal_embedding(self, t: torch.Tensor, dim: int = 256) -> torch.Tensor:
        """Sinusoidal timestep embedding."""
        half_dim = dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
        emb = t.unsqueeze(-1) * emb.unsqueeze(0)
        return torch.cat([emb.sin(), emb.cos()], dim=-1)
    
    def _apply_blocks(self, z, cond, text_ctx, h, w):
        """Apply all GLD blocks once."""
        for block in self.blocks:
            z = block(z, cond, text_ctx, h, w)
        return z
    
    def _recursive_refinement(self, z, cond_base, text_ctx, h, w):
        """
        Full recursive refinement loop.
        
        Returns the refined latent z after T_outer * T_inner applications.
        """
        z_abstract = z.mean(dim=1, keepdim=True).expand_as(z)  # Initial abstract state
        
        step_idx = 0
        for j in range(self.T_outer):
            # Abstract state update (slow H-module)
            z_pooled = z.mean(dim=1, keepdim=True).expand_as(z)
            abstract_input = torch.cat([self.abstract_norm(z), z_pooled], dim=-1)
            z_abstract = z_abstract + torch.tanh(self.abstract_gate) * self.abstract_update(abstract_input)
            
            for i in range(self.T_inner):
                # Add recursion depth information to conditioning
                rec_emb = self.recursion_embed(
                    torch.tensor([step_idx], device=z.device)
                ).expand(z.shape[0], -1)
                cond = cond_base + rec_emb
                
                # Apply shared blocks with abstract state modulation
                z_input = z + z_abstract  # Combine detail + abstract
                z = z + (self._apply_blocks(z_input, cond, text_ctx, h, w) - z) * 0.5  # Damped update
                
                step_idx += 1
        
        return z
    
    def forward(
        self,
        z_t: torch.Tensor,       # [B, C, H, W] - noisy latent
        t: torch.Tensor,          # [B] - timestep (0 to 1)
        text_emb: Optional[torch.Tensor] = None,   # [B, T, cond_dim] - text tokens
        text_global: Optional[torch.Tensor] = None, # [B, cond_dim] - global text embedding
        image_cond: Optional[torch.Tensor] = None,   # [B, C, H, W] - for editing tasks
    ) -> torch.Tensor:
        """
        Forward pass predicting velocity v_theta(z_t, t, c).
        
        For rectified flow: z_t = (1-t) * z_0 + t * epsilon
        Target: v = epsilon - z_0
        """
        B, C, H, W = z_t.shape
        
        # Flatten spatial dims
        z = rearrange(z_t, 'b c h w -> b (h w) c')
        
        # If editing: concatenate condition image (channel-wise before projection)
        if image_cond is not None:
            img_cond_flat = rearrange(image_cond, 'b c h w -> b (h w) c')
            z = z + img_cond_flat  # Additive conditioning preserves spatial correspondence
        
        # Project
        z = self.input_proj(z)
        
        # Build conditioning
        t_emb = self._sinusoidal_embedding(t)
        t_emb = self.time_embed(t_emb)  # [B, cond_dim]
        
        if text_global is not None:
            cond = t_emb + text_global
        else:
            cond = t_emb
        
        # Apply recursive refinement
        if self.training and self.use_ift_training:
            # IFT training: no_grad warmup + 1-step grad
            with torch.no_grad():
                for _ in range(self.T_outer - 1):
                    z = self._recursive_refinement(z, cond, text_emb, H, W)
            # Last step with gradients
            z = self._recursive_refinement(z, cond, text_emb, H, W)
        else:
            # Full recursion (inference or non-IFT training)
            z = self._recursive_refinement(z, cond, text_emb, H, W)
        
        # Output projection
        z = self.out_norm(z)
        v = self.out_proj(z)
        
        # Reshape back to spatial
        v = rearrange(v, 'b (h w) c -> b c h w', h=H, w=W)
        
        return v


# ============================================================================
# Compact VAE (Tiny Decoder inspired by SnapGen)
# ============================================================================

class TinyResBlock(nn.Module):
    """Ultra-compact residual block for tiny decoder."""
    def __init__(self, in_channels: int, out_channels: int = None):
        super().__init__()
        out_channels = out_channels or in_channels
        self.norm1 = nn.GroupNorm(min(8, in_channels), in_channels)
        self.conv1 = DepthwiseSeparableConv2d(in_channels, out_channels, 3)
        self.norm2 = nn.GroupNorm(min(8, out_channels), out_channels)
        self.conv2 = DepthwiseSeparableConv2d(out_channels, out_channels, 3)
        self.skip = nn.Conv2d(in_channels, out_channels, 1, bias=False) if in_channels != out_channels else nn.Identity()
    
    def forward(self, x):
        h = self.conv1(F.silu(self.norm1(x)))
        h = self.conv2(F.silu(self.norm2(h)))
        return self.skip(x) + h


class CompactEncoder(nn.Module):
    """
    Compact image encoder: image -> latent space.
    f=16 spatial compression, C_latent channels.
    
    Uses strided depthwise-separable convolutions for efficiency.
    4 downsampling stages: 256->128->64->32->16 (for 256x256 input)
    """
    def __init__(
        self,
        in_channels: int = 3,
        latent_channels: int = 32,
        base_channels: int = 64,
        num_res_blocks: int = 2,
    ):
        super().__init__()
        channels = [base_channels, base_channels * 2, base_channels * 4, base_channels * 4]
        
        self.stem = nn.Conv2d(in_channels, channels[0], 3, padding=1, bias=False)
        
        self.downs = nn.ModuleList()
        ch_in = channels[0]
        for ch_out in channels:
            blocks = nn.ModuleList()
            # First block handles channel transition
            blocks.append(TinyResBlock(ch_in, ch_out))
            for _ in range(num_res_blocks - 1):
                blocks.append(TinyResBlock(ch_out, ch_out))
            # Downsample with strided conv
            down = nn.Conv2d(ch_out, ch_out, 4, stride=2, padding=1, bias=False)
            self.downs.append(nn.ModuleDict({
                'blocks': blocks,
                'down': down,
            }))
            ch_in = ch_out
        
        # To latent
        self.to_latent = nn.Sequential(
            nn.GroupNorm(8, ch_in),
            nn.SiLU(),
            nn.Conv2d(ch_in, latent_channels * 2, 1, bias=False),  # *2 for mean+logvar
        )
    
    def forward(self, x):
        h = self.stem(x)
        for down_module in self.downs:
            for block in down_module['blocks']:
                h = block(h)
            h = down_module['down'](h)
        
        params = self.to_latent(h)
        mean, logvar = params.chunk(2, dim=1)
        logvar = torch.clamp(logvar, -30.0, 20.0)
        
        return mean, logvar


class TinyDecoder(nn.Module):
    """
    SnapGen-inspired tiny decoder: latent -> image.
    ~1-2M parameters. No attention layers.
    Uses depthwise-separable convolutions + minimal GroupNorm.
    
    4 upsampling stages matching the encoder.
    """
    def __init__(
        self,
        latent_channels: int = 32,
        out_channels: int = 3,
        base_channels: int = 128,
        num_res_blocks: int = 2,
    ):
        super().__init__()
        channels = [base_channels * 2, base_channels * 2, base_channels, base_channels // 2]
        
        self.from_latent = nn.Conv2d(latent_channels, channels[0], 1, bias=False)
        
        self.ups = nn.ModuleList()
        ch_in = channels[0]
        for ch_out in channels:
            blocks = nn.ModuleList()
            for _ in range(num_res_blocks):
                blocks.append(TinyResBlock(ch_in, ch_in))
            # Upsample with channel transition
            up = nn.Sequential(
                nn.Upsample(scale_factor=2, mode='nearest'),
                DepthwiseSeparableConv2d(ch_in, ch_out, 3),
            )
            self.ups.append(nn.ModuleDict({
                'blocks': blocks,
                'up': up,
            }))
            ch_in = ch_out
        
        self.to_image = nn.Sequential(
            nn.GroupNorm(min(8, ch_in), ch_in),
            nn.SiLU(),
            nn.Conv2d(ch_in, out_channels, 3, padding=1),
            nn.Tanh(),  # Output in [-1, 1]
        )
    
    def forward(self, z):
        h = self.from_latent(z)
        for up_module in self.ups:
            for block in up_module['blocks']:
                h = block(h)
            h = up_module['up'](h)
        return self.to_image(h)


class CompactVAE(nn.Module):
    """
    Complete VAE with compact encoder + tiny decoder.
    f=16 compression, configurable latent channels.
    """
    def __init__(
        self,
        in_channels: int = 3,
        latent_channels: int = 32,
        encoder_base_ch: int = 64,
        decoder_base_ch: int = 128,
    ):
        super().__init__()
        self.encoder = CompactEncoder(in_channels, latent_channels, encoder_base_ch)
        self.decoder = TinyDecoder(latent_channels, in_channels, decoder_base_ch)
        self.latent_channels = latent_channels
    
    def encode(self, x):
        mean, logvar = self.encoder(x)
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            z = mean + eps * std
        else:
            z = mean
        return z, mean, logvar
    
    def decode(self, z):
        return self.decoder(z)
    
    def forward(self, x):
        z, mean, logvar = self.encode(x)
        recon = self.decode(z)
        return recon, mean, logvar


# ============================================================================
# Text Conditioner (Lightweight)
# ============================================================================

class SimpleTextEncoder(nn.Module):
    """
    Lightweight text encoder for the standalone prototype.
    In production, this would be replaced by TinyCLIP or a small LM.
    
    For the prototype: simple learned embeddings + small transformer.
    This lets us test the full pipeline without a heavy text encoder.
    """
    def __init__(
        self,
        vocab_size: int = 32000,
        max_length: int = 77,
        dim: int = 256,
        num_layers: int = 4,
        num_heads: int = 4,
    ):
        super().__init__()
        self.dim = dim
        self.token_embed = nn.Embedding(vocab_size, dim)
        self.pos_embed = nn.Embedding(max_length, dim)
        
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=dim, nhead=num_heads, dim_feedforward=dim*4,
            dropout=0.1, activation='gelu', batch_first=True, norm_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.norm = RMSNorm(dim)
        
        # Global pooling projection
        self.global_proj = nn.Sequential(
            nn.Linear(dim, dim),
            nn.SiLU(),
            nn.Linear(dim, dim),
        )
    
    def forward(self, token_ids, attention_mask=None):
        B, T = token_ids.shape
        pos_ids = torch.arange(T, device=token_ids.device).unsqueeze(0).expand(B, -1)
        
        x = self.token_embed(token_ids) + self.pos_embed(pos_ids)
        
        if attention_mask is not None:
            # Convert to transformer mask (True = ignore)
            src_key_padding_mask = ~attention_mask.bool()
        else:
            src_key_padding_mask = None
        
        x = self.transformer(x, src_key_padding_mask=src_key_padding_mask)
        x = self.norm(x)
        
        # Global embedding (mean pool over non-padded tokens)
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).float()
            global_emb = (x * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
        else:
            global_emb = x.mean(dim=1)
        
        global_emb = self.global_proj(global_emb)
        
        return x, global_emb  # [B, T, D], [B, D]


# ============================================================================
# Full LRF Model
# ============================================================================

class LatentRecurrentFlow(nn.Module):
    """
    LatentRecurrentFlow (LRF) - Complete model.
    
    Combines:
    1. CompactVAE for image encoding/decoding
    2. SimpleTextEncoder for text conditioning  
    3. RecursiveLatentCore for denoising
    
    Training modes:
    - 'vae': Train only the VAE
    - 'denoise': Train only the denoising core (freeze VAE)
    - 'e2e': End-to-end fine-tuning
    - 'distill': Consistency distillation from teacher
    """
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__()
        
        config = config or self.default_config()
        self.config = config
        
        # VAE
        self.vae = CompactVAE(
            in_channels=3,
            latent_channels=config['latent_channels'],
            encoder_base_ch=config.get('encoder_base_ch', 64),
            decoder_base_ch=config.get('decoder_base_ch', 128),
        )
        
        # Text encoder
        self.text_encoder = SimpleTextEncoder(
            vocab_size=config.get('vocab_size', 32000),
            max_length=config.get('max_text_length', 77),
            dim=config['cond_dim'],
            num_layers=config.get('text_layers', 4),
            num_heads=config.get('text_heads', 4),
        )
        
        # Denoising core
        self.core = RecursiveLatentCore(
            dim=config['latent_channels'],
            cond_dim=config['cond_dim'],
            num_blocks=config['num_blocks'],
            num_heads=config.get('num_heads', 6),
            head_dim=config.get('head_dim', 64),
            T_inner=config.get('T_inner', 4),
            T_outer=config.get('T_outer', 2),
            ffn_mult=config.get('ffn_mult', 2.67),
            dropout=config.get('dropout', 0.0),
            use_ift_training=config.get('use_ift', True),
        )
        
        # Latent scaling (learnable, stabilizes training)
        self.latent_scale = nn.Parameter(torch.tensor(1.0))
    
    @staticmethod
    def default_config():
        """Default config targeting ~50M params, trainable on 16GB."""
        return {
            'latent_channels': 32,
            'cond_dim': 256,
            'num_blocks': 4,
            'num_heads': 4,
            'head_dim': 64,
            'T_inner': 4,
            'T_outer': 2,
            'ffn_mult': 2.67,
            'dropout': 0.0,
            'use_ift': True,
            'encoder_base_ch': 64,
            'decoder_base_ch': 128,
            'vocab_size': 32000,
            'max_text_length': 77,
            'text_layers': 4,
            'text_heads': 4,
        }
    
    @staticmethod
    def tiny_config():
        """Tiny config for quick testing."""
        return {
            'latent_channels': 16,
            'cond_dim': 128,
            'num_blocks': 2,
            'num_heads': 2,
            'head_dim': 32,
            'T_inner': 2,
            'T_outer': 1,
            'ffn_mult': 2.0,
            'dropout': 0.0,
            'use_ift': False,
            'encoder_base_ch': 32,
            'decoder_base_ch': 64,
            'vocab_size': 32000,
            'max_text_length': 77,
            'text_layers': 2,
            'text_heads': 2,
        }
    
    def encode_image(self, x):
        """Encode image to latent space."""
        z, mean, logvar = self.vae.encode(x)
        return z * self.latent_scale, mean, logvar
    
    def decode_latent(self, z):
        """Decode latent to image."""
        return self.vae.decode(z / self.latent_scale)
    
    def encode_text(self, token_ids, attention_mask=None):
        """Encode text to conditioning vectors."""
        return self.text_encoder(token_ids, attention_mask)
    
    def predict_velocity(self, z_t, t, text_emb=None, text_global=None, image_cond=None):
        """Predict velocity for rectified flow."""
        return self.core(z_t, t, text_emb, text_global, image_cond)
    
    def get_param_groups(self):
        """Return parameter groups for staged training."""
        return {
            'vae_encoder': list(self.vae.encoder.parameters()),
            'vae_decoder': list(self.vae.decoder.parameters()),
            'text_encoder': list(self.text_encoder.parameters()),
            'core': list(self.core.parameters()),
            'latent_scale': [self.latent_scale],
        }
    
    def count_parameters(self):
        """Count parameters per module."""
        counts = {}
        for name, module in [
            ('vae_encoder', self.vae.encoder),
            ('vae_decoder', self.vae.decoder),
            ('text_encoder', self.text_encoder),
            ('core', self.core),
        ]:
            counts[name] = sum(p.numel() for p in module.parameters())
        counts['latent_scale'] = 1
        counts['total'] = sum(counts.values())
        return counts
    
    def forward(self, x=None, token_ids=None, attention_mask=None, **kwargs):
        """Full forward pass for training. See training script for usage."""
        raise NotImplementedError(
            "Use the training pipeline functions instead of calling forward() directly. "
            "See LRFTrainer for VAE training, denoiser training, and distillation."
        )