#!/usr/bin/env python3
"""
Attention Score Model
=====================
Predicts attention/flow score based on market indicators.

This model learns WHEN to pay attention to the market without
hardcoding specific hours or sessions. It uses volume, volatility,
and momentum indicators to identify high-flow moments.

Key Features:
1. Regression output: attention_score (0-3+)
2. Classification output: flow_class (low/medium/high)
3. Dynamic factor-based target (no time leakage)
4. Uses only market-derived features (no hardcoded sessions)

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

import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Optional, Any, Union
from dataclasses import dataclass, field
from pathlib import Path
import joblib
from loguru import logger

try:
    from xgboost import XGBRegressor, XGBClassifier
    HAS_XGBOOST = True
except ImportError:
    HAS_XGBOOST = False
    logger.warning("XGBoost not available")

from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.metrics import accuracy_score, f1_score, classification_report


@dataclass
class AttentionModelConfig:
    """Configuration for the attention score model"""

    # Feature configuration
    feature_names: List[str] = field(default_factory=lambda: [
        'volume_ratio',      # volume / rolling_median(20)
        'volume_z',          # z-score of volume
        'ATR',               # Average True Range
        'ATR_ratio',         # ATR / rolling_median(ATR, 50)
        'CMF',               # Chaikin Money Flow
        'MFI',               # Money Flow Index
        'OBV_delta',         # diff(OBV) / rolling_std(OBV, 20)
        'BB_width',          # (BB_upper - BB_lower) / close
        'displacement',      # (close - open) / ATR
    ])

    # Target configuration
    factor_window: int = 200          # Rolling window for factor calculation
    horizon_bars: int = 3             # Bars ahead for target calculation

    # Classification thresholds
    low_flow_threshold: float = 1.0   # move_multiplier < 1.0 = low flow
    high_flow_threshold: float = 2.0  # move_multiplier >= 2.0 = high flow

    # XGBoost regression parameters
    reg_params: Dict = field(default_factory=lambda: {
        'n_estimators': 200,
        'max_depth': 5,
        'learning_rate': 0.05,
        'subsample': 0.8,
        'colsample_bytree': 0.8,
        'min_child_weight': 10,
        'gamma': 0.1,
        'reg_alpha': 0.1,
        'reg_lambda': 1.0,
        'tree_method': 'hist',
        'random_state': 42
    })

    # XGBoost classification parameters
    clf_params: Dict = field(default_factory=lambda: {
        'n_estimators': 150,
        'max_depth': 4,
        'learning_rate': 0.05,
        'subsample': 0.8,
        'colsample_bytree': 0.8,
        'min_child_weight': 15,
        'gamma': 0.2,
        'reg_alpha': 0.1,
        'reg_lambda': 1.0,
        'tree_method': 'hist',
        'random_state': 42,
        'objective': 'multi:softmax',
        'num_class': 3
    })

    # Validation
    val_split: float = 0.15
    min_train_samples: int = 5000


@dataclass
class AttentionPrediction:
    """Result of attention prediction"""
    attention_score: np.ndarray       # Continuous score (0-3+)
    flow_class: np.ndarray            # Class: 0=low, 1=medium, 2=high
    flow_class_proba: np.ndarray      # Class probabilities

    def to_dict(self) -> Dict[str, Any]:
        return {
            'attention_score': self.attention_score.tolist() if isinstance(self.attention_score, np.ndarray) else [self.attention_score],
            'flow_class': self.flow_class.tolist() if isinstance(self.flow_class, np.ndarray) else [self.flow_class],
            'flow_class_names': ['low_flow', 'medium_flow', 'high_flow']
        }


class AttentionFeatureGenerator:
    """
    Generates attention-specific features from OHLCV data.

    All features are designed to capture market flow/activity
    WITHOUT using time-based features (no hour, session flags, etc.)
    """

    def __init__(self, config: AttentionModelConfig = None):
        self.config = config or AttentionModelConfig()

    def generate_features(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        Generate attention features from OHLCV data.

        Args:
            df: DataFrame with OHLCV columns (open, high, low, close, volume)

        Returns:
            DataFrame with attention features
        """
        features = pd.DataFrame(index=df.index)

        # Normalize column names
        close = df['Close'] if 'Close' in df.columns else df['close']
        high = df['High'] if 'High' in df.columns else df['high']
        low = df['Low'] if 'Low' in df.columns else df['low']
        open_price = df['Open'] if 'Open' in df.columns else df['open']
        volume = df['Volume'] if 'Volume' in df.columns else df.get('volume', pd.Series(1, index=df.index))

        # 1. VOLUME FEATURES
        # Volume ratio: current volume vs rolling median
        vol_median_20 = volume.rolling(20).median()
        features['volume_ratio'] = volume / (vol_median_20 + 1e-10)

        # Volume z-score
        vol_mean_20 = volume.rolling(20).mean()
        vol_std_20 = volume.rolling(20).std()
        features['volume_z'] = (volume - vol_mean_20) / (vol_std_20 + 1e-10)

        # 2. VOLATILITY FEATURES
        # ATR (Average True Range)
        tr1 = high - low
        tr2 = abs(high - close.shift(1))
        tr3 = abs(low - close.shift(1))
        true_range = pd.concat([tr1, tr2, tr3], axis=1).max(axis=1)
        features['ATR'] = true_range.rolling(14).mean()

        # ATR ratio: current ATR vs rolling median
        atr_median_50 = features['ATR'].rolling(50).median()
        features['ATR_ratio'] = features['ATR'] / (atr_median_50 + 1e-10)

        # 3. MONEY FLOW FEATURES
        # Chaikin Money Flow (CMF)
        mf_multiplier = ((close - low) - (high - close)) / (high - low + 1e-10)
        mf_volume = mf_multiplier * volume
        features['CMF'] = mf_volume.rolling(20).sum() / (volume.rolling(20).sum() + 1e-10)

        # Money Flow Index (MFI)
        typical_price = (high + low + close) / 3
        money_flow = typical_price * volume
        positive_flow = money_flow.where(typical_price > typical_price.shift(1), 0)
        negative_flow = money_flow.where(typical_price < typical_price.shift(1), 0)
        positive_sum = positive_flow.rolling(14).sum()
        negative_sum = negative_flow.rolling(14).sum()
        mf_ratio = positive_sum / (negative_sum + 1e-10)
        features['MFI'] = 100 - (100 / (1 + mf_ratio))

        # 4. ORDER FLOW FEATURES
        # OBV delta (normalized)
        obv = (np.sign(close.diff()) * volume).cumsum()
        obv_diff = obv.diff()
        obv_std = obv_diff.rolling(20).std()
        features['OBV_delta'] = obv_diff / (obv_std + 1e-10)

        # 5. PRICE STRUCTURE FEATURES
        # Bollinger Band width (volatility expansion/contraction)
        bb_middle = close.rolling(20).mean()
        bb_std = close.rolling(20).std()
        bb_upper = bb_middle + 2 * bb_std
        bb_lower = bb_middle - 2 * bb_std
        features['BB_width'] = (bb_upper - bb_lower) / (close + 1e-10)

        # Displacement: strength of current candle move
        candle_range = high - low
        features['displacement'] = (close - open_price) / (features['ATR'] + 1e-10)

        # Clean up infinities and NaNs
        features = features.replace([np.inf, -np.inf], np.nan)

        return features

    def compute_target_move_multiplier(
        self,
        df: pd.DataFrame,
        horizon_bars: int = 3,
        factor_window: int = 200
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Compute move_multiplier target for training.

        Formula:
            factor = rolling_median(range, window).shift(1)  # No leakage
            future_range = max(high[t+1:t+horizon+1]) - min(low[t+1:t+horizon+1])
            move_multiplier = future_range / factor

        Args:
            df: DataFrame with OHLCV
            horizon_bars: Number of bars to look ahead
            factor_window: Rolling window for factor calculation

        Returns:
            Tuple of (move_multiplier, valid_mask)
        """
        # Normalize column names
        high = df['High'] if 'High' in df.columns else df['high']
        low = df['Low'] if 'Low' in df.columns else df['low']
        close = df['Close'] if 'Close' in df.columns else df['close']

        n = len(df)

        # Compute dynamic factor with shift(1) to avoid leakage
        candle_range = (high - low).abs()
        factor = candle_range.rolling(factor_window, min_periods=factor_window // 2).median().shift(1)

        # Compute future range
        future_range = np.full(n, np.nan)
        for i in range(n - horizon_bars):
            future_high = high.iloc[i+1:i+1+horizon_bars].max()
            future_low = low.iloc[i+1:i+1+horizon_bars].min()
            future_range[i] = future_high - future_low

        # Compute move_multiplier
        move_multiplier = future_range / (factor.values + 1e-10)

        # Valid mask
        valid_mask = ~np.isnan(move_multiplier) & ~np.isnan(factor.values)

        logger.info(f"Target computed:")
        logger.info(f"  Valid samples: {valid_mask.sum()} / {n}")
        logger.info(f"  Move multiplier range: [{np.nanmin(move_multiplier):.2f}, {np.nanmax(move_multiplier):.2f}]")
        logger.info(f"  Mean move multiplier: {np.nanmean(move_multiplier):.2f}")

        return move_multiplier, valid_mask

    def multiplier_to_class(
        self,
        move_multiplier: np.ndarray,
        low_threshold: float = 1.0,
        high_threshold: float = 2.0
    ) -> np.ndarray:
        """
        Convert move_multiplier to flow class.

        Classes:
            0 = low_flow (move_multiplier < low_threshold)
            1 = medium_flow (low_threshold <= move_multiplier < high_threshold)
            2 = high_flow (move_multiplier >= high_threshold)
        """
        classes = np.zeros(len(move_multiplier), dtype=int)
        classes[(move_multiplier >= low_threshold) & (move_multiplier < high_threshold)] = 1
        classes[move_multiplier >= high_threshold] = 2

        # Log class distribution
        for i, name in enumerate(['low_flow', 'medium_flow', 'high_flow']):
            count = (classes == i).sum()
            pct = count / len(classes) * 100
            logger.info(f"  {name}: {count} ({pct:.1f}%)")

        return classes


class AttentionScoreModel:
    """
    Main attention score model that predicts when to pay attention.

    Outputs:
    1. attention_score (regression): Continuous value 0-3+ indicating
       expected market movement relative to normal
    2. flow_class (classification): Categorical (low/medium/high)

    Usage:
        model = AttentionScoreModel(config)
        model.fit(df_ohlcv)

        prediction = model.predict(df_ohlcv)
        print(f"Attention: {prediction.attention_score}")
        print(f"Flow: {prediction.flow_class}")
    """

    def __init__(self, config: AttentionModelConfig = None):
        self.config = config or AttentionModelConfig()
        self.feature_generator = AttentionFeatureGenerator(self.config)

        self.regressor: Optional[XGBRegressor] = None
        self.classifier: Optional[XGBClassifier] = None

        self.is_fitted = False
        self.training_metrics: Dict[str, float] = {}
        self.feature_importance: Optional[pd.DataFrame] = None

    def fit(
        self,
        df: pd.DataFrame,
        sample_weights: np.ndarray = None
    ) -> 'AttentionScoreModel':
        """
        Train both regression and classification models.

        Args:
            df: DataFrame with OHLCV data (datetime index)
            sample_weights: Optional sample weights

        Returns:
            Self for chaining
        """
        logger.info("="*60)
        logger.info("Training AttentionScoreModel")
        logger.info("="*60)

        # Generate features
        features = self.feature_generator.generate_features(df)

        # Compute targets
        move_multiplier, valid_mask = self.feature_generator.compute_target_move_multiplier(
            df,
            horizon_bars=self.config.horizon_bars,
            factor_window=self.config.factor_window
        )

        # Convert to classification target
        flow_class = self.feature_generator.multiplier_to_class(
            move_multiplier,
            self.config.low_flow_threshold,
            self.config.high_flow_threshold
        )

        # Prepare training data
        X = features[self.config.feature_names].values
        y_reg = move_multiplier
        y_clf = flow_class

        # Apply valid mask
        valid_features = ~np.isnan(X).any(axis=1) & valid_mask
        X_valid = X[valid_features]
        y_reg_valid = y_reg[valid_features]
        y_clf_valid = y_clf[valid_features]

        if sample_weights is not None:
            weights_valid = sample_weights[valid_features]
        else:
            weights_valid = None

        logger.info(f"Training data: {len(X_valid)} samples")

        if len(X_valid) < self.config.min_train_samples:
            raise ValueError(f"Insufficient training data: {len(X_valid)} < {self.config.min_train_samples}")

        # Train/val split (time-based)
        split_idx = int(len(X_valid) * (1 - self.config.val_split))
        X_train, X_val = X_valid[:split_idx], X_valid[split_idx:]
        y_reg_train, y_reg_val = y_reg_valid[:split_idx], y_reg_valid[split_idx:]
        y_clf_train, y_clf_val = y_clf_valid[:split_idx], y_clf_valid[split_idx:]

        if weights_valid is not None:
            weights_train = weights_valid[:split_idx]
        else:
            weights_train = None

        # Train regressor
        logger.info("\nTraining regression model...")
        self.regressor = XGBRegressor(**self.config.reg_params)
        self.regressor.fit(
            X_train, y_reg_train,
            sample_weight=weights_train,
            eval_set=[(X_val, y_reg_val)],
            verbose=False
        )

        # Evaluate regressor
        y_reg_pred = self.regressor.predict(X_val)
        reg_mae = mean_absolute_error(y_reg_val, y_reg_pred)
        reg_rmse = np.sqrt(mean_squared_error(y_reg_val, y_reg_pred))
        reg_r2 = r2_score(y_reg_val, y_reg_pred)

        logger.info(f"  Regression MAE: {reg_mae:.4f}")
        logger.info(f"  Regression RMSE: {reg_rmse:.4f}")
        logger.info(f"  Regression R2: {reg_r2:.4f}")

        # Train classifier
        logger.info("\nTraining classification model...")
        self.classifier = XGBClassifier(**self.config.clf_params)
        self.classifier.fit(
            X_train, y_clf_train,
            sample_weight=weights_train,
            eval_set=[(X_val, y_clf_val)],
            verbose=False
        )

        # Evaluate classifier
        y_clf_pred = self.classifier.predict(X_val)
        clf_acc = accuracy_score(y_clf_val, y_clf_pred)
        clf_f1 = f1_score(y_clf_val, y_clf_pred, average='weighted')

        logger.info(f"  Classification Accuracy: {clf_acc:.4f}")
        logger.info(f"  Classification F1 (weighted): {clf_f1:.4f}")

        # Store metrics
        self.training_metrics = {
            'reg_mae': reg_mae,
            'reg_rmse': reg_rmse,
            'reg_r2': reg_r2,
            'clf_accuracy': clf_acc,
            'clf_f1': clf_f1,
            'n_train': len(X_train),
            'n_val': len(X_val)
        }

        # Compute feature importance
        self._compute_feature_importance()

        self.is_fitted = True
        logger.info("\nAttentionScoreModel training complete")

        return self

    def predict(
        self,
        df: pd.DataFrame,
        features: pd.DataFrame = None
    ) -> AttentionPrediction:
        """
        Predict attention score and flow class.

        Args:
            df: DataFrame with OHLCV data
            features: Optional pre-computed features

        Returns:
            AttentionPrediction with score, class, and probabilities
        """
        if not self.is_fitted:
            raise ValueError("Model not fitted. Call fit() first.")

        # Generate features if not provided
        if features is None:
            features = self.feature_generator.generate_features(df)

        X = features[self.config.feature_names].values

        # Handle NaN by filling with 0 (neutral)
        X = np.nan_to_num(X, nan=0.0)

        # Predict regression
        attention_score = self.regressor.predict(X)
        attention_score = np.maximum(attention_score, 0)  # Floor at 0

        # Predict classification
        flow_class = self.classifier.predict(X)
        flow_class_proba = self.classifier.predict_proba(X)

        return AttentionPrediction(
            attention_score=attention_score,
            flow_class=flow_class,
            flow_class_proba=flow_class_proba
        )

    def predict_single(
        self,
        df_row: pd.DataFrame
    ) -> Tuple[float, int, np.ndarray]:
        """
        Predict for a single row of data.

        Returns:
            Tuple of (attention_score, flow_class, class_probabilities)
        """
        prediction = self.predict(df_row)
        return (
            float(prediction.attention_score[0]),
            int(prediction.flow_class[0]),
            prediction.flow_class_proba[0]
        )

    def _compute_feature_importance(self):
        """Compute and store feature importance from both models."""
        reg_importance = self.regressor.feature_importances_
        clf_importance = self.classifier.feature_importances_

        self.feature_importance = pd.DataFrame({
            'feature': self.config.feature_names,
            'reg_importance': reg_importance,
            'clf_importance': clf_importance,
            'combined': (reg_importance + clf_importance) / 2
        }).sort_values('combined', ascending=False)

        logger.info("\nTop features:")
        for _, row in self.feature_importance.head(5).iterrows():
            logger.info(f"  {row['feature']}: {row['combined']:.4f}")

    def get_training_summary(self) -> Dict[str, Any]:
        """Get summary of training results."""
        return {
            'is_fitted': self.is_fitted,
            'metrics': self.training_metrics,
            'config': {
                'feature_names': self.config.feature_names,
                'factor_window': self.config.factor_window,
                'horizon_bars': self.config.horizon_bars,
                'low_flow_threshold': self.config.low_flow_threshold,
                'high_flow_threshold': self.config.high_flow_threshold
            },
            'feature_importance': self.feature_importance.to_dict() if self.feature_importance is not None else None
        }

    def save(self, path: str):
        """Save model to disk."""
        from dataclasses import asdict
        path = Path(path)
        path.mkdir(parents=True, exist_ok=True)

        # Save models
        joblib.dump(self.regressor, path / 'regressor.joblib')
        joblib.dump(self.classifier, path / 'classifier.joblib')

        # Convert config to dict for pickle compatibility
        config_dict = asdict(self.config)

        # Save metadata (with config as dict to avoid pickle issues)
        metadata = {
            'config_dict': config_dict,
            'training_metrics': self.training_metrics,
            'feature_importance': self.feature_importance.to_dict() if self.feature_importance is not None else None,
            'is_fitted': self.is_fitted
        }
        joblib.dump(metadata, path / 'metadata.joblib')

        logger.info(f"AttentionScoreModel saved to {path}")

    @classmethod
    def load(cls, path: str) -> 'AttentionScoreModel':
        """Load model from disk."""
        path = Path(path)

        # Load metadata
        metadata = joblib.load(path / 'metadata.joblib')

        # Reconstruct config from dict
        if 'config_dict' in metadata:
            config = AttentionModelConfig(**metadata['config_dict'])
        elif 'config' in metadata:
            # Legacy format - config was pickled directly
            config = metadata['config']
        else:
            config = AttentionModelConfig()

        # Create instance
        model = cls(config)
        model.training_metrics = metadata['training_metrics']
        model.is_fitted = metadata['is_fitted']

        if metadata['feature_importance']:
            model.feature_importance = pd.DataFrame(metadata['feature_importance'])

        # Load models
        model.regressor = joblib.load(path / 'regressor.joblib')
        model.classifier = joblib.load(path / 'classifier.joblib')

        logger.info(f"AttentionScoreModel loaded from {path}")
        return model


# Convenience functions for integration with existing pipeline

def generate_attention_features_for_df(df: pd.DataFrame) -> pd.DataFrame:
    """
    Convenience function to generate attention features.

    Args:
        df: DataFrame with OHLCV columns

    Returns:
        DataFrame with attention features added
    """
    generator = AttentionFeatureGenerator()
    features = generator.generate_features(df)
    return pd.concat([df, features], axis=1)


def get_attention_feature_names() -> List[str]:
    """Get list of attention feature names."""
    return AttentionModelConfig().feature_names


if __name__ == "__main__":
    # Test the module
    print("Testing AttentionScoreModel...")

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

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

    # Simulate varying volatility
    volatility = np.where(
        np.random.rand(n) > 0.7,  # 30% high volatility periods
        5.0,
        2.0
    )

    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) * (1 + (volatility > 3).astype(int))
    }, index=dates)

    # Test model
    config = AttentionModelConfig(
        min_train_samples=1000
    )

    model = AttentionScoreModel(config)
    model.fit(df)

    # Test prediction
    print("\nTesting predictions...")
    prediction = model.predict(df.iloc[-100:])

    print(f"Attention scores: min={prediction.attention_score.min():.2f}, "
          f"max={prediction.attention_score.max():.2f}, "
          f"mean={prediction.attention_score.mean():.2f}")

    print(f"\nFlow class distribution:")
    for i, name in enumerate(['low_flow', 'medium_flow', 'high_flow']):
        count = (prediction.flow_class == i).sum()
        print(f"  {name}: {count}")

    # Test save/load
    print("\nTesting save/load...")
    model.save('/tmp/test_attention_model')
    loaded_model = AttentionScoreModel.load('/tmp/test_attention_model')

    pred2 = loaded_model.predict(df.iloc[-10:])
    print(f"Loaded model predictions: {pred2.attention_score}")

    print("\nTest complete!")
