"""CPU stage: reciprocal-rank fusion of e5 / bge / lexical candidates. RRF only decides which candidates survive
the cap; every retriever's own score/rank/gap is kept as a separate feature. Train side also runs the
recall-vs-K probe that picks the cap."""
from multiprocessing import Pool

import numpy as np
import pandas as pd

from common import Progress, get_logger, n_workers, save_json, save_parquet, shard_path, load_json
from config import CFG, EMB_MODELS, WORK
from decide import in_sorted
from prep import load_prep, load_splits
from retrieve_dense import n_shards, queries

log = get_logger("block")
RETR = list(EMB_MODELS) + ["lex"]
G = {}


def _load(r, split, i):
    p = shard_path("lex", split, i, "npz") if r == "lex" else shard_path("dense", split, i, "npz", r)
    z = np.load(p)
    return z["cand"], z["score"].astype(np.float32)


def _block_shard(i):
    split, Q = G["split"], G["Q"]
    lo = i * CFG["query_chunk"]
    frames = []
    for r in RETR:
        cand, score = _load(r, split, i)
        nq, K = cand.shape
        f = pd.DataFrame({"qi": np.repeat(np.arange(nq), K), "c": cand.ravel(),
                          f"{r}_r": np.tile(np.arange(K), nq), f"{r}_s": score.ravel()})
        frames.append(f[f.c >= 0])
    df = frames[0]
    for f in frames[1:]:
        df = df.merge(f, on=["qi", "c"], how="outer")
    k0 = CFG["rrf_k0"]
    df["rrf"] = sum((1.0 / (k0 + 1 + df[f"{r}_r"])).fillna(0) for r in RETR)
    df["support"] = sum(df[f"{r}_r"].notna().astype(np.int8) for r in RETR).astype(np.int8)
    df = df.sort_values(["qi", "rrf"], ascending=[True, False], kind="stable")
    df["rrf_rank"] = df.groupby("qi").cumcount().astype(np.int16)
    df = df[df.rrf_rank < CFG["k_max"]]
    for r in RETR:
        df[f"{r}_gap"] = df.groupby("qi")[f"{r}_s"].transform("max") - df[f"{r}_s"]
        df[f"{r}_r"] = df[f"{r}_r"].fillna(CFG["dense_k"]).astype(np.int16)
    df["q"] = Q[lo + df.qi.values].astype(np.int32)
    df["c"] = df.c.astype(np.int32)
    if split == "train":
        key = df.q.values.astype(np.int64) * G["mult"] + df.c.values
        df["y"] = in_sorted(key, G["gt_key"]).astype(np.int8)
    else:
        df["y"] = np.int8(-1)
    cols = ["q", "c", "rrf", "rrf_rank", "support", "y"] + [f"{r}_{x}" for r in RETR for x in ("r", "s", "gap")]
    for c in cols:
        if df[c].dtype == np.float64:
            df[c] = df[c].astype(np.float32)
    save_parquet(df[cols].reset_index(drop=True), shard_path("block", split, i))
    return i


def run_block(split):
    ns = n_shards(split)
    Q = queries(split)
    G.update(split=split, Q=Q)
    if split == "train":
        sp = load_splits()
        mult = np.int64(1 << 32)
        G["mult"] = mult
        G["gt_key"] = np.sort(sp["gt_q"] * mult + sp["gt_c"])
    todo = [i for i in range(ns) if not shard_path("block", split, i).exists()]
    if todo:
        prog = Progress(log, len(todo), f"block {split} shards")
        with Pool(min(n_workers(), 16)) as pool:
            for _ in pool.imap_unordered(_block_shard, todo):
                prog.step()
    if split == "train":
        recall_probe(ns)


def recall_probe(ns):
    sp = load_splits()
    n_true = sp["n_true"]
    evalq = np.concatenate([sp["g_stack"], sp["g_val"]])
    country = load_prep("train", "s1", ["country"]).country.values
    df = pd.concat([pd.read_parquet(shard_path("block", "train", i),
                                    columns=["q", "rrf_rank", "y"] + [f"{r}_r" for r in RETR]) for i in range(ns)])
    df = df[np.isin(df.q.values, evalq) & (df.y == 1)]
    tot = int(n_true[evalq].sum())
    ctry = pd.Series(country[evalq]).value_counts()
    tot_c = pd.Series(n_true[evalq]).groupby(country[evalq]).sum()
    rep = {"true_pairs": tot, "recall_at_K": {}, "recall_at_K_by_country": {}, "per_retriever_recall@100": {}}
    for K in CFG["k_grid"]:
        m = df.rrf_rank < K
        rep["recall_at_K"][K] = float(m.sum() / tot)
        rep["recall_at_K_by_country"][K] = (df[m].groupby(country[df[m].q.values]).size() / tot_c).round(4).to_dict()
    for r in RETR:
        rep["per_retriever_recall@100"][r] = float((df[f"{r}_r"] < CFG["dense_k"]).sum() / tot)
    full = rep["recall_at_K"][CFG["k_max"]]
    K = next(k for k in CFG["k_grid"] if rep["recall_at_K"][k] >= CFG["k_recall_frac"] * full)
    rep["chosen_K"] = K
    save_json(rep, WORK / "reports" / "blocking_recall.json")
    save_json({"K": int(K)}, WORK / "reports" / "k_choice.json")
    log.info(f"blocking recall: {rep}")