import torch
import torch.nn as nn

class TimeSeriesModel(nn.Module):
    def __init__(self, input_size, output_size, num_layers=5, hidden_size=128, dropout_rate=0.3):
        super(TimeSeriesModel, self).__init__()

        self.output_size = output_size
        
        layers = []
        for i in range(num_layers):
            if i == 0:
                layers.append(nn.Linear(input_size, hidden_size))
            else:
                layers.append(nn.Linear(hidden_size, hidden_size))
            layers.append(nn.BatchNorm1d(hidden_size))
            layers.append(nn.Dropout(dropout_rate))
            layers.append(nn.ReLU())

        self.layers = nn.ModuleList(layers)
        self.output = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        x = torch.squeeze(x, -1)

        for layer in self.layers:
            x = layer(x)
        x = self.output(x)
        return x