import json, sys, os, re
from collections import defaultdict

W = r"C:/Users/fate/AppData/Local/Temp/claude/C--Users-fate-Videos-knowldgeGraph/76bf330d-b5d6-46f3-8f74-51ca0d7b94e2/scratchpad/kgeval"

TYPES = ["CARDINAL","CURR","DATE","EVENT","FAC","GPE","LANGUAGE","LAW","LOC","MONEY",
         "NORP","OCC","ORDINAL","ORG","PERCENT","PERS","PRODUCT","QUANTITY","TIME","UNIT","WEBSITE"]

def read_conll(path, sep_tags):
    """Return list of sentences; each sentence = list of (token, set(tags))."""
    sents, cur = [], []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.rstrip("\n").rstrip("\r")
            if line.strip() == "":
                if cur: sents.append(cur); cur = []
                continue
            if "\t" in line:
                parts = line.split("\t")
                tok, tagstr = parts[0], parts[1] if len(parts) > 1 else "O"
            else:
                parts = line.split(" ")
                tok, tagstr = parts[0], " ".join(parts[1:])
            tags = set(t for t in re.split(sep_tags, tagstr.strip()) if t and t != "O")
            cur.append((tok, tags))
    if cur: sents.append(cur)
    return sents

def spans_for_type(tags_per_tok, t):
    """BIO-decode one type channel -> set of (start, end) inclusive."""
    out, start = set(), None
    for i, tags in enumerate(tags_per_tok):
        b, ii = ("B-"+t) in tags, ("I-"+t) in tags
        if b:
            if start is not None: out.add((start, i-1))
            start = i
        elif ii:
            if start is None: start = i          # I- without B- starts a span (seqeval IOB2-lenient)
        else:
            if start is not None: out.add((start, i-1)); start = None
    if start is not None: out.add((start, len(tags_per_tok)-1))
    return out

def entities(sents):
    """-> set of (sent_idx, type, start, end)"""
    ents = set()
    for si, s in enumerate(sents):
        tags = [t for _, t in s]
        for t in TYPES:
            for (a, b) in spans_for_type(tags, t):
                ents.add((si, t, a, b))
    return ents

def prf(tp, fp, fn):
    p = tp/(tp+fp) if tp+fp else 0.0
    r = tp/(tp+fn) if tp+fn else 0.0
    f = 2*p*r/(p+r) if p+r else 0.0
    return p, r, f

gold = read_conll(f"{W}/gold/nested_test.txt", r"[\s|,;+]+")
print("gold sentences:", len(gold), "tokens:", sum(len(s) for s in gold))

# domain map from CSV, in file order of global_sentence_id first appearance
import csv
dom_by_sent = []
seen = set()
with open(f"{W}/gold/nested_test.csv", encoding="utf-8") as f:
    for row in csv.DictReader(f):
        sid = row["global_sentence_id"]
        if sid not in seen:
            seen.add(sid); dom_by_sent.append(row["Sub_corpus"])
print("domains listed for", len(dom_by_sent), "sentences")

gold_ents = entities(gold)
print("gold entities:", len(gold_ents))

results = {}
for name, sub in [("EXP-001 AraBERTv02 nested", "nested-arabert"),
                  ("EXP-002 MARBERTv2 nested", "nested-marbert")]:
    pred = read_conll(f"{W}/{sub}/test_predictions_conll.txt", r"[|\s,;+]+")
    assert len(pred) == len(gold), (name, len(pred), len(gold))
    mism = sum(1 for a, b in zip(gold, pred) if len(a) != len(b))
    tokmis = sum(1 for a,b in zip(gold,pred) for (ta,_),(tb,_) in zip(a,b) if ta != tb)
    pred_ents = entities(pred)
    tp = len(gold_ents & pred_ents); fp = len(pred_ents - gold_ents); fn = len(gold_ents - pred_ents)
    P, R, F = prf(tp, fp, fn)
    print(f"\n### {name}  (sent len mismatches={mism}, token mismatches={tokmis})")
    print(f"overall micro  P={P:.4f} R={R:.4f} F1={F:.4f}   TP={tp} FP={fp} FN={fn}  pred_ents={len(pred_ents)}")

    # per type
    per_type = {}
    for t in TYPES:
        g = {e for e in gold_ents if e[1] == t}; p_ = {e for e in pred_ents if e[1] == t}
        tp_, fp_, fn_ = len(g & p_), len(p_ - g), len(g - p_)
        pp, rr, ff = prf(tp_, fp_, fn_)
        per_type[t] = dict(support=len(g), pred=len(p_), tp=tp_, fp=fp_, fn=fn_,
                           precision=round(pp,4), recall=round(rr,4), f1=round(ff,4))
    # per domain (recompute P/R too)
    per_dom = {}
    for d in sorted(set(dom_by_sent)):
        idx = {i for i, dd in enumerate(dom_by_sent) if dd == d}
        g = {e for e in gold_ents if e[0] in idx}; p_ = {e for e in pred_ents if e[0] in idx}
        tp_, fp_, fn_ = len(g & p_), len(p_ - g), len(g - p_)
        pp, rr, ff = prf(tp_, fp_, fn_)
        per_dom[d] = dict(sentences=len(idx), support=len(g),
                          precision=round(pp,4), recall=round(rr,4), f1=round(ff,4))
    macro = sum(v["f1"] for v in per_type.values())/len(TYPES)
    results[name] = dict(overall=dict(precision=round(P,4), recall=round(R,4), f1=round(F,4),
                                      tp=tp, fp=fp, fn=fn, gold_entities=len(gold_ents),
                                      pred_entities=len(pred_ents), macro_f1_over_types=round(macro,4)),
                         per_type=per_type, per_domain=per_dom)

with open(f"{W}/pertype_results.json", "w", encoding="utf-8") as f:
    json.dump(dict(gold_sentences=len(gold), gold_tokens=sum(len(s) for s in gold),
                   gold_entities=len(gold_ents), results=results), f, ensure_ascii=False, indent=2)
print("\nwrote pertype_results.json")