#!/usr/bin/env python3
"""Fill the KLD table into README.md and render the two charts, from quant-summary.csv."""
import csv, os, sys
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

W = "/Users/kikocisneros/coco_ppl/aq_amalia"
PUB = f"{W}/publish"
CSV = f"{W}/metrics/quant-summary.csv"
os.makedirs(f"{PUB}/metrics", exist_ok=True)

rows = list(csv.DictReader(open(CSV)))
def f(x, d=None):
    try: return float(x)
    except: return d

# order big->small by size for the table
order = sorted(rows, key=lambda r: -f(r["size_gb"], 0))

# --- markdown KLD table ---
hdr = "| Model | Size GB | PPL | ΔPPL vs F16 | KLD mean | KLD p95 | Top-1 match |\n|---|---:|---:|---:|---:|---:|---:|"
lines = [hdr]
for r in order:
    name = r["name"]
    ppl = r["ppl"]; dppl = r["ppl_delta_vs_f16"]
    km = r["kld_mean"]; kp = r["kld_p95"]; t1 = r["top1_match_pct"]
    if name == "F16":
        lines.append(f"| **F16 (reference)** | {f(r['size_gb']):.1f} | {f(ppl):.2f} | 0.000 | 0.0000 | 0.0000 | 100.0% |")
    else:
        dv = f(dppl, 0); sign = "+" if dv >= 0 else ""
        lines.append(f"| {name} | {f(r['size_gb']):.2f} | {f(ppl):.2f} | {sign}{dv:.3f} | {f(km):.4f} | {f(kp):.4f} | {f(t1):.2f}% |")
table = "\n".join(lines)

readme = open(f"{PUB}/README.md").read()
readme = readme.replace("<!--KLD_TABLE_PLACEHOLDER-->", table)
open(f"{PUB}/README.md", "w").write(readme)
print("README table filled:\n" + table)

# --- chart 1: KLD mean vs size ---
q = [r for r in order if r["name"] != "F16"]
q = sorted(q, key=lambda r: f(r["size_gb"]))
sizes = [f(r["size_gb"]) for r in q]; klds = [f(r["kld_mean"]) for r in q]; names = [r["name"] for r in q]
plt.figure(figsize=(7,4.2))
plt.plot(sizes, klds, "o-", color="#2ca02c", lw=2, ms=7)
for x,y,n in zip(sizes,klds,names):
    plt.annotate(n, (x,y), textcoords="offset points", xytext=(6,6), fontsize=8)
plt.yscale("log"); plt.xlabel("file size (GB)"); plt.ylabel("KLD mean (nats, log) — lower = closer to F16")
plt.title("AMALIA-9B GGUF — fidelity vs size (KLD vs F16, pt, ctx 2048)")
plt.grid(True, which="both", alpha=0.3); plt.tight_layout()
plt.savefig(f"{PUB}/metrics/chart-kld-vs-size.png", dpi=130); plt.close()

# --- chart 2: top-1 match ---
plt.figure(figsize=(7,4.2))
t1s = [f(r["top1_match_pct"]) for r in q]
bars = plt.bar(names, t1s, color="#1f77b4")
for b,v in zip(bars,t1s): plt.text(b.get_x()+b.get_width()/2, v+0.4, f"{v:.1f}", ha="center", fontsize=8)
plt.ylabel("Top-1 token match vs F16 (%)"); plt.ylim(min(t1s)-3, 100.5)
plt.title("AMALIA-9B GGUF — top-1 agreement with F16")
plt.grid(True, axis="y", alpha=0.3); plt.tight_layout()
plt.savefig(f"{PUB}/metrics/chart-top1.png", dpi=130); plt.close()

# copy the csv into publish
import shutil; shutil.copy(CSV, f"{PUB}/metrics/quant-summary.csv")
print("charts + csv written to", f"{PUB}/metrics/")