# -*- coding: utf-8 -*-
import torch
import torch.nn as nn

# Transformer model section
# Transformer models with increased complexity
class Transformer(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(Transformer, self).__init__()

        # Add multiple hidden layers and use the ReLU activation function.
        self.linear1 = nn.Linear(input_dim, 128)
        self.relu1 = nn.ReLU()

        self.linear2 = nn.Linear(128, 64)
        self.relu2 = nn.ReLU()

        self.linear3 = nn.Linear(64, output_dim)

        # BatchNorm layers help accelerate training and stabilize.
        self.batch_norm1 = nn.BatchNorm1d(128)
        self.batch_norm2 = nn.BatchNorm1d(64)

        # Dropout layer, preventing overfitting
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        # The first layer of linear transformation
        x = self.linear1(x)
        x = self.batch_norm1(x)
        x = self.relu1(x)

        # The second layer of linear transformation
        x = self.linear2(x)
        x = self.batch_norm2(x)
        x = self.relu2(x)

        # Dropout layer
        x = self.dropout(x)

        # The final layer output
        x = self.linear3(x)

        return x