import torch.nn as nn
import torch
from NetBlock import Bottleneck

class ConvLSTMCell(nn.Module):

    def __init__(self, input_dim, hidden_dim, channels, higher_channel, out_dim, BN=False):
        """
        Initialize ConvLSTM cell.
        Parameters
        ----------
        input_dim: int
            Number of channels of input tensor.
        hidden_dim: int
            Number of channels of hidden state.
        channels: (int, int, ...)
            input channels for the encoder-decoder network
        """
        super(ConvLSTMCell, self).__init__()

        self.input_dim = input_dim
        self.hidden_dim = channels[-1]

        self.n_layers = len(channels) - 1
        self.out_conv = nn.Sequential(nn.Conv2d(2 * channels[0], out_dim, kernel_size=(3, 3), padding=1))
        self.in_conv = nn.Sequential(
            nn.Conv2d(in_channels=input_dim, out_channels=channels[0], kernel_size=(5, 5), padding=2),
            nn.LeakyReLU(inplace=True)
        )
        if higher_channel != 0:
            self.Merge_Hidden = nn.Sequential(
                nn.Conv2d(in_channels=channels[-1] + higher_channel, out_channels=channels[-1], kernel_size=(3, 3),
                          padding=1, bias=False),
                nn.LeakyReLU(inplace=True)
            )
        self.scale = nn.Parameter(torch.ones(1, dtype=torch.float, requires_grad=True))
        # reverse the channels for the decoder
        reverse_ch = sorted(channels, reverse=True)

        # setting encoder layer and decoder layer for each network level respectively
        for l in range(self.n_layers):
            encoder = nn.Sequential(
                Bottleneck(in_channels=channels[l], out_channels=channels[l + 1], batchnorm=BN),
                nn.MaxPool2d(kernel_size=2),
                Bottleneck(in_channels=channels[l + 1], out_channels=channels[l + 1], batchnorm=BN)
            )
            setattr(self, 'encoder{}'.format(l), encoder)
            upsample = nn.Sequential(
                Bottleneck(in_channels=reverse_ch[l], out_channels=reverse_ch[l + 1], batchnorm=BN),
                nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
                Bottleneck(in_channels=reverse_ch[l + 1], out_channels=reverse_ch[l + 1], batchnorm=BN),
            )
            setattr(self, 'upsample{}'.format(l), upsample)
            merge_conv = nn.Sequential(
                nn.Conv2d(2 * reverse_ch[l + 1], reverse_ch[l + 1], kernel_size=(3, 3), padding=1),
                nn.ReLU(inplace=True)
            )
            setattr(self, 'merge_conv{}'.format(l), merge_conv)

        self.lstm_conv = nn.Sequential(
            nn.Conv2d(2*self.hidden_dim, 4*self.hidden_dim, kernel_size=(3, 3), padding=1),
            nn.LeakyReLU(inplace=True)
        )
        self.lstm_out_conv = nn.Sequential(
            nn.Conv2d(self.hidden_dim, out_channels=self.hidden_dim, kernel_size=(3,3), padding=1),
            nn.LeakyReLU(inplace=True)
        )

    def forward(self, input_tensor, cur_state, higher_h):
        '''

        :param input_tensor:
            input frames or feature maps, size = (b, c, h, w)
        :param cur_state:
            including hidden state and cell state
        '''

        feature_maps = []  # Used to save the feature map generated by the encoder

        # encoding
        x = self.in_conv(input_tensor)
        feature_maps.append(x)
        for l in range(self.n_layers):
            encoder = getattr(self, 'encoder{}'.format(l))
            x = encoder(x)
            feature_maps.append(x)

        # decoding
        feature_maps.reverse()
        x = feature_maps[0]

        h_cur, c_cur = cur_state # h_cur: current hidden state; c_cur: current cell state
        combined = torch.cat([x, h_cur], dim=1)  # concatenate along channel axis
        combined = self.lstm_conv(combined)
        cc_i, cc_f, cc_o, cc_g = torch.split(combined, self.hidden_dim, dim=1)

        # split into 4 tensor(meet with the program of out_channels=4 * self.hidden_dim in line 24)
        i = torch.sigmoid(cc_i)
        f = torch.sigmoid(cc_f)
        o = torch.sigmoid(cc_o)
        g = torch.tanh(cc_g)

        c_next = f * c_cur + i * g
        h_next = o * torch.tanh(c_next)

        x = self.lstm_out_conv(h_next)
        for l in range(len(feature_maps)-1):
            upsample = getattr(self, 'upsample{}'.format(l))
            merge_conv = getattr(self, 'merge_conv{}'.format(l))
            up = upsample(x)
            merge = torch.cat([up, feature_maps[l+1]], dim=1)
            if l < len(feature_maps)-2:
                x = merge_conv(merge)
            else:
                x = self.out_conv(merge)


        return x, (h_next, c_next), None

    def init_hidden(self, batch_size, image_size):
        height, width = image_size
        return (torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),
                torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device))
