from os import name
import torch.nn as nn
import torch.nn.functional as F
import torch
from torch.autograd import Variable
from torch.nn.modules.utils import _pair
 
class ConvRNNCellBase(nn.Module):
    # format string representation
    def __repr__(self):
        s = (
            '{name}({input_channels}), {hidden_channels}, kernel_size={kernel_size}'
            ', stride={stride}'
        )
 
        if self.padding != (0, ) * len(self.padding):
            s += ', padding={padding}'
        if self.dilation != (1, ) * len(self.dilation):
            s += ', dilation={dilation}'
 
        s += ', hidden_kernel_size={hidden_kernel_size}'
        s += ')'
 
        return s.format(name=self.__class__.__name__, **self.__dict__)
 
class ConvLSTMCell(ConvRNNCellBase):
    # based on section 2.1
    def __init__(self,
                 input_channels,
                 hidden_channels,
                 kernel_size=3,
                 stride=1,
                 padding=0,
                 dilation=1,
                 hidden_kernel_size=1,
                 bias=True):
        super(ConvLSTMCell, self).__init__()
        self.input_channels = input_channels
        self.hidden_channels = hidden_channels
 
        self.kernel_size = _pair(kernel_size)
        self.stride = _pair(stride)
        self.padding = _pair(padding)
        self.dilation = _pair(dilation)
 
        self.hidden_kernel_size = _pair(hidden_kernel_size)
 
        hidden_padding = _pair(hidden_kernel_size // 2)
 
        gate_channels = 4 * self.hidden_channels
        self.conv_ih = nn.Conv2d(
            in_channels=self.input_channels,
            out_channels=gate_channels,
            kernel_size=self.kernel_size,
            stride=self.stride,
            padding=self.padding,
            dilation=self.dilation,
            bias=bias)
 
        self.conv_hh = nn.Conv2d(
            in_channels=self.hidden_channels,
            out_channels=gate_channels,
            kernel_size=hidden_kernel_size,
            stride=1,
            padding=hidden_padding,
            dilation=1,
            bias=bias)
 
        self.reset_parameters()
 
    def reset_parameters(self):
        self.conv_ih.reset_parameters()
        self.conv_hh.reset_parameters()
 
    def forward(self, input, hidden):
        # based on equation (4)-(6)
        hx, cx = hidden
        gates = self.conv_ih(input) + self.conv_hh(hx)
 
        ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)
 
        ingate = F.sigmoid(ingate)
        forgetgate = F.sigmoid(forgetgate)
        cellgate = F.tanh(cellgate)
        outgate = F.sigmoid(outgate)
 
        cy = (forgetgate * cx) + (ingate * cellgate)
        hy = outgate * F.tanh(cy)
 
        return hy, cy