import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from gym.spaces import Tuple

from forks.pytorch_a2c_ppo_acktr_gail.a2c_ppo_acktr.distributions import (
    Bernoulli,
    Categorical,
    DiagGaussian,
    DistributionGeneratorTuple,
)
from forks.pytorch_a2c_ppo_acktr_gail.a2c_ppo_acktr.utils import init


def create_output_distribution(action_space, output_size):
    if action_space.__class__.__name__ == "Discrete":
        num_outputs = action_space.n
        dist = Categorical(output_size, num_outputs)
    elif action_space.__class__.__name__ == "Box":
        num_outputs = action_space.shape[0]
        dist = DiagGaussian(output_size, num_outputs)
    elif action_space.__class__.__name__ == "MultiBinary":
        num_outputs = action_space.shape[0]
        dist = Bernoulli(output_size, num_outputs)
    elif action_space.__class__.__name__ == "MultiDiscrete":
        num_outputs = action_space.shape[0]
        dist = DiagGaussian(output_size, num_outputs)

    elif action_space.__class__.__name__ == "Tuple":
        dists = [
            create_output_distribution(space, output_size) for space in action_space
        ]
        # for space in action_space:
        #     print(action_space.__class__.__name__)

        dist = DistributionGeneratorTuple(tuple(dists))
    else:
        raise NotImplementedError
    return dist


class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)


class Policy(nn.Module):
    def __init__(
        self,
        obs_shape,
        action_space,
        base=None,
        base_kwargs=None,
        dist=None,
        obs_space=None,
        symbolic_action_size=0,
    ):
        super(Policy, self).__init__()
        if base_kwargs is None:
            base_kwargs = {}
        if isinstance(base, NNBase):
            self.base = base
        elif base is None:
            if len(obs_shape) == 3:
                base = CNNBase
            elif len(obs_shape) == 1:
                base = MLPBase
            else:
                raise NotImplementedError
            self.base = base(obs_shape[0], **base_kwargs)

        if dist is not None:
            self.dist = dist
        else:
            self.dist = create_output_distribution(action_space, self.base.output_size)

        self.shape = None
        if isinstance(obs_space, Tuple):
            self.shape = (
                obs_space[0].shape,
                obs_space[1].shape[0] + symbolic_action_size,
            )

    @property
    def is_recurrent(self):
        return self.base.is_recurrent

    @property
    def recurrent_hidden_state_size(self):
        """Size of rnn_hx."""
        return self.base.recurrent_hidden_state_size

    def forward(self, inputs, rnn_hxs, masks):
        raise NotImplementedError

    def act(self, inputs, rnn_hxs, masks, deterministic=False):
        # calls CNN or MLP base. Rnn HXS is unused if not recurrent
        # value is critic output, actor features are actor output.
        value, actor_features, rnn_hxs = self.base(inputs, rnn_hxs, masks)
        # dist is output distribution (varies depending on action space)
        dist = self.dist(actor_features)

        if deterministic:
            action = dist.mode()
        else:
            action = dist.sample()

        action_log_probs = dist.log_probs(action)
        dist_entropy = dist.entropy().mean()
        probs = dist.get_probs()
        # print(list(probs[0].cpu()))
        # print(action)
        return value, action, action_log_probs, rnn_hxs, probs

    def get_value(self, inputs, rnn_hxs, masks):
        inputs = self._try_convert(inputs)
        value, _, _ = self.base(inputs, rnn_hxs, masks)
        return value

    def evaluate_actions(self, inputs, rnn_hxs, masks, action):
        # given an sample from the replay buffer, returns the state value,
        # the action log prob (for the chosen action) and the entropy of the distribution
        inputs = self._try_convert(inputs)
        value, actor_features, rnn_hxs = self.base(inputs, rnn_hxs, masks)
        dist = self.dist(actor_features)

        action_log_probs = dist.log_probs(action)
        dist_entropy = dist.entropy().mean()

        return value, action_log_probs, dist_entropy, rnn_hxs

    def _try_convert(self, inputs):
        if self.shape is not None:
            return _unflatten_tuple(self.shape, inputs)
        return inputs


def _unflatten_tuple(shape, _tensor):
    """Assumes tensor is 2 dimensional. converts (n,c*h*w+x) -> ((n,c, h, w),(n, x))"""
    chw = np.prod(shape[0])
    x = shape[1]
    flat, fc = torch.split(_tensor, [chw, x], dim=1)
    return flat.view(_tensor.shape[0], *shape[0]), fc


class NNBase(nn.Module):
    def __init__(self, recurrent, recurrent_input_size, hidden_size):
        super(NNBase, self).__init__()

        self._hidden_size = hidden_size
        self._recurrent = recurrent

        if recurrent:
            self.gru = nn.GRU(recurrent_input_size, hidden_size)
            for name, param in self.gru.named_parameters():
                if "bias" in name:
                    nn.init.constant_(param, 0)
                elif "weight" in name:
                    nn.init.orthogonal_(param)

    @property
    def is_recurrent(self):
        return self._recurrent

    @property
    def recurrent_hidden_state_size(self):
        if self._recurrent:
            return self._hidden_size
        return 1

    @property
    def output_size(self):
        return self._hidden_size

    def _forward_gru(self, x, hxs, masks):
        if x.size(0) == hxs.size(0):
            x, hxs = self.gru(x.unsqueeze(0), (hxs * masks).unsqueeze(0))
            x = x.squeeze(0)
            hxs = hxs.squeeze(0)
        else:
            # x is a (T, N, -1) tensor that has been flatten to (T * N, -1)
            N = hxs.size(0)
            T = int(x.size(0) / N)

            # unflatten
            x = x.view(T, N, x.size(1))

            # Same deal with masks
            masks = masks.view(T, N)

            # Let's figure out which steps in the sequence have a zero for any agent
            # We will always assume t=0 has a zero in it as that makes the logic cleaner
            has_zeros = (masks[1:] == 0.0).any(dim=-1).nonzero().squeeze().cpu()

            # +1 to correct the masks[1:]
            if has_zeros.dim() == 0:
                # Deal with scalar
                has_zeros = [has_zeros.item() + 1]
            else:
                has_zeros = (has_zeros + 1).numpy().tolist()

            # add t=0 and t=T to the list
            has_zeros = [0] + has_zeros + [T]

            hxs = hxs.unsqueeze(0)
            outputs = []
            for i in range(len(has_zeros) - 1):
                # We can now process steps that don't have any zeros in masks together!
                # This is much faster
                start_idx = has_zeros[i]
                end_idx = has_zeros[i + 1]

                rnn_scores, hxs = self.gru(
                    x[start_idx:end_idx], hxs * masks[start_idx].view(1, -1, 1)
                )

                outputs.append(rnn_scores)

            # assert len(outputs) == T
            # x is a (T, N, -1) tensor
            x = torch.cat(outputs, dim=0)
            # flatten
            x = x.view(T * N, -1)
            hxs = hxs.squeeze(0)

        return x, hxs


class CNNBase(NNBase):
    def __init__(
        self,
        num_inputs,
        recurrent=False,
        hidden_size=512,
        fc_size=0,
        deep=False,
        conv=None,
    ):
        """ num inputs is the number of channels """
        super(CNNBase, self).__init__(recurrent, hidden_size, hidden_size)

        init_ = lambda m: init(
            m,
            nn.init.orthogonal_,
            lambda x: nn.init.constant_(x, 0),
            nn.init.calculate_gain("relu"),
        )
        if conv:
            self.main = conv
        else:
            self.main = nn.Sequential(
                init_(nn.Conv2d(num_inputs, 32, 8, stride=4)),
                nn.ReLU(),
                init_(nn.Conv2d(32, 64, 4, stride=2)),
                nn.ReLU(),
                init_(nn.Conv2d(64, 32, 3, stride=1)),
                nn.ReLU(),
                Flatten(),
            )

        if deep:
            self.fc = nn.Sequential(
                init_(nn.Linear(32 * 7 * 7 + fc_size, hidden_size)),
                nn.ReLU(),
                init_(nn.Linear(hidden_size, hidden_size)),
                nn.ReLU(),
            )
        else:
            self.fc = nn.Sequential(
                init_(nn.Linear(32 * 7 * 7 + fc_size, hidden_size)), nn.ReLU()
            )
        self.fc_size = fc_size if fc_size else None

        init_ = lambda m: init(
            m, nn.init.orthogonal_, lambda x: nn.init.constant_(x, 0)
        )

        self.critic_linear = init_(nn.Linear(hidden_size, 1))

        self.train()

    def forward(self, inputs, rnn_hxs, masks):
        if self.fc_size is None:
            x = self.fc(self.main(inputs))
        else:
            cnn, fc = inputs
            fc_in = self.main(cnn / 255)
            # print(
            #     f"dense max: {fc.max()} min: {fc.min()} cnn_out max: {fc_in.max()} min: {fc_in.min()}"
            # )
            x = self.fc(torch.cat((fc_in, fc), dim=1))

        if self.is_recurrent:
            x, rnn_hxs = self._forward_gru(x, rnn_hxs, masks)

        return self.critic_linear(x), x, rnn_hxs


class MLPBase(NNBase):
    def __init__(self, num_inputs, recurrent=False, hidden_size=64):
        super(MLPBase, self).__init__(recurrent, num_inputs, hidden_size)

        if recurrent:
            num_inputs = hidden_size

        init_ = lambda m: init(
            m, nn.init.orthogonal_, lambda x: nn.init.constant_(x, 0), np.sqrt(2)
        )

        self.actor = nn.Sequential(
            init_(nn.Linear(num_inputs, hidden_size)),
            nn.Tanh(),
            init_(nn.Linear(hidden_size, hidden_size)),
            nn.Tanh(),
        )

        self.critic = nn.Sequential(
            init_(nn.Linear(num_inputs, hidden_size)),
            nn.Tanh(),
            init_(nn.Linear(hidden_size, hidden_size)),
            nn.Tanh(),
        )

        self.critic_linear = init_(nn.Linear(hidden_size, 1))

        self.train()

    def forward(self, inputs, rnn_hxs, masks):
        x = inputs

        if self.is_recurrent:
            x, rnn_hxs = self._forward_gru(x, rnn_hxs, masks)

        hidden_critic = self.critic(x)
        hidden_actor = self.actor(x)

        return self.critic_linear(hidden_critic), hidden_actor, rnn_hxs
