import chromadb
from chromadb.utils import embedding_functions

# Model embedding yang sama
emb = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="paraphrase-multilingual-MiniLM-L12-v2"
)

client = chromadb.PersistentClient(path="./chroma_db")

# 15 pertanyaan dan kata kunci dari PDF
UJI = [
    ("Apa arti airworthiness?", "fit for safe flight"),
    ("Apa kepanjangan MRO?", "Maintenance, Repair, and Overhaul"),
    ("Apa tugas line maintenance?", "Routine maintenance tasks"),
    ("Apa itu base maintenance?", "taken out of service"),
    ("Apa tujuan NDT?", "without causing damage"),
    ("Apa tujuan utama aircraft maintenance?", "safety and performance"),
    ("Apa itu pendekatan RCM?", "Reliability-Centered Maintenance"),
    ("Strategi apa yang menggunakan data real-time?", "Condition-Based Maintenance"),
    ("Siapa yang melakukan perawatan pesawat pada masa awal penerbangan?", "Wright brothers"),
    ("Kapan prosedur standar perawatan mulai dikembangkan maskapai?", "1920s"),
    ("Apa layanan yang digunakan maskapai untuk merawat armada?", "MRO services"),
    ("Apa tujuan perawatan pesawat militer?", "mission readiness"),
    ("Teknologi apa yang digunakan untuk pelatihan perawatan?", "Augmented Reality"),
    ("Kendaraan apa yang diteliti untuk inspeksi dan perbaikan?", "Unmanned Aerial Vehicles"),
    ("Material apa yang digunakan Boeing 787 Dreamliner?", "composite materials")
]


# Mencari peringkat pertama yang mengandung kata kunci
def cari_peringkat(koleksi, pertanyaan, kunci, k=5):

    jumlah = min(k, koleksi.count())

    if jumlah == 0:
        return None

    hasil = koleksi.query(
        query_texts=[pertanyaan],
        n_results=jumlah
    )

    # Mengambil daftar dokumen hasil pencarian
    docs = hasil["documents"][0]

    # Membersihkan kata kunci dari spasi dan baris baru
    kunci_bersih = " ".join(kunci.lower().split())

    for peringkat, doc in enumerate(docs, start=1):

        doc_bersih = " ".join(doc.lower().split())

        if kunci_bersih in doc_bersih:
            return peringkat

    return None


# Menghitung Hit@K
def hit_at_k(koleksi, uji, k=3):

    benar = 0

    for pertanyaan, kunci in uji:

        peringkat = cari_peringkat(
            koleksi, pertanyaan, kunci, k
        )

        if peringkat is not None:
            benar += 1

    return benar / len(uji)


# Menghitung MRR@5
def hitung_mrr(koleksi, uji):

    total = 0

    for pertanyaan, kunci in uji:

        peringkat = cari_peringkat(
            koleksi, pertanyaan, kunci, 5
        )

        if peringkat is not None:
            total += 1 / peringkat

    return total / len(uji)


# Membandingkan dua koleksi
for nama in ("pedoman_500", "pedoman_1000"):

    print("\n" + "=" * 50)
    print(f"EVALUASI KOLEKSI: {nama}")

    koleksi = client.get_collection(
        name=nama,
        embedding_function=emb
    )

    if koleksi.count() == 0:
        print("Koleksi kosong!")
        continue

    print("Jumlah pertanyaan:", len(UJI))

    for k in (1, 3, 5):

        nilai = hit_at_k(koleksi, UJI, k)

        print(f"Hit@{k} = {nilai:.2f}")

    mrr = hitung_mrr(koleksi, UJI)

    print(f"MRR@5 = {mrr:.2f}")

    print("\nPERTANYAAN YANG GAGAL DITEMUKAN:")

    for pertanyaan, kunci in UJI:

        peringkat = cari_peringkat(
            koleksi, pertanyaan, kunci, 5
        )

        if peringkat is None:
            print("-", pertanyaan)