import torch
import torch.nn as nn
import torch.nn.functional as F

class RSSMState:
    def __init__(self, stochastic: torch.Tensor, deterministic: torch.Tensor):
        self.stoch = stochastic
        self.deter = deterministic
        
    def get_features(self):
        """Returns the concatenated stochastic and deterministic state."""
        return torch.cat([self.stoch, self.deter], dim=-1)

class RSSM(nn.Module):
    def __init__(self, action_dim: int, embed_dim: int = 256, stoch_dim: int = 32, deter_dim: int = 256, hidden_dim: int = 256):
        super().__init__()
        self.stoch_dim = stoch_dim
        self.deter_dim = deter_dim
        self.action_dim = action_dim
        
        # RNN now takes action_dim + stoch_dim instead of embed_dim
        self.rnn = nn.GRUCell(action_dim + stoch_dim, deter_dim)
        
        # Prior predicts next stochastic state from deterministic state
        self.prior_net = nn.Sequential(
            nn.Linear(deter_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, 2 * stoch_dim)
        )
        
        # Posterior predicts stochastic state from deterministic state + current observation embedding
        self.post_net = nn.Sequential(
            nn.Linear(deter_dim + embed_dim, hidden_dim),
            nn.ELU(),
            nn.Linear(hidden_dim, 2 * stoch_dim)
        )
        
    def initial_state(self, batch_size: int, device: torch.device) -> RSSMState:
        return RSSMState(
            stochastic=torch.zeros(batch_size, self.stoch_dim, device=device),
            deterministic=torch.zeros(batch_size, self.deter_dim, device=device)
        )
        
    def observe_step(self, prev_state: RSSMState, action: torch.Tensor, embed: torch.Tensor) -> tuple:
        """
        Single step of posterior given a real observation embedding.
        Uses the prior action taken to advance the deterministic dynamics.
        """
        rnn_input = torch.cat([prev_state.stoch, action], dim=-1)
        deter = self.rnn(rnn_input, prev_state.deter)
        
        post_stats = self.post_net(torch.cat([deter, embed], dim=-1))
        post_mean, post_std = torch.chunk(post_stats, 2, dim=-1)
        post_std = F.softplus(post_std) + 0.1
        
        prior_stats = self.prior_net(deter)
        prior_mean, prior_std = torch.chunk(prior_stats, 2, dim=-1)
        prior_std = F.softplus(prior_std) + 0.1
        
        stoch = post_mean + post_std * torch.randn_like(post_std)
        new_state = RSSMState(stochastic=stoch, deterministic=deter)
        
        return new_state, (post_mean, post_std), (prior_mean, prior_std)
        
    def imagine_step(self, prev_state: RSSMState, action: torch.Tensor) -> RSSMState:
        """
        Single step of prior imagination without real observations.
        Requires the action intended to be taken.
        """
        rnn_input = torch.cat([prev_state.stoch, action], dim=-1)
        deter = self.rnn(rnn_input, prev_state.deter)
        
        prior_stats = self.prior_net(deter)
        prior_mean, prior_std = torch.chunk(prior_stats, 2, dim=-1)
        prior_std = F.softplus(prior_std) + 0.1
        
        stoch = prior_mean + prior_std * torch.randn_like(prior_std)
        new_state = RSSMState(stochastic=stoch, deterministic=deter)
        
        return new_state