#!/usr/bin/env python3
"""
Checkpoint Comparison Analysis
Analyzes backtest results for all DQN and PPO checkpoints
"""

import json
import sys
from pathlib import Path
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend
import pandas as pd
import numpy as np

def load_results(results_file):
    """Load backtest results from JSON"""
    with open(results_file, 'r') as f:
        return json.load(f)

def filter_valid_results(results):
    """Filter out checkpoints with no trades or invalid metrics"""
    return [r for r in results if r['total_trades'] > 0]

def create_comparison_table(results):
    """Create markdown comparison table"""
    df = pd.DataFrame(results)

    # Separate DQN and PPO
    dqn_df = df[df['model_type'] == 'DQN'].copy()
    ppo_df = df[df['model_type'] == 'PPO'].copy()

    # Sort by Sharpe ratio
    dqn_df_sorted = dqn_df.sort_values('sharpe_ratio', ascending=False).head(10)
    ppo_df_sorted = ppo_df.sort_values('sharpe_ratio', ascending=False).head(10)

    print("\n" + "="*90)
    print("🔵 TOP 10 DQN CHECKPOINTS (Ranked by Sharpe Ratio)")
    print("="*90)
    print(f"{'Epoch':<8} {'Trades':<8} {'Win Rate':<12} {'Sharpe':<10} {'PnL':<12} {'Drawdown':<12} {'Trade Freq':<12}")
    print("-"*90)

    for _, row in dqn_df_sorted.iterrows():
        print(f"{row['epoch']:<8} {row['total_trades']:<8} {row['win_rate']:>10.1f}% {row['sharpe_ratio']:>10.3f} ${row['total_pnl']:>10.2f} {row['max_drawdown']:>11.2%} {row['trade_frequency']:>12.1f}")

    print("\n" + "="*90)
    print("🟢 TOP 10 PPO CHECKPOINTS (Ranked by Sharpe Ratio)")
    print("="*90)
    print(f"{'Epoch':<8} {'Trades':<8} {'Win Rate':<12} {'Sharpe':<10} {'PnL':<12} {'Drawdown':<12} {'Trade Freq':<12}")
    print("-"*90)

    for _, row in ppo_df_sorted.iterrows():
        print(f"{row['epoch']:<8} {row['total_trades']:<8} {row['win_rate']:>10.1f}% {row['sharpe_ratio']:>10.3f} ${row['total_pnl']:>10.2f} {row['max_drawdown']:>11.2%} {row['trade_frequency']:>12.1f}")

    return dqn_df, ppo_df

def plot_epoch_vs_metrics(dqn_df, ppo_df, output_dir):
    """Create plots comparing epoch vs various metrics"""

    # Filter valid results (with trades)
    dqn_valid = dqn_df[dqn_df['total_trades'] > 0]
    ppo_valid = ppo_df[ppo_df['total_trades'] > 0]

    # Create figure with 2x2 subplots
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    fig.suptitle('DQN vs PPO: Checkpoint Performance Analysis', fontsize=16, fontweight='bold')

    # 1. Sharpe Ratio vs Epoch
    ax1 = axes[0, 0]
    ax1.scatter(dqn_valid['epoch'], dqn_valid['sharpe_ratio'],
               alpha=0.6, s=100, c='blue', label='DQN', marker='o')
    ax1.scatter(ppo_valid['epoch'], ppo_valid['sharpe_ratio'],
               alpha=0.6, s=100, c='green', label='PPO', marker='s')
    ax1.axhline(y=0, color='red', linestyle='--', alpha=0.5, label='Break-even')
    ax1.set_xlabel('Epoch', fontsize=12)
    ax1.set_ylabel('Sharpe Ratio', fontsize=12)
    ax1.set_title('Sharpe Ratio vs Training Epoch', fontsize=14, fontweight='bold')
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # 2. Trade Count vs Epoch
    ax2 = axes[0, 1]
    ax2.scatter(dqn_valid['epoch'], dqn_valid['total_trades'],
               alpha=0.6, s=100, c='blue', label='DQN', marker='o')
    ax2.scatter(ppo_valid['epoch'], ppo_valid['total_trades'],
               alpha=0.6, s=100, c='green', label='PPO', marker='s')
    ax2.set_xlabel('Epoch', fontsize=12)
    ax2.set_ylabel('Total Trades', fontsize=12)
    ax2.set_title('Trade Count vs Training Epoch', fontsize=14, fontweight='bold')
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    # 3. Win Rate vs Epoch
    ax3 = axes[1, 0]
    ax3.scatter(dqn_valid['epoch'], dqn_valid['win_rate'],
               alpha=0.6, s=100, c='blue', label='DQN', marker='o')
    ax3.scatter(ppo_valid['epoch'], ppo_valid['win_rate'],
               alpha=0.6, s=100, c='green', label='PPO', marker='s')
    ax3.axhline(y=50, color='red', linestyle='--', alpha=0.5, label='50% Win Rate')
    ax3.set_xlabel('Epoch', fontsize=12)
    ax3.set_ylabel('Win Rate (%)', fontsize=12)
    ax3.set_title('Win Rate vs Training Epoch', fontsize=14, fontweight='bold')
    ax3.legend()
    ax3.grid(True, alpha=0.3)

    # 4. PnL vs Epoch
    ax4 = axes[1, 1]
    ax4.scatter(dqn_valid['epoch'], dqn_valid['total_pnl'],
               alpha=0.6, s=100, c='blue', label='DQN', marker='o')
    ax4.scatter(ppo_valid['epoch'], ppo_valid['total_pnl'],
               alpha=0.6, s=100, c='green', label='PPO', marker='s')
    ax4.axhline(y=0, color='red', linestyle='--', alpha=0.5, label='Break-even')
    ax4.set_xlabel('Epoch', fontsize=12)
    ax4.set_ylabel('Total PnL ($)', fontsize=12)
    ax4.set_title('Total PnL vs Training Epoch', fontsize=14, fontweight='bold')
    ax4.legend()
    ax4.grid(True, alpha=0.3)

    plt.tight_layout()
    output_file = output_dir / 'checkpoint_comparison_plots.png'
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"\n📊 Saved comparison plots to: {output_file}")
    plt.close()

def plot_sharpe_distribution(dqn_df, ppo_df, output_dir):
    """Create box plot comparing Sharpe ratio distributions"""

    dqn_valid = dqn_df[dqn_df['total_trades'] > 0]['sharpe_ratio']
    ppo_valid = ppo_df[ppo_df['total_trades'] > 0]['sharpe_ratio']

    fig, ax = plt.subplots(figsize=(10, 6))

    # Create box plots
    box_data = [dqn_valid, ppo_valid]
    bp = ax.boxplot(box_data, labels=['DQN', 'PPO'], patch_artist=True,
                    showmeans=True, meanline=True)

    # Customize colors
    colors = ['lightblue', 'lightgreen']
    for patch, color in zip(bp['boxes'], colors):
        patch.set_facecolor(color)

    ax.set_ylabel('Sharpe Ratio', fontsize=12)
    ax.set_title('Sharpe Ratio Distribution: DQN vs PPO', fontsize=14, fontweight='bold')
    ax.axhline(y=0, color='red', linestyle='--', alpha=0.5, label='Break-even')
    ax.grid(True, alpha=0.3)
    ax.legend()

    plt.tight_layout()
    output_file = output_dir / 'sharpe_distribution.png'
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"📊 Saved Sharpe distribution plot to: {output_file}")
    plt.close()

def create_statistical_summary(dqn_df, ppo_df):
    """Generate statistical summary"""

    dqn_valid = dqn_df[dqn_df['total_trades'] > 0]
    ppo_valid = ppo_df[ppo_df['total_trades'] > 0]

    print("\n" + "="*90)
    print("📈 STATISTICAL SUMMARY")
    print("="*90)

    print(f"\n{'Metric':<25} {'DQN':>20} {'PPO':>20} {'Winner':>20}")
    print("-"*90)

    metrics = [
        ('Checkpoints with trades', len(dqn_valid), len(ppo_valid)),
        ('Avg Sharpe Ratio', dqn_valid['sharpe_ratio'].mean(), ppo_valid['sharpe_ratio'].mean()),
        ('Max Sharpe Ratio', dqn_valid['sharpe_ratio'].max(), ppo_valid['sharpe_ratio'].max()),
        ('Avg Win Rate (%)', dqn_valid['win_rate'].mean(), ppo_valid['win_rate'].mean()),
        ('Avg Total Trades', dqn_valid['total_trades'].mean(), ppo_valid['total_trades'].mean()),
        ('Avg PnL ($)', dqn_valid['total_pnl'].mean(), ppo_valid['total_pnl'].mean()),
        ('Best PnL ($)', dqn_valid['total_pnl'].max(), ppo_valid['total_pnl'].max()),
    ]

    for name, dqn_val, ppo_val in metrics:
        if name == 'Checkpoints with trades':
            winner = 'DQN' if dqn_val > ppo_val else 'PPO' if ppo_val > dqn_val else 'Tie'
            print(f"{name:<25} {int(dqn_val):>20} {int(ppo_val):>20} {winner:>20}")
        else:
            winner = 'DQN' if dqn_val > ppo_val else 'PPO' if ppo_val > dqn_val else 'Tie'
            print(f"{name:<25} {dqn_val:>20.3f} {ppo_val:>20.3f} {winner:>20}")

    # Best overall checkpoints
    print("\n" + "="*90)
    print("🏆 BEST CHECKPOINTS")
    print("="*90)

    best_dqn = dqn_valid.loc[dqn_valid['sharpe_ratio'].idxmax()]
    best_ppo = ppo_valid.loc[ppo_valid['sharpe_ratio'].idxmax()]

    print(f"\nBest DQN: Epoch {best_dqn['epoch']}")
    print(f"  Sharpe: {best_dqn['sharpe_ratio']:.3f}")
    print(f"  Win Rate: {best_dqn['win_rate']:.1f}%")
    print(f"  Trades: {best_dqn['total_trades']}")
    print(f"  PnL: ${best_dqn['total_pnl']:.2f}")

    print(f"\nBest PPO: Epoch {best_ppo['epoch']}")
    print(f"  Sharpe: {best_ppo['sharpe_ratio']:.3f}")
    print(f"  Win Rate: {best_ppo['win_rate']:.1f}%")
    print(f"  Trades: {best_ppo['total_trades']}")
    print(f"  PnL: ${best_ppo['total_pnl']:.2f}")

def create_markdown_report(dqn_df, ppo_df, output_dir):
    """Create comprehensive markdown report"""

    dqn_valid = dqn_df[dqn_df['total_trades'] > 0]
    ppo_valid = ppo_df[ppo_df['total_trades'] > 0]

    # Get top 10 from each
    dqn_top10 = dqn_valid.sort_values('sharpe_ratio', ascending=False).head(10)
    ppo_top10 = ppo_valid.sort_values('sharpe_ratio', ascending=False).head(10)

    report = []
    report.append("# Checkpoint Backtesting Results")
    report.append(f"\n**Date**: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}")
    report.append(f"**Total Checkpoints Tested**: {len(dqn_df) + len(ppo_df)}")
    report.append(f"**Checkpoints with Valid Trades**: {len(dqn_valid) + len(ppo_valid)}")
    report.append("\n---\n")

    # Executive Summary
    report.append("## Executive Summary")
    report.append(f"\n### DQN Performance")
    report.append(f"- **Best Checkpoint**: Epoch {dqn_valid['sharpe_ratio'].idxmax()}")
    report.append(f"- **Best Sharpe Ratio**: {dqn_valid['sharpe_ratio'].max():.3f}")
    report.append(f"- **Average Sharpe Ratio**: {dqn_valid['sharpe_ratio'].mean():.3f}")
    report.append(f"- **Best PnL**: ${dqn_valid['total_pnl'].max():.2f}")

    report.append(f"\n### PPO Performance")
    report.append(f"- **Best Checkpoint**: Epoch {ppo_valid['sharpe_ratio'].idxmax()}")
    report.append(f"- **Best Sharpe Ratio**: {ppo_valid['sharpe_ratio'].max():.3f}")
    report.append(f"- **Average Sharpe Ratio**: {ppo_valid['sharpe_ratio'].mean():.3f}")
    report.append(f"- **Best PnL**: ${ppo_valid['total_pnl'].max():.2f}")

    # Top 10 DQN
    report.append("\n---\n")
    report.append("## Top 10 DQN Checkpoints")
    report.append("\n| Rank | Epoch | Sharpe | Win Rate | Trades | PnL | Drawdown | Trade Freq |")
    report.append("|------|-------|--------|----------|--------|-----|----------|------------|")

    for rank, (_, row) in enumerate(dqn_top10.iterrows(), 1):
        report.append(f"| {rank} | {row['epoch']} | {row['sharpe_ratio']:.3f} | {row['win_rate']:.1f}% | {row['total_trades']} | ${row['total_pnl']:.2f} | {row['max_drawdown']:.2%} | {row['trade_frequency']:.1f} |")

    # Top 10 PPO
    report.append("\n---\n")
    report.append("## Top 10 PPO Checkpoints")
    report.append("\n| Rank | Epoch | Sharpe | Win Rate | Trades | PnL | Drawdown | Trade Freq |")
    report.append("|------|-------|--------|----------|--------|-----|----------|------------|")

    for rank, (_, row) in enumerate(ppo_top10.iterrows(), 1):
        report.append(f"| {rank} | {row['epoch']} | {row['sharpe_ratio']:.3f} | {row['win_rate']:.1f}% | {row['total_trades']} | ${row['total_pnl']:.2f} | {row['max_drawdown']:.2%} | {row['trade_frequency']:.1f} |")

    # Key Insights
    report.append("\n---\n")
    report.append("## Key Insights")

    # Insight 1: Early vs Late epochs
    dqn_early = dqn_valid[dqn_valid['epoch'] <= 200]
    dqn_late = dqn_valid[dqn_valid['epoch'] > 200]

    report.append("\n### Training Phase Analysis")
    report.append(f"\n**DQN Early Epochs (≤200)**:")
    report.append(f"- Average Sharpe: {dqn_early['sharpe_ratio'].mean():.3f}")
    report.append(f"- Average Trades: {dqn_early['total_trades'].mean():.1f}")
    report.append(f"- Average Win Rate: {dqn_early['win_rate'].mean():.1f}%")

    report.append(f"\n**DQN Late Epochs (>200)**:")
    report.append(f"- Average Sharpe: {dqn_late['sharpe_ratio'].mean():.3f}")
    report.append(f"- Average Trades: {dqn_late['total_trades'].mean():.1f}")
    report.append(f"- Average Win Rate: {dqn_late['win_rate'].mean():.1f}%")

    ppo_early = ppo_valid[ppo_valid['epoch'] <= 200]
    ppo_late = ppo_valid[ppo_valid['epoch'] > 200]

    report.append(f"\n**PPO Early Epochs (≤200)**:")
    report.append(f"- Average Sharpe: {ppo_early['sharpe_ratio'].mean():.3f}")
    report.append(f"- Average Trades: {ppo_early['total_trades'].mean():.1f}")
    report.append(f"- Average Win Rate: {ppo_early['win_rate'].mean():.1f}%")

    report.append(f"\n**PPO Late Epochs (>200)**:")
    report.append(f"- Average Sharpe: {ppo_late['sharpe_ratio'].mean():.3f}")
    report.append(f"- Average Trades: {ppo_late['total_trades'].mean():.1f}")
    report.append(f"- Average Win Rate: {ppo_late['win_rate'].mean():.1f}%")

    # Save report
    report_file = output_dir / 'CHECKPOINT_BACKTEST_REPORT.md'
    with open(report_file, 'w') as f:
        f.write('\n'.join(report))

    print(f"\n📄 Saved markdown report to: {report_file}")

def main():
    if len(sys.argv) < 2:
        print("Usage: python compare_checkpoints.py <results_json_file>")
        sys.exit(1)

    results_file = Path(sys.argv[1])
    if not results_file.exists():
        print(f"Error: Results file not found: {results_file}")
        sys.exit(1)

    # Create output directory
    output_dir = Path(__file__).parent.parent / 'results'
    output_dir.mkdir(exist_ok=True)

    # Load results
    print(f"📖 Loading results from: {results_file}")
    results = load_results(results_file)

    # Filter valid results
    valid_results = filter_valid_results(results)
    print(f"✅ Found {len(valid_results)} checkpoints with valid trades")

    # Create comparison table
    dqn_df, ppo_df = create_comparison_table(valid_results)

    # Statistical summary
    create_statistical_summary(dqn_df, ppo_df)

    # Create plots
    plot_epoch_vs_metrics(dqn_df, ppo_df, output_dir)
    plot_sharpe_distribution(dqn_df, ppo_df, output_dir)

    # Create markdown report
    create_markdown_report(dqn_df, ppo_df, output_dir)

    print("\n✅ Analysis complete!")

if __name__ == '__main__':
    main()