models/eq_adapter.py

"""
Music EQ (Emotional Intelligence) Adapter for TouchGrass.
Detects frustration and adapts responses for music learning context.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Dict, Tuple, List


class MusicEQAdapter(nn.Module):
    """
    Frustration detection adapted for music learning context.
    Music learners get frustrated differently than general users:
    - Finger pain/difficulty ("my fingers hurt", "I can't get this chord")
    - Rhythm frustration ("I keep losing the beat")
    - Progress frustration ("I've been practicing for weeks and still...")
    - Theory overwhelm ("this is too complicated")

    When frustration detected:
    - Simplify explanations automatically
    - Suggest easier alternatives ("try the open G chord instead of barre")
    - Add encouragement naturally
    - Break things into smaller steps
    - Remind them learning music takes time

    4-emotion classification for music context:
    frustrated, confused, excited, confident
    (simpler than general 8-emotion — music context needs fewer)
    """

    # Emotion labels
    EMOTIONS = ["frustrated", "confused", "excited", "confident"]

    # Frustration triggers (keywords/phrases)
    FRUSTRATION_TRIGGERS = [
        "can't", "cannot", "impossible", "too hard", "difficult",
        "fingers hurt", "pain", "hurt", "struggling", "stuck",
        "weeks", "months", "still can't", "giving up", "quit",
        "confused", "don't understand", "too complicated",
        "lost", "overwhelmed", "frustrated", "annoyed",
        "beat", "rhythm", "timing", "off beat",
        "barre", "stretch", "impossible chord",
    ]

    # Encouragement templates for frustrated learners
    ENCOURAGEMENT_TEMPLATES = {
        "frustrated": [
            "I understand this is challenging — learning {instrument} takes time and patience.",
            "Many students struggle with this at first. Let's break it down into smaller steps.",
            "Frustration is normal when learning something new. You're making progress, even if it doesn't feel like it.",
            "Every musician has been where you are. Keep going — it gets easier!",
        ],
        "confused": [
            "Let me explain that in a different way.",
            "I see this is confusing. Here's a simpler approach...",
            "Music theory can be overwhelming. Let's focus on one piece at a time.",
            "That's a great question! Let me break it down step by step.",
        ],
        "excited": [
            "I'm glad you're excited! That enthusiasm will help you learn faster.",
            "Your excitement is contagious! Let's keep that momentum going.",
            "That's the spirit! Music is a wonderful journey.",
        ],
        "confident": [
            "Great confidence! You're on the right track.",
            "Your progress shows you're getting the hang of this.",
            "Keep that confidence — it's key to musical growth.",
        ],
    }

    # Simplification strategies by emotion
    SIMPLIFICATION_STRATEGIES = {
        "frustrated": [
            "suggest_open_chord_alternative",
            "reduce_tempo",
            "break_into_parts",
            "use_easier_tuning",
            "skip_complex_theory",
        ],
        "confused": [
            "use_analogy",
            "show_visual_example",
            "step_by_step",
            "check_prerequisites",
        ],
        "excited": [
            "add_challenge",
            "introduce_next_concept",
            "suggest_creative_exercise",
        ],
        "confident": [
            "maintain_pace",
            "introduce_advanced_topics",
            "suggest_performance_opportunities",
        ],
    }

    def __init__(self, d_model: int, eq_hidden: int = 32):
        """
        Initialize MusicEQAdapter.

        Args:
            d_model: Hidden dimension from base model
            eq_hidden: Hidden dimension for EQ layers (small, lightweight)
        """
        super().__init__()
        self.d_model = d_model
        self.eq_hidden = eq_hidden

        # Frustration detector (binary: frustrated or not)
        self.frustration_detector = nn.Sequential(
            nn.Linear(d_model, eq_hidden),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(eq_hidden, 1),
            nn.Sigmoid()
        )

        # 4-emotion classifier for music context
        self.emotion_classifier = nn.Sequential(
            nn.Linear(d_model, eq_hidden),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(eq_hidden, 4),
        )

        # Simplification gate: modulates response complexity
        # Takes: frustration_score + 4 emotion probs = 5 inputs
        self.simplify_gate = nn.Sequential(
            nn.Linear(5, eq_hidden),
            nn.ReLU(),
            nn.Linear(eq_hidden, d_model),
            nn.Sigmoid()  # Output 0-1 per dimension
        )

        # EQ loss weight (for training)
        self.eq_loss_weight = 0.1

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> Dict[str, torch.Tensor]:
        """
        Forward pass through MusicEQAdapter.

        Args:
            hidden_states: Base model hidden states [batch, seq_len, d_model]
            attention_mask: Attention mask [batch, seq_len]

        Returns:
            Dictionary with emotion predictions and simplification gate
        """
        batch_size, seq_len, d_model = hidden_states.shape

        # Pool hidden states (weighted by attention mask if provided)
        if attention_mask is not None:
            # Mask-based pooling
            mask_expanded = attention_mask.unsqueeze(-1).float()
            pooled = (hidden_states * mask_expanded).sum(dim=1) / mask_expanded.sum(dim=1)
        else:
            pooled = hidden_states.mean(dim=1)  # [batch, d_model]

        # Detect frustration (0-1 score)
        frustration_score = self.frustration_detector(pooled)  # [batch, 1]

        # Classify emotion (4 classes)
        emotion_logits = self.emotion_classifier(pooled)  # [batch, 4]
        emotion_probs = F.softmax(emotion_logits, dim=-1)

        # Compute simplification gate input
        simplify_input = torch.cat([frustration_score, emotion_probs], dim=1)  # [batch, 5]

        # Generate simplification gate (per-dimension modulation)
        simplify_gate = self.simplify_gate(simplify_input)  # [batch, d_model]

        outputs = {
            "frustration_score": frustration_score,
            "emotion_logits": emotion_logits,
            "emotion_probs": emotion_probs,
            "simplify_gate": simplify_gate,
        }

        return outputs

    def detect_frustration(
        self,
        text: str,
        threshold: float = 0.5,
    ) -> Tuple[bool, float, str]:
        """
        Detect frustration in user text (rule-based fallback).

        Args:
            text: User input text
            threshold: Frustration score threshold

        Returns:
            (is_frustrated, score, detected_emotion)
        """
        text_lower = text.lower()

        # Count frustration triggers
        trigger_count = sum(1 for trigger in self.FRUSTRATION_TRIGGERS if trigger in text_lower)

        # Simple scoring
        score = min(1.0, trigger_count / 5.0)  # Normalize to 0-1

        is_frustrated = score >= threshold

        # Determine emotion (simplified rule-based)
        if "confused" in text_lower or "don't understand" in text_lower:
            emotion = "confused"
        elif "excited" in text_lower or "love" in text_lower or "awesome" in text_lower:
            emotion = "excited"
        elif "got it" in text_lower or "understand" in text_lower or "easy" in text_lower:
            emotion = "confident"
        else:
            emotion = "frustrated" if is_frustrated else "neutral"

        return is_frustrated, score, emotion

    def get_encouragement(
        self,
        emotion: str,
        instrument: Optional[str] = None,
        context: Optional[str] = None,
    ) -> str:
        """
        Generate encouragement message based on detected emotion.

        Args:
            emotion: Detected emotion (frustrated, confused, excited, confident)
            instrument: Optional instrument context
            context: Optional specific context (chord, theory, etc)

        Returns:
            Encouragement string
        """
        import random

        if emotion not in self.ENCOURAGEMENT_TEMPLATES:
            emotion = "frustrated"  # Default

        templates = self.ENCOURAGEMENT_TEMPLATES[emotion]
        template = random.choice(templates)

        # Fill in instrument placeholder if present
        if "{instrument}" in template and instrument:
            return template.format(instrument=instrument)
        else:
            return template

    def get_simplification_strategy(
        self,
        emotion: str,
        instrument: Optional[str] = None,
        user_level: str = "beginner",
    ) -> List[str]:
        """
        Get list of simplification strategies to apply.

        Args:
            emotion: Detected emotion
            instrument: Optional instrument context
            user_level: User skill level

        Returns:
            List of strategy names
        """
        strategies = self.SIMPLIFICATION_STRATEGIES.get(emotion, [])

        # Add level-specific strategies
        if user_level == "beginner":
            strategies.append("use_basic_terminology")
            strategies.append("avoid_music_jargon")

        return strategies

    def apply_simplification(
        self,
        response_text: str,
        strategies: List[str],
        emotion: str,
    ) -> str:
        """
        Apply simplification strategies to response text.

        Args:
            response_text: Original response
            strategies: List of strategies to apply
            emotion: Detected emotion

        Returns:
            Simplified response
        """
        simplified = response_text

        for strategy in strategies:
            if strategy == "suggest_open_chord_alternative":
                # Replace barre chords with open alternatives
                simplified = self._replace_barre_with_open(simplified)
            elif strategy == "reduce_tempo":
                # Add tempo suggestion
                if "BPM" in simplified or "tempo" in simplified:
                    simplified += "\n\nTip: Try practicing this at a slower tempo (60-80 BPM) and gradually increase."
            elif strategy == "break_into_parts":
                # Add step-by-step suggestion
                simplified = "Let's break this down:\n\n" + simplified
            elif strategy == "skip_complex_theory":
                # Simplify theory explanations
                simplified = self._simplify_theory(simplified)
            elif strategy == "use_analogy":
                # Add analogies
                simplified = self._add_analogy(simplified)
            elif strategy == "step_by_step":
                # Add numbered steps
                simplified = self._add_numbered_steps(simplified)

        # Prepend encouragement if frustrated
        if emotion == "frustrated":
            encouragement = self.get_encouragation("frustrated")
            simplified = encouragement + "\n\n" + simplified

        return simplified

    def _replace_barre_with_open(self, text: str) -> str:
        """Replace barre chord suggestions with open alternatives."""
        replacements = {
            "F major": "F major (try Fmaj7 or F/C if barre is hard)",
            "B minor": "B minor (try Bm7 or alternative fingering)",
            "barre": "barre (you can also try a partial barre or capo)",
        }
        for original, replacement in replacements.items():
            text = text.replace(original, replacement)
        return text

    def _simplify_theory(self, text: str) -> str:
        """Simplify music theory explanations."""
        # Replace complex terms with simpler explanations
        simplifications = {
            "diatonic": "within the key",
            "chromatic": "all 12 notes",
            "modulation": "changing key",
            "cadence": "ending chord progression",
            "arpeggio": "playing chord notes one at a time",
        }
        for complex_term, simple_term in simplifications.items():
            text = text.replace(complex_term, simple_term)
        return text

    def _add_analogy(self, text: str) -> str:
        """Add musical analogies to explanation."""
        analogy = "\n\nThink of it like this: music is a language — you learn the alphabet (notes), then words (chords), then sentences (progressions)."
        return text + analogy

    def _add_numbered_steps(self, text: str) -> str:
        """Convert paragraph to numbered steps."""
        # Simple implementation: add numbered list if not already
        if "1." not in text and "Step" not in text:
            lines = text.split("\n")
            new_lines = []
            step_num = 1
            for line in lines:
                if line.strip() and not line.strip().startswith(("##", "**", "-", "*")):
                    new_lines.append(f"{step_num}. {line}")
                    step_num += 1
                else:
                    new_lines.append(line)
            return "\n".join(new_lines)
        return text

    def compute_eq_loss(
        self,
        outputs: Dict[str, torch.Tensor],
        emotion_labels: torch.Tensor,
        frustration_labels: torch.Tensor,
    ) -> torch.Tensor:
        """
        Compute EQ training loss.

        Args:
            outputs: Forward pass outputs
            emotion_labels: Ground truth emotion labels [batch]
            frustration_labels: Ground truth frustration labels [batch]

        Returns:
            EQ loss
        """
        # Emotion classification loss
        emotion_logits = outputs["emotion_logits"]
        emotion_loss = F.cross_entropy(emotion_logits, emotion_labels)

        # Frustration detection loss (binary cross-entropy)
        frustration_score = outputs["frustration_score"].squeeze()
        frustration_loss = F.binary_cross_entropy(frustration_score, frustration_labels.float())

        # Combined EQ loss
        eq_loss = emotion_loss + frustration_loss

        return eq_loss * self.eq_loss_weight


def test_eq_adapter():
    """Test the MusicEQAdapter."""
    import torch

    # Create adapter
    d_model = 4096
    adapter = MusicEQAdapter(d_model=d_model, eq_hidden=32)

    # Test input
    batch_size = 2
    seq_len = 20
    hidden_states = torch.randn(batch_size, seq_len, d_model)
    attention_mask = torch.ones(batch_size, seq_len)

    # Forward pass
    outputs = adapter.forward(hidden_states, attention_mask)

    print("Music EQ Adapter outputs:")
    for key, value in outputs.items():
        if isinstance(value, torch.Tensor):
            print(f"  {key}: {value.shape}")
        else:
            print(f"  {key}: {value}")

    # Test frustration detection
    print("\nFrustration detection (rule-based):")
    test_texts = [
        "I've been trying this chord for an hour and I still can't get it",
        "This is so confusing, I don't understand music theory",
        "I'm so excited to learn guitar!",
        "I think I'm getting the hang of this",
    ]
    for text in test_texts:
        is_frustrated, score, emotion = adapter.detect_frustration(text)
        print(f"  '{text[:50]}...' -> frustrated={is_frustrated}, score={score:.2f}, emotion={emotion}")

    # Test encouragement generation
    print("\nEncouragement messages:")
    for emotion in ["frustrated", "confused", "excited", "confident"]:
        msg = adapter.get_encouragement(emotion, instrument="guitar")
        print(f"  {emotion}: {msg[:80]}...")

    # Test simplification
    print("\nSimplification example:")
    original = "To play an F major barre chord, place your index finger across all six strings at the first fret..."
    strategies = ["suggest_open_chord_alternative", "break_into_parts"]
    simplified = adapter.apply_simplification(original, strategies, "frustrated")
    print(f"  Original: {original[:60]}...")
    print(f"  Simplified: {simplified[:80]}...")

    # Test loss computation
    print("\nEQ loss computation:")
    emotion_labels = torch.tensor([0, 2])  # frustrated, excited
    frustration_labels = torch.tensor([1.0, 0.0])  # first frustrated, second not
    eq_loss = adapter.compute_eq_loss(outputs, emotion_labels, frustration_labels)
    print(f"  EQ loss: {eq_loss.item():.4f}")

    print("\nMusic EQ Adapter test complete!")


if __name__ == "__main__":
    test_eq_adapter()