import torch
import torch.nn as nn
from .PositionalEncoding import PositionalEncoding

class AntiCheatPT_256(nn.Module):
    def __init__(
        self,
        feature_dim=44,         # nr of features per tick
        seq_len=256,           # nr of ticks
        nhead=1,                # nr of attention heads
        num_layers=4,           # nr of transformer encoder layers
        dim_feedforward=176,    # hidden size of feedforward network (MLP)
        dropout=0.1             # dropout rate
    ):
        super(AntiCheatPT_256, self).__init__()

        self.positional_encoding = PositionalEncoding(d_model=feature_dim, max_len=seq_len + 1)  # +1 for CLS token

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=feature_dim,
            nhead=nhead,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True
        )   # input shape is (batch, seq_len, d_model)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)

        self.cls_token = nn.Parameter(torch.zeros(1, 1, feature_dim))  # add classification token

        self.fc_out = nn.Sequential(
            nn.Linear(feature_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 1)
        )

    def forward(self, x):
        # x shape: (batch_size, seq_len, feature_dim) = (batch_size, 256, d_model)
        B = x.size(0)

        # add classification token
        cls_tokens = self.cls_token.expand(B, -1, -1).to(x.device)  # (batch_size, 1, d_model)
        x = torch.cat((cls_tokens, x), dim=1)  # -> (batch_size, 257, d_model)

        x = self.positional_encoding(x) # add positional encoding

        x = self.transformer_encoder(x)  # -> (batch_size, 257, d_model)

        cls_output = x[:, 0]  # get output for classification token
        out = self.fc_out(cls_output)  # -> (batch_size, 1)

        return out.squeeze(1)  # -> (batch_size,)