#!/usr/bin/env python3
"""
Train Neural Gating Metamodel and compare with XGBoost Stacking.

This script:
1. Loads the same training data used for XGBoost metamodels
2. Trains Neural Gating version
3. Compares performance metrics

Usage:
    python scripts/train_neural_gating.py --symbols XAUUSD,EURUSD
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import argparse
import numpy as np
import pandas as pd
from pathlib import Path
from datetime import datetime
from loguru import logger
import joblib

# Configure logging
logger.remove()
logger.add(sys.stdout, level="INFO", format="{time:HH:mm:ss} | {level} | {message}")


def main():
    parser = argparse.ArgumentParser(description='Train Neural Gating Metamodel')
    parser.add_argument('--symbols', type=str, default='XAUUSD,EURUSD',
                        help='Comma-separated list of symbols')
    parser.add_argument('--output-dir', type=str, default='models/metamodels_neural',
                        help='Output directory for trained models')
    parser.add_argument('--epochs', type=int, default=100,
                        help='Max training epochs')
    parser.add_argument('--compare', action='store_true',
                        help='Compare with XGBoost metamodel')
    args = parser.parse_args()

    symbols = [s.strip() for s in args.symbols.split(',')]
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    logger.info("=" * 60)
    logger.info("NEURAL GATING METAMODEL TRAINING")
    logger.info("=" * 60)
    logger.info(f"Symbols: {symbols}")
    logger.info(f"Output: {output_dir}")

    # Check PyTorch availability
    try:
        import torch
        logger.info(f"PyTorch version: {torch.__version__}")
        logger.info(f"CUDA available: {torch.cuda.is_available()}")
    except ImportError:
        logger.error("PyTorch not installed! Run: pip install torch")
        return

    from src.models.neural_gating_metamodel import (
        NeuralGatingMetamodelWrapper,
        NeuralGatingConfig
    )

    # Load trainer metadata from XGBoost training
    trainer_metadata_path = Path('models/metamodels/trainer_metadata.joblib')

    if not trainer_metadata_path.exists():
        logger.error("Trainer metadata not found! Train XGBoost metamodels first.")
        return

    trainer_metadata = joblib.load(trainer_metadata_path)
    logger.info(f"Loaded trainer metadata from {trainer_metadata_path}")

    results = {}

    for symbol in symbols:
        logger.info(f"\n{'='*60}")
        logger.info(f"Training Neural Gating for {symbol}")
        logger.info(f"{'='*60}")

        # Load training data from metamodel trainer output
        data_path = Path(f'models/metamodels/{symbol}/training_data.joblib')

        if not data_path.exists():
            # Need to regenerate training data using metamodel trainer
            logger.warning(f"Training data not found for {symbol}, generating via MetamodelTrainer...")

            from src.training.metamodel_trainer import MetamodelTrainer, MetamodelTrainerConfig
            from src.data.database import MySQLConnection

            # Load data from database
            db = MySQLConnection()

            # Map symbol to ticker (database uses C: prefix for forex, X: for crypto)
            ticker_map = {
                'XAUUSD': 'C:XAUUSD',
                'EURUSD': 'C:EURUSD',
                'BTCUSD': 'X:BTCUSD',
                'GBPUSD': 'C:GBPUSD',
                'USDJPY': 'C:USDJPY'
            }
            ticker = ticker_map.get(symbol, f'C:{symbol}')

            # Load 5m and 15m data for OOS period
            query = f"""
                SELECT date_agg as timestamp, open, high, low, close, volume
                FROM tickers_agg_data
                WHERE ticker = '{ticker}'
                AND date_agg >= '2024-01-01' AND date_agg <= '2024-08-31'
                ORDER BY date_agg ASC
            """

            import pandas as pd
            df_raw = pd.read_sql(query, db.engine)

            if len(df_raw) < 5000:
                logger.error(f"Insufficient data for {symbol}: {len(df_raw)} rows")
                continue

            df_raw['timestamp'] = pd.to_datetime(df_raw['timestamp'])
            df_raw.set_index('timestamp', inplace=True)

            # Resample to 5m and 15m
            agg_dict = {
                'open': 'first',
                'high': 'max',
                'low': 'min',
                'close': 'last',
                'volume': 'sum'
            }

            df_5m = df_raw.resample('5min').agg(agg_dict).dropna()
            df_15m = df_raw.resample('15min').agg(agg_dict).dropna()

            logger.info(f"Loaded data: 5m={len(df_5m)}, 15m={len(df_15m)}")

            # Use MetamodelTrainer to generate features
            trainer_config = MetamodelTrainerConfig(
                symbols=[symbol],
                attention_model_path='models/attention',
                base_model_path='models/symbol_timeframe_models',
                output_path=str(output_dir)
            )

            trainer = MetamodelTrainer(trainer_config)

            if not trainer.load_models():
                logger.error("Failed to load models")
                continue

            result = trainer.train_single(df_5m.reset_index(), df_15m.reset_index(), symbol)

            if result.get('status') == 'failed':
                logger.error(f"Training data generation failed: {result.get('reason')}")
                continue

            # Get meta_features and targets from the trained model data
            # We need to extract these from the trainer
            logger.warning("Using XGBoost-generated data - Neural Gating uses same features")

            # Load the just-trained XGBoost data
            xgb_meta_path = Path(str(output_dir)) / symbol / 'training_data.joblib'
            if not xgb_meta_path.exists():
                logger.error(f"Training data not created at {xgb_meta_path}")
                continue

            data = joblib.load(xgb_meta_path)
            meta_features = data['meta_features']
            target_high = data['target_high']
            target_low = data['target_low']

        else:
            data = joblib.load(data_path)
            meta_features = data['meta_features']
            target_high = data['target_high']
            target_low = data['target_low']

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

        # Configure neural gating
        config = NeuralGatingConfig(
            epochs=args.epochs,
            early_stopping_patience=15,
            learning_rate=0.001,
            batch_size=256,
            gating_hidden_dims=[32, 16],
            residual_hidden_dims=[64, 32],
            confidence_hidden_dims=[32, 16],
            dropout=0.2
        )

        # Train
        model = NeuralGatingMetamodelWrapper(symbol, config)
        model.fit(meta_features, target_high, target_low)

        # Save
        model_path = output_dir / symbol
        model.save(str(model_path))

        results[symbol] = model.get_training_summary()

        # Compare with XGBoost if requested
        if args.compare:
            xgb_path = Path(f'models/metamodels/{symbol}')
            if xgb_path.exists():
                from src.models.asset_metamodel import AssetMetamodel
                xgb_model = AssetMetamodel.load(str(xgb_path))
                xgb_summary = xgb_model.get_training_summary()

                logger.info(f"\n{'-'*40}")
                logger.info(f"COMPARISON: Neural Gating vs XGBoost Stacking")
                logger.info(f"{'-'*40}")

                neural_metrics = results[symbol]['metrics']
                xgb_metrics = xgb_summary['metrics']

                logger.info(f"{'Metric':<25} {'Neural':<15} {'XGBoost':<15} {'Winner':<10}")
                logger.info("-" * 65)

                # MAE comparison (lower is better)
                neural_mae = (neural_metrics['mae_high'] + neural_metrics['mae_low']) / 2
                xgb_mae = (xgb_metrics['mae_high'] + xgb_metrics['mae_low']) / 2
                winner = "Neural" if neural_mae < xgb_mae else "XGBoost"
                logger.info(f"{'MAE (avg)':<25} {neural_mae:<15.4f} {xgb_mae:<15.4f} {winner:<10}")

                # R2 comparison (higher is better)
                neural_r2 = (neural_metrics['r2_high'] + neural_metrics['r2_low']) / 2
                xgb_r2 = (xgb_metrics['r2_high'] + xgb_metrics['r2_low']) / 2
                winner = "Neural" if neural_r2 > xgb_r2 else "XGBoost"
                logger.info(f"{'R2 (avg)':<25} {neural_r2:<15.4f} {xgb_r2:<15.4f} {winner:<10}")

                # Confidence accuracy
                neural_conf = neural_metrics['conf_accuracy']
                xgb_conf = xgb_metrics['confidence_accuracy']
                winner = "Neural" if neural_conf > xgb_conf else "XGBoost"
                logger.info(f"{'Confidence Accuracy':<25} {neural_conf:<15.4f} {xgb_conf:<15.4f} {winner:<10}")

                # Improvement over simple average
                neural_imp = (neural_metrics['improvement_high'] + neural_metrics['improvement_low']) / 2
                xgb_imp = xgb_metrics['improvement_over_avg']
                winner = "Neural" if neural_imp > xgb_imp else "XGBoost"
                logger.info(f"{'Improvement over avg':<25} {neural_imp:<15.1f}% {xgb_imp:<15.1f}% {winner:<10}")

    # Save training report
    report_path = output_dir / f'training_report_{datetime.now().strftime("%Y%m%d_%H%M%S")}.md'

    with open(report_path, 'w') as f:
        f.write("# Neural Gating Metamodel Training Report\n\n")
        f.write(f"**Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")

        f.write("## Summary\n\n")
        f.write("| Symbol | MAE High | MAE Low | R2 High | R2 Low | Alpha High | Alpha Low |\n")
        f.write("|--------|----------|---------|---------|--------|------------|----------|\n")

        for symbol, summary in results.items():
            m = summary['metrics']
            f.write(f"| {symbol} | {m['mae_high']:.4f} | {m['mae_low']:.4f} | ")
            f.write(f"{m['r2_high']:.4f} | {m['r2_low']:.4f} | ")
            f.write(f"{m['alpha_high_mean']:.3f} | {m['alpha_low_mean']:.3f} |\n")

        f.write("\n## Gating Analysis\n\n")
        f.write("The alpha values show the learned weight for 5m predictions:\n")
        f.write("- Alpha = 0.5: Equal weight to 5m and 15m\n")
        f.write("- Alpha > 0.5: More weight to 5m (faster timeframe)\n")
        f.write("- Alpha < 0.5: More weight to 15m (slower timeframe)\n\n")

        for symbol, summary in results.items():
            m = summary['metrics']
            f.write(f"### {symbol}\n")
            f.write(f"- Alpha HIGH mean: {m['alpha_high_mean']:.3f}\n")
            f.write(f"- Alpha LOW mean: {m['alpha_low_mean']:.3f}\n")
            f.write(f"- Interpretation: Model weights {'5m more' if m['alpha_high_mean'] > 0.5 else '15m more'} for HIGH, ")
            f.write(f"{'5m more' if m['alpha_low_mean'] > 0.5 else '15m more'} for LOW\n\n")

    logger.info(f"\nReport saved to: {report_path}")
    logger.info("\n" + "=" * 60)
    logger.info("NEURAL GATING TRAINING COMPLETE")
    logger.info("=" * 60)


if __name__ == '__main__':
    main()
