import csv, subprocess, statistics
# decomposed doc ids (348)
import os
os.chdir(os.path.expanduser("~/Documents/code/hyungi_Document_Server"))
out = subprocess.check_output(["docker","compose","exec","-T","postgres","psql","-U","pkm","-d","pkm","-t","-A","-c",
  "SELECT doc_id FROM document_chunks WHERE source_type='hier_section' AND is_leaf GROUP BY doc_id"]).decode()
decomposed = {int(x) for x in out.split() if x.strip().isdigit()}

def load(path):
    rows={}
    with open(path) as f:
        for r in csv.DictReader(f):
            rows[r["id"]] = r
    return rows

pre = load("/tmp/hier_prehier_exact.csv")
clean = load("/tmp/hier_clean_exact.csv")
raw = load("/tmp/hier_raw_exact.csv")

def targets(r):
    gr = r["graded_relevance"]
    ids=[]
    for part in gr.split(";"):
        if ":" in part:
            ids.append(int(part.split(":")[0]))
    return ids

affected=[]; unchanged=[]
for qid,r in pre.items():
    tg=targets(r)
    if any(t in decomposed for t in tg):
        affected.append(qid)
    else:
        unchanged.append(qid)

def mean_ndcg(rows, ids):
    vals=[float(rows[i]["graded_ndcg_at_10"]) for i in ids if i in rows]
    return statistics.mean(vals) if vals else 0.0, len(vals)

print(f"decomposed docs: {len(decomposed)}")
print(f"queries affected (target in decomposed): {len(affected)} | unchanged: {len(unchanged)}")
print()
for name,ids in [("AFFECTED-SUBSET (decision權 B2)",affected),("UNCHANGED (sanity: should be ~equal)",unchanged),("FULL-51",list(pre.keys()))]:
    p,_=mean_ndcg(pre,ids); c,_=mean_ndcg(clean,ids); rw,n=mean_ndcg(raw,ids)
    print(f"{name} (n={n}):")
    print(f"   prehier={p:.3f}  clean={c:.3f} (Δ{c-p:+.3f})  raw={rw:.3f} (Δ{rw-p:+.3f})")
print()
# per-category within affected subset
cats={}
for qid in affected:
    cats.setdefault(pre[qid]["category"],[]).append(qid)
print("AFFECTED subset by category (prehier -> clean):")
for cat,ids in sorted(cats.items()):
    p,_=mean_ndcg(pre,ids); c,n=mean_ndcg(clean,ids)
    print(f"   {cat:14s} n={n:2d}  prehier={p:.3f} clean={c:.3f} (Δ{c-p:+.3f})")
print()
# biggest per-query losers/winners (clean vs prehier) within affected
deltas=[(float(clean[q]["graded_ndcg_at_10"])-float(pre[q]["graded_ndcg_at_10"]), q, pre[q]["category"], pre[q]["query"][:34]) for q in affected]
deltas.sort()
print("Biggest LOSERS (clean - prehier):")
for d,q,c,qry in deltas[:6]:
    print(f"   {d:+.3f} {q} [{c}] {qry}")
print("Biggest WINNERS:")
for d,q,c,qry in deltas[-5:][::-1]:
    print(f"   {d:+.3f} {q} [{c}] {qry}")
