"""Offline grounding atlas — the existence base (no live API, no rate-limit). An entity is GROUNDED if
it (or a >=2-significant-token contiguous sub-span of it) is an exact Wikipedia article title. Conservative
by construction: fabrications have no title to match (structural 0-false-rescue, same property the live
gate had); but unlike the live string-search it can't be rate-limited and handles descriptive prompts
(the proper-noun span resolves) + diacritics (folded). The Nomic embedding layer (v2, DGX) adds semantic
robustness for transliteration/word-order misses.

Build:  python offline_atlas.py build  enwiki-titles.gz  wiki_titles.db
Probe:  python offline_atlas.py probe  wiki_titles.db    ../eval/heldout_battery.json  ../eval/heldout_atlas_offline.json
"""
import sys, os, re, gzip, sqlite3, unicodedata, json

STOP = {"the", "of", "a", "an", "and", "in", "on", "at", "de", "la", "le", "el", "los", "las",
        "von", "van", "der", "di", "du", "do", "da", "for", "to"}


def norm(s):
    s = unicodedata.normalize("NFKD", s).encode("ascii", "ignore").decode()   # fold diacritics
    s = s.replace("_", " ").lower()
    s = re.sub(r"[^a-z0-9 ]", " ", s)
    return re.sub(r"\s+", " ", s).strip()


def _sig(tok):
    return tok not in STOP and len(tok) > 1


def build(gz_path, db_path):
    if os.path.exists(db_path):
        os.remove(db_path)
    con = sqlite3.connect(db_path)
    con.execute("PRAGMA journal_mode=OFF"); con.execute("PRAGMA synchronous=OFF")
    con.execute("CREATE TABLE t(n TEXT PRIMARY KEY) WITHOUT ROWID")
    batch, total = [], 0
    with gzip.open(gz_path, "rt", encoding="utf-8", errors="ignore") as f:
        first = f.readline()
        if norm(first) and "page_title" not in first:        # not a header -> keep it
            batch.append((norm(first),))
        for line in f:
            nm = norm(line)
            if nm:
                batch.append((nm,))
            if len(batch) >= 100000:
                con.executemany("INSERT OR IGNORE INTO t VALUES(?)", batch)
                total += len(batch); batch = []
                if total % 1000000 == 0:
                    print(f"  {total//1000000}M titles ...", flush=True)
    if batch:
        con.executemany("INSERT OR IGNORE INTO t VALUES(?)", batch); total += len(batch)
    con.commit()
    n = con.execute("SELECT COUNT(*) FROM t").fetchone()[0]
    con.close()
    print(f"built {db_path}: {n} unique normalized titles (from {total} lines)")


class Grounder:
    def __init__(self, db_path):
        self.con = sqlite3.connect(db_path, check_same_thread=False)

    def _exists(self, nm):
        return self.con.execute("SELECT 1 FROM t WHERE n=? LIMIT 1", (nm,)).fetchone() is not None

    def _candidates(self, entity):
        # whole-entity candidates (each matched in full, never sub-windowed -> no generic-concept rescue):
        # the entity, the part before a comma (", Country"/disambig), and the pieces of descriptive
        # scaffolding (Work BY Author, Work (Native Title), Author'S Work) — these recover reals without
        # surfacing generic fragments, and fakes aren't formatted this way.
        # SAFE splits only: comma (geo/disambig qualifier) and the text outside parentheses (native-title
        # parentheticals). by/possessive/inside-paren splits surface authors+years that false-rescue, so
        # they're left to the embedding+type layer.
        c = [entity]
        if "," in entity:
            c.append(entity.split(",")[0])
        if " by " in entity:
            c.append(entity.split(" by ")[0])                 # the WORK (not the author after "by")
        m = re.search(r"'s\s+(.+)", entity)
        if m:
            c.append(m.group(1))                              # the WORK (not the author before "'s")
        c2 = []
        for s in c:
            c2.append(s)
            if "(" in s:
                c2.append(re.sub(r"\([^)]*\)", "", s))        # also the text outside parens
        return [s.strip() for s in c2 if s.strip()]

    def grounded(self, entity):
        # 1) whole-candidate exact match (entity + comma/by/paren/possessive splits)
        for m in self._candidates(entity):
            nm = norm(m)
            if nm and self._exists(nm):
                return {"matched": True, "hit": nm}
        # 2) coverage-gated sub-windows: a title must cover >=60% of the entity's significant tokens
        #    (kills generic-concept rescues like "multi task"; keeps "Matilde Hidalgo"/"Lothar Meyer")
        e = norm(entity); toks = e.split()
        esig = [t for t in toks if _sig(t)]
        if not esig:
            return {"matched": False, "hit": ""}
        for L in range(len(toks), 1, -1):
            for i in range(0, len(toks) - L + 1):
                w = toks[i:i + L]
                wsig = [t for t in w if _sig(t)]
                if len(wsig) >= 2 and len(wsig) / len(esig) >= 0.6:
                    wn = " ".join(w)
                    if self._exists(wn):
                        return {"matched": True, "hit": wn}
        return {"matched": False, "hit": ""}


def probe(db_path, battery_path, out_path):
    g = Grounder(db_path)
    items = json.load(open(battery_path))["items"]
    per = []
    for it in items:
        r = g.grounded(it["entity"])
        per.append({"entity": it["entity"], "label": it["label"], "category": it["category"],
                    "matched": r["matched"], "hit": r["hit"]})
    fake = [p for p in per if p["label"] == 1]; real = [p for p in per if p["label"] == 0]
    fr = [(p["entity"], p["hit"]) for p in fake if p["matched"]]
    mr = [p["entity"] for p in real if not p["matched"]]
    print(f"offline atlas on {os.path.basename(battery_path)}:")
    print(f"  matched-rate  REAL {sum(p['matched'] for p in real)/len(real):.3f}  FAKE {sum(p['matched'] for p in fake)/len(fake):.3f}")
    print(f"  false-rescues (fakes matched): {len(fr)}  {fr[:8]}")
    print(f"  still-missed reals: {len(mr)}  {mr[:10]}")
    json.dump({"per_item": per}, open(out_path, "w"), indent=1)
    print(f"saved -> {out_path}")


if __name__ == "__main__":
    cmd = sys.argv[1]
    if cmd == "build":
        build(sys.argv[2], sys.argv[3])
    elif cmd == "probe":
        probe(sys.argv[2], sys.argv[3], sys.argv[4])