import numpy as np
import scipy.signal

import torch
import torch.nn as nn


def combined_shape(length, shape=None):
    if shape is None:
        return (length,)
    return (length, shape) if np.isscalar(shape) else (length, *shape)

def mlp(sizes, activation, output_activation=nn.Identity):
    layers = []
    for j in range(len(sizes)-1):
        act = activation if j < len(sizes)-2 else output_activation
        layers += [nn.Linear(sizes[j], sizes[j+1]), act()]
    return nn.Sequential(*layers)

def count_vars(module):
    return sum([np.prod(p.shape) for p in module.parameters()])

class MLPActor(nn.Module):

    def __init__(self, obs_dim, act_dim, hidden_sizes, activation, act_limit):
        super().__init__()
        pi_sizes = [obs_dim] + list(hidden_sizes) + [act_dim]
        self.pi = mlp(pi_sizes, activation, nn.Tanh)
        self.act_limit = act_limit

    def forward(self, obs):
        # Return output from network scaled to action space limits.
        return self.act_limit * self.pi(obs)

class MLPQFunction(nn.Module):

    def __init__(self, obs_dim, act_dim, hidden_sizes, activation):
        super().__init__()
        self.q = mlp([obs_dim + act_dim] + list(hidden_sizes) + [1], activation)

    def forward(self, obs, act):
        q = self.q(torch.cat([obs, act], dim=-1))
        return torch.squeeze(q, -1) # Critical to ensure q has right shape.

class MLPActorCritic(nn.Module):

    def __init__(self, observation_space, action_space, hidden_sizes=(256,256),
                 activation=nn.ReLU):
        super().__init__()

        obs_dim = observation_space.shape[0]
        act_dim = action_space.shape[0]
        act_limit = action_space.high[0]

        # build policy and value functions
        self.pi = MLPActor(obs_dim, act_dim, hidden_sizes, activation, act_limit)
        self.q1 = MLPQFunction(obs_dim, act_dim, hidden_sizes, activation)
        self.q2 = MLPQFunction(obs_dim, act_dim, hidden_sizes, activation)

    def act(self, obs):
        with torch.no_grad():
            return self.pi(obs).numpy()


class ForwardModel(nn.Module):
    def __init__(self,
                 obs_dim,
                 act_dim,
                 hidden_sizes=(256,256), 
                 activation=nn.ReLU
                 ):
        super().__init__()

        self.mlp = mlp([obs_dim + act_dim] + list(hidden_sizes) + [obs_dim], activation)

    def forward(self, obs, act):
        pred_obs = self.mlp(torch.cat([obs, act], dim=-1))
        return pred_obs

    def compute_loss(self,
                     obs,
                     act,
                     delta_obs,
                     device='cpu', 
                     reduction='elementwise_mean'):
        obs = torch.as_tensor(obs).to(device).float()
        act = torch.as_tensor(act).to(device).float()
        pred_delta = self.forward(obs, act)

        loss = self.loss_function(
            pred_delta, delta_pose, 
            device=device, reduction=reduction
        )
        return loss

    def loss_function(self,
                      pred_delta,
                      delta_pose,
                      device='cpu', 
                      reduction='elementwise_mean'):
        pred_delta = torch.as_tensor(pred_delta).to(device).float()
        delta_pose = torch.as_tensor(delta_pose).to(device).float()
        mse_loss = torch.mean((pred_delta - delta_pose) ** 2, dim=-1)
        if reduction == 'elementwise_mean':
            mse_loss = mse_loss.mean()
        return mse_loss
