import torch
import torch.nn as nn


class ProgrammingClassifier(nn.Module):
    """
    Classifies programming-related intent or domain.
    """

    def __init__(self, input_dim: int = 128, num_classes: int = 5):
        super().__init__()

        # store attributes for engine access
        self.input_dim = input_dim
        self.num_classes = num_classes

        # classifier layer
        self.classifier = nn.Linear(input_dim, num_classes)

        # label names
        self.labels = [
            "general",
            "debugging",
            "architecture",
            "algorithm",
            "deployment"
        ]

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

    def predict_label(self, logits: torch.Tensor) -> str:
        """
        Convert logits to programming domain label
        """
        idx = torch.argmax(logits, dim=-1).item()
        return self.labels[idx]