"""Candidate generation: dense nearest neighbours  ∪  lexical inverted index.

Both channels are kept deliberately. Measured on 3,000 real S1 queries vs 756k
hard negatives from the actual candidate set:

    channel            R@20 overall   non-Latin name   empty address
    lexical / TF-IDF      0.927           0.727            0.574
    dense (e5-small)      0.961           0.910            0.534

and on the 220,683-query dev run against the full 10.3M pool:

    lexical only 0.9839 | dense only 0.9719 | UNION 0.9904

Dense alone is WORSE than lexical alone. Lexical wins on empty addresses, dense
wins hugely on non-Latin names. The union is what scores -- do not drop either.
"""
from __future__ import annotations

import multiprocessing as mp
import numpy as np
import pandas as pd

from text_utils import tokens


def record_text(name, addr, country):
    return f"{name} | {addr} | {country}"


# --------------------------------------------------------------------------- dense
def dense_topk(q_emb, c_emb, cfg, device):
    """Exact top-k cosine over a (possibly memory-mapped) corpus matrix.

    corpus_on_gpu=False streams each corpus shard to the GPU and frees it. That is
    ~15% slower than pinning the whole matrix but it is why this survives a shared
    GPU: model-2 pinned 21 GB and died when another user's job filled the card.
    """
    import torch
    k = cfg["emb_k"]
    dtype = torch.float16 if str(device).startswith("cuda") else torch.float32
    shard = cfg["emb_shard"]
    pin = cfg["emb_corpus_on_gpu"]
    nq = q_emb.shape[0]
    best_s = np.full((nq, k), -1.0, np.float32)
    best_i = np.full((nq, k), -1, np.int64)

    resident = []
    if pin:
        for s in range(0, c_emb.shape[0], shard):
            resident.append((s, torch.from_numpy(np.ascontiguousarray(c_emb[s:s + shard]))
                             .to(dtype).to(device)))

    for si, s in enumerate(range(0, c_emb.shape[0], shard)):
        if pin:
            off, X = resident[si]
        else:
            off = s
            X = torch.from_numpy(np.ascontiguousarray(c_emb[s:s + shard])).to(dtype).to(device)
        XT = X.T.contiguous()
        for qs in range(0, nq, cfg["emb_query_batch"]):
            q = torch.from_numpy(np.ascontiguousarray(q_emb[qs:qs + cfg["emb_query_batch"]])) \
                .to(dtype).to(device)
            sc = q @ XT
            kk = min(k, sc.shape[1])
            v, i = torch.topk(sc, kk, dim=1)
            v = v.float().cpu().numpy()
            i = i.cpu().numpy().astype(np.int64) + off
            sl = slice(qs, qs + len(v))
            cs = np.concatenate([best_s[sl], v], 1)
            ci = np.concatenate([best_i[sl], i], 1)
            take = np.argsort(-cs, axis=1)[:, :k]
            best_s[sl] = np.take_along_axis(cs, take, 1)
            best_i[sl] = np.take_along_axis(ci, take, 1)
            del sc, q
        if not pin:
            del X, XT
            if str(device).startswith("cuda"):
                torch.cuda.empty_cache()
    if pin:
        del resident
        if str(device).startswith("cuda"):
            torch.cuda.empty_cache()
    return best_s, best_i


# --------------------------------------------------------------------------- lexical
def lexical_tokens(name, addr):
    out = []
    for t in set(tokens(name)):
        out.append("n:" + t)
        if len(t) >= 5:
            out.append("p:" + t[:4])
    for t in set(tokens(addr)):
        out.append("a:" + t)
    return list(set(out))


class LexicalIndex:
    def __init__(self, doc_tokens, cfg):
        n = len(doc_tokens)
        lens = np.fromiter((len(x) for x in doc_tokens), dtype=np.int64, count=n)
        docs = np.repeat(np.arange(n, dtype=np.int32), lens)
        toks = pd.Series([t for x in doc_tokens for t in x], dtype="object")
        codes, vocab = pd.factorize(toks)
        df = np.bincount(codes)
        keep = df <= cfg["lex_max_df"]
        mask = keep[codes]
        codes, docs = codes[mask], docs[mask]
        order = np.argsort(codes, kind="stable")
        self.docs = docs[order]
        counts = np.bincount(codes, minlength=len(vocab))
        self.offs = np.concatenate([[0], np.cumsum(counts)])
        self.vocab = {t: i for i, t in enumerate(vocab) if keep[i]}
        idf = np.log(n / np.maximum(df, 1))
        w = np.ones(len(vocab))
        for t, i in self.vocab.items():
            if t.startswith("a:"):
                w[i] = cfg["lex_addr_weight"]
            elif t.startswith("p:"):
                w[i] = cfg["lex_prefix_weight"]
        self.w = idf * w
        self.k = cfg["lex_k"]
        self.df = df
        self.budget = cfg.get("lex_post_budget", 0)

    def query(self, qtoks):
        codes = [self.vocab[t] for t in qtoks if t in self.vocab]
        if not codes:
            return np.empty(0, np.int32), np.empty(0, np.float32)
        codes = np.array(codes)
        if self.budget:
            # Spend the posting budget on the rarest (most informative) tokens
            # first. Without this, raising lex_max_df to keep the address
            # skeleton would make a query like "a:no" scan 1.9M postings.
            order = np.argsort(self.df[codes], kind="stable")
            codes = codes[order]
            cum = np.cumsum(self.offs[codes + 1] - self.offs[codes])
            keep = int(np.searchsorted(cum, self.budget, side="right"))
            codes = codes[:max(keep, 1)]
        starts, ends = self.offs[codes], self.offs[codes + 1]
        lens = ends - starts
        if lens.sum() == 0:
            return np.empty(0, np.int32), np.empty(0, np.float32)
        cand = np.concatenate([self.docs[s:e] for s, e in zip(starts, ends)])
        ws = np.repeat(self.w[codes], lens)
        uniq, inv = np.unique(cand, return_inverse=True)
        sc = np.bincount(inv, weights=ws)
        if len(uniq) > self.k:
            top = np.argpartition(-sc, self.k)[:self.k]
            uniq, sc = uniq[top], sc[top]
        o = np.argsort(-sc)
        return uniq[o], sc[o].astype(np.float32)


_IDX = None


def _lex_worker(chunk):
    return [_IDX.query(q) for q in chunk]


def lexical_topk(index, q_tokens, n_jobs):
    global _IDX
    _IDX = index
    chunks = [q_tokens[i:i + 2000] for i in range(0, len(q_tokens), 2000)]
    if n_jobs > 1 and "fork" in mp.get_all_start_methods():
        with mp.get_context("fork").Pool(n_jobs) as pool:
            res = pool.map(_lex_worker, chunks)
    else:
        res = [_lex_worker(c) for c in chunks]
    return [r for c in res for r in c]


# --------------------------------------------------------------------------- union
def build_candidates(q_rows, emb_s, emb_i, lex_res, q_country, c_country, country_filter,
                     emb2_s=None, emb2_i=None):
    """-> DataFrame [q, c, emb_score, emb_rank, lex_score, lex_rank, (emb2_*), n_keys]."""
    nq, k = emb_i.shape
    e = pd.DataFrame({
        "qpos": np.repeat(np.arange(nq), k),
        "c": emb_i.ravel().astype(np.int64),
        "emb_score": emb_s.ravel(),
        "emb_rank": np.tile(np.arange(k), nq),
    })
    e = e[e["c"] >= 0]
    lq, lc, ls, lr = [], [], [], []
    for i, (c, s) in enumerate(lex_res):
        lq.append(np.full(len(c), i)); lc.append(c); ls.append(s); lr.append(np.arange(len(c)))
    l = pd.DataFrame({
        "qpos": np.concatenate(lq) if lq else np.empty(0, int),
        "c": np.concatenate(lc).astype(np.int64) if lc else np.empty(0, np.int64),
        "lex_score": np.concatenate(ls) if ls else np.empty(0),
        "lex_rank": np.concatenate(lr) if lr else np.empty(0, int),
    })
    cand = e.merge(l, on=["qpos", "c"], how="outer")
    n_keys = cand["emb_score"].notna().astype(int) + cand["lex_score"].notna().astype(int)
    if emb2_i is not None:
        nq2, k2 = emb2_i.shape
        e2 = pd.DataFrame({
            "qpos": np.repeat(np.arange(nq2), k2),
            "c": emb2_i.ravel().astype(np.int64),
            "emb2_score": emb2_s.ravel(),
            "emb2_rank": np.tile(np.arange(k2), nq2),
        })
        e2 = e2[e2["c"] >= 0]
        cand = cand.merge(e2, on=["qpos", "c"], how="outer")
        n_keys = n_keys.reindex(cand.index).fillna(0).astype(int) \
            + cand["emb2_score"].notna().astype(int)
        # left NaN on purpose where e5-large did not retrieve it: LightGBM treats
        # that as missing, which is exactly the signal "the other encoder disagreed"
        cand["emb2_rank"] = cand["emb2_rank"].fillna(99)
    cand["n_keys"] = n_keys
    cand["emb_rank"] = cand["emb_rank"].fillna(99)
    cand["lex_rank"] = cand["lex_rank"].fillna(99)
    cand["lex_score"] = cand["lex_score"].fillna(0.0)
    cand["q"] = np.asarray(q_rows)[cand["qpos"].values]
    if country_filter:
        keep = q_country[cand["q"].values] == c_country[cand["c"].values]
        cand = cand[keep]
    return cand.drop(columns="qpos").reset_index(drop=True)


def fill_missing_emb(cand, q_emb_by_row, c_emb):
    """Cosine for candidates found only by the lexical channel."""
    m = cand["emb_score"].isna().values
    if m.any():
        qe = q_emb_by_row(cand["q"].values[m]).astype(np.float32)
        ce = np.asarray(c_emb[cand["c"].values[m]]).astype(np.float32)
        cand.loc[m, "emb_score"] = (qe * ce).sum(1)
    return cand