import torch
import torch.nn as nn


class LSTMRegressor(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, dropout, output_size=1):
        super(LSTMRegressor, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, dropout=dropout, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h_0 = torch.zeros(self.lstm.num_layers, x.size(0), self.lstm.hidden_size).to(x.device)
        c_0 = torch.zeros(self.lstm.num_layers, x.size(0), self.lstm.hidden_size).to(x.device)

        out, _ = self.lstm(x, (h_0, c_0))
        out = self.fc(out[:, -1, :])
        return out


class LSTMRegressorB(nn.Module):
    def __init__(self, input_size, hidden_size, dropout, sequence_length):
        super(LSTMRegressorB, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers=2, dropout=dropout, batch_first=True)
        fc_list = [nn.Linear(input_size * sequence_length, input_size * sequence_length),
                   nn.Linear(hidden_size * sequence_length, 1)]
        self.fc = nn.Sequential(*fc_list)

    def forward(self, x):
        h_0 = torch.zeros(self.lstm.num_layers, x.size(0), self.lstm.hidden_size).to(x.device)
        c_0 = torch.zeros(self.lstm.num_layers, x.size(0), self.lstm.hidden_size).to(x.device)

        out, _ = self.lstm(x, (h_0, c_0))
        out = out.reshape(out.size(0), -1)
        out = self.fc(out)
        return out


class GRURegressor(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, dropout, output_size=1):
        super(GRURegressor, self).__init__()
        self.gru = nn.GRU(input_size, hidden_size, num_layers, dropout=dropout, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        h_0 = torch.zeros(self.gru.num_layers, x.size(0), self.gru.hidden_size).to(x.device)

        out, _ = self.gru(x, h_0)
        out = self.fc(out[:, -1, :])
        return out