import torch
import torch.nn as nn


class BistableRecurrentCell(nn.Module):
    """Bistable Recurrent Cell Layer (Vecoven et al. 2020)"""

    def __init__(self, input_size, output_dim, *args, **kwargs):
        """Constructor.

        Args:
            input_size (int): Number of input features.
            output_dim (int): Number of output features generated by the cell.
        """

        super(BistableRecurrentCell, self).__init__(*args, **kwargs)

        self.input_size = input_size
        self.output_dim = output_dim
        self.state_size = output_dim

        self.kernelz = nn.Parameter(
            torch.FloatTensor(self.input_size, self.output_dim)
        )
        self.kernelr = nn.Parameter(
            torch.FloatTensor(self.input_size, self.output_dim)
        )
        self.kernelh = nn.Parameter(
            torch.FloatTensor(self.input_size, self.output_dim)
        )

        self.memoryz = nn.Parameter(torch.FloatTensor(self.output_dim, ))
        self.memoryr = nn.Parameter(torch.FloatTensor(self.output_dim, ))

        self.bz = nn.Parameter(torch.FloatTensor(self.output_dim))
        self.br = nn.Parameter(torch.FloatTensor(self.output_dim))

        # Initialise the weights
        nn.init.xavier_uniform_(self.kernelz)
        nn.init.xavier_uniform_(self.kernelr)
        nn.init.xavier_uniform_(self.kernelh)

        nn.init.ones_(self.memoryz)
        nn.init.ones_(self.memoryr)

        nn.init.zeros_(self.bz)
        nn.init.zeros_(self.br)

    def forward(
        self, input_t: torch.Tensor, hidden_state: torch.Tensor = None
    ) -> torch.Tensor:
        """Bistable Recurrent Cell.

        Args:
            input_t (torch.Tensor): Input tensor at time step t.
            hidden_state (torch.Tensor): Hidden state at time time t-1.

        Returns:
            torch.Tensor: Tensor of the new hidden state at time t.
        """

        if hidden_state is None:
            hidden_state = self.hidden_state

        r = torch.tanh(
            torch.matmul(input_t, self.kernelr) + hidden_state * self.memoryr +
            self.br
        ) + 1
        z = torch.sigmoid(
            torch.matmul(input_t, self.kernelz) + hidden_state * self.memoryz +
            self.bz
        )
        output_hidden_state = z * hidden_state + (
            1.0 - z
        ) * torch.tanh(torch.matmul(input_t, self.kernelh) + r * hidden_state)

        return output_hidden_state

    def get_initial_state(self, batch_size=None, dtype=torch.float32):
        """Intialise hidden states.

        Args:
            batch_size (int, optional): Batch size of the input.
                Defaults to None.
            dtype (object, optional): Desired data type.
                Defaults to torch.float32.

        Returns:
            torch.Tensor: Hidden state intialised to 0.
        """
        return torch.zeros((batch_size, self.output_dim), dtype=dtype)
