#!/usr/bin/env python3
"""
Volatility-Biased Self-Attention Module
========================================
Implements attention mechanism that weights timesteps based on volatility/movement.

Key Features:
1. compute_factor_median_range: Rolling median of range with shift(1) for no leakage
2. compute_move_multiplier: delta / factor ratio
3. weight_step: Discrete weight mapping (0, 1, 2, 3)
4. weight_smooth: Softplus smooth weight mapping
5. VolatilityBiasedSelfAttention: PyTorch attention with volatility bias

Theory:
- m < 1 = noise (attention ~0)
- m ~ 2 = 2x normal movement (attention ~2)
- m ~ 3 = 3x normal movement (attention ~3)

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

import numpy as np
import pandas as pd
from typing import Dict, Tuple, Optional, Union
from dataclasses import dataclass
from loguru import logger

# PyTorch imports
try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False
    logger.warning("PyTorch not available - attention models will not work")


@dataclass
class VolatilityAttentionConfig:
    """Configuration for volatility-based attention"""

    # Rolling window for median factor calculation
    factor_window: int = 200

    # Minimum periods for rolling calculation
    min_periods: Optional[int] = None

    # Maximum attention weight
    w_max: float = 3.0

    # Softplus beta (controls transition sharpness)
    beta: float = 4.0

    # Gamma for attention bias scaling
    gamma: float = 1.0

    # Epsilon for numerical stability
    epsilon: float = 1e-12

    # Use step (discrete) or smooth (softplus) weighting
    use_smooth_weights: bool = True

    # Model dimensions (for transformer)
    d_model: int = 64
    n_heads: int = 4
    dropout: float = 0.1

    def __post_init__(self):
        if self.min_periods is None:
            self.min_periods = self.factor_window // 2


# ==============================================================================
# Core Volatility Factor Functions
# ==============================================================================

def compute_factor_median_range(
    df: pd.DataFrame,
    window: int = 200,
    min_periods: Optional[int] = None
) -> pd.Series:
    """
    Compute dynamic factor as rolling median of range with shift(1).

    Uses shift(1) to avoid data leakage - only uses information
    that would have been available at prediction time.

    Args:
        df: DataFrame with High/Low columns (or high/low)
        window: Rolling window size
        min_periods: Minimum periods for rolling (default: window//2)

    Returns:
        Series with dynamic factor for each row
    """
    min_periods = min_periods or window // 2

    # Handle column name variations
    high_col = 'High' if 'High' in df.columns else 'high'
    low_col = 'Low' if 'Low' in df.columns else 'low'

    # Calculate candle range: Delta_t = High_t - Low_t
    rng = (df[high_col] - df[low_col]).abs()

    # Rolling median with shift(1) to avoid leakage
    # f_t = rolling_median(Delta, window).shift(1)
    f = rng.rolling(window=window, min_periods=min_periods).median().shift(1)

    return f


def compute_move_multiplier(
    df: pd.DataFrame,
    factor: pd.Series = None,
    window: int = 200,
    epsilon: float = 1e-12
) -> pd.Series:
    """
    Compute movement multiplier: actual_range / factor.

    Interpretation:
    - m < 1: Movement smaller than typical (noise)
    - m = 1: Typical movement
    - m > 1: Larger than typical movement (signal)
    - m ~ 2: "10 vs 5" scenario
    - m ~ 3: "15 vs 5" scenario

    Args:
        df: DataFrame with High/Low columns
        factor: Pre-computed factor (computed if None)
        window: Rolling window for factor computation
        epsilon: Small value to avoid division by zero

    Returns:
        Series with move multiplier for each row
    """
    if factor is None:
        factor = compute_factor_median_range(df, window=window)

    # Handle column name variations
    high_col = 'High' if 'High' in df.columns else 'high'
    low_col = 'Low' if 'Low' in df.columns else 'low'

    # Calculate actual range
    delta = (df[high_col] - df[low_col]).abs()

    # Compute multiplier: m_t = delta_t / f_t
    m = delta / (factor + epsilon)

    return m


def weight_step(
    m: Union[np.ndarray, pd.Series],
    w_max: int = 3
) -> np.ndarray:
    """
    Discrete step-based weight mapping.

    Maps multiplier to discrete attention weights:
    - m < 1: weight = 0 (noise, ignore)
    - 1 <= m < 2: weight = 1 (normal)
    - 2 <= m < 3: weight = 2 (attention)
    - m >= 3: weight = 3 (high attention)

    Args:
        m: Move multiplier array/series
        w_max: Maximum weight cap

    Returns:
        Array of discrete attention weights (0, 1, 2, 3)
    """
    if isinstance(m, pd.Series):
        m = m.values

    # Floor the multiplier and clip
    w = np.floor(m).clip(0, w_max)

    # Set weight to 0 if m < 1 (noise)
    w[m < 1.0] = 0.0

    return w


def weight_smooth(
    m: Union[np.ndarray, pd.Series],
    w_max: float = 3.0,
    beta: float = 4.0
) -> np.ndarray:
    """
    Smooth softplus-based weight mapping.

    Formula: w = log1p(exp(beta * (m - 1))) / beta

    This creates a smooth transition:
    - m < 1: w approaches 0 (ignore noise)
    - m = 1: w ~ 0 (typical movement, neutral)
    - m > 1: w approaches (m - 1) linearly (attention)

    This is more stable for training as it avoids discontinuities.

    Args:
        m: Move multiplier array/series
        w_max: Maximum weight cap
        beta: Softplus beta parameter (sharpness)

    Returns:
        Array of smooth attention weights
    """
    if isinstance(m, pd.Series):
        m = m.values

    # Apply softplus: log1p(exp(x)) is numerically stable
    x = beta * (m - 1.0)

    # Numerical stability for large x
    w = np.where(
        x > 20,  # For large x, softplus(x) ~ x
        x / beta,
        np.log1p(np.exp(x)) / beta
    )

    # Clip to [0, w_max]
    w = np.clip(w, 0.0, w_max)

    return w


def compute_attention_weights(
    df: pd.DataFrame,
    config: VolatilityAttentionConfig = None,
    normalize: bool = True
) -> np.ndarray:
    """
    Full pipeline to compute volatility-based attention weights.

    Args:
        df: DataFrame with OHLCV data
        config: Configuration object
        normalize: Whether to normalize weights to mean=1

    Returns:
        Array of attention weights
    """
    config = config or VolatilityAttentionConfig()

    # Step 1: Compute factor
    factor = compute_factor_median_range(
        df,
        window=config.factor_window,
        min_periods=config.min_periods
    )

    # Step 2: Compute multiplier
    multiplier = compute_move_multiplier(df, factor, epsilon=config.epsilon)

    # Step 3: Apply weight mapping
    if config.use_smooth_weights:
        weights = weight_smooth(multiplier, w_max=config.w_max, beta=config.beta)
    else:
        weights = weight_step(multiplier, w_max=int(config.w_max))

    # Handle NaN
    nan_mask = np.isnan(weights) | np.isnan(multiplier.values if hasattr(multiplier, 'values') else multiplier)
    weights[nan_mask] = 1.0

    # Normalize
    if normalize:
        valid_mask = ~nan_mask
        if valid_mask.sum() > 0 and weights[valid_mask].mean() > 0:
            weights[valid_mask] = weights[valid_mask] / weights[valid_mask].mean()

    return weights


# ==============================================================================
# PyTorch Attention Module
# ==============================================================================

if HAS_TORCH:

    def build_attention_bias_from_weight(
        w: torch.Tensor,
        gamma: float = 1.0,
        eps: float = 1e-3
    ) -> torch.Tensor:
        """
        Convert volatility weights to attention bias.

        Transforms weights to log-space for additive bias:
        - w=0 -> log(eps) very negative (low attention)
        - w large -> log(w) positive (high attention)

        Args:
            w: (B, T) weights >= 0
            gamma: Scaling factor for bias
            eps: Small value for log stability

        Returns:
            bias: (B, 1, 1, T) for adding to attention scores
        """
        # Convert to log-space
        b = gamma * torch.log(w + eps)

        # Reshape for broadcasting: (B, T) -> (B, 1, 1, T)
        return b[:, None, None, :]


    class VolatilityBiasedSelfAttention(nn.Module):
        """
        Self-Attention with volatility-based bias.

        Standard self-attention with an additional bias term based on volatility:

        scores_{i,j} = (Q_i @ K_j^T) / sqrt(d) + bias_j

        Where bias_j increases for timesteps with high volatility/movement,
        causing the model to "pay more attention" to significant moves.

        Args:
            d_model: Model dimension
            n_heads: Number of attention heads
            dropout: Dropout probability
        """

        def __init__(
            self,
            d_model: int = 64,
            n_heads: int = 4,
            dropout: float = 0.1
        ):
            super().__init__()

            assert d_model % n_heads == 0, "d_model must be divisible by n_heads"

            self.d_model = d_model
            self.n_heads = n_heads
            self.d_head = d_model // n_heads

            # Q, K, V projections
            self.qkv = nn.Linear(d_model, 3 * d_model)
            self.out = nn.Linear(d_model, d_model)
            self.dropout = nn.Dropout(dropout)

            # Learnable gamma for bias scaling
            self.gamma = nn.Parameter(torch.ones(1))

        def forward(
            self,
            x: torch.Tensor,
            attn_weight: torch.Tensor,
            mask: torch.Tensor = None
        ) -> Tuple[torch.Tensor, torch.Tensor]:
            """
            Forward pass with volatility-biased attention.

            Args:
                x: (B, T, D) input sequence
                attn_weight: (B, T) volatility weights per timestep
                mask: Optional attention mask

            Returns:
                output: (B, T, D) attended output
                attn: (B, H, T, T) attention weights
            """
            B, T, D = x.shape

            # Compute Q, K, V
            qkv = self.qkv(x)  # (B, T, 3D)
            q, k, v = qkv.chunk(3, dim=-1)

            # Reshape to multi-head: (B, T, D) -> (B, H, T, d)
            q = q.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
            k = k.view(B, T, self.n_heads, self.d_head).transpose(1, 2)
            v = v.view(B, T, self.n_heads, self.d_head).transpose(1, 2)

            # Compute attention scores: (B, H, T, T)
            scores = (q @ k.transpose(-2, -1)) / (self.d_head ** 0.5)

            # Add volatility bias
            bias = build_attention_bias_from_weight(
                attn_weight,
                gamma=self.gamma.item()
            )
            scores = scores + bias

            # Apply mask if provided
            if mask is not None:
                scores = scores.masked_fill(mask == 0, float('-inf'))

            # Softmax and dropout
            attn = F.softmax(scores, dim=-1)
            attn = self.dropout(attn)

            # Apply attention to values
            y = attn @ v  # (B, H, T, d)

            # Reshape back: (B, H, T, d) -> (B, T, D)
            y = y.transpose(1, 2).contiguous().view(B, T, D)

            return self.out(y), attn


    class VolatilityAttentionBlock(nn.Module):
        """
        Full attention block with volatility bias, feed-forward, and residuals.

        Architecture:
        1. LayerNorm + VolatilityBiasedSelfAttention + Residual
        2. LayerNorm + FeedForward + Residual
        """

        def __init__(
            self,
            d_model: int = 64,
            n_heads: int = 4,
            d_ff: int = 256,
            dropout: float = 0.1
        ):
            super().__init__()

            self.ln1 = nn.LayerNorm(d_model)
            self.attn = VolatilityBiasedSelfAttention(d_model, n_heads, dropout)

            self.ln2 = nn.LayerNorm(d_model)
            self.ff = 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,
            attn_weight: torch.Tensor,
            mask: torch.Tensor = None
        ) -> Tuple[torch.Tensor, torch.Tensor]:
            """
            Forward pass.

            Args:
                x: (B, T, D) input
                attn_weight: (B, T) volatility weights
                mask: Optional mask

            Returns:
                output: (B, T, D)
                attn: (B, H, T, T) attention weights
            """
            # Self-attention with residual
            h, attn = self.attn(self.ln1(x), attn_weight, mask)
            x = x + h

            # Feed-forward with residual
            x = x + self.ff(self.ln2(x))

            return x, attn


    class VolatilityTransformerEncoder(nn.Module):
        """
        Transformer encoder stack with volatility-biased attention.

        Suitable for sequence-to-sequence prediction tasks where
        we want the model to focus more on high-volatility timesteps.
        """

        def __init__(
            self,
            input_dim: int,
            d_model: int = 64,
            n_heads: int = 4,
            n_layers: int = 2,
            d_ff: int = 256,
            dropout: float = 0.1,
            max_seq_len: int = 512
        ):
            super().__init__()

            self.input_projection = nn.Linear(input_dim, d_model)
            self.pos_encoding = nn.Parameter(
                self._get_positional_encoding(max_seq_len, d_model)
            )

            self.layers = nn.ModuleList([
                VolatilityAttentionBlock(d_model, n_heads, d_ff, dropout)
                for _ in range(n_layers)
            ])

            self.ln = nn.LayerNorm(d_model)

        def _get_positional_encoding(
            self,
            max_len: int,
            d_model: int
        ) -> torch.Tensor:
            """Generate sinusoidal positional encoding."""
            pe = torch.zeros(max_len, d_model)
            position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
            div_term = torch.exp(
                torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model)
            )
            pe[:, 0::2] = torch.sin(position * div_term)
            pe[:, 1::2] = torch.cos(position * div_term)
            return pe.unsqueeze(0)  # (1, max_len, d_model)

        def forward(
            self,
            x: torch.Tensor,
            attn_weight: torch.Tensor,
            mask: torch.Tensor = None
        ) -> Tuple[torch.Tensor, list]:
            """
            Forward pass.

            Args:
                x: (B, T, input_dim) input features
                attn_weight: (B, T) volatility weights
                mask: Optional mask

            Returns:
                output: (B, T, d_model)
                attentions: List of attention matrices
            """
            B, T, _ = x.shape

            # Project to d_model and add positional encoding
            x = self.input_projection(x) + self.pos_encoding[:, :T, :]

            attentions = []
            for layer in self.layers:
                x, attn = layer(x, attn_weight, mask)
                attentions.append(attn)

            return self.ln(x), attentions


    class VolatilityRangePredictor(nn.Module):
        """
        Complete model for range prediction using volatility-biased attention.

        Takes OHLCV features and predicts delta_high and delta_low.
        Uses attention mechanism that focuses on high-volatility timesteps.
        """

        def __init__(
            self,
            input_dim: int = 14,
            d_model: int = 64,
            n_heads: int = 4,
            n_layers: int = 2,
            dropout: float = 0.1
        ):
            super().__init__()

            self.encoder = VolatilityTransformerEncoder(
                input_dim=input_dim,
                d_model=d_model,
                n_heads=n_heads,
                n_layers=n_layers,
                dropout=dropout
            )

            # Output heads for high and low predictions
            self.head_high = nn.Sequential(
                nn.Linear(d_model, d_model),
                nn.GELU(),
                nn.Dropout(dropout),
                nn.Linear(d_model, 1)
            )

            self.head_low = nn.Sequential(
                nn.Linear(d_model, d_model),
                nn.GELU(),
                nn.Dropout(dropout),
                nn.Linear(d_model, 1)
            )

        def forward(
            self,
            x: torch.Tensor,
            attn_weight: torch.Tensor,
            mask: torch.Tensor = None
        ) -> Dict[str, torch.Tensor]:
            """
            Forward pass.

            Args:
                x: (B, T, input_dim) input features
                attn_weight: (B, T) volatility weights
                mask: Optional mask

            Returns:
                dict with 'high', 'low' predictions and 'attentions'
            """
            # Encode with volatility-biased attention
            encoded, attentions = self.encoder(x, attn_weight, mask)

            # Use last timestep for prediction (or could use pooling)
            last_hidden = encoded[:, -1, :]  # (B, d_model)

            # Predict high and low deltas
            pred_high = self.head_high(last_hidden).squeeze(-1)  # (B,)
            pred_low = self.head_low(last_hidden).squeeze(-1)    # (B,)

            return {
                'high': pred_high,
                'low': pred_low,
                'attentions': attentions
            }


# ==============================================================================
# Test Module
# ==============================================================================

if __name__ == "__main__":
    print("Testing Volatility Attention Module")
    print("=" * 60)

    # Create sample OHLCV data
    np.random.seed(42)
    n = 1000

    dates = pd.date_range('2025-01-01', periods=n, freq='5min')
    price = 2650 + np.cumsum(np.random.randn(n) * 0.5)

    # Simulate varying volatility
    volatility = np.where(
        (dates.hour >= 13) & (dates.hour < 16),  # London/NY overlap
        5.0,  # High volatility
        2.0   # Normal volatility
    )

    df = pd.DataFrame({
        'open': price,
        'high': price + np.abs(np.random.randn(n)) * volatility,
        'low': price - np.abs(np.random.randn(n)) * volatility,
        'close': price + np.random.randn(n) * 0.5,
        'volume': np.random.randint(100, 1000, n)
    }, index=dates)

    # Test factor computation
    print("\n1. Testing compute_factor_median_range...")
    factor = compute_factor_median_range(df, window=100)
    print(f"   Factor mean: {factor.mean():.4f}")
    print(f"   Factor std: {factor.std():.4f}")

    # Test multiplier
    print("\n2. Testing compute_move_multiplier...")
    multiplier = compute_move_multiplier(df, factor)
    print(f"   Multiplier mean: {multiplier.mean():.2f}")
    print(f"   Multiplier > 2x: {(multiplier > 2).sum()} samples")
    print(f"   Multiplier > 3x: {(multiplier > 3).sum()} samples")

    # Test step weights
    print("\n3. Testing weight_step...")
    w_step = weight_step(multiplier)
    print(f"   Weight distribution:")
    print(f"   - w=0: {(w_step == 0).sum()}")
    print(f"   - w=1: {(w_step == 1).sum()}")
    print(f"   - w=2: {(w_step == 2).sum()}")
    print(f"   - w=3: {(w_step == 3).sum()}")

    # Test smooth weights
    print("\n4. Testing weight_smooth...")
    w_smooth = weight_smooth(multiplier, beta=4.0)
    print(f"   Smooth weight range: [{w_smooth.min():.3f}, {w_smooth.max():.3f}]")
    print(f"   Smooth weight mean: {np.nanmean(w_smooth):.3f}")

    # Test full pipeline
    print("\n5. Testing compute_attention_weights...")
    config = VolatilityAttentionConfig(
        factor_window=100,
        w_max=3.0,
        beta=4.0,
        use_smooth_weights=True
    )
    weights = compute_attention_weights(df, config)
    print(f"   Final weights mean: {weights.mean():.3f}")
    print(f"   High attention (w>1.5): {(weights > 1.5).sum()} samples")

    # Test PyTorch module if available
    if HAS_TORCH:
        print("\n6. Testing VolatilityBiasedSelfAttention (PyTorch)...")

        B, T, D = 4, 50, 64
        x = torch.randn(B, T, D)
        attn_w = torch.from_numpy(weights[:T]).float().unsqueeze(0).expand(B, -1)

        model = VolatilityBiasedSelfAttention(d_model=D, n_heads=4)
        output, attn = model(x, attn_w)

        print(f"   Input shape: {x.shape}")
        print(f"   Output shape: {output.shape}")
        print(f"   Attention shape: {attn.shape}")

        print("\n7. Testing VolatilityRangePredictor...")

        input_dim = 14
        x = torch.randn(B, T, input_dim)

        predictor = VolatilityRangePredictor(
            input_dim=input_dim,
            d_model=64,
            n_heads=4,
            n_layers=2
        )

        result = predictor(x, attn_w)
        print(f"   High prediction shape: {result['high'].shape}")
        print(f"   Low prediction shape: {result['low'].shape}")
        print(f"   Number of attention layers: {len(result['attentions'])}")

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