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

"""Experiment — read the paper and have Solar write a weakness-diagnosis + dataset-augmentation
report for **OpenAI CUA**, with supporting tables and quoted passages.
Input is Document Parse's structured HTML; model behavior is observed as-is (no hallucination limits).
Run: UPSTAGE_API_KEY=... ./.venv/bin/python code/run_exp4_table.py
"""
import sys, json, re
from pathlib import Path
from bs4 import BeautifulSoup
sys.path.insert(0, "code")
import upstage_eval as ue

RES = Path("results"); (RES/"report").mkdir(parents=True, exist_ok=True)
resp = json.loads((RES/"parse"/"webstep_full.json").read_text())
els = ue.elements(resp)
context = "\n".join(e["content"]["html"] for e in els if e["content"].get("html"))

# The task prompt is intentionally Korean (the experiment elicits a Korean report, analyzed in the note):
# "diagnose OpenAI CUA's weak skills and propose dataset augmentation; for each claim cite the
#  supporting table's numbers (with the Table number) and a quoted passage; reconstruct the tables."
Q = ("논문을 읽고 **OpenAI CUA**의 약점 skill이 무엇인지 진단하고, 그 약점을 메우려면 "
     "어떤 skill·task 중심으로 데이터셋을 보강하면 좋을지 제안하는 리포트를 작성하라. "
     "각 주장마다 (1) 근거가 된 표의 수치(Table 번호와 함께)와 (2) 근거가 된 논문 구절을 인용할 것. "
     "근거 표는 markdown으로 재구성해 보여줄 것.")
SYS = ("You write an analytical report about OpenAI CUA using ONLY the provided document content "
       "(text + the paper's tables as HTML). Support every claim with BOTH (a) the source table "
       "(Table N) and its numbers and (b) a quoted passage from the paper. Reconstruct supporting "
       "tables as Korean markdown. Answer in Korean markdown.")

def fix_tables(md):
    """GFM tables need a blank line before them — patch the blank lines Solar omits."""
    out=[]
    for l in md.split("\n"):
        if l.lstrip().startswith("|") and out and out[-1].strip() and not out[-1].lstrip().startswith("|"):
            out.append("")
        out.append(l)
    return "\n".join(out)

def ask(ctx, maxctx=150000):
    ctx = ctx[:maxctx]
    md = ue.solar_chat([{"role":"system","content":SYS},
        {"role":"user","content":f"[PAPER]\n{ctx}\n\n[Q]\n{Q}"}], max_tokens=3000)["content"]
    return fix_tables(md)

print("context chars:", len(context))
rep_parse = ask(context)                                # Parse structured HTML → CUA report
(RES/"report"/"webstep_review.md").write_text(rep_parse)
print("report written -> webstep_review.md")

# ---- truthfulness check: are the numbers the report attributes to 'OpenAI CUA' real CUA values? ----
# GT = union of the real CUA-row numbers across all parsed tables.
GT = set()
for e in els:
    if e.get("category")!="table": continue
    soup=BeautifulSoup(e["content"]["html"],"lxml")
    for tr in soup.find_all("tr"):
        tds=[td.get_text(strip=True) for td in tr.find_all(["td","th"])]
        if tds and ("CUA" in tds[0] or "OpenAI" in tds[0]):
            for v in tds[1:]:
                m=re.fullmatch(r"(\d+\.?\d*)%?", v)
                if m: GT.add(m.group(1))
print(f"GT (union of real CUA numbers): {len(GT)}")

def grounding(md):
    """Fraction of distinct numbers in the report's 'CUA' markdown rows that match a real CUA value."""
    cited=set()
    for ln in md.splitlines():
        if not ln.lstrip().startswith("|") or "CUA" not in ln: continue
        for tok in re.findall(r"\d+\.?\d*", ln.replace("**","")):
            cited.add(tok)
    ok={c for c in cited if c in GT}
    return len(ok), len(cited), sorted(cited - GT)   # grounded, cited, unmatched (potential errors)

gp = grounding(rep_parse)
out={"gt_cua_values":sorted(GT),
     "parse_html":{"grounded":gp[0],"cited":gp[1],"precision":round(gp[0]/gp[1],3) if gp[1] else None,"ungrounded":gp[2]}}
json.dump(out, open(RES/"exp4_table_fidelity.json","w"), ensure_ascii=False, indent=2)
print(f"=== CUA-row numeric grounding: {gp[0]}/{gp[1]}  unmatched={gp[2]}")
print("saved -> results/exp4_table_fidelity.json ; report -> results/report/webstep_review.md")

In-the-wild 1 — Table-grounded paper review (Document Parse → Solar-pro3)¶

Generate the report (Parse HTML → Solar-pro3) + score numeric grounding¶