#!/usr/bin/env python3
"""
Equity & Leverage Optimizer for ML v11.2 Long Strategy
Finds optimal combination of:
- Equity % per trade (25%, 50%, 75%, 100%)
- Leverage (5x, 10x, 15x, 20x)
- Scale behavior on consecutive signals

User Question: "should we use 100% equity each trade and then full blast with 20x 
when we get 2 buy signals in a row?"
"""

import pandas as pd
from datetime import datetime
from itertools import product

# Trade data from ML v11.2 backtest
TRADES = [
    # (trade_num, entry_time, exit_time, entry_price, exit_price, pnl_pct)
    (1, "2025-12-01 14:00", "2025-12-01 16:30", 127.26, 123.69, -2.81),
    (2, "2025-12-02 23:00", "2025-12-03 05:40", 139.45, 141.85, 1.72),
    (3, "2025-12-03 03:10", "2025-12-03 05:40", 139.77, 141.85, 1.49),
    (4, "2025-12-04 07:35", "2025-12-04 20:10", 144.09, 140.05, -2.80),
    (5, "2025-12-07 17:00", "2025-12-07 18:15", 130.99, 133.09, 1.60),
    (6, "2025-12-08 03:20", "2025-12-08 07:45", 133.64, 135.78, 1.60),
    (7, "2025-12-09 05:05", "2025-12-09 16:40", 133.90, 135.67, 1.32),
    (8, "2025-12-09 16:10", "2025-12-09 16:40", 133.15, 135.67, 1.89),
    (9, "2025-12-10 07:15", "2025-12-10 21:10", 139.24, 141.41, 1.56),
    (10, "2025-12-10 08:55", "2025-12-10 21:10", 139.12, 141.41, 1.65),
    (11, "2025-12-11 13:35", "2025-12-11 19:15", 131.41, 133.79, 1.81),
    (12, "2025-12-11 16:25", "2025-12-11 19:15", 131.95, 133.79, 1.39),
    (13, "2025-12-13 08:45", "2025-12-15 00:25", 132.82, 129.07, -2.82),
    (14, "2025-12-13 23:40", "2025-12-15 00:25", 132.76, 129.07, -2.78),
    (15, "2025-12-15 02:45", "2025-12-15 03:15", 130.96, 133.06, 1.60),
    (16, "2025-12-15 23:20", "2025-12-16 01:20", 125.93, 127.95, 1.60),
    (17, "2025-12-16 11:25", "2025-12-16 11:45", 126.71, 128.74, 1.60),
    (18, "2025-12-16 15:30", "2025-12-17 15:50", 128.70, 130.24, 1.20),
    (19, "2025-12-17 12:40", "2025-12-17 15:50", 127.67, 130.24, 2.01),
    (20, "2025-12-18 10:50", "2025-12-18 14:30", 123.58, 125.56, 1.60),
    (21, "2025-12-19 17:10", "2025-12-22 01:15", 126.22, 128.11, 1.50),
    (22, "2025-12-19 20:15", "2025-12-22 01:15", 125.96, 128.11, 1.71),
    (23, "2025-12-22 04:30", "2025-12-22 15:30", 126.06, 128.21, 1.71),
    (24, "2025-12-22 05:25", "2025-12-22 15:30", 126.32, 128.21, 1.50),
    (25, "2025-12-23 01:20", "2025-12-24 03:40", 126.10, 121.98, -3.27),
    (26, "2025-12-23 12:50", "2025-12-24 03:40", 124.89, 121.98, -2.33),
    (27, "2025-12-24 11:40", "2025-12-25 16:45", 122.07, 124.04, 1.61),
    (28, "2025-12-25 15:10", "2025-12-25 16:45", 122.09, 124.04, 1.60),
    (29, "2025-12-26 10:30", "2025-12-26 14:50", 123.03, 125.00, 1.60),
    (30, "2025-12-27 18:50", "2025-12-29 01:20", 123.34, 125.69, 1.91),
    (31, "2025-12-29 00:00", "2025-12-29 01:20", 124.08, 125.69, 1.30),
    (32, "2025-12-30 07:15", "2025-12-30 17:15", 123.82, 125.81, 1.61),
    (33, "2025-12-31 00:25", "2025-12-31 12:40", 124.72, 126.72, 1.60),
    (34, "2025-12-31 23:30", "2026-01-02 00:25", 124.90, 126.91, 1.61),
    (35, "2026-01-01 00:15", "2026-01-02 00:25", 124.92, 126.91, 1.59),
    (36, "2026-01-02 15:50", "2026-01-02 17:30", 128.81, 130.88, 1.61),
    (37, "2026-01-03 10:40", "2026-01-03 23:35", 131.33, 133.34, 1.53),
    (38, "2026-01-03 14:25", "2026-01-03 23:35", 131.14, 133.34, 1.68),
    (39, "2026-01-04 22:20", "2026-01-05 02:15", 134.23, 136.57, 1.74),
    (40, "2026-01-05 01:15", "2026-01-05 02:15", 134.60, 136.57, 1.46),
    (41, "2026-01-06 10:50", "2026-01-06 15:10", 138.24, 140.46, 1.61),
    (42, "2026-01-06 21:45", "2026-01-06 22:15", 138.89, 141.12, 1.61),
    (43, "2026-01-07 04:00", "2026-01-07 16:15", 140.00, 136.08, -2.80),
    (44, "2026-01-08 16:30", "2026-01-08 17:40", 135.09, 137.26, 1.61),
]

INITIAL_CAPITAL = 1400.0
COMMISSION_PCT = 0.05  # 0.05% per trade (0.1% round trip)


def parse_time(t):
    return datetime.strptime(t, "%Y-%m-%d %H:%M")


def group_trades_by_exit():
    """Group trades that exit at the same time (concurrent positions)."""
    trades_df = []
    for t in TRADES:
        trades_df.append({
            'trade_num': t[0],
            'entry_time': parse_time(t[1]),
            'exit_time': parse_time(t[2]),
            'entry_price': t[3],
            'exit_price': t[4],
            'pnl_pct': t[5]
        })
    
    df = pd.DataFrame(trades_df)
    df = df.sort_values(['entry_time', 'trade_num'])
    
    # Group by exit time
    groups = []
    for exit_time, group in df.groupby('exit_time'):
        group = group.sort_values('entry_time')
        groups.append({
            'exit_time': exit_time,
            'trades': group.to_dict('records'),
            'count': len(group)
        })
    
    groups.sort(key=lambda x: x['exit_time'])
    return groups


def simulate_strategy(equity_pct, leverage, scale_on_consecutive=True, scale_leverage_boost=1.0):
    """
    Simulate the strategy with given parameters.
    
    Args:
        equity_pct: Percentage of equity to use per trade (0.25, 0.50, 0.75, 1.00)
        leverage: Leverage multiplier (5, 10, 15, 20)
        scale_on_consecutive: If True, scale into consecutive signals
        scale_leverage_boost: Multiplier for leverage on 2nd signal (e.g., 2.0 = double leverage)
    
    Returns:
        dict with simulation results
    """
    groups = group_trades_by_exit()
    
    equity = INITIAL_CAPITAL
    peak_equity = INITIAL_CAPITAL
    max_drawdown_pct = 0
    
    equity_curve = [INITIAL_CAPITAL]
    trade_results = []
    
    for group in groups:
        trades = group['trades']
        num_concurrent = len(trades)
        
        if num_concurrent == 1:
            # Single trade - use base equity % and leverage
            trade = trades[0]
            position_size = equity * equity_pct
            effective_leverage = leverage
            
            # Calculate P&L
            gross_pnl_pct = trade['pnl_pct']
            leveraged_pnl_pct = gross_pnl_pct * effective_leverage
            commission_cost = position_size * effective_leverage * (COMMISSION_PCT * 2) / 100
            net_pnl = (position_size * leveraged_pnl_pct / 100) - commission_cost
            
            equity += net_pnl
            
            trade_results.append({
                'trade_num': trade['trade_num'],
                'concurrent': 1,
                'position': 1,
                'equity_used_pct': equity_pct * 100,
                'leverage': effective_leverage,
                'gross_pnl_pct': gross_pnl_pct,
                'net_pnl': net_pnl,
                'equity_after': equity
            })
            
        else:
            # Multiple concurrent trades
            if scale_on_consecutive:
                # SCALING MODE: First trade = base, second trade = boosted
                # Allocate equity between trades
                total_allocation = equity_pct
                
                for i, trade in enumerate(trades):
                    if i == 0:
                        # First entry - base parameters
                        alloc = total_allocation / num_concurrent
                        effective_leverage = leverage
                    else:
                        # Consecutive entry - potentially boosted
                        alloc = total_allocation / num_concurrent
                        effective_leverage = leverage * scale_leverage_boost
                    
                    position_size = equity * alloc
                    
                    gross_pnl_pct = trade['pnl_pct']
                    leveraged_pnl_pct = gross_pnl_pct * effective_leverage
                    commission_cost = position_size * effective_leverage * (COMMISSION_PCT * 2) / 100
                    net_pnl = (position_size * leveraged_pnl_pct / 100) - commission_cost
                    
                    equity += net_pnl
                    
                    trade_results.append({
                        'trade_num': trade['trade_num'],
                        'concurrent': num_concurrent,
                        'position': i + 1,
                        'equity_used_pct': alloc * 100,
                        'leverage': effective_leverage,
                        'gross_pnl_pct': gross_pnl_pct,
                        'net_pnl': net_pnl,
                        'equity_after': equity
                    })
            else:
                # NON-SCALING MODE: Skip additional entries
                trade = trades[0]
                position_size = equity * equity_pct
                effective_leverage = leverage
                
                gross_pnl_pct = trade['pnl_pct']
                leveraged_pnl_pct = gross_pnl_pct * effective_leverage
                commission_cost = position_size * effective_leverage * (COMMISSION_PCT * 2) / 100
                net_pnl = (position_size * leveraged_pnl_pct / 100) - commission_cost
                
                equity += net_pnl
                
                trade_results.append({
                    'trade_num': trade['trade_num'],
                    'concurrent': num_concurrent,
                    'position': 1,
                    'equity_used_pct': equity_pct * 100,
                    'leverage': effective_leverage,
                    'gross_pnl_pct': gross_pnl_pct,
                    'net_pnl': net_pnl,
                    'equity_after': equity
                })
        
        # Track equity curve and drawdown
        equity_curve.append(equity)
        peak_equity = max(peak_equity, equity)
        drawdown_pct = ((peak_equity - equity) / peak_equity) * 100
        max_drawdown_pct = max(max_drawdown_pct, drawdown_pct)
    
    # Calculate final metrics
    total_pnl = equity - INITIAL_CAPITAL
    roi_pct = (total_pnl / INITIAL_CAPITAL) * 100
    
    # Risk/Reward ratio
    risk_reward = roi_pct / max_drawdown_pct if max_drawdown_pct > 0 else float('inf')
    
    return {
        'equity_pct': equity_pct,
        'leverage': leverage,
        'scale_on_consecutive': scale_on_consecutive,
        'scale_leverage_boost': scale_leverage_boost,
        'final_equity': equity,
        'total_pnl': total_pnl,
        'roi_pct': roi_pct,
        'max_drawdown_pct': max_drawdown_pct,
        'risk_reward': risk_reward,
        'trades': trade_results
    }


def run_optimization():
    """Run full grid search across equity %, leverage, and scaling options."""
    
    equity_levels = [0.25, 0.50, 0.75, 1.00]
    leverage_levels = [5, 10, 15, 20]
    scale_options = [False, True]
    scale_boost_options = [1.0, 1.5, 2.0]  # 1.0 = no boost, 2.0 = double on 2nd signal
    
    results = []
    
    print("=" * 100)
    print("EQUITY & LEVERAGE OPTIMIZER - ML v11.2 Long Strategy")
    print("=" * 100)
    print(f"Initial Capital: ${INITIAL_CAPITAL:,.2f}")
    print(f"Commission: {COMMISSION_PCT}% per trade")
    print(f"Total Trades: {len(TRADES)}")
    print("=" * 100)
    
    # Run all combinations
    for equity_pct in equity_levels:
        for leverage in leverage_levels:
            # Non-scaling mode
            result = simulate_strategy(equity_pct, leverage, scale_on_consecutive=False)
            results.append(result)
            
            # Scaling mode with different boost levels
            for boost in scale_boost_options:
                result = simulate_strategy(equity_pct, leverage, scale_on_consecutive=True, scale_leverage_boost=boost)
                results.append(result)
    
    # Sort by ROI
    results.sort(key=lambda x: x['roi_pct'], reverse=True)
    
    # Find best by different criteria
    best_roi = max(results, key=lambda x: x['roi_pct'])
    best_risk_reward = max(results, key=lambda x: x['risk_reward'])
    safest = min([r for r in results if r['roi_pct'] > 0], key=lambda x: x['max_drawdown_pct'])
    
    # Print top 15 by ROI
    print("\n📊 TOP 15 CONFIGURATIONS BY ROI:")
    print("-" * 100)
    print(f"{'Rank':<5} {'Equity%':<10} {'Leverage':<10} {'Scale':<8} {'Boost':<8} {'Final $':<12} {'ROI %':<10} {'Max DD%':<10} {'R/R':<8}")
    print("-" * 100)
    
    for i, r in enumerate(results[:15], 1):
        scale_str = "Yes" if r['scale_on_consecutive'] else "No"
        boost_str = f"{r['scale_leverage_boost']:.1f}x" if r['scale_on_consecutive'] else "-"
        print(f"{i:<5} {r['equity_pct']*100:<10.0f} {r['leverage']:<10}x {scale_str:<8} {boost_str:<8} ${r['final_equity']:<11,.2f} {r['roi_pct']:<10.2f} {r['max_drawdown_pct']:<10.2f} {r['risk_reward']:<8.2f}")
    
    # Print analysis sections
    print("\n" + "=" * 100)
    print("🎯 BEST BY CATEGORY:")
    print("=" * 100)
    
    print(f"\n💰 HIGHEST ROI:")
    print(f"   Equity: {best_roi['equity_pct']*100:.0f}%")
    print(f"   Leverage: {best_roi['leverage']}x")
    print(f"   Scale: {'Yes' if best_roi['scale_on_consecutive'] else 'No'}")
    print(f"   Boost: {best_roi['scale_leverage_boost']:.1f}x")
    print(f"   Final: ${best_roi['final_equity']:,.2f} ({best_roi['roi_pct']:.2f}% ROI)")
    print(f"   Max DD: {best_roi['max_drawdown_pct']:.2f}%")
    print(f"   Risk/Reward: {best_roi['risk_reward']:.2f}")
    
    print(f"\n⚖️ BEST RISK/REWARD:")
    print(f"   Equity: {best_risk_reward['equity_pct']*100:.0f}%")
    print(f"   Leverage: {best_risk_reward['leverage']}x")
    print(f"   Scale: {'Yes' if best_risk_reward['scale_on_consecutive'] else 'No'}")
    print(f"   Boost: {best_risk_reward['scale_leverage_boost']:.1f}x")
    print(f"   Final: ${best_risk_reward['final_equity']:,.2f} ({best_risk_reward['roi_pct']:.2f}% ROI)")
    print(f"   Max DD: {best_risk_reward['max_drawdown_pct']:.2f}%")
    print(f"   Risk/Reward: {best_risk_reward['risk_reward']:.2f}")
    
    print(f"\n🛡️ SAFEST (Lowest DD with positive ROI):")
    print(f"   Equity: {safest['equity_pct']*100:.0f}%")
    print(f"   Leverage: {safest['leverage']}x")
    print(f"   Scale: {'Yes' if safest['scale_on_consecutive'] else 'No'}")
    print(f"   Boost: {safest['scale_leverage_boost']:.1f}x")
    print(f"   Final: ${safest['final_equity']:,.2f} ({safest['roi_pct']:.2f}% ROI)")
    print(f"   Max DD: {safest['max_drawdown_pct']:.2f}%")
    print(f"   Risk/Reward: {safest['risk_reward']:.2f}")
    
    # Answer user's specific question
    print("\n" + "=" * 100)
    print("❓ USER QUESTION: 100% equity + 20x leverage with 2x boost on consecutive signals?")
    print("=" * 100)
    
    # Find this specific configuration
    user_config = next((r for r in results if r['equity_pct'] == 1.0 and r['leverage'] == 20 
                        and r['scale_on_consecutive'] and r['scale_leverage_boost'] == 2.0), None)
    
    if user_config:
        print(f"\n   Your proposed setup (100% equity, 20x, scale with 2x boost):")
        print(f"   Final: ${user_config['final_equity']:,.2f}")
        print(f"   ROI: {user_config['roi_pct']:.2f}%")
        print(f"   Max DD: {user_config['max_drawdown_pct']:.2f}%")
        print(f"   Risk/Reward: {user_config['risk_reward']:.2f}")
    
    # Compare to alternatives
    print("\n   📊 COMPARISON TO ALTERNATIVES:")
    
    # Same equity, same leverage, no scaling
    no_scale = next((r for r in results if r['equity_pct'] == 1.0 and r['leverage'] == 20 
                     and not r['scale_on_consecutive']), None)
    if no_scale:
        print(f"\n   100% equity, 20x, NO scaling:")
        print(f"   Final: ${no_scale['final_equity']:,.2f} ({no_scale['roi_pct']:.2f}% ROI)")
        print(f"   Max DD: {no_scale['max_drawdown_pct']:.2f}%")
    
    # Half leverage, with scaling
    half_lev = next((r for r in results if r['equity_pct'] == 1.0 and r['leverage'] == 10 
                     and r['scale_on_consecutive'] and r['scale_leverage_boost'] == 2.0), None)
    if half_lev:
        print(f"\n   100% equity, 10x base (20x on consecutive), scaling:")
        print(f"   Final: ${half_lev['final_equity']:,.2f} ({half_lev['roi_pct']:.2f}% ROI)")
        print(f"   Max DD: {half_lev['max_drawdown_pct']:.2f}%")
        print(f"   Risk/Reward: {half_lev['risk_reward']:.2f}")
    
    # Heat map visualization
    print("\n" + "=" * 100)
    print("📈 ROI HEAT MAP (100% Equity, Scaling ON, varying Leverage & Boost):")
    print("=" * 100)
    print(f"{'Leverage':<12}", end="")
    for boost in scale_boost_options:
        print(f"Boost {boost:.1f}x      ", end="")
    print()
    print("-" * 60)
    
    for lev in leverage_levels:
        print(f"{lev}x          ", end="")
        for boost in scale_boost_options:
            r = next((r for r in results if r['equity_pct'] == 1.0 and r['leverage'] == lev 
                      and r['scale_on_consecutive'] and r['scale_leverage_boost'] == boost), None)
            if r:
                print(f"{r['roi_pct']:>6.0f}%       ", end="")
        print()
    
    # Drawdown heat map
    print("\n" + "=" * 100)
    print("⚠️ MAX DRAWDOWN HEAT MAP (100% Equity, Scaling ON):")
    print("=" * 100)
    print(f"{'Leverage':<12}", end="")
    for boost in scale_boost_options:
        print(f"Boost {boost:.1f}x      ", end="")
    print()
    print("-" * 60)
    
    for lev in leverage_levels:
        print(f"{lev}x          ", end="")
        for boost in scale_boost_options:
            r = next((r for r in results if r['equity_pct'] == 1.0 and r['leverage'] == lev 
                      and r['scale_on_consecutive'] and r['scale_leverage_boost'] == boost), None)
            if r:
                print(f"{r['max_drawdown_pct']:>6.1f}%       ", end="")
        print()
    
    # Final recommendation
    print("\n" + "=" * 100)
    print("🏆 RECOMMENDATION:")
    print("=" * 100)
    
    # Find sweet spot: good ROI with acceptable drawdown
    balanced = sorted([r for r in results if r['max_drawdown_pct'] < 60], 
                      key=lambda x: x['risk_reward'], reverse=True)[:3]
    
    print("\nTOP 3 BALANCED OPTIONS (ROI with <60% max drawdown):")
    for i, r in enumerate(balanced, 1):
        scale_str = "Yes" if r['scale_on_consecutive'] else "No"
        boost_str = f"{r['scale_leverage_boost']:.1f}x" if r['scale_on_consecutive'] else "-"
        print(f"\n#{i}: {r['equity_pct']*100:.0f}% equity, {r['leverage']}x leverage, Scale: {scale_str}, Boost: {boost_str}")
        print(f"    Final: ${r['final_equity']:,.2f} | ROI: {r['roi_pct']:.1f}% | MaxDD: {r['max_drawdown_pct']:.1f}% | R/R: {r['risk_reward']:.2f}")
    
    return results


if __name__ == "__main__":
    run_optimization()