import torch
from torch import nn

class LSTMRegressor(nn.Module):
    def __init__(self, input_size = 24, hidden_size = 64, num_layers = 2, dropout = 0.2):
        super().__init__()

        self.lstm = nn.LSTM(input_size = input_size,
                            hidden_size = hidden_size,
                            num_layers = num_layers,
                            dropout= dropout,
                            batch_first = True)

        self.fc = nn.Linear(in_features = hidden_size,
                            out_features = 1)

    def forward(self, x):
        out, _ = self.lstm(x)
        last_time_step = out[:, -1, :]
        out= self.fc(last_time_step)
        return out