import math

import torch
import torch.nn as nn

from NormalizingFlows.made_submodules.mades import GaussianMadeBN


class MaskedAutoregressiveFlow(nn.Module):
    """
        Implements a Masked Autoregressive Flow, which is a stack of mades such that the random numbers which drive made i
        are generated by made i-1. The first made is driven by standard gaussian noise. In the current implementation, all
        mades are of the same type. If there is only one made in the stack, then it's equivalent to a single made.
        """
    def __init__(self, in_size, hidden_sizes, n_mades, batch_norm=True):
        super(MaskedAutoregressiveFlow, self).__init__()
        self.bn = batch_norm
        self.models = nn.ModuleList([])
        order = lambda ind: 'default' if ind == 0 else 'reverse'
        for i in range(n_mades):
            self.models.append(GaussianMadeBN(in_size, hidden_sizes, order(i)))

    def forward(self, x):
        out = x
        log_det_du_dx = 0.
        for i, layer in enumerate(self.models):
            log_det_inverse, u = layer(out)
            log_det_du_dx += log_det_inverse
            out = u

        log_probs = torch.sum(-0.5 * (math.log(2 * math.pi) + u**2), dim=1) + log_det_du_dx
        return log_probs, u

    def reverse(self, u):
        with torch.no_grad():
            for i, layer in enumerate(reversed(self.models)):
                u = layer.reverse(u)
        return u








