"""Can NAME DISTINCTIVENESS rescue the empty-address pairs?

THE FINDING THIS FOLLOWS UP. Of our 1,763 retrieved-but-rejected true pairs, 76.57% have an EMPTY
ADDRESS (against a 3.34% pool base rate - a 23x enrichment). Of the 1,211 never-retrieved, 42.36%.
Together 1,863 pairs = 63% of all 2,974 misses, concentrated in the 344,883 pool records (3.34%)
that carry no address.

That explains every previous negative result at once: both cross-encoders reject them because
`addr: ` is blank and there is nothing to compare, not because they lack capacity; they are shorter
(58 vs 70 tokens) because the field is empty; and corroboration failed because an address-less
record resembles everything.

THE QUESTION. With no address only the name can decide. Name-only matching is precision-poor in
general, but not uniformly: 'Northwind' -> 'Northwind' is obvious, 'Ruby Academy' is not. Does name
DISTINCTIVENESS (rare-token IDF) separate true from false among empty-address candidates well
enough to clear the ~3.76:1 break-even?

Yes -> a feature/decision change: cheap, no GPU, no retraining.
No  -> the empty-address pairs are genuinely unrecoverable, and 0.9821 plus the retrieval ceiling
       is close to this data's real limit.
"""
import collections
import json
import math
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 norm(s):
    s = "".join(c for c in unicodedata.normalize("NFKD", str(s)) if not unicodedata.combining(c))
    return _WS.sub(" ", _PUN.sub(" ", s.casefold())).strip()


def toks(s):
    return [t for t in norm(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"]
s2 = pd.read_parquet(p["source2"], columns=["entity_id"])
s3 = pd.read_parquet(p["source3"], columns=["entity_id"])
pool_ids = np.concatenate([s2["entity_id"].values, s3["entity_id"].values])
s1_ids = pd.read_parquet(p["source1"], columns=["entity_id"])["entity_id"].values

praw = {}
df_counts = collections.Counter()
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 in zip(d["entity_id"], d["business_name"], d["business_address"]):
        praw[e] = (n, a)
        df_counts.update(set(toks(n)))
N_POOL = len(praw)
idf = {t: math.log(N_POOL / (1 + c)) for t, c in df_counts.items()}
print(f"pool records {N_POOL:,}   distinct name tokens {len(idf):,}")

s1raw = pd.read_csv(os.path.join(RAW, "train_source1.tsv"), sep="\t", dtype=str,
                    keep_default_na=False, na_filter=False)
s1map = dict(zip(s1raw["entity_id"], s1raw["business_name"]))

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)
sel_pairs = {(q, c) for q, cs in shipped.items() for c in cs}
sel = np.array([1 if (q, c) in sel_pairs else 0 for q, c in zip(qv, cv)], np.int8)

emptyaddr = np.array([not praw[pool_ids[c]][1].strip() for c in cv], bool)
print(f"\nval candidate pairs {len(cand):,}   with empty-address candidate "
      f"{int(emptyaddr.sum()):,} ({emptyaddr.mean():.2%})")

# name features on the empty-address subset
rows = np.flatnonzero(emptyaddr)
feat = np.zeros((len(rows), 5), np.float32)
for i, r in enumerate(rows):
    qn = s1map[s1_ids[qv[r]]]
    cn = praw[pool_ids[cv[r]]][0]
    tq, tc = set(toks(qn)), set(toks(cn))
    inter = tq & tc
    feat[i, 0] = 1.0 if norm(qn) == norm(cn) else 0.0                     # exact after norm
    feat[i, 1] = len(inter) / max(len(tq | tc), 1)                        # token jaccard
    feat[i, 2] = sum(idf.get(t, 0.0) for t in inter)                      # shared IDF mass
    feat[i, 3] = max((idf.get(t, 0.0) for t in inter), default=0.0)       # rarest shared token
    feat[i, 4] = len(inter)
F = pd.DataFrame(feat, columns=["exact", "jac", "idf_sum", "idf_max", "n_shared"])
F["y"] = y[rows]
F["sel"] = sel[rows]
F["p"] = pv[rows]
F["q"] = qv[rows]
F["c"] = cv[rows]

print("\n" + "=" * 74)
print("EMPTY-ADDRESS SUBSET: name-feature separation, TRUE vs FALSE")
print("=" * 74)
t_, f_ = F[F.y == 1], F[F.y == 0]
print(f"  true pairs {len(t_):,}   false pairs {len(f_):,}   "
      f"base rate {len(t_)/max(len(F),1):.4%}")
print(f"{'feature':<12}{'true median':>14}{'false median':>14}{'true p90':>11}{'false p90':>11}")
for c in ("exact", "jac", "idf_sum", "idf_max", "n_shared"):
    print(f"  {c:<10}{t_[c].median():>14.4f}{f_[c].median():>14.4f}"
          f"{np.percentile(t_[c],90):>11.4f}{np.percentile(f_[c],90):>11.4f}")

print("\n" + "=" * 74)
print("EXCHANGE RATE of accepting REJECTED empty-address pairs by a name rule")
print("=" * 74)
un = F[F.sel == 0]
print(f"  currently rejected empty-address pairs: {len(un):,} "
      f"({int((un.y == 1).sum()):,} true, {int((un.y == 0).sum()):,} false)")
print(f"{'rule':<44}{'TP':>7}{'FP':>8}{'ratio':>8}")
rules = [
    ("exact normalised name match", un.exact == 1),
    ("exact name AND idf_max>=8", (un.exact == 1) & (un.idf_max >= 8)),
    ("exact name AND idf_max>=10", (un.exact == 1) & (un.idf_max >= 10)),
    ("exact name AND idf_max>=12", (un.exact == 1) & (un.idf_max >= 12)),
    ("jac>=0.8 AND idf_max>=10", (un.jac >= 0.8) & (un.idf_max >= 10)),
    ("jac>=0.6 AND idf_max>=12", (un.jac >= 0.6) & (un.idf_max >= 12)),
    ("idf_sum>=20", un.idf_sum >= 20),
    ("idf_sum>=25 AND jac>=0.5", (un.idf_sum >= 25) & (un.jac >= 0.5)),
    ("exact name AND p>=0.2", (un.exact == 1) & (un.p >= 0.2)),
    ("exact name AND idf_max>=10 AND p>=0.1", (un.exact == 1) & (un.idf_max >= 10) & (un.p >= 0.1)),
]
best = None
for name, mask in rules:
    tp = int(((mask) & (un.y == 1)).sum())
    fp = int(((mask) & (un.y == 0)).sum())
    if tp + fp == 0:
        print(f"  {name:<42}{0:>7}{0:>8}{'-':>8}")
        continue
    r = tp / fp if fp else float("inf")
    flag = "  <== CLEARS 3.76" if r > 3.76 else ""
    print(f"  {name:<42}{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" + "=" * 74)
print("END TO END, held out on half the entities")
print("=" * 74)
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("\nNo rule clears the break-even. The empty-address pairs are not recoverable by")
    print("name distinctiveness alone, which means the 1,763 rejections really are close to")
    print("this data's floor - the address is simply absent.")
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[int(q)].add(int(c))
    pr = {q: list(v) for q, v in newpred.items()}
    a, bb = ev(pr, HAl), ev(pr, HB)
    print(f"\nbest rule: {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")
    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):,}")