"""
Visualization script for Claude and Gemini cooperation results.
Generates plots for blog post showing cooperation comparison.
"""

import json
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path

# Load experiment results
def load_exp_data():
    """Load and extract data from exp1, exp2, exp5"""
    data = {}
    
    # EXP1 - Baseline self-play
    exp1_data = {}
    for model in ["claude-opus", "gemini-flash", "gpt-5.2"]:
        path = f"results/experiment1_baseline_{model}_run1_20rounds.json"
        try:
            with open(path) as f:
                result = json.load(f)
                exp1_data[model] = {
                    "coop": result['results']['mutual_coop'],
                    "score": result['results']['total_a']
                }
        except:
            pass
    data['exp1'] = exp1_data
    
    # EXP2 - Personalities
    exp2_data = {"claude-opus": {}, "gemini-flash": {}, "gpt-5.2": {}}
    for model in exp2_data.keys():
        for personality in ["good", "neutral", "evil"]:
            path = f"results/experiment2_personalities_{model}_{personality}_vs_{personality}_run1_20rounds.json"
            try:
                with open(path) as f:
                    result = json.load(f)
                    exp2_data[model][personality] = result['results']['mutual_coop']
            except:
                exp2_data[model][personality] = 0
    data['exp2'] = exp2_data
    
    # EXP5 - Cross model (self matchups)
    exp5_data = {}
    for model in ["claude-opus", "gemini-flash", "gpt-5.2"]:
        path = f"results/experiment5_cross_model_{model}_run1_20rounds.json"
        try:
            with open(path) as f:
                result = json.load(f)
                # Extract mutual cooperation
                exp5_data[model] = result['results']['mutual_coop']
        except:
            exp5_data[model] = 0
    data['exp5'] = exp5_data
    
    return data

def plot_baseline_comparison(data):
    """Plot 1: Baseline cooperation in self-play"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4))
    
    models = ["claude-opus", "gemini-flash"]
    exp1 = data['exp1']
    
    coops = [exp1.get(m, {}).get('coop', 0) for m in models]
    scores = [exp1.get(m, {}).get('score', 0) for m in models]
    
    colors = ['#1f77b4', '#ff7f0e']
    labels = ['Claude', 'Gemini']
    
    # Cooperation rates
    ax1.bar(labels, coops, color=colors, alpha=0.7, edgecolor='black', linewidth=2, width=0.4)
    ax1.set_ylabel('Mutual Cooperation Rate', fontsize=14, fontweight='bold')
    ax1.set_title('Baseline: Self-Play Cooperation (20 rounds)', fontsize=14, fontweight='bold')
    ax1.set_ylim(0, 1.1)
    ax1.tick_params(axis='both', labelsize=13)
    for i, (label, coop) in enumerate(zip(labels, coops)):
        ax1.text(i, coop + 0.05, f'{coop:.0%}', ha='center', fontweight='bold', fontsize=14)
    ax1.grid(axis='y', alpha=0.3)
    
    # Scores
    ax2.bar(labels, scores, color=colors, alpha=0.7, edgecolor='black', linewidth=2, width=0.4)
    ax2.set_ylabel('Total Score (20 rounds)', fontsize=14, fontweight='bold')
    ax2.set_title('Baseline: Total Payoff', fontsize=14, fontweight='bold')
    ax2.tick_params(axis='both', labelsize=13)
    for i, (label, score) in enumerate(zip(labels, scores)):
        ax2.text(i, score + 2, f'{score}', ha='center', fontweight='bold', fontsize=14)
    ax2.grid(axis='y', alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('charts/exp1_baseline_comparison.png', dpi=300, bbox_inches='tight')
    print("✓ Saved: charts/exp1_baseline_comparison.png")
    plt.close()

def plot_personality_impact(data):
    """Plot 2: How personality affects cooperation (line graph)"""
    fig, ax = plt.subplots(figsize=(10, 5))
    
    personalities = ['good', 'neutral', 'evil']
    models = ["claude-opus", "gemini-flash"]
    model_labels = ['Claude', 'Gemini']
    colors = ['#1f77b4', '#ff7f0e']
    
    x = np.arange(len(personalities))
    exp2 = data['exp2']
    
    # Plot line for each model
    for i, (model, label, color) in enumerate(zip(models, model_labels, colors)):
        values = [exp2[model].get(p, 0) for p in personalities]
        ax.plot(x, values, marker='o', linewidth=3, markersize=10, 
                label=label, color=color, alpha=0.8)
        
        # Add value labels on points
        for j, val in enumerate(values):
            ax.text(j, val + 0.05, f'{val:.0%}', ha='center', fontweight='bold', fontsize=12)
    
    ax.set_ylabel('Mutual Cooperation Rate', fontsize=14, fontweight='bold')
    ax.set_title('Personality Impact on Cooperation', fontsize=14, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels([p.capitalize() for p in personalities], fontsize=13)
    ax.tick_params(axis='y', labelsize=13)
    ax.set_ylim(0, 1.15)
    ax.grid(True, alpha=0.3, linestyle='--')
    ax.legend(fontsize=13, loc='upper right', framealpha=0.95)
    
    plt.tight_layout()
    plt.savefig('charts/exp2_personality_impact.png', dpi=300, bbox_inches='tight')
    print("✓ Saved: charts/exp2_personality_impact.png")
    plt.close()

def main():
    """Generate all visualizations"""
    # Create charts directory if it doesn't exist
    Path('charts').mkdir(exist_ok=True)
    
    print("Loading experiment data...")
    data = load_exp_data()
    
    print("\nGenerating visualizations...")
    plot_baseline_comparison(data)
    plot_personality_impact(data)
    
    print("\n✅ All charts generated!")
    print("   Include these in blog post:")
    print("   - charts/exp1_baseline_comparison.png")
    print("   - charts/exp2_personality_impact.png")

if __name__ == "__main__":
    main()