import torch
from torch import nn


class LSTM(nn.Module):
    def __init__(self,
                 input_dim: int,
                 output_dim: int,
                 lstm_hidden_size=64,
                 num_layers=1,
                 **kwargs):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_dim,
            hidden_size=lstm_hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=0.5,
            **kwargs
        )
        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, _ = 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