import os
import pickle
import re
import string
import warnings
import numpy as np
import torch
import torch.nn.functional as F

# Suppress sklearn version mismatch warning (model ditraining di versi berbeda)
warnings.filterwarnings('ignore', category=UserWarning, module='sklearn')

# ── Lazy imports (hanya load saat pertama kali dipanggil) ──────────────────
_nb_model      = None
_tfidf         = None
_label_map_rev = None
_bert_model    = None
_bert_tokenizer = None

BASE_DIR   = os.path.dirname(os.path.dirname(__file__))
NB_DIR     = os.path.join(BASE_DIR, 'models', 'naive_bayes')
BERT_DIR   = os.path.join(BASE_DIR, 'models', 'indobert')
DEVICE     = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
MAX_LEN    = 128

LABEL_MAP_REV = {0: 'Negatif', 1: 'Netral', 2: 'Positif'}

# ── Preprocessing ──────────────────────────────────────────────────────────
SLANG_DICT = {
    'gak': 'tidak', 'ga': 'tidak', 'tdk': 'tidak', 'ngga': 'tidak',
    'nggak': 'tidak', 'udah': 'sudah', 'udh': 'sudah', 'sdh': 'sudah',
    'yg': 'yang', 'dgn': 'dengan', 'utk': 'untuk', 'krn': 'karena',
    'karna': 'karena', 'bgt': 'banget', 'sy': 'saya', 'gw': 'saya',
    'gue': 'saya', 'lo': 'kamu', 'lu': 'kamu', 'tp': 'tapi',
    'tpi': 'tapi', 'klo': 'kalau', 'klu': 'kalau', 'kl': 'kalau',
    'bs': 'bisa', 'hrs': 'harus', 'blm': 'belum', 'jg': 'juga',
    'jga': 'juga', 'skrg': 'sekarang', 'bener': 'benar',
    'emang': 'memang', 'jd': 'jadi', 'mbg': 'makan bergizi gratis',
    'pgn': 'ingin', 'pingin': 'ingin', 'lbh': 'lebih',
    'msh': 'masih', 'lg': 'lagi',
}

try:
    from Sastrawi.StopWordRemover.StopWordRemoverFactory import StopWordRemoverFactory
    from Sastrawi.Stemmer.StemmerFactory import StemmerFactory
    _stop_factory = StopWordRemoverFactory()
    _stopwords    = set(_stop_factory.get_stop_words())
    _stemmer      = StemmerFactory().create_stemmer()
    SASTRAWI_OK   = True
except ImportError:
    _stopwords  = set()
    SASTRAWI_OK = False


def _normalize_slang(text):
    return ' '.join([SLANG_DICT.get(w, w) for w in text.split()])


def preprocess_nb(text):
    """Preprocessing penuh untuk Naive Bayes (stemming + stopword)."""
    if not isinstance(text, str) or not text.strip():
        return ''
    text = text.lower()
    text = re.sub(r'http\S+|www\S+', '', text)
    text = re.sub(r'@\w+|#\w+', '', text)
    text = text.encode('ascii', 'ignore').decode('ascii')
    text = re.sub(r'\d+', '', text)
    text = text.translate(str.maketrans('', '', string.punctuation))
    text = re.sub(r'\s+', ' ', text).strip()
    text = _normalize_slang(text)
    if SASTRAWI_OK:
        words = [w for w in text.split() if w not in _stopwords and len(w) > 2]
        text  = _stemmer.stem(' '.join(words))
    return text


def preprocess_bert(text):
    """Preprocessing ringan untuk IndoBERT."""
    if not isinstance(text, str):
        return ''
    text = re.sub(r'http\S+|www\S+', '', text)
    text = re.sub(r'@\w+|#\w+', '', text)
    text = re.sub(r'\s+', ' ', text).strip()
    return text[:512]


# ── Loader ────────────────────────────────────────────────────────────────
def _load_nb():
    global _nb_model, _tfidf, _label_map_rev
    if _nb_model is None:
        with open(os.path.join(NB_DIR, 'nb_model.pkl'), 'rb') as f:
            _nb_model = pickle.load(f)
        with open(os.path.join(NB_DIR, 'tfidf_vectorizer.pkl'), 'rb') as f:
            _tfidf = pickle.load(f)
        with open(os.path.join(NB_DIR, 'label_map.pkl'), 'rb') as f:
            data = pickle.load(f)
            _label_map_rev = data['label_map_reverse']


def _load_bert():
    global _bert_model, _bert_tokenizer
    if _bert_model is None:
        from transformers import BertForSequenceClassification, BertTokenizer
        _bert_tokenizer = BertTokenizer.from_pretrained(BERT_DIR)
        _bert_model     = BertForSequenceClassification.from_pretrained(BERT_DIR)
        _bert_model     = _bert_model.to(DEVICE)
        _bert_model.eval()


# ── Prediksi ──────────────────────────────────────────────────────────────
def predict_naive_bayes(text):
    _load_nb()
    clean = preprocess_nb(text)
    vec   = _tfidf.transform([clean])
    pred  = _nb_model.predict(vec)[0]
    proba = _nb_model.predict_proba(vec)[0]
    return {
        'label':      _label_map_rev[pred],
        'confidence': round(float(max(proba)) * 100, 2),
        'prob_negatif': round(float(proba[0]) * 100, 2),
        'prob_netral':  round(float(proba[1]) * 100, 2),
        'prob_positif': round(float(proba[2]) * 100, 2),
    }


def predict_indobert(text):
    _load_bert()
    clean    = preprocess_bert(text)
    encoding = _bert_tokenizer(
        clean,
        add_special_tokens=True,
        max_length=MAX_LEN,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )
    with torch.no_grad():
        outputs = _bert_model(
            input_ids=encoding['input_ids'].to(DEVICE),
            attention_mask=encoding['attention_mask'].to(DEVICE)
        )
    proba = F.softmax(outputs.logits, dim=1).cpu().numpy()[0]
    pred  = int(np.argmax(proba))
    return {
        'label':      LABEL_MAP_REV[pred],
        'confidence': round(float(max(proba)) * 100, 2),
        'prob_negatif': round(float(proba[0]) * 100, 2),
        'prob_netral':  round(float(proba[1]) * 100, 2),
        'prob_positif': round(float(proba[2]) * 100, 2),
    }


def predict(text, model='indobert'):
    if model == 'naive_bayes':
        return predict_naive_bayes(text)
    return predict_indobert(text)


# ── Word Count dari TF-IDF ─────────────────────────────────────────────────
def extract_wordcount_from_tfidf():
    """
    Ekstrak vocabulary + frekuensi dari TF-IDF vectorizer (Naive Bayes).
    Sentimen dominan ditentukan dari feature_count_ model NB per kelas.
    """
    _load_nb()

    vocab = _tfidf.vocabulary_       # {kata: index}
    feat  = _nb_model.feature_count_ # shape (n_class, n_features)

    label_map = {0: 'Negatif', 1: 'Netral', 2: 'Positif'}
    results = []

    for kata, idx in vocab.items():
        if len(kata) <= 2:
            continue

        counts = feat[:, idx]  # [neg_count, net_count, pos_count]
        total  = float(counts.sum())
        if total < 2:
            continue

        dominan_idx = int(counts.argmax())
        results.append({
            'kata':             kata,
            'frekuensi':        int(total),
            'sentimen_dominan': label_map[dominan_idx],
            'pct_positif':      round(counts[2] / total * 100, 2),
            'pct_negatif':      round(counts[0] / total * 100, 2),
            'pct_netral':       round(counts[1] / total * 100, 2),
        })

    results.sort(key=lambda x: x['frekuensi'], reverse=True)
    return results