"""THE DECISIVE TEST: does GLOBAL name uniqueness rescue the empty-address pairs?

WHY. My earlier rescue test scored each pair in isolation: 'exact normalised name match' on
empty-address candidates gave 442 TP / 3,074 FP = 0.14:1 against a ~3.76:1 break-even, and I
concluded the pairs were unrecoverable. That test never asked whether any OTHER S1 entity was
competing for the record - which is global information, and is what `new plan.md` Step 1/Step 4
and LIFTOVER R3 both propose (reverse search, competing-owner features).

Measured on the raw data: of the empty-address pool records whose core name matches at least one
S1, 55% (train) / 59% (test) match EXACTLY ONE S1. For those the assignment is forced: there is no
competitor. That is a completely different question from 'is this pair similar', and it is the most
plausible explanation for a 0.99066 leaderboard score.

The competitor count is computed over ALL 2.2M train S1 entities, not just the 20k val slice,
because a competitor can be anywhere in the data. That is the whole point.
"""
import collections
import json
import os
import re
import sys
import unicodedata

import joblib
import numpy as np
import pandas as pd

sys.path.insert(0, "src")
import config as C
import decide as DC
from evaluate import macro_f05

RAW = os.environ.get("RAW_DIR", "/scratch/user4/dai/dataset_unstop/dataset/train")
_WS = re.compile(r"\s+")
_PUN = re.compile(r"[^\w\s]", re.UNICODE)
LEGAL = set("inc incorporated llc ltd limited corp corporation co company lp llp pllc pc pa plc "
            "trust holdings pvt private pte sa sarl sas sasu sci eurl snc".split())


def core(s):
    s = "".join(c for c in unicodedata.normalize("NFKD", str(s)) if not unicodedata.combining(c))
    s = _WS.sub(" ", _PUN.sub(" ", s.casefold())).strip()
    return " ".join(t for t in s.split() if t not in LEGAL)


b = joblib.load("work/results_A/model_bundle.joblib")
MM, THR = b["miss_mass"], b["param"]
CACHE = "work/val_scored_20000.parquet"
cand = pd.read_parquet(CACHE)
truths = {int(k): set(v) for k, v in json.load(open(CACHE + ".truth.json")).items()}
ids = sorted(truths)
shipped = DC.predict_sets(cand, "q", "c", "p", b["method"], THR, MM, True)
m0, _ = macro_f05(shipped, truths, ids)
print(f"baseline F0.5 {m0['f05']:.5f}  P {m0['precision']:.4f}  R {m0['recall']:.4f}")

p = C.PATHS["train"]
s1p = pd.read_parquet(p["source1"], columns=["entity_id"])
s1_ids = s1p["entity_id"].values
pool_ids = np.concatenate([
    pd.read_parquet(p["source2"], columns=["entity_id"])["entity_id"].values,
    pd.read_parquet(p["source3"], columns=["entity_id"])["entity_id"].values])

s1raw = pd.read_csv(os.path.join(RAW, "train_source1.tsv"), sep="\t", dtype=str,
                    keep_default_na=False, na_filter=False)
print(f"all S1 entities (the competitor universe): {len(s1raw):,}")
s1_core = {}
name_count = collections.Counter()
for e, n, c in zip(s1raw["entity_id"], s1raw["business_name"], s1raw["country"]):
    k = core(n) + "\x00" + c
    s1_core[e] = k
    name_count[k] += 1

praw = {}
for src in (2, 3):
    d = pd.read_csv(os.path.join(RAW, f"train_source{src}.tsv"), sep="\t", dtype=str,
                    keep_default_na=False, na_filter=False)
    for e, n, a, c in zip(d["entity_id"], d["business_name"], d["business_address"],
                          d["country"]):
        praw[e] = (core(n) + "\x00" + c, not a.strip())

qv = cand["q"].to_numpy()
cv = cand["c"].to_numpy()
pv = cand["p"].to_numpy()
y = np.array([1 if c in truths.get(q, ()) else 0 for q, c in zip(qv, cv)], np.int8)
selp = {(q, c) for q, cs in shipped.items() for c in cs}
sel = np.array([1 if (q, c) in selp else 0 for q, c in zip(qv, cv)], np.int8)

rows = []
for i in range(len(cand)):
    ck, ea = praw[pool_ids[cv[i]]]
    qk = s1_core[s1_ids[qv[i]]]
    rows.append((ea, ck == qk and ck.split("\x00")[0] != "", name_count.get(ck, 0)))
F = pd.DataFrame(rows, columns=["empty_addr", "name_eq", "n_competitors"])
F["y"] = y
F["sel"] = sel
F["p"] = pv
F["q"] = qv
F["c"] = cv

ea = F[F.empty_addr]
print(f"\nempty-address candidate pairs in val: {len(ea):,}")
un = ea[ea.sel == 0]
print(f"currently REJECTED: {len(un):,} ({int((un.y==1).sum()):,} true, "
      f"{int((un.y==0).sum()):,} false)")

print("\n" + "=" * 78)
print("EXACT NAME MATCH, split by HOW MANY S1 ENTITIES COMPETE for that record")
print("=" * 78)
m = un[un.name_eq]
print(f"  rejected empty-address pairs with an exact core-name match: {len(m):,}")
print(f"{'n_competitors':>14}{'TP':>8}{'FP':>9}{'ratio':>9}   (needs > 3.76)")
for lo, hi, lab in [(1, 1, "1 (forced)"), (2, 2, "2"), (3, 5, "3-5"), (6, 10, "6-10"),
                    (11, 10 ** 9, "11+")]:
    s = m[(m.n_competitors >= lo) & (m.n_competitors <= hi)]
    tp, fp = int((s.y == 1).sum()), int((s.y == 0).sum())
    if tp + fp == 0:
        continue
    r = tp / fp if fp else float("inf")
    flag = "   <== CLEARS" if r > 3.76 else ""
    print(f"{lab:>14}{tp:>8,}{fp:>9,}{r:>9.2f}{flag}")

print("\n" + "=" * 78)
print("RULES combining uniqueness with the pairwise score")
print("=" * 78)
print(f"{'rule':<52}{'TP':>7}{'FP':>8}{'ratio':>8}")
rules = [
    ("name_eq AND n_competitors==1", (un.name_eq) & (un.n_competitors == 1)),
    ("name_eq AND n_comp==1 AND p>=0.05", (un.name_eq) & (un.n_competitors == 1) & (un.p >= 0.05)),
    ("name_eq AND n_comp==1 AND p>=0.20", (un.name_eq) & (un.n_competitors == 1) & (un.p >= 0.20)),
    ("name_eq AND n_comp==1 AND p>=0.40", (un.name_eq) & (un.n_competitors == 1) & (un.p >= 0.40)),
    ("name_eq AND n_comp<=2", (un.name_eq) & (un.n_competitors <= 2)),
    ("name_eq AND n_comp<=2 AND p>=0.20", (un.name_eq) & (un.n_competitors <= 2) & (un.p >= 0.20)),
]
best = None
for name, mask in rules:
    tp, fp = int((mask & (un.y == 1)).sum()), int((mask & (un.y == 0)).sum())
    if tp + fp == 0:
        print(f"  {name:<50}{0:>7}{0:>8}{'-':>8}")
        continue
    r = tp / fp if fp else float("inf")
    flag = "   <== CLEARS" if r > 3.76 else ""
    print(f"  {name:<50}{tp:>7,}{fp:>8,}{r:>8.2f}{flag}")
    if r > 3.76 and (best is None or tp > best[1]):
        best = (name, tp, fp, mask)

print("\n" + "=" * 78)
print("END TO END, held out on half the entities")
print("=" * 78)
h = np.array([hash(("half", int(q))) % 2 for q in ids])
HA = set(np.array(ids)[h == 0].tolist())
HAl = sorted(HA)
HB = [q for q in ids if q not in HA]


def ev(pr, idl):
    return macro_f05({q: pr.get(q, []) for q in idl}, truths, idl)[0]["f05"]


bA, bB = ev(shipped, HAl), ev(shipped, HB)
print(f"baseline  half A {bA:.5f}   half B {bB:.5f}")
if best is None:
    print("\nNothing clears the break-even even with global uniqueness. Then the empty-address")
    print("pairs really are unrecoverable and 0.99066 must come from somewhere else entirely.")
else:
    name, tp, fp, mask = best
    add = un[mask]
    newpred = {q: set(v) for q, v in shipped.items()}
    for q, c in zip(add["q"], add["c"]):
        newpred.setdefault(int(q), set()).add(int(c))
    pr = {q: list(v) for q, v in newpred.items()}
    a, bb = ev(pr, HAl), ev(pr, HB)
    print(f"\nbest: {name}   (+{len(add):,} pairs, {tp} TP / {fp} FP)")
    print(f"  half A {a:.5f}  ({a-bA:+.5f})")
    print(f"  half B {bb:.5f}  ({bb-bB:+.5f})   <-- HELD OUT, the honest number")
    dup = collections.Counter()
    for q, cs in newpred.items():
        for c in cs:
            dup[c] += 1
    print(f"  one-owner violations: {sum(1 for v in dup.values() if v > 1):,}")

print("\n" + "=" * 78)
print("SCALE-UP: how many such pairs exist across ALL of val, not just rejected ones?")
print("=" * 78)
allm = ea[(ea.name_eq) & (ea.n_competitors == 1)]
print(f"  empty-address + exact name + unique competitor: {len(allm):,} pairs")
print(f"    already selected : {int((allm.sel == 1).sum()):,}")
print(f"    true             : {int((allm.y == 1).sum()):,}")
print(f"  -> per entity that is {len(allm)/len(ids):.4f}; scaled to 1.73M test entities "
      f"about {len(allm)/len(ids)*1732544:,.0f} pairs")