Blame for backend/src/predict.py

"""Inference module for emotion prediction using our trained model."""

import math
import os
from pathlib import Path
from typing import Optional, Tuple
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

from .config import config
from .visualize import EmotionVisualizer


# Sentinel: "use the predictor's configured prior", distinct from None ("no adjustment").
_INSTANCE_PRIOR = object()


class EmotionPredictor:
    """
    Predict emotions from text using our custom trained model.

    Trained on dair-ai/emotion dataset with 16,000+ samples.
    Emotions: sadness, joy, love, anger, fear, surprise
    """

    def __init__(
        self,
        model_path: Path = None,
        enable_viz: bool = True,
        target_sarcasm_prior: Optional[float] = config.target_sarcasm_prior,
        sarcasm_threshold: Optional[float] = config.sarcasm_threshold,
        train_sarcasm_prior: Optional[float] = None,
    ):
        """
        Initialize the emotion predictor.

        Args:
            model_path: Path to trained model directory.
                       Defaults to models/emotion_classifier/final
            enable_viz: Whether to enable chart visualization
            target_sarcasm_prior: Target sarcasm prevalence in deployment text (0-1)
            sarcasm_threshold: Optional decision threshold for sarcasm class (0-1)
            train_sarcasm_prior: Optional override for sarcasm prevalence seen in training
        """
        if not model_path:
            # Dynamically fetch the latest modified folder in models/ (or default to emotion_classifier/final)
            models_dir = config.model_dir
            subdirs = [d for d in models_dir.iterdir() if d.is_dir() and "emotion_classifier" in d.name]
            if subdirs:
                latest_dir = max(subdirs, key=lambda d: d.stat().st_mtime)
                # If there's a final folder inside it, use it
                if (latest_dir / "final").exists():
                    self.model_path = latest_dir / "final"
                else:
                    self.model_path = latest_dir
            else:
                self.model_path = config.model_dir / "emotion_classifier" / "final"
        else:
            self.model_path = Path(model_path)
            
        self.device = self._get_device()
        self.enable_viz = enable_viz
        self.target_sarcasm_prior = self._validate_probability(target_sarcasm_prior, "target_sarcasm_prior")
        self.sarcasm_threshold = self._validate_probability(sarcasm_threshold, "sarcasm_threshold", allow_none=True)
        self.user_train_sarcasm_prior = self._validate_probability(
            train_sarcasm_prior,
            "train_sarcasm_prior",
            allow_none=True,
        )

        # Initialize visualizer if enabled
        if self.enable_viz:
            self.visualizer = EmotionVisualizer()

        self._load_model()

    @staticmethod
    def _validate_probability(value: Optional[float], name: str, allow_none: bool = True) -> Optional[float]:
        """Validate probability-like arguments."""
        if value is None:
            if allow_none:
                return None
            raise ValueError(f"{name} cannot be None")

        value = float(value)
        if not 0.0 < value < 1.0:
            raise ValueError(f"{name} must be in the open interval (0, 1), got {value}")
        return value

    def _get_device(self):
        """Get the best available device."""
        if torch.cuda.is_available():
            return torch.device("cuda")
        elif torch.backends.mps.is_available():
            return torch.device("mps")
        return torch.device("cpu")

    def _load_model(self):
        """Load model and tokenizer, with optional Hugging Face Hub fallback.

        If the local model path doesn't exist, checks the HF_MODEL_REPO env var
        (e.g. "username/eumora-emotion-classifier") and downloads from HF Hub.
        Use HF_TOKEN for private repos.
        """
        if self.model_path.exists():
            source = str(self.model_path)
        else:
            hf_repo = os.environ.get("HF_MODEL_REPO")
            if hf_repo:
                source = hf_repo
                print(f"Local model not found. Loading from HuggingFace Hub: {hf_repo}")
            else:
                raise FileNotFoundError(
                    f"❌ No trained model found at {self.model_path}\n"
                    f"   Either run 'python main.py train' or set HF_MODEL_REPO env var."
                )

        hf_token = os.environ.get("HF_TOKEN") or None
        print(f"Loading model from: {source}")
        self.tokenizer = AutoTokenizer.from_pretrained(source, token=hf_token)
        self.model = AutoModelForSequenceClassification.from_pretrained(source, token=hf_token)
        self.model.to(self.device)
        self.model.eval()

        # Get label mappings from model config
        self.id2label = {int(k): v for k, v in self.model.config.id2label.items()}
        self.label2id = {k: int(v) for k, v in self.model.config.label2id.items()}
        self.sarcasm_idx = self.label2id.get("sarcasm")
        self.train_sarcasm_prior = self._resolve_train_sarcasm_prior()

        # Checkpoints from train_v2.py carry a temperature and a per-class prior
        # shift folded into the classifier weights. The sarcasm prior shift and
        # the 8% tripwire below were hand-tuned for the older, uncalibrated
        # model; applied on top of folded calibration they correct twice.
        # Measured on a calibrated checkpoint they changed 3% of labels and
        # shifted every sarcasm logit by -1.735, which /api/score then averages
        # into the track vector. Legacy checkpoints keep the old behaviour.
        calibration = getattr(self.model.config, "calibration", None) or {}
        self.calibration_folded = bool(calibration.get("folded"))

        if self.sarcasm_idx is None and (self.target_sarcasm_prior is not None or self.sarcasm_threshold is not None):
            print("⚠️ Loaded model has no 'sarcasm' class. Prior adjustment and sarcasm threshold are disabled.")

    def _resolve_train_sarcasm_prior(self) -> Optional[float]:
        """Resolve training sarcasm prior from explicit override, model metadata, or config fallback."""
        if self.user_train_sarcasm_prior is not None:
            return self.user_train_sarcasm_prior

        model_prior = getattr(self.model.config, "sarcasm_train_prior", None)
        if model_prior is not None:
            try:
                return self._validate_probability(model_prior, "model.sarcasm_train_prior", allow_none=False)
            except ValueError:
                pass

        if self.sarcasm_idx is not None:
            return self._validate_probability(
                config.assumed_train_sarcasm_prior,
                "config.assumed_train_sarcasm_prior",
                allow_none=False,
            )

        return None

    def _compute_sarcasm_logit_shift(self, target_prior=_INSTANCE_PRIOR) -> float:
        """Compute prior-shift logit adjustment for sarcasm vs non-sarcasm."""
        if getattr(self, "calibration_folded", False):
            return 0.0
        if target_prior is _INSTANCE_PRIOR:
            target = self.target_sarcasm_prior
        else:
            target = self._validate_probability(target_prior, "target_sarcasm_prior")
        if self.sarcasm_idx is None or target is None or self.train_sarcasm_prior is None:
            return 0.0

        target_odds = target / (1.0 - target)
        train_odds = self.train_sarcasm_prior / (1.0 - self.train_sarcasm_prior)
        return math.log(target_odds) - math.log(train_odds)

    def _apply_sarcasm_prior_adjustment(self, logits: torch.Tensor,
                                        target_prior=_INSTANCE_PRIOR) -> Tuple[torch.Tensor, float]:
        """Shift sarcasm logit to better match deployment prevalence."""
        shift = self._compute_sarcasm_logit_shift(target_prior)
        if self.sarcasm_idx is None or shift == 0.0:
            return logits, 0.0

        adjusted = logits.clone()
        adjusted[self.sarcasm_idx] = adjusted[self.sarcasm_idx] + shift
        return adjusted, shift

    def _apply_sarcasm_threshold(self, probs: torch.Tensor, pred_idx: int) -> int:
        """Apply optional one-vs-rest sarcasm thresholding to final class decision."""
        if self.sarcasm_idx is None or self.sarcasm_threshold is None:
            return pred_idx

        sarcasm_prob = probs[self.sarcasm_idx].item()
        if sarcasm_prob >= self.sarcasm_threshold:
            return self.sarcasm_idx

        if pred_idx == self.sarcasm_idx:
            non_sarcasm_probs = probs.clone()
            non_sarcasm_probs[self.sarcasm_idx] = -1.0
            return torch.argmax(non_sarcasm_probs).item()

        return pred_idx

    def _postprocess(self, logits: torch.Tensor, target_sarcasm_prior=_INSTANCE_PRIOR):
        """One row of logits -> (pred_idx, confidence, rounded probabilities, sarcasm shift).

        Shared by predict() and predict_probs_many(), so batched and single
        predictions go through identical adjustments and rounding.
        """
        adjusted_logits, sarcasm_logit_shift = self._apply_sarcasm_prior_adjustment(
            logits, target_sarcasm_prior)

        probs = torch.softmax(adjusted_logits, dim=-1)
        pred_idx = torch.argmax(probs).item()

        # Hardcode 8% sarcasm tripwire (legacy, uncalibrated checkpoints only)
        if (
            not self.calibration_folded
            and self.sarcasm_idx is not None
            and probs[self.sarcasm_idx].item() >= 0.08
        ):
            if pred_idx != self.sarcasm_idx:
                # Swap the max probability with sarcasm so it shows up as the dominant emotion visually
                original_max_val = probs[pred_idx].item()
                sarcasm_val = probs[self.sarcasm_idx].item()
                probs[pred_idx] = sarcasm_val
                probs[self.sarcasm_idx] = original_max_val
            pred_idx = self.sarcasm_idx

        confidence = probs[pred_idx].item()

        # Ensure probabilities sum to exactly 100% when displayed
        raw_probs = [probs[i].item() for i in range(len(probs))]

        # Convert to percentages and round, then normalize to ensure 100% total
        percent_probs = [p * 100 for p in raw_probs]
        rounded_percents = [round(p, 1) for p in percent_probs]

        # Adjust largest value to make sum exactly 100.0%
        total_percent = sum(rounded_percents)
        if total_percent != 100.0:
            max_idx = rounded_percents.index(max(rounded_percents))
            rounded_percents[max_idx] += (100.0 - total_percent)

        # Convert back to probability format (0-1 range)
        emotion_probs = {
            self.id2label[i]: round(rounded_percents[i] / 100.0, 4)
            for i in range(len(probs))
        }
        return pred_idx, confidence, emotion_probs, sarcasm_logit_shift

    def predict_probs_many(self, texts: list) -> list:
        """Probabilities and confidence for several texts in one forward pass.

        What lyric scoring needs, without the explanation and chart. One padded
        batch per track instead of one pass per stanza: 2.0x faster on two CPU
        threads, measured over 48 stanzas from the gold songs, which matters on
        a Space capped at 2 CPUs where DeBERTa is the scoring bottleneck.
        """
        if not texts:
            return []
        inputs = self.tokenizer(
            list(texts),
            return_tensors="pt",
            truncation=True,
            max_length=config.max_length,
            padding=True,
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        with torch.no_grad():
            logits = self.model(**inputs).logits.detach().cpu()
        out = []
        for row in logits:
            _, confidence, probs, _ = self._postprocess(row)
            out.append({"probabilities": probs, "confidence": round(confidence, 4)})
        return out

    def predict_raw(self, text: str) -> dict:
        """Plain softmax over the classifier: what a Hugging Face
        text-classification pipeline, and so the HF inference router, returns.
        No sarcasm prior shift, no tripwire, no rounding. The Vercel emotion
        route falls back to this when the router fails, and it must get the
        same numbers it would have got from the router."""
        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=config.max_length)
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        with torch.no_grad():
            probs = torch.softmax(self.model(**inputs).logits[0].float(), dim=-1).cpu().tolist()
        return {self.id2label[i]: p for i, p in enumerate(probs)}

    def predict(self, text: str, create_chart: bool = False, show_chart: bool = True,
                target_sarcasm_prior=_INSTANCE_PRIOR) -> dict:
        """
        Predict emotion from text.

        Args:
            text: Input text (lyrics, sentence, etc.)
            create_chart: Whether to generate a visualization chart
            show_chart: Whether to display the chart (only if create_chart=True)
            target_sarcasm_prior: Override the deployment sarcasm prior for this
                call only (None disables the adjustment). The API serves
                concurrent requests from one predictor, so a per-request prior
                must not be set on the instance: it would leak into every
                request that runs after it, including lyric scoring.

        Returns:
            dict with emotion, confidence, probabilities, explanation, and optional chart_path
        """
        # Tokenize
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            max_length=config.max_length,
            padding=True
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        # Predict
        with torch.no_grad():
            outputs = self.model(**inputs)

        logits = outputs.logits[0].detach().cpu()
        pred_idx, confidence, emotion_probs, sarcasm_logit_shift = self._postprocess(
            logits, target_sarcasm_prior)
        emotion = self.id2label[pred_idx]

        # Get music context
        music_context = config.emotion_to_music_mood.get(emotion, {
            "mood": emotion, "energy": "medium", "valence": "neutral"
        })

        result = {
            "emotion": emotion,
            "confidence": round(confidence, 4),
            "probabilities": emotion_probs,
            "music_context": music_context,
            "explanation": self._generate_explanation(emotion, confidence, emotion_probs, music_context),
            "calibration": {
                "target_sarcasm_prior": self.target_sarcasm_prior,
                "train_sarcasm_prior": self.train_sarcasm_prior,
                "sarcasm_logit_shift": round(sarcasm_logit_shift, 4),
                "sarcasm_threshold": self.sarcasm_threshold,
                "sarcasm_probability": round(emotion_probs.get("sarcasm", 0.0), 4),
            },
        }

        # Generate visualization if requested
        if create_chart and self.enable_viz:
            try:
                chart_path = self.visualizer.create_emotion_bar_chart(
                    emotion_probs, text, show_chart=show_chart, primary_emotion=emotion
                )
                result["chart_path"] = str(chart_path)
                print(f"Chart saved to: {chart_path}")
            except Exception as e:
                print(f"Could not create chart: {e}")
                result["chart_path"] = None

        return result

    def _generate_explanation(self, emotion: str, confidence: float,
                               probs: dict, music_context: dict) -> str:
        """Generate XAI-style explanation for the prediction."""
        confidence_level = "high" if confidence > 0.7 else "moderate" if confidence > 0.4 else "low"

        # Find secondary emotion
        sorted_emotions = sorted(probs.items(), key=lambda x: x[1], reverse=True)
        secondary = sorted_emotions[1] if len(sorted_emotions) > 1 else None

        # Emotion descriptors
        descriptors = {
            "sadness": "melancholic and sorrowful themes",
            "joy": "uplifting and celebratory content",
            "love": "romantic and affectionate sentiments",
            "anger": "intense and confrontational language",
            "fear": "anxious and uncertain undertones",
            "surprise": "unexpected and wonder-filled expressions",
            "neutral": "balanced and observational language",
            "sarcasm": "ironic or intentionally contradictory phrasing",
        }

        descriptor = descriptors.get(emotion, f"{emotion} emotional markers")

        explanation = (
            f"Detected {descriptor} with {confidence_level} confidence ({confidence:.1%}). "
            f"Suggests {music_context.get('mood', emotion)} music with "
            f"{music_context.get('energy', 'medium')} energy."
        )

        if secondary and secondary[1] > 0.15:
            explanation += f" Secondary: {secondary[0]} ({secondary[1]:.1%})."

        return explanation

    def predict_batch(self, texts: list) -> list:
        """Predict emotions for multiple texts."""
        return [self.predict(text) for text in texts]

    def analyze_song(self, title: str, artist: str, lyrics: str, create_chart: bool = True) -> dict:
        """Full song analysis with metadata and optional visualization."""
        prediction = self.predict(lyrics, create_chart=False)  # We'll create a detailed chart instead

        result = {
            "song": {"title": title, "artist": artist},
            "analysis": prediction,
            "tags": self._generate_tags(prediction),
        }

        # Generate detailed visualization for song analysis
        if create_chart and self.enable_viz:
            try:
                chart_path = self.visualizer.create_detailed_analysis_chart(
                    prediction, f"{title} by {artist}"
                )
                result["chart_path"] = str(chart_path)
                print(f"📊 Detailed analysis chart saved to: {chart_path}")
            except Exception as e:
                print(f"⚠️ Could not create detailed chart: {e}")
                result["chart_path"] = None

        return result

    def predict_with_visualization(self, text: str, chart_type: str = "simple") -> dict:
        """
        Predict with automatic visualization.

        Args:
            text: Input text
            chart_type: Type of chart ('simple', 'detailed')

        Returns:
            Prediction result with chart
        """
        if chart_type == "detailed":
            result = self.predict(text, create_chart=False)
            if self.enable_viz:
                try:
                    chart_path = self.visualizer.create_detailed_analysis_chart(result, text)
                    result["chart_path"] = str(chart_path)
                except Exception as e:
                    print(f"⚠️ Could not create detailed chart: {e}")
                    result["chart_path"] = None
        else:
            result = self.predict(text, create_chart=True, show_chart=True)

        return result

    def _generate_tags(self, prediction: dict) -> list:
        """Generate recommendation tags from prediction."""
        tags = []
        emotion = prediction["emotion"]
        music_ctx = prediction["music_context"]

        tags.append(f"emotion:{emotion}")
        tags.append(f"mood:{music_ctx.get('mood', emotion)}")
        tags.append(f"energy:{music_ctx.get('energy', 'medium')}")
        tags.append(f"valence:{music_ctx.get('valence', 'neutral')}")

        # Activity suggestions
        activity_map = {
            "joy": ["party", "workout", "celebration"],
            "sadness": ["reflection", "rainy-day", "comfort"],
            "love": ["romantic", "date-night", "slow-dance"],
            "anger": ["workout", "release", "intensity"],
            "fear": ["thriller", "suspense", "atmospheric"],
            "surprise": ["discovery", "adventure", "exploration"],
        }

        activities = activity_map.get(emotion, ["general"])
        tags.extend([f"activity:{a}" for a in activities])

        return tags


def demo():
    """Demo of the trained emotion predictor with visualizations."""
    print("\n" + "=" * 60)
    print("🎵 EUMORA - Emotion Analysis Demo (Custom Trained Model)")
    print("=" * 60)

    try:
        predictor = EmotionPredictor(enable_viz=True)
    except FileNotFoundError as e:
        print(f"\n{e}")
        return

    test_samples = [
        ("Happy lyrics", "I feel so alive today, everything is wonderful and bright!"),
        ("Sad lyrics", "My heart is broken, tears falling like rain in the dark night"),
        ("Angry lyrics", "I hate this, you betrayed me, I want to scream at the world!"),
        ("Love lyrics", "You are my everything, I want to hold you forever my darling"),
        ("Fear lyrics", "Something is watching me in the shadows, I'm scared to move"),
        ("Surprise lyrics", "I can't believe it happened! This is incredible, wow!"),
    ]

    print(f"\nAnalyzing {len(test_samples)} samples...\n")
    print("-" * 60)

    # Create comparison charts
    all_results = []

    for label, text in test_samples:
        result = predictor.predict(text, create_chart=False)  # Individual charts disabled for comparison
        all_results.append(result)

        print(f"\n📝 {label}:")
        print(f"   \"{text[:50]}...\"")
        print(f"   🎭 Emotion: {result['emotion'].upper()} ({result['confidence']:.1%})")
        print(f"   🎸 Context: {result['music_context']}")

        # Show probability distribution (text-based)
        sorted_probs = sorted(result['probabilities'].items(), key=lambda x: x[1], reverse=True)
        print(f"   📊 Distribution:")
        for emo, prob in sorted_probs[:3]:
            bar = "█" * int(prob * 20)
            print(f"      {emo:>10}: {bar:<20} {prob:.1%}")

        print("-" * 60)

    # Create comparison visualization
    print(f"\n📊 Generating comparison chart...")
    try:
        titles = [label for label, _ in test_samples]
        comparison_path = predictor.visualizer.create_comparison_chart(
            all_results, titles, show_chart=True
        )
        print(f"📊 Comparison chart saved to: {comparison_path}")
    except Exception as e:
        print(f"⚠️ Could not create comparison chart: {e}")

    print("\n✅ Demo complete!")


if __name__ == "__main__":
    demo()