from __future__ import annotations

import torch
from schema import ACTIONS, VOCABULARY_SIZE
from torch import nn


class KernelMindPolicy(nn.Module):
    def __init__(self, dimensions: int = 32) -> None:
        super().__init__()
        self.token_embedding = nn.Embedding(VOCABULARY_SIZE, dimensions)
        self.class_token = nn.Parameter(torch.zeros(1, 1, dimensions))
        self.position_embedding = nn.Parameter(torch.randn(1, 8, dimensions) * 0.02)
        layer = nn.TransformerEncoderLayer(
            d_model=dimensions,
            nhead=4,
            dim_feedforward=64,
            dropout=0.0,
            batch_first=True,
            activation="gelu",
            norm_first=True,
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=2)
        self.action_heads = nn.Linear(dimensions, 3 * len(ACTIONS))

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        embedded = self.token_embedding(tokens)
        class_token = self.class_token.expand(len(tokens), -1, -1)
        sequence = torch.cat([class_token, embedded], dim=1)
        hidden = self.encoder(sequence + self.position_embedding)
        return self.action_heads(hidden[:, 0]).reshape(len(tokens), 3, len(ACTIONS))


def parameter_count(model: nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters())