import os
import nltk
import pandas as pd
import pickle
from collections import Counter, defaultdict

try:
    nltk.data.find('tokenizers/punkt')
except LookupError:
    print("[INFO] Sedang mendownload data NLTK 'punkt'...")
    nltk.download('punkt')
    nltk.download('punkt_tab')

class ModelNgram:
    """
    Statistical Language Model berbasis N-gram (Bigram).
    Dilengkapi dengan fitur Model Persistence (Simpan & Muat).
    """
    
    def __init__(self):
        self.bigram_counts = defaultdict(Counter)
        self.unigram_counts = Counter()
        self.total_words = 0

    def latih(self, path_corpus):
        print(f"Sedang melatih model dari: {os.path.basename(path_corpus)}...")
        
        if not os.path.exists(path_corpus):
            print(f"WARNING: File {path_corpus} tidak ditemukan! Melewati...")
            return

        text = ""

        # --- LOGIKA BACA FILE (TXT / CSV) ---
        if path_corpus.endswith('.csv'):
            try:
                df = pd.read_csv(path_corpus)
                if len(df.columns) <= 1: 
                    df = pd.read_csv(path_corpus, sep=';')
                
                target_col = None
                possible_names = ['text', 'kalimat', 'sentence', 'content', 'teks', 'jawaban', 'reply']
                
                for col in df.columns:
                    if col.lower() in possible_names:
                        target_col = col
                        break
                
                if target_col:
                    text = " ".join(df[target_col].astype(str).tolist()).lower()
                else:
                    return
            except Exception as e:
                print(f"ERROR Membaca CSV: {e}")
                return

        else:
            try:
                with open(path_corpus, 'r', encoding='utf-8') as f:
                    text = f.read().lower()
            except UnicodeDecodeError:
                with open(path_corpus, 'r', encoding='latin-1') as f:
                    text = f.read().lower()

        if not text: return

        # --- TOKENISASI ---
        tokens = nltk.word_tokenize(text)
        self.total_words += len(tokens)
        self.unigram_counts.update(tokens)

        # Hitung Bigram
        for i in range(len(tokens) - 1):
            w1 = tokens[i]
            w2 = tokens[i+1]
            self.bigram_counts[w1][w2] += 1
            
        print(f"-> Selesai. Total kata sekarang: {self.total_words}")

    def get_probabilitas(self, kata_sebelum, kata_target):
        """Menghitung probabilitas P(target | sebelum)"""
        kata_sebelum = kata_sebelum.lower()
        kata_target = kata_target.lower()
        
        count_prev = self.unigram_counts[kata_sebelum]
        if count_prev == 0: return 1e-10
            
        count_pair = self.bigram_counts[kata_sebelum][kata_target]
        prob = count_pair / count_prev
        return prob if prob > 0 else 1e-10

    # --- FUNGSI MENGAMBIL KANDIDAT TERATAS---
    def get_top_candidates(self, prev_word, top_k=5):
        """
        Mengembalikan daftar (kata, probabilitas) terbaik yang mungkin muncul setelah prev_word.
        Digunakan oleh Beam Search dan Nucleus Sampling.
        """
        prev_word = prev_word.lower()
        
        # Jika kata sebelumnya tidak ada di kamus, kembalikan list kosong
        if prev_word not in self.bigram_counts:
            return []
            
        # Ambil semua kemungkinan kata selanjutnya
        next_words_counter = self.bigram_counts[prev_word]
        total_occurrence = self.unigram_counts[prev_word]
        
        candidates = []
        for word, count in next_words_counter.items():
            prob = count / total_occurrence
            candidates.append((word, prob))
            
        # Urutkan dari probabilitas tertinggi
        candidates.sort(key=lambda x: x[1], reverse=True)
        
        # Ambil top_k saja
        return candidates[:top_k]

    # =========================================================================
    # MODEL PERSISTENCE - Simpan & Muat Model
    # =========================================================================
    
    def simpan_model(self, path):
        """
        Simpan model yang sudah ditraining ke file menggunakan pickle.
        
        Args:
            path: Path file untuk menyimpan model (contoh: 'model/ngram_model.pkl')
            
        Returns:
            True jika berhasil, False jika gagal
        """
        try:
            # Buat direktori jika belum ada
            dir_path = os.path.dirname(path)
            if dir_path and not os.path.exists(dir_path):
                os.makedirs(dir_path)
            
            # Data yang akan disimpan
            model_data = {
                'bigram_counts': dict(self.bigram_counts),
                'unigram_counts': dict(self.unigram_counts),
                'total_words': self.total_words,
                'version': '1.0'  # Untuk kompatibilitas masa depan
            }
            
            with open(path, 'wb') as f:
                pickle.dump(model_data, f)
            
            # Hitung ukuran file
            size_mb = os.path.getsize(path) / (1024 * 1024)
            print(f"[SIMPAN] Model berhasil disimpan ke: {path}")
            print(f"[SIMPAN] Ukuran file: {size_mb:.2f} MB")
            print(f"[SIMPAN] Total kata: {self.total_words:,}")
            print(f"[SIMPAN] Vocab size: {len(self.unigram_counts):,}")
            return True
            
        except Exception as e:
            print(f"[ERROR] Gagal menyimpan model: {e}")
            return False
    
    def muat_model(self, path):
        """
        Muat model dari file pickle.
        
        Args:
            path: Path file model yang sudah disimpan
            
        Returns:
            True jika berhasil, False jika gagal
        """
        try:
            if not os.path.exists(path):
                print(f"[MUAT] File tidak ditemukan: {path}")
                return False
            
            with open(path, 'rb') as f:
                model_data = pickle.load(f)
            
            # Restore data ke object
            self.bigram_counts = defaultdict(Counter, {
                k: Counter(v) for k, v in model_data['bigram_counts'].items()
            })
            self.unigram_counts = Counter(model_data['unigram_counts'])
            self.total_words = model_data['total_words']
            
            print(f"[MUAT] Model berhasil dimuat dari: {path}")
            print(f"[MUAT] Total kata: {self.total_words:,}")
            print(f"[MUAT] Vocab size: {len(self.unigram_counts):,}")
            return True
            
        except Exception as e:
            print(f"[ERROR] Gagal memuat model: {e}")
            return False
    
    @staticmethod
    def ada_model_tersimpan(path):
        """Cek apakah file model sudah ada."""
        return os.path.exists(path)


# --- TEST MODEL PERSISTENCE ---
if __name__ == "__main__":
    print("=== TEST MODEL PERSISTENCE ===\n")
    
    model = ModelNgram()
    path_model = "model/ngram_trained.pkl"
    
    # Cek apakah sudah ada model tersimpan
    if ModelNgram.ada_model_tersimpan(path_model):
        print("[TEST] Model tersimpan ditemukan, memuat...")
        model.muat_model(path_model)
    else:
        print("[TEST] Tidak ada model tersimpan, training dari corpus...")
        model.latih("data/corpus_indonesia.txt")
        model.simpan_model(path_model)
    
    # Test fungsi
    print(f"\nTest get_top_candidates('saya'):")
    print(model.get_top_candidates('saya', top_k=5))