import numpy as np
import torch
import torch.nn as nn
from torch.nn import init
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.autograd import Variable
from functools import reduce
from operator import mul


def init_linear(linear_layer):
    init.xavier_uniform(linear_layer.weight.data, np.sqrt(2))
    linear_layer.bias.data.fill_(0.0)


def init_linear_in_sequential(sequential_layer):
    for l in sequential_layer:
        if isinstance(l, nn.Linear) or isinstance(l, nn.Conv2d):
            init_linear(l)


def init_layers(linear_layers):
    for l in linear_layers:
        init_linear(l)


class CnnToRecurrent(nn.Module):

    def __init__(self, input_dim, num_actions, convs, hiddens, lstm_size=256):
        super(CnnToRecurrent, self).__init__()

        if not isinstance(input_dim, list):
            raise ValueError('Input size must be a list of sizes')
        if not isinstance(hiddens, list):
            raise ValueError
        if not isinstance(convs, list):
            raise ValueError
        if num_actions <= 0:
            raise ValueError('num_actions must be larger than 0')

        conv_layers = []
        in_channels = input_dim[0]
        for out_channels, kernel_size, stride, padding in convs:
            conv_layers += [nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding), nn.ReLU(inplace=True)]
            in_channels = out_channels

        self.conv_layers = nn.Sequential(*conv_layers)

        fc_layers = []
        conv_flat_size = self._get_conv_output(input_dim)
        sizes = [conv_flat_size] + hiddens
        for i in range(len(sizes) - 1):
            fc_layers += [nn.Linear(sizes[i], sizes[i + 1]), nn.ReLU()]
        self.fc_layers = nn.Sequential(*fc_layers)

        self.lstm = nn.LSTMCell(hiddens[-1], lstm_size)

        self.value_layer = nn.Linear(lstm_size, 1)
        self.policy_layer = nn.Linear(lstm_size, num_actions)

        self.num_actions = num_actions

        self._init_layers()

    def _init_layers(self):
        sequentials = [self.fc_layers]
        parameterised_layers = []
        for s in sequentials:
            for l in s:
                if isinstance(l, nn.Linear) or isinstance(l, nn.Conv2d):
                    parameterised_layers.append(l)
        init_layers(parameterised_layers)

        init.xavier_uniform(self.lstm.weight_ih.data, 1)
        init.xavier_uniform(self.lstm.weight_hh.data, 1)
        self.lstm.bias_ih.data.fill_(0)
        self.lstm.bias_hh.data.fill_(0)

        init.xavier_uniform(self.policy_layer.weight.data, 1)
        self.policy_layer.bias.data.fill_(0)
        init.xavier_uniform(self.value_layer.weight.data, 1)
        self.value_layer.bias.data.fill_(0)

    def forward(self, x, hx, cx):
        dims = x.size()
        x = x.view(dims[0], dims[3], dims[1], dims[2]).float()
        x /= 255.
        x = self.conv_layers(x)
        x = x.view(1, -1)
        x = self.fc_layers(x)
        hx, cx = self.lstm(x, (hx, cx))
        x = hx

        action_logits = self.policy_layer(x)
        value = self.value_layer(x)

        return action_logits.squeeze(), value.squeeze(), hx, cx

    def _get_conv_output(self, shape):
        batch_size = 1
        x = Variable(torch.rand(batch_size, *shape))
        output_feat = self.conv_layers(x)
        n_size = output_feat.data.view(batch_size, -1).size(1)
        return n_size

    def sync_parameters(self, shared_model):
        self.load_state_dict(shared_model.state_dict())


def cnn_to_recurrent(input_dim, num_actions, convs, hiddens, lstm_size):
    """This model takes as input an observation and returns values of all actions.
    Parameters
    ----------
    input_dim: [int, int, int]
        list of the dimensions of the input
        (num_channels, height, width)
    num_actions: int
        number of possible actions
    convs: [(int, int int)]
        list of convolutional layers in form of
        (out_channels, kernel_size, stride, padding)
    hiddens: [int]
        list of sizes of hidden layers
    lstm_size: int
        size of the hidden size of the LSTM
    Returns
    -------
    policy: nn.Module
        The A3C policy
    """
    return CnnToRecurrent(input_dim, num_actions, convs, hiddens, lstm_size)


class CnnToMlp(nn.Module):

    def __init__(self, input_dim, num_actions, convs, hiddens):
        super(CnnToMlp, self).__init__()

        if not isinstance(input_dim, list):
            raise ValueError('Input size must be a list of sizes')
        if not isinstance(hiddens, list):
            raise ValueError
        if not isinstance(convs, list):
            raise ValueError
        if num_actions <= 0:
            raise ValueError('num_actions must be larger than 0')

        conv_layers = []
        in_channels = input_dim[0]
        for out_channels, kernel_size, stride in convs:
            conv_layers += [nn.Conv2d(in_channels, out_channels, kernel_size, stride), nn.ReLU(inplace=True)]
            in_channels = out_channels

        self.conv_layers = nn.Sequential(*conv_layers)

        fc_layers = []
        conv_flat_size = self._get_conv_output(input_dim)
        sizes = [conv_flat_size] + hiddens
        for i in range(len(sizes) - 1):
            fc_layers += [nn.Linear(sizes[i], sizes[i + 1]), nn.ReLU(inplace=True)]
        self.fc_layers = nn.Sequential(*fc_layers)

        self.value_layer = nn.Linear(hiddens[-1], 1)
        self.policy_layer = nn.Linear(hiddens[-1], num_actions)

        self.num_actions = num_actions

        self._init_layers()

    def _init_layers(self):
        sequentials = [self.conv_layers, self.fc_layers]
        parameterised_layers = []
        for s in sequentials:
            for l in s:
                if isinstance(l, nn.Linear) or isinstance(l, nn.Conv2d):
                    parameterised_layers.append(l)
        init_layers(parameterised_layers)

        init.xavier_uniform(self.policy_layer.weight.data, 1)
        self.policy_layer.bias.data.fill_(0)
        init.xavier_uniform(self.value_layer.weight.data, 1)
        self.value_layer.bias.data.fill_(0)

    def forward(self, x, hx, cx):
        dims = x.size()

        x = x.view(dims[0], dims[3], dims[1], dims[2]).float() / 255.
        x = self.conv_layers(x)
        x = x.view(1, -1)
        x = self.fc_layers(x)

        action_logits = self.policy_layer(x)
        value = self.value_layer(x)

        return action_logits.squeeze(), value.squeeze(), hx, cx

    def _get_conv_output(self, shape):
        batch_size = 1
        x = Variable(torch.rand(batch_size, *shape))
        output_feat = self.conv_layers(x)
        n_size = output_feat.data.view(batch_size, -1).size(1)

        return n_size

    def sync_parameters(self, shared_model):
        self.load_state_dict(shared_model.state_dict())


def cnn_to_mlp(input_dim, num_actions, convs, hiddens):
    """This model takes as input an observation and returns values of all actions.
    Parameters
    ----------
    input_dim: [int, int, int]
        list of the dimensions of the input
        (num_channels, height, width)
    num_actions: int
        number of possible actions
    convs: [(int, int int)]
        list of convolutional layers in form of
        (out_channels, kernel_size, stride)
    hiddens: [int]
        list of sizes of hidden layers
    Returns
    -------
    policy: nn.Module
        The A2C policy
    """
    return CnnToMlp(input_dim, num_actions, convs, hiddens)


class Mlp(nn.Module):
    def __init__(self, input_dim, num_actions, hiddens):
        super(Mlp, self).__init__()

        if not isinstance(input_dim, list):
            raise ValueError('Input size must be a list of sizes')
        if not isinstance(hiddens, list):
            raise ValueError('hiddens size must be a list of sizes')
        if num_actions <= 0:
            raise ValueError('num_actions must be larger than 0')

        self.num_actions = num_actions

        layers = []
        self.input_flat_size = reduce(mul, input_dim, 1)
        sizes = [self.input_flat_size] + hiddens
        for i in range(len(sizes)-1):
            layers += [nn.Linear(sizes[i], sizes[i+1]), nn.ReLU(inplace=True)]

        self.net = nn.Sequential(*layers)
        self.policy_layer = nn.Linear(sizes[-1], num_actions)
        self.value_layer = nn.Linear(sizes[-1], 1)

        self._init_layers()

    def _init_layers(self):
        parameterised_layers = [self.value_layer, self.policy_layer]
        for l in self.net.modules():
            if isinstance(l, nn.Linear):
                parameterised_layers.append(l)
        init_layers(parameterised_layers)

        init.xavier_uniform(self.policy_layer.weight.data, 1)
        self.policy_layer.bias.data.fill_(0)
        init.xavier_uniform(self.value_layer.weight.data, 1)
        self.value_layer.bias.data.fill_(0)

    def forward(self, x, hx, cx):
        # x = x.view(-1, self.input_flat_size)
        x = self.net(x)

        policy_logit = self.policy_layer(x)
        value = self.value_layer(x)
        return policy_logit.squeeze(), value.squeeze(), hx, cx

    def sync_parameters(self, shared_net):
        self.load_state_dict(shared_net.state_dict())


def mlp(input_dim, num_actions, hiddens):
    """This model takes as input an observation and returns values of all actions.
    Parameters
    ----------
    input_dim: [int]
        list of the dimensions of the input
    num_actions: int
        number of possible actions
    hiddens: [int]
        list of sizes of hidden layers
    Returns
    -------
    q_func: nn.Module
        q_function for DQN algorithm.
    """
    model = Mlp(input_dim, num_actions, hiddens)
    return model
