# Class to define the network architecture of the models
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.optim import Adam


class VanillaLSTM(nn.Module):
    def __init__(
        self, input_dim=1, hidden_dim=64, output_dim=1, num_layers=2, dropout=0.2
    ):
        super(VanillaLSTM, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers

        self.lstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout,
        )
        self.fc = nn.Linear(in_features=hidden_dim, out_features=output_dim)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_()
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).requires_grad_()

        out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach()))
        out = self.fc(out[:, -1, :])

        return out

class VAE(nn.Module):

    def __init__(self, seq_len=48, n_features=1, hidden_dim=64, latent_dim=16, dropout=0.3):
        super(VAE, self).__init__()
        self.seq_len = seq_len
        self.hidden_dim = hidden_dim

        # Encoder
        self.enc_lstm = nn.LSTM(
            input_size=n_features,
            hidden_size=hidden_dim,
            batch_first=True
        )
        self.enc_dropout = nn.Dropout(p=dropout)
        self.fc_mu = nn.Linear(hidden_dim, latent_dim)
        self.fc_var = nn.Linear(hidden_dim, latent_dim)

        # Decoder
        self.fc_upsample = nn.Linear(latent_dim, seq_len * hidden_dim)
        self.dec_dropout = nn.Dropout(p=dropout)
        self.dec_lstm = nn.LSTM(
            input_size=hidden_dim,
            hidden_size=hidden_dim,
            batch_first=True
        )
        self.fc_out = nn.Linear(hidden_dim, n_features)

    def reparameterize(self, mu, log_var):
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        # Encode
        _, (h_enc, c_enc) = self.enc_lstm(x)
        h_enc = h_enc.squeeze(0)  # shape: (batch_size, hidden_dim)
        h_enc = self.enc_dropout(h_enc)
        mu, log_var = self.fc_mu(h_enc), self.fc_var(h_enc)

        # Reparameterize at latent space
        z = self.reparameterize(mu, log_var)

        # Decode
        z = self.fc_upsample(z)
        z = z.view(-1, self.seq_len, self.hidden_dim)
        decoded, _ = self.dec_lstm(z)
        dec_out = self.dec_dropout(decoded)
        out = self.fc_out(dec_out)

        return out, mu, log_var

class Transformer(nn.Module):

    def __init__(self, input_dim=1, model_dim=64, num_layers=2, num_heads=4, dropout=0.2):
        super(Transformer, self).__init__()
        self.model_dim = model_dim
        self.num_layers = num_layers

        self.embedding = nn.Linear(input_dim, model_dim)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=model_dim,
            nhead=num_heads,
            dropout=dropout,
            dim_feedforward=2*model_dim, # 128
            batch_first=True
        )
        encoder_norm = nn.LayerNorm(model_dim)

        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layer,
            num_layers=num_layers,
            norm=encoder_norm
        )

        decoder_layer = nn.TransformerDecoderLayer(
            d_model=model_dim,
            nhead=num_heads,
            dropout=dropout,
            dim_feedforward=2*model_dim, # 128
            batch_first=True
        )
        decoder_norm = nn.LayerNorm(model_dim)

        self.transformer_decoder = nn.TransformerDecoder(
            decoder_layer,
            num_layers=num_layers,
            norm=decoder_norm
        )
        self.output = nn.Linear(model_dim, input_dim)

    def forward(self, x):
        embed_x = self.embedding(x)
        enc_out = self.transformer_encoder(embed_x)
        dec_out = self.transformer_decoder(embed_x, enc_out)
        out = self.output(dec_out)
        return out