#!/usr/bin/env python3
"""
Visualize experiment results.
"""

import matplotlib.pyplot as plt
import pandas as pd

# Create results dataframe
results = [
    ['Actual trades (baseline)', 39, 0.3846, 8.73, 340.38, 319.93],
    ['No filter', 31, 0.3226, -90.58, -2807.89, 217.24],
    ['Time limit', 15, 0.2667, -116.75, -1751.18, 183.98],
    ['Momentum only', 6, 0.1667, -156.18, -937.08, 249.80]
]

df = pd.DataFrame(results, columns=['Experiment', 'Trades', 'WR', 'Avg PnL', 'Total PnL', 'Avg|PnL|'])

# Create visualizations
fig, axes = plt.subplots(2, 2, figsize=(16, 12))

# 1. Win Rate
ax1 = axes[0, 0]
ax1.barh(df['Experiment'], df['WR'] * 100, alpha=0.7, color='steelblue')
ax1.set_title('Win Rate Comparison', fontsize=14, fontweight='bold')
ax1.set_xlabel('Win Rate (%)')
ax1.set_xlim(0, 50)
ax1.grid(True, alpha=0.3, axis='x')
for i, wr in enumerate(df['WR']):
    ax1.text(wr * 100 + 0.5, i, f'{wr:.1%}', va='center', fontweight='bold')

# 2. Avg PnL
ax2 = axes[0, 1]
colors = ['green' if pnl > 0 else 'red' for pnl in df['Avg PnL']]
ax2.barh(df['Experiment'], df['Avg PnL'], alpha=0.7, color=colors)
ax2.axvline(x=0, color='black', linewidth=2)
ax2.set_title('Average PnL per Trade', fontsize=14, fontweight='bold')
ax2.set_xlabel('Avg PnL (RUB)')
ax2.grid(True, alpha=0.3, axis='x')
for i, pnl in enumerate(df['Avg PnL']):
    ax2.text(pnl + (0.5 if pnl > 0 else -0.5), i, f'{pnl:.2f}', va='center', fontweight='bold')

# 3. Trades vs Avg PnL
ax3 = axes[1, 0]
scatter = ax3.scatter(df['Trades'], df['Avg PnL'], s=df['Trades']*20, alpha=0.7)
for i, exp in enumerate(df['Experiment']):
    ax3.text(df['Trades'].iloc[i], df['Avg PnL'].iloc[i] + 5, exp, ha='center')
ax3.axhline(y=0, color='black', linewidth=2)
ax3.set_title('Trades vs Average PnL', fontsize=14, fontweight='bold')
ax3.set_xlabel('Number of Trades')
ax3.set_ylabel('Avg PnL (RUB)')
ax3.grid(True, alpha=0.3)

# 4. Total PnL
ax4 = axes[1, 1]
colors = ['green' if pnl > 0 else 'red' for pnl in df['Total PnL']]
ax4.barh(df['Experiment'], df['Total PnL'], alpha=0.7, color=colors)
ax4.axvline(x=0, color='black', linewidth=2)
ax4.set_title('Total PnL', fontsize=14, fontweight='bold')
ax4.set_xlabel('Total PnL (RUB)')
ax4.grid(True, alpha=0.3, axis='x')
for i, pnl in enumerate(df['Total PnL']):
    ax4.text(pnl + (50 if pnl > 0 else -50), i, f'{pnl:,.0f}', va='center', fontweight='bold')

plt.tight_layout()
plt.savefig('/home/ai/projects/AI_Strategy/experiments/visualizations.png', dpi=150)
print("✅ Saved: experiments/visualizations.png")