#!/usr/bin/env python3
"""
Enhanced test script for incremental backtester using real BTC data
with comprehensive visualization and analysis features.

ENHANCED FEATURES:
- Stop Loss/Take Profit Visualization: Different colors and markers for exit types
  * Green triangles (^): Buy entries
  * Blue triangles (v): Strategy exits  
  * Dark red X: Stop loss exits (prominent markers)
  * Gold stars (*): Take profit exits
  * Gray squares: End-of-day exits

- Portfolio Tracking: Combined USD + BTC value calculation
  * Real-time portfolio value based on current BTC price
  * Separate tracking of USD balance and BTC holdings
  * Portfolio composition visualization

- Three-Panel Analysis:
  1. Price chart with trading signals and exit types
  2. Portfolio value over time with profit/loss zones
  3. Portfolio composition (USD vs BTC value breakdown)

- Comprehensive Data Export:
  * CSV: Individual trades with entry/exit details
  * JSON: Complete performance statistics
  * CSV: Portfolio value tracking over time
  * PNG: Multi-panel visualization charts

- Performance Analysis:
  * Exit type breakdown and performance
  * Win/loss distribution analysis
  * Best/worst trade identification
  * Detailed trade-by-trade logging
"""

import os
import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime
from typing import Dict, List
import warnings
import json
warnings.filterwarnings('ignore')

# Add the project root to the path
sys.path.insert(0, os.path.abspath('.'))

from cycles.IncStrategies.inc_backtester import IncBacktester, BacktestConfig
from cycles.IncStrategies.random_strategy import IncRandomStrategy
from cycles.IncStrategies.metatrend_strategy import IncMetaTrendStrategy
from cycles.utils.storage import Storage
from cycles.utils.data_utils import aggregate_to_minutes


def save_trades_to_csv(trades: List[Dict], filename: str) -> None:
    """Save trades to CSV file in the same format as existing trades file."""
    if not trades:
        print("No trades to save")
        return
    
    # Convert trades to the exact format of the existing file
    formatted_trades = []
    
    for trade in trades:
        # Create entry row (buy signal)
        entry_row = {
            'entry_time': trade['entry_time'],
            'exit_time': '',  # Empty for entry row
            'entry_price': trade['entry'],
            'exit_price': '',  # Empty for entry row
            'profit_pct': 0.0,  # 0 for entry
            'type': 'BUY',
            'fee_usd': trade.get('entry_fee_usd', 10.0)  # Default fee if not available
        }
        formatted_trades.append(entry_row)
        
        # Create exit row (sell signal)
        exit_type = trade.get('type', 'META_TREND_EXIT_SIGNAL')
        if exit_type == 'STRATEGY_EXIT':
            exit_type = 'META_TREND_EXIT_SIGNAL'
        elif exit_type == 'STOP_LOSS':
            exit_type = 'STOP_LOSS'
        elif exit_type == 'TAKE_PROFIT':
            exit_type = 'TAKE_PROFIT'
        elif exit_type == 'EOD':
            exit_type = 'EOD'
        
        exit_row = {
            'entry_time': trade['entry_time'],
            'exit_time': trade['exit_time'],
            'entry_price': trade['entry'],
            'exit_price': trade['exit'],
            'profit_pct': trade['profit_pct'],
            'type': exit_type,
            'fee_usd': trade.get('exit_fee_usd', trade.get('total_fees_usd', 10.0))
        }
        formatted_trades.append(exit_row)
    
    # Convert to DataFrame and save
    trades_df = pd.DataFrame(formatted_trades)
    
    # Ensure the columns are in the exact same order
    column_order = ['entry_time', 'exit_time', 'entry_price', 'exit_price', 'profit_pct', 'type', 'fee_usd']
    trades_df = trades_df[column_order]
    
    # Save with same formatting
    trades_df.to_csv(filename, index=False)
    print(f"Saved {len(formatted_trades)} trade signals ({len(trades)} complete trades) to: {filename}")
    
    # Print summary for comparison
    buy_signals = len([t for t in formatted_trades if t['type'] == 'BUY'])
    sell_signals = len(formatted_trades) - buy_signals
    print(f"  - Buy signals: {buy_signals}")
    print(f"  - Sell signals: {sell_signals}")
    
    # Show exit type breakdown
    exit_types = {}
    for trade in formatted_trades:
        if trade['type'] != 'BUY':
            exit_type = trade['type']
            exit_types[exit_type] = exit_types.get(exit_type, 0) + 1
    
    if exit_types:
        print(f"  - Exit types: {exit_types}")


def save_stats_to_json(stats: Dict, filename: str) -> None:
    """Save statistics to JSON file."""
    # Convert any datetime objects to strings for JSON serialization
    stats_copy = stats.copy()
    for key, value in stats_copy.items():
        if isinstance(value, pd.Timestamp):
            stats_copy[key] = value.isoformat()
        elif isinstance(value, dict):
            for k, v in value.items():
                if isinstance(v, pd.Timestamp):
                    value[k] = v.isoformat()
    
    with open(filename, 'w') as f:
        json.dump(stats_copy, f, indent=2, default=str)
    print(f"Saved statistics to: {filename}")


def calculate_portfolio_over_time(data: pd.DataFrame, trades: List[Dict], initial_usd: float, debug: bool = False) -> pd.DataFrame:
    """Calculate portfolio value over time with proper USD + BTC tracking."""
    print("Calculating portfolio value over time...")
    
    # Create portfolio tracking with detailed state
    portfolio_data = data[['close']].copy()
    portfolio_data['portfolio_value'] = initial_usd
    portfolio_data['usd_balance'] = initial_usd
    portfolio_data['btc_balance'] = 0.0
    portfolio_data['position'] = 0  # 0 = cash, 1 = in position
    
    if not trades:
        return portfolio_data
    
    # Initialize state
    current_usd = initial_usd
    current_btc = 0.0
    in_position = False
    
    # Sort trades by entry time
    sorted_trades = sorted(trades, key=lambda x: x['entry_time'])
    trade_idx = 0
    
    print(f"Processing {len(sorted_trades)} trades across {len(portfolio_data)} data points...")
    
    for i, (timestamp, row) in enumerate(portfolio_data.iterrows()):
        current_price = row['close']
        
        # Check if we need to execute any trades at this timestamp
        while trade_idx < len(sorted_trades):
            trade = sorted_trades[trade_idx]
            
            # Check for entry
            if trade['entry_time'] <= timestamp and not in_position:
                # Execute buy order
                entry_price = trade['entry']
                current_btc = current_usd / entry_price
                current_usd = 0.0
                in_position = True
                if debug:
                    print(f"Entry {trade_idx + 1}: Buy at ${entry_price:.2f}, BTC: {current_btc:.6f}")
                break
                
            # Check for exit
            elif trade['exit_time'] <= timestamp and in_position:
                # Execute sell order
                exit_price = trade['exit']
                current_usd = current_btc * exit_price
                current_btc = 0.0
                in_position = False
                exit_type = trade.get('type', 'STRATEGY_EXIT')
                if debug:
                    print(f"Exit {trade_idx + 1}: {exit_type} at ${exit_price:.2f}, USD: ${current_usd:.2f}")
                trade_idx += 1
                break
            else:
                break
        
        # Calculate total portfolio value (USD + BTC value)
        btc_value = current_btc * current_price
        total_value = current_usd + btc_value
        
        # Update portfolio data
        portfolio_data.iloc[i, portfolio_data.columns.get_loc('portfolio_value')] = total_value
        portfolio_data.iloc[i, portfolio_data.columns.get_loc('usd_balance')] = current_usd
        portfolio_data.iloc[i, portfolio_data.columns.get_loc('btc_balance')] = current_btc
        portfolio_data.iloc[i, portfolio_data.columns.get_loc('position')] = 1 if in_position else 0
    
    return portfolio_data


def create_comprehensive_plot(data: pd.DataFrame, trades: List[Dict], portfolio_data: pd.DataFrame,
                            strategy_name: str, save_path: str) -> None:
    """Create comprehensive plot with price, trades, and portfolio value."""
    
    print(f"Creating comprehensive plot with {len(data)} data points and {len(trades)} trades...")
    
    # Create figure with subplots
    fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(16, 16), 
                                       gridspec_kw={'height_ratios': [2, 1, 1]})
    
    # Plot 1: Price action with trades
    ax1.plot(data.index, data['close'], label='BTC Price', color='black', linewidth=1.5)
    
    # Plot trades with different markers for different exit types
    if trades:
        entry_times = [trade['entry_time'] for trade in trades]
        entry_prices = [trade['entry'] for trade in trades]
        
        # Separate exits by type
        strategy_exits = []
        stop_loss_exits = []
        take_profit_exits = []
        eod_exits = []
        
        for trade in trades:
            exit_type = trade.get('type', 'STRATEGY_EXIT')
            exit_data = (trade['exit_time'], trade['exit'])
            
            if exit_type == 'STOP_LOSS':
                stop_loss_exits.append(exit_data)
            elif exit_type == 'TAKE_PROFIT':
                take_profit_exits.append(exit_data)
            elif exit_type == 'EOD':
                eod_exits.append(exit_data)
            else:
                strategy_exits.append(exit_data)
        
        # Plot entry points (green triangles)
        ax1.scatter(entry_times, entry_prices, color='darkgreen', marker='^', 
                  s=100, label=f'Buy ({len(entry_times)})', zorder=6, alpha=0.9, edgecolors='white', linewidth=1)
        
        # Plot different types of exits with distinct styling
        if strategy_exits:
            exit_times, exit_prices = zip(*strategy_exits)
            ax1.scatter(exit_times, exit_prices, color='blue', marker='v', 
                      s=100, label=f'Strategy Exit ({len(strategy_exits)})', zorder=5, alpha=0.8, edgecolors='white', linewidth=1)
        
        if stop_loss_exits:
            exit_times, exit_prices = zip(*stop_loss_exits)
            ax1.scatter(exit_times, exit_prices, color='darkred', marker='X', 
                      s=150, label=f'Stop Loss ({len(stop_loss_exits)})', zorder=7, alpha=1.0, edgecolors='white', linewidth=2)
        
        if take_profit_exits:
            exit_times, exit_prices = zip(*take_profit_exits)
            ax1.scatter(exit_times, exit_prices, color='gold', marker='*', 
                      s=150, label=f'Take Profit ({len(take_profit_exits)})', zorder=6, alpha=0.9, edgecolors='black', linewidth=1)
        
        if eod_exits:
            exit_times, exit_prices = zip(*eod_exits)
            ax1.scatter(exit_times, exit_prices, color='gray', marker='s', 
                      s=80, label=f'End of Day ({len(eod_exits)})', zorder=5, alpha=0.8, edgecolors='white', linewidth=1)
        
        # Print exit type summary
        print(f"Exit types: Strategy={len(strategy_exits)}, Stop Loss={len(stop_loss_exits)}, "
              f"Take Profit={len(take_profit_exits)}, EOD={len(eod_exits)}")
    
    ax1.set_title(f'{strategy_name} - BTC Trading Signals (Q1 2023)', fontsize=16, fontweight='bold')
    ax1.set_ylabel('Price (USD)', fontsize=12)
    ax1.legend(loc='upper left', fontsize=10)
    ax1.grid(True, alpha=0.3)
    
    # Plot 2: Portfolio value over time
    ax2.plot(portfolio_data.index, portfolio_data['portfolio_value'], 
             label='Total Portfolio Value', color='blue', linewidth=2)
    ax2.axhline(y=portfolio_data['portfolio_value'].iloc[0], color='gray', 
                linestyle='--', alpha=0.7, label='Initial Value')
    
    # Add profit/loss shading
    initial_value = portfolio_data['portfolio_value'].iloc[0]
    profit_mask = portfolio_data['portfolio_value'] > initial_value
    loss_mask = portfolio_data['portfolio_value'] < initial_value
    
    ax2.fill_between(portfolio_data.index, portfolio_data['portfolio_value'], initial_value,
                     where=profit_mask, color='green', alpha=0.2, label='Profit Zone')
    ax2.fill_between(portfolio_data.index, portfolio_data['portfolio_value'], initial_value,
                     where=loss_mask, color='red', alpha=0.2, label='Loss Zone')
    
    ax2.set_title('Portfolio Value Over Time (USD + BTC)', fontsize=14, fontweight='bold')
    ax2.set_ylabel('Portfolio Value (USD)', fontsize=12)
    ax2.legend(loc='upper left', fontsize=10)
    ax2.grid(True, alpha=0.3)
    
    # Plot 3: Portfolio composition (USD vs BTC value)
    usd_values = portfolio_data['usd_balance']
    btc_values = portfolio_data['btc_balance'] * portfolio_data['close']
    
    ax3.fill_between(portfolio_data.index, 0, usd_values, 
                     color='green', alpha=0.6, label='USD Balance')
    ax3.fill_between(portfolio_data.index, usd_values, usd_values + btc_values, 
                     color='orange', alpha=0.6, label='BTC Value')
    
    # Mark position periods
    position_mask = portfolio_data['position'] == 1
    if position_mask.any():
        ax3.fill_between(portfolio_data.index, 0, portfolio_data['portfolio_value'],
                        where=position_mask, color='orange', alpha=0.2, label='In Position')
    
    ax3.set_title('Portfolio Composition (USD vs BTC)', fontsize=14, fontweight='bold')
    ax3.set_ylabel('Value (USD)', fontsize=12)
    ax3.set_xlabel('Date', fontsize=12)
    ax3.legend(loc='upper left', fontsize=10)
    ax3.grid(True, alpha=0.3)
    
    # Format x-axis for all plots
    for ax in [ax1, ax2, ax3]:
        ax.xaxis.set_major_locator(mdates.WeekdayLocator())
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%m-%d'))
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
    
    # Save plot
    plt.tight_layout()
    plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    print(f"Comprehensive plot saved to: {save_path}")


def compare_with_existing_trades(new_trades_file: str, existing_trades_file: str = "results/trades_15min(15min)_ST3pct.csv") -> None:
    """Compare the new incremental trades with existing strategy trades."""
    try:
        if not os.path.exists(existing_trades_file):
            print(f"Existing trades file not found: {existing_trades_file}")
            return
            
        print(f"\n📊 COMPARING WITH EXISTING STRATEGY:")
        
        # Load both files
        new_df = pd.read_csv(new_trades_file)
        existing_df = pd.read_csv(existing_trades_file)
        
        # Count signals
        new_buy_signals = len(new_df[new_df['type'] == 'BUY'])
        new_sell_signals = len(new_df[new_df['type'] != 'BUY'])
        
        existing_buy_signals = len(existing_df[existing_df['type'] == 'BUY'])
        existing_sell_signals = len(existing_df[existing_df['type'] != 'BUY'])
        
        print(f"📈 SIGNAL COMPARISON:")
        print(f"  Incremental Strategy:")
        print(f"    - Buy signals: {new_buy_signals}")
        print(f"    - Sell signals: {new_sell_signals}")
        print(f"  Existing Strategy:")
        print(f"    - Buy signals: {existing_buy_signals}")
        print(f"    - Sell signals: {existing_sell_signals}")
        
        # Compare exit types
        new_exit_types = new_df[new_df['type'] != 'BUY']['type'].value_counts().to_dict()
        existing_exit_types = existing_df[existing_df['type'] != 'BUY']['type'].value_counts().to_dict()
        
        print(f"\n🎯 EXIT TYPE COMPARISON:")
        print(f"  Incremental Strategy: {new_exit_types}")
        print(f"  Existing Strategy: {existing_exit_types}")
        
        # Calculate profit comparison
        new_profits = new_df[new_df['type'] != 'BUY']['profit_pct'].sum()
        existing_profits = existing_df[existing_df['type'] != 'BUY']['profit_pct'].sum()
        
        print(f"\n💰 PROFIT COMPARISON:")
        print(f"  Incremental Strategy: {new_profits*100:.2f}% total")
        print(f"  Existing Strategy: {existing_profits*100:.2f}% total")
        print(f"  Difference: {(new_profits - existing_profits)*100:.2f}%")
        
    except Exception as e:
        print(f"Error comparing trades: {e}")


def test_single_strategy():
    """Test a single strategy and create comprehensive analysis."""
    print("\n" + "="*60)
    print("TESTING SINGLE STRATEGY")
    print("="*60)
    
    # Create storage instance
    storage = Storage()
    
    # Create backtester configuration using 3 months of data
    config = BacktestConfig(
        data_file="btcusd_1-min_data.csv",
        start_date="2025-01-01",
        end_date="2025-05-01",
        initial_usd=10000,
        stop_loss_pct=0.03,  # 3% stop loss to match existing
        take_profit_pct=0.0
    )
    
    # Create strategy
    strategy = IncMetaTrendStrategy(
        name="metatrend",
        weight=1.0,
        params={
            "timeframe": "15min",
            "enable_logging": False
        }
    )
    
    print(f"Testing strategy: {strategy.name}")
    print(f"Strategy timeframe: {strategy.params.get('timeframe', '15min')}")
    print(f"Stop loss: {config.stop_loss_pct*100:.1f}%")
    print(f"Date range: {config.start_date} to {config.end_date}")
    
    # Run backtest
    print(f"\n🚀 Running backtest...")
    backtester = IncBacktester(config, storage)
    result = backtester.run_single_strategy(strategy)
    
    # Print results
    print(f"\n📊 RESULTS:")
    print(f"Strategy: {strategy.__class__.__name__}")
    profit = result['final_usd'] - result['initial_usd']
    print(f"Total Profit: ${profit:.2f} ({result['profit_ratio']*100:.2f}%)")
    print(f"Total Trades: {result['n_trades']}")
    print(f"Win Rate: {result['win_rate']*100:.2f}%")
    print(f"Max Drawdown: {result['max_drawdown']*100:.2f}%")
    print(f"Average Trade: {result['avg_trade']*100:.2f}%")
    print(f"Total Fees: ${result['total_fees_usd']:.2f}")
    
    # Create results directory
    os.makedirs("results", exist_ok=True)
    
    # Save trades in the same format as existing file
    if result['trades']:
        # Create filename matching the existing format
        timeframe = strategy.params.get('timeframe', '15min')
        stop_loss_pct = int(config.stop_loss_pct * 100)
        trades_filename = f"results/trades_incremental_{timeframe}({timeframe})_ST{stop_loss_pct}pct.csv"
        save_trades_to_csv(result['trades'], trades_filename)
        
        # Compare with existing trades
        compare_with_existing_trades(trades_filename)
    
    # Save statistics to JSON
    stats_filename = f"results/incremental_stats_{config.start_date}_{config.end_date}.json"
    save_stats_to_json(result, stats_filename)
    
    # Load and aggregate data for plotting
    print(f"\n📈 CREATING COMPREHENSIVE ANALYSIS...")
    data = storage.load_data("btcusd_1-min_data.csv", config.start_date, config.end_date)
    print(f"Loaded {len(data)} minute-level data points")
    
    # Aggregate to strategy timeframe using existing data_utils
    timeframe_minutes = 15  # Match strategy timeframe
    print(f"Aggregating to {timeframe_minutes}-minute bars using data_utils...")
    aggregated_data = aggregate_to_minutes(data, timeframe_minutes)
    print(f"Aggregated to {len(aggregated_data)} bars")
    
    # Calculate portfolio value over time
    portfolio_data = calculate_portfolio_over_time(aggregated_data, result['trades'], config.initial_usd, debug=False)
    
    # Save portfolio data to CSV
    portfolio_filename = f"results/incremental_portfolio_{config.start_date}_{config.end_date}.csv"
    portfolio_data.to_csv(portfolio_filename)
    print(f"Saved portfolio data to: {portfolio_filename}")
    
    # Create comprehensive plot
    plot_path = f"results/incremental_comprehensive_{config.start_date}_{config.end_date}.png"
    create_comprehensive_plot(aggregated_data, result['trades'], portfolio_data, 
                            "Incremental MetaTrend Strategy", plot_path)
    
    return result


def main():
    """Main test function."""
    print("🚀 Starting Comprehensive Incremental Backtester Test (Q1 2023)")
    print("=" * 80)
    
    try:
        # Test single strategy
        result = test_single_strategy()
        
        print("\n" + "="*80)
        print("✅ TEST COMPLETED SUCCESSFULLY!")
        print("="*80)
        print(f"📁 Check the 'results/' directory for:")
        print(f"   - Trading plot: incremental_comprehensive_q1_2023.png")
        print(f"   - Trades data: trades_incremental_15min(15min)_ST3pct.csv")
        print(f"   - Statistics: incremental_stats_2025-01-01_2025-05-01.json")
        print(f"   - Portfolio data: incremental_portfolio_2025-01-01_2025-05-01.csv")
        print(f"📊 Strategy processed {result['data_points_processed']} data points")
        print(f"🎯 Strategy warmup: {'✅ Complete' if result['warmup_complete'] else '❌ Incomplete'}")
        
        # Show some trade details
        if result['n_trades'] > 0:
            print(f"\n📈 DETAILED TRADE ANALYSIS:")
            print(f"First trade: {result.get('first_trade', {}).get('entry_time', 'N/A')}")
            print(f"Last trade: {result.get('last_trade', {}).get('exit_time', 'N/A')}")
            
            # Analyze trades by exit type
            trades = result['trades']
            
            # Group trades by exit type
            exit_types = {}
            for trade in trades:
                exit_type = trade.get('type', 'STRATEGY_EXIT')
                if exit_type not in exit_types:
                    exit_types[exit_type] = []
                exit_types[exit_type].append(trade)
            
            print(f"\n📊 EXIT TYPE ANALYSIS:")
            for exit_type, type_trades in exit_types.items():
                profits = [trade['profit_pct'] for trade in type_trades]
                avg_profit = np.mean(profits) * 100
                win_rate = len([p for p in profits if p > 0]) / len(profits) * 100
                
                print(f"  {exit_type}:")
                print(f"    Count: {len(type_trades)}")
                print(f"    Avg Profit: {avg_profit:.2f}%")
                print(f"    Win Rate: {win_rate:.1f}%")
                
                if exit_type == 'STOP_LOSS':
                    avg_loss = np.mean([p for p in profits if p <= 0]) * 100
                    print(f"    Avg Loss: {avg_loss:.2f}%")
            
            # Overall profit distribution
            all_profits = [trade['profit_pct'] for trade in trades]
            winning_trades = [p for p in all_profits if p > 0]
            losing_trades = [p for p in all_profits if p <= 0]
            
            print(f"\n📈 OVERALL PROFIT DISTRIBUTION:")
            if winning_trades:
                print(f"Winning trades: {len(winning_trades)} (avg: {np.mean(winning_trades)*100:.2f}%)")
                print(f"Best trade: {max(winning_trades)*100:.2f}%")
            if losing_trades:
                print(f"Losing trades: {len(losing_trades)} (avg: {np.mean(losing_trades)*100:.2f}%)")
                print(f"Worst trade: {min(losing_trades)*100:.2f}%")
        
        return True
        
    except Exception as e:
        print(f"\n❌ Error during testing: {e}")
        import traceback
        traceback.print_exc()
        return False


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