"""TF-IDF + MLP baseline classifier."""

import numpy as np
import torch
from scipy.sparse import spmatrix
from torch import nn
from torch.utils.data import Dataset


class TfidfDataset(Dataset):
    """Wraps a sparse TF-IDF matrix + integer labels as dense float32 tensors."""

    def __init__(self, features: spmatrix, labels: np.ndarray):
        self.features = features
        self.labels = torch.as_tensor(labels, dtype=torch.long)

    def __len__(self) -> int:
        return self.features.shape[0]

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        row = self.features[idx].toarray().astype(np.float32).squeeze(0)
        return torch.from_numpy(row), self.labels[idx]


class SimpleMLP(nn.Module):
    """Input(TF-IDF) -> hidden -> hidden -> num_classes."""

    def __init__(self, input_size: int, hidden_size: int, num_classes: int, dropout: float = 0.3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)