import torch
from torch import nn


class TicketClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.dropout = nn.Dropout(p=0.5)
        self.linear = nn.Linear(4, 3)

    def forward(self, inputs):
        return self.linear(self.dropout(inputs))


class_names = ["standard", "priority", "urgent"]
model = TicketClassifier()

with torch.no_grad():
    model.linear.weight.copy_(
        torch.tensor(
            [
                [0.4, -0.2, 0.1, 0.3],
                [-0.1, 0.5, 0.2, -0.2],
                [0.2, 0.1, -0.1, 0.6],
            ],
            dtype=torch.float32,
        )
    )
    model.linear.bias.copy_(torch.tensor([0.0, 0.1, -0.1]))

features = torch.tensor(
    [
        [0.9, 0.1, 0.2, 0.3],
        [0.1, 0.8, 0.4, 0.2],
        [0.2, 0.2, 0.1, 0.9],
    ],
    dtype=torch.float32,
)

model.eval()

with torch.inference_mode():
    inference_mode_active = torch.is_inference_mode_enabled()
    logits = model(features)
    probabilities = torch.softmax(logits, dim=1)
    predicted_indexes = probabilities.argmax(dim=1)
    predicted_labels = [class_names[index] for index in predicted_indexes.tolist()]

print(f"model training mode: {model.training}")
print(f"dropout training mode: {model.dropout.training}")
print(f"inference mode active: {inference_mode_active}")
print(f"logits shape: {tuple(logits.shape)}")
print(f"probabilities shape: {tuple(probabilities.shape)}")
print(f"predicted labels: {', '.join(predicted_labels)}")
print(f"logits require grad: {logits.requires_grad}")