import json, csv, re
from collections import defaultdict, Counter
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):
    sents, cur = [], []
    for line in open(path, encoding="utf-8"):
        line = line.rstrip("\r\n")
        if line.strip() == "":
            if cur: sents.append(cur); cur = []
            continue
        if "\t" in line: tok, tagstr = (line.split("\t")+["O"])[:2]
        else:
            p = line.split(" "); tok, tagstr = p[0], " ".join(p[1:])
        cur.append((tok, set(t for t in re.split(sep, tagstr.strip()) if t and t != "O")))
    if cur: sents.append(cur)
    return sents

# --- reconstruct sentences from CSV with their domain ---
by_sid = defaultdict(list); dom_of = {}
for row in csv.DictReader(open(f"{W}/gold/nested_test.csv", encoding="utf-8")):
    sid = row["global_sentence_id"]
    by_sid[sid].append((int(row["token_position"]), row["token"]))
    dom_of[sid] = row["Sub_corpus"]
csv_sents = {}
for sid, toks in by_sid.items():
    key = tuple(t for _, t in sorted(toks))
    csv_sents.setdefault(key, []).append(sid)
print("csv sentences:", len(by_sid), "unique token-tuples:", len(csv_sents))
print("csv domain counts:", Counter(dom_of.values()))

gold = read_conll(f"{W}/gold/nested_test.txt", r"[\s|,;+]+")
# map each gold sentence -> domain by token-tuple lookup
pool = {k: list(v) for k, v in csv_sents.items()}
dom_by_idx, unmatched = [], 0
for s in gold:
    key = tuple(t for t, _ in s)
    if key in pool and pool[key]:
        dom_by_idx.append(dom_of[pool[key].pop(0)])
    else:
        dom_by_idx.append(None); unmatched += 1
print("unmatched gold sentences:", unmatched)
print("mapped domain counts:", Counter(dom_by_idx))

def spans(tags, t):
    out, st = set(), None
    for i, tg in enumerate(tags):
        b, ii = ("B-"+t) in tg, ("I-"+t) in tg
        if b:
            if st is not None: out.add((st, i-1))
            st = i
        elif ii:
            if st is None: st = i
        else:
            if st is not None: out.add((st, i-1)); st = None
    if st is not None: out.add((st, len(tags)-1))
    return out

def ents(sents):
    e = set()
    for si, s in enumerate(sents):
        tg = [x for _, x in s]
        for t in TYPES:
            for a, b in spans(tg, t): e.add((si, t, a, b))
    return e

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
    return p, r, (2*p*r/(p+r) if p+r else 0.0)

G = ents(gold)
pub = {m: json.load(open(f"{W}/{m}/test_metrics.json", encoding="utf-8")) for m in ("nested-arabert","nested-marbert")}
out = {}
for name, sub in [("AraBERTv02","nested-arabert"), ("MARBERTv2","nested-marbert")]:
    P = ents(read_conll(f"{W}/{sub}/test_predictions_conll.txt", r"[|\s,;+]+"))
    pubmap = {d["domain"]: d["f1"] for d in pub[sub]["per_domain"]}
    rows = {}
    print(f"\n=== {name} per-domain (recomputed vs published) ===")
    print(f"{'DOMAIN':<20}{'sent':>6}{'ents':>6} | {'P':>7}{'R':>7}{'F1':>7} | {'pubF1':>7}{'diff':>8}")
    for d in sorted(set(x for x in dom_by_idx if x)):
        idx = {i for i, dd in enumerate(dom_by_idx) if dd == d}
        g = {e for e in G if e[0] in idx}; p_ = {e for e in P if e[0] in idx}
        tp, fp, fn = len(g & p_), len(p_ - g), len(g - p_)
        pp, rr, ff = prf(tp, fp, fn)
        rows[d] = dict(sentences=len(idx), entities=len(g), precision=round(pp,4),
                       recall=round(rr,4), f1=round(ff,4), published_f1=round(pubmap.get(d,0),4))
        print(f"{d:<20}{len(idx):>6}{len(g):>6} | {pp:>7.4f}{rr:>7.4f}{ff:>7.4f} | {pubmap.get(d,0):>7.4f}{ff-pubmap.get(d,0):>+8.4f}")
    out[name] = rows
json.dump(out, open(f"{W}/perdomain_results.json","w",encoding="utf-8"), ensure_ascii=False, indent=2)
print("\nwrote perdomain_results.json")