model/predict.py

"""
AES-Feedback Inference Pipeline
Combines IndoBERT scoring + IndoSBERT coherence + Feedback Engine.

Pair Encoding Mode (BARU):
- Jika model adalah best_model_pair/, predictor auto-load kunci_jawaban
- predict() secara otomatis lookup kunci_jawaban berdasarkan id_soal
- Tokenizer menerima (jawaban_siswa, kunci_jawaban) sebagai pair input
- Mode single-text (model lama) tetap berfungsi tanpa perubahan

Soal Matching (BARU v2):
- predict() menerima parameter `soal` untuk auto-match id_soal & kunci_jawaban
- predict() menerima parameter `kunci_jawaban` untuk override manual
- SoalMatcher menggunakan IndoSBERT untuk match input soal dengan dataset prompts

Usage:
    python -m model.predict
"""
import os
import pandas as pd
import torch
from typing import Optional
from transformers import AutoTokenizer, AutoModelForSequenceClassification

from model.config import (
    SAVED_MODELS_DIR, SAVED_MODELS_DIR_PAIR, MAX_SEQ_LENGTH,
    NUM_LABELS, INDOBERT_MODEL_NAME, SCORE_RANGES, DATASET_PATH_PAIR,
)
from model.semantic_analyzer import SemanticAnalyzer
from model.feedback_engine import FeedbackEngine
from model.soal_matcher import SoalMatcher


class AESPredictor:
    """
    Complete AES prediction pipeline.
    
    Loads:
    1. Fine-tuned IndoBERT for score prediction
    2. IndoSBERT for coherence analysis
    3. Rule-based feedback engine
    4. SoalMatcher for soal-to-prompt matching (BARU v2)
    
    Pair Encoding:
    - Jika model_dir adalah best_model_pair/, mapping kunci_jawaban
      otomatis di-load dari dataset_indonesia_pair.csv
    - predict() secara otomatis lookup kunci_jawaban berdasarkan id_soal
    """

    def __init__(self, model_dir=None):
        """
        Initialize the prediction pipeline.
        
        Args:
            model_dir: Path to the fine-tuned IndoBERT model directory.
                       Defaults to saved_models/best_model.
                       Jika best_model_pair/, auto-enable pair encoding.
        """
        if model_dir is None:
            model_dir = os.path.join(SAVED_MODELS_DIR, "best_model")

        self.pair_encoding = "best_model_pair" in model_dir

        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        print(f"Loading IndoBERT from: {model_dir}")
        if not os.path.exists(model_dir):
            raise FileNotFoundError(
                f"Model not found at {model_dir}. "
                "Please run 'python -m model.train' first."
            )
        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
        self.scoring_model = AutoModelForSequenceClassification.from_pretrained(
            model_dir, num_labels=NUM_LABELS
        )
        self.scoring_model.to(self.device)
        self.scoring_model.eval()
        print("IndoBERT loaded successfully.")

        self.kunci_jawaban_map = {}
        if self.pair_encoding:
            if os.path.exists(DATASET_PATH_PAIR):
                df = pd.read_csv(DATASET_PATH_PAIR)
                for soal in df['id_soal'].unique():
                    kunci = df[df['id_soal'] == soal]['kunci_jawaban'].iloc[0]
                    self.kunci_jawaban_map[soal] = kunci
                    print(f"  Kunci jawaban loaded for {soal}: {kunci[:60]}...")
            else:
                print(f"  [WARNING] Dataset pair tidak ditemukan di {DATASET_PATH_PAIR}")

        self.semantic_analyzer = SemanticAnalyzer()
        self.feedback_engine = FeedbackEngine()
        self.soal_matcher = SoalMatcher()

        print("[OK] AES Prediction pipeline ready!")

    def predict_score_raw(self, essay_text, text_pair=None):
        """
        Predict raw class using fine-tuned IndoBERT.

        Returns class 0-4 (not yet normalized or denormalized).
        
        Args:
            essay_text: Essay text string (jawaban_siswa)
            text_pair: Optional paired text (kunci_jawaban) for pair encoding.
                       Jika disediakan, tokenizer akan encode (essay, text_pair).
            
        Returns:
            int: Predicted class (0-4)
        """
        if text_pair is not None:
            encoding = self.tokenizer(
                essay_text,
                text_pair=text_pair,
                max_length=MAX_SEQ_LENGTH,
                padding="max_length",
                truncation=True,
                return_tensors="pt",
            )
        else:
            encoding = self.tokenizer(
                essay_text,
                max_length=MAX_SEQ_LENGTH,
                padding="max_length",
                truncation=True,
                return_tensors="pt",
            )

        input_ids = encoding["input_ids"].to(self.device)
        attention_mask = encoding["attention_mask"].to(self.device)

        with torch.no_grad():
            outputs = self.scoring_model(
                input_ids=input_ids,
                attention_mask=attention_mask
            )
            pred = torch.argmax(outputs.logits, dim=-1).item()

        return pred

    def normalized_from_class(self, class_id):
        """
        Convert class 0-4 to normalized score [0, 1].
        
        Maps classes evenly across [0, 1]:
            0 → 0.0,  1 → 0.25,  2 → 0.5,  3 → 0.75,  4 → 1.0
        """
        return class_id / 4.0

    def denormalize(self, normalized_score, id_soal):
        """
        Convert normalized score [0, 1] back to original scale.
        
        Args:
            normalized_score: Float in [0, 1]
            id_soal: Question ID (Q01/Q02/Q03/Q04)
            
        Returns:
            int: Score in original scale (1-3 for Q01/Q03, 1-4 for Q02/Q04)
        """
        min_score, max_score = SCORE_RANGES.get(id_soal, (1, 5))
        return int(normalized_score * (max_score - min_score) + min_score + 0.5)

    def predict(self, essay_text, id_soal="Q01", soal=None, kunci_jawaban=None):
        """
        Full prediction: score + coherence analysis + formative feedback.
        
        Args:
            essay_text: Essay text string (jawaban_siswa)
            id_soal: Question ID (Q01/Q02/Q03/Q04) for score denormalization
            soal: Optional soal text for auto-matching (BARU v2)
            kunci_jawaban: Optional kunci_jawaban override (BARU v2)
            
        Returns:
            dict with:
                - overall_score (int, denormalized to original scale)
                - normalized_score (float, [0, 1])
                - max_score (int, max possible score for this question)
                - id_soal (str, echoed back)
                - coherence (dict with coherence metrics)
                - feedback (dict with 3 aspect feedbacks)
                - matched (bool, whether soal was matched) [BARU v2]
        """
        # BARU v2: Soal matching logic
        matched = False
        if soal is not None:
            match = self.soal_matcher.match(soal)
            if match is not None:
                id_soal = match.id_soal
                if kunci_jawaban is None:
                    kunci_jawaban = match.kunci_jawaban
                matched = True

        # 1. Predict raw class
        text_pair = None
        if kunci_jawaban is not None:
            text_pair = kunci_jawaban
        elif self.pair_encoding:
            text_pair = self.kunci_jawaban_map.get(id_soal)
            if text_pair is None:
                print(f"  [WARNING] Kunci jawaban tidak ditemukan untuk {id_soal}")

        class_id = self.predict_score_raw(essay_text, text_pair=text_pair)

        normalized = self.normalized_from_class(class_id)

        _, max_score = SCORE_RANGES.get(id_soal, (1, 5))
        overall_score = self.denormalize(normalized, id_soal)

        coherence = self.semantic_analyzer.analyze(essay_text)
        feedback = self.feedback_engine.generate(essay_text, overall_score, coherence)

        return {
            "overall_score": overall_score,
            "normalized_score": round(normalized, 4),
            "max_score": max_score,
            "id_soal": id_soal,
            "coherence": coherence,
            "feedback": feedback,
            "matched": matched,
        }


def main():
    """Interactive demo for testing the prediction pipeline."""
    print("=" * 60)
    print("  AES-Feedback: Prediction Demo")
    print("=" * 60)

    predictor = AESPredictor()

    sample_essay = (
        "Pendidikan kewarganegaraan sangat penting bagi siswa karena melalui "
        "pendidikan ini siswa dapat memahami hak dan kewajiban sebagai warga negara. "
        "Pertama, siswa belajar tentang nilai-nilai Pancasila yang menjadi dasar "
        "negara Indonesia. Kedua, siswa memahami pentingnya toleransi dalam "
        "kehidupan bermasyarakat. Misalnya, di sekolah kita harus menghormati "
        "teman yang berbeda agama dan suku. Selain itu, pendidikan kewarganegaraan "
        "juga mengajarkan tentang demokrasi, sehingga siswa dapat berpartisipasi "
        "secara aktif dalam kehidupan berbangsa. Oleh karena itu, mata pelajaran "
        "PKN harus terus diajarkan di sekolah."
    )

    print("\n[Sample Essay]:")
    print(f"  {sample_essay[:200]}...")
    print()

    result = predictor.predict(sample_essay, id_soal="Q01")

    print(f"\n[Overall Score]: {result['overall_score']}/{result['max_score']} (normalized: {result['normalized_score']})")
    print(f"\n[Coherence Level]: {result['coherence']['coherence_level']}")
    print(f"   Adjacent Coherence: {result['coherence']['adjacent_coherence']:.4f}")

    fb = result["feedback"]
    print(f"\n[Argument Structure [{fb['argument_structure']['level']}]]:")
    print(f"   {fb['argument_structure']['feedback']}")

    print(f"\n[Reasoning [{fb['reasoning']['level']}]]:")
    print(f"   {fb['reasoning']['feedback']}")

    print(f"\n[Evidence Use [{fb['evidence_use']['level']}]]:")
    print(f"   {fb['evidence_use']['feedback']}")

    print("-" * 60)
    print("Masukkan id_soal (Q01/Q02/Q03/Q04) dan essay (pisahkan dengan '|'):")
    print("  Contoh: Q01 | Pendidikan kewarganegaraan sangat penting...\n")

    while True:
        raw = input("Soal|Essay> ").strip()
        if raw.lower() == "quit":
            break
        if not raw:
            continue

        if "|" in raw:
            id_soal, essay = [x.strip() for x in raw.split("|", 1)]
        else:
            id_soal = "Q01"
            essay = raw

        result = predictor.predict(essay, id_soal=id_soal)
        print(f"\n  Score: {result['overall_score']}/{result['max_score']} [{result['id_soal']}]")
        fb = result["feedback"]
        print(f"  Argument: [{fb['argument_structure']['level']}] {fb['argument_structure']['feedback']}")
        print(f"  Reasoning: [{fb['reasoning']['level']}] {fb['reasoning']['feedback']}")
        print(f"  Evidence: [{fb['evidence_use']['level']}] {fb['evidence_use']['feedback']}")
        print()


if __name__ == "__main__":
    main()