#!/usr/bin/env python3
"""
Align Strategy Timing for Fair Comparison
=========================================

This script aligns the timing between original and incremental strategies
by removing early trades from the original strategy that occur before
the incremental strategy starts trading (warmup period).
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime
import json

def load_trade_files():
    """Load both strategy trade files."""
    
    print("📊 LOADING TRADE FILES")
    print("=" * 60)
    
    # Load original strategy trades
    original_file = "../results/trades_15min(15min)_ST3pct.csv"
    incremental_file = "../results/trades_incremental_15min(15min)_ST3pct.csv"
    
    print(f"Loading original trades: {original_file}")
    original_df = pd.read_csv(original_file)
    original_df['entry_time'] = pd.to_datetime(original_df['entry_time'])
    original_df['exit_time'] = pd.to_datetime(original_df['exit_time'])
    
    print(f"Loading incremental trades: {incremental_file}")
    incremental_df = pd.read_csv(incremental_file)
    incremental_df['entry_time'] = pd.to_datetime(incremental_df['entry_time'])
    incremental_df['exit_time'] = pd.to_datetime(incremental_df['exit_time'])
    
    print(f"Original trades: {len(original_df)} total")
    print(f"Incremental trades: {len(incremental_df)} total")
    
    return original_df, incremental_df

def find_alignment_point(original_df, incremental_df):
    """Find the point where both strategies should start for fair comparison."""
    
    print(f"\n🕐 FINDING ALIGNMENT POINT")
    print("=" * 60)
    
    # Find when incremental strategy starts trading
    incremental_start = incremental_df[incremental_df['type'] == 'BUY']['entry_time'].min()
    print(f"Incremental strategy first trade: {incremental_start}")
    
    # Find original strategy trades before this point
    original_buys = original_df[original_df['type'] == 'BUY']
    early_trades = original_buys[original_buys['entry_time'] < incremental_start]
    
    print(f"Original trades before incremental start: {len(early_trades)}")
    
    if len(early_trades) > 0:
        print(f"First original trade: {original_buys['entry_time'].min()}")
        print(f"Last early trade: {early_trades['entry_time'].max()}")
        print(f"Time gap: {incremental_start - original_buys['entry_time'].min()}")
        
        # Show the early trades that will be excluded
        print(f"\n📋 EARLY TRADES TO EXCLUDE:")
        for i, trade in early_trades.iterrows():
            print(f"  {trade['entry_time']} - ${trade['entry_price']:.0f}")
    
    return incremental_start

def align_strategies(original_df, incremental_df, alignment_time):
    """Align both strategies to start at the same time."""
    
    print(f"\n⚖️ ALIGNING STRATEGIES")
    print("=" * 60)
    
    # Filter original strategy to start from alignment time
    aligned_original = original_df[original_df['entry_time'] >= alignment_time].copy()
    
    # Incremental strategy remains the same (already starts at alignment time)
    aligned_incremental = incremental_df.copy()
    
    print(f"Original trades after alignment: {len(aligned_original)}")
    print(f"Incremental trades: {len(aligned_incremental)}")
    
    # Reset indices for clean comparison
    aligned_original = aligned_original.reset_index(drop=True)
    aligned_incremental = aligned_incremental.reset_index(drop=True)
    
    return aligned_original, aligned_incremental

def calculate_aligned_performance(aligned_original, aligned_incremental):
    """Calculate performance metrics for aligned strategies."""
    
    print(f"\n💰 CALCULATING ALIGNED PERFORMANCE")
    print("=" * 60)
    
    def calculate_strategy_performance(df, strategy_name):
        """Calculate performance for a single strategy."""
        
        # Filter to complete trades (buy + sell pairs)
        buy_signals = df[df['type'] == 'BUY'].copy()
        sell_signals = df[df['type'].str.contains('EXIT|EOD', na=False)].copy()
        
        print(f"\n{strategy_name}:")
        print(f"  Buy signals: {len(buy_signals)}")
        print(f"  Sell signals: {len(sell_signals)}")
        
        if len(buy_signals) == 0:
            return {
                'final_value': 10000,
                'total_return': 0.0,
                'trade_count': 0,
                'win_rate': 0.0,
                'avg_trade': 0.0
            }
        
        # Calculate performance using same logic as comparison script
        initial_usd = 10000
        current_usd = initial_usd
        
        for i, buy_trade in buy_signals.iterrows():
            # Find corresponding sell trade
            sell_trades = sell_signals[sell_signals['entry_time'] == buy_trade['entry_time']]
            if len(sell_trades) == 0:
                continue
                
            sell_trade = sell_trades.iloc[0]
            
            # Calculate trade performance
            entry_price = buy_trade['entry_price']
            exit_price = sell_trade['exit_price']
            profit_pct = sell_trade['profit_pct']
            
            # Apply profit/loss
            current_usd *= (1 + profit_pct)
        
        total_return = ((current_usd - initial_usd) / initial_usd) * 100
        
        # Calculate trade statistics
        profits = sell_signals['profit_pct'].values
        winning_trades = len(profits[profits > 0])
        win_rate = (winning_trades / len(profits)) * 100 if len(profits) > 0 else 0
        avg_trade = np.mean(profits) * 100 if len(profits) > 0 else 0
        
        print(f"  Final value: ${current_usd:,.0f}")
        print(f"  Total return: {total_return:.1f}%")
        print(f"  Win rate: {win_rate:.1f}%")
        print(f"  Average trade: {avg_trade:.2f}%")
        
        return {
            'final_value': current_usd,
            'total_return': total_return,
            'trade_count': len(profits),
            'win_rate': win_rate,
            'avg_trade': avg_trade,
            'profits': profits.tolist()
        }
    
    # Calculate performance for both strategies
    original_perf = calculate_strategy_performance(aligned_original, "Aligned Original")
    incremental_perf = calculate_strategy_performance(aligned_incremental, "Incremental")
    
    # Compare performance
    print(f"\n📊 PERFORMANCE COMPARISON:")
    print("=" * 60)
    print(f"Original (aligned):  ${original_perf['final_value']:,.0f} ({original_perf['total_return']:+.1f}%)")
    print(f"Incremental:         ${incremental_perf['final_value']:,.0f} ({incremental_perf['total_return']:+.1f}%)")
    
    difference = incremental_perf['total_return'] - original_perf['total_return']
    print(f"Difference:          {difference:+.1f}%")
    
    if abs(difference) < 5:
        print("✅ Performance is now closely aligned!")
    elif difference > 0:
        print("📈 Incremental strategy outperforms after alignment")
    else:
        print("📉 Original strategy still outperforms")
    
    return original_perf, incremental_perf

def save_aligned_results(aligned_original, aligned_incremental, original_perf, incremental_perf):
    """Save aligned results for further analysis."""
    
    print(f"\n💾 SAVING ALIGNED RESULTS")
    print("=" * 60)
    
    # Save aligned trade files
    aligned_original.to_csv("../results/trades_original_aligned.csv", index=False)
    aligned_incremental.to_csv("../results/trades_incremental_aligned.csv", index=False)
    
    print("Saved aligned trade files:")
    print("  - ../results/trades_original_aligned.csv")
    print("  - ../results/trades_incremental_aligned.csv")
    
    # Save performance comparison
    comparison_results = {
        'alignment_analysis': {
            'original_performance': original_perf,
            'incremental_performance': incremental_perf,
            'performance_difference': incremental_perf['total_return'] - original_perf['total_return'],
            'trade_count_difference': incremental_perf['trade_count'] - original_perf['trade_count'],
            'win_rate_difference': incremental_perf['win_rate'] - original_perf['win_rate']
        },
        'timestamp': datetime.now().isoformat()
    }
    
    with open("../results/aligned_performance_comparison.json", "w") as f:
        json.dump(comparison_results, f, indent=2)
    
    print("  - ../results/aligned_performance_comparison.json")

def create_aligned_visualization(aligned_original, aligned_incremental):
    """Create visualization of aligned strategies."""
    
    print(f"\n📊 CREATING ALIGNED VISUALIZATION")
    print("=" * 60)
    
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 10))
    
    # Get buy signals for plotting
    orig_buys = aligned_original[aligned_original['type'] == 'BUY']
    inc_buys = aligned_incremental[aligned_incremental['type'] == 'BUY']
    
    # Plot 1: Trade timing comparison
    ax1.scatter(orig_buys['entry_time'], orig_buys['entry_price'], 
               alpha=0.7, label='Original (Aligned)', color='blue', s=40)
    ax1.scatter(inc_buys['entry_time'], inc_buys['entry_price'], 
               alpha=0.7, label='Incremental', color='red', s=40)
    ax1.set_title('Aligned Strategy Trade Timing Comparison')
    ax1.set_xlabel('Date')
    ax1.set_ylabel('Entry Price ($)')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # Plot 2: Cumulative performance
    def calculate_cumulative_returns(df):
        """Calculate cumulative returns over time."""
        buy_signals = df[df['type'] == 'BUY'].copy()
        sell_signals = df[df['type'].str.contains('EXIT|EOD', na=False)].copy()
        
        cumulative_returns = []
        current_value = 10000
        dates = []
        
        for i, buy_trade in buy_signals.iterrows():
            sell_trades = sell_signals[sell_signals['entry_time'] == buy_trade['entry_time']]
            if len(sell_trades) == 0:
                continue
                
            sell_trade = sell_trades.iloc[0]
            current_value *= (1 + sell_trade['profit_pct'])
            
            cumulative_returns.append(current_value)
            dates.append(sell_trade['exit_time'])
        
        return dates, cumulative_returns
    
    orig_dates, orig_returns = calculate_cumulative_returns(aligned_original)
    inc_dates, inc_returns = calculate_cumulative_returns(aligned_incremental)
    
    if orig_dates:
        ax2.plot(orig_dates, orig_returns, label='Original (Aligned)', color='blue', linewidth=2)
    if inc_dates:
        ax2.plot(inc_dates, inc_returns, label='Incremental', color='red', linewidth=2)
    
    ax2.set_title('Aligned Strategy Cumulative Performance')
    ax2.set_xlabel('Date')
    ax2.set_ylabel('Portfolio Value ($)')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('../results/aligned_strategy_comparison.png', dpi=300, bbox_inches='tight')
    print("Visualization saved: ../results/aligned_strategy_comparison.png")

def main():
    """Main alignment function."""
    
    print("🚀 ALIGNING STRATEGY TIMING FOR FAIR COMPARISON")
    print("=" * 80)
    
    try:
        # Load trade files
        original_df, incremental_df = load_trade_files()
        
        # Find alignment point
        alignment_time = find_alignment_point(original_df, incremental_df)
        
        # Align strategies
        aligned_original, aligned_incremental = align_strategies(
            original_df, incremental_df, alignment_time
        )
        
        # Calculate aligned performance
        original_perf, incremental_perf = calculate_aligned_performance(
            aligned_original, aligned_incremental
        )
        
        # Save results
        save_aligned_results(aligned_original, aligned_incremental, 
                           original_perf, incremental_perf)
        
        # Create visualization
        create_aligned_visualization(aligned_original, aligned_incremental)
        
        print(f"\n✅ ALIGNMENT COMPLETED SUCCESSFULLY!")
        print("=" * 80)
        print("The strategies are now aligned for fair comparison.")
        print("Check the results/ directory for aligned trade files and analysis.")
        
        return True
        
    except Exception as e:
        print(f"\n❌ Error during alignment: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)