"""
⚡ ULTRON Systems // 1.1 Billion Parameter Financial Conformer Foundation Model
Engineered for NVIDIA A100 (80GB) Tensor Core Accelerators:
1. Reversible Instance Normalization (RevIN)
2. Dilated Multi-Scale Overlapping 1D Conv Tokenizer (kernels 3, 7, 15, 31, 63, 127)
3. Grouped Query Attention (GQA) with RoPE and FlashAttention-2
4. Conformer Depthwise Separable Conv Modules in every block
5. SwiGLU Gated Feed-Forward Networks
6. Multi-Horizon Dual Forecasting + Masked Candlestick Pre-Training Heads
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple, Optional

from config_1b import Model1BConfig

class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization (standard in frontier LLMs)."""
    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: torch.Tensor) -> torch.Tensor:
        variance = x.pow(2).mean(-1, keepdim=True)
        return x * torch.rsqrt(variance + self.eps) * self.weight

class SwiGLUFFN(nn.Module):
    """SwiGLU Gated Feed-Forward Network."""
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.05):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up   = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU: (SiLU(x * W_gate) * (x * W_up)) * W_down
        return self.dropout(self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)))

class MultiScaleOverlappingTokenizer1B(nn.Module):
    """
    Overlapping 1D Convolution Tokenizer with 6 multi-scale receptive fields:
    3m, 7m, 15m, 31m, 63m, 127m.
    Compresses long sequence (e.g. 2048 bars) into 512 patch tokens with stride 4.
    """
    def __init__(self, in_features: int, d_model: int, kernels: List[int] = [3, 7, 15, 31, 63, 127], stride: int = 4):
        super().__init__()
        self.stride = stride
        branch_dim = d_model // len(kernels)
        self.branches = nn.ModuleList()
        
        for k in kernels:
            pad = k // 2
            self.branches.append(
                nn.Sequential(
                    nn.Conv1d(in_features, branch_dim, kernel_size=k, stride=stride, padding=pad),
                    nn.SiLU()
                )
            )
        concat_dim = branch_dim * len(kernels)
        self.proj = nn.Linear(concat_dim, d_model, bias=False)
        self.norm = RMSNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [Batch, SeqLen, InFeatures] -> [Batch, InFeatures, SeqLen]
        x_trans = x.transpose(1, 2)
        outs = []
        min_len = None
        for b in self.branches:
            out = b(x_trans)
            if min_len is None or out.shape[2] < min_len:
                min_len = out.shape[2]
            outs.append(out)
            
        aligned = [o[:, :, :min_len] for o in outs]
        fused = torch.cat(aligned, dim=1).transpose(1, 2)
        return self.norm(self.proj(fused))

class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_seq_len: int = 4096):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        t = torch.arange(max_seq_len, dtype=torch.float32)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)

    def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        seq_len = q.shape[2]
        cos = self.cos_cached[:seq_len, :].unsqueeze(0).unsqueeze(0)
        sin = self.sin_cached[:seq_len, :].unsqueeze(0).unsqueeze(0)
        
        def rotate_half(x):
            x1 = x[..., :x.shape[-1] // 2]
            x2 = x[..., x.shape[-1] // 2:]
            return torch.cat((-x2, x1), dim=-1)
            
        q_rot = (q * cos) + (rotate_half(q) * sin)
        k_rot = (k * cos) + (rotate_half(k) * sin)
        return q_rot, k_rot

class ConformerDepthwiseConv1B(nn.Module):
    """Conformer Depthwise Convolutional Sub-Layer."""
    def __init__(self, d_model: int, kernel_size: int = 7, dropout: float = 0.05):
        super().__init__()
        self.norm = RMSNorm(d_model)
        self.pw_in = nn.Linear(d_model, 2 * d_model, bias=False)
        self.dw_conv = nn.Conv1d(
            d_model, d_model, 
            kernel_size=kernel_size, 
            padding=kernel_size // 2, 
            groups=d_model,
            bias=False
        )
        self.act = nn.SiLU()
        self.pw_out = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        res = x
        x = self.norm(x)
        # Pointwise GLU expansion
        x = self.pw_in(x)
        x, gate = x.chunk(2, dim=-1)
        x = x * F.silu(gate)
        # Depthwise 1D conv
        x = x.transpose(1, 2)
        x = self.dw_conv(x)
        x = self.act(x).transpose(1, 2)
        # Pointwise projection
        x = self.dropout(self.pw_out(x))
        return res + x

class Conformer1BBlock(nn.Module):
    """
    1B Conformer Layer:
    Pre-RMSNorm -> GQA FlashAttention-2 with RoPE -> Depthwise Conv -> SwiGLU FFN
    """
    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int, d_ff: int, conv_kernel: int = 7, dropout: float = 0.05):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = d_model // n_heads
        self.num_queries_per_kv = n_heads // n_kv_heads
        
        # Attention
        self.norm_attn = RMSNorm(d_model)
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.out_proj = nn.Linear(d_model, d_model, bias=False)
        self.rope = RotaryEmbedding(self.head_dim)
        
        # Conformer Convolution module
        self.conv_module = ConformerDepthwiseConv1B(d_model, kernel_size=conv_kernel, dropout=dropout)
        
        # SwiGLU FFN
        self.norm_ffn = RMSNorm(d_model)
        self.ffn = SwiGLUFFN(d_model, d_ff, dropout=dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Step 1: Grouped Query Attention with FlashAttention-2
        res = x
        x_norm = self.norm_attn(x)
        B, T, _ = x_norm.shape
        
        q = self.q_proj(x_norm).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x_norm).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x_norm).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        
        # RoPE
        q, k = self.rope(q, k)
        
        # Expand KV heads for Grouped Query Attention (GQA)
        if self.num_queries_per_kv > 1:
            k = k.repeat_interleave(self.num_queries_per_kv, dim=1)
            v = v.repeat_interleave(self.num_queries_per_kv, dim=1)
            
        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, self.d_model)
        x = res + self.out_proj(attn_out)
        
        # Step 2: Conformer Convolution
        x = self.conv_module(x)
        
        # Step 3: SwiGLU FFN
        x = x + self.ffn(self.norm_ffn(x))
        return x

class FinancialConformer1B(nn.Module):
    """
    1.1 Billion Parameter Financial Conformer Foundation Model.
    """
    def __init__(self, cfg: Model1BConfig):
        super().__init__()
        self.cfg = cfg
        
        # Tokenizer
        self.tokenizer = MultiScaleOverlappingTokenizer1B(
            in_features=cfg.in_features,
            d_model=cfg.d_model,
            kernels=cfg.conv_kernels,
            stride=cfg.patch_stride
        )
        
        # 20 Deep Conformer Blocks
        self.blocks = nn.ModuleList([
            Conformer1BBlock(
                d_model=cfg.d_model,
                n_heads=cfg.n_heads,
                n_kv_heads=cfg.n_kv_heads,
                d_ff=cfg.d_ff,
                conv_kernel=cfg.conv_depthwise_kernel,
                dropout=cfg.dropout
            ) for _ in range(cfg.n_layers)
        ])
        
        self.final_norm = RMSNorm(cfg.d_model)
        
        # Multi-Horizon Return Trajectory Head: [1m, 3m, 5m, 15m, 30m, 60m, 1d]
        self.return_head = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.d_model // 2),
            nn.SiLU(),
            nn.Linear(cfg.d_model // 2, cfg.num_horizons)
        )
        
        # Directional Up/Down Probability Logits
        self.dir_head = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.d_model // 2),
            nn.SiLU(),
            nn.Linear(cfg.d_model // 2, cfg.num_horizons)
        )

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        # x: [Batch, SeqLen=2048, InFeatures=8]
        tokens = self.tokenizer(x) # [Batch, 512, 2048]
        
        for block in self.blocks:
            tokens = block(tokens)
            
        tokens = self.final_norm(tokens)
        context = tokens[:, -1, :] # Final temporal causal token [Batch, 2048]
        
        pred_returns = self.return_head(context)
        pred_logits = self.dir_head(context)
        return pred_returns, pred_logits

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

if __name__ == "__main__":
    cfg = Model1BConfig()
    model = FinancialConformer1B(cfg)
    total_params = model.count_parameters()
    print(f"⚡ [ULTRON] Financial Conformer 1B Initialized:")
    print(f"   • Total Trainable Parameters: {total_params:,} ({total_params / 1e9:.3f} Billion)")