#!/usr/bin/env python3
"""
Experimenting with quality filters for MoERegression.

Usage:
    python experiment_quality_filters.py

This script:
1. Simulates trades with different quality filters
2. Evaluates which filters improve win rate and PnL
3. Compares results to actual trades
4. Generates actionable insights
"""

import sys
sys.path.insert(0, '/home/ai/projects/AI_Strategy')

import pandas as pd
import numpy as np
from db.connection import get_connection
# from config import HORIZONS, MAX_BARS_RATIO  # Not needed for this experiment

# Experiment configs
EXPERIMENT_CONFIGS = {
    'no_filter': {
        'name': 'No filters (baseline)',
        'params': {}
    },
    'flat_only': {
        'name': 'Flat filter only',
        'params': {
            'chopper_threshold': 0.618,
            'vol_ratio_threshold': 0.8
        }
    },
    'momentum_only': {
        'name': 'Momentum filter only',
        'params': {
            'min_momentum_5': 0.02,
            'min_momentum_10': 0.04,
            'max_volatility': 2.0
        }
    },
    'flat_momentum': {
        'name': 'Flat + Momentum',
        'params': {
            'chopper_threshold': 0.618,
            'vol_ratio_threshold': 0.8,
            'min_momentum_5': 0.02,
            'min_momentum_10': 0.04
        }
    },
    'liquidity_only': {
        'name': 'Liquidity filter only',
        'params': {
            'min_volume_ratio': 1.2,
            'min_liquidity': 1000000
        }
    },
    'time_limit': {
        'name': 'Time limit filter',
        'params': {
            'max_bars_to_tp': 48
        }
    },
    'composite': {
        'name': 'Composite filter',
        'params': {
            'chopper_threshold': 0.618,
            'vol_ratio_threshold': 0.8,
            'min_momentum_5': 0.02,
            'min_momentum_10': 0.04,
            'min_volume_ratio': 1.2,
            'max_bars_to_tp': 48
        }
    }
}


def load_trade_data():
    """Load actual trades from database."""
    query = """
    SELECT
        ticker,
        direction,
        entry_price,
        exit_price,
        entry_time,
        exit_time,
        close_reason,
        sl_price,
        tp_price,
        volume,
        pnl,
        pnl_pct,
        atr_entry,
        created_at
    FROM trades_closed_regression
    ORDER BY entry_time
    """

    with get_connection() as conn:
        df = pd.read_sql(query, conn)

    print(f"✅ Loaded {len(df)} actual trades")
    return df


def load_candidate_signals():
    """Load trade data for analysis."""
    query = """
    SELECT
        ticker,
        direction,
        entry_price,
        atr_entry,
        pnl
    FROM trades_closed_regression
    ORDER BY entry_time
    """

    with get_connection() as conn:
        signals = pd.read_sql(query, conn)

    # We'll simulate OHLC data using entry_price, sl_price, tp_price
    # This is a simplification for the experiment
    signals['High'] = signals['entry_price'] * 1.02
    signals['Low'] = signals['entry_price'] * 0.98
    signals['Close'] = signals['entry_price']
    signals['Volume'] = 1000  # Simplified

    print(f"✅ Loaded {len(signals)} trades")
    return signals


def engineer_features(df):
    """Engineer features for simulation."""
    df = df.copy()

    # Use entry_price as current price
    df['current_price'] = df['entry_price']

    # Simple features
    df['SMA_5'] = df['current_price'].rolling(5).mean()
    df['SMA_10'] = df['current_price'].rolling(10).mean()
    df['SMA_20'] = df['current_price'].rolling(20).mean()

    # Momentum
    df['momentum_5'] = (df['current_price'] - df['Close'].shift(4)) / df['Close'].shift(4)
    df['momentum_10'] = (df['current_price'] - df['Close'].shift(9)) / df['Close'].shift(9)

    # Volatility
    df['atr_pct'] = df['atr_entry'] / df['current_price']

    # Volume ratio (simplified)
    df['volume_ratio'] = 1.0  # Cannot calculate without actual volume history

    # Price range
    df['high_low_range'] = (df['High'] - df['Low']) / df['current_price']

    # Choppiness (simplified)
    df['chopper'] = (df['current_price'] - df['Low'].rolling(20).min()) / (df['High'].rolling(20).max() - df['Low'].rolling(20).min())

    # RSI (simplified, 14 period)
    delta = df['current_price'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(14).mean()
    rs = gain / loss
    df['RSI_14'] = 100 - (100 / (1 + rs))

    # Remove NaN
    df.dropna(inplace=True)

    return df


def check_flat_market(df, i, params):
    """Check if market is in flat (consolidation) condition."""
    chopper_threshold = params.get('chopper_threshold', 0.618)
    vol_ratio_threshold = params.get('vol_ratio_threshold', 0.8)

    is_flat = (
        (df['chopper'].iloc[i] < chopper_threshold) and
        (df['volume_ratio'].iloc[i] < vol_ratio_threshold)
    )

    return is_flat


def check_momentum(df, i, params):
    """Check if momentum is favorable."""
    min_momentum_5 = params.get('min_momentum_5', 0.0)
    min_momentum_10 = params.get('min_momentum_10', 0.0)

    has_momentum_5 = df['momentum_5'].iloc[i] > min_momentum_5
    has_momentum_10 = df['momentum_10'].iloc[i] > min_momentum_10

    return has_momentum_5 and has_momentum_10


def check_liquidity(df, i, params):
    """Check if liquidity is sufficient."""
    min_volume_ratio = params.get('min_volume_ratio', 1.0)
    min_liquidity = params.get('min_liquidity', 1000000)

    has_volume = df['Volume'].iloc[i] >= min_liquidity
    has_volume_ratio = df['volume_ratio'].iloc[i] >= min_volume_ratio

    return has_volume and has_volume_ratio


def check_time_limit(df, i, params):
    """Check if target is reachable within time limit."""
    max_bars = params.get('max_bars_to_tp', 48)

    # Simple check: if pred_peak is too far in the future
    is_reachable = df['bars_to_peak'].iloc[i] <= max_bars

    return is_reachable


def simulate_trades(df, params, direction='LONG'):
    """Simulate trades with given filter parameters."""
    trades = []

    # Filter candidates
    mask = pd.Series(True, index=df.index)

    if 'chopper_threshold' in params:
        mask &= (df['chopper'] < params['chopper_threshold'])
    if 'vol_ratio_threshold' in params:
        mask &= (df['volume_ratio'] < params['vol_ratio_threshold'])
    if 'min_momentum_5' in params:
        mask &= (df['momentum_5'] > params['min_momentum_5'])
    if 'min_momentum_10' in params:
        mask &= (df['momentum_10'] > params['min_momentum_10'])
    if 'min_volume_ratio' in params:
        mask &= (df['volume_ratio'] >= params['min_volume_ratio'])
    if 'max_bars_to_tp' in params:
        # Simplified time limit: check if momentum is sufficient
        mask &= (df['momentum_5'] >= 0.02)

    df_filtered = df[mask]

    # Simulate trades on filtered data
    for i in range(len(df_filtered)):
        if df_filtered['direction'].iloc[i] != direction:
            continue

        # Use existing entry price
        entry_price = df_filtered['entry_price'].iloc[i]

        # Simulate close (use actual exit price from trade)
        pnl = df_filtered['pnl'].iloc[i]

        trades.append({
            'entry_price': entry_price,
            'exit_price': entry_price + pnl,  # Simplified exit
            'pnl': pnl,
            'atr': df_filtered['atr_entry'].iloc[i],
            'chopper': df_filtered['chopper'].iloc[i],
            'momentum_5': df_filtered['momentum_5'].iloc[i],
            'volume_ratio': df_filtered['volume_ratio'].iloc[i]
        })

    return pd.DataFrame(trades)


def analyze_experiment(trades_df, experiment_name, direction='LONG'):
    """Analyze simulated trades."""
    if len(trades_df) == 0:
        return None

    total_trades = len(trades_df)
    tp_trades = (trades_df['pnl'] > 0).sum()
    sl_trades = (trades_df['pnl'] <= 0).sum()

    win_rate = tp_trades / total_trades
    avg_pnl = trades_df['pnl'].mean()
    total_pnl = trades_df['pnl'].sum()
    avg_abs_pnl = trades_df['pnl'].abs().mean()

    # Calculate quality metrics
    avg_tp_pnl = trades_df[trades_df['pnl'] > 0]['pnl'].mean()
    avg_sl_pnl = trades_df[trades_df['pnl'] <= 0]['pnl'].mean()
    win_loss_ratio = avg_tp_pnl / avg_sl_pnl if avg_sl_pnl != 0 else float('inf')

    metrics = {
        'experiment_name': experiment_name,
        'direction': direction,
        'total_trades': total_trades,
        'win_rate': win_rate,
        'avg_pnl': avg_pnl,
        'total_pnl': total_pnl,
        'avg_abs_pnl': avg_abs_pnl,
        'avg_tp_pnl': avg_tp_pnl,
        'avg_sl_pnl': avg_sl_pnl,
        'win_loss_ratio': win_loss_ratio
    }

    return metrics


def main():
    print("="*80)
    print("🔬 EXPERIMENT: Quality Filters for MoERegression")
    print("="*80)
    print(f"\n📅 Date: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"📦 Backup: trades_closed_regression_backup (50 records)")
    print()

    # Load data
    print("1️⃣ Loading data...")
    df_trades = load_trade_data()
    df_candidate_signals = load_candidate_signals()

    # Engineer features for candidate signals
    print("\n2️⃣ Engineering features...")
    df_candidate_signals = engineer_features(df_candidate_signals)

    # Run experiments
    print("\n3️⃣ Running experiments...")
    results = []

    for exp_name, exp_config in EXPERIMENT_CONFIGS.items():
        print(f"\n   Running: {exp_name}")
        print(f"   Config: {exp_config['params']}")

        # Simulate LONG trades
        trades_long = simulate_trades(df_candidate_signals, exp_config['params'], direction='LONG')
        metrics_long = analyze_experiment(trades_long, exp_name, direction='LONG')

        # Simulate SHORT trades
        trades_short = simulate_trades(df_candidate_signals, exp_config['params'], direction='SHORT')
        metrics_short = analyze_experiment(trades_short, exp_name, direction='SHORT')

        if metrics_long:
            results.append(metrics_long)
            print(f"      LONG:  {len(trades_long)} trades, WR: {metrics_long['win_rate']:.1%}, Avg PnL: {metrics_long['avg_pnl']:.2f}")
        else:
            print(f"      LONG:  0 trades (no candidates)")

        if metrics_short:
            results.append(metrics_short)
            print(f"      SHORT: {len(trades_short)} trades, WR: {metrics_short['win_rate']:.1%}, Avg PnL: {metrics_short['avg_pnl']:.2f}")
        else:
            print(f"      SHORT: 0 trades (no candidates)")

    # Load actual performance for comparison
    print("\n4️⃣ Comparing with actual trades...")
    long_trades = df_trades[df_trades['direction'] == 'LONG']

    actual_metrics = {
        'experiment_name': 'Actual trades (baseline)',
        'direction': 'LONG',
        'total_trades': len(long_trades),
        'win_rate': (long_trades['pnl'] > 0).sum() / len(long_trades) if len(long_trades) > 0 else 0,
        'avg_pnl': long_trades['pnl'].mean() if len(long_trades) > 0 else 0,
        'total_pnl': long_trades['pnl'].sum() if len(long_trades) > 0 else 0,
        'avg_abs_pnl': long_trades['pnl'].abs().mean() if len(long_trades) > 0 else 0,
        'avg_tp_pnl': long_trades[long_trades['pnl'] > 0]['pnl'].mean() if len(long_trades) > 0 else 0,
        'avg_sl_pnl': long_trades[long_trades['pnl'] <= 0]['pnl'].abs().mean() if len(long_trades) > 0 else 0,
        'win_loss_ratio': long_trades[long_trades['pnl'] > 0]['pnl'].mean() / abs(long_trades[long_trades['pnl'] <= 0]['pnl'].mean()) if len(long_trades) > 0 and long_trades[long_trades['pnl'] <= 0]['pnl'].mean() != 0 else float('inf')
    }
    results.append(actual_metrics)

    # Display results
    print("\n" + "="*80)
    print("📊 EXPERIMENT RESULTS")
    print("="*80)

    # Sort by win rate (descending)
    results_df = pd.DataFrame(results)
    results_df = results_df.sort_values(['experiment_name', 'total_trades'], ascending=[True, False])

    # Format output
    print(f"\n{'Experiment':<30} {'Trades':<8} {'WR':<10} {'Avg PnL':<12} {'Total PnL':<12} {'Avg |PnL|':<12}")
    print("-"*80)

    for _, row in results_df.iterrows():
        print(f"{row['experiment_name']:<30} {row['total_trades']:<8} "
              f"{row['win_rate']:<10.2%} {row['avg_pnl']:>10,.2f} "
              f"{row['total_pnl']:>10,.2f} {row['avg_abs_pnl']:>10,.2f}")

    # Identify best experiments
    print("\n" + "="*80)
    print("🏆 TOP 3 EXPERIMENTS")
    print("="*80)

    best_by_wr = results_df.nlargest(3, 'win_rate')
    best_by_pnl = results_df.nlargest(3, 'avg_pnl')

    print(f"\nBy Win Rate:")
    for _, row in best_by_wr.iterrows():
        print(f"  {row['experiment_name']:<30} WR: {row['win_rate']:.1%}, Trades: {row['total_trades']}")

    print(f"\nBy Average PnL:")
    for _, row in best_by_pnl.iterrows():
        print(f"  {row['experiment_name']:<30} Avg PnL: {row['avg_pnl']:.2f}, Trades: {row['total_trades']}")

    # Recommendations
    print("\n" + "="*80)
    print("💡 RECOMMENDATIONS")
    print("="*80)

    best_long_wr = results_df[(results_df['direction'] == 'LONG') & (results_df['experiment_name'] != 'Actual trades (baseline)')].nlargest(1, 'win_rate')
    best_long_pnl = results_df[(results_df['direction'] == 'LONG')].nlargest(1, 'avg_pnl')

    if not best_long_wr.empty:
        print(f"\n✅ BEST LONG FILTER: {best_long_wr.iloc[0]['experiment_name']}")
        print(f"   Win Rate: {best_long_wr.iloc[0]['win_rate']:.1%}")
        print(f"   Trades: {best_long_wr.iloc[0]['total_trades']}")

    if not best_long_pnl.empty:
        print(f"\n✅ BEST LONG PnL: {best_long_pnl.iloc[0]['experiment_name']}")
        print(f"   Avg PnL: {best_long_pnl.iloc[0]['avg_pnl']:.2f}")
        print(f"   Trades: {best_long_pnl.iloc[0]['total_trades']}")

    print("\n" + "="*80)
    print("✅ Experiment complete")
    print("="*80)

    # Save results
    output_file = '/home/ai/projects/AI_Strategy/experiments/results.csv'
    results_df.to_csv(output_file, index=False)
    print(f"\n💾 Results saved to: {output_file}")


if __name__ == '__main__':
    main()