import json
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

J = {
 "pref": json.load(open("/mnt/localssd/pref_results2.json")),
 "iqa":  json.load(open("/mnt/localssd/pref_results_iqa.json")),
 "panel":json.load(open("/mnt/localssd/pref_results_panel.json")),
 "aes":  json.load(open("/mnt/localssd/pref_results_aes.json")),
 "qal":  json.load(open("/mnt/localssd/pref_results_qalign.json")),
}
# (display, src, kind, key, type, lower_better)
REG = [
 ("ImageReward","pref","models","ImageReward-v1.0","prompt-pref",False),
 ("PickScore","pref","models","PickScore_v1","prompt-pref",False),
 ("NIQE","panel","metrics","niqe","unsupervised-IQA",True),
 ("BRISQUE","panel","metrics","brisque","unsupervised-IQA",True),
 ("PI","panel","metrics","pi","unsupervised-IQA",True),
 ("IL-NIQE","panel","metrics","ilniqe","unsupervised-IQA",True),
 ("CLIP-IQA","iqa","metrics","clipiqa","learned-IQA",False),
 ("CLIP-IQA+","iqa","metrics","clipiqa+","learned-IQA",False),
 ("NIMA(AVA)","panel","metrics","nima","learned-IQA",False),
 ("PaQ-2-PiQ","panel","metrics","paq2piq","learned-IQA",False),
 ("HyperIQA","panel","metrics","hyperiqa","learned-IQA",False),
 ("DBCNN","panel","metrics","dbcnn","learned-IQA",False),
 ("TReS","panel","metrics","tres","learned-IQA",False),
 ("ARNIQA","panel","metrics","arniqa","learned-IQA",False),
 ("QualiCLIP","panel","metrics","qualiclip","learned-IQA",False),
 ("MUSIQ","panel","metrics","musiq","learned-IQA",False),
 ("MANIQA","panel","metrics","maniqa","learned-IQA",False),
 ("TOPIQ-NR","panel","metrics","topiq_nr","learned-IQA",False),
 ("Q-Align-qual","qal","metrics","q_align_quality","learned-IQA",False),
 ("LAION-aes","iqa","metrics","laion_aes","aesthetic",False),
 ("Aes-v2.5","aes","metrics","aes_v2.5","aesthetic",False),
 ("Q-Align-aes","qal","metrics","q_align_aesthetic","aesthetic",False),
]
def summ(src,kind,key,dom):
    return J[src]["domains"][dom][kind][key]["summary"]

rows=[]
for disp,src,kind,key,typ,lb in REG:
    ps=summ(src,kind,key,"pixel"); ls=summ(src,kind,key,"latent")
    rows.append(dict(disp=disp,typ=typ,lb=lb,
        pix_wr=ps["win_rate"],pix_p=ps["binom_p_two_sided"],
        lat_wr=ls["win_rate"],lat_p=ls["binom_p_two_sided"]))

# ---- text table ----
def stars(p): return "***" if p<1e-3 else "**" if p<1e-2 else "*" if p<5e-2 else "ns"
print("="*104)
print(f"{'scorer':<14}{'type':<18}{'LB':<4}{'PIXEL win%':>11} {'p':>9} {'':>3}   {'LATENT win%':>12} {'p':>9} {'':>3}")
print("-"*104)
for r in rows:
    print(f"{r['disp']:<14}{r['typ']:<18}{('Y' if r['lb'] else ''):<4}"
          f"{r['pix_wr']:>10.1f} {r['pix_p']:>9.1e} {stars(r['pix_p']):>3}   "
          f"{r['lat_wr']:>11.1f} {r['lat_p']:>9.1e} {stars(r['lat_p']):>3}")
print("="*104)

# ---- summary counts ----
def classify(wr,p):
    if p>=0.05: return "tie"
    return "gan_better" if wr>50 else "gan_worse"
for dom,wk,pk in [("PIXEL","pix_wr","pix_p"),("LATENT","lat_wr","lat_p")]:
    c={"gan_better":0,"tie":0,"gan_worse":0}
    for r in rows: c[classify(r[wk],r[pk])]+=1
    print(f"{dom}: of {len(rows)} scorers -> GAN significantly BETTER={c['gan_better']}, "
          f"TIE(ns)={c['tie']}, GAN significantly WORSE={c['gan_worse']}")

json.dump({"rows":rows}, open("/mnt/localssd/pref_master_table.json","w"), indent=2)

# ---- figure: horizontal grouped bars ----
TCOL={"prompt-pref":"#4a3aa7","unsupervised-IQA":"#008300","learned-IQA":"#256abf","aesthetic":"#eb6834"}
BLUE="#2a78d6"; ORANGE="#eb6834"; INK="#0b0b0b"; INK2="#52514e"; SURF="#fcfcfb"; GRID="#e6e5e2"
plt.rcParams.update({"figure.facecolor":SURF,"axes.facecolor":SURF,"text.color":INK,
  "axes.labelcolor":INK,"axes.edgecolor":INK2,"xtick.color":INK2,"ytick.color":INK2,
  "font.size":10,"axes.spines.top":False,"axes.spines.right":False})
# order rows grouped by type (already grouped in REG order); build y positions with gaps between groups
order=rows
ypos=[]; y=0; labels=[]; sep=[]; lasttyp=None
for r in order:
    if lasttyp is not None and r['typ']!=lasttyp:
        y-=0.7; sep.append(y+0.35)
    ypos.append(y); labels.append(r['disp']+("  (LB)" if r['lb'] else "")); y-=1; lasttyp=r['typ']
ypos=np.array(ypos); h=0.38
fig,ax=plt.subplots(figsize=(11,12))
for r,yp in zip(order,ypos):
    ax.barh(yp+h/1.6, r['pix_wr'], height=h, color=BLUE, zorder=3)
    ax.barh(yp-h/1.6, r['lat_wr'], height=h, color=ORANGE, zorder=3)
    ax.text(r['pix_wr']+0.8, yp+h/1.6, f"{r['pix_wr']:.0f}{'' if r['pix_p']<0.05 else ' ns'}",
            va="center", ha="left", fontsize=8, color=INK,
            fontweight="bold" if r['pix_p']<0.05 else "normal")
    ax.text(r['lat_wr']+0.8, yp-h/1.6, f"{r['lat_wr']:.0f}{'' if r['lat_p']<0.05 else ' ns'}",
            va="center", ha="left", fontsize=8, color=INK,
            fontweight="bold" if r['lat_p']<0.05 else "normal")
ax.axvline(50, ls="--", lw=1.6, color=INK2, zorder=2)
ax.text(50,ypos.max()+1.2,"50% = tie",ha="center",va="bottom",color=INK2,fontsize=9)
ax.set_yticks(ypos); ax.set_yticklabels(labels, fontsize=9)
# type group labels on right
seen=set()
for r,yp in zip(order,ypos):
    if r['typ'] not in seen:
        seen.add(r['typ'])
        ax.text(103, yp+0.5, r['typ'], color=TCOL[r['typ']], fontsize=10, fontweight="bold", va="top")
ax.set_xlim(0,108); ax.set_xlabel("GAN win-rate  (% of 1065 matched pairs where GAN is better, polarity-correct)")
ax.set_ylim(ypos.min()-1, ypos.max()+2)
from matplotlib.patches import Patch
ax.legend(handles=[Patch(color=BLUE,label="PIXEL (DeCo baseline vs gan_dinov2_text@25k)"),
                   Patch(color=ORANGE,label="LATENT (SANA L1 vs L2+GAN)")],
          frameon=False, loc="lower right", fontsize=9)
ax.set_title("GAN vs no-GAN: image-only quality/aesthetic + preference battery (22 scorer-heads, n=1065/domain)\n"
             "bars right of 50% = GAN preferred; 'ns' = not significant (binomial p>=0.05); LB = lower-better metric (polarity-flipped)",
             fontsize=11, fontweight="bold")
ax.xaxis.grid(True,color=GRID,zorder=0); ax.set_axisbelow(True)
fig.tight_layout()
fig.savefig("/mnt/localssd/pref_iqa_panel.png", dpi=150, facecolor=SURF, bbox_inches="tight")
print("saved /mnt/localssd/pref_iqa_panel.png")