"""Few-shot (q, a) demonstration retrieval (Deep Past 1st pattern).

For each test row:
- Embed test_q with AfriE5
- Find top-3 most-similar train (Q, A) pairs (by Q similarity, SAME-SUBSET ONLY)
- Prepend as "Example 1: Q: ...\nA: ...\n\nExample 2: ...\n\n..." before original prompt
"""
import os, json, csv, pathlib, sys, time
os.environ["HF_HOME"] = "/mnt/msrh/Magic_submission/hf_cache"
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "4")  # use GPU 4 to avoid pre_rationale on 0-3

import numpy as np
import torch
from sentence_transformers import SentenceTransformer

WORK = pathlib.Path("/mnt/msrh/Magic_submission")
TRAIN_CSV = WORK/"data/Train.csv"
VAL_CSV   = WORK/"data/Val.csv"
TEST_CSV  = WORK/"data/Test.csv"
TEST_ORIG = WORK/"LF/data/msrh_rag_test_k3_AfriE5_TV.json"
OUT_DIR   = WORK/"LF/data"   # Direct write to LF/data/ — output is fewshot test json
OUT_DIR.mkdir(exist_ok=True, parents=True)

# Read train+val (TV pool) and test
import pandas as pd
train = pd.read_csv(TRAIN_CSV).reset_index(drop=True)
val   = pd.read_csv(VAL_CSV).reset_index(drop=True)
test  = pd.read_csv(TEST_CSV).reset_index(drop=True)
tv_pool = pd.concat([train, val], ignore_index=True)
print(f"TV pool={len(tv_pool)} Test={len(test)}")

MODEL = "/mnt/msrh/Magic_submission/hub/AfriE5-Large-instruct"
print(f"[{time.strftime('%H:%M:%S')}] loading {MODEL} on GPU {os.environ.get('CUDA_VISIBLE_DEVICES')}")
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: "

def encode_q(qs):
    qs = [QUERY_INST + str(q) for q in qs]
    return model.encode(qs, batch_size=32, normalize_embeddings=True, show_progress_bar=False, convert_to_numpy=True)
def encode_p(ps):
    return model.encode([str(p) for p in ps], batch_size=32, normalize_embeddings=True, show_progress_bar=False, convert_to_numpy=True)

# Per-subset top-3 train (Q,A) for each test (need SAME subset for language consistency)
print("Encoding TV pool questions...")
t0 = time.time()
demos = {}  # test_ID -> list of {Q, A}
K = 3

for subset in sorted(test["subset"].unique()):
    pool_sub = tv_pool[tv_pool["subset"] == subset].reset_index(drop=True)
    test_sub = test[test["subset"] == subset].reset_index(drop=True)
    if len(pool_sub) == 0 or len(test_sub) == 0: continue
    print(f"  {subset}: pool={len(pool_sub)} test={len(test_sub)}")
    
    pool_emb = encode_q(pool_sub["input"].tolist())
    test_emb = encode_q(test_sub["input"].tolist())
    sims = test_emb @ pool_emb.T  # cosine since normalized
    
    # For each test row, get top-K train indices (exclude self if same ID exists in pool)
    for i, t_row in test_sub.iterrows():
        s = sims[i].copy()
        # Exclude any pool row with same ID (test rows shouldn't be in TV pool but defensive)
        for j, p_row in pool_sub.iterrows():
            if p_row["ID"] == t_row["ID"]:
                s[j] = -1.0
        top_idx = np.argsort(-s)[:K]
        demos[t_row["ID"]] = [
            {"Q": str(pool_sub.iloc[int(j)]["input"]), "A": str(pool_sub.iloc[int(j)]["output"])}
            for j in top_idx
        ]

print(f"  encoded in {time.time()-t0:.1f}s")
print(f"  demos built for {len(demos)} test rows")

# Build new test JSON: original prompt + few-shot demos prepended in "Examples" block
INSTRUCTION = ("Answer the following maternal/sexual/reproductive health question in {LANG}. "
               "Use the retrieved contexts as your primary sources — copy exact phrasing where the contexts already address the question. "
               "Be concise and factually accurate.")

new_test_path = OUT_DIR / "msrh_rag_test_k3_AfriE5_TV_fewshot_k3.json"
n = 0
with open(TEST_ORIG) as fin, open(new_test_path, "w") as fout:
    for r_idx, ln in enumerate(fin):
        r = json.loads(ln)
        test_row = test.iloc[r_idx]
        test_id = test_row["ID"]
        orig_content = r["messages"][0]["content"]
        # Get demos for this row
        dem_list = demos.get(test_id, [])
        if dem_list:
            ex_block = "\n\n".join([f"Example {j+1}:\nQ: {d['Q']}\nA: {d['A']}" for j, d in enumerate(dem_list)])
            # Insert ex_block right after the instruction (before contexts)
            idx = orig_content.find("Context 1:")
            if idx > 0:
                before = orig_content[:idx].rstrip()
                after = orig_content[idx:]
                new_content = f"{before}\n\nFew-shot examples (similar questions from training):\n{ex_block}\n\n{after}"
            else:
                new_content = orig_content
        else:
            new_content = orig_content
        r["messages"] = [{"role": "user", "content": new_content}]
        fout.write(json.dumps(r, ensure_ascii=False) + "\n")
        n += 1

print(f"\nwrote {new_test_path} ({n} rows, demos used: {sum(1 for d in demos.values() if d)})")