import math
import re
from collections import Counter


class UKWMSAI_RAG:
  """Sistem RAG Sederhana berbasis Pure Python untuk topik Kecerdasan Buatan

  di Teknik Informatika Universitas Katolik Widya Mandala Surabaya (UKWMS).
  """

  def __init__(self, knowledge_base):
    print("[INFO] Menginisialisasi Basis Pengetahuan UKWMS - AI...")
    self.documents = knowledge_base
    self.vocab = self._build_vocab(self.documents)
    # Pra-penghitungan vektor TF-IDF untuk dokumen lokal
    self.doc_vectors = [
        self._compute_tfidf(doc, self.vocab, self.documents)
        for doc in self.documents
    ]

  def _tokenize(self, text):
    """Membersihkan teks dan memecahnya menjadi kata-kata (token)."""
    return re.findall(r"\b\w+\b", text.lower())

  def _build_vocab(self, docs):
    """Membangun kamus kata unik dari seluruh dokumen."""
    vocab = set()
    for doc in docs:
      vocab.update(self._tokenize(doc))
    return sorted(list(vocab))

  def _compute_tfidf(self, text, vocab, all_docs):
    """Menghitung bobot TF-IDF secara manual tanpa library eksternal."""
    tokens = self._tokenize(text)
    tf = Counter(tokens)
    total_tokens = len(tokens) if tokens else 1

    vector = {}
    for word in vocab:
      # Term Frequency (TF)
      term_freq = tf[word] / total_tokens

      # Inverse Document Frequency (IDF)
      doc_count = sum(1 for d in all_docs if word in self._tokenize(d))
      idf = (
          math.log(len(all_docs) / (doc_count + 1)) + 1 if doc_count > 0 else 0
      )

      vector[word] = term_freq * idf
    return vector

  def _cosine_similarity(self, vec1, vec2):
    """Menghitung kemiripan kosinus antara dua vektor TF-IDF."""
    intersection = set(vec1.keys()).intersection(set(vec2.keys()))
    numerator = sum(vec1[x] * vec2[x] for x in intersection)

    sum1 = sum(val ** 2 for val in vec1.values())
    sum2 = sum(val ** 2 for val in vec2.values())
    denominator = math.sqrt(sum1) * math.sqrt(sum2)

    if denominator == 0:
      return 0.0
    return numerator / denominator

  def retrieve(self, query, top_k=1):
    """Mencari dokumen paling relevan menggunakan Cosine Similarity."""
    query_vector = self._compute_tfidf(query, self.vocab, self.documents)
    scores = []

    for i, doc_vec in enumerate(self.doc_vectors):
      score = self._cosine_similarity(query_vector, doc_vec)
      scores.append((score, self.documents[i]))

    # Urutkan berdasarkan skor tertinggi
    scores.sort(key=lambda x: x[0], reverse=True)
    return [doc for score, doc in scores[:top_k]]

  def generate_answer(self, query):
    """Tahap Generation: Menggabungkan hasil retrieval dengan format jawaban RAG."""
    retrieved_docs = self.retrieve(query, top_k=1)
    context = (
        retrieved_docs[0]
        if retrieved_docs
        else "Maaf, informasi tidak ditemukan."
    )

    # Template Generator Extractive RAG
    answer = (
        f"==================================================\n"
        f"HASIL GENERASI RAG - INFORMATIKA UKWMS (KECERDASAN BUATAN)\n"
        f"==================================================\n"
        f"Pertanyaan : {query}\n\n"
        f"Konteks Relevan (Retrieved):\n-> {context}\n\n"
        f"Jawaban Sistem:\n"
        f"Berdasarkan kurikulum dan data akademik Prodi Informatika UKWMS, "
        f"fokus pada bidang Kecerdasan Buatan mencakup informasi berikut:\n"
        f'"{context}"\n'
        f"=================================================="
    )
    return answer, retrieved_docs


# ==========================================
# EKSEKUSI PROGRAM
# ==========================================
if __name__ == "__main__":
  # 1. Dataset / Basis Pengetahuan Lokal UKWMS - Kecerdasan Buatan
  ukwms_ai_kb = [
      (
          "Program Studi Informatika Fakultas Teknik UKWMS menyediakan"
          " peminatan khusus Data Science and Artificial Intelligence (AI)"
          " untuk membekali mahasiswa dengan kemampuan analisis data dan"
          " pembelajaran mesin."
      ),
      (
          "Mata kuliah inti dalam bidang Kecerdasan Buatan di Informatika UKWMS"
          " meliputi Machine Learning, Deep Learning, Natural Language"
          " Processing (NLP), Computer Vision, dan Data Mining."
      ),
      (
          "Laboratorium Komputasi dan Sistem Cerdas di Teknik Informatika"
          " UKWMS Kampus Kalijudan Surabaya sering digunakan mahasiswa untuk"
          " mengembangkan proyek akhir berbasis kecerdasan buatan dan IoT."
      ),
      (
          "Penelitian dosen dan mahasiswa Informatika UKWMS di bidang AI"
          " banyak berfokus pada pengenalan pola citra medis, sistem"
          " rekomendasi, dan analisis sentimen bahasa Indonesia."
      ),
      (
          "Lulusan peminatan Kecerdasan Buatan dari Informatika UKWMS bergelar"
          " Sarjana Komputer (S.Kom.) dengan profisiensi tinggi sebagai AI"
          " Engineer, Data Scientist, dan Machine Learning Practitioner."
      ),
  ]

  # 2. Inisialisasi Sistem RAG
  rag = UKWMSAI_RAG(ukwms_ai_kb)

  # 3. Contoh Pengujian Pertanyaan Tugas
  sample_query = (
      "Apa saja mata kuliah kecerdasan buatan yang dipelajari di informatika"
      " ukwms?"
  )

  # 4. Menjalankan Pipeline
  final_answer, sources = rag.generate_answer(sample_query)
  print("\n" + final_answer)