import matplotlib.pyplot as plt
import numpy as np

# --- 1. Data Definition ---
# Benchmarks and their sample counts
# Note: TextVQA uses 1250 samples as per user specification (val set subset or specific config)
bench_keys = ['DocVQA', 'InfoVQA', 'TextVQA', 'V*', 'RW-Lite']
samples = {
    'DocVQA': 5349,
    'InfoVQA': 701,
    'TextVQA': 1250,
    'V*': 191,
    'RW-Lite': 960
}

# Raw Data: {Label: {'accs': [Doc, Info, Text, V*, RW], 'times': [Doc, Info, Text, V*, RW]}}
# Data extracted from provided experimental results table

raw_data_baseline = {
    'Base\nMax-256': {
        'accs': [81.2, 54.0, 75.5, 53.4, 27.6],
        'times': [1330, 567, 938, 33, 292]
    },
    'Base\nMax-576': {
        'accs': [92.0, 70.1, 81.1, 64.4, 35.4],
        'times': [1990, 904, 1567, 56, 532]
    },
    'Base\nMax-1024': {
        'accs': [93.9, 76.5, 82.8, 70.2, 37.5],
        'times': [3066, 1423, 2471, 94, 896]
    },
    'Base\nMax-2048': {
        'accs': [94.7, 80.4, 82.9, 72.3, 39.3],
        'times': [6053, 2047, 2600, 213, 1925]
    }
}

raw_data_ours = {
    'Gated\nMax-256': {
        'accs': [89.1, 69.6, 81.8, 67.5, 37.3],
        # Note: RW-Lite time estimated as 435s based on avg ~1.5x slowdown from baseline (292s)
        'times': [1919, 880, 1285, 53, 435]
    },
    'Gated\nMax-576': {
        'accs': [93.8, 78.6, 82.8, 80.1, 43.4],
        'times': [2430, 1277, 1878, 85, 813]
    },
    'Gated\nMax-1024': {
        'accs': [94.6, 80.8, 83.9, 82.2, 46.6],
        'times': [3481, 1801, 2677, 130, 1420]
    }
}


# --- 2. Data Processing ---
def calculate_metrics(raw_data):
    avg_speeds = []
    avg_accs = []
    labels = []

    for label, data in raw_data.items():
        accs = data['accs']
        times = data['times']

        # Calculate samples/second for each benchmark
        task_speeds = []
        for i, task in enumerate(bench_keys):
            n_samples = samples[task]
            t_seconds = times[i]
            if t_seconds > 0:
                task_speeds.append(n_samples / t_seconds)
            else:
                task_speeds.append(0)

        # Weighted Average Speed (Average of speeds)
        avg_speed = np.mean(task_speeds)

        # Average Accuracy
        avg_acc = np.mean(accs)

        avg_speeds.append(avg_speed)
        avg_accs.append(avg_acc)
        labels.append(label)

    return avg_speeds, avg_accs, labels


base_x, base_y, base_labels = calculate_metrics(raw_data_baseline)
ours_x, ours_y, ours_labels = calculate_metrics(raw_data_ours)

# --- 3. Plotting ---
plt.figure(figsize=(10, 7))

# Plot Baseline Curve
plt.plot(base_x, base_y, marker='o', linestyle='--', color='#7f8c8d', label='Baseline', linewidth=2, markersize=9,
         markerfacecolor='white', markeredgewidth=2)

# Plot Ours Curve
plt.plot(ours_x, ours_y, marker='*', linestyle='-', color='#e74c3c', label='Ours-Gated', linewidth=2, markersize=14,
         markerfacecolor='#e74c3c')

# Annotate Points (Baseline)
for x, y, label in zip(base_x, base_y, base_labels):
    offset = (0, -20)
    if '2048' in label:
        offset = (0, 10)
        ha_align = 'right'
    else:
        ha_align = 'center'

    plt.annotate(label, (x, y), xytext=offset, textcoords='offset points',
                 ha=ha_align, fontsize=9, color='#555', weight='normal',
                 bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="none", alpha=0.7))

# Annotate Points (Ours)
for x, y, label in zip(ours_x, ours_y, ours_labels):
    offset = (0, 10)
    if '1024' in label:
        offset = (10, 0)
        ha_align = 'left'
    else:
        ha_align = 'center'

    plt.annotate(label, (x, y), xytext=offset, textcoords='offset points',
                 ha=ha_align, fontsize=9, fontweight='bold', color='#c0392b',
                 bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="none", alpha=0.7))

# --- 4. Styling ---
plt.title('Accuracy vs. Efficiency Trade-off\n(Average across DocVQA, InfoVQA, TextVQA, V*, RW-Lite)', fontsize=14,
          pad=15)
plt.xlabel('Average Throughput (Samples / Second) \n← Slower | Faster →', fontsize=12, fontweight='bold')
plt.ylabel('Average Accuracy (%)', fontsize=12, fontweight='bold')
plt.grid(True, linestyle=':', alpha=0.6)
plt.legend(fontsize=11, loc='lower right')

# Add Arrows for interpretation
plt.annotate('', xy=(max(base_x) * 0.95, min(base_y) * 1.02), xytext=(max(base_x) * 0.85, min(base_y) * 1.02),
             arrowprops=dict(arrowstyle="->", color='green', lw=1.5))
plt.text(max(base_x) * 0.90, min(base_y) * 1.01, 'Higher Efficiency', color='green', ha='center', fontsize=9)

plt.annotate('', xy=(min(base_x) * 1.05, max(ours_y) * 0.98), xytext=(min(base_x) * 1.05, max(ours_y) * 0.90),
             arrowprops=dict(arrowstyle="->", color='green', lw=1.5))
plt.text(min(base_x) * 1.06, max(ours_y) * 0.94, 'Higher Accuracy', color='green', rotation=90, va='center', fontsize=9)

plt.tight_layout()
plt.show()
# plt.savefig('acc_efficiency_tradeoff_updated.png', dpi=300)
# print("Plot saved as 'acc_efficiency_tradeoff_updated.png'")