import matplotlib.pyplot as plt
import numpy as np

# Benchmarks included: DocVQA, ChartQA, OCRBench, InfoVQA, TextVQA, V*, POPE
models = ['Baseline-576', '+Ours (Full)', '+Ours-Gate', 'Baseline-1024']
colors = ['#95a5a6', '#e74c3c', '#2ecc71', '#3498db']  # Grey, Red, Green, Blue
markers = ['o', '^', '*', 's']

# Data Structure: {Benchmark: (Score_List, Time_List)}
# Order matches 'models' list
data = {
    'DocVQA': ([92.0, 93.6, 93.4, 93.9], [586, 806, 722, 799]),
    'ChartQA': ([82.6, 85.6, 83.9, 84.4], [227, 396, 266, 355]),
    'OCRBench': ([81.5, 84.1, 83.5, 83.1], [125, 193, 155, 172]),
    'InfoVQA': ([70.2, 78.7, 77.2, 76.6], [251, 454, 355, 375]),
    'TextVQA': ([81.1, 83.1, 82.7, 82.8], [436, 806, 506, 678]),
    'V*': ([64.4, 76.4, 77.5, 70.2], [15, 29, 26, 24]),
    'POPE': ([86.7, 88.6, 87.4, 86.7], [675, 1299, 788, 1115])
}

# Lists to store aggregated data
avg_abs_scores = []
avg_rel_times = []

# Calculate Averages for each model
num_models = len(models)
for i in range(num_models):
    total_score = 0
    total_rel_time = 0
    count = 0

    for bench, (scores, times) in data.items():
        # Sum absolute scores
        total_score += scores[i]

        # Sum relative time (Current Time / Baseline-576 Time)
        base_time = times[0]
        total_rel_time += times[i] / base_time

        count += 1

    avg_abs_scores.append(total_score / count)
    avg_rel_times.append(total_rel_time / count)

# Create Plot
plt.figure(figsize=(11, 8))
plt.grid(True, linestyle='--', alpha=0.6)

# Plot Points
for i, model in enumerate(models):
    plt.scatter(avg_rel_times[i], avg_abs_scores[i],
                color=colors[i], s=250, label=model, marker=markers[i], zorder=10, edgecolors='k')

    # Dynamic annotation placement to avoid overlap
    xytext_offset = (0, 15)
    if 'Baseline-1024' in model:
        xytext_offset = (0, -25)  # Move label down

    plt.annotate(f"{model}\n({avg_rel_times[i]:.2f}x Time, {avg_abs_scores[i]:.1f}% Acc)",
                 (avg_rel_times[i], avg_abs_scores[i]),
                 xytext=xytext_offset, textcoords='offset points', ha='center',
                 fontsize=10, fontweight='bold',
                 bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.7, ec="none"))

# Axis Limits (optional, for aesthetics)
plt.xlim(0.9, 1.8)
plt.ylim(78, 86)  # Adjust based on data range to zoom in on differences

# Axis Labels
plt.xlabel('Relative Time Cost (Normalized to Baseline-576)\n← Faster | Slower →', fontsize=12, fontweight='bold')
plt.ylabel('Average Absolute Accuracy (%)\nHigher is Better ↑', fontsize=12, fontweight='bold')
plt.title('Performance vs. Efficiency Trade-off\n(Average across 7 Benchmarks)', fontsize=14, pad=15)

# Reference Lines
# Vertical line at Baseline-576 time (1.0)
plt.axvline(x=1.0, color='#95a5a6', linestyle=':', linewidth=2, label='Baseline Speed')

# Legend
plt.legend(loc='lower right', fontsize=11, frameon=True, framealpha=0.9, shadow=True)

plt.tight_layout()
plt.show()
# plt.savefig('trade_off_plot_abs.png')
# print("Plot saved as 'trade_off_plot_abs.png'")