"""
Attention Score Provider
========================

Loads and serves attention scores from trained Level 0 models.
These scores are used as features (F51, F52) in Level 1 predictions.

Author: ML Pipeline
Version: 1.0.0
Created: 2026-01-25
"""

import sys
from pathlib import Path
from typing import Dict, Optional, Tuple
import numpy as np
import pandas as pd
from loguru import logger
import joblib

# Import attention model classes
import importlib.util

_models_dir = Path(__file__).parent.parent / 'models'
_attention_module_path = _models_dir / 'attention_score_model.py'

# Use consistent module name for joblib pickle compatibility
_module_name = "models.attention_score_model"
if _module_name not in sys.modules and _attention_module_path.exists():
    spec = importlib.util.spec_from_file_location(_module_name, _attention_module_path)
    attention_module = importlib.util.module_from_spec(spec)
    sys.modules[_module_name] = attention_module
    spec.loader.exec_module(attention_module)
    AttentionScoreModel = attention_module.AttentionScoreModel
    HAS_ATTENTION_MODEL = True
elif _module_name in sys.modules:
    attention_module = sys.modules[_module_name]
    AttentionScoreModel = attention_module.AttentionScoreModel
    HAS_ATTENTION_MODEL = True
else:
    AttentionScoreModel = None
    HAS_ATTENTION_MODEL = False
    logger.warning("AttentionScoreModel not available")


class AttentionProvider:
    """
    Provides attention scores for predictions.

    Loads trained attention models and computes scores
    to be used as features in Level 1+ models.

    Usage:
        provider = AttentionProvider('models/attention')
        provider.load()

        # Get scores for prediction
        score, flow_class = provider.get_attention_features(df, 'XAUUSD', '5m')
    """

    def __init__(self, models_dir: str = 'models/attention'):
        """
        Initialize attention provider.

        Args:
            models_dir: Directory containing trained attention models
        """
        self.models_dir = Path(models_dir)
        self.models: Dict[str, 'AttentionScoreModel'] = {}
        self.metadata: Dict = {}
        self._loaded = False

    def load(self) -> bool:
        """
        Load all attention models from disk.

        Returns:
            True if at least one model was loaded
        """
        if not HAS_ATTENTION_MODEL:
            logger.warning("AttentionScoreModel class not available")
            return False

        if not self.models_dir.exists():
            logger.warning(f"Attention models directory not found: {self.models_dir}")
            return False

        # Load metadata
        metadata_path = self.models_dir / 'trainer_metadata.joblib'
        if metadata_path.exists():
            self.metadata = joblib.load(metadata_path)
            logger.info(f"Loaded attention trainer metadata")

        # Load individual models
        loaded = 0
        for model_dir in self.models_dir.iterdir():
            if not model_dir.is_dir():
                continue

            # Skip non-model directories
            if not model_dir.name.endswith('_attention'):
                continue

            try:
                model = AttentionScoreModel.load(str(model_dir))
                key = model_dir.name  # e.g., 'XAUUSD_5m_attention'
                self.models[key] = model
                loaded += 1
                logger.debug(f"Loaded attention model: {key}")
            except Exception as e:
                logger.warning(f"Failed to load {model_dir.name}: {e}")

        self._loaded = loaded > 0

        if self._loaded:
            logger.info(f"✅ Loaded {loaded} attention models from {self.models_dir}")
        else:
            logger.warning(f"No attention models loaded from {self.models_dir}")

        return self._loaded

    @property
    def is_loaded(self) -> bool:
        """Check if models are loaded."""
        return self._loaded

    @property
    def available_models(self) -> list:
        """List of available model keys."""
        return list(self.models.keys())

    def get_model_key(self, symbol: str, timeframe: str) -> str:
        """Generate model key from symbol and timeframe."""
        return f"{symbol}_{timeframe}_attention"

    def has_model(self, symbol: str, timeframe: str) -> bool:
        """Check if model exists for symbol/timeframe."""
        key = self.get_model_key(symbol, timeframe)
        return key in self.models

    def get_attention_score(
        self,
        df: pd.DataFrame,
        symbol: str,
        timeframe: str
    ) -> np.ndarray:
        """
        Get attention scores for a DataFrame.

        Args:
            df: OHLCV DataFrame with datetime index
            symbol: Trading symbol (e.g., 'XAUUSD')
            timeframe: Timeframe (e.g., '5m')

        Returns:
            Array of attention scores (1.0 = average, >2.0 = high flow)
        """
        key = self.get_model_key(symbol, timeframe)

        if key not in self.models:
            # Return neutral scores if model not available
            logger.debug(f"No attention model for {key}, returning neutral scores")
            return np.ones(len(df))

        try:
            prediction = self.models[key].predict(df)
            return prediction.attention_score
        except Exception as e:
            logger.warning(f"Attention prediction failed for {key}: {e}")
            return np.ones(len(df))

    def get_attention_class(
        self,
        df: pd.DataFrame,
        symbol: str,
        timeframe: str
    ) -> np.ndarray:
        """
        Get attention flow class for a DataFrame.

        Args:
            df: OHLCV DataFrame
            symbol: Trading symbol
            timeframe: Timeframe

        Returns:
            Array of flow classes (0=low, 1=medium, 2=high)
        """
        key = self.get_model_key(symbol, timeframe)

        if key not in self.models:
            # Return medium flow if model not available
            return np.ones(len(df), dtype=int)

        try:
            prediction = self.models[key].predict(df)
            return prediction.flow_class
        except Exception as e:
            logger.warning(f"Attention class prediction failed for {key}: {e}")
            return np.ones(len(df), dtype=int)

    def get_attention_features(
        self,
        df: pd.DataFrame,
        symbol: str,
        timeframe: str
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Get both attention score and class as features.

        This is the main method for integrating attention into Level 1+ models.

        Args:
            df: OHLCV DataFrame
            symbol: Trading symbol
            timeframe: Timeframe

        Returns:
            Tuple of (attention_score, flow_class) arrays
        """
        key = self.get_model_key(symbol, timeframe)

        if key not in self.models:
            # Return neutral values if model not available
            n = len(df)
            return np.ones(n), np.ones(n, dtype=int)

        try:
            prediction = self.models[key].predict(df)
            return prediction.attention_score, prediction.flow_class
        except Exception as e:
            logger.warning(f"Attention features failed for {key}: {e}")
            n = len(df)
            return np.ones(n), np.ones(n, dtype=int)

    def get_latest_attention(
        self,
        df: pd.DataFrame,
        symbol: str,
        timeframe: str
    ) -> Dict:
        """
        Get attention info for the latest bar.

        Useful for real-time signal generation.

        Args:
            df: OHLCV DataFrame (at least 200 bars for feature calculation)
            symbol: Trading symbol
            timeframe: Timeframe

        Returns:
            Dict with attention_score, flow_class, flow_label, is_high_flow
        """
        score, flow_class = self.get_attention_features(df, symbol, timeframe)

        latest_score = float(score[-1]) if len(score) > 0 else 1.0
        latest_class = int(flow_class[-1]) if len(flow_class) > 0 else 1

        flow_labels = {0: 'low_flow', 1: 'medium_flow', 2: 'high_flow'}

        return {
            'attention_score': latest_score,
            'flow_class': latest_class,
            'flow_label': flow_labels.get(latest_class, 'medium_flow'),
            'is_high_flow': latest_class == 2,
            'is_tradeable': latest_score >= 1.0  # Only trade when attention >= 1
        }

    def should_trade(
        self,
        df: pd.DataFrame,
        symbol: str,
        timeframe: str,
        min_score: float = 1.0
    ) -> bool:
        """
        Check if current market conditions are favorable for trading.

        Args:
            df: OHLCV DataFrame
            symbol: Trading symbol
            timeframe: Timeframe
            min_score: Minimum attention score to trade

        Returns:
            True if attention score >= min_score
        """
        attention = self.get_latest_attention(df, symbol, timeframe)
        return attention['attention_score'] >= min_score


# Global singleton instance
_attention_provider: Optional[AttentionProvider] = None


def get_attention_provider(models_dir: str = None) -> AttentionProvider:
    """
    Get or create the global attention provider instance.

    Args:
        models_dir: Directory containing attention models

    Returns:
        AttentionProvider instance
    """
    global _attention_provider

    if _attention_provider is None:
        # Default path relative to ml-engine
        if models_dir is None:
            models_dir = Path(__file__).parent.parent.parent / 'models' / 'attention'

        _attention_provider = AttentionProvider(str(models_dir))
        _attention_provider.load()

    return _attention_provider
