"""Heading recovery from PDFs using font size / weight, scored against the PDF's own outline (bookmarks)."""
import re, sys, collections, pypdf, pdfplumber

PDFS = ["1706.03762", "2005.11401", "2201.11903", "2210.03629", "2312.10997"]
DIR = sys.argv[1]

def norm(s):
    s = re.sub(r"^\s*(?:[A-Z]|\d+)(?:\.\d+)*\.?\s+", "", s.strip())  # drop "2.1 " / "A "
    return re.sub(r"[^a-z0-9]", "", s.lower())

def outline(path):
    r = pypdf.PdfReader(path); out = []
    def f(o):
        for x in o:
            if isinstance(x, list): f(x)
            else: out.append(x.title)
    f(r.outline); return out

def lines(path):
    res = []
    with pdfplumber.open(path) as pdf:
        for pn, page in enumerate(pdf.pages):
            words = page.extract_words(extra_attrs=["size", "fontname"], keep_blank_chars=False)
            rows = collections.defaultdict(list)
            for w in words: rows[round(w["top"] / 3)].append(w)
            for k in sorted(rows):
                ws = sorted(rows[k], key=lambda w: w["x0"])
                text = " ".join(w["text"] for w in ws)
                size = max(w["size"] for w in ws)
                bold = all(re.search(r"bold|bd|black|heavy|-b$|bx|medi|semib|demi", w["fontname"], re.I) for w in ws)
                res.append(dict(text=text, size=round(size, 1), bold=bold, page=pn, nwords=len(ws), x0=ws[0]["x0"]))
    return res

def body_size(ls):
    c = collections.Counter()
    for l in ls: c[l["size"]] += len(l["text"])
    return c.most_common(1)[0][0]

NUM = re.compile(r"^(?:\d+(?:\.\d+)*\.?|[A-Z](?:\.\d+)*\.?)\s+[A-Z]")
def short(l): return 1 <= l["nwords"] <= 14 and len(l["text"]) < 100 and not l["text"].rstrip().endswith((".", ",", ";"))

METHODS = {
    "A. larger than body font": lambda l, b: l["size"] >= b + 0.5 and short(l),
    "B. larger or bold": lambda l, b: (l["size"] >= b + 0.5 or l["bold"]) and short(l),
    "C. bold/larger + numbered or capitalised": lambda l, b: (l["size"] >= b + 0.5 or l["bold"]) and short(l) and (NUM.match(l["text"]) is not None or l["text"][:1].isupper()),
    "D. bold/larger + numbered pattern": lambda l, b: (l["size"] >= b + 0.5 or l["bold"]) and short(l) and NUM.match(l["text"]) is not None,
    "E. any line with numbered pattern (no font)": lambda l, b: short(l) and NUM.match(l["text"]) is not None,
}

def match(p, t): return p == t or (len(p) >= 8 and t.startswith(p))

tot = {m: dict(tp=0, pred=0, hit=0, truth=0) for m in METHODS}
print(f"{'pdf':12} {'pages':>5} {'outline':>7} {'body pt':>7}")
for pid in PDFS:
    path = f"{DIR}/{pid}.pdf"
    truth = [norm(t) for t in outline(path)]; ls = lines(path); b = body_size(ls)
    print(f"{pid:12} {len(set(l['page'] for l in ls)):>5} {len(truth):>7} {b:>7}")
    for m, fn in METHODS.items():
        pred = [norm(l["text"]) for l in ls if fn(l, b)]
        tp = sum(any(match(p, t) for t in truth) for p in pred)
        hit = sum(any(match(p, t) for p in pred) for t in truth)
        for k, v in dict(tp=tp, pred=len(pred), hit=hit, truth=len(truth)).items(): tot[m][k] += v
        print(f"   {m:44} pred={len(pred):4} precision={tp/max(len(pred),1):.2f} recall={hit/len(truth):.2f}")
print("\nTOTAL (all five PDFs pooled)")
print(f"{'method':46} {'pred':>5} {'prec':>6} {'recall':>7} {'F1':>6}")
for m, t in tot.items():
    p = t["tp"] / max(t["pred"], 1); r = t["hit"] / t["truth"]; f1 = 2 * p * r / (p + r) if p + r else 0
    print(f"{m:46} {t['pred']:>5} {p:6.2f} {r:7.2f} {f1:6.2f}")