"""Build train JSON with top-3 AfriE5-similar same-subset (Q,A) demos prepended.
Source: msrh_rag_train_afrie5_TV_k3.json (36501 rows = Train+Val).
Per row: find top-3 same-subset OTHER (Q,A) pairs, exclude self by ID.
"""
import os, json, time, pathlib
os.environ["HF_HOME"] = "/mnt/msrh/Magic_submission/hf_cache"
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "4")
import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer

WORK = pathlib.Path("/mnt/msrh/Magic_submission")
TRAIN_CSV = WORK/"data/Train.csv"
VAL_CSV   = WORK/"data/Val.csv"
SRC_JSON  = WORK/"LF/data/msrh_rag_train_afrie5_TV_k3.json"
OUT_JSON  = WORK/"LF/data/msrh_rag_train_afrie5_TV_k3_fewshot.json"

train = pd.read_csv(TRAIN_CSV).reset_index(drop=True)
val   = pd.read_csv(VAL_CSV).reset_index(drop=True)
tv_pool = pd.concat([train, val], ignore_index=True)
print(f"TV pool size: {len(tv_pool)}")

MODEL = "/mnt/msrh/Magic_submission/hub/AfriE5-Large-instruct"
print(f"[{time.strftime('%H:%M:%S')}] loading {MODEL}")
model = SentenceTransformer(MODEL, device="cuda", trust_remote_code=True)
QUERY_INST = "Instruct: Given a maternal/sexual/reproductive health question, retrieve relevant passages that answer the question\nQuery: "

K = 3
demos_by_id = {}  # ID -> [(Q, A), ...]

for subset in sorted(tv_pool["subset"].unique()):
    pool_sub = tv_pool[tv_pool["subset"] == subset].reset_index(drop=True)
    print(f"[{time.strftime('%H:%M:%S')}] subset={subset} n={len(pool_sub)}")
    pool_emb = model.encode([QUERY_INST + str(q) for q in pool_sub["input"]],
                            batch_size=32, normalize_embeddings=True,
                            show_progress_bar=False, convert_to_numpy=True)
    sims = pool_emb @ pool_emb.T
    np.fill_diagonal(sims, -1.0)
    top_idx = np.argsort(-sims, axis=1)[:, :K]
    for i, row in pool_sub.iterrows():
        demos_by_id[row["ID"]] = [
            (str(pool_sub.iloc[int(j)]["input"]), str(pool_sub.iloc[int(j)]["output"]))
            for j in top_idx[i]
        ]
print(f"[{time.strftime('%H:%M:%S')}] demos built for {len(demos_by_id)} rows")

# The source train JSON has rows in the same order as tv_pool (train then val per build script)
n_inj = 0
n_skip = 0
with open(SRC_JSON) as fin, open(OUT_JSON, "w") as fout:
    for i, ln in enumerate(fin):
        r = json.loads(ln)
        pool_row = tv_pool.iloc[i]
        dem = demos_by_id.get(pool_row["ID"], [])
        if dem:
            content = r["messages"][0]["content"]
            idx = content.find("Context 1:")
            if idx > 0:
                before = content[:idx].rstrip()
                after = content[idx:]
                ex_block = "\n\n".join([f"Example {k+1}:\nQ: {q}\nA: {a}" for k, (q, a) in enumerate(dem)])
                new_content = f"{before}\n\nFew-shot examples (similar questions from training):\n{ex_block}\n\n{after}"
                r["messages"][0]["content"] = new_content
                n_inj += 1
            else:
                n_skip += 1
        else:
            n_skip += 1
        fout.write(json.dumps(r, ensure_ascii=False) + "\n")
print(f"\nInjected demos in {n_inj} rows, skipped {n_skip}")
print(f"wrote {OUT_JSON}")