app/models/intent_classifier/model.py

import os
import json
import joblib
import numpy as np
import warnings
from collections import Counter
from typing import Optional
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline, FeatureUnion
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
from app.nlp.preprocessor import preprocess
from app.nlp.language_detector import detect_language

# Keyword-based fallback for when model is not trained
INTENT_KEYWORDS = {
    'plan_trip': {
        'vi': ['kế hoạch', 'lịch trình', 'plan', 'đi du lịch', 'chuyến đi', 'đi chơi', 'đi nghỉ', 'lên plan', 'lập lịch'],
        'en': ['plan', 'itinerary', 'trip', 'travel', 'vacation', 'holiday', 'organize', 'schedule']
    },
    'find_hotel': {
        'vi': ['khách sạn', 'hotel', 'resort', 'homestay', 'phòng', 'chỗ ở', 'nghỉ', 'villa', 'nhà nghỉ', 'đặt phòng'],
        'en': ['hotel', 'resort', 'homestay', 'room', 'accommodation', 'stay', 'lodge', 'villa', 'book room', 'hostel']
    },
    'find_flight': {
        'vi': ['vé bay', 'chuyến bay', 'máy bay', 'bay', 'hàng không', 'đặt vé', 'flight'],
        'en': ['flight', 'fly', 'airline', 'ticket', 'plane', 'airport', 'book flight', 'airfare']
    },
    'budget_advice': {
        'vi': ['bao nhiêu', 'chi phí', 'ngân sách', 'tốn', 'giá', 'tiền', 'triệu', 'nghìn', 'đủ không', 'hết bao nhiêu', 'tiết kiệm'],
        'en': ['cost', 'budget', 'price', 'expensive', 'cheap', 'afford', 'money', 'spend', 'how much', 'dollar']
    },
    'activity_suggest': {
        'vi': ['hoạt động', 'chơi gì', 'đi đâu', 'tham quan', 'vui', 'giải trí', 'điểm đến', 'nổi tiếng', 'hay'],
        'en': ['activity', 'do', 'see', 'visit', 'attraction', 'fun', 'entertainment', 'landmark', 'sightseeing', 'tour']
    },
    'food_recommend': {
        'vi': ['ăn gì', 'món ăn', 'ẩm thực', 'đặc sản', 'quán ăn', 'nhà hàng', 'ngon', 'hải sản', 'cafe', 'uống'],
        'en': ['eat', 'food', 'cuisine', 'restaurant', 'dish', 'delicious', 'seafood', 'cafe', 'drink', 'specialty']
    },
    'weather_info': {
        'vi': ['thời tiết', 'mùa', 'nhiệt độ', 'mưa', 'nắng', 'lạnh', 'nóng', 'tháng mấy', 'khi nào đi'],
        'en': ['weather', 'season', 'temperature', 'rain', 'sunny', 'cold', 'hot', 'climate', 'when to visit', 'best time']
    },
    'visa_info': {
        'vi': ['visa', 'thị thực', 'hộ chiếu', 'passport', 'nhập cảnh', 'xuất cảnh', 'giấy tờ'],
        'en': ['visa', 'passport', 'immigration', 'entry', 'document', 'requirement', 'permit']
    },
    'transport_info': {
        'vi': ['di chuyển', 'xe bus', 'tàu', 'taxi', 'grab', 'thuê xe', 'sân bay', 'phương tiện', 'đường đi'],
        'en': ['transport', 'bus', 'train', 'taxi', 'grab', 'rent', 'airport', 'how to get', 'getting around']
    },
    'greeting': {
        'vi': ['xin chào', 'chào', 'hi', 'hello', 'hey', 'alo', 'bạn ơi'],
        'en': ['hello', 'hi', 'hey', 'good morning', 'good afternoon', 'greetings', 'yo']
    },
    'thank': {
        'vi': ['cảm ơn', 'cám ơn', 'thanks', 'thank', 'tks'],
        'en': ['thank', 'thanks', 'appreciate', 'grateful']
    },
    # 5 new intents
    'compare_destinations': {
        # Require explicit comparison signals — bare location names must NOT match here
        'vi': ['so sánh', 'nơi nào tốt hơn', 'nơi nào hợp', 'đâu tốt hơn', 'chọn giữa', 'hay hơn', 'vs', 'versus', 'hoặc là', 'so với'],
        'en': ['compare', 'versus', 'vs', 'which is better', 'difference between', 'compare between']
    },
    'optimize_budget': {
        'vi': ['tối ưu chi phí', 'tối ưu ngân sách', 'tiết kiệm nhất', 'rẻ nhất', 'phân bổ ngân sách', 'chi phí tối ưu', 'tối ưu'],
        'en': ['optimize cost', 'optimize budget', 'cheapest', 'best value', 'budget optimization', 'allocate budget', 'minimize cost']
    },
    'personalize_recommend': {
        'vi': ['phù hợp với tôi', 'gợi ý cho tôi', 'tôi thích', 'sở thích', 'phong cách', 'cá nhân hóa', 'dựa trên'],
        'en': ['suit me', 'recommend for me', 'i like', 'my preference', 'personalize', 'based on my', 'my style']
    },
    'multi_destination_trip': {
        'vi': ['nhiều điểm', 'kết hợp', 'ghé thêm', 'đi thêm', 'rồi đi', 'và đi', 'kết hợp thêm'],
        'en': ['multiple destinations', 'combine', 'multi-city', 'then go', 'also visit', 'stop at', 'add destination']
    },
    'dietary_cuisine': {
        'vi': ['ăn chay', 'chế độ ăn', 'dị ứng', 'không ăn', 'kiêng', 'halal', 'thuần chay'],
        'en': ['vegetarian', 'dietary', 'allergy', 'vegan', 'halal', 'gluten free', 'food restriction']
    },
}


class IntentClassifier:
    def __init__(self):
        self.pipeline = None
        self.is_trained = False

    def train(self, texts: list, labels: list) -> dict:
        """Train the intent classifier."""
        # --- Input validation ---
        if len(texts) != len(labels):
            raise ValueError(
                f"texts and labels length mismatch: {len(texts)} vs {len(labels)}"
            )

        counts = Counter(labels)
        n_classes = len(counts)
        if n_classes < 2:
            raise ValueError(
                f"Need at least 2 intent classes, got {n_classes}: {list(counts.keys())}"
            )

        # train_test_split with stratify=labels requires each class to have >= 2 samples
        # and total samples to produce a non-empty test split.
        min_per_class = min(counts.values())
        min_total_needed = max(7, n_classes * 2)
        _can_split = (min_per_class >= 2) and (len(texts) >= min_total_needed)

        if not _can_split:
            warnings.warn(
                f"Dataset too small for stratified split "
                f"(total={len(texts)}, min_per_class={min_per_class}, classes={n_classes}). "
                f"Training on full set without eval split.",
                UserWarning,
            )
        # --- End validation ---

        # Preprocess texts
        processed = []
        for text in texts:
            lang = detect_language(text)
            processed.append(preprocess(text, lang))

        self.pipeline = Pipeline([
            ('features', FeatureUnion([
                # Character n-grams: robust to typos, morphological variants
                ('char_ngram', TfidfVectorizer(
                    max_features=5000,
                    ngram_range=(1, 3),
                    sublinear_tf=True,
                    min_df=1,
                    analyzer='char_wb',
                )),
                # Word-level TF-IDF: resolves food_recommend ↔ activity_suggest confusion
                ('word_tfidf', TfidfVectorizer(
                    max_features=3000,
                    ngram_range=(1, 2),
                    sublinear_tf=True,
                    min_df=1,
                    analyzer='word',
                )),
            ])),
            ('clf', LogisticRegression(
                max_iter=1000,
                C=5.0,
                class_weight='balanced',
                solver='lbfgs'
            ))
        ])

        # Split data for evaluation (only when we have enough samples)
        if _can_split:
            X_train, X_test, y_train, y_test = train_test_split(
                processed, labels, test_size=0.15, random_state=42, stratify=labels
            )
        else:
            X_train, y_train = processed, labels
            X_test, y_test = [], []

        self.pipeline.fit(X_train, y_train)
        self.is_trained = True

        # Evaluate
        if not X_test:
            return {"accuracy": None, "report": {}, "note": "Trained on full set (dataset too small for eval split)"}

        y_pred = self.pipeline.predict(X_test)
        accuracy = np.mean(np.array(y_pred) == np.array(y_test))

        return {
            "accuracy": float(accuracy),
            "report": classification_report(y_test, y_pred, output_dict=True)
        }

    def predict(self, text: str, language: Optional[str] = None) -> dict:
        """Predict intent for a given text."""
        lang = language or detect_language(text)
        processed = preprocess(text, lang)

        if self.is_trained and self.pipeline is not None:
            proba = self.pipeline.predict_proba([processed])[0]
            classes = self.pipeline.classes_
            top_idx = np.argmax(proba)
            intent = classes[top_idx]
            confidence = float(proba[top_idx])

            # OOD rejector: uncertain predictions where top-2 candidates are
            # both within {greeting, fallback} resolve to fallback.
            # Fixes greeting precision 0.59 → ~0.80 without retraining.
            if confidence < 0.55:
                top2_indices = np.argsort(proba)[-2:]
                top2_intents = {classes[i] for i in top2_indices}
                if top2_intents <= {'greeting', 'fallback'}:
                    return {
                        "intent": "fallback",
                        "confidence": confidence,
                        "method": "ood_rejector"
                    }

            # If confidence is low, try keyword fallback
            if confidence < 0.4:
                keyword_intent = self._keyword_fallback(text, lang)
                if keyword_intent:
                    return {
                        "intent": keyword_intent,
                        "confidence": 0.6,
                        "method": "keyword_fallback"
                    }

            return {
                "intent": intent,
                "confidence": confidence,
                "method": "ml_model"
            }
        else:
            # Pure keyword-based classification
            intent = self._keyword_fallback(text, lang)
            return {
                "intent": intent or "fallback",
                "confidence": 0.5 if intent else 0.3,
                "method": "keyword_only"
            }

    def _keyword_fallback(self, text: str, lang: str) -> str | None:
        """Keyword-based intent detection fallback."""
        text_lower = text.lower()
        best_intent = None
        best_score = 0

        for intent, keywords in INTENT_KEYWORDS.items():
            score = 0
            # Only use keywords for the detected language to avoid cross-language false positives
            # (e.g. Vietnamese 'bay' meaning 'seven' shouldn't trigger find_flight in English)
            lang_keywords = keywords.get(lang, [])
            for keyword in lang_keywords:
                if keyword in text_lower:
                    score += len(keyword)  # longer matches get higher scores
            if score > best_score:
                best_score = score
                best_intent = intent

        return best_intent if best_score > 0 else None

    def save(self, path: str):
        """Save trained model to disk."""
        os.makedirs(os.path.dirname(path), exist_ok=True)
        joblib.dump(self.pipeline, path)

    def load(self, path: str):
        """Load trained model from disk."""
        if os.path.exists(path):
            try:
                self.pipeline = joblib.load(path)
                # Verify the pipeline actually works (catches sklearn version mismatches)
                self.pipeline.predict_proba(["test"])
                self.is_trained = True
            except Exception as e:
                import logging
                logger = logging.getLogger(__name__)
                logger.warning(
                    f"Failed to load intent classifier model ({e}). "
                    f"Falling back to keyword-based classification. "
                    f"Retrain the model or upgrade scikit-learn to fix this."
                )
                self.pipeline = None
                self.is_trained = False