import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import RegularPolygon
from matplotlib.path import Path
from matplotlib.projections.polar import PolarAxes
from matplotlib.projections import register_projection
from matplotlib.spines import Spine
from matplotlib.transforms import Affine2D

# ==========================================
# 1. DỮ LIỆU
# ==========================================
metrics = ['FID', 'IS', 'KID', 'Precision', 'Recall']
metric_types = ['lower', 'higher', 'lower', 'higher', 'higher']

data = {
    'MSE (Baseline)':   [16.6858, 7.7798, 0.007464, 0.7855, 0.6386],
    'AWWL (Ours)':      [16.6214, 7.9472, 0.007229, 0.7805, 0.6440], 
    'Huber':            [20.1373, 7.6236, 0.008887, 0.7840, 0.6133],
    'L1':               [19.7064, 7.6952, 0.008971, 0.7891, 0.6213],
    'Perceptual':       [20.0242, 7.7098, 0.008738, 0.7905, 0.6096],
    'SNR-Weighted':     [19.8316, 7.5798, 0.009146, 0.7866, 0.6153],
}

# --- BẢNG MÀU "ELEGANT ACADEMIC" ---
# AWWL: Đỏ Đô (Burgundy) - Sang trọng, nổi bật
# MSE: Đen Than (Charcoal) - Cơ bản, chắc chắn
# Others: Xám xanh/Xám tím (Slate/Muted tones) - Làm nền tinh tế
styles = {
    'AWWL (Ours)':    {'color': '#800000', 'ls': '-', 'lw': 2.2, 'marker': 'D', 'ms': 5, 'alpha': 1.0, 'zorder': 10}, 
    'MSE (Baseline)': {'color': '#2F4F4F', 'ls': '--', 'lw': 1.8, 'marker': 'o', 'ms': 4, 'alpha': 0.9, 'zorder': 9}, 
    'Huber':          {'color': '#708090', 'ls': ':', 'lw': 1.0, 'marker': '', 'ms': 0, 'alpha': 0.6, 'zorder': 5}, 
    'L1':             {'color': '#778899', 'ls': ':', 'lw': 1.0, 'marker': '', 'ms': 0, 'alpha': 0.6, 'zorder': 4}, 
    'Perceptual':     {'color': '#696969', 'ls': ':', 'lw': 1.0, 'marker': '', 'ms': 0, 'alpha': 0.6, 'zorder': 3}, 
    'SNR-Weighted':   {'color': '#A9A9A9', 'ls': ':', 'lw': 1.0, 'marker': '', 'ms': 0, 'alpha': 0.6, 'zorder': 2}, 
}

# ==========================================
# 2. XỬ LÝ DỮ LIỆU
# ==========================================
baseline_vals = data['MSE (Baseline)']

def normalize_data(raw_values):
    norm_vals = []
    for val, base, mtype in zip(raw_values, baseline_vals, metric_types):
        if mtype == 'higher':
            norm_vals.append(val / base) 
        else:
            norm_vals.append(base / val) 
    return norm_vals

# ==========================================
# 3. SETUP TRỤC RADAR
# ==========================================
def radar_factory(num_vars, frame='polygon'):
    theta = np.linspace(0, 2*np.pi, num_vars, endpoint=False)

    class RadarAxes(PolarAxes):
        name = 'radar'
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.set_theta_zero_location('N')

        def fill(self, *args, closed=True, **kwargs):
            return super().fill(closed=closed, *args, **kwargs)

        def plot(self, *args, **kwargs):
            lines = super().plot(*args, **kwargs)
            for line in lines:
                self._close_line(line)

        def _close_line(self, line):
            x, y = line.get_data()
            if x[0] != x[-1]:
                x = np.concatenate((x, [x[0]]))
                y = np.concatenate((y, [y[0]]))
                line.set_data(x, y)

        def set_varlabels(self, labels):
            self.set_thetagrids(np.degrees(theta), labels, fontsize=10, fontweight='bold', fontfamily='serif')

        def _gen_axes_patch(self):
            return RegularPolygon((0.5, 0.5), num_vars, radius=.5, edgecolor="k")

        def _gen_axes_spines(self):
            if frame == 'circle': return super()._gen_axes_spines()
            spine = Spine(axes=self, spine_type='circle', path=Path.unit_regular_polygon(num_vars))
            spine.set_transform(Affine2D().scale(.5).translate(.5, .5) + self.transAxes)
            return {'polar': spine}

    register_projection(RadarAxes)
    return theta

# ==========================================
# 4. VẼ BIỂU ĐỒ (TINH CHỈNH THẨM MỸ)
# ==========================================
# Dùng font Times New Roman (hoặc tương tự) cho cảm giác học thuật
plt.rcParams.update({
    'font.family': 'serif',
    'font.serif': ['Times New Roman', 'DejaVu Serif'],
    'font.size': 12,
    'text.usetex': False # Tắt latex engine nếu không cài, dùng matplotlib math mode
})

N = len(metrics)
theta = radar_factory(N, frame='polygon')

fig, ax = plt.subplots(figsize=(8, 7), subplot_kw=dict(projection='radar'))
fig.subplots_adjust(top=0.88, bottom=0.12)

# Grid lines mờ và mảnh hơn
ax.grid(color='#AAAAAA', linestyle=':', linewidth=0.5, alpha=0.7)

# Vẽ vòng tròn tham chiếu MSE = 1.0 (Nét liền mảnh, màu đen)
ax.plot(theta, [1.0]*N, color='black', linestyle='-', linewidth=0.6, alpha=0.4)

# Vẽ các đường dữ liệu
for label, raw_vals in data.items():
    norm_vals = normalize_data(raw_vals)
    s = styles[label]
    
    # Plot line
    ax.plot(theta, norm_vals, label=label, 
            color=s['color'], ls=s['ls'], lw=s['lw'], 
            marker=s['marker'], ms=s['ms'], alpha=s['alpha'], zorder=s['zorder'])
    
    # Tô màu cho AWWL: Màu đỏ nhạt, trong suốt
    if label == 'AWWL (Ours)':
        ax.fill(theta, norm_vals, color=s['color'], alpha=0.08)

# Cấu hình trục
ax.set_varlabels(metrics)

# Zoom range (0.8 -> 1.05)
ax.set_ylim(0.8, 1.04) 

# Nhãn trục radial (nhỏ, màu xám)
ax.set_rgrids([0.85, 0.90, 0.95, 1.0], labels=['0.85', '0.90', '0.95', '1.0'], angle=0, fontsize=8, color='#555555')

# Title (Đơn giản, Font serif đậm)
plt.title("Relative Performance Overview\n(Normalized to MSE Baseline = 1.0)", y=1.08, fontsize=13, fontweight='bold')

# Legend (Tinh tế hơn: Không khung viền, nằm ngang)
legend = ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.08),
          ncol=3, frameon=False, fontsize=10)

plt.tight_layout()
plt.savefig("radar_chart_elegant.png", dpi=300, bbox_inches='tight')
print("✅ Saved radar_chart_elegant.png")
plt.show()