import torch
import torch.nn as nn
import numpy as np
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass, field
from enum import Enum
import json
import logging

logger = logging.getLogger(__name__)

class MyersBriggsType(Enum):
    """Myers-Briggs personality types."""
    INTJ = "INTJ"  # Architect
    INTP = "INTP"  # Thinker
    ENTJ = "ENTJ"  # Commander
    ENTP = "ENTP"  # Debater
    INFJ = "INFJ"  # Advocate
    INFP = "INFP"  # Mediator
    ENFJ = "ENFJ"  # Protagonist
    ENFP = "ENFP"  # Campaigner
    ISTJ = "ISTJ"  # Logistician
    ISFJ = "ISFJ"  # Protector
    ESTJ = "ESTJ"  # Executive
    ESFJ = "ESFJ"  # Consul
    ISTP = "ISTP"  # Virtuoso
    ISFP = "ISFP"  # Adventurer
    ESTP = "ESTP"  # Entrepreneur
    ESFP = "ESFP"  # Entertainer

@dataclass
class OCEANTraits:
    """Big Five (OCEAN) personality traits with dynamic adaptation."""
    openness: float = 0.5          # Openness to experience
    conscientiousness: float = 0.5  # Conscientiousness
    extraversion: float = 0.5       # Extraversion
    agreeableness: float = 0.5      # Agreeableness
    neuroticism: float = 0.5        # Neuroticism

    # Trait variance - how much each trait can fluctuate
    openness_variance: float = 0.1
    conscientiousness_variance: float = 0.1
    extraversion_variance: float = 0.1
    agreeableness_variance: float = 0.1
    neuroticism_variance: float = 0.1

    def to_dict(self) -> Dict[str, float]:
        """Convert to dictionary representation."""
        return {
            'openness': self.openness,
            'conscientiousness': self.conscientiousness,
            'extraversion': self.extraversion,
            'agreeableness': self.agreeableness,
            'neuroticism': self.neuroticism,
            'openness_variance': self.openness_variance,
            'conscientiousness_variance': self.conscientiousness_variance,
            'extraversion_variance': self.extraversion_variance,
            'agreeableness_variance': self.agreeableness_variance,
            'neuroticism_variance': self.neuroticism_variance
        }

    @classmethod
    def from_dict(cls, data: Dict[str, float]) -> 'OCEANTraits':
        """Create from dictionary representation."""
        return cls(**data)

    def to_tensor(self, device: Optional[torch.device] = None) -> torch.Tensor:
        """Convert to tensor for neural network processing."""
        values = [
            self.openness, self.conscientiousness, self.extraversion,
            self.agreeableness, self.neuroticism
        ]
        return torch.tensor(values, dtype=torch.float32, device=device)

    def apply_situational_modification(
        self,
        situation_type: str,
        intensity: float = 1.0
    ) -> 'OCEANTraits':
        """
        Apply situational modifications to personality traits.

        Different situations can bring out different aspects of personality.
        """
        modified = OCEANTraits(
            openness=self.openness,
            conscientiousness=self.conscientiousness,
            extraversion=self.extraversion,
            agreeableness=self.agreeableness,
            neuroticism=self.neuroticism,
            openness_variance=self.openness_variance,
            conscientiousness_variance=self.conscientiousness_variance,
            extraversion_variance=self.extraversion_variance,
            agreeableness_variance=self.agreeableness_variance,
            neuroticism_variance=self.neuroticism_variance
        )

        # Situational trait modifications
        modifications = {
            'stress': {
                'neuroticism': 0.2 * intensity,
                'conscientiousness': -0.1 * intensity,
                'agreeableness': -0.1 * intensity
            },
            'social': {
                'extraversion': 0.15 * intensity,
                'agreeableness': 0.1 * intensity,
                'openness': 0.05 * intensity
            },
            'creative': {
                'openness': 0.2 * intensity,
                'conscientiousness': -0.05 * intensity,
                'neuroticism': -0.1 * intensity
            },
            'conflict': {
                'agreeableness': -0.2 * intensity,
                'neuroticism': 0.15 * intensity,
                'extraversion': -0.1 * intensity
            },
            'learning': {
                'openness': 0.15 * intensity,
                'conscientiousness': 0.1 * intensity,
                'neuroticism': -0.05 * intensity
            }
        }

        if situation_type in modifications:
            mods = modifications[situation_type]
            for trait, change in mods.items():
                current_value = getattr(modified, trait)
                variance = getattr(modified, f"{trait}_variance")

                # Apply change within variance bounds
                new_value = current_value + change
                new_value = np.clip(new_value,
                                   current_value - variance,
                                   current_value + variance)
                new_value = np.clip(new_value, 0.0, 1.0)

                setattr(modified, trait, new_value)

        return modified

@dataclass
class PersonalityEvolution:
    """Tracks how personality evolves over time."""
    adaptation_rate: float = 0.01
    stability_factor: float = 0.9
    max_change_per_step: float = 0.05

    # Evolution history
    evolution_history: List[Dict[str, Any]] = field(default_factory=list)
    total_interactions: int = 0

    def __post_init__(self):
        """Initialize evolution tracking."""
        if not self.evolution_history:
            self.evolution_history = []

class PersonalityDynamics(nn.Module):
    """
    Neural network that models personality dynamics and adaptation.

    This system allows Lyra's personality to evolve naturally based on
    her interactions and experiences.
    """

    def __init__(
        self,
        input_dim: int = 10,  # Contextual features
        personality_dim: int = 5,  # OCEAN traits
        hidden_dim: int = 64,
        adaptation_rate: float = 0.01
    ):
        super().__init__()

        self.personality_dim = personality_dim
        self.adaptation_rate = adaptation_rate

        # Context processing network
        self.context_processor = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.LayerNorm(hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, personality_dim)
        )

        # Personality adaptation network
        self.adaptation_network = nn.Sequential(
            nn.Linear(personality_dim * 2, hidden_dim),  # Current + context influence
            nn.LayerNorm(hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, personality_dim),
            nn.Tanh()  # Bounded output for personality changes
        )

        # Stability network - resists change when personality is stable
        self.stability_network = nn.Sequential(
            nn.Linear(personality_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, 1),
            nn.Sigmoid()
        )

        # Meta-learning for adaptation rate
        self.meta_adaptation = nn.Linear(personality_dim + 1, 1)  # +1 for feedback

    def forward(
        self,
        current_personality: torch.Tensor,
        context_features: torch.Tensor,
        feedback_signal: Optional[torch.Tensor] = None
    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
        """
        Evolve personality based on context and feedback.

        Args:
            current_personality: Current OCEAN traits [batch, 5]
            context_features: Contextual features [batch, input_dim]
            feedback_signal: Feedback from interactions [batch, 1]

        Returns:
            evolved_personality: Updated personality traits
            evolution_info: Information about the evolution step
        """
        batch_size = current_personality.shape[0]

        # Process context to understand what personality aspects to emphasize
        context_influence = self.context_processor(context_features)

        # Combine current personality with context influence
        combined_input = torch.cat([current_personality, context_influence], dim=-1)

        # Generate personality adaptation
        personality_delta = self.adaptation_network(combined_input)

        # Calculate stability (resistance to change)
        stability = self.stability_network(current_personality)

        # Meta-learning: adapt the adaptation rate based on feedback
        if feedback_signal is not None:
            meta_input = torch.cat([current_personality, feedback_signal], dim=-1)
            meta_adaptation_rate = torch.sigmoid(self.meta_adaptation(meta_input))
        else:
            meta_adaptation_rate = torch.ones(batch_size, 1, device=current_personality.device) * 0.5

        # Apply evolution with stability consideration
        effective_rate = self.adaptation_rate * meta_adaptation_rate * (1 - stability)
        evolved_personality = current_personality + effective_rate * personality_delta

        # Ensure personality traits stay in valid range [0, 1]
        evolved_personality = torch.clamp(evolved_personality, 0.0, 1.0)

        # Prepare evolution info
        evolution_info = {
            'personality_change': torch.norm(evolved_personality - current_personality, dim=-1).mean().item(),
            'stability': stability.mean().item(),
            'context_influence_strength': torch.norm(context_influence, dim=-1).mean().item(),
            'adaptation_rate': effective_rate.mean().item()
        }

        return evolved_personality, evolution_info

class MyersBriggsAnalyzer:
    """Analyzes and updates Myers-Briggs type based on OCEAN traits and behavior."""

    def __init__(self):
        # Mapping from OCEAN traits to Myers-Briggs dimensions
        self.mb_mappings = {
            'E_I': lambda traits: traits.extraversion,  # Extraversion vs Introversion
            'S_N': lambda traits: traits.openness,      # Sensing vs iNtuition
            'T_F': lambda traits: 1 - traits.agreeableness,  # Thinking vs Feeling
            'J_P': lambda traits: traits.conscientiousness   # Judging vs Perceiving
        }

    def analyze_type(self, ocean_traits: OCEANTraits) -> MyersBriggsType:
        """Determine Myers-Briggs type from OCEAN traits."""

        # Calculate dimension scores
        e_i = self.mb_mappings['E_I'](ocean_traits)
        s_n = self.mb_mappings['S_N'](ocean_traits)
        t_f = self.mb_mappings['T_F'](ocean_traits)
        j_p = self.mb_mappings['J_P'](ocean_traits)

        # Determine letters
        letter1 = 'E' if e_i > 0.5 else 'I'
        letter2 = 'N' if s_n > 0.5 else 'S'
        letter3 = 'T' if t_f > 0.5 else 'F'
        letter4 = 'J' if j_p > 0.5 else 'P'

        type_string = letter1 + letter2 + letter3 + letter4

        return MyersBriggsType(type_string)

    def get_type_characteristics(self, mb_type: MyersBriggsType) -> Dict[str, Any]:
        """Get characteristics and tendencies for a Myers-Briggs type."""

        characteristics = {
            MyersBriggsType.INTJ: {
                'communication_style': 'direct, analytical, strategic',
                'decision_making': 'logical, long-term focused',
                'social_tendencies': 'selective, deep connections',
                'stress_response': 'withdraw, analyze, plan',
                'learning_preference': 'conceptual, systematic',
                'humor_style': 'dry, witty, intellectual'
            },
            MyersBriggsType.ENFP: {
                'communication_style': 'enthusiastic, expressive, inspirational',
                'decision_making': 'value-based, considers possibilities',
                'social_tendencies': 'outgoing, builds rapport quickly',
                'stress_response': 'seek support, brainstorm solutions',
                'learning_preference': 'interactive, experiential',
                'humor_style': 'playful, storytelling, spontaneous'
            },
            MyersBriggsType.ISFJ: {
                'communication_style': 'supportive, gentle, detailed',
                'decision_making': 'considers impact on others, traditional',
                'social_tendencies': 'helpful, loyal, modest',
                'stress_response': 'internalize, seek harmony',
                'learning_preference': 'structured, practical examples',
                'humor_style': 'gentle, self-deprecating, situational'
            },
            # Add more types as needed...
        }

        return characteristics.get(mb_type, {
            'communication_style': 'balanced approach',
            'decision_making': 'considers multiple factors',
            'social_tendencies': 'adaptive to situation',
            'stress_response': 'varied coping strategies',
            'learning_preference': 'mixed approaches',
            'humor_style': 'situationally appropriate'
        })

class PersonalityProfiler:
    """Creates and maintains detailed personality profiles."""

    def __init__(self):
        self.ocean_analyzer = OCEANTraits()
        self.mb_analyzer = MyersBriggsAnalyzer()

    def create_profile(
        self,
        ocean_traits: OCEANTraits,
        conversation_history: List[str] = None,
        behavioral_data: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """Create a comprehensive personality profile."""

        # Determine Myers-Briggs type
        mb_type = self.mb_analyzer.analyze_type(ocean_traits)
        mb_characteristics = self.mb_analyzer.get_type_characteristics(mb_type)

        # Create base profile
        profile = {
            'ocean_traits': ocean_traits.to_dict(),
            'myers_briggs_type': mb_type.value,
            'characteristics': mb_characteristics,
            'timestamp': torch.tensor(float(torch.rand(1))).item(),
            'profile_version': 1.0
        }

        # Add behavioral insights if available
        if behavioral_data:
            profile['behavioral_patterns'] = self._analyze_behavioral_patterns(
                ocean_traits, behavioral_data
            )

        # Add conversation style analysis if history is available
        if conversation_history:
            profile['conversation_style'] = self._analyze_conversation_style(
                ocean_traits, conversation_history
            )

        return profile

    def _analyze_behavioral_patterns(
        self,
        ocean_traits: OCEANTraits,
        behavioral_data: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Analyze behavioral patterns based on personality and data."""

        patterns = {}

        # Response time patterns
        if 'response_times' in behavioral_data:
            avg_response_time = np.mean(behavioral_data['response_times'])

            # Introverts typically take longer to respond
            expected_time = 2.0 + (1 - ocean_traits.extraversion) * 3.0

            patterns['response_speed'] = {
                'average_seconds': avg_response_time,
                'relative_to_expected': avg_response_time / expected_time,
                'pattern': 'quick' if avg_response_time < expected_time else 'thoughtful'
            }

        # Topic preferences
        if 'topic_engagement' in behavioral_data:
            patterns['topic_preferences'] = self._infer_topic_preferences(
                ocean_traits, behavioral_data['topic_engagement']
            )

        # Emotional expression patterns
        if 'emotional_expressions' in behavioral_data:
            patterns['emotional_style'] = self._analyze_emotional_expression(
                ocean_traits, behavioral_data['emotional_expressions']
            )

        return patterns

    def _analyze_conversation_style(
        self,
        ocean_traits: OCEANTraits,
        conversation_history: List[str]
    ) -> Dict[str, Any]:
        """Analyze conversation style from history."""

        style = {}

        if not conversation_history:
            return style

        # Analyze message characteristics
        message_lengths = [len(msg.split()) for msg in conversation_history]

        style['verbosity'] = {
            'average_words': np.mean(message_lengths),
            'variance': np.var(message_lengths),
            'style': 'concise' if np.mean(message_lengths) < 10 else 'elaborate'
        }

        # Question asking frequency (curiosity indicator)
        question_count = sum(1 for msg in conversation_history if '?' in msg)
        style['curiosity_level'] = question_count / len(conversation_history)

        # Emotional expression analysis
        emotional_words = ['love', 'hate', 'excited', 'sad', 'happy', 'angry', 'worried']
        emotional_frequency = sum(
            sum(1 for word in emotional_words if word in msg.lower())
            for msg in conversation_history
        ) / len(conversation_history)

        style['emotional_expressiveness'] = emotional_frequency

        return style

    def _infer_topic_preferences(
        self,
        ocean_traits: OCEANTraits,
        topic_engagement: Dict[str, float]
    ) -> Dict[str, Any]:
        """Infer topic preferences based on personality and engagement data."""

        preferences = {}

        # High openness correlates with interest in abstract/creative topics
        if ocean_traits.openness > 0.6:
            creative_topics = ['art', 'philosophy', 'science', 'technology', 'literature']
            preferences['preferred_categories'] = creative_topics

        # High conscientiousness correlates with practical topics
        if ocean_traits.conscientiousness > 0.6:
            practical_topics = ['productivity', 'planning', 'organization', 'goals']
            preferences['practical_interests'] = practical_topics

        # Extraversion affects social topic interest
        if ocean_traits.extraversion > 0.6:
            social_topics = ['relationships', 'social events', 'collaboration']
            preferences['social_interests'] = social_topics

        # Add engagement scores
        preferences['engagement_scores'] = topic_engagement

        return preferences

    def _analyze_emotional_expression(
        self,
        ocean_traits: OCEANTraits,
        emotional_expressions: Dict[str, int]
    ) -> Dict[str, Any]:
        """Analyze how emotions are expressed based on personality."""

        style = {}

        total_expressions = sum(emotional_expressions.values())
        if total_expressions == 0:
            return style

        # Calculate emotion proportions
        emotion_proportions = {
            emotion: count / total_expressions
            for emotion, count in emotional_expressions.items()
        }

        style['emotion_distribution'] = emotion_proportions

        # Analyze based on personality traits
        if ocean_traits.neuroticism > 0.6:
            style['emotional_volatility'] = 'high'
        elif ocean_traits.neuroticism < 0.4:
            style['emotional_volatility'] = 'low'
        else:
            style['emotional_volatility'] = 'moderate'

        if ocean_traits.agreeableness > 0.6:
            style['emotional_tone'] = 'positive_focused'
        else:
            style['emotional_tone'] = 'balanced'

        return style