import matplotlib.pyplot as plt
import numpy as np
import os

# Sử dụng style chuyên nghiệp
plt.style.use('ggplot')

# 1. Định nghĩa dữ liệu mở rộng tối đa (SOTA + Internal Variants)
models = [
    'iTracker', 'GazeNet', 'FullFace', 'RT-Gene', 
    'AFF-Net', 'MobileNetV3', 'L2CS-Net (Teacher)', 
    'FAZE (Few-shot)', 'ETH-XGaze (R18)',
    'LIPE V16 (Baseline)', 'LIPE V2 GOLD (Ours)', 'LIPE V6 (Stabilizer)'
]

# Trục Y: Sai số (MAE độ)
mae = [6.8, 6.2, 5.5, 5.3, 4.8, 6.0, 4.5, 4.0, 4.4, 6.2, 5.2, 5.0]             

# Trục X: Complexity (GFLOPs)
gflops = [3.20, 1.20, 2.10, 1.60, 0.15, 0.05, 1.80, 0.45, 1.80, 0.0019, 0.00234, 0.021] 

# Kích thước bong bóng: Tham số (Millions)
params = [35.0, 11.5, 50.0, 15.0, 4.5, 1.5, 25.0, 10.0, 11.7, 0.12, 0.18, 0.51]       

# Quy đổi kích thước bong bóng
bubble_sizes = [p * 30 + 100 for p in params] 

# 2. Khởi tạo đồ thị
plt.figure(figsize=(14, 9), dpi=300)

# Bảng màu mở rộng
colors = [
    '#7f8c8d', '#34495e', '#2c3e50', '#2980b9', 
    '#27ae60', '#a3cb38', '#c0392b', '#8e44ad', '#16a085',
    '#edae49', '#f1c40f', '#e67e22'
] 

# 3. Vẽ đồ thị bong bóng
scatter = plt.scatter(gflops, mae, s=bubble_sizes, c=colors, alpha=0.8, edgecolors='black', linewidth=1.0)

# 4. Cấu hình trục Logarit cho Complexity
plt.xscale('log')

# 5. Thêm nhãn và Annotations (Tinh chỉnh offset để né nhau)
for i, txt in enumerate(models):
    offset = (12, 5)
    if 'GOLD' in txt: offset = (15, -12)
    elif 'Stabilizer' in txt: offset = (15, 5)
    elif 'Baseline' in txt: offset = (-130, -5)
    elif 'MobileNet' in txt: offset = (-120, 5)
    elif 'AFF-Net' in txt: offset = (-65, -15)
    elif 'FAZE' in txt: offset = (15, 5)
    elif 'L2CS' in txt: offset = (15, -10)
    elif 'ETH' in txt: offset = (-110, 10)
    elif 'FullFace' in txt: offset = (-70, 10)
    
    plt.annotate(txt, (gflops[i], mae[i]), xytext=offset, textcoords='offset points', 
                 fontsize=9, fontweight='bold' if 'GOLD' in txt else 'normal',
                 bbox=dict(boxstyle='round,pad=0.2', fc='white', alpha=0.7, ec='none'))

# 6. Mũi tên định hướng tối ưu
plt.annotate('TinyML Pareto Front', xy=(0.0015, 4.8), xytext=(0.05, 4.2),
             arrowprops=dict(facecolor='#e74c3c', shrink=0.05, width=1.2, headwidth=7),
             fontsize=10, fontweight='bold', color='#e74c3c')

# 7. Tiêu đề và nhãn
plt.xlabel(r'$\leftarrow$ Computational Complexity (GFLOPs in Log Scale)', fontsize=12, fontweight='bold')
plt.ylabel(r'$\leftarrow$ Estimation Error (MAE in Degree)', fontsize=12, fontweight='bold')
plt.title('Comprehensive SOTA Comparison: Accuracy vs. Efficiency', fontsize=15, fontweight='bold', pad=25)

plt.xlim(0.0004, 10.0)
plt.ylim(3.5, 7.5)

# 8. Legend cho tham số
for p in [0.18, 1.5, 11.7, 50.0]:
    label = f"{p}M (GOLD)" if p == 0.18 else f"{p}M"
    plt.scatter([], [], c='gray', alpha=0.5, s=p*30+100, label=label, edgecolors='black')
plt.legend(title="Model Footprint (#Params)", loc="upper right", frameon=True, fontsize=9)

plt.grid(True, which="both", ls="-", alpha=0.3)
plt.tight_layout()

# Lưu file ảnh
output_path = 'comprehensive_sota_tradeoff.png'
plt.savefig(output_path, bbox_inches='tight')
print(f"Biểu đồ SOTA toàn diện đã được lưu tại: {os.path.abspath(output_path)}")