src/ml/classifier.py

"""Wrapper AI untuk dipakai aplikasi: prediksi kategori + deteksi spam.

Menggantikan dropdown kategori manual (revisi software.md line 62) dan menjalankan
verifikasi kelayakan otomatis (revisi panduan.md UC-04). Hasil prediksi inilah
yang nanti dicek/dikonfirmasi oleh admin.
"""

from __future__ import annotations

from dataclasses import dataclass

import os
import joblib

from src.constants import (
    CATEGORY_MODEL_PATH,
    LABEL_LAYAK,
    LABEL_SPAM,
    SPAM_MODEL_PATH,
)


@dataclass
class AIResult:
    """Hasil analisis AI atas sebuah laporan."""

    kategori: str | None        # kategori prediksi (None jika terdeteksi spam)
    kategori_confidence: float  # 0..1
    spam_label: str             # LABEL_LAYAK / LABEL_SPAM
    spam_confidence: float      # 0..1 keyakinan terhadap label spam yang diberikan
    is_spam: bool
    prioritas: str | None = "Rendah"
    rekomendasi_ai: str | None = None
    tanggungjawab: str | None = "Pemerintah Kabupaten"


class Classifier:
    """Memuat model terlatih sekali, lalu melayani prediksi berulang."""

    def __init__(self, spam_model=None, category_model=None):
        self.spam_model = spam_model
        self.category_model = category_model

    @classmethod
    def load(cls) -> "Classifier":
        """Muat model dari disk. Latih otomatis bila belum ada."""
        if not SPAM_MODEL_PATH.exists() or not CATEGORY_MODEL_PATH.exists():
            from src.ml.train import train_and_save

            train_and_save()
        spam_model = joblib.load(SPAM_MODEL_PATH)
        category_model = joblib.load(CATEGORY_MODEL_PATH)
        return cls(spam_model=spam_model, category_model=category_model)

    def predict(self, text: str, attachment_path: str | None = None) -> AIResult:
        """Analisis satu laporan: deteksi spam lalu (jika layak) klasifikasi kategori."""
        text = (text or "").strip()
        if not text:
            return AIResult(None, 0.0, LABEL_SPAM, 1.0, True, "Rendah", None)

        # Coba gunakan OpenAI API jika key tersedia
        if os.environ.get("OPENAI_API_KEY"):
            try:
                from src.ml.openai_client import analyze_report
                openai_res = analyze_report(text, attachment_path)
                if openai_res:
                    return AIResult(
                        kategori=openai_res.kategori,
                        kategori_confidence=openai_res.kategori_confidence,
                        spam_label=LABEL_SPAM if openai_res.is_spam else LABEL_LAYAK,
                        spam_confidence=openai_res.spam_confidence,
                        is_spam=openai_res.is_spam,
                        prioritas=openai_res.urgensi,
                        rekomendasi_ai=openai_res.rekomendasi,
                        tanggungjawab=openai_res.tanggungjawab
                    )
            except Exception as e:
                print(f"[Warning] Gagal menjalankan analisis OpenAI, fallback ke model lokal: {e}")

        # 1) Deteksi spam secara lokal
        spam_proba = self.spam_model.predict_proba([text])[0]
        spam_classes = list(self.spam_model.classes_)  # mis. [0, 1]
        p_spam = spam_proba[spam_classes.index(1)]
        is_spam = p_spam >= 0.5
        spam_label = LABEL_SPAM if is_spam else LABEL_LAYAK
        spam_confidence = p_spam if is_spam else (1.0 - p_spam)

        # 2) Klasifikasi kategori (hanya bila laporan dianggap layak)
        kategori = None
        kategori_conf = 0.0
        if not is_spam:
            cat_proba = self.category_model.predict_proba([text])[0]
            best_idx = cat_proba.argmax()
            kategori = self.category_model.classes_[best_idx]
            kategori_conf = float(cat_proba[best_idx])

        return AIResult(
            kategori=kategori,
            kategori_confidence=round(kategori_conf, 4),
            spam_label=spam_label,
            spam_confidence=round(float(spam_confidence), 4),
            is_spam=bool(is_spam),
            prioritas="Rendah",
            rekomendasi_ai="Laporan Anda telah kami terima dan akan segera diproses.",
            tanggungjawab="Pemerintah Kabupaten"
        )