import matplotlib.pyplot as plt
import numpy as np

# ==========================================
# 1. CẤU HÌNH DỮ LIỆU (THEO KẾT QUẢ CỦA BẠN)
# ==========================================

# --- Baseline MSE ---
mse_fid = 16.6858
mse_is = 7.7798

# --- Set 1: Tác động của Alpha (Dùng nhóm p=2.0 để thấy rõ xu hướng xấu đi) ---
# Dữ liệu từ bảng: a0.5 p2.0, a0.8 p2.0, a0.95 p2.0
# (Nhóm này cho thấy rõ ràng: Tăng alpha -> FID tăng vọt)
alphas = [0.5, 0.8, 0.95]
fid_alpha = [18.70, 22.65, 27.95]
is_alpha = [7.72, 7.40, 6.99]

# --- Set 2: Tác động của Power (Tại nhóm tốt nhất a=0.2) ---
# Dữ liệu từ bảng: a0.2 p0.0, a0.2 p0.2, a0.2 p0.5, a0.2 p1.0
power_labels = ['Static\n(p=0.0)', 'Gentle\n(p=0.2)', 'Balanced\n(p=0.5)', 'Linear\n(p=1.0)']
# Lưu ý: Sắp xếp theo thứ tự p tăng dần
fid_power = [16.5495, 16.6813, 16.6058, 16.6214] 
is_power = [7.8111, 7.8560, 7.8218, 7.9472] # p=1.0 IS cao vọt

# ==========================================
# 2. CẤU HÌNH VẼ BIỂU ĐỒ
# ==========================================
plt.rcParams.update({
    "font.family": "serif",
    "font.size": 12,
    "axes.labelsize": 12,
    "axes.titlesize": 14,
    "xtick.labelsize": 10,
    "ytick.labelsize": 10,
    "legend.fontsize": 9,
    "figure.figsize": (14, 6) # Rộng hơn chút cho thoáng
})

fig, (ax1, ax2) = plt.subplots(1, 2)

# -------------------------------------------------------
# BIỂU ĐỒ 1: Impact of Alpha (Line Plot)
# -------------------------------------------------------
color_fid = 'tab:red'
color_is = 'tab:blue'

ax1.set_title(r"(a) Impact of Weight Balance ($\alpha$ with $p=2.0$)")
ax1.set_xlabel(r"Alpha Value ($\alpha$)")
ax1.set_ylabel(r"FID Score ($\downarrow$)", color=color_fid, fontweight='bold')

# Vẽ FID
line1 = ax1.plot(alphas, fid_alpha, marker='o', linestyle='-', color=color_fid, linewidth=2, label='FID')
ax1.tick_params(axis='y', labelcolor=color_fid)
ax1.grid(True, linestyle='--', alpha=0.3)
ax1.set_xticks(alphas) # Chỉ hiện các mốc có dữ liệu

# Vẽ Baseline MSE (FID)
ax1.axhline(y=mse_fid, color='gray', linestyle='--', alpha=0.7, label='MSE Baseline')

# Trục phụ cho IS
ax1_twin = ax1.twinx()
ax1_twin.set_ylabel(r"Inception Score ($\uparrow$)", color=color_is, fontweight='bold')
line2 = ax1_twin.plot(alphas, is_alpha, marker='s', linestyle='--', color=color_is, linewidth=2, label='IS')
ax1_twin.tick_params(axis='y', labelcolor=color_is)

# Legend
lines = line1 + line2
labs = [l.get_label() for l in lines]
ax1.legend(lines, labs, loc='upper left') 
# Note: Đặt upper left vì đồ thị đi lên sang phải

# -------------------------------------------------------
# BIỂU ĐỒ 2: Impact of Power (Bar Chart) - Zoomed In
# -------------------------------------------------------
x = np.arange(len(power_labels))
width = 0.35

ax2.set_title(r"(b) Impact of Adaptive Schedule ($p$ with $\alpha=0.2$)")
ax2.set_ylabel("Score")

# Vẽ cột
rects1 = ax2.bar(x - width/2, fid_power, width, label=r'FID ($\downarrow$)', color='#ff9999', edgecolor='black', alpha=0.9)
rects2 = ax2.bar(x + width/2, is_power, width, label=r'IS ($\uparrow$)', color='#66b3ff', edgecolor='black', alpha=0.9)

# Đường Baseline MSE
ax2.axhline(y=mse_fid, color='red', linestyle='-', linewidth=1.5, label='MSE FID')
ax2.axhline(y=mse_is, color='blue', linestyle='-', linewidth=1.5, label='MSE IS')

# Setup trục
ax2.set_xticks(x)
ax2.set_xticklabels(power_labels)

# --- QUAN TRỌNG: ZOOM VÀO VÙNG DỮ LIỆU ---
# Vì số liệu rất sát nhau (16.5 - 16.7), ta phải cắt trục Y
# FID ~ 16.6, IS ~ 7.8 -> Set range từ 7 đến 17 là đẹp
ax2.set_ylim(7.0, 17.5) 

ax2.legend(loc='upper center', ncol=2, framealpha=0.9)
ax2.grid(True, axis='y', linestyle='--', alpha=0.3)

# Hàm gắn nhãn số liệu lên cột
def autolabel(rects, is_fid=False):
    for rect in rects:
        height = rect.get_height()
        # Nếu là FID thì in đậm nếu thấp hơn MSE
        weight = 'bold' if (is_fid and height < mse_fid) or (not is_fid and height > mse_is) else 'normal'
        ax2.annotate(f'{height:.2f}',
                    xy=(rect.get_x() + rect.get_width() / 2, height),
                    xytext=(0, 3),
                    textcoords="offset points",
                    ha='center', va='bottom', fontsize=9, fontweight=weight)

autolabel(rects1, is_fid=True)
autolabel(rects2, is_fid=False)

plt.tight_layout()
plt.savefig("ablation_study_plot_final.png", dpi=300)
print("✅ Saved ablation_study_plot_final.png")
plt.show()