"""Small CNN for 32-way insect acoustic species classification."""
import torch
from torch import nn

class InsectCNN(nn.Module):
    def __init__(self, num_classes=32):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(),
            nn.AdaptiveAvgPool2d((4, 4)),
        )
        self.classifier = nn.Sequential(nn.Flatten(), nn.Dropout(0.25),
                                        nn.Linear(64*4*4, 128), nn.ReLU(),
                                        nn.Dropout(0.25), nn.Linear(128, num_classes))
    def forward(self, x):
        return self.classifier(self.features(x))