Lyra/lyra/personality/adaptation.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
from datetime import datetime, timedelta

from .matrix import PersonalityMatrix, PersonalityTrait
from .traits import OCEANTraits

logger = logging.getLogger(__name__)

class PersonalityAdapter(nn.Module):
    """
    Advanced personality adaptation system that helps Lyra adapt her personality
    in real-time based on conversation context, user preferences, and social dynamics.
    """

    def __init__(
        self,
        personality_matrix: PersonalityMatrix,
        adaptation_strength: float = 0.3,
        memory_length: int = 50
    ):
        super().__init__()

        self.personality_matrix = personality_matrix
        self.adaptation_strength = adaptation_strength
        self.memory_length = memory_length

        # Adaptation networks
        self.context_analyzer = nn.Sequential(
            nn.Linear(512, 256),  # Context embedding input
            nn.LayerNorm(256),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 64)
        )

        self.user_preference_network = nn.Sequential(
            nn.Linear(64 + 15, 128),  # Context + personality features
            nn.LayerNorm(128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 15),  # Output personality adjustments
            nn.Tanh()
        )

        # Social dynamics understanding
        self.social_dynamics_analyzer = nn.Sequential(
            nn.Linear(32, 64),  # Social context features
            nn.ReLU(),
            nn.Linear(64, 32),
            nn.ReLU(),
            nn.Linear(32, 10),  # Social adjustment factors
            nn.Sigmoid()
        )

        # Conversation memory for learning user preferences
        self.conversation_memory = []
        self.user_preference_cache = {}

        # Adaptation history for analysis
        self.adaptation_history = []

    def forward(
        self,
        context_embedding: torch.Tensor,
        user_id: Optional[str] = None,
        social_context: Optional[Dict[str, Any]] = None,
        conversation_history: Optional[List[str]] = None
    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
        """
        Adapt personality for current context and user.

        Args:
            context_embedding: Current conversation context
            user_id: ID of current user
            social_context: Social context information
            conversation_history: Recent conversation for preference learning

        Returns:
            adapted_personality_weights: Personality adjustments for response
            adaptation_info: Information about adaptations made
        """
        batch_size = context_embedding.shape[0]
        device = context_embedding.device

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

        # Get base personality
        base_personality = self._get_base_personality_features().to(device)
        base_personality = base_personality.unsqueeze(0).repeat(batch_size, 1)

        # User-specific adaptations
        user_adaptations = torch.zeros_like(base_personality)
        if user_id:
            user_adaptations = self._get_user_adaptations(
                user_id, context_features, conversation_history
            )

        # Social context adaptations
        social_adaptations = torch.zeros(batch_size, 10, device=device)
        if social_context:
            social_features = self._extract_social_features(social_context, device)
            social_adaptations = self.social_dynamics_analyzer(social_features)

        # Combine base personality with context and user preferences
        combined_input = torch.cat([context_features, base_personality], dim=1)
        personality_adjustments = self.user_preference_network(combined_input)

        # Apply social adaptations
        social_influence = social_adaptations.mean(dim=1, keepdim=True)
        personality_adjustments = personality_adjustments * (0.7 + 0.6 * social_influence)

        # Apply user-specific adaptations
        final_adjustments = (
            self.adaptation_strength * personality_adjustments +
            0.3 * user_adaptations
        )

        # Ensure reasonable adaptation bounds
        final_adjustments = torch.clamp(final_adjustments, -0.5, 0.5)

        # Store adaptation for learning
        adaptation_info = self._record_adaptation(
            user_id, final_adjustments, context_features, social_context
        )

        return final_adjustments, adaptation_info

    def _get_base_personality_features(self) -> torch.Tensor:
        """Get current base personality as feature vector."""
        ocean = self.personality_matrix.ocean_traits.to_tensor()

        custom_traits = torch.tensor([
            trait.value for trait in self.personality_matrix.custom_traits.values()
        ], dtype=torch.float32)

        return torch.cat([ocean, custom_traits])

    def _get_user_adaptations(
        self,
        user_id: str,
        context_features: torch.Tensor,
        conversation_history: Optional[List[str]]
    ) -> torch.Tensor:
        """Get personality adaptations specific to this user."""
        device = context_features.device
        batch_size = context_features.shape[0]

        # Initialize with zero adaptations
        adaptations = torch.zeros(batch_size, 15, device=device)

        # Check if we have learned preferences for this user
        if user_id in self.user_preference_cache:
            user_prefs = self.user_preference_cache[user_id]

            # Apply learned preferences
            for trait_idx, adjustment in enumerate(user_prefs.get('trait_preferences', [])):
                if trait_idx < 15:
                    adaptations[:, trait_idx] = adjustment

        # Learn from conversation history if available
        if conversation_history:
            learned_adaptations = self._learn_from_conversation(
                user_id, conversation_history, context_features
            )
            adaptations = 0.7 * adaptations + 0.3 * learned_adaptations

        return adaptations

    def _extract_social_features(
        self,
        social_context: Dict[str, Any],
        device: torch.device
    ) -> torch.Tensor:
        """Extract features from social context."""
        features = torch.zeros(1, 32, device=device)

        # Group conversation indicators
        if social_context.get('is_group_conversation', False):
            features[0, 0] = 1.0
            features[0, 1] = social_context.get('group_size', 0) / 20.0  # Normalize

        # Formality level
        formality = social_context.get('formality_level', 0.5)
        features[0, 2] = formality

        # Emotional tone of conversation
        emotional_tone = social_context.get('emotional_tone', {})
        for i, emotion in enumerate(['positive', 'negative', 'neutral', 'excited', 'calm']):
            if i < 5:
                features[0, 3 + i] = emotional_tone.get(emotion, 0.0)

        # Topic category
        topic = social_context.get('topic_category', 'general')
        topic_mapping = {
            'technical': 8, 'casual': 9, 'emotional': 10, 'creative': 11,
            'professional': 12, 'academic': 13, 'social': 14, 'personal': 15
        }
        if topic in topic_mapping:
            features[0, topic_mapping[topic]] = 1.0

        # Conflict or disagreement present
        if social_context.get('has_conflict', False):
            features[0, 16] = 1.0

        # User's apparent expertise level
        expertise = social_context.get('user_expertise_level', 0.5)
        features[0, 17] = expertise

        # Time pressure
        time_pressure = social_context.get('time_pressure', 0.0)
        features[0, 18] = time_pressure

        # Cultural context features
        cultural_context = social_context.get('cultural_context', {})
        features[0, 19] = cultural_context.get('directness_preference', 0.5)
        features[0, 20] = cultural_context.get('hierarchy_awareness', 0.5)

        return features

    def _learn_from_conversation(
        self,
        user_id: str,
        conversation_history: List[str],
        context_features: torch.Tensor
    ) -> torch.Tensor:
        """Learn user preferences from conversation patterns."""
        device = context_features.device
        batch_size = context_features.shape[0]

        adaptations = torch.zeros(batch_size, 15, device=device)

        if len(conversation_history) < 3:
            return adaptations

        # Analyze conversation patterns
        conversation_analysis = self._analyze_conversation_patterns(conversation_history)

        # Update user preference cache
        if user_id not in self.user_preference_cache:
            self.user_preference_cache[user_id] = {
                'trait_preferences': [0.0] * 15,
                'conversation_count': 0,
                'satisfaction_history': [],
                'adaptation_success': {}
            }

        user_cache = self.user_preference_cache[user_id]

        # Learn trait preferences based on conversation success
        if conversation_analysis['engagement_level'] > 0.7:
            # High engagement - strengthen current personality settings
            current_traits = self._get_base_personality_features()
            for i in range(min(15, len(current_traits))):
                learning_rate = 0.05
                user_cache['trait_preferences'][i] = (
                    0.95 * user_cache['trait_preferences'][i] +
                    learning_rate * (current_traits[i].item() - 0.5)
                )

        elif conversation_analysis['engagement_level'] < 0.3:
            # Low engagement - try different personality approach
            for i in range(15):
                # Slightly push toward opposite direction
                current_adjustment = user_cache['trait_preferences'][i]
                user_cache['trait_preferences'][i] = current_adjustment * 0.9

        # Apply learned preferences to adaptations
        for i, pref in enumerate(user_cache['trait_preferences']):
            adaptations[:, i] = pref

        user_cache['conversation_count'] += 1

        return adaptations

    def _analyze_conversation_patterns(self, conversation_history: List[str]) -> Dict[str, float]:
        """Analyze conversation patterns to infer user preferences and engagement."""
        if not conversation_history:
            return {'engagement_level': 0.5}

        # Simple heuristic analysis (in a real system, this would be more sophisticated)
        total_length = sum(len(msg.split()) for msg in conversation_history)
        avg_length = total_length / len(conversation_history)

        # Question frequency (indicates engagement)
        question_count = sum(1 for msg in conversation_history if '?' in msg)
        question_ratio = question_count / len(conversation_history)

        # Emotional indicators
        positive_words = ['good', 'great', 'awesome', 'love', 'excellent', 'amazing', 'perfect']
        negative_words = ['bad', 'terrible', 'hate', 'awful', 'worst', 'horrible']

        positive_count = sum(
            sum(1 for word in positive_words if word in msg.lower())
            for msg in conversation_history
        )
        negative_count = sum(
            sum(1 for word in negative_words if word in msg.lower())
            for msg in conversation_history
        )

        # Calculate engagement level
        engagement_score = 0.5  # Base engagement

        # Longer messages indicate engagement
        if avg_length > 10:
            engagement_score += 0.2
        elif avg_length < 3:
            engagement_score -= 0.2

        # Questions indicate engagement
        engagement_score += question_ratio * 0.3

        # Emotional valence
        if positive_count > negative_count:
            engagement_score += 0.2
        elif negative_count > positive_count:
            engagement_score -= 0.2

        engagement_score = np.clip(engagement_score, 0.0, 1.0)

        return {
            'engagement_level': engagement_score,
            'avg_message_length': avg_length,
            'question_ratio': question_ratio,
            'emotional_valence': (positive_count - negative_count) / max(1, len(conversation_history))
        }

    def _record_adaptation(
        self,
        user_id: Optional[str],
        adaptations: torch.Tensor,
        context_features: torch.Tensor,
        social_context: Optional[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """Record adaptation for analysis and learning."""
        adaptation_record = {
            'timestamp': datetime.now().isoformat(),
            'user_id': user_id,
            'adaptations': adaptations[0].detach().cpu().numpy().tolist(),
            'context_strength': torch.norm(context_features).item(),
            'social_context_type': social_context.get('topic_category', 'general') if social_context else 'general'
        }

        self.adaptation_history.append(adaptation_record)

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

        # Prepare return info
        adaptation_info = {
            'adaptation_magnitude': torch.norm(adaptations).item(),
            'primary_adaptations': self._identify_primary_adaptations(adaptations[0]),
            'user_specific': user_id is not None,
            'social_context_present': social_context is not None
        }

        return adaptation_info

    def _identify_primary_adaptations(self, adaptations: torch.Tensor) -> Dict[str, float]:
        """Identify the main personality adaptations being made."""
        trait_names = [
            'openness', 'conscientiousness', 'extraversion', 'agreeableness', 'neuroticism',
            'humor_level', 'sarcasm_tendency', 'empathy_level', 'curiosity', 'playfulness',
            'intellectualism', 'spontaneity', 'supportiveness', 'assertiveness', 'creativity'
        ]

        # Find adaptations with magnitude > 0.1
        significant_adaptations = {}
        for i, adaptation in enumerate(adaptations):
            if abs(adaptation.item()) > 0.1 and i < len(trait_names):
                significant_adaptations[trait_names[i]] = adaptation.item()

        return significant_adaptations

    def learn_from_feedback(
        self,
        user_id: str,
        feedback_score: float,
        conversation_context: str,
        adaptations_made: torch.Tensor
    ):
        """
        Learn from user feedback about personality adaptations.

        This helps Lyra understand which personality adaptations work well
        with different users and contexts.
        """
        if user_id not in self.user_preference_cache:
            return

        user_cache = self.user_preference_cache[user_id]

        # Record satisfaction
        user_cache['satisfaction_history'].append({
            'feedback_score': feedback_score,
            'adaptations': adaptations_made.detach().cpu().numpy().tolist(),
            'context': conversation_context,
            'timestamp': datetime.now().isoformat()
        })

        # Keep only recent history
        if len(user_cache['satisfaction_history']) > 50:
            user_cache['satisfaction_history'] = user_cache['satisfaction_history'][-25:]

        # Update adaptation success tracking
        adaptation_key = self._hash_adaptations(adaptations_made)
        if adaptation_key not in user_cache['adaptation_success']:
            user_cache['adaptation_success'][adaptation_key] = {
                'success_count': 0,
                'total_count': 0,
                'avg_feedback': 0.0
            }

        success_data = user_cache['adaptation_success'][adaptation_key]
        success_data['total_count'] += 1
        success_data['avg_feedback'] = (
            (success_data['avg_feedback'] * (success_data['total_count'] - 1) + feedback_score) /
            success_data['total_count']
        )

        if feedback_score > 0.6:
            success_data['success_count'] += 1

        logger.info(f"Updated adaptation learning for user {user_id}: "
                   f"feedback={feedback_score:.2f}, adaptations={adaptation_key}")

    def _hash_adaptations(self, adaptations: torch.Tensor) -> str:
        """Create a hash key for adaptation patterns."""
        # Quantize adaptations to reduce sensitivity
        quantized = torch.round(adaptations * 10) / 10
        return str(quantized.detach().cpu().numpy().tolist())

    def get_adaptation_analytics(self) -> Dict[str, Any]:
        """Get analytics about personality adaptations."""
        if not self.adaptation_history:
            return {'status': 'no_data'}

        recent_adaptations = [
            a for a in self.adaptation_history
            if datetime.fromisoformat(a['timestamp']) > datetime.now() - timedelta(hours=24)
        ]

        analytics = {
            'total_adaptations': len(self.adaptation_history),
            'recent_adaptations': len(recent_adaptations),
            'unique_users': len(set(
                a['user_id'] for a in self.adaptation_history
                if a['user_id'] is not None
            )),
            'avg_adaptation_magnitude': np.mean([
                np.linalg.norm(a['adaptations']) for a in recent_adaptations
            ]) if recent_adaptations else 0.0,
            'most_adapted_traits': self._get_most_adapted_traits(),
            'user_preference_learning': {
                user_id: {
                    'conversation_count': data['conversation_count'],
                    'adaptation_success_rate': len([
                        s for s in data['adaptation_success'].values()
                        if s['success_count'] / max(1, s['total_count']) > 0.6
                    ]) / max(1, len(data['adaptation_success']))
                }
                for user_id, data in self.user_preference_cache.items()
            }
        }

        return analytics

    def _get_most_adapted_traits(self) -> Dict[str, float]:
        """Get traits that are adapted most frequently."""
        trait_names = [
            'openness', 'conscientiousness', 'extraversion', 'agreeableness', 'neuroticism',
            'humor_level', 'sarcasm_tendency', 'empathy_level', 'curiosity', 'playfulness',
            'intellectualism', 'spontaneity', 'supportiveness', 'assertiveness', 'creativity'
        ]

        trait_adaptations = {name: [] for name in trait_names}

        for adaptation_record in self.adaptation_history:
            for i, adaptation in enumerate(adaptation_record['adaptations']):
                if i < len(trait_names):
                    trait_adaptations[trait_names[i]].append(abs(adaptation))

        return {
            name: np.mean(adaptations) if adaptations else 0.0
            for name, adaptations in trait_adaptations.items()
        }

    def reset_user_adaptations(self, user_id: str):
        """Reset learned adaptations for a specific user."""
        if user_id in self.user_preference_cache:
            del self.user_preference_cache[user_id]
            logger.info(f"Reset personality adaptations for user {user_id}")

    def export_personality_insights(self) -> Dict[str, Any]:
        """Export insights about personality adaptation patterns."""
        return {
            'adaptation_history': self.adaptation_history[-100:],  # Recent history
            'user_preferences': {
                user_id: {
                    'trait_preferences': data['trait_preferences'],
                    'conversation_count': data['conversation_count'],
                    'avg_satisfaction': np.mean([
                        s['feedback_score'] for s in data['satisfaction_history']
                    ]) if data['satisfaction_history'] else 0.0
                }
                for user_id, data in self.user_preference_cache.items()
            },
            'analytics': self.get_adaptation_analytics()
        }