import numpy as np
import gymnasium as gym

class EntangledQRLAgent:
    """
    Quantum-inspired RL agent with entangled action representation.
    Example: Key-Door gridworld where collecting key and opening door are entangled.
    """
    def __init__(self, state_dim, action_dim, horizon=2):
        # Entangled Q-function: Q(s₁, a₁, s₂, a₂) for action pairs
        # This captures correlations between actions at different time steps
        self.entangled_q = np.zeros((state_dim, action_dim, state_dim, action_dim))
        
        # Shared encoder for entangled representation
        self.shared_encoder = self._init_encoder(state_dim)
        
        # Policy heads for different time steps (but sharing encoder)
        self.policy_head_t1 = self._init_policy_head()
        self.policy_head_t2 = self._init_policy_head()
        
        self.horizon = horizon
        self.learning_rate = 0.1
        self.gamma = 0.99
    
    def _init_encoder(self, state_dim):
        """Shared encoder that produces entangled latent state"""
        # In practice, this would be a neural network
        # For simplicity, we use a linear transformation
        return np.random.randn(state_dim, state_dim) * 0.1
    
    def _init_policy_head(self):
        """Policy head that reads from shared encoder"""
        return np.random.randn(state_dim, 4) * 0.1  # 4 actions: up, down, left, right
    
    def encode_state(self, state):
        """Encode state using shared encoder (creates entangled representation)"""
        return np.dot(self.shared_encoder, state)
    
    def act_entangled(self, state1, state2=None):
        """
        Select entangled action pair (a₁, a₂).
        The actions are correlated through the shared encoder.
        """
        z1 = self.encode_state(state1)
        
        if state2 is None:
            # For first action, use policy head
            logits = np.dot(z1, self.policy_head_t1)
            action1 = self._sample_action(logits)
            return action1
        else:
            # For second action, consider entanglement with first action
            z2 = self.encode_state(state2)
            
            # Entangled Q-value: depends on both states and their correlation
            # Q(s₁, a₁, s₂, a₂) = f(encoder(s₁), encoder(s₂), correlation)
            correlation = np.dot(z1, z2)  # Quantum correlation
            
            # Select action pair that maximizes entangled Q-value
            best_pair = None
            best_value = -np.inf
            
            for a1 in range(4):
                for a2 in range(4):
                    # Entangled value depends on correlation
                    value = self.entangled_q[state1, a1, state2, a2] + correlation
                    if value > best_value:
                        best_value = value
                        best_pair = (a1, a2)
            
            return best_pair
    
    def _sample_action(self, logits):
        """Sample action from logits (softmax policy)"""
        exp_logits = np.exp(logits - np.max(logits))
        probs = exp_logits / np.sum(exp_logits)
        return np.random.choice(len(probs), p=probs)
    
    def update_entangled(self, s1, a1, s2, a2, reward, next_s1, next_s2, done):
        """
        Update entangled Q-values using quantum-inspired TD learning.
        Updates capture correlations between action pairs.
        """
        if done:
            target = reward
        else:
            # Bellman update for entangled Q-function
            # Q*(s₁, a₁, s₂, a₂) = r + γ max_{a₁', a₂'} Q*(s₁', a₁', s₂', a₂')
            next_max = np.max(self.entangled_q[next_s1, :, next_s2, :])
            target = reward + self.gamma * next_max
        
        # TD error
        current_q = self.entangled_q[s1, a1, s2, a2]
        td_error = target - current_q
        
        # Update with quantum-inspired correlation term
        correlation = np.dot(self.encode_state(s1), self.encode_state(s2))
        self.entangled_q[s1, a1, s2, a2] += self.learning_rate * (td_error + 0.1 * correlation)
        
        # Also update shared encoder to strengthen entanglement
        self._update_encoder(s1, s2, td_error)
    
    def _update_encoder(self, s1, s2, td_error):
        """Update shared encoder to strengthen correlations"""
        # Gradient-based update (simplified)
        z1 = self.encode_state(s1)
        z2 = self.encode_state(s2)
        gradient = td_error * np.outer(z1, z2)
        self.shared_encoder += 0.01 * gradient

# Example: Key-Door Gridworld
class KeyDoorGridworld:
    """Gridworld where agent must collect key before opening door"""
    def __init__(self, size=5):
        self.size = size
        self.agent_pos = (0, 0)
        self.key_pos = (2, 2)
        self.door_pos = (4, 4)
        self.has_key = False
        self.door_open = False
    
    def reset(self):
        self.agent_pos = (0, 0)
        self.has_key = False
        self.door_open = False
        return self._get_state()
    
    def _get_state(self):
        """State includes position, key status, door status"""
        return (self.agent_pos[0] * self.size + self.agent_pos[1], 
                int(self.has_key), int(self.door_open))
    
    def step(self, action):
        """Actions: 0=up, 1=down, 2=left, 3=right"""
        # Move agent
        moves = [(-1,0), (1,0), (0,-1), (0,1)]
        new_pos = (self.agent_pos[0] + moves[action][0],
                  self.agent_pos[1] + moves[action][1])
        
        if 0 <= new_pos[0] < self.size and 0 <= new_pos[1] < self.size:
            self.agent_pos = new_pos
        
        # Check for key
        if self.agent_pos == self.key_pos and not self.has_key:
            self.has_key = True
            reward = 10
        
        # Check for door
        elif self.agent_pos == self.door_pos:
            if self.has_key and not self.door_open:
                self.door_open = True
                reward = 50  # Big reward for opening door with key
            elif not self.has_key:
                reward = -10  # Penalty for trying without key
            else:
                reward = 0
        else:
            reward = -1  # Small penalty for each step
        
        done = self.door_open
        return self._get_state(), reward, done