#!/usr/bin/env python3
"""
Hierarchical Attention Network for Multi-Timeframe Synthesis
=============================================================
Implements a hierarchical attention mechanism that processes features
from multiple timeframes and learns to weight their contributions.

Architecture:
1. Per-TF Attention: Self-attention within each timeframe
2. Cross-TF Attention: Attention across timeframes
3. TF Weighting: Learned weights per timeframe
4. Unified Representation: Final merged representation

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

import math
from typing import Dict, List, Optional, Tuple, Any

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


@dataclass
class HierarchicalAttentionConfig:
    """Configuration for Hierarchical Attention Network."""
    d_model: int = 128
    n_heads: int = 4
    d_k: int = 32
    d_v: int = 32
    d_ff: int = 256
    n_layers: int = 2
    dropout: float = 0.1
    max_seq_len: int = 256
    timeframes: Tuple[str, ...] = ('5m', '15m', '1h', '4h')
    use_cross_tf_attention: bool = True
    use_learnable_tf_weights: bool = True


class ScaledDotProductAttention(nn.Module):
    """Scaled dot-product attention mechanism."""

    def __init__(self, d_k: int, dropout: float = 0.1):
        super().__init__()
        self.scale = math.sqrt(d_k)
        self.dropout = nn.Dropout(dropout)

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Compute scaled dot-product attention.

        Args:
            q: Query (batch, n_heads, seq_len_q, d_k)
            k: Key (batch, n_heads, seq_len_k, d_k)
            v: Value (batch, n_heads, seq_len_k, d_v)
            mask: Optional attention mask

        Returns:
            output: Attended values
            attention_weights: Attention distribution
        """
        scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale

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

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

        output = torch.matmul(attention_weights, v)

        return output, attention_weights


class MultiHeadAttentionTF(nn.Module):
    """Multi-head attention for a single timeframe."""

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

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_k
        self.d_v = d_v

        # Linear projections
        self.w_q = nn.Linear(d_model, n_heads * d_k, bias=False)
        self.w_k = nn.Linear(d_model, n_heads * d_k, bias=False)
        self.w_v = nn.Linear(d_model, n_heads * d_v, bias=False)
        self.w_o = nn.Linear(n_heads * d_v, d_model, bias=False)

        self.attention = ScaledDotProductAttention(d_k, dropout)
        self.dropout = nn.Dropout(dropout)

        self._init_weights()

    def _init_weights(self):
        """Initialize weights with Xavier uniform."""
        for module in [self.w_q, self.w_k, self.w_v, self.w_o]:
            nn.init.xavier_uniform_(module.weight)

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

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

        Returns:
            output: Attended output (batch, seq_len, d_model)
            attention_weights: Attention weights (batch, n_heads, seq_len, seq_len)
        """
        batch_size, seq_len, _ = x.shape

        # Project to Q, K, V
        q = self.w_q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        k = self.w_k(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        v = self.w_v(x).view(batch_size, seq_len, self.n_heads, self.d_v).transpose(1, 2)

        # Apply attention
        attended, attention_weights = self.attention(q, k, v, mask)

        # Concatenate heads and project
        attended = attended.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
        output = self.w_o(attended)
        output = self.dropout(output)

        return output, attention_weights


class CrossTimeframeAttention(nn.Module):
    """
    Cross-timeframe attention module.

    Allows each timeframe to attend to features from other timeframes,
    learning dependencies and relationships between different time scales.
    """

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

        self.d_model = d_model
        self.n_heads = n_heads
        self.n_timeframes = n_timeframes
        self.d_k = d_model // n_heads

        # Projections for cross-attention
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

        self.scale = math.sqrt(self.d_k)
        self.dropout = nn.Dropout(dropout)

        # Layer norm for each timeframe
        self.layer_norms = nn.ModuleList([
            nn.LayerNorm(d_model) for _ in range(n_timeframes)
        ])

        self._init_weights()

    def _init_weights(self):
        """Initialize weights."""
        for module in [self.q_proj, self.k_proj, self.v_proj, self.out_proj]:
            nn.init.xavier_uniform_(module.weight)
            if module.bias is not None:
                nn.init.zeros_(module.bias)

    def forward(
        self,
        tf_representations: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        """
        Apply cross-timeframe attention.

        Args:
            tf_representations: Dictionary of {timeframe: tensor}
                Each tensor has shape (batch, d_model)

        Returns:
            Updated representations with cross-TF information
        """
        tf_names = list(tf_representations.keys())
        batch_size = list(tf_representations.values())[0].shape[0]

        # Stack all TF representations: (batch, n_tf, d_model)
        stacked = torch.stack([tf_representations[tf] for tf in tf_names], dim=1)

        # Project to Q, K, V
        q = self.q_proj(stacked)  # (batch, n_tf, d_model)
        k = self.k_proj(stacked)
        v = self.v_proj(stacked)

        # Reshape for multi-head attention
        q = q.view(batch_size, len(tf_names), self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(batch_size, len(tf_names), self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(batch_size, len(tf_names), self.n_heads, self.d_k).transpose(1, 2)

        # Compute attention scores
        scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale
        attention_weights = F.softmax(scores, dim=-1)
        attention_weights = self.dropout(attention_weights)

        # Apply attention
        attended = torch.matmul(attention_weights, v)

        # Reshape back
        attended = attended.transpose(1, 2).contiguous().view(batch_size, len(tf_names), -1)

        # Output projection
        output = self.out_proj(attended)

        # Add residual and apply layer norm
        result = {}
        for i, tf in enumerate(tf_names):
            residual = tf_representations[tf]
            updated = self.layer_norms[i](residual + output[:, i])
            result[tf] = updated

        return result


class TimeframeEncoder(nn.Module):
    """
    Encoder for a single timeframe.

    Processes sequence of features and produces a fixed-size representation.
    """

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

        self.d_model = d_model

        # Input projection
        self.input_proj = nn.Linear(input_dim, d_model)

        # Positional encoding
        self.pos_encoding = nn.Parameter(torch.randn(1, max_seq_len, d_model) * 0.02)

        # Transformer encoder layers
        self.layers = nn.ModuleList()
        for _ in range(n_layers):
            self.layers.append(nn.ModuleDict({
                'attention': MultiHeadAttentionTF(d_model, n_heads, d_model // n_heads, d_model // n_heads, dropout),
                'norm1': nn.LayerNorm(d_model),
                'ff': nn.Sequential(
                    nn.Linear(d_model, d_ff),
                    nn.GELU(),
                    nn.Dropout(dropout),
                    nn.Linear(d_ff, d_model),
                    nn.Dropout(dropout)
                ),
                'norm2': nn.LayerNorm(d_model)
            }))

        # Global pooling
        self.global_pool = nn.Sequential(
            nn.AdaptiveAvgPool1d(1),
        )

        # Output projection
        self.output_proj = nn.Linear(d_model, d_model)

    def forward(
        self,
        x: torch.Tensor,
        mask: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Encode timeframe features.

        Args:
            x: Input features (batch, seq_len, input_dim)
            mask: Optional attention mask

        Returns:
            representation: Fixed-size representation (batch, d_model)
            sequence_output: Full sequence output (batch, seq_len, d_model)
        """
        batch_size, seq_len, _ = x.shape

        # Project input
        x = self.input_proj(x)

        # Add positional encoding
        x = x + self.pos_encoding[:, :seq_len, :]

        # Apply transformer layers
        for layer in self.layers:
            # Self-attention with residual
            attended, _ = layer['attention'](x, mask)
            x = layer['norm1'](x + attended)

            # Feed-forward with residual
            ff_out = layer['ff'](x)
            x = layer['norm2'](x + ff_out)

        sequence_output = x

        # Global pooling for representation
        pooled = self.global_pool(x.transpose(1, 2)).squeeze(-1)  # (batch, d_model)

        # Output projection
        representation = self.output_proj(pooled)

        return representation, sequence_output


class HierarchicalAttention(nn.Module):
    """
    Hierarchical Attention Network for Multi-Timeframe Synthesis.

    Architecture:
    1. TimeframeEncoder for each TF -> per-TF representations
    2. Cross-TF Attention -> enriched representations
    3. Learnable TF weights -> weighted combination
    4. Unified representation for downstream tasks

    Input: Dictionary of {timeframe: features}
    Output: Unified representation tensor
    """

    def __init__(
        self,
        config: HierarchicalAttentionConfig,
        input_dims: Optional[Dict[str, int]] = None
    ):
        """
        Initialize Hierarchical Attention.

        Args:
            config: Model configuration
            input_dims: Dictionary of {timeframe: feature_dim}
        """
        super().__init__()

        self.config = config
        self.timeframes = config.timeframes

        # Default input dimensions if not provided
        default_dim = 64
        self.input_dims = input_dims or {tf: default_dim for tf in self.timeframes}

        # Per-timeframe encoders
        self.tf_encoders = nn.ModuleDict({
            tf: TimeframeEncoder(
                input_dim=self.input_dims.get(tf, default_dim),
                d_model=config.d_model,
                n_heads=config.n_heads,
                d_ff=config.d_ff,
                n_layers=config.n_layers,
                dropout=config.dropout,
                max_seq_len=config.max_seq_len
            )
            for tf in self.timeframes
        })

        # Cross-TF attention
        if config.use_cross_tf_attention:
            self.cross_tf_attention = CrossTimeframeAttention(
                d_model=config.d_model,
                n_heads=config.n_heads,
                n_timeframes=len(self.timeframes),
                dropout=config.dropout
            )
        else:
            self.cross_tf_attention = None

        # Learnable TF weights
        if config.use_learnable_tf_weights:
            self.tf_weights = nn.Parameter(torch.ones(len(self.timeframes)))
        else:
            self.register_buffer('tf_weights', torch.ones(len(self.timeframes)))

        # Final fusion layer
        self.fusion = nn.Sequential(
            nn.Linear(config.d_model * len(self.timeframes), config.d_model * 2),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.d_model * 2, config.d_model),
            nn.LayerNorm(config.d_model)
        )

        # Output layer
        self.output_proj = nn.Linear(config.d_model, config.d_model)

        logger.info(f"Initialized HierarchicalAttention with {len(self.timeframes)} timeframes")

    def forward(
        self,
        tf_features: Dict[str, torch.Tensor],
        return_tf_representations: bool = False
    ) -> torch.Tensor:
        """
        Forward pass of hierarchical attention.

        Args:
            tf_features: Dictionary of {timeframe: features}
                Each features tensor has shape (batch, seq_len, feature_dim)
            return_tf_representations: If True, also return per-TF representations

        Returns:
            unified_representation: (batch, d_model)
            tf_representations (optional): Dictionary of per-TF representations
        """
        # Encode each timeframe
        tf_representations = {}
        tf_sequences = {}

        for tf in self.timeframes:
            if tf not in tf_features:
                raise ValueError(f"Missing features for timeframe: {tf}")

            features = tf_features[tf]

            # Handle numpy arrays
            if not isinstance(features, torch.Tensor):
                features = torch.tensor(features, dtype=torch.float32)

            # Ensure 3D: (batch, seq_len, features)
            if features.dim() == 2:
                features = features.unsqueeze(0)

            representation, sequence = self.tf_encoders[tf](features)
            tf_representations[tf] = representation
            tf_sequences[tf] = sequence

        # Apply cross-TF attention if enabled
        if self.cross_tf_attention is not None:
            tf_representations = self.cross_tf_attention(tf_representations)

        # Apply learnable weights
        weights = F.softmax(self.tf_weights, dim=0)

        weighted_reps = []
        for i, tf in enumerate(self.timeframes):
            weighted_reps.append(tf_representations[tf] * weights[i])

        # Concatenate all representations
        concatenated = torch.cat(weighted_reps, dim=-1)

        # Apply fusion
        fused = self.fusion(concatenated)

        # Output projection
        unified = self.output_proj(fused)

        if return_tf_representations:
            return unified, tf_representations
        return unified

    def get_tf_contributions(self) -> Dict[str, float]:
        """
        Get the learned contribution of each timeframe.

        Returns:
            Dictionary of {timeframe: weight}
        """
        weights = F.softmax(self.tf_weights, dim=0)
        return {tf: float(weights[i]) for i, tf in enumerate(self.timeframes)}

    def get_attention_weights(
        self,
        tf_features: Dict[str, torch.Tensor]
    ) -> Dict[str, torch.Tensor]:
        """
        Get attention weights from each timeframe encoder.

        Args:
            tf_features: Dictionary of {timeframe: features}

        Returns:
            Dictionary of attention weights per timeframe
        """
        attention_weights = {}

        for tf in self.timeframes:
            if tf not in tf_features:
                continue

            features = tf_features[tf]
            if not isinstance(features, torch.Tensor):
                features = torch.tensor(features, dtype=torch.float32)
            if features.dim() == 2:
                features = features.unsqueeze(0)

            # Get attention from first layer
            encoder = self.tf_encoders[tf]
            x = encoder.input_proj(features)
            x = x + encoder.pos_encoding[:, :features.shape[1], :]

            _, attn = encoder.layers[0]['attention'](x)
            attention_weights[tf] = attn

        return attention_weights


class HierarchicalAttentionWithAlignment(HierarchicalAttention):
    """
    Extended Hierarchical Attention with alignment/conflict handling.

    Incorporates alignment and conflict scores into the attention mechanism.
    """

    def __init__(
        self,
        config: HierarchicalAttentionConfig,
        input_dims: Optional[Dict[str, int]] = None
    ):
        super().__init__(config, input_dims)

        # Alignment/conflict processing
        self.alignment_proj = nn.Sequential(
            nn.Linear(2, config.d_model // 4),
            nn.GELU(),
            nn.Linear(config.d_model // 4, config.d_model // 4)
        )

        # Modify fusion to include alignment features
        self.fusion_with_alignment = nn.Sequential(
            nn.Linear(config.d_model * len(self.timeframes) + config.d_model // 4, config.d_model * 2),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.d_model * 2, config.d_model),
            nn.LayerNorm(config.d_model)
        )

    def forward(
        self,
        tf_features: Dict[str, torch.Tensor],
        alignment: Optional[torch.Tensor] = None,
        conflict: Optional[torch.Tensor] = None,
        return_tf_representations: bool = False
    ) -> torch.Tensor:
        """
        Forward pass with alignment/conflict.

        Args:
            tf_features: Dictionary of {timeframe: features}
            alignment: Alignment scores (batch,) or (batch, seq_len)
            conflict: Conflict scores (batch,) or (batch, seq_len)
            return_tf_representations: If True, also return per-TF representations

        Returns:
            unified_representation: (batch, d_model)
        """
        # Encode each timeframe
        tf_representations = {}

        for tf in self.timeframes:
            if tf not in tf_features:
                raise ValueError(f"Missing features for timeframe: {tf}")

            features = tf_features[tf]
            if not isinstance(features, torch.Tensor):
                features = torch.tensor(features, dtype=torch.float32)
            if features.dim() == 2:
                features = features.unsqueeze(0)

            representation, _ = self.tf_encoders[tf](features)
            tf_representations[tf] = representation

        # Apply cross-TF attention if enabled
        if self.cross_tf_attention is not None:
            tf_representations = self.cross_tf_attention(tf_representations)

        # Apply learnable weights
        weights = F.softmax(self.tf_weights, dim=0)

        weighted_reps = []
        for i, tf in enumerate(self.timeframes):
            weighted_reps.append(tf_representations[tf] * weights[i])

        # Concatenate all representations
        concatenated = torch.cat(weighted_reps, dim=-1)

        # Process alignment/conflict if provided
        if alignment is not None and conflict is not None:
            batch_size = concatenated.shape[0]

            # Handle sequence alignment/conflict by taking mean
            if alignment.dim() > 1:
                alignment = alignment.mean(dim=-1)
            if conflict.dim() > 1:
                conflict = conflict.mean(dim=-1)

            # Ensure proper shape
            alignment = alignment.view(batch_size, 1)
            conflict = conflict.view(batch_size, 1)

            alignment_features = torch.cat([alignment, conflict], dim=-1)
            alignment_encoded = self.alignment_proj(alignment_features)

            # Concatenate with TF features
            concatenated = torch.cat([concatenated, alignment_encoded], dim=-1)

            # Apply fusion with alignment
            fused = self.fusion_with_alignment(concatenated)
        else:
            # Use regular fusion
            fused = self.fusion(concatenated)

        # Output projection
        unified = self.output_proj(fused)

        if return_tf_representations:
            return unified, tf_representations
        return unified


if __name__ == "__main__":
    # Test HierarchicalAttention
    torch.manual_seed(42)

    # Create configuration
    config = HierarchicalAttentionConfig(
        d_model=128,
        n_heads=4,
        d_ff=256,
        n_layers=2,
        dropout=0.1,
        timeframes=('5m', '15m', '1h', '4h')
    )

    # Define input dimensions for each timeframe
    input_dims = {
        '5m': 64,
        '15m': 48,
        '1h': 40,
        '4h': 32
    }

    # Create model
    model = HierarchicalAttention(config, input_dims)

    # Create sample input
    batch_size = 8
    tf_features = {
        '5m': torch.randn(batch_size, 100, 64),
        '15m': torch.randn(batch_size, 33, 48),
        '1h': torch.randn(batch_size, 8, 40),
        '4h': torch.randn(batch_size, 2, 32)
    }

    # Forward pass
    output = model(tf_features)

    print(f"Output shape: {output.shape}")
    print(f"TF contributions: {model.get_tf_contributions()}")

    # Test with alignment
    model_with_alignment = HierarchicalAttentionWithAlignment(config, input_dims)

    alignment = torch.rand(batch_size)
    conflict = torch.rand(batch_size)

    output_aligned = model_with_alignment(tf_features, alignment, conflict)
    print(f"Output with alignment shape: {output_aligned.shape}")

    print("\nHierarchicalAttention test complete!")
