Lyra/lyra/emotions/expressions.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import random
import numpy as np
from typing import Dict, List, Any, Optional, Tuple
import re
import logging

from .system import EmotionalState

logger = logging.getLogger(__name__)

class EmotionalExpressionEngine(nn.Module):
    """
    Advanced system for expressing emotions naturally in text responses.

    This engine translates Lyra's internal emotional state into human-like
    emotional expression patterns, including word choice, punctuation,
    formatting, and even intentional typos when emotionally excited.
    """

    def __init__(
        self,
        vocab_size: int = 50000,
        expression_dim: int = 128,
        device: Optional[torch.device] = None
    ):
        super().__init__()

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

        # Emotion-to-expression mapping network
        self.emotion_mapper = nn.Sequential(
            nn.Linear(19, 64),  # 19 emotional dimensions
            nn.LayerNorm(64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, expression_dim)
        )

        # Expression style generators
        self.punctuation_generator = nn.Linear(expression_dim, 10)  # Different punctuation styles
        self.emphasis_generator = nn.Linear(expression_dim, 8)     # Emphasis patterns
        self.word_choice_generator = nn.Linear(expression_dim, 12) # Word choice modifications

        # Emotional vocabulary mappings
        self.emotional_vocabularies = self._initialize_emotional_vocabularies()

        # Expression patterns for different emotions
        self.expression_patterns = self._initialize_expression_patterns()

        # Human-like inconsistency parameters
        self.typo_probability = nn.Parameter(torch.tensor(0.02))
        self.excitement_threshold = nn.Parameter(torch.tensor(0.7))

        self.to(self.device)

    def _initialize_emotional_vocabularies(self) -> Dict[str, Dict[str, List[str]]]:
        """Initialize emotion-specific vocabularies."""
        return {
            'joy': {
                'intensifiers': ['absolutely', 'totally', 'completely', 'incredibly', 'amazingly'],
                'exclamations': ['wow', 'yay', 'awesome', 'fantastic', 'brilliant'],
                'adjectives': ['wonderful', 'amazing', 'fantastic', 'incredible', 'delightful'],
                'expressions': ['I love this!', 'This is so cool!', 'Amazing!', 'Wonderful!']
            },
            'sadness': {
                'softeners': ['I guess', 'maybe', 'perhaps', 'I suppose'],
                'expressions': ['I feel sad about this', 'This makes me a bit down', 'That\'s disappointing'],
                'adjectives': ['disappointing', 'unfortunate', 'sad', 'melancholy'],
                'hesitations': ['well...', 'I mean...', 'it\'s just that...']
            },
            'anger': {
                'intensifiers': ['absolutely', 'completely', 'totally', 'seriously'],
                'expressions': ['That\'s not okay', 'I don\'t like this', 'This is frustrating'],
                'exclamations': ['No way!', 'Seriously?!', 'Come on!'],
                'adjectives': ['frustrating', 'annoying', 'ridiculous', 'unacceptable']
            },
            'fear': {
                'hesitations': ['I\'m not sure...', 'Maybe...', 'I worry that...'],
                'expressions': ['I\'m a bit worried', 'This concerns me', 'I\'m nervous about this'],
                'softeners': ['perhaps', 'possibly', 'might be']
            },
            'surprise': {
                'exclamations': ['Oh!', 'Wow!', 'Really?!', 'No way!', 'Seriously?!'],
                'expressions': ['I didn\'t expect that!', 'That\'s surprising!', 'Whoa!'],
                'questions': ['Really?', 'Are you serious?', 'Wait, what?']
            },
            'curiosity': {
                'questions': ['How does that work?', 'What do you think?', 'Tell me more!'],
                'expressions': ['I\'m curious about...', 'I wonder...', 'That\'s interesting...'],
                'intensifiers': ['really', 'very', 'quite', 'particularly']
            },
            'love': {
                'warmers': ['I really care about', 'I love how', 'I adore', 'I cherish'],
                'expressions': ['You\'re wonderful', 'I appreciate you', 'That means a lot'],
                'softeners': ['sweetly', 'gently', 'warmly', 'tenderly']
            },
            'pride': {
                'expressions': ['I\'m proud of', 'That\'s impressive', 'Well done!', 'I accomplished'],
                'intensifiers': ['really', 'truly', 'genuinely', 'absolutely']
            }
        }

    def _initialize_expression_patterns(self) -> Dict[str, Dict[str, Any]]:
        """Initialize expression patterns for different emotions."""
        return {
            'joy': {
                'punctuation': ['!', '!!', '! :)', '! 😊'],
                'capitalization': 0.2,  # 20% chance of enthusiastic caps
                'repetition': ['so so', 'really really', 'very very'],
                'typo_reduction': 0.5,  # Less typos when happy
                'exclamation_frequency': 0.4
            },
            'sadness': {
                'punctuation': ['.', '...', '. :('],
                'capitalization': 0.0,  # No enthusiastic caps
                'hesitations': ['um...', 'well...', 'I guess...'],
                'typo_increase': 1.2,  # Slightly more typos when sad
                'trailing_offs': ['...', '..', '.']
            },
            'anger': {
                'punctuation': ['!', '!!', '!!!'],
                'capitalization': 0.3,  # More caps when angry
                'repetition': ['absolutely', 'totally'],
                'emphasis': ['*frustrated*', '*annoyed*'],
                'exclamation_frequency': 0.6
            },
            'fear': {
                'punctuation': ['.', '...', '?'],
                'hesitations': ['I think...', 'Maybe...', 'I\'m not sure...'],
                'questions': 0.3,  # More questions when uncertain
                'softening': 0.4
            },
            'surprise': {
                'punctuation': ['!', '?!', '!!', '?!!'],
                'capitalization': 0.4,  # High caps for surprise
                'exclamations': ['Whoa!', 'Oh!', 'Wow!'],
                'questions': 0.5
            },
            'curiosity': {
                'punctuation': ['?', '!', '...'],
                'questions': 0.6,  # High question frequency
                'thinking': ['Hmm...', 'I wonder...', 'Interesting...'],
                'exploration': ['What if...', 'How about...', 'Maybe...']
            }
        }

    def forward(
        self,
        text: str,
        emotional_state: EmotionalState,
        intensity_multiplier: float = 1.0,
        context: Optional[str] = None
    ) -> Tuple[str, Dict[str, Any]]:
        """
        Apply emotional expression to text based on current emotional state.

        Args:
            text: Base text to emotionally express
            emotional_state: Current emotional state
            intensity_multiplier: Multiplier for emotional expression intensity
            context: Context for expression (formal, casual, etc.)

        Returns:
            expressed_text: Text with emotional expression applied
            expression_info: Information about applied expressions
        """
        # Convert emotional state to tensor
        emotion_tensor = emotional_state.to_tensor(self.device).unsqueeze(0)

        # Generate expression features
        expression_features = self.emotion_mapper(emotion_tensor)

        # Generate expression parameters
        punctuation_weights = torch.softmax(self.punctuation_generator(expression_features), dim=-1)
        emphasis_weights = torch.softmax(self.emphasis_generator(expression_features), dim=-1)
        word_choice_weights = torch.softmax(self.word_choice_generator(expression_features), dim=-1)

        # Get dominant emotion for pattern selection
        dominant_emotion, emotion_intensity = emotional_state.get_dominant_emotion()

        # Apply expression modifications
        expressed_text = text
        expression_info = {'modifications': []}

        # Apply emotional vocabulary
        expressed_text, vocab_mods = self._apply_emotional_vocabulary(
            expressed_text, dominant_emotion, emotion_intensity * intensity_multiplier
        )
        expression_info['modifications'].extend(vocab_mods)

        # Apply punctuation patterns
        expressed_text, punct_mods = self._apply_punctuation_patterns(
            expressed_text, dominant_emotion, emotion_intensity * intensity_multiplier
        )
        expression_info['modifications'].extend(punct_mods)

        # Apply emphasis patterns
        expressed_text, emphasis_mods = self._apply_emphasis_patterns(
            expressed_text, dominant_emotion, emotion_intensity * intensity_multiplier
        )
        expression_info['modifications'].extend(emphasis_mods)

        # Apply human-like inconsistencies
        expressed_text, inconsistency_mods = self._apply_human_inconsistencies(
            expressed_text, emotional_state, intensity_multiplier
        )
        expression_info['modifications'].extend(inconsistency_mods)

        # Apply contextual adjustments
        if context:
            expressed_text, context_mods = self._apply_contextual_adjustments(
                expressed_text, context, emotional_state
            )
            expression_info['modifications'].extend(context_mods)

        expression_info.update({
            'dominant_emotion': dominant_emotion,
            'emotion_intensity': emotion_intensity,
            'total_modifications': len(expression_info['modifications']),
            'expression_strength': intensity_multiplier
        })

        return expressed_text, expression_info

    def _apply_emotional_vocabulary(
        self,
        text: str,
        emotion: str,
        intensity: float
    ) -> Tuple[str, List[str]]:
        """Apply emotion-specific vocabulary modifications."""
        modifications = []

        if emotion not in self.emotional_vocabularies:
            return text, modifications

        vocab = self.emotional_vocabularies[emotion]

        # Apply with probability based on intensity
        application_prob = min(0.8, intensity * 0.6)

        # Add emotional expressions
        if 'expressions' in vocab and random.random() < application_prob:
            if random.random() < 0.3:  # 30% chance to add expression
                expression = random.choice(vocab['expressions'])
                if random.random() < 0.5:
                    text = expression + ' ' + text
                else:
                    text = text + ' ' + expression
                modifications.append(f"Added expression: {expression}")

        # Modify intensifiers
        if 'intensifiers' in vocab and intensity > 0.6:
            intensifiers = ['very', 'really', 'quite', 'pretty']
            for intensifier in intensifiers:
                if intensifier in text and random.random() < application_prob:
                    new_intensifier = random.choice(vocab['intensifiers'])
                    text = text.replace(intensifier, new_intensifier, 1)
                    modifications.append(f"Replaced '{intensifier}' with '{new_intensifier}'")

        # Add exclamations for high-energy emotions
        if 'exclamations' in vocab and intensity > 0.7:
            if random.random() < application_prob * 0.5:
                exclamation = random.choice(vocab['exclamations'])
                text = text + ' ' + exclamation
                modifications.append(f"Added exclamation: {exclamation}")

        return text, modifications

    def _apply_punctuation_patterns(
        self,
        text: str,
        emotion: str,
        intensity: float
    ) -> Tuple[str, List[str]]:
        """Apply emotion-specific punctuation patterns."""
        modifications = []

        if emotion not in self.expression_patterns:
            return text, modifications

        patterns = self.expression_patterns[emotion]

        # Modify ending punctuation
        if 'punctuation' in patterns and intensity > 0.5:
            # Find sentences ending with basic punctuation
            sentences = re.split(r'[.!?]+', text)
            if len(sentences) > 1:  # Has ending punctuation
                new_punct = random.choice(patterns['punctuation'])
                # Replace last punctuation
                text = re.sub(r'[.!?]+$', new_punct, text.strip())
                modifications.append(f"Changed ending punctuation to: {new_punct}")

        # Add trailing offs for sad emotions
        if 'trailing_offs' in patterns and intensity > 0.6:
            if random.random() < 0.3:
                trailing = random.choice(patterns['trailing_offs'])
                if not text.endswith(('...', '..', '!')):
                    text = text.rstrip('.!?') + trailing
                    modifications.append(f"Added trailing: {trailing}")

        # Increase exclamation frequency
        if 'exclamation_frequency' in patterns:
            freq = patterns['exclamation_frequency'] * intensity
            if random.random() < freq and '.' in text:
                text = text.replace('.', '!', 1)
                modifications.append("Changed period to exclamation")

        return text, modifications

    def _apply_emphasis_patterns(
        self,
        text: str,
        emotion: str,
        intensity: float
    ) -> Tuple[str, List[str]]:
        """Apply emphasis patterns like caps, repetition, etc."""
        modifications = []

        if emotion not in self.expression_patterns:
            return text, modifications

        patterns = self.expression_patterns[emotion]

        # Apply capitalization
        if 'capitalization' in patterns and intensity > 0.6:
            cap_prob = patterns['capitalization'] * intensity
            words = text.split()
            for i, word in enumerate(words):
                if random.random() < cap_prob and len(word) > 3:
                    words[i] = word.upper()
                    modifications.append(f"Capitalized: {word}")
            text = ' '.join(words)

        # Apply repetition
        if 'repetition' in patterns and intensity > 0.7:
            if random.random() < 0.2:
                repetitions = patterns['repetition']
                for rep in repetitions:
                    if rep in text:
                        text = text.replace(rep, rep + ' ' + rep.split()[-1], 1)
                        modifications.append(f"Added repetition: {rep}")
                        break

        # Add emphasis markers
        if 'emphasis' in patterns and intensity > 0.6:
            if random.random() < 0.3:
                emphasis = random.choice(patterns['emphasis'])
                text = text + ' ' + emphasis
                modifications.append(f"Added emphasis: {emphasis}")

        return text, modifications

    def _apply_human_inconsistencies(
        self,
        text: str,
        emotional_state: EmotionalState,
        intensity: float
    ) -> Tuple[str, List[str]]:
        """Apply human-like inconsistencies like typos, hesitations."""
        modifications = []

        # Emotional typos (more when excited or upset)
        arousal = emotional_state.get_emotional_arousal()
        base_typo_prob = float(self.typo_probability)

        # Increase typo probability with high arousal
        if arousal > float(self.excitement_threshold):
            typo_prob = base_typo_prob * (1 + arousal)
        else:
            typo_prob = base_typo_prob * 0.5

        # Apply typos
        if random.random() < typo_prob * intensity:
            text, typo_mod = self._add_realistic_typo(text)
            if typo_mod:
                modifications.append(typo_mod)

        # Add hesitations for uncertain emotions
        if emotional_state.fear > 0.6 or emotional_state.emotional_clarity < 0.5:
            if random.random() < 0.3:
                hesitations = ['um...', 'well...', 'I mean...', 'like...']
                hesitation = random.choice(hesitations)
                text = hesitation + ' ' + text
                modifications.append(f"Added hesitation: {hesitation}")

        # Add thinking markers for curiosity
        if emotional_state.curiosity > 0.7:
            if random.random() < 0.4:
                thinking_markers = ['Hmm...', 'Let me think...', 'Interesting...']
                marker = random.choice(thinking_markers)
                text = marker + ' ' + text
                modifications.append(f"Added thinking marker: {marker}")

        return text, modifications

    def _add_realistic_typo(self, text: str) -> Tuple[str, Optional[str]]:
        """Add realistic typos that humans might make when emotional."""
        words = text.split()
        if len(words) < 2:
            return text, None

        # Common typo patterns
        typo_patterns = [
            ('the', 'teh'),
            ('and', 'adn'),
            ('you', 'yuo'),
            ('that', 'taht'),
            ('this', 'tihs'),
            ('really', 'realy'),
            ('because', 'becuase'),
            ('definitely', 'definately'),
            ('probably', 'probaly')
        ]

        # Try to apply a typo
        for original, typo in typo_patterns:
            if original in words:
                idx = words.index(original)
                words[idx] = typo
                return ' '.join(words), f"Added typo: {original} -> {typo}"

        # Letter swapping typo
        target_word_idx = random.randint(0, len(words) - 1)
        word = words[target_word_idx]

        if len(word) > 3:
            # Swap two adjacent letters
            pos = random.randint(0, len(word) - 2)
            word_list = list(word)
            word_list[pos], word_list[pos + 1] = word_list[pos + 1], word_list[pos]
            words[target_word_idx] = ''.join(word_list)
            return ' '.join(words), f"Letter swap typo in: {word}"

        return text, None

    def _apply_contextual_adjustments(
        self,
        text: str,
        context: str,
        emotional_state: EmotionalState
    ) -> Tuple[str, List[str]]:
        """Apply contextual adjustments based on conversation context."""
        modifications = []

        # Formal context - reduce emotional expression
        if context in ['formal', 'professional', 'academic']:
            # Remove excessive punctuation
            text = re.sub(r'!{2,}', '!', text)
            text = re.sub(r'\?{2,}', '?', text)
            modifications.append("Reduced punctuation for formal context")

            # Remove casual expressions
            casual_expressions = ['wow', 'yay', 'awesome', 'cool', 'omg']
            for expr in casual_expressions:
                if expr.lower() in text.lower():
                    text = re.sub(r'\b' + expr + r'\b', '', text, flags=re.IGNORECASE)
                    modifications.append(f"Removed casual expression: {expr}")

        # Casual context - enhance emotional expression
        elif context in ['casual', 'friendly', 'personal']:
            # Allow more emotional expression
            if emotional_state.joy > 0.7 and random.random() < 0.3:
                text = text + ' 😊'
                modifications.append("Added emoji for casual context")

        # Crisis/support context - adjust for empathy
        elif context in ['support', 'crisis', 'emotional']:
            # Add empathetic language
            if emotional_state.empathy_level > 0.6:
                empathy_phrases = [
                    "I understand how you feel.",
                    "That sounds really difficult.",
                    "I'm here for you."
                ]
                if random.random() < 0.4:
                    phrase = random.choice(empathy_phrases)
                    text = phrase + ' ' + text
                    modifications.append(f"Added empathy: {phrase}")

        return text, modifications

    def analyze_emotional_expression(self, text: str) -> Dict[str, Any]:
        """Analyze the emotional expression in a given text."""
        analysis = {
            'detected_emotions': [],
            'expression_intensity': 0.0,
            'punctuation_analysis': {},
            'vocabulary_analysis': {},
            'inconsistencies': []
        }

        # Punctuation analysis
        exclamations = len(re.findall(r'!+', text))
        questions = len(re.findall(r'\?+', text))
        ellipses = len(re.findall(r'\.{2,}', text))

        analysis['punctuation_analysis'] = {
            'exclamations': exclamations,
            'questions': questions,
            'ellipses': ellipses,
            'caps_words': len(re.findall(r'\b[A-Z]{2,}\b', text))
        }

        # Detect emotions from vocabulary
        for emotion, vocab in self.emotional_vocabularies.items():
            emotion_score = 0
            for category, words in vocab.items():
                for word in words:
                    if word.lower() in text.lower():
                        emotion_score += 1

            if emotion_score > 0:
                analysis['detected_emotions'].append({
                    'emotion': emotion,
                    'score': emotion_score,
                    'confidence': min(1.0, emotion_score / 3.0)
                })

        # Overall expression intensity
        intensity_indicators = exclamations + questions + ellipses
        analysis['expression_intensity'] = min(1.0, intensity_indicators / 5.0)

        # Detect potential typos
        common_typos = ['teh', 'adn', 'yuo', 'taht', 'tihs', 'realy', 'becuase', 'definately']
        for typo in common_typos:
            if typo in text.lower():
                analysis['inconsistencies'].append(f"Possible typo: {typo}")

        return analysis

    def get_expression_statistics(self) -> Dict[str, Any]:
        """Get statistics about emotional expression patterns."""
        return {
            'current_typo_probability': float(self.typo_probability),
            'excitement_threshold': float(self.excitement_threshold),
            'available_emotions': list(self.emotional_vocabularies.keys()),
            'expression_patterns': list(self.expression_patterns.keys()),
            'total_vocabulary_entries': sum(
                len(vocab) for emotion_vocab in self.emotional_vocabularies.values()
                for vocab in emotion_vocab.values()
            )
        }

    def calibrate_expression_intensity(self, feedback_history: List[Dict[str, Any]]):
        """Calibrate expression intensity based on user feedback."""
        if not feedback_history:
            return

        # Analyze feedback for different expression intensities
        positive_feedback = [f for f in feedback_history if f.get('rating', 0) > 0.6]
        negative_feedback = [f for f in feedback_history if f.get('rating', 0) < 0.4]

        # Adjust typo probability based on feedback
        if len(positive_feedback) > len(negative_feedback):
            # More positive feedback - can be more expressive
            self.typo_probability.data *= 1.02
        else:
            # More negative feedback - tone down expression
            self.typo_probability.data *= 0.98

        # Clamp typo probability
        self.typo_probability.data = torch.clamp(self.typo_probability.data, 0.001, 0.1)

        logger.info(f"Calibrated expression intensity - typo probability: {float(self.typo_probability):.4f}")

    def create_emotional_response_template(self, emotion: str, intensity: float) -> str:
        """Create a template response showing how this emotion would be expressed."""
        if emotion not in self.emotional_vocabularies:
            return "I'm not sure how to express that emotion."

        vocab = self.emotional_vocabularies[emotion]
        template_parts = []

        # Add emotional expression
        if 'expressions' in vocab:
            expression = random.choice(vocab['expressions'])
            template_parts.append(expression)

        # Add appropriate punctuation
        if emotion in self.expression_patterns:
            patterns = self.expression_patterns[emotion]
            if 'punctuation' in patterns:
                punct = random.choice(patterns['punctuation'])
                if template_parts:
                    template_parts[-1] = template_parts[-1].rstrip('.!?') + punct

        return ' '.join(template_parts) if template_parts else f"*feeling {emotion}*"