text: '''
  =====
  Distributed by: Notre Dame SCAI Lab (MIT Liscense)
  - Associated publication:
  url: http://aimsciences.org//article/id/3a9f3d14-3421-4947-a45f-a9cc74edd097
  doi: https://dx.doi.org/10.3934/fods.2020019
  github: https://github.com/zabaras/deep-turbulence
  =====
  '''
  import sys, time
  sys.path.append(".")
  # sys.path.append("../..")
  
  from nn.modules.actNorm import ActNorm
  from nn.modules.glowConv import InvertibleConv1x1, InvertibleConv1x1LU
  from nn.modules.flowAffine import LSTMAffineCouplingLayer, AffineCouplingLayer
  from nn.modules.flowUtils import Squeeze, CheckerSqueeze, Split
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import numpy as np
  
  class AffineCouplingBlock(nn.Module):
      """An invertible affine coupling block consisting of an activation
      normalization, conditional affine coupling layer and a 1x1 convolution.
      See Fig. 6 of paper: https://arxiv.org/abs/2006.04731
  
      :param in_features: Number of input feature channels
      :type in_features: int
      :param cond_features: Number of conditional feature channels
      :type cond_features: int
      :param train_sampling: train by sampling (1x1 conv. is inverted for the forward pass), defaults to True
      :type train_sampling: bool, optional
      :param LUdecompose: Use the LU decomposition for 1x1 convolution, defaults to False
      :type LUdecompose: bool, optional
      """
      def __init__(self, in_features, cond_features, 
          train_sampling=True, LUdecompose=False):
          """Constructor method
          """ 
          super(AffineCouplingBlock, self).__init__()
          self.norm = ActNorm(in_features)
          if LUdecompose:
              self.conv = InvertibleConv1x1LU(in_features, 
                  train_sampling=train_sampling)
          else:
              self.conv = InvertibleConv1x1(in_features,
                  train_sampling=train_sampling)
  
          self.coupling = AffineCouplingLayer(in_features, cond_features)
  
      def forward(self, x, cond):
          """Forward pass
  
          :param x:  [B, in_features, H, W] input feature tensor
          :type x: torch.Tensor
          :param cond: [B, cond_features, H, W] conditional feature tensor
          :type cond: torch.Tensor
          :returns: 
              - y: [B, in_features, H, W] Output feature tensor
              - logdet: log determinate of affine block
          :rtype: (torch.Tensor, torch.Tensor)
          """
          out, logdet1 = self.norm(x)
          out, logdet2 = self.conv(out)
          y, logdet3 = self.coupling(out, cond)
  
          return y, logdet3 + logdet2 + logdet1
  
      def reverse(self, y, cond):
          """Backward pass
  
          :param y:  [B, in_features, H, W] input feature tensor
          :type y: torch.Tensor
          :param cond: [B, cond_features, H, W] conditional feature tensor
          :type cond: torch.Tensor
          :returns: 
              - x: [B, in_features, H, W] Output feature tensor
              - logdet: log determinate of affine block
          :rtype: (torch.Tensor, torch.Tensor)
          """
          out, logdet1 = self.coupling.reverse(y, cond)
          out, logdet2 = self.conv.reverse(out)
          x, logdet3 = self.norm.reverse(out)
          return x, logdet3 + logdet2 + logdet1
  
  class UnNormedAffineCouplingBlock(nn.Module):
      """An unnormalized invertible affine coupling block consisting 
      of an conditional affine coupling layer and a 1x1 convolution.
      See Fig. 6 of paper: https://arxiv.org/abs/2006.04731
  
      :param in_features: Number of input feature channels
      :type in_features: int
      :param cond_features: Number of conditional feature channels
      :type cond_features: int
      :param train_sampling: train by sampling (1x1 conv. is inverted for the forward pass), defaults to True
      :type train_sampling: bool, optional
      :param LUdecompose: Use the LU decomposition for 1x1 convolution, defaults to False
      :type LUdecompose: bool, optional
      """
      def __init__(self, in_features, cond_features, 
          train_sampling=True, LUdecompose=False):
          """Constructor method
          """ 
          super(UnNormedAffineCouplingBlock, self).__init__()
          if LUdecompose:
              self.conv = InvertibleConv1x1LU(in_features, 
                  train_sampling=train_sampling)
          else:
              self.conv = InvertibleConv1x1(in_features,
                  train_sampling=train_sampling)
  
          self.coupling = AffineCouplingLayer(in_features, cond_features)
  
      def forward(self, x, cond):
          """Forward pass
  
          :param x:  [B, in_features, H, W] input feature tensor
          :type x: torch.Tensor
          :param cond: [B, cond_features, H, W] conditional feature tensor
          :type cond: torch.Tensor
          :returns: 
              - y: [B, in_features, H, W] Output feature tensor
              - logdet: log determinate of affine block
          :rtype: (torch.Tensor, torch.Tensor)
          """
          out, logdet1 = self.conv(x)
          y, logdet2 = self.coupling(out, cond)
          return y, logdet1 + logdet2
  
      def reverse(self, y, cond):
          """Backward pass
  
          :param y:  [B, in_features, H, W] input feature tensor
          :type y: torch.Tensor
          :param cond: [B, cond_features, H, W] conditional feature tensor
          :type cond: torch.Tensor
          :returns: 
              - x: [B, in_features, H, W] Output feature tensor
              - logdet: log determinate of affine block
          :rtype: (torch.Tensor, torch.Tensor)
          """
          out, logdet1 = self.coupling.reverse(y, cond)
          x, logdet2 = self.conv.reverse(out)
          return x, logdet1 + logdet2
  
  class LSTMCouplingBlock(nn.Module):
      """An LSTM invertible affine coupling block consisting 
      of an activation normalization, lstm affine coupling layer and a 1x1 convolution.
      See Fig. 6 of paper: https://arxiv.org/abs/2006.04731
  
      :param in_features: Number of input feature channels
      :type in_features: int
      :param cond_features: Number of conditional feature channels
      :type cond_features: int
      :param rec_features: Number of recurrent feature channels in :class:`nn.modules.convLSTM.LSTMAffineCouplingLayer`
      :type rec_features: int
      :param train_sampling: train by sampling (1x1 conv. is inverted for the forward pass), defaults to True
      :type train_sampling: bool, optional
      :param LUdecompose: Use the LU decomposition for 1x1 convolution, defaults to False
      :type LUdecompose: bool, optional
      """
      def __init__(self, in_features, cond_features, rec_features,
          train_sampling=True, LUdecompose=False):
          """Constructor method
          """ 
          super(LSTMCouplingBlock, self).__init__()
          self.norm = ActNorm(in_features)
          self.norm2 = ActNorm(in_features)
          if LUdecompose:
              self.conv = InvertibleConv1x1LU(in_features, 
                  train_sampling=train_sampling)
          else:
              self.conv = InvertibleConv1x1(in_features,
                  train_sampling=train_sampling)
          
          self.coupling = LSTMAffineCouplingLayer(in_features, cond_features, rec_features)
  
      def forward(self, x, cond, rec_states=None):
          """Forward pass
  
          :param x:  [B, in_features, H, W] input feature tensor
          :type x: torch.Tensor
          :param cond: [B, cond_features, H, W] conditional feature tensor
          :type cond: torch.Tensor
          :param rec_states: tuple of LSTM states (hidden state, cell state), defaults to None
          :type rec_states: tuple, optional
          :returns: 
              - y: [B, in_features, H, W] Output feature tensor
              - logdet: log determinate of affine block
              - states_out: tuple of LSTM (cell, hidden) states
          :rtype: (torch.Tensor, torch.Tensor, tuple)
          """
          out, logdet1 = self.norm(x)
          out, logdet2 = self.conv(out)
          y, logdet3, out_states = self.coupling(out, cond, rec_states)
          return y, logdet3 + logdet2 + logdet1, out_states
  
      def reverse(self, y, cond, rec_states=None):
          """Backward pass
  
          :param y:  [B, in_features, H, W] input feature tensor
          :type y: torch.Tensor
          :param cond: [B, cond_features, H, W] conditional feature tensor
          :type cond: torch.Tensor
          :param rec_states: tuple of LSTM states (hidden state, cell state), defaults to None
          :type rec_states: tuple, optional
          :returns: 
              - x: [B, in_features, H, W] Output feature tensor
              - logdet: log determinate of affine block
              - states_out: tuple of LSTM (cell, hidden) states
          :rtype: (torch.Tensor, torch.Tensor, tuple)
          """
          out, logdet1, out_states = self.coupling.reverse(y, cond, rec_states)
          out, logdet2 = self.conv.reverse(out)
          x, logdet3 = self.norm.reverse(out)
          return x, logdet3+logdet2+logdet1, out_states
  
  class LSTMFLowBlock(nn.Module):
      """LSTM Conditional Flow block considting of an a stack of
      invertible affine layers.
      See Fig. 6 of paper: https://arxiv.org/abs/2006.04731
  
      :param in_features: Number of input feature channels
      :type in_features: int
      :param cond_features: Number of conditional feature channels
      :type cond_features: int
      :param rec_features: Number of recurrent feature channels in LSTM
      :type rec_features: int
      :param n_layers: Number of affine layers in block
      :type n_layers: int
      :param factor: Squeeze factor to reduce the dimensionality of the feature map, defaults to 2
      :type factor: int, optional
      :param train_sampling: Train by sampling (1x1 conv. is inverted for the forward pass), defaults to True
      :type train_sampling: bool, optional
      :param LUdecompose: Use the LU decomposition for 1x1 convolution, defaults to False
      :type LUdecompose: bool, optional
      :param do_split: Perform a split at the end of block, defaults to True. Refer to :class:`nn.modules.flowUtils.Split`
      :type do_split: bool, optional
      :param squeeze_type: Type of squeeze method; 0 = :class:`nn.modules.flowUtils.CheckerSqueeze`, 
          1 = :class:`nn.modules.flowUtils.Squeeze`, defaults to 0
      :type squeeze_type: int, optional
      """
      def __init__(self, in_features, cond_features, rec_features, 
          n_layers, factor=2, LUdecompose=True, train_sampling=False, 
          do_split=True, squeeze_type=0):
          """Constructor method
          """ 
          super(LSTMFLowBlock, self).__init__()
          self.do_split = do_split
          self.n_layers = n_layers
  
          if(squeeze_type == 0):
              self.squeeze = CheckerSqueeze(factor)
          else:
              self.squeeze = Squeeze(factor)
          in_features = in_features * factor ** 2 # Features after squeeze
  
          self.revlayers = nn.Sequential()
          for i in range(n_layers-1):
              # No actnorm or conv first flow layer
              if(i==0):
                  self.revlayers.add_module('affine_layer{}'.format(i+1), 
                      UnNormedAffineCouplingBlock(in_features, cond_features,
                      LUdecompose=LUdecompose, train_sampling=train_sampling))
              else:
                  self.revlayers.add_module('affine_layer{}'.format(i+1), 
                      AffineCouplingBlock(in_features, cond_features,
                      LUdecompose=LUdecompose, train_sampling=train_sampling))
  
          self.revlayers.add_module('affine_layer{}'.format(n_layers), 
                  LSTMCouplingBlock(in_features, cond_features, rec_features,
                  LUdecompose=LUdecompose, train_sampling=train_sampling))
  
          # Split off variables into latent params
          if do_split:
              self.split = Split(in_features)
  
      def forward(self, x, cond, rec_states, return_eps=False):
          """Forward pass
  
          :param x:  [B, in_features, H, W] input feature tensor
          :type x: torch.Tensor
          :param cond: [B, cond_features, H, W] conditional feature tensor
          :type cond: torch.Tensor
          :param rec_states: tuple of LSTM states (hidden state, cell state), defaults to None
          :type rec_states: tuple, optional
          :param return_eps: Return samples from latent densities, defaults to False
          :type return_eps: bool, optional
          :returns: 
              - y: [B, in_features, H, W] Output feature tensor
              - logdet: [B] log determinate of block + the log-likelihood of the latent variables
              - out_states: tuple of LSTM (cell, hidden) states
              - eps: [B, 2*in_features, H//2, W//2] (for factor=2) tensor of sampled latent variables from unit gaussian
          :rtype: (torch.Tensor, torch.Tensor, tuple, torch.Tensor)
          """
          out_states = []
          logdet = 0.
          x = self.squeeze(x)
  
          # Loop through affine layers
          for i, revlayer in enumerate(self.revlayers._modules.values()):
              # If last flow block in forward its an LSTM layer
              if(i==self.n_layers-1):
                  if rec_states is None:
                      x, dlogdet, state_out = revlayer.forward(x, cond)
                  else:
                      x, dlogdet, state_out = revlayer.forward(x, cond, rec_states)
                  out_states = state_out
              else:
                  x, dlogdet = revlayer.forward(x, cond)
  
              logdet = logdet + dlogdet
          
          if self.do_split:
              x, log_prob_prior, eps = self.split(x, return_eps=return_eps)
              logdet = logdet + log_prob_prior
              return x, logdet, out_states, eps
          else:
              return x, logdet, out_states, None
  
      def reverse(self, y, cond, rec_states, eps=None):
          """Backward pass
  
          :param y:  [B, in_features, H, W] input feature tensor
          :type y: torch.Tensor
          :param cond: [B, cond_features, H, W] conditional feature tensor
          :type cond: torch.Tensor
          :param rec_states: tuple of LSTM states (hidden state, cell state), defaults to None
          :type rec_states: tuple, optional
          :param eps: [B, 2*in_features, H//2, W//2] (for factor=2) tensor of sampled latent variables from unit gaussian, defaults to None
          :type eps: torch.Tensor, optional
          :returns: 
              - x: [B, in_features, H, W] Output feature tensor
              - logdet: log determinate block
              - out_states: tuple of LSTM (cell, hidden) states
          :rtype: (torch.Tensor, torch.Tensor, tuple)
          """
          # eps is not used if this revblock does not split in the end
          # focus on sampling
          out_states = []
          logdet = 0.
          if self.do_split:
              y, log_prob_prior = self.split.reverse(y, eps)
              logdet = logdet + log_prob_prior
  
          # Loop through affine layers
          for i, revlayer in reversed(list(enumerate(self.revlayers._modules.values()))):
              # If first flow block in reverse it is an LSTM cell
              if(i==self.n_layers-1):
                  if rec_states is None:
                      y, dlogdet, state_out = revlayer.reverse(y, cond)
                  else:
                      y, dlogdet, state_out = revlayer.reverse(y, cond, rec_states)
                  out_states = state_out
              else:
                  y, dlogdet = revlayer.reverse(y, cond)
              logdet = logdet + dlogdet
              
          return self.squeeze.reverse(y), logdet, out_states
  
  if __name__ == '__main__':
      import time
  
      test_flow = LSTMFLowBlock(4, 2, 4, 1, factor=2, do_split=True, squeeze_type=0)
      # couple = AffineCouplingBlock(4, 4, conv_ksize=3)
  
      x_in =  torch.randn(1, 4, 8, 8)
      h_in =  torch.randn(1, 2, 4, 4)
  
      tic = time.time()
      y_out, log_det, out_states, eps = test_flow.forward(x_in, h_in, None, return_eps=True)
      print(time.time()-tic)
      print(eps.size())
      tic = time.time()
      x_rev, log_det, out_states = test_flow.reverse(y_out, h_in, None, eps=eps)
      print(time.time()-tic)
  
      # x_in =  torch.randn(1, 4, 4, 4)
      # y_out, log_det = couple.forward(x_in, h_in)
      # x_rev, log_det = couple.reverse(y_out, h_in)
      print(torch.sum(torch.abs(x_in[0,0]-x_rev[0,0])))
      print(x_rev[0,0])
      print(x_in[0,0])
  