import sys, os, json, glob, re, copy
from pathlib import Path
# Run from the repo root (move up if the notebook was opened inside notebooks/)
if Path.cwd().name == "notebooks": os.chdir("..")
sys.path.insert(0, "code")
import upstage_eval as ue
import metrics as M
from bs4 import BeautifulSoup
assert os.environ.get("UPSTAGE_API_KEY"), "export UPSTAGE_API_KEY first"
RES, FIG = Path("results"), Path("figures")
for d in [RES, RES/"parse", RES/"ie", RES/"report", FIG]: d.mkdir(parents=True, exist_ok=True)
FORCE = False  # set True to re-call the API (ignore cache)
def cached_parse(path, name, cats=("table",)):
    out = RES/"parse"/f"{name}.json"
    if out.exists() and not FORCE: return json.loads(out.read_text())
    resp = ue.document_parse(path, base64_categories=cats)
    out.write_text(json.dumps(resp, ensure_ascii=False)); return resp

EXP1B = {"table_page":"data/exp1b_paper/table_page.pdf",
         "chart_page":"data/exp1b_paper/chart_page.pdf",
         "diagram_page":"data/exp1b_paper/diagram_page.pdf"}
for name, path in EXP1B.items():
    resp = cached_parse(path, f"1b_{name}", cats=("table","figure","chart","equation"))
    ue.save_crops(resp, FIG/f"1b_{name}_crops")
    ue.draw_bboxes(path, resp, FIG/f"1b_{name}_overlay.png")
    print(f"{name:14s} {ue.categories(resp)} -> saved overlay + crops")

"""Quantitative scoring for the figure/table experiment:
 - Track A localization recall + Track B chart series/type stats
 - table TEDS (parse table HTML vs gt/exp1_tables/*.html)
Run: ./.venv/bin/python code/score_exp1.py
"""
import sys, json, re
from pathlib import Path
from bs4 import BeautifulSoup
sys.path.insert(0, "code")
import upstage_eval as ue
import metrics as M

RES = Path("results"); out = {}

# ---------- Track A: localization recall ----------
loc = json.load(open("gt/exp1b_localization.json"))
cats = ["table","figure","chart","caption","equation"]
agg = {c: [0,0] for c in cats}   # [detected_correct, true]
per_page = {}
for pg, v in loc.items():
    if not isinstance(v, dict) or "true_counts" not in v: continue
    det, tru = v["detected"], v["true_counts"]
    row = {}
    for c in cats:
        t = tru.get(c,0); d = det.get(c,0)
        if t==0 and d==0: continue
        hit = min(d,t)                       # recall: does detection cover the true count?
        agg[c][0]+=hit; agg[c][1]+=t
        row[c] = f"{hit}/{t}" + (f" (+{d-t} over-detected)" if d>t else "")
    per_page[pg]=row
prim = ["table","figure","chart"]
prim_hit = sum(agg[c][0] for c in prim); prim_tot = sum(agg[c][1] for c in prim)
out["localization"] = {
  "per_page": per_page,
  "primary_recall": f"{prim_hit}/{prim_tot}",
  "caption_recall": f"{agg['caption'][0]}/{agg['caption'][1]}",
  "note": "primary=table+chart+figure. caption: missed on diagram_page (detected as paragraph)."
}

# ---------- Track B: chart series/type ----------
cg = json.load(open("gt/exp1b_chart_values.json")).get("charts_VERIFY",{})
els = ue.elements(json.loads((RES/"parse"/"1b_chart_page.json").read_text()))
charts = [e for e in els if e.get("category")=="chart"]
rows=[]; series_lost=0; type_wrong=0
for e in charts:
    h=e["content"]["html"]
    pred_series = len(BeautifulSoup(h,"lxml").select("tbody tr"))
    m=re.search(r"Chart Type:\s*([a-zA-Z]+)", h); pred_type=m.group(1) if m else "?"
    g=cg.get(f"el{e['id']}",{}); true_series=g.get("true_series"); true_type=g.get("true_type")
    lost = (true_series or 0) > pred_series
    twrong = true_type and pred_type.lower() not in true_type.lower()
    series_lost+=lost; type_wrong+=twrong
    rows.append({"id":e["id"],"pred_type":pred_type,"true_type":true_type,
                 "pred_series":pred_series,"true_series":true_series,
                 "series_lost":lost,"type_wrong":bool(twrong)})
out["chart"] = {"charts":rows,
  "series_lost": f"{series_lost}/{len(charts)}",
  "type_wrong":  f"{type_wrong}/{len(charts)}",
  "note": "all 4 charts are 2-series dot plots → collapsed to a single item_01 series, type misclassified as line/bar."}

# ---------- table TEDS ----------
def parse_tables(doc):
    els=ue.elements(json.loads((RES/"parse"/f"{doc}.json").read_text()))
    return [e["content"]["html"] for e in els if e.get("category")=="table"]
def gt_tables(doc):
    html=Path(f"gt/exp1_tables/{doc}.html").read_text()
    return [str(t) for t in BeautifulSoup(html,"lxml").find_all("table")]  # split multiple <table>s

teds_res=[]
for doc in ["complex_table_1","complex_table_2","korean_table"]:
    P, G = parse_tables(doc), gt_tables(doc)
    n=min(len(P),len(G))
    scores=[M.teds(P[i], G[i]) for i in range(n)]
    circular=all(P[i].strip()==G[i].strip() for i in range(n))  # is GT verbatim the parse output?
    teds_res.append({"doc":doc,"n_tables":n,
                     "teds_per_table":[round(s,3) for s in scores],
                     "teds_mean":round(sum(scores)/n,3) if n else None,
                     "gt_is_parse_output":circular})
out["teds"]={"by_doc":teds_res,
  "VALID": not any(t["gt_is_parse_output"] for t in teds_res),
  "note": "GT HTML is identical to the Document Parse output → TEDS=1.0 is circular (meaningless). "
          "A valid TEDS needs an independent GT built by hand from the source PDF."}

json.dump(out, open(RES/"scores_exp1b.json","w"), ensure_ascii=False, indent=2)
print("=== Localization ===")
print(" primary(table+chart+figure):", out["localization"]["primary_recall"],
      "| caption:", out["localization"]["caption_recall"])
for pg,r in per_page.items(): print("  ",pg,r)
print("=== Chart === series_lost:", out["chart"]["series_lost"], "| type_wrong:", out["chart"]["type_wrong"])
print("=== TEDS ===")
for t in teds_res: print("  ",t["doc"], "mean", t["teds_mean"], t["teds_per_table"])
print("saved -> results/scores_exp1b.json")

Model Test 1 — Document Parse: figure/table, two tracks¶

1. Parse + crops + bbox overlay (Track A input)¶

2. Scoring — detection recall + chart series/type (+ table TEDS for reference)¶