import chromadb
from chromadb.utils import embedding_functions

# Panggil DB dan Model
emb = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="paraphrase-multilingual-MiniLM-L12-v2")
client = chromadb.PersistentClient(path="./chroma_db")
koleksi = client.get_or_create_collection("pedoman", embedding_function=emb)

# Daftar (Pertanyaan, Kata Kunci Jawaban yang HARUS ada di chunk)
UJI = [
    ("Apa syarat cuti akademik?", "cuti akademik"),
    ("Berapa lama cuti akademik bisa diajukan?", "dua semester"),
    ("Berapa batas kehadiran untuk ujian?", "75 persen"),
    ("Bolehkah berhenti kuliah sementara?", "cuti akademik"),
    ("Bagaimana cara mengajukan beasiswa?", "beasiswa"),
    ("Apa sanksi jika telat bayar kuliah?", "terlambat melakukan pembayaran"),
    ("Kapan batas pembayaran biaya kuliah?", "dua minggu sebelum"),
    ("Apa yang harus disimpan setelah bayar?", "bukti pembayaran")
]

def hit_at_k(uji, k=3):
    benar = 0
    for pertanyaan, kunci in uji:
        # Ambil k hasil teratas
        hasil = koleksi.query(query_texts=[pertanyaan], n_results=k)["documents"][0]
        # Cek apakah kunci ada di salah satu hasil
        if any(kunci.lower() in d.lower() for d in hasil):
            benar += 1
    return benar / len(uji)

print(f"Total pertanyaan uji: {len(UJI)}\n")
print("Mengevaluasi Retrieval...")
for k in (1, 3, 5):
    skor = hit_at_k(UJI, k)
    print(f"Hit@{k} = {skor:.2f}  (Benar {int(skor*len(UJI))} dari {len(UJI)} pertanyaan)")