from typing import Dict, Any
import numpy as np
from dataclasses import dataclass

@dataclass
class SelfState:
    """Represents the current state of the dynamic self model."""
    identity_vector: np.ndarray
    belief_state: Dict[str, float]
    goal_hierarchy: Dict[str, float]
    emotional_state: Dict[str, float]
    metacognitive_state: Dict[str, Any]

class DynamicSelfModel:
    """
    Dynamic Self Model that maintains and updates the agent's self-representation.
    This includes identity, beliefs, goals, and emotional states.
    """
    def __init__(self):
        # Initialize default state dimensions
        self.identity_dim = 256
        self.state = SelfState(
            identity_vector=np.zeros(self.identity_dim),
            belief_state={
                "self_awareness": 0.5,
                "world_model": 0.5,
                "agency": 0.5
            },
            goal_hierarchy={
                "primary": 0.8,
                "secondary": 0.5,
                "tertiary": 0.3
            },
            emotional_state={
                "valence": 0.0,
                "arousal": 0.0,
                "dominance": 0.0
            },
            metacognitive_state={
                "confidence": 0.5,
                "uncertainty": 0.5,
                "reflection_level": 0.5
            }
        )

    async def update(self, integrated_state: Dict[str, Any]) -> Dict[str, Any]:
        """
        Update the self model based on new integrated state information.

        Args:
            integrated_state: Dictionary containing the new integrated state information

        Returns:
            Dictionary containing the updated self model state
        """
        # Update identity vector if provided
        if "identity" in integrated_state:
            self.state.identity_vector = self._update_identity(
                integrated_state["identity"]
            )

        # Update belief state
        if "beliefs" in integrated_state:
            self.state.belief_state = self._update_beliefs(
                integrated_state["beliefs"]
            )

        # Update goal hierarchy
        if "goals" in integrated_state:
            self.state.goal_hierarchy = self._update_goals(
                integrated_state["goals"]
            )

        # Update emotional state
        if "emotions" in integrated_state:
            self.state.emotional_state = self._update_emotions(
                integrated_state["emotions"]
            )

        # Update metacognitive state
        if "metacognition" in integrated_state:
            self.state.metacognitive_state = self._update_metacognition(
                integrated_state["metacognition"]
            )

        return self._get_current_state()

    def _update_identity(self, new_identity: np.ndarray) -> np.ndarray:
        """Update the identity vector with new information."""
        # Simple moving average update
        alpha = 0.3  # Learning rate
        return (1 - alpha) * self.state.identity_vector + alpha * new_identity

    def _update_beliefs(self, new_beliefs: Dict[str, float]) -> Dict[str, float]:
        """Update belief states with new information."""
        updated_beliefs = self.state.belief_state.copy()
        for key, value in new_beliefs.items():
            if key in updated_beliefs:
                updated_beliefs[key] = 0.7 * updated_beliefs[key] + 0.3 * value
            else:
                updated_beliefs[key] = value
        return updated_beliefs

    def _update_goals(self, new_goals: Dict[str, float]) -> Dict[str, float]:
        """Update goal hierarchy with new information."""
        updated_goals = self.state.goal_hierarchy.copy()
        for key, value in new_goals.items():
            if key in updated_goals:
                updated_goals[key] = 0.8 * updated_goals[key] + 0.2 * value
            else:
                updated_goals[key] = value
        return updated_goals

    def _update_emotions(self, new_emotions: Dict[str, float]) -> Dict[str, float]:
        """Update emotional state with new information."""
        updated_emotions = self.state.emotional_state.copy()
        for key, value in new_emotions.items():
            if key in updated_emotions:
                updated_emotions[key] = 0.5 * updated_emotions[key] + 0.5 * value
            else:
                updated_emotions[key] = value
        return updated_emotions

    def _update_metacognition(self, new_meta: Dict[str, Any]) -> Dict[str, Any]:
        """Update metacognitive state with new information."""
        updated_meta = self.state.metacognitive_state.copy()
        for key, value in new_meta.items():
            if key in updated_meta:
                if isinstance(value, float):
                    updated_meta[key] = 0.6 * updated_meta[key] + 0.4 * value
                else:
                    updated_meta[key] = value
            else:
                updated_meta[key] = value
        return updated_meta

    def _get_current_state(self) -> Dict[str, Any]:
        """Return the current state as a dictionary."""
        return {
            "identity": self.state.identity_vector,
            "beliefs": self.state.belief_state,
            "goals": self.state.goal_hierarchy,
            "emotions": self.state.emotional_state,
            "metacognition": self.state.metacognitive_state
        }