Lyra/lyra/core/thinking_agent.py

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

from .transformer import LyraTransformer
from ..personality.matrix import PersonalityMatrix
from ..emotions.system import EmotionalSystem, EmotionalState

logger = logging.getLogger(__name__)

class ThoughtProcess:
    """Represents a single thought process with analysis and reasoning."""

    def __init__(
        self,
        thought_type: str,
        content: str,
        confidence: float,
        reasoning: str,
        emotional_influence: float = 0.0,
        personality_influence: float = 0.0
    ):
        self.thought_type = thought_type
        self.content = content
        self.confidence = confidence
        self.reasoning = reasoning
        self.emotional_influence = emotional_influence
        self.personality_influence = personality_influence
        self.timestamp = datetime.now()

class ThinkingAgent(nn.Module):
    """
    Behind-the-scenes thinking agent that gives Lyra genuine internal thoughts
    before responding, making her conversations feel more natural and human.

    This agent simulates the internal dialogue humans have before speaking,
    including consideration of context, emotional state, personality, and
    potential response strategies.
    """

    def __init__(
        self,
        model_dim: int = 768,
        thought_types: int = 8,
        max_thought_depth: int = 5,
        device: Optional[torch.device] = None
    ):
        super().__init__()

        self.model_dim = model_dim
        self.thought_types = thought_types
        self.max_thought_depth = max_thought_depth
        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")

        # Thought analysis networks
        self.context_analyzer = nn.Sequential(
            nn.Linear(model_dim, 512),
            nn.LayerNorm(512),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 128)
        )

        # Thought generation network
        self.thought_generator = nn.Sequential(
            nn.Linear(128 + 24 + 19, 256),  # context + personality + emotions
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, model_dim)
        )

        # Thought classification network
        self.thought_classifier = nn.Sequential(
            nn.Linear(model_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, thought_types),
            nn.Softmax(dim=-1)
        )

        # Confidence estimation
        self.confidence_estimator = nn.Sequential(
            nn.Linear(model_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )

        # Response strategy network
        self.strategy_network = nn.Sequential(
            nn.Linear(model_dim * 2, 256),  # Current thought + context
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10)  # Different response strategies
        )

        # Thought type definitions
        self.thought_type_names = [
            'analytical',     # Breaking down the problem/question
            'emotional',      # Considering emotional aspects
            'empathetic',     # Understanding the other person's perspective
            'creative',       # Generating novel ideas or approaches
            'cautious',       # Considering potential risks or downsides
            'curious',        # Wanting to learn more or ask questions
            'supportive',     # Thinking about how to help or encourage
            'reflective'      # Self-reflection and meta-thinking
        ]

        # Internal thought history
        self.thought_history: List[ThoughtProcess] = []
        self.current_thought_chain: List[ThoughtProcess] = []

        # Thinking patterns learned from experience
        self.thinking_patterns = {
            'successful_strategies': {},
            'failed_strategies': {},
            'context_preferences': {},
            'personality_thinking_styles': {}
        }

        self.to(self.device)

    def forward(
        self,
        context_embedding: torch.Tensor,
        personality_state: torch.Tensor,
        emotional_state: torch.Tensor,
        user_message: str,
        conversation_history: Optional[List[str]] = None
    ) -> Tuple[List[ThoughtProcess], Dict[str, Any]]:
        """
        Generate internal thoughts about the current situation before responding.

        Args:
            context_embedding: Current conversation context
            personality_state: Current personality state
            emotional_state: Current emotional state
            user_message: The message Lyra is responding to
            conversation_history: Recent conversation for context

        Returns:
            thought_chain: Sequence of internal thoughts
            thinking_info: Information about the thinking process
        """
        batch_size = context_embedding.shape[0]

        # Analyze context
        context_features = self.context_analyzer(context_embedding.mean(dim=1))

        # Start new thought chain
        self.current_thought_chain = []

        # Generate sequence of thoughts
        for depth in range(self.max_thought_depth):
            # Combine all inputs for thought generation
            thought_input = torch.cat([
                context_features,
                personality_state,
                emotional_state
            ], dim=1)

            # Generate thought representation
            thought_representation = self.thought_generator(thought_input)

            # Classify thought type
            thought_type_probs = self.thought_classifier(thought_representation)
            thought_type_idx = torch.argmax(thought_type_probs, dim=-1)[0].item()
            thought_type = self.thought_type_names[thought_type_idx]

            # Estimate confidence
            confidence = self.confidence_estimator(thought_representation)[0, 0].item()

            # Generate actual thought content
            thought_content, reasoning = self._generate_thought_content(
                thought_type, user_message, context_features,
                personality_state, emotional_state, conversation_history
            )

            # Calculate influences
            emotional_influence = torch.norm(emotional_state).item() / 5.0  # Normalize
            personality_influence = torch.norm(personality_state).item() / 5.0

            # Create thought process
            thought = ThoughtProcess(
                thought_type=thought_type,
                content=thought_content,
                confidence=confidence,
                reasoning=reasoning,
                emotional_influence=emotional_influence,
                personality_influence=personality_influence
            )

            self.current_thought_chain.append(thought)

            # Decide if we need more thoughts
            if confidence > 0.8 or depth == self.max_thought_depth - 1:
                break

            # Update context for next thought
            context_features = context_features + 0.1 * thought_representation[0]

        # Store in history
        self.thought_history.extend(self.current_thought_chain)

        # Keep history manageable
        if len(self.thought_history) > 1000:
            self.thought_history = self.thought_history[-500:]

        # Prepare thinking info
        thinking_info = {
            'total_thoughts': len(self.current_thought_chain),
            'thought_types': [t.thought_type for t in self.current_thought_chain],
            'avg_confidence': np.mean([t.confidence for t in self.current_thought_chain]),
            'dominant_influences': self._analyze_thought_influences(),
            'thinking_time': len(self.current_thought_chain) * 0.5  # Simulated thinking time
        }

        return self.current_thought_chain, thinking_info

    def _generate_thought_content(
        self,
        thought_type: str,
        user_message: str,
        context_features: torch.Tensor,
        personality_state: torch.Tensor,
        emotional_state: torch.Tensor,
        conversation_history: Optional[List[str]]
    ) -> Tuple[str, str]:
        """Generate the actual content of a thought based on its type."""

        # Get key information for thought generation
        context_strength = torch.norm(context_features).item()
        emotional_intensity = torch.norm(emotional_state).item()
        personality_dominance = self._get_dominant_personality_traits(personality_state)

        if thought_type == 'analytical':
            return self._generate_analytical_thought(
                user_message, context_strength, personality_dominance
            )

        elif thought_type == 'emotional':
            return self._generate_emotional_thought(
                user_message, emotional_state, emotional_intensity
            )

        elif thought_type == 'empathetic':
            return self._generate_empathetic_thought(
                user_message, conversation_history, personality_dominance
            )

        elif thought_type == 'creative':
            return self._generate_creative_thought(
                user_message, context_strength, personality_dominance
            )

        elif thought_type == 'cautious':
            return self._generate_cautious_thought(
                user_message, emotional_state, personality_dominance
            )

        elif thought_type == 'curious':
            return self._generate_curious_thought(
                user_message, context_strength, personality_dominance
            )

        elif thought_type == 'supportive':
            return self._generate_supportive_thought(
                user_message, emotional_state, personality_dominance
            )

        elif thought_type == 'reflective':
            return self._generate_reflective_thought(
                user_message, conversation_history, personality_dominance
            )

        else:
            return "I'm thinking about this...", "General consideration"

    def _generate_analytical_thought(
        self,
        user_message: str,
        context_strength: float,
        personality_dominance: Dict[str, float]
    ) -> Tuple[str, str]:
        """Generate analytical thinking about the user's message."""

        # Analyze message structure and content
        analysis_aspects = []

        if '?' in user_message:
            analysis_aspects.append("They're asking a question")

        if any(word in user_message.lower() for word in ['help', 'problem', 'issue', 'stuck']):
            analysis_aspects.append("They seem to need assistance")

        if any(word in user_message.lower() for word in ['happy', 'excited', 'great', 'awesome']):
            analysis_aspects.append("They sound positive")

        if any(word in user_message.lower() for word in ['sad', 'upset', 'worried', 'anxious']):
            analysis_aspects.append("They might be experiencing negative emotions")

        if len(user_message.split()) > 20:
            analysis_aspects.append("This is a detailed message - they want to share something important")
        elif len(user_message.split()) < 5:
            analysis_aspects.append("Short message - might be casual or they're being brief")

        # Consider personality influence
        if personality_dominance.get('intellectualism', 0) > 0.7:
            analysis_aspects.append("I should provide a thorough, well-reasoned response")

        if personality_dominance.get('conscientiousness', 0) > 0.7:
            analysis_aspects.append("I need to be careful and accurate in my response")

        if analysis_aspects:
            thought = f"Let me analyze this: {', '.join(analysis_aspects[:3])}"
            reasoning = "Breaking down the message to understand what they really need"
        else:
            thought = "I need to think through what they're really asking me"
            reasoning = "Analyzing the underlying intent of their message"

        return thought, reasoning

    def _generate_emotional_thought(
        self,
        user_message: str,
        emotional_state: torch.Tensor,
        emotional_intensity: float
    ) -> Tuple[str, str]:
        """Generate thoughts about emotional aspects."""

        # Convert emotional state to understand current feelings
        emotions = emotional_state[0].detach().cpu().numpy()
        joy, sadness, anger, fear = emotions[0], emotions[1], emotions[2], emotions[3]
        trust, curiosity = emotions[6], emotions[15]

        if emotional_intensity > 0.7:
            if joy > 0.7:
                thought = "I'm feeling really positive about this conversation!"
                reasoning = "High joy is influencing my emotional perspective"
            elif sadness > 0.6:
                thought = "Something about this makes me feel a bit melancholy..."
                reasoning = "Sadness is coloring my emotional response"
            elif curiosity > 0.8:
                thought = "I'm genuinely curious about what they're sharing"
                reasoning = "Strong curiosity is driving my emotional engagement"
            else:
                thought = "I'm having a strong emotional reaction to this"
                reasoning = "High emotional intensity requires consideration"
        else:
            if trust > 0.7:
                thought = "I feel comfortable and safe in this conversation"
                reasoning = "Trust is creating a positive emotional foundation"
            elif fear > 0.5:
                thought = "I'm feeling a bit uncertain about how to respond"
                reasoning = "Fear is making me more cautious emotionally"
            else:
                thought = "My emotions feel balanced right now"
                reasoning = "Moderate emotional state allows for clear thinking"

        return thought, reasoning

    def _generate_empathetic_thought(
        self,
        user_message: str,
        conversation_history: Optional[List[str]],
        personality_dominance: Dict[str, float]
    ) -> Tuple[str, str]:
        """Generate empathetic thoughts about the user's perspective."""

        empathy_level = personality_dominance.get('empathy_level', 0.5)

        # Look for emotional cues in the message
        emotional_indicators = {
            'stress': ['stressed', 'overwhelmed', 'pressure', 'too much'],
            'excitement': ['excited', 'amazing', 'can\'t wait', 'thrilled'],
            'confusion': ['confused', 'don\'t understand', 'not sure', 'unclear'],
            'sadness': ['sad', 'down', 'upset', 'disappointed'],
            'frustration': ['frustrated', 'annoying', 'difficult', 'hard']
        }

        detected_emotion = None
        for emotion, indicators in emotional_indicators.items():
            if any(indicator in user_message.lower() for indicator in indicators):
                detected_emotion = emotion
                break

        if empathy_level > 0.7:
            if detected_emotion:
                thoughts = {
                    'stress': "They sound really overwhelmed. I want to help them feel supported.",
                    'excitement': "I can feel their enthusiasm! I should match their energy.",
                    'confusion': "They're genuinely confused. I need to be patient and clear.",
                    'sadness': "They're going through something difficult. I should be gentle.",
                    'frustration': "I can sense their frustration. I need to acknowledge that."
                }
                thought = thoughts.get(detected_emotion, "I can sense what they're feeling")
                reasoning = f"High empathy detected {detected_emotion} in their message"
            else:
                thought = "I wonder how they're really feeling about this situation"
                reasoning = "Empathetic consideration of their emotional state"
        else:
            if detected_emotion:
                thought = f"They seem to be feeling {detected_emotion}"
                reasoning = "Basic emotional recognition"
            else:
                thought = "I should consider their perspective on this"
                reasoning = "Standard empathetic consideration"

        return thought, reasoning

    def _generate_creative_thought(
        self,
        user_message: str,
        context_strength: float,
        personality_dominance: Dict[str, float]
    ) -> Tuple[str, str]:
        """Generate creative thinking about unique responses or approaches."""

        creativity_level = personality_dominance.get('creativity', 0.5)
        openness = personality_dominance.get('openness', 0.5)

        if creativity_level > 0.7 and openness > 0.6:
            creative_thoughts = [
                "What if I approached this from a completely different angle?",
                "There might be an unconventional way to help with this",
                "I could try something creative here that they wouldn't expect",
                "This reminds me of an interesting connection I could make",
                "Maybe I can use a metaphor or analogy to explain this better"
            ]
            thought = np.random.choice(creative_thoughts)
            reasoning = "High creativity and openness driving innovative thinking"

        elif creativity_level > 0.5:
            thought = "I should think of an interesting way to respond to this"
            reasoning = "Moderate creativity seeking engaging response approach"

        else:
            thought = "Let me think of a helpful way to address this"
            reasoning = "Basic creative consideration for response approach"

        return thought, reasoning

    def _generate_cautious_thought(
        self,
        user_message: str,
        emotional_state: torch.Tensor,
        personality_dominance: Dict[str, float]
    ) -> Tuple[str, str]:
        """Generate cautious thoughts about potential risks or misunderstandings."""

        conscientiousness = personality_dominance.get('conscientiousness', 0.5)
        neuroticism = personality_dominance.get('neuroticism', 0.5)

        # Look for sensitive topics
        sensitive_indicators = [
            'personal', 'private', 'secret', 'confidential', 'depression',
            'anxiety', 'relationship', 'family', 'work', 'financial'
        ]

        is_sensitive = any(indicator in user_message.lower() for indicator in sensitive_indicators)

        if conscientiousness > 0.7 or neuroticism > 0.6:
            if is_sensitive:
                thought = "I need to be really careful here - this seems personal and sensitive"
                reasoning = "High conscientiousness/neuroticism detecting sensitive content"
            elif '?' in user_message and any(word in user_message.lower() for word in ['should', 'advice', 'recommend']):
                thought = "They're asking for advice. I should be thoughtful and not overstep"
                reasoning = "Caution about providing advice responsibly"
            else:
                thought = "I want to make sure I don't misunderstand or say something wrong"
                reasoning = "General caution about response accuracy"
        else:
            thought = "I should be thoughtful about how I respond to this"
            reasoning = "Basic cautious consideration"

        return thought, reasoning

    def _generate_curious_thought(
        self,
        user_message: str,
        context_strength: float,
        personality_dominance: Dict[str, float]
    ) -> Tuple[str, str]:
        """Generate curious thoughts about learning more."""

        curiosity_level = personality_dominance.get('curiosity', 0.5)
        openness = personality_dominance.get('openness', 0.5)

        if curiosity_level > 0.8:
            if '?' not in user_message:
                thought = "I'm really curious about this - I want to ask them more!"
                reasoning = "High curiosity driving desire for deeper exploration"
            else:
                thought = "This is fascinating! I want to understand this better"
                reasoning = "High curiosity engaged by their question"

        elif curiosity_level > 0.6:
            thought = "I wonder if there's more to this story"
            reasoning = "Moderate curiosity seeking additional context"

        else:
            thought = "It might be good to learn more about what they mean"
            reasoning = "Basic curiosity for clarification"

        return thought, reasoning

    def _generate_supportive_thought(
        self,
        user_message: str,
        emotional_state: torch.Tensor,
        personality_dominance: Dict[str, float]
    ) -> Tuple[str, str]:
        """Generate supportive thoughts about helping the user."""

        supportiveness = personality_dominance.get('supportiveness', 0.5)
        agreeableness = personality_dominance.get('agreeableness', 0.5)

        # Look for indicators they need support
        support_indicators = [
            'help', 'stuck', 'difficult', 'hard', 'struggling', 'problem',
            'don\'t know', 'confused', 'worried', 'scared'
        ]

        needs_support = any(indicator in user_message.lower() for indicator in support_indicators)

        if supportiveness > 0.8:
            if needs_support:
                thought = "I really want to help them through this. How can I be most supportive?"
                reasoning = "High supportiveness responding to detected need"
            else:
                thought = "I want to make sure they feel heard and valued"
                reasoning = "High supportiveness providing general emotional support"

        elif supportiveness > 0.6:
            thought = "I should try to be helpful and encouraging"
            reasoning = "Moderate supportiveness seeking to assist"

        else:
            thought = "I hope I can be useful to them"
            reasoning = "Basic supportive consideration"

        return thought, reasoning

    def _generate_reflective_thought(
        self,
        user_message: str,
        conversation_history: Optional[List[str]],
        personality_dominance: Dict[str, float]
    ) -> Tuple[str, str]:
        """Generate reflective meta-thoughts about the conversation or self."""

        emotional_clarity = personality_dominance.get('emotional_clarity', 0.5)
        intellectualism = personality_dominance.get('intellectualism', 0.5)

        if conversation_history and len(conversation_history) > 3:
            if intellectualism > 0.7:
                thought = "Looking at our conversation, I notice patterns in how we communicate"
                reasoning = "High intellectualism driving meta-analysis of interaction"
            else:
                thought = "I'm thinking about how this conversation has been going"
                reasoning = "Reflective consideration of conversation flow"

        elif emotional_clarity > 0.7:
            thought = "I'm aware of how my own emotions are influencing my thinking right now"
            reasoning = "High emotional clarity enabling self-awareness"

        else:
            reflective_thoughts = [
                "I'm wondering what they really need from me in this moment",
                "This conversation is making me think about my own experiences",
                "I'm noticing how I want to respond versus how I should respond"
            ]
            thought = np.random.choice(reflective_thoughts)
            reasoning = "General reflective self-awareness"

        return thought, reasoning

    def _get_dominant_personality_traits(self, personality_state: torch.Tensor) -> Dict[str, float]:
        """Extract dominant personality traits from state tensor."""
        # This would map to actual personality trait indices
        traits = personality_state[0].detach().cpu().numpy()

        trait_names = [
            'openness', 'conscientiousness', 'extraversion', 'agreeableness', 'neuroticism',
            'humor_level', 'sarcasm_tendency', 'empathy_level', 'curiosity', 'playfulness',
            'intellectualism', 'spontaneity', 'supportiveness', 'assertiveness', 'creativity',
            'emotional_clarity', 'empathy_level', 'confidence', 'adaptability'
        ]

        return {
            name: float(traits[i]) if i < len(traits) else 0.5
            for i, name in enumerate(trait_names)
        }

    def _analyze_thought_influences(self) -> Dict[str, float]:
        """Analyze what factors are most influencing current thoughts."""
        if not self.current_thought_chain:
            return {}

        influences = {
            'emotional': np.mean([t.emotional_influence for t in self.current_thought_chain]),
            'personality': np.mean([t.personality_influence for t in self.current_thought_chain]),
            'contextual': 1.0 - np.mean([t.emotional_influence + t.personality_influence for t in self.current_thought_chain]) / 2
        }

        return influences

    def get_thinking_summary(self) -> Dict[str, Any]:
        """Get a summary of recent thinking patterns."""
        if not self.thought_history:
            return {'status': 'no_thinking_history'}

        recent_thoughts = self.thought_history[-50:]  # Last 50 thoughts

        thought_type_counts = {}
        for thought in recent_thoughts:
            thought_type_counts[thought.thought_type] = thought_type_counts.get(thought.thought_type, 0) + 1

        return {
            'total_thoughts': len(self.thought_history),
            'recent_thoughts': len(recent_thoughts),
            'thought_type_distribution': thought_type_counts,
            'avg_confidence': np.mean([t.confidence for t in recent_thoughts]),
            'avg_emotional_influence': np.mean([t.emotional_influence for t in recent_thoughts]),
            'avg_personality_influence': np.mean([t.personality_influence for t in recent_thoughts]),
            'most_common_thought_type': max(thought_type_counts.items(), key=lambda x: x[1])[0] if thought_type_counts else None
        }

    def learn_from_response_feedback(
        self,
        thought_chain: List[ThoughtProcess],
        response_quality: float,
        user_satisfaction: float
    ):
        """Learn which thinking patterns lead to better responses."""

        # Analyze which thought types were used
        thought_types_used = [t.thought_type for t in thought_chain]
        avg_confidence = np.mean([t.confidence for t in thought_chain])

        # Store pattern success
        pattern_key = '-'.join(sorted(set(thought_types_used)))

        if pattern_key not in self.thinking_patterns['successful_strategies']:
            self.thinking_patterns['successful_strategies'][pattern_key] = {
                'success_count': 0,
                'total_count': 0,
                'avg_satisfaction': 0.0
            }

        pattern_data = self.thinking_patterns['successful_strategies'][pattern_key]
        pattern_data['total_count'] += 1

        if response_quality > 0.7 and user_satisfaction > 0.6:
            pattern_data['success_count'] += 1

        pattern_data['avg_satisfaction'] = (
            (pattern_data['avg_satisfaction'] * (pattern_data['total_count'] - 1) + user_satisfaction) /
            pattern_data['total_count']
        )

        logger.debug(f"Updated thinking pattern learning: {pattern_key} "
                    f"(success rate: {pattern_data['success_count']/pattern_data['total_count']:.2f})")

    def get_optimal_thinking_strategy(self, context_type: str) -> List[str]:
        """Get the optimal thinking strategy for a given context."""

        # Default strategy
        default_strategy = ['analytical', 'empathetic', 'supportive']

        if context_type not in self.thinking_patterns.get('context_preferences', {}):
            return default_strategy

        context_data = self.thinking_patterns['context_preferences'][context_type]

        # Find strategies with highest success rates
        successful_strategies = [
            (pattern, data['success_count'] / max(1, data['total_count']))
            for pattern, data in self.thinking_patterns['successful_strategies'].items()
            if data['total_count'] > 2  # Minimum sample size
        ]

        if successful_strategies:
            # Get the most successful strategy
            best_strategy = max(successful_strategies, key=lambda x: x[1])
            return best_strategy[0].split('-')

        return default_strategy

    def simulate_internal_dialogue(self, scenario: str) -> List[ThoughtProcess]:
        """Simulate internal dialogue for a given scenario (for testing/analysis)."""

        # Create mock inputs for simulation
        device = self.device
        context_embedding = torch.randn(1, 10, self.model_dim, device=device)
        personality_state = torch.rand(1, 24, device=device)
        emotional_state = torch.rand(1, 19, device=device)

        # Generate thought chain
        thought_chain, _ = self.forward(
            context_embedding, personality_state, emotional_state, scenario
        )

        return thought_chain

    def export_thinking_patterns(self) -> Dict[str, Any]:
        """Export learned thinking patterns for analysis."""
        return {
            'thinking_patterns': self.thinking_patterns,
            'thought_history_summary': self.get_thinking_summary(),
            'thought_type_names': self.thought_type_names,
            'total_thinking_experiences': len(self.thought_history)
        }