"""Credit: mostly based on Ilya's excellent implementation here: https://github.com/ikostrikov/pytorch-flows"""
import numpy as np
import torch
import torch.nn as nn
from torch.nn import functional as F


class InverseAutoregressiveFlow(nn.Module):
    """Inverse Autoregressive Flows with LSTM-type update. One block.
  
  Eq 11-14 of https://arxiv.org/abs/1606.04934
  """

    def __init__(self, num_input, num_hidden, num_context):
        super().__init__()
        self.made = MADE(num_input=num_input, num_output=num_input * 2,
                         num_hidden=num_hidden, num_context=num_context)
        # init such that sigmoid(s) is close to 1 for stability
        self.sigmoid_arg_bias = nn.Parameter(torch.ones(num_input) * 2)
        self.sigmoid = nn.Sigmoid()
        self.log_sigmoid = nn.LogSigmoid()

    def forward(self, input, context=None):
        m, s = torch.chunk(self.made(input, context), chunks=2, dim=-1)
        s = s + self.sigmoid_arg_bias
        sigmoid = self.sigmoid(s)
        z = sigmoid * input + (1 - sigmoid) * m
        return z, -self.log_sigmoid(s)


class FlowSequential(nn.Sequential):
    """Forward pass."""

    def forward(self, input, context=None):
        total_log_prob = torch.zeros_like(input, device=input.device)
        for block in self._modules.values():
            input, log_prob = block(input, context)
            total_log_prob += log_prob
        return input, total_log_prob


class MaskedLinear(nn.Module):
    """Linear layer with some input-output connections masked."""

    def __init__(self, in_features, out_features, mask, context_features=None, bias=True):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features, bias)
        self.register_buffer("mask", mask)
        if context_features is not None:
            self.cond_linear = nn.Linear(context_features, out_features, bias=False)

    def forward(self, input, context=None):
        output = F.linear(input, self.mask * self.linear.weight, self.linear.bias)
        if context is None:
            return output
        else:
            return output + self.cond_linear(context)


class MADE(nn.Module):
    """Implements MADE: Masked Autoencoder for Distribution Estimation.

  Follows https://arxiv.org/abs/1502.03509

  This is used to build MAF: Masked Autoregressive Flow (https://arxiv.org/abs/1705.07057).
  """

    def __init__(self, num_input, num_output, num_hidden, num_context):
        super().__init__()
        # m corresponds to m(k), the maximum degree of a node in the MADE paper
        self._m = []
        self._masks = []
        self._build_masks(num_input, num_output, num_hidden, num_layers=3)
        self._check_masks()
        modules = []
        self.input_context_net = MaskedLinear(num_input, num_hidden, self._masks[0], num_context)
        modules.append(nn.ReLU())
        modules.append(MaskedLinear(num_hidden, num_hidden, self._masks[1], context_features=None))
        modules.append(nn.ReLU())
        modules.append(MaskedLinear(num_hidden, num_output, self._masks[2], context_features=None))
        self.net = nn.Sequential(*modules)

    def _build_masks(self, num_input, num_output, num_hidden, num_layers):
        """Build the masks according to Eq 12 and 13 in the MADE paper."""
        rng = np.random.RandomState(0)
        # assign input units a number between 1 and D
        self._m.append(np.arange(1, num_input + 1))
        for i in range(1, num_layers + 1):
            # randomly assign maximum number of input nodes to connect to
            if i == num_layers:
                # assign output layer units a number between 1 and D
                m = np.arange(1, num_input + 1)
                assert num_output % num_input == 0, "num_output must be multiple of num_input"
                self._m.append(np.hstack([m for _ in range(num_output // num_input)]))
            else:
                # assign hidden layer units a number between 1 and D-1
                self._m.append(rng.randint(1, num_input, size=num_hidden))
                # self._m.append(np.arange(1, num_hidden + 1) % (num_input - 1) + 1)
            if i == num_layers:
                mask = self._m[i][None, :] > self._m[i - 1][:, None]
            else:
                # input to hidden & hidden to hidden
                mask = self._m[i][None, :] >= self._m[i - 1][:, None]
            # need to transpose for torch linear layer, shape (num_output, num_input)
            self._masks.append(torch.from_numpy(mask.astype(np.float32).T))

    def _check_masks(self):
        """Check that the connectivity matrix between layers is lower triangular."""
        # (num_input, num_hidden)
        prev = self._masks[0].t()
        for i in range(1, len(self._masks)):
            # num_hidden is second axis
            prev = prev @ self._masks[i].t()
        final = prev.numpy()
        num_input = self._masks[0].shape[1]
        num_output = self._masks[-1].shape[0]
        assert final.shape == (num_input, num_output)
        if num_output == num_input:
            assert np.triu(final).all() == 0
        else:
            for submat in np.split(final,
                                   indices_or_sections=num_output // num_input,
                                   axis=1):
                assert np.triu(submat).all() == 0

    def forward(self, input, context=None):
        # first hidden layer receives input and context
        hidden = self.input_context_net(input, context)
        # rest of the network is conditioned on both input and context
        return self.net(hidden)


class Reverse(nn.Module):
    """ An implementation of a reversing layer from
  Density estimation using Real NVP
  (https://arxiv.org/abs/1605.08803).

  From https://github.com/ikostrikov/pytorch-flows/blob/master/main.py
  """

    def __init__(self, num_input):
        super(Reverse, self).__init__()
        self.perm = np.array(np.arange(0, num_input)[::-1])
        self.inv_perm = np.argsort(self.perm)

    def forward(self, inputs, context=None, mode='forward'):
        if mode == "forward":
            return inputs[:, :, self.perm], torch.zeros_like(inputs, device=inputs.device)
        elif mode == "inverse":
            return inputs[:, :, self.inv_perm], torch.zeros_like(inputs, device=inputs.device)
        else:
            raise ValueError("Mode must be one of {forward, inverse}.")
