# model_neural.py
# Neural Language Model wrapper menggunakan HuggingFace Transformers
# Model: GPT-2 Indonesian (cahya/gpt2-small-indonesian-522M)
# 
# ANTI-HALLUCINATION STRATEGIES:
# 1. StoppingCriteria - berhenti di punctuation
# 2. Blacklist filtering - hapus kata halusinasi
# 3. Strict max_new_tokens - batasi panjang ketat

import torch
import re
from transformers import GPT2LMHeadModel, GPT2Tokenizer, StoppingCriteria, StoppingCriteriaList
import warnings

warnings.filterwarnings("ignore", category=FutureWarning)

# =============================================================================
# BLACKLIST KATA HALUSINASI - Kata yang sering muncul tanpa diminta
# =============================================================================
HALLUCINATION_BLACKLIST = {
    # Lokasi fiktif
    "amerika", "serikat", "indonesia", "jepang", "china", "eropa", 
    "jakarta", "surabaya", "bandung", "yogyakarta",
    
    # Konteks religius
    "surga", "neraka", "tuhan", "allah", "dewa", "malaikat",
    
    # Entitas politik
    "presiden", "menteri", "gubernur", "bupati", "walikota",
    
    # Kata sambung berlebihan
    "kemudian", "selanjutnya", "akhirnya", "oleh",
    
    # Noise tokens
    "``", "''", "--", "...",
}


# =============================================================================
# STOPPING CRITERIA - Berhenti di Punctuation
# =============================================================================
class StopOnPunctuation(StoppingCriteria):
    """Custom stopping criteria yang berhenti di tanda baca."""
    
    def __init__(self, tokenizer, stop_chars=[".", "!", "?", ","]):
        self.tokenizer = tokenizer
        self.stop_ids = set()
        
        # Encode setiap stop character untuk dapat ID-nya
        for char in stop_chars:
            try:
                ids = tokenizer.encode(char, add_special_tokens=False)
                self.stop_ids.update(ids)
            except:
                pass
    
    def __call__(self, input_ids, scores, **kwargs):
        # Cek apakah token terakhir adalah stop token
        if len(input_ids[0]) > 0:
            last_token = input_ids[0, -1].item()
            if last_token in self.stop_ids:
                return True
        return False


# =============================================================================
# MODEL NEURAL CLASS
# =============================================================================
class ModelNeural:
    """
    Wrapper untuk Neural Language Model (GPT-2 Indonesian).
    Dilengkapi dengan anti-hallucination strategies.
    """
    
    def __init__(self, model_name="cahya/gpt2-small-indonesian-522M"):
        """
        Inisialisasi model GPT-2 Indonesian.
        Download otomatis dari HuggingFace jika belum ada.
        """
        print(f"[NEURAL] Memuat model: {model_name}...")
        
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        print(f"[NEURAL] Menggunakan device: {self.device}")
        
        # Load tokenizer dan model
        self.tokenizer = GPT2Tokenizer.from_pretrained(model_name)
        self.model = GPT2LMHeadModel.from_pretrained(model_name)
        self.model.to(self.device)
        self.model.eval()
        
        # Set pad token (GPT-2 tidak punya default pad token)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
            self.model.config.pad_token_id = self.model.config.eos_token_id
        
        # Setup stopping criteria
        self.stopping_criteria = StoppingCriteriaList([
            StopOnPunctuation(self.tokenizer)
        ])
        
        print(f"[NEURAL] Model siap! Vocab size: {len(self.tokenizer)}")
    
    def _clean_output(self, text, input_word_count):
        """
        Post-processing untuk membersihkan output dari halusinasi.
        
        Args:
            text: Raw output dari model
            input_word_count: Jumlah kata dalam input (untuk batasan panjang)
            
        Returns:
            String yang sudah dibersihkan
        """
        if not text:
            return ""
        
        # Step 1: Potong di punctuation pertama
        match = re.search(r'[.!?,;:]', text)
        if match:
            text = text[:match.start()]
        
        # Step 2: Batasi panjang (input + 2 kata maksimal)
        max_words = input_word_count + 2
        words = text.split()[:max_words]
        
        # Step 3: Filter kata halusinasi
        filtered_words = []
        for w in words:
            if w.lower() not in HALLUCINATION_BLACKLIST:
                filtered_words.append(w)
        
        return " ".join(filtered_words).strip()
    
    def generate_beam(self, input_text, num_beams=5, max_new_tokens=5):
        """
        Generate teks menggunakan Beam Search dengan anti-hallucination.
        
        Args:
            input_text: Teks input (prompt)
            num_beams: Jumlah beam untuk pencarian
            max_new_tokens: Maksimum token baru (default: 5 untuk Gloss-to-Text)
            
        Returns:
            String hasil generasi (bersih, tanpa halusinasi)
        """
        input_word_count = len(input_text.split())
        
        # Tokenize input
        inputs = self.tokenizer(input_text, return_tensors="pt").to(self.device)
        input_length = inputs.input_ids.shape[1]
        
        # Generate dengan beam search + stopping criteria
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                num_beams=num_beams,
                early_stopping=True,
                do_sample=False,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
                no_repeat_ngram_size=2,
                stopping_criteria=self.stopping_criteria,  # ANTI-HALLUCINATION
            )
        
        # Decode output
        generated_tokens = outputs[0][input_length:]
        raw_result = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
        
        # Clean output
        clean_result = self._clean_output(raw_result, input_word_count)
        
        return clean_result
    
    def generate_nucleus(self, input_text, top_p=0.7, temperature=0.6, max_new_tokens=5):
        """
        Generate teks menggunakan Nucleus Sampling dengan anti-hallucination.
        
        Args:
            input_text: Teks input (prompt)
            top_p: Cumulative probability threshold (tuned: 0.7)
            temperature: Suhu sampling (tuned: 0.6 - lebih konservatif)
            max_new_tokens: Maksimum token baru (default: 5)
            
        Returns:
            String hasil generasi (bersih, tanpa halusinasi)
        """
        input_word_count = len(input_text.split())
        
        # Tokenize input
        inputs = self.tokenizer(input_text, return_tensors="pt").to(self.device)
        input_length = inputs.input_ids.shape[1]
        
        # Generate dengan nucleus sampling + stopping criteria
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=True,
                top_p=top_p,
                temperature=temperature,
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
                no_repeat_ngram_size=2,
                stopping_criteria=self.stopping_criteria,  # ANTI-HALLUCINATION
            )
        
        # Decode output
        generated_tokens = outputs[0][input_length:]
        raw_result = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
        
        # Clean output
        clean_result = self._clean_output(raw_result, input_word_count)
        
        return clean_result
    
    def get_continuation(self, input_text, method="beam", **kwargs):
        """
        Interface umum untuk generate continuation.
        
        Args:
            input_text: Teks input
            method: "beam" atau "nucleus"
            **kwargs: Parameter tambahan untuk method yang dipilih
            
        Returns:
            String hasil generasi
        """
        if method == "beam":
            return self.generate_beam(input_text, **kwargs)
        elif method == "nucleus":
            return self.generate_nucleus(input_text, **kwargs)
        else:
            raise ValueError(f"Method tidak dikenal: {method}")


# --- TEST ---
if __name__ == "__main__":
    print("=== TEST MODEL NEURAL (ANTI-HALLUCINATION) ===\n")
    
    # Inisialisasi model
    model = ModelNeural()
    
    # Test cases - kasus yang sebelumnya bermasalah
    test_inputs = [
        "Kami sedang belajar",      # Sebelumnya: "...di Amerika Serikat"
        "Siapa yang datang",         # Sebelumnya: "...ke surga"
        "Besok saya olahraga",       # Sebelumnya: "...dan pelatih saya"
        "Nama saya",
        "Terima kasih",
    ]
    
    print("\n--- Hasil Generasi (Dengan Anti-Hallucination) ---")
    for text in test_inputs:
        print(f"\nInput: '{text}'")
        
        beam_out = model.generate_beam(text, num_beams=5, max_new_tokens=5)
        print(f"  Beam Search : {text} {beam_out}".strip())
        
        nuc_out = model.generate_nucleus(text, top_p=0.7, temperature=0.6, max_new_tokens=5)
        print(f"  Nucleus     : {text} {nuc_out}".strip())

