#!/usr/bin/env python3
"""
Attention Model Trainer
=======================
Trains attention score models for multiple symbols and timeframes.

This trainer:
1. Loads OHLCV data from database
2. Generates attention features
3. Trains regression + classification models
4. Saves models per symbol/timeframe
5. Generates training reports

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

import sys
from pathlib import Path
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field
import json

import numpy as np
import pandas as pd
from loguru import logger
import joblib

# Direct import to avoid loading models/__init__.py (which has circular imports)
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:
    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)
else:
    attention_module = sys.modules[_module_name]

AttentionScoreModel = attention_module.AttentionScoreModel
AttentionModelConfig = attention_module.AttentionModelConfig
AttentionFeatureGenerator = attention_module.AttentionFeatureGenerator
AttentionPrediction = attention_module.AttentionPrediction


@dataclass
class AttentionTrainerConfig:
    """Configuration for attention model training"""

    # Symbols to train
    symbols: List[str] = field(default_factory=lambda: [
        'XAUUSD', 'EURUSD', 'BTCUSD', 'GBPUSD', 'USDJPY'
    ])

    # Timeframes to train
    timeframes: List[str] = field(default_factory=lambda: ['5m', '15m'])

    # Training period
    train_years: float = 5.0
    holdout_years: float = 1.0  # Last year for validation

    # Model configuration
    model_config: AttentionModelConfig = field(default_factory=AttentionModelConfig)

    # Output directory
    output_dir: str = 'models/attention'


@dataclass
class AttentionTrainingResult:
    """Result of training attention model for one symbol/timeframe"""
    symbol: str
    timeframe: str
    reg_mae: float
    reg_rmse: float
    reg_r2: float
    clf_accuracy: float
    clf_f1: float
    n_train: int
    n_val: int
    class_distribution: Dict[str, int]
    feature_importance: Dict[str, float]
    train_date_range: str
    holdout_date_range: str


class AttentionModelTrainer:
    """
    Trainer for attention score models.

    Trains one model per symbol/timeframe combination.

    Usage:
        trainer = AttentionModelTrainer(config)
        results = trainer.train_all(data_dict)
        trainer.save('models/attention/')
    """

    def __init__(self, config: AttentionTrainerConfig = None):
        self.config = config or AttentionTrainerConfig()
        self.models: Dict[str, AttentionScoreModel] = {}
        self.results: Dict[str, AttentionTrainingResult] = {}

        logger.info("AttentionModelTrainer initialized")
        logger.info(f"  Symbols: {self.config.symbols}")
        logger.info(f"  Timeframes: {self.config.timeframes}")

    def _get_model_key(self, symbol: str, timeframe: str) -> str:
        """Generate unique key for model."""
        return f"{symbol}_{timeframe}_attention"

    def _split_train_holdout(
        self,
        df: pd.DataFrame
    ) -> Tuple[pd.DataFrame, pd.DataFrame]:
        """
        Split data into training and holdout sets.

        Args:
            df: DataFrame with datetime index

        Returns:
            Tuple of (train_df, holdout_df)
        """
        if isinstance(df.index, pd.DatetimeIndex):
            max_date = df.index.max()
            min_date = df.index.min()
        else:
            raise ValueError("DataFrame must have DatetimeIndex")

        # Holdout = last N years
        holdout_start = max_date - timedelta(days=self.config.holdout_years * 365)

        # Training = everything before holdout (up to train_years)
        train_start = holdout_start - timedelta(days=self.config.train_years * 365)
        train_start = max(train_start, min_date)

        train_mask = (df.index >= train_start) & (df.index < holdout_start)
        holdout_mask = df.index >= holdout_start

        train_df = df[train_mask].copy()
        holdout_df = df[holdout_mask].copy()

        logger.info(f"Data split:")
        logger.info(f"  Training: {train_start.strftime('%Y-%m-%d')} to {holdout_start.strftime('%Y-%m-%d')} ({len(train_df)} samples)")
        logger.info(f"  Holdout: {holdout_start.strftime('%Y-%m-%d')} to {max_date.strftime('%Y-%m-%d')} ({len(holdout_df)} samples)")

        return train_df, holdout_df

    def train_single(
        self,
        df: pd.DataFrame,
        symbol: str,
        timeframe: str
    ) -> AttentionTrainingResult:
        """
        Train attention model for a single symbol/timeframe.

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

        Returns:
            AttentionTrainingResult with metrics
        """
        key = self._get_model_key(symbol, timeframe)
        logger.info(f"\n{'='*60}")
        logger.info(f"Training {key}")
        logger.info(f"{'='*60}")

        # Split train/holdout
        train_df, holdout_df = self._split_train_holdout(df)

        if len(train_df) < self.config.model_config.min_train_samples:
            logger.warning(f"Insufficient training data: {len(train_df)}")
            return None

        # Create and train model
        model = AttentionScoreModel(self.config.model_config)
        model.fit(train_df)

        # Evaluate on holdout
        holdout_features = model.feature_generator.generate_features(holdout_df)
        holdout_prediction = model.predict(holdout_df, holdout_features)

        # Compute holdout targets for comparison
        move_multiplier, valid_mask = model.feature_generator.compute_target_move_multiplier(
            holdout_df,
            horizon_bars=self.config.model_config.horizon_bars,
            factor_window=self.config.model_config.factor_window
        )

        # Calculate holdout metrics
        valid_indices = valid_mask & ~np.isnan(holdout_features[self.config.model_config.feature_names].values).any(axis=1)

        if valid_indices.sum() > 0:
            y_true_reg = move_multiplier[valid_indices]
            y_pred_reg = holdout_prediction.attention_score[valid_indices]

            from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score

            holdout_mae = mean_absolute_error(y_true_reg, y_pred_reg)
            holdout_rmse = np.sqrt(mean_squared_error(y_true_reg, y_pred_reg))
            holdout_r2 = r2_score(y_true_reg, y_pred_reg)

            logger.info(f"\nHoldout performance:")
            logger.info(f"  MAE: {holdout_mae:.4f}")
            logger.info(f"  RMSE: {holdout_rmse:.4f}")
            logger.info(f"  R2: {holdout_r2:.4f}")
        else:
            holdout_mae = model.training_metrics['reg_mae']
            holdout_rmse = model.training_metrics['reg_rmse']
            holdout_r2 = model.training_metrics['reg_r2']

        # Store model
        self.models[key] = model

        # Class distribution
        class_names = ['low_flow', 'medium_flow', 'high_flow']
        class_dist = {
            name: int((holdout_prediction.flow_class == i).sum())
            for i, name in enumerate(class_names)
        }

        # Feature importance
        feat_imp = dict(zip(
            model.feature_importance['feature'].tolist(),
            model.feature_importance['combined'].tolist()
        ))

        # Create result
        result = AttentionTrainingResult(
            symbol=symbol,
            timeframe=timeframe,
            reg_mae=holdout_mae,
            reg_rmse=holdout_rmse,
            reg_r2=holdout_r2,
            clf_accuracy=model.training_metrics['clf_accuracy'],
            clf_f1=model.training_metrics['clf_f1'],
            n_train=model.training_metrics['n_train'],
            n_val=model.training_metrics['n_val'],
            class_distribution=class_dist,
            feature_importance=feat_imp,
            train_date_range=f"{train_df.index.min()} to {train_df.index.max()}",
            holdout_date_range=f"{holdout_df.index.min()} to {holdout_df.index.max()}"
        )

        self.results[key] = result

        return result

    def train_all(
        self,
        data_dict: Dict[str, Dict[str, pd.DataFrame]]
    ) -> Dict[str, AttentionTrainingResult]:
        """
        Train attention models for all symbols and timeframes.

        Args:
            data_dict: Dictionary structured as {symbol: {timeframe: df}}
                       e.g., {'XAUUSD': {'5m': df_5m, '15m': df_15m}}

        Returns:
            Dictionary of all AttentionTrainingResult objects
        """
        all_results = {}

        for symbol in self.config.symbols:
            if symbol not in data_dict:
                logger.warning(f"No data for symbol: {symbol}")
                continue

            for timeframe in self.config.timeframes:
                if timeframe not in data_dict[symbol]:
                    logger.warning(f"No {timeframe} data for {symbol}")
                    continue

                df = data_dict[symbol][timeframe]
                result = self.train_single(df, symbol, timeframe)

                if result:
                    key = self._get_model_key(symbol, timeframe)
                    all_results[key] = result

        return all_results

    def predict(
        self,
        df: pd.DataFrame,
        symbol: str,
        timeframe: str
    ) -> AttentionPrediction:
        """
        Generate attention prediction for a symbol/timeframe.

        Args:
            df: DataFrame with OHLCV
            symbol: Symbol
            timeframe: Timeframe

        Returns:
            AttentionPrediction
        """
        key = self._get_model_key(symbol, timeframe)

        if key not in self.models:
            raise ValueError(f"No trained model for {key}")

        return self.models[key].predict(df)

    def get_attention_score(
        self,
        df: pd.DataFrame,
        symbol: str,
        timeframe: str
    ) -> np.ndarray:
        """
        Get just the attention score (for use as feature in other models).

        Args:
            df: DataFrame with OHLCV
            symbol: Symbol
            timeframe: Timeframe

        Returns:
            Array of attention scores
        """
        prediction = self.predict(df, symbol, timeframe)
        return prediction.attention_score

    def get_training_summary(self) -> pd.DataFrame:
        """Get summary of all trained models as DataFrame."""
        rows = []
        for key, result in self.results.items():
            rows.append({
                'model_key': key,
                'symbol': result.symbol,
                'timeframe': result.timeframe,
                'reg_mae': result.reg_mae,
                'reg_rmse': result.reg_rmse,
                'reg_r2': result.reg_r2,
                'clf_accuracy': result.clf_accuracy,
                'clf_f1': result.clf_f1,
                'n_train': result.n_train,
                'n_val': result.n_val,
                'low_flow_pct': result.class_distribution.get('low_flow', 0) /
                               (sum(result.class_distribution.values()) + 1) * 100,
                'high_flow_pct': result.class_distribution.get('high_flow', 0) /
                                (sum(result.class_distribution.values()) + 1) * 100
            })

        return pd.DataFrame(rows)

    def save(self, path: str = None):
        """Save all models and metadata to disk."""
        path = Path(path or self.config.output_dir)
        path.mkdir(parents=True, exist_ok=True)

        # Save each model
        for key, model in self.models.items():
            model_path = path / key
            model.save(str(model_path))

        # Save trainer metadata
        metadata = {
            'config': {
                'symbols': self.config.symbols,
                'timeframes': self.config.timeframes,
                'train_years': self.config.train_years,
                'holdout_years': self.config.holdout_years
            },
            'model_keys': list(self.models.keys()),
            'results': {
                k: {
                    'symbol': v.symbol,
                    'timeframe': v.timeframe,
                    'reg_mae': v.reg_mae,
                    'reg_r2': v.reg_r2,
                    'clf_accuracy': v.clf_accuracy,
                    'class_distribution': v.class_distribution
                }
                for k, v in self.results.items()
            },
            'trained_at': datetime.now().isoformat()
        }
        joblib.dump(metadata, path / 'trainer_metadata.joblib')

        # Save summary CSV
        summary = self.get_training_summary()
        if not summary.empty:
            summary.to_csv(path / 'training_summary.csv', index=False)

        logger.info(f"Saved {len(self.models)} attention models to {path}")

    @classmethod
    def load(cls, path: str) -> 'AttentionModelTrainer':
        """Load trainer and all models from disk."""
        path = Path(path)

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

        # Create config
        config = AttentionTrainerConfig(
            symbols=metadata['config']['symbols'],
            timeframes=metadata['config']['timeframes'],
            train_years=metadata['config']['train_years'],
            holdout_years=metadata['config']['holdout_years']
        )

        trainer = cls(config)

        # Load models
        for key in metadata['model_keys']:
            model_path = path / key
            if model_path.exists():
                trainer.models[key] = AttentionScoreModel.load(str(model_path))

        logger.info(f"Loaded {len(trainer.models)} attention models from {path}")
        return trainer


def generate_attention_training_report(
    trainer: AttentionModelTrainer,
    output_path: str
) -> str:
    """
    Generate markdown report of attention model training.

    Args:
        trainer: Trained AttentionModelTrainer
        output_path: Path to save report

    Returns:
        Path to generated report
    """
    report = f"""# Attention Score Model Training Report

**Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

## Configuration

- **Symbols:** {', '.join(trainer.config.symbols)}
- **Timeframes:** {', '.join(trainer.config.timeframes)}
- **Training Years:** {trainer.config.train_years}
- **Holdout Years:** {trainer.config.holdout_years}

## Training Results Summary

| Model | Symbol | TF | MAE | RMSE | R2 | Clf Acc | High Flow % |
|-------|--------|-----|-----|------|-----|---------|-------------|
"""

    for key, result in trainer.results.items():
        total_samples = sum(result.class_distribution.values())
        high_pct = result.class_distribution.get('high_flow', 0) / (total_samples + 1) * 100

        report += f"| {key} | {result.symbol} | {result.timeframe} | "
        report += f"{result.reg_mae:.4f} | {result.reg_rmse:.4f} | {result.reg_r2:.4f} | "
        report += f"{result.clf_accuracy:.2%} | {high_pct:.1f}% |\n"

    report += """

## Feature Importance (Combined)

"""

    # Show top features per model
    for key, result in trainer.results.items():
        report += f"\n### {key}\n\n"
        report += "| Feature | Importance |\n|---------|------------|\n"

        sorted_features = sorted(result.feature_importance.items(), key=lambda x: -x[1])
        for feat, imp in sorted_features[:5]:
            report += f"| {feat} | {imp:.4f} |\n"

    report += """

## Class Distribution

| Model | Low Flow | Medium Flow | High Flow |
|-------|----------|-------------|-----------|
"""

    for key, result in trainer.results.items():
        low = result.class_distribution.get('low_flow', 0)
        med = result.class_distribution.get('medium_flow', 0)
        high = result.class_distribution.get('high_flow', 0)
        total = low + med + high

        report += f"| {key} | {low} ({low/total*100:.1f}%) | {med} ({med/total*100:.1f}%) | {high} ({high/total*100:.1f}%) |\n"

    report += """

## Interpretation

- **attention_score**: Continuous value (0-3+). Higher = expect larger market movement.
  - < 1.0: Low flow, avoid trading
  - 1.0-2.0: Medium flow, standard setups
  - > 2.0: High flow, best opportunities

- **flow_class**: Categorical classification
  - 0 = low_flow: Expect move_multiplier < 1.0
  - 1 = medium_flow: Expect move_multiplier 1.0-2.0
  - 2 = high_flow: Expect move_multiplier >= 2.0

## Usage

```python
from training.attention_trainer import AttentionModelTrainer

# Load trained models
trainer = AttentionModelTrainer.load('models/attention/')

# Get attention score for new data
attention = trainer.get_attention_score(df_ohlcv, 'XAUUSD', '5m')

# Use as feature in other models
df['attention_score'] = attention
```

---
*Report generated by AttentionModelTrainer*
"""

    with open(output_path, 'w') as f:
        f.write(report)

    logger.info(f"Report saved to {output_path}")
    return output_path


if __name__ == "__main__":
    # Test the trainer
    print("Testing AttentionModelTrainer...")

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

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

    volatility = np.where(np.random.rand(n) > 0.7, 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)

    # Create data dict
    data_dict = {
        'XAUUSD': {'5m': df}
    }

    # Test trainer
    config = AttentionTrainerConfig(
        symbols=['XAUUSD'],
        timeframes=['5m'],
        train_years=1.0,
        holdout_years=0.5
    )
    config.model_config.min_train_samples = 1000

    trainer = AttentionModelTrainer(config)
    results = trainer.train_all(data_dict)

    print("\nTraining summary:")
    print(trainer.get_training_summary())

    # Test save/load
    print("\nTesting save/load...")
    trainer.save('/tmp/test_attention_trainer')
    loaded_trainer = AttentionModelTrainer.load('/tmp/test_attention_trainer')

    # Test prediction
    attention = loaded_trainer.get_attention_score(df.iloc[-100:], 'XAUUSD', '5m')
    print(f"\nAttention scores: {attention[:5]}")

    print("\nTest complete!")
