"""Pair features for the LightGBM scorer. All features are country-agnostic."""
import multiprocessing as mp
import numpy as np
import pandas as pd
from rapidfuzz import process, fuzz
from rapidfuzz.distance import JaroWinkler, Levenshtein

from text_utils import abbrev_aware_match, token_jaccard, numbers


def _cp(a, b, scorer, workers):
    return process.cpdist(a, b, scorer=scorer, workers=workers).astype(np.float32)


_A = _B = None


def _py_feats(rng):
    lo, hi = rng
    out = np.empty((hi - lo, 6), np.float32)
    for j, i in enumerate(range(lo, hi)):
        a, b = _A[i], _B[i]
        out[j, 0] = abbrev_aware_match(a["name"], b["name"])
        out[j, 1] = abbrev_aware_match(a["addr"], b["addr"])
        out[j, 2] = token_jaccard(a["name"], b["name"])
        out[j, 3] = token_jaccard(a["addr"], b["addr"])
        na, nb = numbers(a["nums"]), numbers(b["nums"])
        out[j, 4] = len(na & nb) / len(na | nb) if (na or nb) else np.nan
        la = {x for x in na if len(x) >= 4}
        lb = {x for x in nb if len(x) >= 4}
        out[j, 5] = (1.0 if la & lb else 0.0) if (la and lb) else np.nan
    return out


def compute_features(cand, s1v, s23v, n_jobs=8):
    """cand: DataFrame with q (S1 row), c (S23 row) and blocking columns.
    s1v / s23v: dict of numpy arrays per column (name_core_abbr, address_abbrev_norm, ...)."""
    q, c = cand["q"].values, cand["c"].values
    f = pd.DataFrame(index=cand.index)
    na, nb = s1v["name_core_abbr"][q], s23v["name_core_abbr"][c]
    fa, fb = s1v["name_norm"][q], s23v["name_norm"][c]
    aa, ab = s1v["address_abbrev_norm"][q], s23v["address_abbrev_norm"][c]
    w = n_jobs

    f["name_jw"] = _cp(na, nb, JaroWinkler.normalized_similarity, w)
    f["name_lev"] = _cp(na, nb, Levenshtein.normalized_similarity, w)
    f["name_tsort"] = _cp(na, nb, fuzz.token_sort_ratio, w) / 100
    f["name_tset"] = _cp(na, nb, fuzz.token_set_ratio, w) / 100
    f["name_partial"] = _cp(na, nb, fuzz.partial_ratio, w) / 100
    f["fullname_ratio"] = _cp(fa, fb, fuzz.ratio, w) / 100
    f["addr_ratio"] = _cp(aa, ab, fuzz.ratio, w) / 100
    f["addr_tset"] = _cp(aa, ab, fuzz.token_set_ratio, w) / 100
    f["addr_partial"] = _cp(aa, ab, fuzz.partial_ratio, w) / 100

    global _A, _B
    _A = [{"name": x, "addr": y, "nums": z} for x, y, z in
          zip(na, aa, s1v["nums"][q])]
    _B = [{"name": x, "addr": y, "nums": z} for x, y, z in
          zip(nb, ab, s23v["nums"][c])]
    n = len(cand)
    step = max(1, n // (n_jobs * 8) + 1)
    ranges = [(i, min(i + step, n)) for i in range(0, n, step)]
    if n_jobs > 1 and n > 50_000 and "fork" in mp.get_all_start_methods():
        with mp.get_context("fork").Pool(n_jobs) as pool:
            parts = pool.map(_py_feats, ranges)
    else:
        parts = [_py_feats(r) for r in ranges]
    pf = np.concatenate(parts) if parts else np.empty((0, 6), np.float32)
    for j, name in enumerate(["name_abbrev", "addr_abbrev", "name_jacc", "addr_jacc",
                              "num_jacc", "longnum_match"]):
        f[name] = pf[:, j]
    _A = _B = None

    la, lb = s1v["legal_suffix"][q], s23v["legal_suffix"][c]
    f["suffix_both"] = ((la != "") & (lb != "")).astype(np.int8)
    f["suffix_eq"] = np.where(f["suffix_both"] == 1, (la == lb).astype(np.float32), np.nan)
    f["same_country"] = (s1v["country_norm"][q] == s23v["country_norm"][c]).astype(np.int8)
    f["name_len_a"] = np.fromiter((len(x) for x in na), np.int32, n)
    f["name_len_b"] = np.fromiter((len(x) for x in nb), np.int32, n)
    f["addr_empty_any"] = ((aa == "") | (ab == "")).astype(np.int8)
    f["is_s3"] = np.char.startswith(s23v["entity_id"][c].astype(str), "S3-").astype(np.int8)

    # blocking meta + within-query context
    # (raw blocking columns emb_score/emb_rank/lex_score/lex_rank/n_keys stay in `cand`)
    g = cand.groupby("q")
    f["emb_gap"] = cand["emb_score"].values - g["emb_score"].transform("max").values
    f["lex_gap"] = cand["lex_score"].values - g["lex_score"].transform("max").values
    f["n_cands"] = g["c"].transform("size").values
    f["name_jw_gap"] = f["name_jw"].values - f.groupby(cand["q"].values)["name_jw"].transform("max").values
    return f