#!/usr/bin/env python3
"""
Attention Extractor Module
===========================
Provides utilities to extract and visualize attention weights
from Price-Focused Attention models.

This module enables:
1. Extracting attention weights from trained models
2. Producing attention heatmaps
3. Analyzing attention patterns
4. Persisting attention scores to database

Key Features:
- Extract attention from any layer/head
- Generate matplotlib heatmaps
- Compute attention statistics
- Save to PostgreSQL database

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

import json
from datetime import datetime
from typing import List, Optional, Dict, Any, Tuple, Union
from dataclasses import dataclass, field

import numpy as np
import torch
import torch.nn as nn

try:
    import matplotlib.pyplot as plt
    import matplotlib.colors as mcolors
    HAS_MATPLOTLIB = True
except ImportError:
    HAS_MATPLOTLIB = False

try:
    from sqlalchemy import create_engine, text
    from sqlalchemy.engine import Engine
    HAS_SQLALCHEMY = True
except ImportError:
    HAS_SQLALCHEMY = False


@dataclass
class AttentionScores:
    """
    Container for extracted attention scores.

    Attributes:
        scores: Attention weights tensor (batch, n_heads, seq_len, seq_len)
        layer_idx: Which layer the scores came from
        head_idx: Optional head index if single-head extraction
        sequence_len: Length of the input sequence
        n_heads: Number of attention heads
        metadata: Additional metadata (symbol, timestamp, etc.)
    """
    scores: np.ndarray
    layer_idx: int
    head_idx: Optional[int] = None
    sequence_len: int = 0
    n_heads: int = 8
    metadata: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        if self.sequence_len == 0:
            self.sequence_len = self.scores.shape[-1]

    @property
    def shape(self) -> Tuple[int, ...]:
        """Get shape of attention scores."""
        return self.scores.shape

    def mean_attention(self) -> np.ndarray:
        """
        Get mean attention across all heads.

        Returns:
            Mean attention of shape (batch, seq_len, seq_len)
        """
        if len(self.scores.shape) == 4:
            return self.scores.mean(axis=1)
        return self.scores

    def head_attention(self, head_idx: int) -> np.ndarray:
        """
        Get attention for a specific head.

        Args:
            head_idx: Head index (0 to n_heads - 1)

        Returns:
            Attention for that head of shape (batch, seq_len, seq_len)
        """
        if len(self.scores.shape) != 4:
            raise ValueError("Scores must have 4 dimensions: (batch, heads, seq, seq)")
        return self.scores[:, head_idx, :, :]

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            'scores': self.scores.tolist(),
            'layer_idx': self.layer_idx,
            'head_idx': self.head_idx,
            'sequence_len': self.sequence_len,
            'n_heads': self.n_heads,
            'metadata': self.metadata
        }


class AttentionExtractor:
    """
    Utility class for extracting and analyzing attention weights.

    This class provides methods to:
    1. Extract attention weights from PriceFocusedAttention models
    2. Generate attention heatmaps for visualization
    3. Compute attention statistics
    4. Save attention data to database

    Example:
        >>> extractor = AttentionExtractor()
        >>> scores = extractor.get_attention_scores(model, input_data)
        >>> extractor.visualize_attention(scores, save_path='attention.png')
        >>> extractor.save_to_database(scores, symbol='XAUUSD', timestamp=datetime.now())
    """

    def __init__(
        self,
        db_connection_string: Optional[str] = None,
        default_colormap: str = 'viridis'
    ) -> None:
        """
        Initialize the AttentionExtractor.

        Args:
            db_connection_string: Optional database connection string
                                 Format: postgresql://user:pass@host:port/db
            default_colormap: Default matplotlib colormap for heatmaps
        """
        self.db_connection_string = db_connection_string
        self.default_colormap = default_colormap
        self._engine: Optional[Engine] = None

        if db_connection_string and HAS_SQLALCHEMY:
            self._engine = create_engine(db_connection_string)

    def get_attention_scores(
        self,
        model: nn.Module,
        x: torch.Tensor,
        layer_idx: int = -1,
        head_idx: Optional[int] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> AttentionScores:
        """
        Extract attention scores from a model.

        Args:
            model: PriceFocusedAttention model or compatible model
            x: Input tensor of shape (batch, seq_len, features)
            layer_idx: Layer index to extract from (-1 for last layer)
            head_idx: Optional specific head to extract (None for all heads)
            metadata: Optional metadata to attach to the scores

        Returns:
            AttentionScores object containing the extracted weights
        """
        # Ensure model is in evaluation mode
        was_training = model.training
        model.eval()

        with torch.no_grad():
            # Forward pass to get attention weights
            if hasattr(model, 'forward'):
                # Most attention models return (output, attentions)
                output, attentions = model(x, return_all_attentions=True)
            else:
                raise ValueError("Model must have a forward method")

        # Restore training mode if needed
        if was_training:
            model.train()

        # Get specified layer's attention
        if not attentions:
            raise ValueError("Model did not return attention weights")

        n_layers = len(attentions)
        if layer_idx < 0:
            layer_idx = n_layers + layer_idx

        if layer_idx < 0 or layer_idx >= n_layers:
            raise ValueError(f"layer_idx {layer_idx} out of range [0, {n_layers})")

        attention_tensor = attentions[layer_idx]

        # Convert to numpy
        scores_np = attention_tensor.cpu().numpy()

        # Extract specific head if requested
        if head_idx is not None:
            scores_np = scores_np[:, head_idx:head_idx+1, :, :]

        return AttentionScores(
            scores=scores_np,
            layer_idx=layer_idx,
            head_idx=head_idx,
            n_heads=attention_tensor.shape[1],
            metadata=metadata or {}
        )

    def visualize_attention(
        self,
        scores: Union[AttentionScores, np.ndarray, torch.Tensor],
        head_idx: int = 0,
        batch_idx: int = 0,
        title: Optional[str] = None,
        save_path: Optional[str] = None,
        figsize: Tuple[int, int] = (10, 8),
        colormap: Optional[str] = None,
        show_colorbar: bool = True,
        x_labels: Optional[List[str]] = None,
        y_labels: Optional[List[str]] = None
    ) -> Optional[Any]:
        """
        Visualize attention weights as a heatmap.

        Args:
            scores: AttentionScores object, numpy array, or tensor
            head_idx: Attention head to visualize (default: 0)
            batch_idx: Batch sample to visualize (default: 0)
            title: Optional plot title
            save_path: If provided, save figure to this path
            figsize: Figure size (width, height)
            colormap: Matplotlib colormap name
            show_colorbar: Whether to show colorbar
            x_labels: Optional labels for x-axis (key positions)
            y_labels: Optional labels for y-axis (query positions)

        Returns:
            Matplotlib figure object if matplotlib is available, None otherwise
        """
        if not HAS_MATPLOTLIB:
            print("matplotlib not available. Cannot visualize attention.")
            return None

        # Extract numpy array from scores
        if isinstance(scores, AttentionScores):
            attention = scores.scores
        elif isinstance(scores, torch.Tensor):
            attention = scores.cpu().numpy()
        else:
            attention = scores

        # Get specific batch and head
        if len(attention.shape) == 4:
            attention_2d = attention[batch_idx, head_idx, :, :]
        elif len(attention.shape) == 3:
            attention_2d = attention[batch_idx, :, :]
        else:
            attention_2d = attention

        # Create figure
        fig, ax = plt.subplots(figsize=figsize)

        # Create heatmap
        cmap = colormap or self.default_colormap
        im = ax.imshow(attention_2d, cmap=cmap, aspect='auto')

        # Add colorbar
        if show_colorbar:
            cbar = plt.colorbar(im, ax=ax)
            cbar.set_label('Attention Weight')

        # Set title
        if title:
            ax.set_title(title)
        else:
            ax.set_title(f'Attention Heatmap (Head {head_idx})')

        # Set axis labels
        ax.set_xlabel('Key Position')
        ax.set_ylabel('Query Position')

        # Add custom tick labels if provided
        if x_labels is not None:
            ax.set_xticks(range(len(x_labels)))
            ax.set_xticklabels(x_labels, rotation=45, ha='right')

        if y_labels is not None:
            ax.set_yticks(range(len(y_labels)))
            ax.set_yticklabels(y_labels)

        plt.tight_layout()

        # Save if path provided
        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')

        return fig

    def visualize_all_heads(
        self,
        scores: Union[AttentionScores, np.ndarray],
        batch_idx: int = 0,
        save_path: Optional[str] = None,
        figsize_per_head: Tuple[int, int] = (4, 3)
    ) -> Optional[Any]:
        """
        Visualize attention from all heads in a grid.

        Args:
            scores: AttentionScores object or numpy array
            batch_idx: Batch sample to visualize
            save_path: Optional path to save the figure
            figsize_per_head: Size of each subplot

        Returns:
            Matplotlib figure object
        """
        if not HAS_MATPLOTLIB:
            print("matplotlib not available.")
            return None

        # Extract numpy array
        if isinstance(scores, AttentionScores):
            attention = scores.scores
            n_heads = scores.n_heads
        else:
            attention = scores
            n_heads = attention.shape[1]

        # Calculate grid dimensions
        n_cols = min(4, n_heads)
        n_rows = (n_heads + n_cols - 1) // n_cols

        fig, axes = plt.subplots(
            n_rows, n_cols,
            figsize=(figsize_per_head[0] * n_cols, figsize_per_head[1] * n_rows)
        )
        axes = np.atleast_2d(axes)

        for head_idx in range(n_heads):
            row = head_idx // n_cols
            col = head_idx % n_cols
            ax = axes[row, col]

            attention_2d = attention[batch_idx, head_idx, :, :]
            im = ax.imshow(attention_2d, cmap=self.default_colormap, aspect='auto')
            ax.set_title(f'Head {head_idx}')
            ax.set_xlabel('Key')
            ax.set_ylabel('Query')

        # Hide unused subplots
        for idx in range(n_heads, n_rows * n_cols):
            row = idx // n_cols
            col = idx % n_cols
            axes[row, col].axis('off')

        plt.suptitle('Attention Patterns Across All Heads')
        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')

        return fig

    def compute_attention_statistics(
        self,
        scores: Union[AttentionScores, np.ndarray]
    ) -> Dict[str, Any]:
        """
        Compute statistics about attention patterns.

        Args:
            scores: AttentionScores object or numpy array

        Returns:
            Dictionary with attention statistics
        """
        if isinstance(scores, AttentionScores):
            attention = scores.scores
        else:
            attention = scores

        # Flatten for statistics
        flat = attention.flatten()

        # Compute per-head statistics
        head_stats = []
        for h in range(attention.shape[1]):
            head_attn = attention[:, h, :, :]
            head_stats.append({
                'head': h,
                'mean': float(head_attn.mean()),
                'std': float(head_attn.std()),
                'max': float(head_attn.max()),
                'entropy': float(self._compute_entropy(head_attn))
            })

        # Compute diagonal attention (self-attention strength)
        batch_size, n_heads, seq_len, _ = attention.shape
        diagonal_attention = np.zeros((batch_size, n_heads, seq_len))
        for i in range(seq_len):
            diagonal_attention[:, :, i] = attention[:, :, i, i]

        return {
            'global': {
                'mean': float(flat.mean()),
                'std': float(flat.std()),
                'max': float(flat.max()),
                'min': float(flat.min())
            },
            'per_head': head_stats,
            'diagonal_attention_mean': float(diagonal_attention.mean()),
            'sparsity': float((attention < 0.01).sum() / attention.size)
        }

    def _compute_entropy(self, attention: np.ndarray) -> float:
        """Compute entropy of attention distribution."""
        # Ensure attention is a probability distribution
        attention = attention.clip(min=1e-10)
        # Normalize if not already
        attention = attention / attention.sum(axis=-1, keepdims=True)
        # Compute entropy
        entropy = -np.sum(attention * np.log(attention + 1e-10), axis=-1)
        return float(entropy.mean())

    def save_to_database(
        self,
        scores: AttentionScores,
        symbol: str,
        timestamp: datetime,
        table_name: str = 'ml.attention_scores',
        additional_data: Optional[Dict[str, Any]] = None
    ) -> bool:
        """
        Save attention scores to PostgreSQL database.

        Args:
            scores: AttentionScores object
            symbol: Trading symbol (e.g., 'XAUUSD')
            timestamp: Timestamp of the prediction
            table_name: Database table name (default: ml.attention_scores)
            additional_data: Additional data to store

        Returns:
            True if save was successful, False otherwise
        """
        if not HAS_SQLALCHEMY or self._engine is None:
            print("SQLAlchemy not available or no database connection configured")
            return False

        try:
            # Prepare data for insertion
            data = {
                'symbol': symbol,
                'timestamp': timestamp,
                'layer_idx': scores.layer_idx,
                'head_idx': scores.head_idx,
                'sequence_len': scores.sequence_len,
                'n_heads': scores.n_heads,
                'scores_json': json.dumps(scores.scores.tolist()),
                'metadata': json.dumps(scores.metadata),
                'statistics': json.dumps(self.compute_attention_statistics(scores)),
                'created_at': datetime.now()
            }

            # Add additional data
            if additional_data:
                data['additional_data'] = json.dumps(additional_data)

            # Create insert query
            columns = ', '.join(data.keys())
            placeholders = ', '.join([f':{k}' for k in data.keys()])
            query = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"

            with self._engine.connect() as conn:
                conn.execute(text(query), data)
                conn.commit()

            return True

        except Exception as e:
            print(f"Error saving to database: {e}")
            return False

    def load_from_database(
        self,
        symbol: str,
        start_timestamp: datetime,
        end_timestamp: Optional[datetime] = None,
        table_name: str = 'ml.attention_scores',
        limit: int = 100
    ) -> List[AttentionScores]:
        """
        Load attention scores from database.

        Args:
            symbol: Trading symbol
            start_timestamp: Start of time range
            end_timestamp: End of time range (default: now)
            table_name: Database table name
            limit: Maximum number of records

        Returns:
            List of AttentionScores objects
        """
        if not HAS_SQLALCHEMY or self._engine is None:
            print("SQLAlchemy not available or no database connection")
            return []

        end_timestamp = end_timestamp or datetime.now()

        query = f"""
        SELECT symbol, timestamp, layer_idx, head_idx, sequence_len,
               n_heads, scores_json, metadata
        FROM {table_name}
        WHERE symbol = :symbol
          AND timestamp >= :start_timestamp
          AND timestamp <= :end_timestamp
        ORDER BY timestamp DESC
        LIMIT :limit
        """

        try:
            with self._engine.connect() as conn:
                result = conn.execute(text(query), {
                    'symbol': symbol,
                    'start_timestamp': start_timestamp,
                    'end_timestamp': end_timestamp,
                    'limit': limit
                })

                scores_list = []
                for row in result:
                    scores = AttentionScores(
                        scores=np.array(json.loads(row.scores_json)),
                        layer_idx=row.layer_idx,
                        head_idx=row.head_idx,
                        sequence_len=row.sequence_len,
                        n_heads=row.n_heads,
                        metadata=json.loads(row.metadata) if row.metadata else {}
                    )
                    scores_list.append(scores)

                return scores_list

        except Exception as e:
            print(f"Error loading from database: {e}")
            return []


def extract_attention_for_analysis(
    model: nn.Module,
    x: torch.Tensor,
    normalize: bool = True
) -> Dict[str, np.ndarray]:
    """
    Extract attention weights for detailed analysis.

    Convenience function that extracts attention from all layers
    and computes useful summary statistics.

    Args:
        model: PriceFocusedAttention model
        x: Input tensor (batch, seq_len, features)
        normalize: Whether to normalize attention weights

    Returns:
        Dictionary with:
        - 'all_layers': List of attention arrays per layer
        - 'mean_attention': Mean attention across all layers
        - 'last_layer': Attention from the last layer
        - 'head_importance': Importance score per head
    """
    extractor = AttentionExtractor()

    # Get number of layers by doing a forward pass
    model.eval()
    with torch.no_grad():
        _, all_attentions = model(x, return_all_attentions=True)

    n_layers = len(all_attentions)

    # Extract attention from each layer
    layer_attentions = []
    for layer_idx in range(n_layers):
        scores = extractor.get_attention_scores(model, x, layer_idx=layer_idx)
        layer_attentions.append(scores.scores)

    # Stack and compute mean
    stacked = np.stack(layer_attentions, axis=0)  # (n_layers, batch, heads, seq, seq)
    mean_attention = stacked.mean(axis=0)  # (batch, heads, seq, seq)

    # Compute head importance (variance of attention patterns)
    head_importance = []
    for h in range(mean_attention.shape[1]):
        head_attn = mean_attention[:, h, :, :]
        importance = float(head_attn.std())
        head_importance.append(importance)

    # Normalize importance
    head_importance = np.array(head_importance)
    if normalize and head_importance.sum() > 0:
        head_importance = head_importance / head_importance.sum()

    return {
        'all_layers': layer_attentions,
        'mean_attention': mean_attention,
        'last_layer': layer_attentions[-1],
        'head_importance': head_importance.tolist(),
        'n_layers': n_layers,
        'n_heads': mean_attention.shape[1]
    }
