from torch import nn


class LSTMModel(nn.Module):
    def __init__(self,
                 input_dim: int,
                 output_dim: int,
                 cnn_channels: int,
                 cnn_kernel_size: int,
                 embed_dim: int,
                 lstm_hidden_size=64,
                 num_layers=1):
        super().__init__()
        self.cnn = nn.Sequential(
            nn.Conv1d(input_dim, cnn_channels, kernel_size=cnn_kernel_size, padding=1),
            nn.BatchNorm1d(cnn_channels),
            nn.ReLU(),
            nn.Conv1d(cnn_channels, embed_dim, kernel_size=cnn_kernel_size, padding=1),
            nn.BatchNorm1d(embed_dim),
            nn.ReLU()
        )
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=lstm_hidden_size,
            num_layers=num_layers,
            batch_first=True,
        )
        self.head = nn.Sequential(
            nn.Linear(lstm_hidden_size, output_dim)  # Output is a scalar Δt
        )

    def forward(self, x):
        # x: (batch_size, seq_len, input_size)
        x = x.permute(0, 2, 1)
        # x: (batch_size, input_size, seq_len)
        x = self.cnn(x)
        # x: (batch_size, embed_dim, seq_len)
        x = x.permute(0, 2, 1)
        # x: (batch_size, seq_len, embed_dim)
        x, _ = self.lstm(x)
        # x: (batch_size, seq_len, lstm_hidden_size)
        x = x[:, -1, :]  # Get the last time step
        # x: (batch_size, lstm_hidden_size)
        x = self.head(x)
        # x: (batch_size, output_size)
        return x