"""Маленькая свёрточная сеть (CNN) для картинок 28x28.

Свёртка (Conv) ищет простые узоры: кусочки линий, углы, дуги.
Следующие свёртки собирают из них узоры сложнее: зигзаг, лучи, круг.
MaxPool уменьшает картинку вдвое и оставляет самое сильное.
В конце полносвязные слои (Linear) решают, какой это класс.
"""
import torch
from torch import nn


def conv_block(c_in, c_out, dropout):
    """Две свёртки 3x3 + BatchNorm + ReLU, потом уменьшение в 2 раза."""
    return nn.Sequential(
        nn.Conv2d(c_in, c_out, 3, padding=1), nn.BatchNorm2d(c_out), nn.ReLU(),
        nn.Conv2d(c_out, c_out, 3, padding=1), nn.BatchNorm2d(c_out), nn.ReLU(),
        nn.MaxPool2d(2), nn.Dropout(dropout),
    )


class SpellCNN(nn.Module):
    def __init__(self, n_classes=6):
        super().__init__()
        self.features = nn.Sequential(
            conv_block(1, 32, 0.1),    # 1x28x28 -> 32x14x14
            conv_block(32, 64, 0.2),   # 32x14x14 -> 64x7x7
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 7 * 7, 128), nn.ReLU(), nn.Dropout(0.3),
            nn.Linear(128, n_classes),
        )

    def forward(self, x):
        return self.classifier(self.features(x))


def count_params(model):
    return sum(p.numel() for p in model.parameters())


def save(model, path, classes, meta=None):
    torch.save({"state_dict": model.state_dict(), "classes": classes, "meta": meta or {}}, path)


def load(path, device="cpu"):
    ckpt = torch.load(path, map_location=device, weights_only=False)
    model = SpellCNN(len(ckpt["classes"]))
    model.load_state_dict(ckpt["state_dict"])
    model.eval()
    return model.to(device), ckpt["classes"]