"""
⚡ ULTRON Systems // Advanced Convolutional Transformer Architecture for OHLCV
Implements:
1. Reversible Instance Normalization (RevIN)
2. Multi-Scale Overlapping 1D Convolutional Tokenizer (kernels 3, 7, 15, 31)
3. Rotary Positional Embeddings (RoPE)
4. Deep Conformer Blocks (Self-Attention + Depthwise Separable 1D Conv)
5. Multi-Horizon Dual Forecasting Head (Return Magnitude + Directional Probability)
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple

from config import ModelConfig

class RevIN(nn.Module):
    """
    Reversible Instance Normalization to eliminate non-stationary distribution drift.
    Normalizes input sequences per-instance and per-feature.
    """
    def __init__(self, num_features: int, eps: float = 1e-5, affine: bool = True):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.affine = affine
        if self.affine:
            self.affine_weight = nn.Parameter(torch.ones(num_features))
            self.affine_bias = nn.Parameter(torch.zeros(num_features))

    def forward(self, x: torch.Tensor, mode: str = "norm") -> torch.Tensor:
        # x: [Batch, SeqLen, Features]
        if mode == "norm":
            self.mean = torch.mean(x, dim=1, keepdim=True).detach()
            self.stdev = torch.sqrt(torch.var(x, dim=1, keepdim=True, unbiased=False) + self.eps).detach()
            x_norm = (x - self.mean) / self.stdev
            if self.affine:
                x_norm = x_norm * self.affine_weight + self.affine_bias
            return x_norm
        elif mode == "denorm":
            x_denorm = x
            if self.affine:
                x_denorm = (x_denorm - self.affine_bias) / (self.affine_weight + self.eps)
            return x_denorm * self.stdev + self.mean
        else:
            raise ValueError(f"Unknown RevIN mode: {mode}")

class MultiScaleOverlappingConvTokenizer(nn.Module):
    """
    Extracts multi-frequency candlestick representations via overlapping 1D convolutions.
    Operates over overlapping temporal windows with kernel sizes [3, 7, 15, 31].
    Stride S < K ensures heavy overlap between adjacent token receptive fields.
    """
    def __init__(self, in_features: int, d_model: int, kernels: List[int] = [3, 7, 15, 31], stride: int = 4):
        super().__init__()
        self.stride = stride
        branch_dim = d_model // len(kernels)
        self.branches = nn.ModuleList()
        
        for k in kernels:
            # Symmetric causal/same padding
            pad = k // 2
            self.branches.append(
                nn.Sequential(
                    nn.Conv1d(in_features, branch_dim, kernel_size=k, stride=stride, padding=pad),
                    nn.GELU()
                )
            )
            
        # Projection layer to fuse multi-scale features into exact d_model
        concat_dim = branch_dim * len(kernels)
        self.proj = nn.Linear(concat_dim, d_model)
        self.norm = nn.LayerNorm(d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [Batch, SeqLen, InFeatures] -> [Batch, InFeatures, SeqLen]
        x_trans = x.transpose(1, 2)
        
        branch_outs = []
        min_len = None
        for branch in self.branches:
            out = branch(x_trans)
            if min_len is None or out.shape[2] < min_len:
                min_len = out.shape[2]
            branch_outs.append(out)
            
        # Align temporal length across all branches
        aligned = [b[:, :, :min_len] for b in branch_outs]
        fused = torch.cat(aligned, dim=1) # [Batch, concat_dim, Patches]
        
        # Transpose back: [Batch, Patches, concat_dim]
        fused = fused.transpose(1, 2)
        return self.norm(self.proj(fused))

class RotaryEmbedding(nn.Module):
    """Rotary Position Embeddings (RoPE) for sequence tokens."""
    def __init__(self, dim: int, max_seq_len: int = 2048):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        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) # [1, 1, seq_len, dim]
        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 ConformerConvModule(nn.Module):
    """
    Depthwise Separable Convolution block to model fine-grained micro-structure
    and local volatility gradients.
    """
    def __init__(self, d_model: int, kernel_size: int = 7, dropout: float = 0.1):
        super().__init__()
        self.norm = nn.LayerNorm(d_model)
        # Pointwise expansion
        self.pointwise1 = nn.Conv1d(d_model, 2 * d_model, kernel_size=1)
        # Depthwise conv
        self.depthwise = nn.Conv1d(
            d_model, d_model, 
            kernel_size=kernel_size, 
            stride=1, 
            padding=kernel_size // 2, 
            groups=d_model
        )
        self.bn = nn.BatchNorm1d(d_model)
        self.act = nn.SiLU()
        # Pointwise projection
        self.pointwise2 = nn.Conv1d(d_model, d_model, kernel_size=1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [Batch, Patches, d_model]
        res = x
        x = self.norm(x).transpose(1, 2)
        # Pointwise expansion with Gated Linear Unit
        x = self.pointwise1(x)
        x = F.glu(x, dim=1) # [Batch, d_model, Patches]
        # Depthwise
        x = self.depthwise(x)
        x = self.bn(x)
        x = self.act(x)
        x = self.pointwise2(x)
        x = self.dropout(x)
        return res + x.transpose(1, 2)

class ConformerBlock(nn.Module):
    """
    Hybrid Convolution + FlashAttention Conformer Block.
    """
    def __init__(self, d_model: int, n_heads: int = 4, d_ff: int = 256, conv_kernel: int = 7, dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        
        # 1. Multi-Head Self-Attention with RoPE
        self.ln_attn = nn.LayerNorm(d_model)
        self.q_proj = nn.Linear(d_model, d_model, bias=False)
        self.k_proj = nn.Linear(d_model, d_model, bias=False)
        self.v_proj = nn.Linear(d_model, d_model, bias=False)
        self.out_proj = nn.Linear(d_model, d_model)
        self.rope = RotaryEmbedding(self.head_dim)
        self.attn_dropout = dropout
        
        # 2. Conformer Convolution Module
        self.conv_module = ConformerConvModule(d_model, kernel_size=conv_kernel, dropout=dropout)
        
        # 3. Feed-Forward SwiGLU Network
        self.ln_ffn = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Step 1: Pre-LN Self Attention with FlashAttention-2 / scaled_dot_product_attention
        res = x
        x_norm = self.ln_attn(x)
        B, T, C = 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_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x_norm).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        
        # Apply RoPE
        q, k = self.rope(q, k)
        
        # Native PyTorch FlashAttention / SDPA
        attn_out = F.scaled_dot_product_attention(
            q, k, v, 
            dropout_p=self.attn_dropout if self.training else 0.0,
            is_causal=True # Causal temporal mask ensures no future leakage
        )
        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, C)
        x = res + self.out_proj(attn_out)
        
        # Step 2: Conformer Depthwise Convolution
        x = self.conv_module(x)
        
        # Step 3: Feed-Forward
        x = x + self.ffn(self.ln_ffn(x))
        return x

class OHLCVTransformer(nn.Module):
    """
    End-to-End Advanced Convolutional Transformer for Multi-Horizon OHLCV Forecasting.
    """
    def __init__(self, cfg: ModelConfig):
        super().__init__()
        self.cfg = cfg
        self.revin = RevIN(num_features=cfg.in_features) if cfg.use_revin else None
        
        # Multi-scale overlapping tokenizer
        self.tokenizer = MultiScaleOverlappingConvTokenizer(
            in_features=cfg.in_features,
            d_model=cfg.d_model,
            kernels=cfg.conv_kernels,
            stride=cfg.patch_stride
        )
        
        # Conformer Layers
        self.blocks = nn.ModuleList([
            ConformerBlock(
                d_model=cfg.d_model,
                n_heads=cfg.n_heads,
                d_ff=cfg.d_ff,
                conv_kernel=cfg.conv_depthwise_kernel,
                dropout=cfg.dropout
            ) for _ in range(cfg.n_layers)
        ])
        
        self.final_norm = nn.LayerNorm(cfg.d_model)
        
        # Multi-Horizon Heads:
        # 1. Return Regression Head: predicts expected log return for [1m, 3m, 5m, 10m, 15m]
        self.reg_head = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.d_ff),
            nn.GELU(),
            nn.Dropout(cfg.dropout),
            nn.Linear(cfg.d_ff, cfg.num_horizons)
        )
        
        # 2. Directional Classification Head: predicts binary direction probability
        self.dir_head = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.d_ff),
            nn.GELU(),
            nn.Dropout(cfg.dropout),
            nn.Linear(cfg.d_ff, cfg.num_horizons)
        )

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        x: [Batch, SeqLen, InFeatures]
        Returns:
            pred_returns: [Batch, NumHorizons] (expected future log returns)
            pred_logits:  [Batch, NumHorizons] (logits for P(return > 0))
        """
        # RevIN normalisation
        if self.revin is not None:
            x = self.revin(x, mode="norm")
            
        # Overlapping 1D conv tokenization
        tokens = self.tokenizer(x) # [Batch, Patches, d_model]
        
        # Transformer processing
        for block in self.blocks:
            tokens = block(tokens)
            
        tokens = self.final_norm(tokens)
        
        # Global aggregation: use final token (most recent causal context)
        context = tokens[:, -1, :] # [Batch, d_model]
        
        pred_returns = self.reg_head(context)
        pred_logits = self.dir_head(context)
        
        return pred_returns, pred_logits