import math
import random
import time

# =============================================================================
# BLACKLIST KATA - Mencegah output halusinasi seperti "Nama ilmiah"
# =============================================================================
BLACKLIST = {
    "ilmiah", "desa", "kota", "lain", "tersebut", "lahir", 
    "``", "''", "(", ")", ":", ";", "--", "...",
    "adalah", "merupakan", "yaitu", "yakni",  # Kata definisi Wikipedia
}


# =============================================================================
# CLASS 1: DECODING UNTUK N-GRAM MODEL (STATISTICAL)
# =============================================================================
class Decoding:
    """Decoding strategies untuk Statistical N-gram Model."""
    
    def __init__(self, model):
        self.model = model

    # --- ALGORITMA 1: BEAM SEARCH (Deterministik / Pasti) ---
    def beam_search(self, words, beam_width=3):
        # Inisialisasi beam dengan (score, [list_kata])
        # Score 0 (log probability)
        beams = [(0, [])] 
        
        for word in words:
            new_beams = []
            for score, path in beams:
                # 1. Tambahkan kata asli dari glosa (tanpa diubah)
                # Probabilitas dianggap 1.0 (log 0) karena ini input wajib
                new_beams.append((score, path + [word]))
                
                # 2. Coba sisipkan kata penghubung (Preposition/Conjunction)
                # Ambil kata terakhir dari path untuk konteks
                last_word = path[-1] if path else "<s>"
                
                # Cek probabilitas kata penghubung yang mungkin muncul setelah last_word
                candidates = self.model.get_top_candidates(last_word, top_k=5)
                
                # FILTER BLACKLIST
                candidates = [(c, p) for c, p in candidates if c.lower() not in BLACKLIST]
                
                for candidate, prob in candidates:
                    # Log probability (biar angkanya tidak terlalu kecil)
                    if prob > 0:
                        log_prob = math.log(prob)
                        # Tambahkan kandidat + kata asli
                        # Contoh: "SAYA" -> (sisip "sedang") -> "MAKAN"
                        new_beams.append((score + log_prob, path + [candidate, word]))

            # Pruning: Ambil K beam terbaik
            # Urutkan berdasarkan score terbesar (descending)
            beams = sorted(new_beams, key=lambda x: x[0], reverse=True)[:beam_width]
            
        # Ambil path terbaik di akhir
        best_score, best_path = beams[0]
        return " ".join(best_path)

    # --- ALGORITMA 2: NUCLEUS SAMPLING (Stokastik / Variatif) ---
    def nucleus_sampling(self, words, top_p=0.5):  # Tuned: 0.9 -> 0.5
        output_path = []
        
        for i, word in enumerate(words):
            # 1. Konteks kata sebelumnya
            if output_path:
                last_word = output_path[-1]
            else:
                last_word = "<s>"
            
            # 2. Ambil semua kemungkinan kata sambung setelah last_word
            # Kita ambil agak banyak (top 20) untuk disampling
            candidates = self.model.get_top_candidates(last_word, top_k=20)
            
            # FILTER BLACKLIST
            candidates = [(c, p) for c, p in candidates if c.lower() not in BLACKLIST]
            
            # Jika tidak ada kandidat (probabilitas 0 semua), lanjut ke kata glosa
            if not candidates:
                output_path.append(word)
                continue
                
            # 3. Urutkan kandidat dari probabilitas tertinggi
            candidates.sort(key=lambda x: x[1], reverse=True)
            
            # 4. Cumulative Probability Cutoff (Inti Nucleus Sampling)
            cumulative_prob = 0.0
            nucleus_set = []
            probs_set = []
            
            for cand, prob in candidates:
                cumulative_prob += prob
                nucleus_set.append(cand)
                probs_set.append(prob)
                # Jika total probabilitas sudah > top_p, berhenti ambil kandidat
                if cumulative_prob >= top_p:
                    break
            
            # 5. Sampling (Pilih satu secara acak berbobot dari nucleus_set)
            # Ambil satu kata sambung dari nucleus set
            if nucleus_set:
                chosen_connector = random.choices(nucleus_set, weights=probs_set, k=1)[0]
                
                # Cek probabilitas transisi langsung vs transisi lewat konektor
                # Ini simplifikasi agar Nucleus tidak "nyampah" terlalu banyak
                prob_direct = self.model.get_probabilitas(last_word, word)
                prob_via_conn = self.model.get_probabilitas(last_word, chosen_connector)
                
                # Jika probabilitas konektor cukup signifikan dibanding langsung, sisipkan
                if prob_via_conn > (prob_direct * 0.5): # Threshold heuristik
                    output_path.append(chosen_connector)
            
            # Masukkan kata glosa asli (Wajib masuk)
            output_path.append(word)
            
        return " ".join(output_path)


# =============================================================================
# CLASS 2: DECODING UNTUK NEURAL MODEL (GPT-2)
# =============================================================================
class NeuralDecoding:
    """
    Decoding strategies untuk Neural Language Model (GPT-2).
    
    Untuk task Gloss-to-Text, Neural LM berfungsi sebagai:
    1. Text normalizer (menambah kata penghubung minimal)
    2. Grammar scorer (memilih urutan kata terbaik)
    
    BUKAN untuk text continuation/story generation.
    """
    
    def __init__(self, model_neural):
        """
        Args:
            model_neural: Instance dari ModelNeural (model_neural.py)
        """
        self.model = model_neural
    
    def beam_search(self, words, beam_width=5, max_tokens=5):
        """
        Beam Search menggunakan Neural LM.
        
        Untuk Gloss-to-Text: hanya menambah 1-2 kata penghubung jika perlu.
        
        Args:
            words: List kata dari glosa yang sudah dinormalisasi
            beam_width: Jumlah beam (default 5)
            max_tokens: Maksimum token baru yang digenerate (dibatasi kecil)
            
        Returns:
            String hasil generasi (pendek, sesuai target)
        """
        # Gabungkan kata menjadi prompt
        prompt = " ".join(words)
        
        # Jika input sudah cukup pendek, langsung return saja
        # (untuk kasus seperti "Halo", "Selamat pagi")
        if len(words) <= 2:
            return prompt
        
        # Generate menggunakan Neural Model (dengan max_tokens kecil)
        continuation = self.model.generate_beam(
            input_text=prompt,
            num_beams=beam_width,
            max_new_tokens=max_tokens
        )
        
        # Gabungkan prompt + continuation (minimal)
        full_output = f"{prompt} {continuation}".strip()
        
        # TRUNCATION: Potong agar output tidak terlalu panjang
        # Hitung berapa kata yang diharapkan (input + 1-3 kata)
        max_words = len(words) + 3
        output_words = full_output.split()[:max_words]
        
        # Potong di titik/koma pertama jika ada
        result = []
        for w in output_words:
            if "." in w or "," in w:
                # Ambil bagian sebelum punctuation
                clean = w.split(".")[0].split(",")[0]
                if clean:
                    result.append(clean)
                break
            result.append(w)
        
        return " ".join(result) if result else prompt
    
    def nucleus_sampling(self, words, top_p=0.9, temperature=0.7, max_tokens=5):
        """
        Nucleus Sampling menggunakan Neural LM.
        
        Untuk Gloss-to-Text: menambah variasi kata penghubung.
        
        Args:
            words: List kata dari glosa yang sudah dinormalisasi
            top_p: Cumulative probability threshold
            temperature: Sampling temperature
            max_tokens: Maksimum token baru yang digenerate
            
        Returns:
            String hasil generasi (pendek, sesuai target)
        """
        # Gabungkan kata menjadi prompt
        prompt = " ".join(words)
        
        # Jika input sudah cukup pendek, langsung return
        if len(words) <= 2:
            return prompt
        
        # Generate menggunakan Neural Model
        continuation = self.model.generate_nucleus(
            input_text=prompt,
            top_p=top_p,
            temperature=temperature,
            max_new_tokens=max_tokens
        )
        
        # Gabungkan prompt + continuation
        full_output = f"{prompt} {continuation}".strip()
        
        # TRUNCATION: Batasi panjang output
        max_words = len(words) + 3
        output_words = full_output.split()[:max_words]
        
        # Potong di punctuation pertama
        result = []
        for w in output_words:
            if "." in w or "," in w:
                clean = w.split(".")[0].split(",")[0]
                if clean:
                    result.append(clean)
                break
            result.append(w)
        
        return " ".join(result) if result else prompt