"""Lab 3 - Ukur Hit@k dan MRR untuk satu atau beberapa koleksi.

Pakai:
    python evaluasi_retrieval.py 500 1000 1500
Hasil juga disimpan ke hasil_lab3.csv (rincian per pertanyaan) dan hasil_lab3.txt (ringkasan).
"""
import csv
import sys
from umum import ambil_embedding, buka_klien, ambil_koleksi, nama_koleksi
from uji import UJI

KS = (1, 3, 5)
K_MAKS = max(KS)

ukuran_list = [int(a) for a in sys.argv[1:]] or [800]
if len(UJI) < 15:
    print(f"PERINGATAN: baru {len(UJI)} pertanyaan uji; tugas meminta 15. Lengkapi uji.py.\n")

emb = ambil_embedding()
klien = buka_klien()
keluaran, baris_csv = [], []

def cetak(s=""):
    print(s)
    keluaran.append(s)

ringkasan = {}
for ukuran in ukuran_list:
    koleksi = ambil_koleksi(klien, nama_koleksi(ukuran), emb)

    # cek: apakah tiap frasa kunci memang ada di suatu chunk? (menangkap salah ketik)
    semua = " ".join(koleksi.get(include=["documents"])["documents"]).lower()
    for q, kunci in UJI:
        if kunci.lower() not in semua:
            cetak(f"  [!] koleksi {ukuran}: frasa '{kunci}' tidak ditemukan di chunk mana pun "
                  f"(salah ketik, atau terbelah di batas chunk) - pertanyaan: {q}")

    hit = {k: 0 for k in KS}
    rr_total = 0.0
    gagal = []
    for q, kunci in UJI:
        docs = koleksi.query(query_texts=[q], n_results=K_MAKS)["documents"][0]
        peringkat = next((i + 1 for i, d in enumerate(docs) if kunci.lower() in d.lower()), None)
        for k in KS:
            if peringkat and peringkat <= k:
                hit[k] += 1
        rr_total += (1 / peringkat) if peringkat else 0
        if not peringkat:
            gagal.append(q)
        baris_csv.append([ukuran, q, kunci, peringkat or "tidak muncul"])
    n = len(UJI)
    ringkasan[ukuran] = ({k: hit[k] / n for k in KS}, rr_total / n, gagal)

cetak()
cetak(f"{'ukuran chunk':>13} " + " ".join(f"Hit@{k:<3}" for k in KS) + f" {'MRR':>6}")
for ukuran, (h, mrr, _) in ringkasan.items():
    cetak(f"{ukuran:>13} " + " ".join(f"{h[k]:>6.2f}" for k in KS) + f" {mrr:>6.2f}")
cetak()
for ukuran, (_, _, gagal) in ringkasan.items():
    cetak(f"Pertanyaan yang gagal masuk top-{K_MAKS} pada ukuran {ukuran} ({len(gagal)}):")
    for q in gagal:
        cetak(f"   - {q}")

with open("hasil_lab3.csv", "w", newline="", encoding="utf-8") as f:
    w = csv.writer(f)
    w.writerow(["ukuran_chunk", "pertanyaan", "frasa_kunci", "peringkat_chunk_benar"])
    w.writerows(baris_csv)
with open("hasil_lab3.txt", "w", encoding="utf-8") as f:
    f.write("\n".join(keluaran))
print("\nDisimpan: hasil_lab3.csv dan hasil_lab3.txt")