r"""
《Convolutional Refine LSTM》
"""
import warnings

import torch
from torch import nn

__all__ = ["ConvRLSTM"]


class ConvRLSTM(nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int, kernel_size: int):
        r"""
        :param in_channels:       输入的通道数
        :param hidden_channels:   隐藏层通道数
        :param kernel_size:       卷积核尺寸
        """
        warnings.warn("一次改动太多不好写，暂时放弃", DeprecationWarning)
        super(ConvRLSTM, self).__init__()
        self.in_channels = in_channels
        self.hidden_channels = hidden_channels

        padding = kernel_size // 2

        self.conv_x = nn.Conv2d(
            in_channels=in_channels, out_channels=hidden_channels * 4,
            kernel_size=kernel_size, stride=1, padding=padding
        )

        self.conv_m = nn.Conv2d(
            in_channels=hidden_channels, out_channels=hidden_channels * 4,
            kernel_size=kernel_size, stride=1, padding=padding
        )

    def forward(self, x, m, c):
        x_concat = self.conv_x(x)
        m_concat = self.conv_m(m)

        x_concat = torch.layer_norm(x_concat, x_concat.shape[1:])
        m_concat = torch.layer_norm(m_concat, m_concat.shape[1:])

        f_x, u_x, r_x, o_x = torch.split(x_concat, self.hidden_channels, dim=1)
        f_m, u_m, r_m, o_m = torch.split(m_concat, self.hidden_channels, dim=1)

        f = torch.sigmoid(f_x + f_m)
        u = torch.tanh(u_x + u_m)
        r = torch.sigmoid(r_x + r_m)
        g = r * (1 - torch.pow(1 - f, 2)) + (1 - r) * torch.pow(f, 2)
        c = g * c + (1 - g) * u
        o = torch.sigmoid(o_x + o_m)
        m = o * torch.tanh(c)

        return m, c
