#!/usr/bin/env python3
"""
Price-Focused Attention Architecture
=====================================
Implements self-attention mechanism focused purely on price movements
(returns and derivatives) without any temporal bias.

This module provides a transformer-based architecture that:
1. Operates on sequences of returns (not raw prices)
2. Is completely time-agnostic (no real timestamps)
3. Uses only price variation information
4. Learns to attend to significant price movements

Key Features:
- Self-attention on return sequences
- Multi-head attention with 8 heads, d_model=256
- 4 transformer encoder layers
- Feed-forward dimension d_ff=1024
- Learnable positional encoding (sequence position only)

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

from typing import Tuple, List, Optional, Dict, Any
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from .multi_head_attention import MultiHeadAttention, create_causal_mask
from .positional_encoding import LearnablePositionalEncoding


@dataclass
class PriceAttentionConfig:
    """
    Configuration for Price-Focused Attention model.

    Default architecture:
    - d_model: 256 (model dimension)
    - n_heads: 8 (attention heads)
    - d_k: 32 (key/query dimension per head)
    - n_layers: 4 (transformer encoder layers)
    - d_ff: 1024 (feed-forward dimension)
    - dropout: 0.1
    - max_seq_len: 512
    """
    # Model dimensions
    d_model: int = 256
    n_heads: int = 8
    d_k: int = 32
    d_v: int = 32

    # Architecture
    n_layers: int = 4
    d_ff: int = 1024
    max_seq_len: int = 512

    # Regularization
    dropout: float = 0.1
    attention_dropout: float = 0.1

    # Input features
    input_features: int = 4  # Default: returns, log_returns, range_pct, velocity

    # Activation
    activation: str = "gelu"  # "relu" or "gelu"

    # Pre-layer norm (more stable training)
    pre_norm: bool = True


class FeedForwardBlock(nn.Module):
    """
    Position-wise Feed-Forward Network.

    Implements: FFN(x) = max(0, xW1 + b1)W2 + b2
    With optional GELU activation instead of ReLU.

    Args:
        d_model: Input and output dimension
        d_ff: Hidden layer dimension
        dropout: Dropout probability
        activation: Activation function ("relu" or "gelu")
    """

    def __init__(
        self,
        d_model: int = 256,
        d_ff: int = 1024,
        dropout: float = 0.1,
        activation: str = "gelu"
    ) -> None:
        super().__init__()

        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

        if activation == "gelu":
            self.activation = nn.GELU()
        elif activation == "relu":
            self.activation = nn.ReLU()
        else:
            raise ValueError(f"Unknown activation: {activation}")

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

        Args:
            x: Input tensor (batch, seq_len, d_model)

        Returns:
            Output tensor (batch, seq_len, d_model)
        """
        x = self.linear1(x)
        x = self.activation(x)
        x = self.dropout(x)
        x = self.linear2(x)
        x = self.dropout(x)
        return x


class PriceAttentionEncoderLayer(nn.Module):
    """
    Single layer of the Price-Focused Attention encoder.

    Architecture (pre-norm):
        x -> LayerNorm -> MultiHeadAttention -> Dropout -> + x
        x -> LayerNorm -> FeedForward -> Dropout -> + x

    Args:
        config: PriceAttentionConfig object
    """

    def __init__(self, config: PriceAttentionConfig) -> None:
        super().__init__()

        self.config = config

        # Layer normalization
        self.norm1 = nn.LayerNorm(config.d_model)
        self.norm2 = nn.LayerNorm(config.d_model)

        # Self-attention
        self.self_attention = MultiHeadAttention(
            d_model=config.d_model,
            n_heads=config.n_heads,
            d_k=config.d_k,
            d_v=config.d_v,
            dropout=config.attention_dropout
        )

        # Feed-forward
        self.feed_forward = FeedForwardBlock(
            d_model=config.d_model,
            d_ff=config.d_ff,
            dropout=config.dropout,
            activation=config.activation
        )

        # Dropout for residual connections
        self.dropout = nn.Dropout(config.dropout)

        self.pre_norm = config.pre_norm

    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        return_attention: bool = True
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        """
        Forward pass of encoder layer.

        Args:
            x: Input tensor (batch, seq_len, d_model)
            mask: Optional attention mask
            return_attention: Whether to return attention weights

        Returns:
            output: Output tensor (batch, seq_len, d_model)
            attention_weights: Attention weights if return_attention is True
        """
        if self.pre_norm:
            # Pre-norm: LayerNorm before attention/ff
            # Self-attention block
            residual = x
            x = self.norm1(x)
            attn_output, attention_weights = self.self_attention(
                x, x, x, mask=mask, return_attention=return_attention
            )
            x = residual + self.dropout(attn_output)

            # Feed-forward block
            residual = x
            x = self.norm2(x)
            ff_output = self.feed_forward(x)
            x = residual + ff_output
        else:
            # Post-norm: LayerNorm after attention/ff
            # Self-attention block
            attn_output, attention_weights = self.self_attention(
                x, x, x, mask=mask, return_attention=return_attention
            )
            x = self.norm1(x + self.dropout(attn_output))

            # Feed-forward block
            ff_output = self.feed_forward(x)
            x = self.norm2(x + ff_output)

        return x, attention_weights


class PriceFocusedAttention(nn.Module):
    """
    Price-Focused Self-Attention model.

    This model implements a transformer encoder that operates on sequences
    of price returns and derivatives. It is designed to be completely
    time-agnostic, focusing only on price variation patterns.

    The model does NOT use:
    - Real timestamps
    - Session indicators
    - Time-of-day features
    - Day-of-week features

    The model ONLY uses:
    - Returns (price changes)
    - Log returns
    - Range percentages
    - Price velocity/momentum
    - Other price-derived features

    Args:
        config: PriceAttentionConfig object (optional, uses defaults if None)
        input_features: Number of input features per timestep
                       Overrides config.input_features if provided

    Example:
        >>> config = PriceAttentionConfig(d_model=256, n_heads=8, n_layers=4)
        >>> model = PriceFocusedAttention(config, input_features=4)
        >>> x = torch.randn(32, 100, 4)  # (batch, seq_len, features)
        >>> output, attentions = model(x)
        >>> print(output.shape)  # (32, 100, 256)
    """

    def __init__(
        self,
        config: Optional[PriceAttentionConfig] = None,
        input_features: Optional[int] = None
    ) -> None:
        super().__init__()

        self.config = config or PriceAttentionConfig()

        # Override input features if provided
        if input_features is not None:
            self.config.input_features = input_features

        # Input projection: features -> d_model
        self.input_projection = nn.Linear(
            self.config.input_features,
            self.config.d_model
        )

        # Learnable positional encoding (time-agnostic)
        self.positional_encoding = LearnablePositionalEncoding(
            d_model=self.config.d_model,
            max_seq_len=self.config.max_seq_len,
            dropout=self.config.dropout
        )

        # Encoder layers
        self.encoder_layers = nn.ModuleList([
            PriceAttentionEncoderLayer(self.config)
            for _ in range(self.config.n_layers)
        ])

        # Final layer normalization (for pre-norm architecture)
        if self.config.pre_norm:
            self.final_norm = nn.LayerNorm(self.config.d_model)
        else:
            self.final_norm = nn.Identity()

        # Initialize weights
        self._init_weights()

    def _init_weights(self) -> None:
        """Initialize model weights using Xavier initialization."""
        for name, param in self.named_parameters():
            if 'weight' in name and param.dim() > 1:
                nn.init.xavier_uniform_(param)
            elif 'bias' in name:
                nn.init.zeros_(param)

    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
        return_all_attentions: bool = False
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """
        Forward pass of the Price-Focused Attention model.

        Args:
            x: Input tensor of shape (batch, seq_len, input_features)
               Contains return-based features (NO timestamps!)
            mask: Optional attention mask of shape (batch, 1, seq_len, seq_len)
                  or (batch, 1, 1, seq_len) for causal masking
            return_all_attentions: If True, returns attention from all layers

        Returns:
            output: Encoded output of shape (batch, seq_len, d_model)
            attentions: List of attention weight tensors, one per layer
                       Each has shape (batch, n_heads, seq_len, seq_len)
        """
        batch_size, seq_len, _ = x.shape

        # Project input features to d_model
        x = self.input_projection(x)

        # Add positional encoding (sequence position only, no timestamps!)
        x = self.positional_encoding(x)

        # Pass through encoder layers
        all_attentions = []
        for layer in self.encoder_layers:
            x, attention_weights = layer(
                x, mask=mask, return_attention=True
            )
            if return_all_attentions and attention_weights is not None:
                all_attentions.append(attention_weights)

        # Final layer normalization
        x = self.final_norm(x)

        # If not returning all attentions, just return the last layer's
        if not return_all_attentions and all_attentions:
            all_attentions = [all_attentions[-1]]

        return x, all_attentions

    def get_attention_scores(
        self,
        x: torch.Tensor,
        layer_idx: int = -1
    ) -> torch.Tensor:
        """
        Get attention scores for a specific layer.

        Args:
            x: Input tensor (batch, seq_len, input_features)
            layer_idx: Layer index to extract attention from
                      -1 means last layer (default)

        Returns:
            Attention scores of shape (batch, n_heads, seq_len, seq_len)
        """
        _, all_attentions = self.forward(x, return_all_attentions=True)

        if not all_attentions:
            raise ValueError("No attention scores available")

        return all_attentions[layer_idx]

    def encode_sequence(
        self,
        x: torch.Tensor,
        pooling: str = "last"
    ) -> torch.Tensor:
        """
        Encode a sequence to a fixed-size representation.

        Args:
            x: Input tensor (batch, seq_len, input_features)
            pooling: Pooling method - "last", "first", "mean", or "max"

        Returns:
            Encoded representation of shape (batch, d_model)
        """
        output, _ = self.forward(x, return_all_attentions=False)

        if pooling == "last":
            return output[:, -1, :]
        elif pooling == "first":
            return output[:, 0, :]
        elif pooling == "mean":
            return output.mean(dim=1)
        elif pooling == "max":
            return output.max(dim=1)[0]
        else:
            raise ValueError(f"Unknown pooling method: {pooling}")


class PriceAttentionEncoder(nn.Module):
    """
    Complete Price-Focused Attention Encoder with prediction heads.

    This is a higher-level wrapper that adds task-specific output heads
    to the base PriceFocusedAttention model.

    Supports multiple output modes:
    - "sequence": Full sequence output (batch, seq_len, d_model)
    - "last": Last position output (batch, d_model)
    - "mean": Mean pooled output (batch, d_model)

    Args:
        config: PriceAttentionConfig object
        input_features: Number of input features
        output_dim: Output dimension (default: d_model)
        output_mode: Output pooling mode
    """

    def __init__(
        self,
        config: Optional[PriceAttentionConfig] = None,
        input_features: int = 4,
        output_dim: Optional[int] = None,
        output_mode: str = "last"
    ) -> None:
        super().__init__()

        self.config = config or PriceAttentionConfig()
        self.output_mode = output_mode

        # Base attention model
        self.attention = PriceFocusedAttention(
            config=self.config,
            input_features=input_features
        )

        # Output projection if needed
        if output_dim is not None and output_dim != self.config.d_model:
            self.output_projection = nn.Linear(self.config.d_model, output_dim)
        else:
            self.output_projection = nn.Identity()
            output_dim = self.config.d_model

        self.output_dim = output_dim

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

        Args:
            x: Input tensor (batch, seq_len, input_features)
            mask: Optional attention mask

        Returns:
            output: Encoded output, shape depends on output_mode
            attentions: List of attention weight tensors
        """
        # Get full sequence encoding
        encoded, attentions = self.attention(x, mask=mask, return_all_attentions=True)

        # Apply pooling based on output_mode
        if self.output_mode == "sequence":
            output = encoded
        elif self.output_mode == "last":
            output = encoded[:, -1, :]
        elif self.output_mode == "first":
            output = encoded[:, 0, :]
        elif self.output_mode == "mean":
            output = encoded.mean(dim=1)
        elif self.output_mode == "max":
            output = encoded.max(dim=1)[0]
        else:
            raise ValueError(f"Unknown output_mode: {self.output_mode}")

        # Apply output projection
        output = self.output_projection(output)

        return output, attentions


def compute_return_features(
    prices: torch.Tensor,
    epsilon: float = 1e-8
) -> torch.Tensor:
    """
    Compute return-based features from price sequence.

    This function converts raw OHLC prices to return-based features
    suitable for the Price-Focused Attention model.

    Args:
        prices: Price tensor of shape (batch, seq_len, 4)
                Expected order: [open, high, low, close]
        epsilon: Small value for numerical stability

    Returns:
        Features tensor of shape (batch, seq_len, 4):
        - returns: (close_t - close_{t-1}) / close_{t-1}
        - log_returns: log(close_t / close_{t-1})
        - range_pct: (high - low) / close
        - velocity: close - open / close
    """
    batch_size, seq_len, _ = prices.shape

    open_price = prices[:, :, 0]
    high = prices[:, :, 1]
    low = prices[:, :, 2]
    close = prices[:, :, 3]

    # Simple returns
    returns = torch.zeros_like(close)
    returns[:, 1:] = (close[:, 1:] - close[:, :-1]) / (close[:, :-1] + epsilon)

    # Log returns
    log_returns = torch.zeros_like(close)
    log_returns[:, 1:] = torch.log((close[:, 1:] + epsilon) / (close[:, :-1] + epsilon))

    # Range percentage
    range_pct = (high - low) / (close + epsilon)

    # Velocity (intra-candle movement)
    velocity = (close - open_price) / (close + epsilon)

    # Stack features
    features = torch.stack([returns, log_returns, range_pct, velocity], dim=-1)

    return features
