import torch, torch.nn as nn, math

class SinusoidalPE(nn.Module):
    def __init__(self, d_model, max_len=512):
        super().__init__()
        pe = torch.zeros(max_len, d_model); pos = torch.arange(0, max_len).unsqueeze(1).float()
        div = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(pos * div); pe[:, 1::2] = torch.cos(pos * div)
        self.register_buffer("pe", pe.unsqueeze(0))
    def forward(self, x): return x + self.pe[:, :x.size(1)]

class TransSLR(nn.Module):
    def __init__(self, feat_dim=225, d_model=128, nhead=4, num_layers=2,
                 ffn_dim=512, dropout=0.6, num_classes=30):
        super().__init__()
        self.proj = nn.Linear(feat_dim, d_model); self.pos_enc = SinusoidalPE(d_model)
        enc_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead,
                        dim_feedforward=ffn_dim, dropout=dropout, batch_first=True)
        self.encoder = nn.TransformerEncoder(enc_layer, num_layers=num_layers)
        self.gap = nn.AdaptiveAvgPool1d(1); self.drop = nn.Dropout(dropout)
        self.head = nn.Linear(d_model, num_classes)
    def forward(self, x):
        x = self.proj(x); x = self.pos_enc(x); x = self.encoder(x)
        x = self.gap(x.transpose(1,2)).squeeze(-1)
        return self.head(self.drop(x))
    def get_embeddings(self, x):
        x = self.proj(x); x = self.pos_enc(x); x = self.encoder(x)
        return self.gap(x.transpose(1,2)).squeeze(-1)