"""AlphaForge - Complete Quantitative Trading System

Usage:
    python main.py --mode train --tickers SPY QQQ AAPL MSFT
    python main.py --mode backtest --start 2020-01-01 --end 2024-01-01
    python main.py --mode live --config config.yaml
"""
import argparse
import numpy as np
import pandas as pd
import torch
import warnings
warnings.filterwarnings('ignore')

from market_data import MarketDataPipeline
from alpha_model import AlphaEnsemble
from sentiment_model import SentimentAlphaModel
from volatility_model import VolatilityEngine
from portfolio_optimizer import PortfolioOptimizer
from options_pricer import MLOptionsPricer
from backtest_engine import BacktestEngine, compute_information_coefficient, RegimeDetector


def parse_args():
    parser = argparse.ArgumentParser(description='AlphaForge Quant System')
    parser.add_argument('--mode', type=str, default='train',
                        choices=['train', 'backtest', 'live', 'options'])
    parser.add_argument('--tickers', type=str, nargs='+',
                        default=['SPY','QQQ','AAPL','MSFT','GOOGL','AMZN','META','NVDA','TSLA','JPM'])
    parser.add_argument('--start', type=str, default='2020-01-01')
    parser.add_argument('--end', type=str, default='2024-01-01')
    parser.add_argument('--lookback', type=int, default=60)
    parser.add_argument('--horizon', type=int, default=5)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--device', type=str, default='cpu')
    parser.add_argument('--initial_capital', type=float, default=1_000_000)
    parser.add_argument('--output', type=str, default='results/')
    return parser.parse_args()


def train_alpha_model(args):
    """Train the multi-asset alpha model"""
    print("=" * 60)
    print("ALPHA FORGE - Multi-Asset Alpha Model Training")
    print("=" * 60)
    
    # Fetch data
    pipeline = MarketDataPipeline(args.tickers, args.start, args.end)
    data = pipeline.fetch_data()
    
    # Create features
    features_df = pipeline.create_feature_matrix()
    X, y, tickers, dates = pipeline.create_sequences(
        features_df, lookback=args.lookback, forecast_horizon=args.horizon
    )
    
    print(f"\nDataset: {len(X)} samples, {X.shape[2]} features, seq_len={args.lookback}")
    
    # Train/val/test split (time-based)
    n = len(X)
    train_end = int(n * 0.7)
    val_end = int(n * 0.85)
    
    X_train, y_train = X[:train_end], y[:train_end]
    X_val, y_val = X[train_end:val_end], y[train_end:val_end]
    X_test, y_test = X[val_end:], y[val_end:]
    
    print(f"Train: {len(X_train)}, Val: {len(X_val)}, Test: {len(X_test)}")
    
    # Train ensemble
    ensemble = AlphaEnsemble(
        input_size=X.shape[2],
        seq_len=args.lookback,
        device=args.device
    )
    
    metrics = ensemble.fit(
        X_train, y_train,
        X_val, y_val,
        epochs=args.epochs,
        batch_size=64,
        lr=1e-4
    )
    
    # Test predictions
    test_pred = ensemble.predict(X_test)
    test_ic = compute_information_coefficient(
        pd.Series(test_pred),
        pd.Series(y_test),
        by_date=False
    )
    
    print(f"\nTest IC: {test_ic['mean_ic']:.4f}")
    print(f"LSTM final val IC: {metrics['lstm']['val_ic'][-1]:.4f}")
    print(f"Transformer final val IC: {metrics['transformer']['val_ic'][-1]:.4f}")
    
    # Save model
    torch.save(ensemble.lstm.state_dict(), f"{args.output}/lstm_model.pt")
    torch.save(ensemble.transformer.state_dict(), f"{args.output}/transformer_model.pt")
    
    return ensemble, metrics, test_ic


def run_backtest(args):
    """Run full pipeline backtest"""
    print("=" * 60)
    print("ALPHA FORGE - Full Pipeline Backtest")
    print("=" * 60)
    
    # Fetch data
    pipeline = MarketDataPipeline(args.tickers, args.start, args.end)
    data = pipeline.fetch_data()
    features_df = pipeline.create_feature_matrix()
    
    X, y, tickers_arr, dates = pipeline.create_sequences(
        features_df, lookback=args.lookback, forecast_horizon=args.horizon
    )
    
    # Split
    n = len(X)
    train_end = int(n * 0.7)
    val_end = int(n * 0.85)
    
    X_train, y_train = X[:train_end], y[:train_end]
    X_test, y_test = X[val_end:], y[val_end:]
    dates_test = dates[val_end:]
    tickers_test = tickers_arr[val_end:]
    
    # Train alpha model
    print("\n[1/4] Training Alpha Model...")
    ensemble = AlphaEnsemble(input_size=X.shape[2], seq_len=args.lookback, device=args.device)
    ensemble.fit(X_train, y_train, epochs=30, batch_size=64, lr=1e-4)
    
    # Generate predictions
    alpha_pred = ensemble.predict(X_test)
    
    # Build prediction DataFrame
    pred_df = pd.DataFrame({
        'date': dates_test,
        'ticker': tickers_test,
        'predicted_return': alpha_pred,
        'actual_return': y_test
    })
    
    # Compute IC
    ic_metrics = compute_information_coefficient(
        pred_df['predicted_return'],
        pred_df['actual_return'],
        by_date=True
    )
    print(f"Mean IC: {ic_metrics['mean_ic']:.4f} +/- {ic_metrics['ic_std']:.4f}")
    print(f"IC IR: {ic_metrics['ic_ir']:.4f}")
    
    # Train volatility model
    print("\n[2/4] Training Volatility Model...")
    vol_engine = VolatilityEngine()
    
    # Build returns matrix for covariance
    returns_dict = {}
    for ticker in args.tickers:
        if ticker in data:
            close = data[ticker]['Close'].values.flatten()
            returns_dict[ticker] = pd.Series(
                np.log(close[1:] / close[:-1]),
                index=data[ticker].index[1:]
            )
    returns_df = pd.DataFrame(returns_dict).fillna(0)
    
    # Fit GARCH for each ticker
    for ticker in args.tickers:
        if ticker in returns_df.columns:
            vol_engine.fit_garch(returns_df[ticker], ticker)
    
    # Portfolio optimization and backtest
    print("\n[3/4] Running Portfolio Optimization...")
    
    # Get unique test dates
    test_dates = pd.to_datetime(pred_df['date'].unique())
    test_dates = sorted(test_dates)
    
    # Rebalance every 5 days
    rebalance_dates = test_dates[::5]
    
    optimizer = PortfolioOptimizer(
        max_weight=0.25,
        risk_aversion=2.0,
        transaction_cost=0.0003,
        turnover_penalty=0.001
    )
    
    weights_history = []
    
    for rebalance_date in rebalance_dates:
        # Get predictions for this date
        day_preds = pred_df[pred_df['date'] == rebalance_date]
        
        if len(day_preds) < 3:
            continue
        
        # Build mu vector
        mu = day_preds.set_index('ticker')['predicted_return'].reindex(args.tickers).fillna(0).values
        
        # Build covariance matrix
        try:
            Sigma = vol_engine.build_covariance_matrix(returns_df, rebalance_date)
            Sigma = Sigma.reindex(index=args.tickers, columns=args.tickers).fillna(0)
            Sigma = Sigma.values
        except:
            Sigma = np.eye(len(args.tickers)) * 0.04
        
        # Optimize
        result = optimizer.optimize_max_sharpe(mu, Sigma)
        
        weights_row = pd.Series(result['weights'], index=args.tickers)
        weights_row.name = rebalance_date
        weights_history.append(weights_row)
    
    weights_df = pd.DataFrame(weights_history)
    
    # Build returns for backtest
    backtest_returns = returns_df.reindex(weights_df.index).fillna(0)
    
    # Run backtest
    print("\n[4/4] Running Backtest...")
    engine = BacktestEngine(
        initial_capital=args.initial_capital,
        transaction_cost=0.0003,
        slippage=0.0001
    )
    
    metrics = engine.run_backtest(
        backtest_returns,
        weights_df,
        rebalance_dates=weights_df.index
    )
    
    # Regime detection
    if 'SPY' in returns_df.columns:
        regime_detector = RegimeDetector()
        spy_returns = returns_df['SPY'].reindex(weights_df.index).fillna(0)
        regimes = regime_detector.detect_regimes(spy_returns)
        regime_stats = regime_detector.get_regime_stats(spy_returns)
        print("\nRegime Statistics:")
        print(regime_stats.to_string())
    
    # Print results
    print("\n" + "=" * 60)
    print("BACKTEST RESULTS")
    print("=" * 60)
    print(f"Total Return:      {metrics['total_return']*100:.2f}%")
    print(f"Annualized Return: {metrics['annualized_return']*100:.2f}%")
    print(f"Volatility:        {metrics['volatility']*100:.2f}%")
    print(f"Sharpe Ratio:      {metrics['sharpe_ratio']:.3f}")
    print(f"Sortino Ratio:     {metrics['sortino_ratio']:.3f}")
    print(f"Max Drawdown:      {metrics['max_drawdown']*100:.2f}%")
    print(f"Calmar Ratio:      {metrics['calmar_ratio']:.3f}")
    print(f"Win Rate:          {metrics['win_rate']*100:.1f}%")
    print(f"Alpha:             {metrics['alpha']*100:.2f}%")
    print(f"Beta:              {metrics['beta']:.3f}")
    print(f"Information Ratio: {metrics['information_ratio']:.3f}")
    print(f"Avg Turnover:      {metrics['avg_turnover']*100:.2f}%")
    print(f"Total Costs:       ${metrics['total_transaction_costs']:,.2f}")
    print(f"Final Capital:     ${metrics['final_capital']:,.2f}")
    print(f"Trades:            {metrics['n_trades']}")
    
    # Save results
    import os
    os.makedirs(args.output, exist_ok=True)
    
    results = {
        'metrics': metrics,
        'ic_metrics': ic_metrics,
        'equity_curve': engine.get_equity_curve().to_dict(),
        'weights': weights_df.to_dict()
    }
    
    import json
    with open(f"{args.output}/backtest_results.json", 'w') as f:
        json.dump({k: v for k, v in results.items() if k != 'weights'}, f, indent=2, default=str)
    
    weights_df.to_csv(f"{args.output}/weights_history.csv")
    
    print(f"\nResults saved to {args.output}/")
    
    return metrics, engine


def train_options_model(args):
    """Train ML options pricing model"""
    print("=" * 60)
    print("ALPHA FORGE - Options Pricing Model")
    print("=" * 60)
    
    pricer = MLOptionsPricer(device=args.device)
    
    # Generate synthetic training data
    print("Generating synthetic option data...")
    train_df = pricer.generate_synthetic_options(n_samples=50000)
    val_df = pricer.generate_synthetic_options(n_samples=10000)
    
    X_train = pricer.prepare_features(train_df)
    y_train = train_df['price'].values
    X_val = pricer.prepare_features(val_df)
    y_val = val_df['price'].values
    
    print(f"Training samples: {len(X_train)}, Validation: {len(X_val)}")
    
    # Train
    metrics = pricer.fit(X_train, y_train, X_val, y_val, epochs=100, batch_size=256)
    
    # Test on a few examples
    test_df = pricer.generate_synthetic_options(n_samples=5)
    X_test = pricer.prepare_features(test_df)
    
    ml_prices = pricer.predict(X_test)
    bs_prices = []
    for i in range(len(test_df)):
        if test_df['option_type'].iloc[i] == 'call':
            p = pricer.bs.call_price(
                test_df['S'].iloc[i], test_df['K'].iloc[i],
                test_df['T'].iloc[i], test_df['r'].iloc[i],
                test_df['sigma_hist'].iloc[i]
            )
        else:
            p = pricer.bs.put_price(
                test_df['S'].iloc[i], test_df['K'].iloc[i],
                test_df['T'].iloc[i], test_df['r'].iloc[i],
                test_df['sigma_hist'].iloc[i]
            )
        bs_prices.append(p)
    
    print("\nSample Predictions:")
    print(f"{'True':>10} {'ML':>10} {'BS':>10} {'ML Err%':>10} {'BS Err%':>10}")
    for i in range(len(test_df)):
        true_p = test_df['price'].iloc[i]
        ml_err = abs(ml_prices[i] - true_p) / true_p * 100
        bs_err = abs(bs_prices[i] - true_p) / true_p * 100
        print(f"{true_p:>10.2f} {ml_prices[i]:>10.2f} {bs_prices[i]:>10.2f} {ml_err:>10.2f} {bs_err:>10.2f}")
    
    # Save
    import os
    os.makedirs(args.output, exist_ok=True)
    torch.save(pricer.model.state_dict(), f"{args.output}/options_model.pt")
    
    return pricer, metrics


def main():
    args = parse_args()
    
    if args.mode == 'train':
        train_alpha_model(args)
    elif args.mode == 'backtest':
        run_backtest(args)
    elif args.mode == 'options':
        train_options_model(args)
    else:
        print("Live mode not implemented in this version")


if __name__ == '__main__':
    main()