import torch
import torch.nn as nn


class LSTMClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, dropout, output_size=1):
        super(LSTMClassifier, 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, :])
        out = torch.sigmoid_(out)
        return out


class LSTMClassifierB(nn.Module):
    def __init__(self, input_size, hidden_size, dropout, num_layers):
        super(LSTMClassifierB, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, dropout=dropout, batch_first=True)

        fc_layers = []
        input_dim = hidden_size

        fc_layers.append(nn.Linear(input_dim, input_dim))
        fc_layers.append(nn.ReLU())
        fc_layers.append(nn.Linear(input_dim, 1))  # Output layer for binary classification
        self.fc = nn.Sequential(*fc_layers)

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

        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out[:, -1, :])  # Only take the output from the last time step
        out = torch.sigmoid(out)  # Apply sigmoid activation
        return out


class LSTMClassifierC(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, sequence_length):
        super(LSTMClassifierC, self).__init__()
        self.hidden_size = hidden_size         # Number of features in the hidden state
        self.num_layers = num_layers           # Number of recurrent layers in the LSTM
        self.sequence_length = sequence_length # Length of the input sequences

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

        # Single linear layer after flattening the LSTM output
        self.fc = nn.Linear(hidden_size * sequence_length, 1)

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

        out, _ = self.lstm(x, (h0, c0))
        out = out.reshape(out.size(0), -1)  # Flatten the sequence
        out = self.fc(out)  # Pass through the single linear layer
        out = torch.sigmoid(out)  # Apply sigmoid activation
        return out