#!/usr/bin/env python3
"""
VBP CNN Encoder - 1D Convolutional Neural Network with Self-Attention
======================================================================
Implements a 1D CNN architecture with self-attention for encoding
volatility patterns in financial time series.

Architecture:
- Multi-scale Conv1D layers with different kernel sizes
- BatchNorm + ReLU + Dropout after each conv
- Self-attention layer to capture long-range dependencies
- Global pooling for sequence-to-vector encoding

Key Design Choices:
1. Multiple kernel sizes (3, 5, 7) to capture patterns at different scales
2. Increasing filter counts (32, 64, 128) for hierarchical features
3. Self-attention after CNN to weight important timesteps
4. Residual connections for gradient flow

Author: ML-Specialist (NEXUS v4.0)
Version: 1.0.0
Created: 2026-01-25
"""

import math
from typing import Tuple, Optional, List, Dict
from dataclasses import dataclass, field

import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger


@dataclass
class CNNEncoderConfig:
    """Configuration for the CNN encoder."""

    # Input configuration
    input_channels: int = 1  # Number of input feature channels

    # Conv layer configurations [filters, kernel_size]
    conv_configs: List[Tuple[int, int]] = field(default_factory=lambda: [
        (32, 3),
        (64, 5),
        (128, 7)
    ])

    # Attention configuration
    attention_heads: int = 4
    attention_dropout: float = 0.1

    # General settings
    dropout: float = 0.3
    use_attention: bool = True
    use_residual: bool = True

    # Output dimension
    output_dim: int = 128


class Conv1DBlock(nn.Module):
    """
    1D Convolutional block with BatchNorm, ReLU, and Dropout.

    Applies:
    1. Conv1D with same padding
    2. BatchNorm
    3. ReLU activation
    4. Dropout

    Args:
        in_channels: Number of input channels
        out_channels: Number of output channels
        kernel_size: Convolution kernel size
        dropout: Dropout probability
        use_residual: Whether to add residual connection
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        dropout: float = 0.3,
        use_residual: bool = True
    ):
        super().__init__()

        self.use_residual = use_residual and (in_channels == out_channels)

        # Padding to maintain sequence length
        padding = kernel_size // 2

        self.conv = nn.Conv1d(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=kernel_size,
            padding=padding,
            bias=False
        )
        self.bn = nn.BatchNorm1d(out_channels)
        self.dropout = nn.Dropout(dropout)

        # Projection for residual if dimensions don't match
        if use_residual and in_channels != out_channels:
            self.residual_proj = nn.Conv1d(in_channels, out_channels, kernel_size=1)
            self.use_residual = True
        else:
            self.residual_proj = None

        self._init_weights()

    def _init_weights(self):
        """Initialize weights with Kaiming initialization."""
        nn.init.kaiming_normal_(self.conv.weight, mode='fan_out', nonlinearity='relu')
        nn.init.ones_(self.bn.weight)
        nn.init.zeros_(self.bn.bias)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass.

        Args:
            x: Input tensor of shape (batch, channels, seq_len)

        Returns:
            Output tensor of shape (batch, out_channels, seq_len)
        """
        identity = x

        out = self.conv(x)
        out = self.bn(out)
        out = F.relu(out)
        out = self.dropout(out)

        if self.use_residual:
            if self.residual_proj is not None:
                identity = self.residual_proj(identity)
            out = out + identity

        return out


class SelfAttention1D(nn.Module):
    """
    Self-Attention layer for 1D sequences.

    Applies multi-head self-attention to capture long-range dependencies
    in the sequence after CNN encoding.

    Args:
        embed_dim: Embedding dimension (same as number of channels)
        num_heads: Number of attention heads
        dropout: Attention dropout probability
    """

    def __init__(
        self,
        embed_dim: int,
        num_heads: int = 4,
        dropout: float = 0.1
    ):
        super().__init__()

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"

        # Linear projections for Q, K, V
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)

        # Output projection
        self.out_proj = nn.Linear(embed_dim, embed_dim)

        # Dropout
        self.dropout = nn.Dropout(dropout)

        # Layer norm for pre-attention
        self.norm = nn.LayerNorm(embed_dim)

        self._init_weights()

    def _init_weights(self):
        """Initialize weights."""
        nn.init.xavier_uniform_(self.q_proj.weight)
        nn.init.xavier_uniform_(self.k_proj.weight)
        nn.init.xavier_uniform_(self.v_proj.weight)
        nn.init.xavier_uniform_(self.out_proj.weight)

        nn.init.zeros_(self.q_proj.bias)
        nn.init.zeros_(self.k_proj.bias)
        nn.init.zeros_(self.v_proj.bias)
        nn.init.zeros_(self.out_proj.bias)

    def forward(
        self,
        x: torch.Tensor,
        return_attention: bool = False
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Forward pass.

        Args:
            x: Input tensor of shape (batch, seq_len, embed_dim)
            return_attention: Whether to return attention weights

        Returns:
            output: Attended output of shape (batch, seq_len, embed_dim)
            attention_weights: Optional attention weights (batch, heads, seq_len, seq_len)
        """
        batch_size, seq_len, _ = x.shape

        # Layer norm
        x_norm = self.norm(x)

        # Project to Q, K, V
        q = self.q_proj(x_norm)  # (batch, seq_len, embed_dim)
        k = self.k_proj(x_norm)
        v = self.v_proj(x_norm)

        # Reshape for multi-head attention
        # (batch, seq_len, num_heads, head_dim) -> (batch, num_heads, seq_len, head_dim)
        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # Scaled dot-product attention
        scale = math.sqrt(self.head_dim)
        scores = torch.matmul(q, k.transpose(-2, -1)) / scale  # (batch, heads, seq, seq)

        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)

        # Apply attention to values
        attended = torch.matmul(attention_weights, v)  # (batch, heads, seq, head_dim)

        # Reshape back
        attended = attended.transpose(1, 2).contiguous().view(batch_size, seq_len, self.embed_dim)

        # Output projection with residual
        output = self.out_proj(attended) + x

        if return_attention:
            return output, attention_weights
        return output, None


class VolatilityCNN(nn.Module):
    """
    1D CNN with Self-Attention for Volatility Pattern Encoding.

    Architecture:
    1. Multi-scale Conv1D blocks with increasing filters
    2. Self-attention layer for global context
    3. Global pooling (mean + max)
    4. Final projection to output dimension

    Input: (batch, seq_len, features)
    Output: (batch, output_dim) encoded representation

    Args:
        input_dim: Number of input features
        config: CNN configuration
    """

    def __init__(
        self,
        input_dim: int,
        config: Optional[CNNEncoderConfig] = None
    ):
        super().__init__()

        self.config = config or CNNEncoderConfig()
        self.input_dim = input_dim

        # Input projection to standard dimension
        first_channels = self.config.conv_configs[0][0]
        self.input_proj = nn.Linear(input_dim, first_channels)

        # Build convolutional layers
        self.conv_blocks = nn.ModuleList()
        in_channels = first_channels

        for out_channels, kernel_size in self.config.conv_configs:
            block = Conv1DBlock(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                dropout=self.config.dropout,
                use_residual=self.config.use_residual
            )
            self.conv_blocks.append(block)
            in_channels = out_channels

        self.final_channels = in_channels

        # Self-attention layer
        if self.config.use_attention:
            self.attention = SelfAttention1D(
                embed_dim=self.final_channels,
                num_heads=self.config.attention_heads,
                dropout=self.config.attention_dropout
            )
        else:
            self.attention = None

        # Final layer norm
        self.final_norm = nn.LayerNorm(self.final_channels)

        # Output projection
        # After global pooling: concat of mean and max = 2 * final_channels
        pooled_dim = 2 * self.final_channels
        self.output_proj = nn.Sequential(
            nn.Linear(pooled_dim, self.config.output_dim),
            nn.ReLU(),
            nn.Dropout(self.config.dropout),
            nn.Linear(self.config.output_dim, self.config.output_dim)
        )

        self._init_output_weights()

    def _init_output_weights(self):
        """Initialize output projection weights."""
        for module in self.output_proj:
            if isinstance(module, nn.Linear):
                nn.init.xavier_uniform_(module.weight)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)

    def forward(
        self,
        x: torch.Tensor,
        return_attention: bool = False
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Forward pass.

        Args:
            x: Input tensor of shape (batch, seq_len, features)
            return_attention: Whether to return attention weights

        Returns:
            encoded: Encoded representation (batch, output_dim)
            attention_weights: Optional attention weights
        """
        batch_size, seq_len, _ = x.shape

        # Project input features
        x = self.input_proj(x)  # (batch, seq_len, first_channels)

        # Transpose for Conv1D: (batch, channels, seq_len)
        x = x.transpose(1, 2)

        # Apply conv blocks
        for conv_block in self.conv_blocks:
            x = conv_block(x)

        # Transpose back: (batch, seq_len, channels)
        x = x.transpose(1, 2)

        # Apply self-attention
        attention_weights = None
        if self.attention is not None:
            x, attention_weights = self.attention(x, return_attention=return_attention)

        # Layer norm
        x = self.final_norm(x)

        # Global pooling: combine mean and max
        mean_pool = x.mean(dim=1)  # (batch, channels)
        max_pool = x.max(dim=1)[0]  # (batch, channels)
        pooled = torch.cat([mean_pool, max_pool], dim=-1)  # (batch, 2*channels)

        # Output projection
        encoded = self.output_proj(pooled)  # (batch, output_dim)

        return encoded, attention_weights

    def get_feature_maps(self, x: torch.Tensor) -> List[torch.Tensor]:
        """
        Get intermediate feature maps from each conv block.

        Useful for visualization and debugging.

        Args:
            x: Input tensor of shape (batch, seq_len, features)

        Returns:
            List of feature maps from each conv block
        """
        feature_maps = []

        # Project input
        x = self.input_proj(x)
        x = x.transpose(1, 2)

        # Get feature maps from each block
        for conv_block in self.conv_blocks:
            x = conv_block(x)
            feature_maps.append(x.transpose(1, 2))  # Store in (batch, seq, channels) format

        return feature_maps


class MultiScaleVolatilityCNN(nn.Module):
    """
    Multi-Scale CNN that processes input at different temporal scales.

    Applies separate CNN branches with different dilation rates to capture
    patterns at multiple time scales, then fuses the representations.

    Args:
        input_dim: Number of input features
        config: CNN configuration
        scales: List of dilation rates for multi-scale processing
    """

    def __init__(
        self,
        input_dim: int,
        config: Optional[CNNEncoderConfig] = None,
        scales: List[int] = None
    ):
        super().__init__()

        self.config = config or CNNEncoderConfig()
        self.scales = scales or [1, 2, 4]  # Different dilation rates
        self.input_dim = input_dim

        # Create a CNN branch for each scale
        self.branches = nn.ModuleList()
        for scale in self.scales:
            branch = self._create_branch(input_dim, scale)
            self.branches.append(branch)

        # Fusion layer
        branch_output_dim = self.config.output_dim
        fused_dim = branch_output_dim * len(self.scales)

        self.fusion = nn.Sequential(
            nn.Linear(fused_dim, self.config.output_dim * 2),
            nn.ReLU(),
            nn.Dropout(self.config.dropout),
            nn.Linear(self.config.output_dim * 2, self.config.output_dim)
        )

        # Self-attention for temporal weighting
        if self.config.use_attention:
            self.temporal_attention = SelfAttention1D(
                embed_dim=self.config.output_dim,
                num_heads=self.config.attention_heads,
                dropout=self.config.attention_dropout
            )
        else:
            self.temporal_attention = None

    def _create_branch(self, input_dim: int, dilation: int) -> nn.Module:
        """Create a single CNN branch with given dilation."""
        layers = []

        # Input projection
        first_channels = self.config.conv_configs[0][0]
        layers.append(nn.Linear(input_dim, first_channels))
        layers.append(nn.ReLU())

        # Store as sequential for the linear part
        input_proj = nn.Sequential(*layers)

        # Conv layers with dilation
        conv_layers = []
        in_channels = first_channels

        for out_channels, kernel_size in self.config.conv_configs:
            # Adjust padding for dilation
            padding = ((kernel_size - 1) * dilation) // 2

            conv = nn.Conv1d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                padding=padding,
                dilation=dilation,
                bias=False
            )
            bn = nn.BatchNorm1d(out_channels)

            conv_layers.extend([conv, bn, nn.ReLU(), nn.Dropout(self.config.dropout)])
            in_channels = out_channels

        conv_sequential = nn.Sequential(*conv_layers)

        # Output projection
        output_proj = nn.Linear(in_channels, self.config.output_dim)

        return nn.ModuleDict({
            'input_proj': input_proj,
            'conv': conv_sequential,
            'output_proj': output_proj
        })

    def forward(
        self,
        x: torch.Tensor,
        return_attention: bool = False
    ) -> Tuple[torch.Tensor, Optional[Dict[str, torch.Tensor]]]:
        """
        Forward pass.

        Args:
            x: Input tensor of shape (batch, seq_len, features)
            return_attention: Whether to return attention info

        Returns:
            encoded: Encoded representation (batch, output_dim)
            attention_info: Optional dict with attention weights per scale
        """
        batch_size, seq_len, _ = x.shape
        branch_outputs = []

        for branch in self.branches:
            # Input projection
            h = branch['input_proj'](x)  # (batch, seq_len, channels)

            # Conv (needs channel-first format)
            h = h.transpose(1, 2)
            h = branch['conv'](h)
            h = h.transpose(1, 2)

            # Global pooling
            h_pooled = h.mean(dim=1)  # (batch, channels)

            # Output projection
            h_out = branch['output_proj'](h_pooled)  # (batch, output_dim)
            branch_outputs.append(h_out)

        # Concatenate branch outputs
        fused = torch.cat(branch_outputs, dim=-1)  # (batch, num_scales * output_dim)

        # Fusion
        encoded = self.fusion(fused)  # (batch, output_dim)

        attention_info = None
        if return_attention:
            attention_info = {'branch_outputs': branch_outputs}

        return encoded, attention_info


if __name__ == "__main__":
    # Test the CNN encoder module
    print("Testing VBP CNN Encoder")
    print("=" * 60)

    # Test configuration
    config = CNNEncoderConfig(
        input_channels=1,
        conv_configs=[(32, 3), (64, 5), (128, 7)],
        attention_heads=4,
        dropout=0.3,
        use_attention=True,
        output_dim=128
    )

    batch_size = 8
    seq_len = 100
    input_dim = 32  # Number of features from feature engineering

    # Create sample input
    x = torch.randn(batch_size, seq_len, input_dim)

    print(f"\n1. Testing VolatilityCNN...")
    print(f"   Input shape: {x.shape}")

    model = VolatilityCNN(input_dim=input_dim, config=config)
    print(f"   Model parameters: {sum(p.numel() for p in model.parameters()):,}")

    # Forward pass
    encoded, attention = model(x, return_attention=True)
    print(f"   Output shape: {encoded.shape}")
    if attention is not None:
        print(f"   Attention shape: {attention.shape}")

    # Test feature maps
    print(f"\n2. Testing feature maps extraction...")
    feature_maps = model.get_feature_maps(x)
    for i, fm in enumerate(feature_maps):
        print(f"   Layer {i+1} feature map: {fm.shape}")

    print(f"\n3. Testing MultiScaleVolatilityCNN...")
    multi_scale_model = MultiScaleVolatilityCNN(
        input_dim=input_dim,
        config=config,
        scales=[1, 2, 4]
    )
    print(f"   Model parameters: {sum(p.numel() for p in multi_scale_model.parameters()):,}")

    encoded_ms, attn_info = multi_scale_model(x, return_attention=True)
    print(f"   Output shape: {encoded_ms.shape}")

    print(f"\n4. Testing gradient flow...")
    target = torch.randn(batch_size, config.output_dim)
    loss = F.mse_loss(encoded, target)
    loss.backward()

    # Check gradients
    grad_norms = []
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norms.append((name, param.grad.norm().item()))

    print(f"   Total parameters with gradients: {len(grad_norms)}")
    print(f"   Sample gradient norms:")
    for name, norm in grad_norms[:5]:
        print(f"     {name}: {norm:.6f}")

    print(f"\n5. Testing inference mode...")
    model.eval()
    with torch.no_grad():
        encoded_eval, _ = model(x)
    print(f"   Inference output shape: {encoded_eval.shape}")

    print("\n" + "=" * 60)
    print("All CNN encoder tests passed!")
