# -*- coding: utf-8 -*-
"""
Tái tạo bi-encoder và bộ đọc trích xuất ĐÃ FINE-TUNE theo bài (bộ thực nghiệm
25/10/2026), kiểm tra nhanh trên tập test, rồi đẩy lên HF Hub cho Space dùng.

Bộ thực nghiệm không lưu trọng số (run_reading.py: save_strategy="no", không có
save_pretrained; run_experiments.py: --save-model tắt), nên phải huấn luyện lại.

Cấu hình (--config, giống tên trong ketqua_ablation_v2/):
    silver_sinh  gold + silver sinh mới       F1 trích xuất top-5 0,745 (tốt nhất) -> mặc định
    chinh        gold + silver người duyệt    cấu hình chính của bài (0,679)
    gold         chỉ gold                     0,717

Chạy trên vast.ai: xem run_vast_train.sh (cài thư viện, chạy nền, ghi log, bật Vast Monitor).

    export HF_TOKEN=hf_...            # token WRITE của tài khoản phamluan
    python tools/train_models.py --data data --user phamluan --config silver_sinh \
        --update-space phamluan/namboqa-api --data-repo phamluan/namboqa-private

Xong việc, script:
  1. đẩy <user>/namboqa-bi-encoder-<cfg> và <user>/namboqa-xlmr-reader-<cfg> (private)
  2. tính embedding corpus -> emb_*.npy, đẩy vào dataset private (Space khỏi mã hoá lại)
  3. đánh giá trên 695 câu test bằng ĐÚNG code của demo (engine.py) và in so sánh với bài
  4. (--update-space) đặt RETRIEVER_MODEL / READER_MODEL cho Space rồi restart

Cách chia (chia_tap.chia_theo_tinh, seed 42), nạp silver (chia_tap.nap_silver) và
siêu tham số y hệt nambo_run/code/run_experiments.py và run_reading.py.
"""

import argparse
import json
import os
import random
import re
import string
import sys
import time
import unicodedata
from collections import Counter
from pathlib import Path

import numpy as np
import torch

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
import chia_tap  # noqa: E402
from vm import vm_finish, vm_log, vm_start  # noqa: E402

CONFIGS = {
    "silver_sinh": dict(no_silver=False, nguon="sinh"),
    "chinh": dict(no_silver=False, nguon=None),
    "gold": dict(no_silver=True, nguon=None),
}
# Số liệu bài báo để đối chiếu (seed 42 / trung bình 3 seed)
PAPER = {
    "silver_sinh": {"R@5 hybrid": 0.882, "F1 oracle": 0.906, "F1 top-5": 0.745},
    "chinh": {"R@5 hybrid": 0.881, "F1 oracle": 0.901, "F1 top-5": 0.679, "F1 top-3": 0.717},
    "gold": {"R@5 hybrid": 0.882, "F1 oracle": 0.903, "F1 top-5": 0.717, "F1 top-3": 0.734},
}


RUN = "namboqa_demo"


def load_jsonl(p):
    return [json.loads(l) for l in open(p, encoding="utf-8") if l.strip()]


# chấm điểm: chép từ run_reading.py
def norm_ans(s):
    s = unicodedata.normalize("NFC", (s or "").lower().strip())
    s = "".join(c for c in s if c not in string.punctuation)
    s = re.sub(r"\b(là|của|tại|ở|vào|năm|the|a|an)\b", " ", s)
    return re.sub(r"\s+", " ", s).strip()


def f1_score(pred, gold):
    pt, gt = norm_ans(pred).split(), norm_ans(gold).split()
    if not pt or not gt:
        return float(pt == gt)
    same = sum((Counter(pt) & Counter(gt)).values())
    if same == 0:
        return 0.0
    p, r = same / len(pt), same / len(gt)
    return 2 * p * r / (p + r)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True, help="thư mục có qa_final, passages, places_clean, silver_final")
    ap.add_argument("--user", required=True, help="tài khoản HF, vd phamluan")
    ap.add_argument("--config", default="silver_sinh", choices=list(CONFIGS))
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--public", action="store_true", help="đẩy mô hình dạng public")
    ap.add_argument("--skip-retriever", action="store_true")
    ap.add_argument("--skip-reader", action="store_true")
    ap.add_argument("--skip-push", action="store_true", help="chỉ train + đánh giá, không đẩy lên Hub")
    ap.add_argument("--update-space", default=None, help="vd phamluan/namboqa-api")
    ap.add_argument("--data-repo", default=None, help="dataset private để đẩy emb_*.npy, vd phamluan/namboqa-private")
    ap.add_argument("--out", default="trained")
    a = ap.parse_args()
    cfg = CONFIGS[a.config]
    global RUN
    run = RUN = f"namboqa_demo_{a.config}_s{a.seed}"
    vm_start(run, config=vars(a))

    random.seed(a.seed); np.random.seed(a.seed); torch.manual_seed(a.seed)
    dev = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
    print("Thiết bị:", dev, torch.cuda.get_device_name(0) if dev == "cuda" else "")
    d, out = Path(a.data), Path(a.out)
    out.mkdir(exist_ok=True)
    passages = load_jsonl(d / "passages.jsonl")
    pid2text = {p["id"]: p["text"] for p in passages}
    qa = load_jsonl(d / "qa_final.jsonl")
    for q in qa:
        q["gold"] = [g for g in q["gold_passage_ids"] if g in pid2text]
    qa = [q for q in qa if q["gold"]]
    train_qa, _, test_qa, tk = chia_tap.chia_theo_tinh(qa, chia_tap.tinh_cua(d))
    test_qa, _ = chia_tap.tach_danh_gia(test_qa)
    silver = []
    if not cfg["no_silver"]:
        silver, tk_sv = chia_tap.nap_silver(d, qa, tk["train_p"], pid2text, nguon=cfg["nguon"])
        print("Silver:", tk_sv)
    print(f"Cấu hình {a.config}: train gold {len(train_qa)} câu / {len(tk['train_p'])} địa điểm, "
          f"silver {len(silver)}, test {len(test_qa)}")
    private = not a.public
    tag = a.config.replace("_", "-")
    retr_repo = f"{a.user}/namboqa-bi-encoder-{tag}"
    read_repo = f"{a.user}/namboqa-xlmr-reader-{tag}"
    emb_file = out / f"emb_{re.sub(r'[^\w.-]', '_', retr_repo)}.npy"
    res = {"config": a.config, "seed": a.seed}

    # ------------------------------------------------ bi-encoder (run_experiments.py)
    if not a.skip_retriever:
        from sentence_transformers import InputExample, SentenceTransformer, losses
        from torch.utils.data import DataLoader
        samples = [InputExample(texts=[q["question"], pid2text[q["gold"][0]]])
                   for q in train_qa + silver if q["type"] != "unanswerable"]
        print(f"Fine-tune bi-encoder: {len(samples)} cặp")
        ft = SentenceTransformer("bkai-foundation-models/vietnamese-bi-encoder", device=dev)
        dl = DataLoader(samples, shuffle=True, batch_size=16, drop_last=True)
        t0 = time.time()
        ft.fit(train_objectives=[(dl, losses.MultipleNegativesRankingLoss(ft))],
               epochs=3, warmup_steps=int(0.1 * len(dl) * 3), show_progress_bar=True)
        res["bi_encoder_phut"] = (time.time() - t0) / 60
        vm_log(run, bi_encoder_phut=res["bi_encoder_phut"], bi_encoder_cap=len(samples))
        ft.save(str(out / f"bi_encoder_{tag}"))
        emb = ft.encode([p["text"] for p in passages], batch_size=64,
                        normalize_embeddings=True, convert_to_numpy=True)
        np.save(emb_file, emb.astype(np.float32))
        if not a.skip_push:
            ft.push_to_hub(retr_repo, private=private, exist_ok=True)
            print("Đã đẩy", retr_repo)
        del ft
        if dev == "cuda":
            torch.cuda.empty_cache()

    # ------------------------------------------------ extractive (run_reading.py)
    if not a.skip_reader:
        from datasets import Dataset
        from transformers import (AutoModelForQuestionAnswering, AutoTokenizer, Trainer,
                                  TrainerCallback, TrainingArguments, default_data_collator)
        base = "deepset/xlm-roberta-base-squad2"
        tok = AutoTokenizer.from_pretrained(base)
        model = AutoModelForQuestionAnswering.from_pretrained(base)
        ex, miss = [], 0
        for q in train_qa + silver:
            if q["type"] == "unanswerable":
                ex.append({"question": q["question"], "context": pid2text[q["gold"][0]],
                           "answer": "", "start": 0, "impossible": True})
                continue
            ctx = " ".join(pid2text[g] for g in q["gold"])
            pos = ctx.find(q["answer"])
            if pos < 0:
                miss += 1
                continue
            ex.append({"question": q["question"], "context": ctx,
                       "answer": q["answer"], "start": pos, "impossible": False})
        print(f"Mẫu extractive: {len(ex)} (bỏ {miss} câu không định vị được đáp án)")

        def prep(b):
            enc = tok(b["question"], b["context"], truncation="only_second", max_length=384,
                      padding="max_length", return_offsets_mapping=True)
            st, en = [], []
            for i, off in enumerate(enc["offset_mapping"]):
                if b["impossible"][i]:
                    st.append(0); en.append(0); continue
                s, e = b["start"][i], b["start"][i] + len(b["answer"][i])
                sid = enc.sequence_ids(i)
                ts = te = 0
                for j, (x, y) in enumerate(off):
                    if sid[j] != 1:
                        continue
                    if x <= s < y: ts = j
                    if x < e <= y: te = j
                st.append(ts); en.append(te)
            enc["start_positions"], enc["end_positions"] = st, en
            enc.pop("offset_mapping")
            return enc

        class _VMCb(TrainerCallback):
            def on_log(self, args, state, control, logs=None, **kw):
                if state.is_world_process_zero and logs:
                    vm_log(run, step=state.global_step,
                           **{k: v for k, v in logs.items() if isinstance(v, (int, float))})

        ds = Dataset.from_list(ex).map(prep, batched=True, remove_columns=list(ex[0].keys()))
        # MPS (Mac 16 GB) hết bộ nhớ ở batch 8: batch nhỏ + tích luỹ, batch hiệu dụng vẫn 8
        bs, acc = (2, 4) if dev == "mps" else (8, 1)
        targs = TrainingArguments(output_dir=str(out / "_tmp"), num_train_epochs=10,
                                  per_device_train_batch_size=bs, gradient_accumulation_steps=acc,
                                  learning_rate=3e-5, weight_decay=0.01, seed=a.seed,
                                  logging_steps=20, save_strategy="no", report_to=[])
        t0 = time.time()
        Trainer(model=model, args=targs, train_dataset=ds, data_collator=default_data_collator,
                callbacks=[_VMCb()]).train()
        res["reader_phut"] = (time.time() - t0) / 60
        res["reader_mau"] = len(ex)
        vm_log(run, reader_phut=res["reader_phut"], reader_mau=len(ex))
        model.save_pretrained(out / f"reader_{tag}"); tok.save_pretrained(out / f"reader_{tag}")
        if not a.skip_push:
            model.push_to_hub(read_repo, private=private); tok.push_to_hub(read_repo, private=private)
            print("Đã đẩy", read_repo)

    # ------------------------------------------------ đánh giá bằng đúng code của demo
    if not (a.skip_retriever or a.skip_reader):
        res.update(evaluate(d, out, tag, emb_file, test_qa, dev))
        ref = PAPER[a.config]
        print("\n===== Kiểm tra trên tập test (code demo) so với bài =====")
        for k, v in res.items():
            if isinstance(v, float) and not k.endswith("phut"):
                print(f"  {k:28s} {v:.3f}" + (f"   (bài: {ref[k]:.3f})" if k in ref else ""))
        vm_log(run, **{k.replace(" ", "_").replace("@", ""): v for k, v in res.items()
                       if isinstance(v, (int, float))})
    (out / f"eval_{tag}.json").write_text(json.dumps(res, ensure_ascii=False, indent=1), encoding="utf-8")

    # ------------------------------------------------ cập nhật dataset + Space
    if not a.skip_push and (a.data_repo or a.update_space):
        from huggingface_hub import HfApi
        api = HfApi()
        if a.data_repo and emb_file.exists():
            api.upload_file(path_or_fileobj=str(emb_file), path_in_repo=emb_file.name,
                            repo_id=a.data_repo, repo_type="dataset")
            print("Đã đẩy", emb_file.name, "vào", a.data_repo)
        if a.update_space:
            if not a.skip_retriever:
                api.add_space_variable(a.update_space, "RETRIEVER_MODEL", retr_repo)
            if not a.skip_reader:
                api.add_space_variable(a.update_space, "READER_MODEL", read_repo)
            api.restart_space(a.update_space, factory_reboot=True)
            print("Đã đặt biến và khởi động lại", a.update_space)
    vm_finish(run, "done")


def evaluate(d, out, tag, emb_file, test_qa, dev):
    """Chạy đúng Engine của demo: hybrid α=0,5 + XLM-R đọc 3 đoạn (như Space), và
    XLM-R trên đoạn tham chiếu (oracle) để so với bài."""
    os.environ["RETRIEVER_MODEL"] = str(out / f"bi_encoder_{tag}")
    from engine import Engine
    from sentence_transformers import SentenceTransformer
    from transformers import AutoModelForQuestionAnswering, AutoTokenizer
    e = Engine(d)
    e.load_fast()
    e.encoder = SentenceTransformer(str(out / f"bi_encoder_{tag}"), device=dev)
    e.emb = np.load(emb_file).astype(np.float32)
    e.status["dense"] = "ready"
    e.qa_tok = AutoTokenizer.from_pretrained(out / f"reader_{tag}")
    e.qa_model = AutoModelForQuestionAnswering.from_pretrained(out / f"reader_{tag}").to(dev).eval()
    e._torch = torch
    e.status["extractive"] = "ready"

    hit5, f1_ora, ref_ora = [], [], []
    f1_k = {k: [] for k in (1, 3, 5)}
    ref_k = {k: [] for k in (1, 3, 5)}
    for q in test_qa:
        hits, _ = e.retrieve(q["question"], "hybrid", 5)
        ids = [e.ids[i] for i, _ in hits]
        hit5.append(bool(set(q["gold"]) & set(ids)))
        if q["type"] not in ("factoid", "unanswerable"):
            continue
        ora_ctx = " ".join(e.txts[e.pid2idx[g]] for g in q["gold"])[:3000]
        r = e.extract(q["question"], ora_ctx)
        pred_ora = r["answer"]          # "" nghĩa là mô hình từ chối
        for k in (1, 3, 5):
            ctx = " ".join(e.txts[i] for i, _ in hits[:k])[:3000]
            rk = e.extract(q["question"], ctx)
            pk = rk["answer"]
            if q["type"] == "factoid":
                f1_k[k].append(f1_score(pk, q["answer"]))
            else:
                ref_k[k].append(not pk)
        if q["type"] == "factoid":
            f1_ora.append(f1_score(pred_ora, q["answer"]))
        else:
            ref_ora.append(not pred_ora)
    res = {"R@5 hybrid": float(np.mean(hit5)), "F1 oracle": float(np.mean(f1_ora)),
           "Từ chối đúng oracle": float(np.mean(ref_ora))}
    for k in (1, 3, 5):
        res[f"F1 top-{k}"] = float(np.mean(f1_k[k]))
        res[f"Từ chối đúng top-{k}"] = float(np.mean(ref_k[k]))
    return res


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        vm_finish(RUN, "stopped"); raise
    except BaseException as ex:
        vm_finish(RUN, "failed", ex); raise