"""Phase 1.4 — cross-banner product matching with confidence tiers.

Tiers (calibrated on UPC ground truth Metro<->Super C, src/matching/calibrate_thresholds.py, then frozen):
  A  exact GS1-valid GTIN equality (only feeds whose external_id is a UPC: metro, superc — 100 % check-digit valid)
  B  brand key equal + package size equal (0.1-unit bucket) + identical cleaned token sets + identical numeric tokens
     (labelled-set precision 0.913, recall 0.984)
  C  brand key equal + size equal + token_sort_ratio >= 0.90 + identical numeric tokens, not already B
     (labelled-set precision 0.845 cumulative)
  D  (brand missing on one side) size equal + identical full-name token sets (brand tokens kept on both sides) +
     identical numeric tokens + brand tokens of the branded side present in the brandless name
     (validation round 1 showed that a 0.92 sort-ratio rule let 'spaghetti' ~ 'spaghettini' through: 3/60 errors)
Excluded from matching: listings without a parsable package size (circular items are mostly size-less) and
multi-product circular offers ("X ou Y", "assortis", "sélectionnés", comma lists). A first version using
rapidfuzz token_set_ratio was rejected because a name that is a token-subset of another scores 1.0
("beurre végétal salé" ~ "beurre végétal non salé"); see data/processed/match_validation_sample_v0_tokenset_rejected.csv.
Pairs are formed only across different banners (same company allowed, flagged `same_company`); same-banner links
between the storefront feed and the circular feed are recorded separately (`link_type = 'channel'`).
Outputs: data/processed/product_matches.parquet (pairs), match_groups_{tier}.parquet (union-find groups), and a
stratified validation sample (data/processed/match_validation_sample.csv) for manual + LLM annotation.
"""
from __future__ import annotations
import re, unicodedata, itertools, hashlib
import numpy as np, pandas as pd
from rapidfuzz import fuzz, process
from src.common import path, get_logger, ledger, seed_everything, SEED, SETTINGS
from src.cleaning.units import size_residual

log = get_logger("match")
_ALNUM = re.compile(r"[^a-z0-9]+")
_SIZE_TOK = re.compile(r"\b\d+(?:[.,]\d+)?\s*(?:x|×)?\s*\d*(?:[.,]\d+)?\s*(?:kg|g|mg|lb|lbs|oz|ml|cl|l|un|unites?|ea|each|pk|ch)\b")
_NUM = re.compile(r"\d+(?:[.,]\d+)?")
STOP = {"de","du","des","la","le","les","l","d","et","a","au","aux","en","pour","the","of","and","with","un","une",
        "sans","avec","ou","or","pc","chaque","each","marque","brand"}
GTIN_FEEDS = {"metro", "superc"}
MULTI_RE = re.compile(r"\b(ou|or|assortis?|assorted|selectionne?s?|varietes?|au choix|et plus|et autres)\b|,|/", re.I)

def strip_accents(s: str) -> str:
    return "".join(c for c in unicodedata.normalize("NFKD", s or "") if not unicodedata.combining(c))

def brand_key(b: str) -> str:
    k = _ALNUM.sub("", strip_accents(b or "").lower())
    return k if len(k) >= 2 else ""

def name_key(name: str, brand: str = "") -> str:
    s = strip_accents(name or "").lower().replace(" ", " ")
    s = _SIZE_TOK.sub(" ", s); s = _ALNUM.sub(" ", s)
    drop = set(STOP)
    if brand:
        drop |= set(_ALNUM.sub(" ", strip_accents(brand).lower()).split())
    toks = [t for t in s.split() if t and t not in drop]
    return " ".join(sorted(toks))

def num_tokens(nk: str) -> frozenset:
    return frozenset(x.replace(",", ".") for x in _NUM.findall(nk))

def size_bucket(qty: float, base: str, pack=1) -> str:
    """Package key: total quantity in base unit AND multipack count ('12x355 ml' != '6x710 ml')."""
    if pd.isna(qty) or base is None: return ""
    pk = 1 if pd.isna(pack) else int(pack)
    return f"{round(float(qty), 1):g}|{base}|x{pk}"

def prepare(prod: pd.DataFrame, ret: pd.DataFrame) -> pd.DataFrame:
    excl = set(SETTINGS["scope"]["exclude_sources"])
    p = prod[~prod.source.isin(excl)].copy()
    p = p[p.name.fillna("").str.len() > 0]
    p = p.merge(ret[["feed_id", "banner", "company", "channel"]], left_on="source", right_on="feed_id", how="left")
    p["bkey"] = p.brand.fillna("").map(brand_key)
    resid = p.size_label.fillna("").map(size_residual)
    p["nkey"] = [(name_key(n, b) + " " + r).strip() for n, b, r in zip(p.name.fillna(""), p.brand.fillna(""), resid)]
    p["nkey"] = p.nkey.map(lambda k: " ".join(sorted(k.split())))
    p["nkey_full"] = [(name_key(n, "") + " " + r).strip() for n, r in zip(p.name.fillna(""), resid)]
    p["nkey_full"] = p.nkey_full.map(lambda k: " ".join(sorted(k.split())))
    p["ntoks"] = p.nkey.map(num_tokens)
    p["skey"] = [size_bucket(q, b, k) for q, b, k in zip(p.size_qty, p.size_base, p.size_pack)]
    p["gtin"] = np.where(p.source.isin(GTIN_FEEDS) & p.gtin_valid, p.external_id.str.lstrip("0"), None)
    p["multi_offer"] = p.name.fillna("").map(lambda n: bool(MULTI_RE.search(strip_accents(n))))
    return p

def tier_a(p: pd.DataFrame) -> pd.DataFrame:
    g = p[p.gtin.notna()]
    pairs = []
    for gt, grp in g.groupby("gtin"):
        if len(grp) < 2: continue
        for i, j in itertools.combinations(grp.index, 2):
            if grp.at[i, "banner"] == grp.at[j, "banner"] and grp.at[i, "channel"] == grp.at[j, "channel"]: continue
            pairs.append((grp.at[i, "uid"], grp.at[j, "uid"], "A", 1.0, True, True, 1.0))
    return pd.DataFrame(pairs, columns=["uid_a", "uid_b", "tier", "score", "brand_match", "size_match", "name_sim"])

def _block_pairs(block: pd.DataFrame, thr: float) -> list:
    out = []
    if len(block) < 2 or len(block) > 4000: return out
    names = block.nkey.tolist(); uids = block.uid.tolist(); banners = block.banner.tolist(); chans = block.channel.tolist()
    ntoks = block.ntoks.tolist(); tsets = [frozenset(n.split()) for n in names]
    sim = process.cdist(names, names, scorer=fuzz.token_sort_ratio, workers=-1) / 100.0
    iu = np.triu_indices(len(block), k=1)
    for i, j in zip(*iu):
        s = sim[i, j]
        if s < thr: continue
        if ntoks[i] != ntoks[j]: continue
        if banners[i] == banners[j] and chans[i] == chans[j]: continue
        out.append((uids[i], uids[j], float(s), tsets[i] == tsets[j]))
    return out

def tier_bc(p: pd.DataFrame) -> pd.DataFrame:
    rows = []
    q = p[(p.bkey != "") & (p.skey != "") & (~p.multi_offer)]
    for (bk, sk), block in q.groupby(["bkey", "skey"]):
        for ua, ub, s, eq in _block_pairs(block, 0.90):
            rows.append((ua, ub, "B" if eq else "C", s, True, True, s))
    return pd.DataFrame(rows, columns=["uid_a", "uid_b", "tier", "score", "brand_match", "size_match", "name_sim"])

def tier_d(p: pd.DataFrame) -> pd.DataFrame:
    """Brandless feeds: block on size, compare full names (brand kept in the name), require the branded side's brand
    tokens to appear in the brandless name."""
    rows = []
    q = p[(p.skey != "") & (~p.multi_offer)]
    for sk, block in q.groupby("skey"):
        bl = block[block.bkey == ""]; br = block[block.bkey != ""]
        if len(bl) == 0 or len(br) == 0 or len(block) > 6000: continue
        sim = process.cdist(bl.nkey_full.tolist(), br.nkey_full.tolist(), scorer=fuzz.token_sort_ratio, workers=-1) / 100.0
        ii, jj = np.where(sim >= 0.92)
        for i, j in zip(ii, jj):
            a, b = bl.iloc[i], br.iloc[j]
            if a.banner == b.banner and a.channel == b.channel: continue
            if set(a.nkey_full.split()) != set(b.nkey_full.split()): continue
            if num_tokens(a.nkey_full) != num_tokens(b.nkey_full): continue
            btoks = [t for t in _ALNUM.sub(" ", strip_accents(b.brand).lower()).split() if len(t) >= 3]
            if btoks and not all(t in a.nkey_full.split() for t in btoks): continue
            rows.append((a.uid, b.uid, "D", float(sim[i, j]), False, True, float(sim[i, j])))
    return pd.DataFrame(rows, columns=["uid_a", "uid_b", "tier", "score", "brand_match", "size_match", "name_sim"])

def union_groups(pairs: pd.DataFrame, tiers: set[str]) -> pd.DataFrame:
    parent = {}
    def find(x):
        while parent.setdefault(x, x) != x:
            parent[x] = parent[parent[x]]; x = parent[x]
        return x
    for ua, ub in pairs.loc[pairs.tier.isin(tiers), ["uid_a", "uid_b"]].itertuples(index=False):
        ra, rb = find(ua), find(ub)
        if ra != rb: parent[ra] = rb
    rows = [(hashlib.md5(find(u).encode()).hexdigest()[:12], u) for u in parent]
    return pd.DataFrame(rows, columns=["group_id", "uid"])

def main() -> None:
    seed_everything()
    prod = pd.read_parquet(path("processed") / "products.parquet")
    ret = pd.read_parquet(path("processed") / "retailers.parquet")
    p = prepare(prod, ret)
    log.info("matching universe: %d listings (active+inactive, excl. %s); with size %d; multi-product offers excluded %d",
             len(p), SETTINGS["scope"]["exclude_sources"], (p.skey != "").sum(), p.multi_offer.sum())
    parts = [tier_a(p), tier_bc(p), tier_d(p)]
    pairs = pd.concat(parts, ignore_index=True)
    # keep the best tier per unordered pair
    pairs["key"] = [tuple(sorted(x)) for x in zip(pairs.uid_a, pairs.uid_b)]
    order = {"A": 0, "B": 1, "C": 2, "D": 3}
    pairs["rank"] = pairs.tier.map(order)
    pairs = pairs.sort_values(["key", "rank"]).drop_duplicates("key").drop(columns=["key", "rank"])
    meta = p.set_index("uid")[["banner", "company", "channel", "source", "category", "active"]]
    for side in ("a", "b"):
        m = meta.loc[pairs[f"uid_{side}"]].reset_index(drop=True); m.columns = [f"{c}_{side}" for c in m.columns]
        pairs = pd.concat([pairs.reset_index(drop=True), m], axis=1)
    pairs["same_company"] = pairs.company_a == pairs.company_b
    pairs["link_type"] = np.where(pairs.banner_a == pairs.banner_b, "channel", "cross_banner")
    pairs["both_active"] = (pairs.active_a == 1) & (pairs.active_b == 1)
    pairs.to_parquet(path("processed") / "product_matches.parquet", index=False)
    log.info("pairs: %s ; cross-banner by tier: %s", len(pairs),
             pairs[pairs.link_type == "cross_banner"].tier.value_counts().to_dict())
    for name, tiers in {"A": {"A"}, "AB": {"A", "B"}, "ABC": {"A", "B", "C"}, "ABCD": {"A", "B", "C", "D"}}.items():
        g = union_groups(pairs[pairs.link_type == "cross_banner"], tiers)
        g = g.merge(meta[["banner"]], left_on="uid", right_index=True)
        nb = g.groupby("group_id").banner.nunique()
        g = g[g.group_id.isin(nb[nb >= 2].index)]
        g.to_parquet(path("processed") / f"match_groups_{name}.parquet", index=False)
        log.info("groups tiers %-4s: %d groups, %d listings, median banners/group %.0f, max %d", name,
                 g.group_id.nunique(), len(g), nb[nb >= 2].median(), nb.max())
    # Food-Ka's own links for comparison
    fk = pd.read_parquet(path("interim") / "raw_product_links.parquet")
    fkpairs = set()
    for gid, grp in fk.groupby("group_id"):
        for a, b in itertools.combinations(sorted(grp.uid), 2): fkpairs.add((a, b))
    ours = set(tuple(sorted(x)) for x in zip(pairs.uid_a, pairs.uid_b))
    log.info("Food-Ka product_links pairs: %d ; recovered by our tiers A-D: %d (%.1f%%)", len(fkpairs),
             len(fkpairs & ours), 100 * len(fkpairs & ours) / max(len(fkpairs), 1))
    # stratified validation sample: 40 per tier (A,B,C,D) × link_type cross_banner, with listing text for annotators
    rng = np.random.default_rng(SEED)
    cb = pairs[pairs.link_type == "cross_banner"]
    samp = cb.groupby("tier", group_keys=False).apply(lambda g: g.sample(min(len(g), 60), random_state=SEED))
    txt = prod.set_index("uid")[["name", "brand", "size_label", "price", "source"]]
    for side in ("a", "b"):
        t = txt.loc[samp[f"uid_{side}"]].reset_index(drop=True); t.columns = [f"{c}_{side}_txt" for c in t.columns]
        samp = pd.concat([samp.reset_index(drop=True), t], axis=1)
    samp = samp.sample(frac=1, random_state=SEED).reset_index(drop=True)
    samp.insert(0, "sample_id", range(1, len(samp) + 1))
    samp["label_llm"] = ""; samp["label_human"] = ""; samp["note"] = ""
    samp.to_csv(path("processed") / "match_validation_sample.csv", index=False)
    log.info("validation sample: %d pairs (%s)", len(samp), samp.tier.value_counts().to_dict())
    ledger("1", "match_products", f"{len(pairs)} pairs; sample {len(samp)}")

if __name__ == "__main__":
    main()