import torch
from torch import nn


# NEW MODEL WITH SMALL ALTERATIONS
# The Model is partially adapted from the Timeseries Clustering repository developed by Tejas Lodaya:
# https://github.com/tejaslodaya/timeseries-clustering-vae/blob/master/vrae/vrae.py
class Encoder(nn.Module):
    """Encoder module of the Variational Autoencoder."""

    def __init__(
        self,
        NUM_FEATURES: int,
        hidden_size_layer_1: int,
        hidden_size_layer_2: int,
        dropout_encoder: float,
    ):
        """
        Initialize the Encoder module.

        Parameters
        ----------
        NUM_FEATURES : int
            Number of input features.
        hidden_size_layer_1 : int
            Size of the first hidden layer.
        hidden_size_layer_2 : int
            Size of the second hidden layer.
        dropout_encoder : float
            Dropout rate for regularization.
        """
        super(Encoder, self).__init__()

        self.input_size = NUM_FEATURES
        self.hidden_size = hidden_size_layer_1
        self.hidden_size_2 = hidden_size_layer_2
        self.n_layers = 2
        self.dropout = dropout_encoder
        self.bidirectional = True

        self.encoder_rnn = nn.GRU(
            input_size=self.input_size,
            hidden_size=self.hidden_size,
            num_layers=self.n_layers,
            bias=True,
            batch_first=True,
            dropout=self.dropout,
            bidirectional=self.bidirectional,
        )  # UNRELEASED!

        self.hidden_factor = (2 if self.bidirectional else 1) * self.n_layers

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        """
        Forward pass of the Encoder module.

        Parameters
        ----------
        inputs : torch.Tensor
            Input tensor of shape (batch_size, sequence_length, num_features).

        Returns
        -------
        torch.Tensor:
            Encoded representation tensor of shape (batch_size, hidden_size_layer_1 * 4).
        """
        outputs_1, hidden_1 = self.encoder_rnn(inputs)  # UNRELEASED!
        hidden = torch.cat(
            (
                hidden_1[0, ...],
                hidden_1[1, ...],
                hidden_1[2, ...],
                hidden_1[3, ...],
            ),
            1,
        )
        return hidden


class Lambda(nn.Module):
    """Lambda module for computing the latent space parameters."""

    def __init__(
        self,
        ZDIMS: int,
        hidden_size_layer_1: int,
        softplus: bool,
    ):
        """
        Initialize the Lambda module.

        Parameters
        ----------
        ZDIMS : int
            Size of the latent space.
        hidden_size_layer_1 : int
            Size of the first hidden layer.
        hidden_size_layer_2 : int, deprecated
            Size of the second hidden layer.
        softplus : bool
            Whether to use softplus activation for logvar.
        """
        super(Lambda, self).__init__()

        self.hid_dim = hidden_size_layer_1 * 4
        self.latent_length = ZDIMS
        self.softplus = softplus

        self.hidden_to_mean = nn.Linear(self.hid_dim, self.latent_length)
        self.hidden_to_logvar = nn.Linear(self.hid_dim, self.latent_length)

        if self.softplus is True:
            print(
                "Using a softplus activation to ensures that the variance is parameterized as non-negative and activated by a smooth function"
            )
            self.softplus_fn = nn.Softplus()

    def forward(
        self,
        hidden: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Forward pass of the Lambda module.

        Parameters
        ----------
        hidden : torch.Tensor
            Hidden representation tensor of shape (batch_size, hidden_size_layer_1 * 4).

        Returns
        -------
        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
            Latent space tensor, mean tensor, logvar tensor.
        """
        self.mean = self.hidden_to_mean(hidden)
        if self.softplus is True:
            self.logvar = self.softplus_fn(self.hidden_to_logvar(hidden))
        else:
            self.logvar = self.hidden_to_logvar(hidden)

        if self.training:
            std = torch.exp(0.5 * self.logvar)
            eps = torch.randn_like(std)
            return eps.mul(std).add_(self.mean), self.mean, self.logvar
        else:
            return self.mean, self.mean, self.logvar


class Decoder(nn.Module):
    """Decoder module of the Variational Autoencoder."""

    def __init__(
        self,
        TEMPORAL_WINDOW: int,
        ZDIMS: int,
        NUM_FEATURES: int,
        hidden_size_rec: int,
        dropout_rec: float,
    ):
        """
        Initialize the Decoder module.

        Parameters
        ----------
        TEMPORAL_WINDOW : int
            Size of the temporal window.
        ZDIMS : int
            Size of the latent space.
        NUM_FEATURES : int
            Number of input features.
        hidden_size_rec : int
            Size of the recurrent hidden layer.
        dropout_rec : float
            Dropout rate for regularization.
        """
        super(Decoder, self).__init__()

        self.num_features = NUM_FEATURES
        self.sequence_length = TEMPORAL_WINDOW
        self.hidden_size = hidden_size_rec
        self.latent_length = ZDIMS
        self.n_layers = 1
        self.dropout = dropout_rec
        self.bidirectional = True

        self.rnn_rec = nn.GRU(
            self.latent_length,
            hidden_size=self.hidden_size,
            num_layers=self.n_layers,
            bias=True,
            batch_first=True,
            dropout=self.dropout,
            bidirectional=self.bidirectional,
        )

        self.hidden_factor = (2 if self.bidirectional else 1) * self.n_layers  # NEW

        self.latent_to_hidden = nn.Linear(
            self.latent_length,
            self.hidden_size * self.hidden_factor,
        )  # NEW
        self.hidden_to_output = nn.Linear(
            self.hidden_size * (2 if self.bidirectional else 1),
            self.num_features,
        )

    def forward(
        self,
        inputs: torch.Tensor,
        z: torch.Tensor,
    ) -> torch.Tensor:
        """
        Forward pass of the Decoder module.

        Parameters
        ----------
        inputs : torch.Tensor
            Input tensor of shape (batch_size, seq_len, ZDIMS).
        z : torch.Tensor
            Latent space tensor of shape (batch_size, ZDIMS).

        Returns
        -------
        torch.Tensor:
            Decoded output tensor of shape (batch_size, seq_len, NUM_FEATURES).
        """
        batch_size = inputs.size(0)  # NEW
        hidden = self.latent_to_hidden(z)  # NEW
        hidden = hidden.view(self.hidden_factor, batch_size, self.hidden_size)  # NEW
        decoder_output, _ = self.rnn_rec(inputs, hidden)
        prediction = self.hidden_to_output(decoder_output)
        return prediction


class Decoder_Future(nn.Module):
    """Decoder module for predicting future sequences."""

    def __init__(
        self,
        TEMPORAL_WINDOW: int,
        ZDIMS: int,
        NUM_FEATURES: int,
        FUTURE_STEPS: int,
        hidden_size_pred: int,
        dropout_pred: float,
    ):
        """
        Initialize the Decoder_Future module.

        Parameters
        ----------
        TEMPORAL_WINDOW : int
            Size of the temporal window.
        ZDIMS : int
            Size of the latent space.
        NUM_FEATURES : int
            Number of input features.
        FUTURE_STEPS : int
            Number of future steps to predict.
        hidden_size_pred : int
            Size of the prediction hidden layer.
        dropout_pred : float
            Dropout rate for regularization.
        """
        super(Decoder_Future, self).__init__()

        self.num_features = NUM_FEATURES
        self.future_steps = FUTURE_STEPS
        self.sequence_length = TEMPORAL_WINDOW
        self.hidden_size = hidden_size_pred
        self.latent_length = ZDIMS
        self.n_layers = 1
        self.dropout = dropout_pred
        self.bidirectional = True

        self.rnn_pred = nn.GRU(
            self.latent_length,
            hidden_size=self.hidden_size,
            num_layers=self.n_layers,
            bias=True,
            batch_first=True,
            dropout=self.dropout,
            bidirectional=self.bidirectional,
        )

        self.hidden_factor = (2 if self.bidirectional else 1) * self.n_layers  # NEW

        self.latent_to_hidden = nn.Linear(
            self.latent_length,
            self.hidden_size * self.hidden_factor,
        )
        self.hidden_to_output = nn.Linear(self.hidden_size * 2, self.num_features)

    def forward(
        self,
        inputs: torch.Tensor,
        z: torch.Tensor,
    ) -> torch.Tensor:
        """
        Forward pass of the Decoder_Future module.

        Parameters
        ----------
        inputs : torch.Tensor
            Input tensor of shape (batch_size, seq_len, ZDIMS).
        z : torch.Tensor
            Latent space tensor of shape (batch_size, ZDIMS).

        Returns
        -------
        torch.Tensor:
            Predicted future tensor of shape (batch_size, FUTURE_STEPS, NUM_FEATURES).
        """
        batch_size = inputs.size(0)
        hidden = self.latent_to_hidden(z)
        hidden = hidden.view(self.hidden_factor, batch_size, self.hidden_size)
        inputs = inputs[:, : self.future_steps, :]
        decoder_output, _ = self.rnn_pred(inputs, hidden)
        prediction = self.hidden_to_output(decoder_output)
        return prediction


class RNN_VAE(nn.Module):
    """Variational Autoencoder module."""

    def __init__(
        self,
        TEMPORAL_WINDOW: int,
        ZDIMS: int,
        NUM_FEATURES: int,
        FUTURE_DECODER: bool,
        FUTURE_STEPS: int,
        hidden_size_layer_1: int,
        hidden_size_layer_2: int,
        hidden_size_rec: int,
        hidden_size_pred: int,
        dropout_encoder: float,
        dropout_rec: float,
        dropout_pred: float,
        softplus: bool,
    ):
        """
        Initialize the VAE module.

        Parameters
        ----------
        TEMPORAL_WINDOW : int
            Size of the temporal window.
        ZDIMS : int
            Size of the latent space.
        NUM_FEATURES : int
            Number of input features.
        FUTURE_DECODER : bool
            Whether to include a future decoder.
        FUTURE_STEPS : int
            Number of future steps to predict.
        hidden_size_layer_1 : int
            Size of the first hidden layer.
        hidden_size_layer_2 : int
            Size of the second hidden layer.
        hidden_size_rec : int
            Size of the recurrent hidden layer.
        hidden_size_pred : int
            Size of the prediction hidden layer.
        dropout_encoder : float
            Dropout rate for encoder.
        """
        super(RNN_VAE, self).__init__()

        self.FUTURE_DECODER = FUTURE_DECODER
        self.seq_len = int(TEMPORAL_WINDOW / 2)
        self.encoder = Encoder(
            NUM_FEATURES,
            hidden_size_layer_1,
            hidden_size_layer_2,
            dropout_encoder,
        )
        self.lmbda = Lambda(
            ZDIMS,
            hidden_size_layer_1,
            softplus,
        )
        self.decoder = Decoder(
            self.seq_len,
            ZDIMS,
            NUM_FEATURES,
            hidden_size_rec,
            dropout_rec,
        )
        if FUTURE_DECODER:
            self.decoder_future = Decoder_Future(
                self.seq_len,
                ZDIMS,
                NUM_FEATURES,
                FUTURE_STEPS,
                hidden_size_pred,
                dropout_pred,
            )

    def forward(self, seq: torch.Tensor) -> tuple:
        """
        Forward pass of the VAE.

        Parameters
        ----------
        seq : torch.Tensor
            Input sequence tensor of shape (batch_size, seq_len, NUM_FEATURES).

        Returns
        -------
        Tuple containing:
            - If FUTURE_DECODER is True:
                - prediction (torch.Tensor): Reconstructed input sequence tensor.
                - future (torch.Tensor): Predicted future sequence tensor.
                - z (torch.Tensor): Latent representation tensor.
                - mu (torch.Tensor): Mean of the latent distribution tensor.
                - logvar (torch.Tensor): Log variance of the latent distribution tensor.
            - If FUTURE_DECODER is False:
                - prediction (torch.Tensor): Reconstructed input sequence tensor.
                - z (torch.Tensor): Latent representation tensor.
                - mu (torch.Tensor): Mean of the latent distribution tensor.
                - logvar (torch.Tensor): Log variance of the latent distribution tensor.
        """
        # Encode input sequence
        h_n = self.encoder(seq)

        # Compute the latent state via reparametrization trick
        z, mu, logvar = self.lmbda(h_n)
        ins = z.unsqueeze(2).repeat(1, 1, self.seq_len)
        ins = ins.permute(0, 2, 1)

        # Predict the future of the sequence from the latent state
        prediction = self.decoder(ins, z)

        if self.FUTURE_DECODER:
            future = self.decoder_future(ins, z)
            return prediction, future, z, mu, logvar
        else:
            return prediction, z, mu, logvar
