"""
AIFinder Inference Module
Load the trained model and predict AI provider
"""

import joblib
import numpy as np

from config import MODEL_DIR


class AIFinder:
    def __init__(self, model_dir=MODEL_DIR):
        self.models = joblib.load(f"{model_dir}/rf_4provider.joblib")
        self.pipeline = joblib.load(f"{model_dir}/pipeline_4provider.joblib")
        self.le = joblib.load(f"{model_dir}/enc_4provider.joblib")

    def predict(self, text):
        """Predict the provider for a given text"""
        X = self.pipeline.transform([text])
        proba = np.mean([m.predict_proba(X) for m in self.models], axis=0)
        pred_idx = np.argmax(proba[0])
        return self.le.classes_[pred_idx]

    def predict_proba(self, text):
        """Get prediction probabilities"""
        X = self.pipeline.transform([text])
        proba = np.mean([m.predict_proba(X) for m in self.models], axis=0)
        return dict(zip(self.le.classes_, proba[0]))

    def predict_with_confidence(self, text):
        """Predict with confidence score"""
        proba = self.predict_proba(text)
        provider = max(proba, key=proba.get)
        confidence = proba[provider]
        return provider, confidence


if __name__ == "__main__":
    finder = AIFinder()

    test_texts = [
        "AI is like a really smart robot helper.",
        "Yes, coding is one of my stronger skills!",
        "A lot, depending on what you need.",
    ]

    for text in test_texts:
        provider, conf = finder.predict_with_confidence(text)
        print(f"Text: {text[:50]}...")
        print(f"Provider: {provider} (confidence: {conf:.2f})")
        print()