#!/usr/bin/env python3
"""
PVA (Price Variation Attention) Trainer
========================================
Training pipeline for the PVA strategy model.

This module provides comprehensive training functionality:
- Single symbol training
- Walk-forward validation
- Model versioning and saving
- Training metrics and logging

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

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

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

try:
    import torch
    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

from .model import PVAModel, PVAConfig, PVAPrediction
from .feature_engineering import PVAFeatureEngineer, PVAFeatureConfig
from ...data import TrainingDataLoader, TrainingDataConfig


@dataclass
class PVATrainerConfig:
    """Configuration for PVA model training."""

    # Data loading
    timeframe: str = '5m'
    batch_size: int = 64

    # Sequence parameters
    sequence_length: int = 100
    target_horizon: int = 12

    # Training parameters
    encoder_epochs: int = 50
    encoder_learning_rate: float = 1e-4
    early_stopping_patience: int = 10

    # Walk-forward validation
    walk_forward_splits: int = 5
    walk_forward_test_size: float = 0.2
    walk_forward_gap: int = 0
    min_train_size: int = 10000

    # Validation
    val_ratio: float = 0.15

    # Model saving
    model_base_dir: str = "models/pva"
    save_intermediate: bool = True

    # Feature engineering
    return_periods: List[int] = field(default_factory=lambda: [1, 5, 10, 20])
    volatility_window: int = 20
    stats_window: int = 50

    # Model architecture
    d_model: int = 256
    n_heads: int = 8
    n_layers: int = 4
    d_ff: int = 1024
    dropout: float = 0.1

    # XGBoost parameters
    xgb_n_estimators: int = 200
    xgb_max_depth: int = 6
    xgb_learning_rate: float = 0.05

    # Device
    device: str = "cpu"


@dataclass
class TrainingMetrics:
    """Metrics from a training run."""

    symbol: str
    timeframe: str
    version: str
    train_samples: int
    val_samples: int
    encoder_train_loss: float
    encoder_val_loss: float
    direction_accuracy: float
    magnitude_mse: float
    magnitude_mae: float
    training_time_seconds: float
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            'symbol': self.symbol,
            'timeframe': self.timeframe,
            'version': self.version,
            'train_samples': self.train_samples,
            'val_samples': self.val_samples,
            'encoder_train_loss': self.encoder_train_loss,
            'encoder_val_loss': self.encoder_val_loss,
            'direction_accuracy': self.direction_accuracy,
            'magnitude_mse': self.magnitude_mse,
            'magnitude_mae': self.magnitude_mae,
            'training_time_seconds': self.training_time_seconds,
            'timestamp': self.timestamp
        }


@dataclass
class WalkForwardResult:
    """Results from walk-forward validation."""

    symbol: str
    n_folds: int
    fold_metrics: List[Dict[str, Any]]
    avg_direction_accuracy: float
    std_direction_accuracy: float
    avg_magnitude_mae: float
    std_magnitude_mae: float
    total_training_time: float

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            'symbol': self.symbol,
            'n_folds': self.n_folds,
            'fold_metrics': self.fold_metrics,
            'avg_direction_accuracy': self.avg_direction_accuracy,
            'std_direction_accuracy': self.std_direction_accuracy,
            'avg_magnitude_mae': self.avg_magnitude_mae,
            'std_magnitude_mae': self.std_magnitude_mae,
            'total_training_time': self.total_training_time
        }


class PVATrainer:
    """
    Trainer for PVA (Price Variation Attention) models.

    Handles the complete training pipeline:
    1. Data loading from PostgreSQL
    2. Feature engineering
    3. Model training (encoder + XGBoost head)
    4. Validation and metrics
    5. Model versioning and saving

    Usage:
        trainer = PVATrainer()

        # Train model for a single symbol
        model, metrics = trainer.train('XAUUSD', '2023-01-01', '2024-12-31')

        # Walk-forward validation
        results = trainer.walk_forward_train('XAUUSD', n_folds=5)

        # Save trained model
        trainer.save_model(model, 'XAUUSD', 'v1.0.0')
    """

    def __init__(
        self,
        config: Optional[PVATrainerConfig] = None,
        data_loader: Optional[TrainingDataLoader] = None
    ):
        """
        Initialize the PVA trainer.

        Args:
            config: Trainer configuration
            data_loader: Data loader for training data (creates new if None)
        """
        self.config = config or PVATrainerConfig()
        self.data_loader = data_loader or TrainingDataLoader()

        # Initialize feature engineer
        feature_config = PVAFeatureConfig(
            return_periods=self.config.return_periods,
            volatility_window=self.config.volatility_window,
            stats_window=self.config.stats_window,
            sequence_length=self.config.sequence_length
        )
        self.feature_engineer = PVAFeatureEngineer(feature_config)

        # Training history
        self._training_history: List[TrainingMetrics] = []

        logger.info("PVATrainer initialized")
        logger.info(f"  Timeframe: {self.config.timeframe}")
        logger.info(f"  Sequence length: {self.config.sequence_length}")
        logger.info(f"  Target horizon: {self.config.target_horizon}")

    def train(
        self,
        symbol: str,
        start_date: str,
        end_date: str,
        model_config: Optional[PVAConfig] = None
    ) -> Tuple[PVAModel, TrainingMetrics]:
        """
        Train a PVA model for a single symbol.

        Args:
            symbol: Trading symbol (e.g., 'XAUUSD')
            start_date: Training start date (YYYY-MM-DD)
            end_date: Training end date (YYYY-MM-DD)
            model_config: Model configuration (uses trainer config if None)

        Returns:
            Tuple of (trained model, training metrics)
        """
        import time
        start_time = time.time()

        logger.info(f"Training PVA model for {symbol}")
        logger.info(f"  Date range: {start_date} to {end_date}")

        # Step 1: Load data
        logger.info("Loading training data...")
        df = self.data_loader.get_training_data(
            symbol=symbol,
            start_date=start_date,
            end_date=end_date,
            timeframe=self.config.timeframe,
            include_features=False  # We'll compute our own features
        )

        if df.empty:
            raise ValueError(f"No data found for {symbol} in date range")

        logger.info(f"  Loaded {len(df):,} rows")

        # Step 2: Prepare training data
        logger.info("Preparing features and sequences...")
        X_train, y_train, X_val, y_val = self.feature_engineer.prepare_training_data(
            df,
            target_horizon=self.config.target_horizon,
            target_type='return',
            seq_length=self.config.sequence_length,
            val_ratio=self.config.val_ratio
        )

        n_features = X_train.shape[2]
        logger.info(f"  Train samples: {len(X_train):,}")
        logger.info(f"  Val samples: {len(X_val):,}")
        logger.info(f"  Features: {n_features}")

        # Step 3: Create model
        if model_config is None:
            model_config = PVAConfig(
                input_features=n_features,
                sequence_length=self.config.sequence_length,
                d_model=self.config.d_model,
                n_heads=self.config.n_heads,
                n_layers=self.config.n_layers,
                d_ff=self.config.d_ff,
                dropout=self.config.dropout,
                xgb_n_estimators=self.config.xgb_n_estimators,
                xgb_max_depth=self.config.xgb_max_depth,
                xgb_learning_rate=self.config.xgb_learning_rate,
                device=self.config.device
            )

        model = PVAModel(model_config)

        # Step 4: Train encoder
        logger.info("Training attention encoder...")
        encoder_history = model.fit_encoder(
            X_train, y_train, X_val, y_val,
            epochs=self.config.encoder_epochs,
            batch_size=self.config.batch_size,
            learning_rate=self.config.encoder_learning_rate,
            early_stopping_patience=self.config.early_stopping_patience
        )

        # Step 5: Train XGBoost head
        logger.info("Training XGBoost prediction head...")
        xgb_metrics = model.fit_xgboost(X_train, y_train, X_val, y_val)

        # Step 6: Compile metrics
        training_time = time.time() - start_time

        metrics = TrainingMetrics(
            symbol=symbol,
            timeframe=self.config.timeframe,
            version=datetime.now().strftime('%Y%m%d_%H%M%S'),
            train_samples=len(X_train),
            val_samples=len(X_val),
            encoder_train_loss=encoder_history['train_loss'][-1],
            encoder_val_loss=encoder_history['val_loss'][-1] if encoder_history['val_loss'] else 0,
            direction_accuracy=xgb_metrics.get('val_direction_accuracy', 0),
            magnitude_mse=xgb_metrics.get('val_magnitude_mse', 0),
            magnitude_mae=xgb_metrics.get('val_magnitude_mae', 0),
            training_time_seconds=training_time
        )

        self._training_history.append(metrics)

        logger.info(f"Training complete in {training_time:.1f}s")
        logger.info(f"  Direction accuracy: {metrics.direction_accuracy:.4f}")
        logger.info(f"  Magnitude MAE: {metrics.magnitude_mae:.6f}")

        return model, metrics

    def validate(
        self,
        model: PVAModel,
        val_data: pd.DataFrame
    ) -> Dict[str, float]:
        """
        Validate a trained model on held-out data.

        Args:
            model: Trained PVA model
            val_data: Validation DataFrame with OHLCV data

        Returns:
            Dictionary of validation metrics
        """
        from sklearn.metrics import accuracy_score, mean_squared_error, mean_absolute_error

        logger.info(f"Validating model on {len(val_data)} rows...")

        # Prepare validation data
        X_val, y_val = self.feature_engineer.prepare_training_data(
            val_data,
            target_horizon=self.config.target_horizon,
            target_type='return',
            seq_length=self.config.sequence_length,
            val_ratio=0.0  # Use all data for validation
        )

        if len(X_val) == 0:
            logger.warning("No valid validation samples")
            return {'error': 'No valid samples'}

        # Make predictions
        predictions = model.predict(X_val)

        # Extract values
        pred_directions = np.array([p.direction for p in predictions])
        pred_magnitudes = np.array([p.magnitude for p in predictions])
        pred_confidences = np.array([p.confidence for p in predictions])

        # True values
        true_directions = (y_val > 0).astype(int)
        true_magnitudes = np.abs(y_val)

        # Binary direction predictions
        pred_direction_binary = (pred_directions > 0).astype(int)

        # Compute metrics
        metrics = {
            'direction_accuracy': accuracy_score(true_directions, pred_direction_binary),
            'magnitude_mse': mean_squared_error(true_magnitudes, pred_magnitudes),
            'magnitude_mae': mean_absolute_error(true_magnitudes, pred_magnitudes),
            'avg_confidence': float(np.mean(pred_confidences)),
            'n_samples': len(X_val)
        }

        # Directional return metric (profit factor proxy)
        signed_returns = pred_directions * y_val
        metrics['avg_directional_return'] = float(np.mean(signed_returns))
        metrics['sharpe_proxy'] = float(np.mean(signed_returns) / (np.std(signed_returns) + 1e-10))

        logger.info(f"Validation metrics: {metrics}")
        return metrics

    def walk_forward_train(
        self,
        symbol: str,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
        n_folds: int = 5
    ) -> WalkForwardResult:
        """
        Perform walk-forward validation training.

        Walk-forward validation trains multiple models on expanding/sliding
        windows of data, validating on subsequent periods. This provides
        a more realistic estimate of model performance.

        Args:
            symbol: Trading symbol
            start_date: Start date (defaults to 3 years ago)
            end_date: End date (defaults to today)
            n_folds: Number of walk-forward folds

        Returns:
            WalkForwardResult with metrics for each fold
        """
        import time
        from datetime import timedelta

        total_start = time.time()

        if start_date is None:
            start_date = (datetime.now() - timedelta(days=365 * 3)).strftime('%Y-%m-%d')
        if end_date is None:
            end_date = datetime.now().strftime('%Y-%m-%d')

        logger.info(f"Walk-forward training for {symbol}")
        logger.info(f"  Date range: {start_date} to {end_date}")
        logger.info(f"  Folds: {n_folds}")

        # Load all data
        df = self.data_loader.get_training_data(
            symbol=symbol,
            start_date=start_date,
            end_date=end_date,
            timeframe=self.config.timeframe,
            include_features=False
        )

        if df.empty:
            raise ValueError(f"No data found for {symbol}")

        n_samples = len(df)
        logger.info(f"  Total samples: {n_samples:,}")

        # Calculate fold sizes
        step_size = n_samples // (n_folds + 1)
        test_size = int(step_size * self.config.walk_forward_test_size)

        if step_size < self.config.min_train_size:
            logger.warning(f"Step size ({step_size}) < min train size ({self.config.min_train_size})")
            n_folds = max(1, n_samples // self.config.min_train_size - 1)
            step_size = n_samples // (n_folds + 1)
            test_size = int(step_size * self.config.walk_forward_test_size)
            logger.info(f"  Adjusted to {n_folds} folds")

        # Train models for each fold
        fold_metrics = []
        fold_models = []

        for fold_idx in range(n_folds):
            logger.info(f"\n{'=' * 60}")
            logger.info(f"FOLD {fold_idx + 1}/{n_folds}")
            logger.info(f"{'=' * 60}")

            # Calculate indices
            train_start = 0  # Expanding window
            train_end = (fold_idx + 1) * step_size
            val_start = train_end + self.config.walk_forward_gap
            val_end = min(val_start + test_size, n_samples)

            if val_end > n_samples or (train_end - train_start) < self.config.min_train_size:
                logger.warning(f"Skipping fold {fold_idx + 1}: insufficient data")
                continue

            # Split data
            train_df = df.iloc[train_start:train_end].copy()
            val_df = df.iloc[val_start:val_end].copy()

            logger.info(f"  Train: {len(train_df):,} samples")
            logger.info(f"  Val: {len(val_df):,} samples")

            # Prepare training data
            X_train, y_train = self.feature_engineer.prepare_training_data(
                train_df,
                target_horizon=self.config.target_horizon,
                target_type='return',
                seq_length=self.config.sequence_length,
                val_ratio=0.0
            )

            # Prepare validation data
            X_val, y_val = self.feature_engineer.prepare_training_data(
                val_df,
                target_horizon=self.config.target_horizon,
                target_type='return',
                seq_length=self.config.sequence_length,
                val_ratio=0.0
            )

            if len(X_train) == 0 or len(X_val) == 0:
                logger.warning(f"Skipping fold {fold_idx + 1}: no valid sequences")
                continue

            # Create and train model
            n_features = X_train.shape[2]
            model_config = PVAConfig(
                input_features=n_features,
                sequence_length=self.config.sequence_length,
                d_model=self.config.d_model,
                n_heads=self.config.n_heads,
                n_layers=self.config.n_layers,
                d_ff=self.config.d_ff,
                dropout=self.config.dropout,
                xgb_n_estimators=self.config.xgb_n_estimators,
                xgb_max_depth=self.config.xgb_max_depth,
                xgb_learning_rate=self.config.xgb_learning_rate,
                device=self.config.device
            )

            model = PVAModel(model_config)

            # Split train data for encoder validation
            split_idx = int(len(X_train) * 0.85)
            X_train_enc, X_val_enc = X_train[:split_idx], X_train[split_idx:]
            y_train_enc, y_val_enc = y_train[:split_idx], y_train[split_idx:]

            # Train encoder
            fold_start = time.time()
            encoder_history = model.fit_encoder(
                X_train_enc, y_train_enc, X_val_enc, y_val_enc,
                epochs=self.config.encoder_epochs,
                batch_size=self.config.batch_size,
                learning_rate=self.config.encoder_learning_rate,
                early_stopping_patience=self.config.early_stopping_patience
            )

            # Train XGBoost
            xgb_metrics = model.fit_xgboost(X_train, y_train, X_val, y_val)
            fold_time = time.time() - fold_start

            # Validate on test set
            val_metrics = self._compute_fold_metrics(model, X_val, y_val)

            fold_result = {
                'fold': fold_idx + 1,
                'train_size': len(X_train),
                'val_size': len(X_val),
                'encoder_train_loss': encoder_history['train_loss'][-1],
                'encoder_val_loss': encoder_history['val_loss'][-1] if encoder_history['val_loss'] else 0,
                'direction_accuracy': val_metrics['direction_accuracy'],
                'magnitude_mse': val_metrics['magnitude_mse'],
                'magnitude_mae': val_metrics['magnitude_mae'],
                'directional_return': val_metrics['avg_directional_return'],
                'training_time': fold_time
            }

            fold_metrics.append(fold_result)
            fold_models.append(model)

            logger.info(f"Fold {fold_idx + 1} complete:")
            logger.info(f"  Direction accuracy: {fold_result['direction_accuracy']:.4f}")
            logger.info(f"  Magnitude MAE: {fold_result['magnitude_mae']:.6f}")

            # Save intermediate model if configured
            if self.config.save_intermediate:
                model_dir = Path(self.config.model_base_dir) / symbol / f"fold_{fold_idx + 1}"
                model.save(str(model_dir))

        # Compute aggregate metrics
        if not fold_metrics:
            raise RuntimeError("No folds completed successfully")

        direction_accuracies = [f['direction_accuracy'] for f in fold_metrics]
        magnitude_maes = [f['magnitude_mae'] for f in fold_metrics]

        total_time = time.time() - total_start

        result = WalkForwardResult(
            symbol=symbol,
            n_folds=len(fold_metrics),
            fold_metrics=fold_metrics,
            avg_direction_accuracy=float(np.mean(direction_accuracies)),
            std_direction_accuracy=float(np.std(direction_accuracies)),
            avg_magnitude_mae=float(np.mean(magnitude_maes)),
            std_magnitude_mae=float(np.std(magnitude_maes)),
            total_training_time=total_time
        )

        logger.info(f"\n{'=' * 60}")
        logger.info("WALK-FORWARD VALIDATION COMPLETE")
        logger.info(f"{'=' * 60}")
        logger.info(f"  Folds completed: {result.n_folds}")
        logger.info(f"  Avg direction accuracy: {result.avg_direction_accuracy:.4f} +/- {result.std_direction_accuracy:.4f}")
        logger.info(f"  Avg magnitude MAE: {result.avg_magnitude_mae:.6f} +/- {result.std_magnitude_mae:.6f}")
        logger.info(f"  Total time: {total_time:.1f}s")

        return result

    def _compute_fold_metrics(
        self,
        model: PVAModel,
        X_val: np.ndarray,
        y_val: np.ndarray
    ) -> Dict[str, float]:
        """Compute metrics for a single fold."""
        from sklearn.metrics import accuracy_score, mean_squared_error, mean_absolute_error

        predictions = model.predict(X_val)

        pred_directions = np.array([p.direction for p in predictions])
        pred_magnitudes = np.array([p.magnitude for p in predictions])

        true_directions = (y_val > 0).astype(int)
        true_magnitudes = np.abs(y_val)
        pred_direction_binary = (pred_directions > 0).astype(int)

        metrics = {
            'direction_accuracy': accuracy_score(true_directions, pred_direction_binary),
            'magnitude_mse': mean_squared_error(true_magnitudes, pred_magnitudes),
            'magnitude_mae': mean_absolute_error(true_magnitudes, pred_magnitudes),
            'n_samples': len(X_val)
        }

        signed_returns = pred_directions * y_val
        metrics['avg_directional_return'] = float(np.mean(signed_returns))

        return metrics

    def save_model(
        self,
        model: PVAModel,
        symbol: str,
        version: str
    ) -> str:
        """
        Save a trained model with version tracking.

        Args:
            model: Trained PVA model
            symbol: Symbol the model was trained on
            version: Version string (e.g., 'v1.0.0')

        Returns:
            Path where model was saved
        """
        model_dir = Path(self.config.model_base_dir) / symbol / version
        model_dir.mkdir(parents=True, exist_ok=True)

        # Save model
        model.save(str(model_dir))

        # Save training metadata
        metadata = {
            'symbol': symbol,
            'version': version,
            'timeframe': self.config.timeframe,
            'sequence_length': self.config.sequence_length,
            'target_horizon': self.config.target_horizon,
            'feature_config': {
                'return_periods': self.config.return_periods,
                'volatility_window': self.config.volatility_window,
                'stats_window': self.config.stats_window
            },
            'trainer_config': {
                'd_model': self.config.d_model,
                'n_heads': self.config.n_heads,
                'n_layers': self.config.n_layers,
                'd_ff': self.config.d_ff
            },
            'saved_at': datetime.now().isoformat()
        }

        metadata_path = model_dir / 'metadata.json'
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=2)

        # Save feature names
        feature_names_path = model_dir / 'feature_names.json'
        with open(feature_names_path, 'w') as f:
            json.dump(self.feature_engineer.get_feature_names(), f)

        logger.info(f"Model saved to {model_dir}")
        return str(model_dir)

    def load_model(
        self,
        symbol: str,
        version: str,
        device: str = "cpu"
    ) -> PVAModel:
        """
        Load a saved model.

        Args:
            symbol: Symbol the model was trained on
            version: Version string
            device: Device to load model to

        Returns:
            Loaded PVAModel
        """
        model_dir = Path(self.config.model_base_dir) / symbol / version
        return PVAModel.load(str(model_dir), device=device)

    def get_training_history(self) -> List[Dict[str, Any]]:
        """Get history of all training runs."""
        return [m.to_dict() for m in self._training_history]


def train_pva_model(
    symbol: str,
    start_date: str,
    end_date: str,
    save_path: Optional[str] = None
) -> Tuple[PVAModel, Dict[str, Any]]:
    """
    Convenience function to train a PVA model.

    Args:
        symbol: Trading symbol
        start_date: Training start date
        end_date: Training end date
        save_path: Optional path to save model

    Returns:
        Tuple of (trained model, metrics dict)
    """
    trainer = PVATrainer()
    model, metrics = trainer.train(symbol, start_date, end_date)

    if save_path:
        trainer.save_model(model, symbol, datetime.now().strftime('%Y%m%d_%H%M%S'))

    return model, metrics.to_dict()


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

    # Note: This requires database connection
    # For standalone testing, we'll create synthetic data

    np.random.seed(42)
    n = 10000

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

    df = pd.DataFrame({
        'open': price,
        'high': price + np.abs(np.random.randn(n)) * 2,
        'low': price - np.abs(np.random.randn(n)) * 2,
        'close': price + np.random.randn(n) * 0.5,
        'volume': np.random.randint(100, 1000, n)
    }, index=dates)

    # Create trainer with smaller config for testing
    config = PVATrainerConfig(
        sequence_length=50,
        target_horizon=6,
        encoder_epochs=5,
        d_model=64,
        n_heads=4,
        n_layers=2,
        walk_forward_splits=3,
        min_train_size=1000,
        val_ratio=0.15
    )

    # Create mock data loader
    class MockDataLoader:
        def get_training_data(self, **kwargs):
            return df.copy()

    trainer = PVATrainer(config=config, data_loader=MockDataLoader())

    # Test feature engineering
    print("\nTesting feature engineering...")
    X, y = trainer.feature_engineer.prepare_training_data(
        df, target_horizon=6, seq_length=50
    )
    print(f"  X shape: {X.shape}")
    print(f"  y shape: {y.shape}")

    # Test walk-forward (abbreviated)
    print("\nTesting walk-forward validation (abbreviated)...")
    result = trainer.walk_forward_train(
        symbol='TEST',
        n_folds=2
    )
    print(f"  Folds: {result.n_folds}")
    print(f"  Avg accuracy: {result.avg_direction_accuracy:.4f}")

    print("\nTest complete!")
