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

# Danh sách các folder bạn muốn so sánh
# (Tên folder phải khớp với cái bạn đã chạy)
experiments = [
    {"path": "ablat_best", "label": "AWWL (Best: a0.8, p2.0)", "color": "red"},
    {"path": "ablat_static", "label": "Static Loss (p0.0)", "color": "gray"},
    {"path": "ablat_alpha05", "label": "Balanced (a0.5)", "color": "blue"},
    {"path": "ablat_power1", "label": "Linear Decay (p1.0)", "color": "orange"},
]

# Hàm làm mượt (Smoothing) vì loss nhảy lung tung khó nhìn
def smooth(scalars, weight=0.95):  
    last = scalars[0]  
    smoothed = list()  
    for point in scalars:  
        smoothed_val = last * weight + (1 - weight) * point  
        smoothed.append(smoothed_val)  
        last = smoothed_val  
    return smoothed

plt.figure(figsize=(10, 6))

for exp in experiments:
    json_path = os.path.join(exp["path"], "loss_history.json")
    
    if os.path.exists(json_path):
        with open(json_path, "r") as f:
            data = json.load(f)
            losses = data["losses"]
            
            # Chỉ vẽ nếu có dữ liệu
            if len(losses) > 0:
                # Làm mượt loss để biểu đồ đẹp như trong báo
                smoothed_losses = smooth(losses, weight=0.99) # 0.99 là rất mượt
                
                # Tạo trục X (Steps)
                steps = range(len(smoothed_losses))
                
                plt.plot(steps, smoothed_losses, label=exp["label"], color=exp["color"], linewidth=1.5)
                print(f"Loaded {exp['label']}: {len(losses)} steps")
    else:
        print(f"Warning: File not found {json_path}")

plt.title("Training Loss Convergence Comparison")
plt.xlabel("Training Steps")
plt.ylabel("Loss (Smoothed)")
plt.legend()
plt.grid(True, alpha=0.3)
plt.yscale("log") # Dùng scale log nếu chênh lệch lớn, hoặc bỏ dòng này nếu muốn linear
plt.tight_layout()

# Lưu biểu đồ ra ảnh
plt.savefig("loss_comparison_chart.png", dpi=300)
print("✅ Chart saved to loss_comparison_chart.png")
plt.show()