import torch

class ParameterisedPolicy(torch.nn.Module):
    """
    REINFORCE RL agent class. Returns action when the ParameterisedPolicy.act(observation) method is used.
    observation is a gym state vector.
    obs_len - length of the state vector
    act_space_len - length of the action vector
    
    """
    def __init__(self, obs_len=8, act_space_len=2):
        super().__init__()
        self.obs_len = obs_len
        self.act_space_len = act_space_len
        self.lin_1 = torch.nn.Linear(self.obs_len, 256)
        self.rel_1 = torch.nn.ReLU()
        
        self.lin_2 = torch.nn.Linear(256, 128)
        self.rel_2 = torch.nn.ReLU()
            
        self.lin_3 = torch.nn.Linear(128, self.act_space_len)
        
        self.lin_4 = torch.nn.Linear(128, self.act_space_len)
        self.elu = torch.nn.ELU()
        
        
    def forward(self, x):
        x = self.lin_1(x)
        x = self.rel_1(x)
        
        x = self.lin_2(x)
        x = self.rel_2(x)
        
        mu = self.lin_3(x)
        
        x = self.lin_4(x)
        sigma = self.elu(x) + 1.000001
        
        return mu, sigma
    
    def act(self, observation):
        """
        Method returns action when gym state vector is passed.
        """
        (mus, sigmas) = self.forward(torch.tensor(observation, dtype=torch.float32))
        m = torch.distributions.normal.Normal(mus, sigmas)
        action = m.sample().detach().numpy()
        
        return action