"""Transformers-compatible wrapper for harpertoken/tone.

The model is a plain two-conv CNN over log-mel spectrograms, not a
transformer. This subclasses PreTrainedModel so the weights serialize honestly
as config.json plus model.safetensors and load through AutoModel. Importing
this module registers the architecture.
"""

import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel


class ToneConfig(PretrainedConfig):
    model_type = "tone_cnn"

    def __init__(self, num_labels=10, dropout=0.25, **kwargs):
        super().__init__(**kwargs)
        self.num_labels = num_labels
        self.dropout = dropout


LABELS = ["yes", "no", "up", "down", "left", "right", "on", "off", "stop", "go"]


class ToneCNN(PreTrainedModel):
    config_class = ToneConfig
    base_model_prefix = "cnn"

    # Transformers 5 reads this mapping while finalising a load. Nothing here
    # is tied, but the attribute has to exist or from_pretrained raises
    # AttributeError before the weights are placed.
    all_tied_weights_keys = {}

    def __init__(self, config):
        super().__init__(config)
        self.cnn = nn.ModuleDict(
            {
                "conv1": nn.Conv2d(1, 32, 3, padding=1),
                "conv2": nn.Conv2d(32, 64, 3, padding=1),
                "fc1": nn.Linear(64 * 8 * 8, 128),
                "fc2": nn.Linear(128, config.num_labels),
            }
        )
        self.pool = nn.MaxPool2d(2, 2)
        self.adapt = nn.AdaptiveAvgPool2d((8, 8))
        self.drop = nn.Dropout(config.dropout)

    def forward(self, mel):
        x = self.pool(torch.relu(self.cnn["conv1"](mel)))
        x = self.pool(torch.relu(self.cnn["conv2"](x)))
        x = self.adapt(x)
        x = x.view(x.size(0), -1)
        return self.cnn["fc2"](self.drop(torch.relu(self.cnn["fc1"](x))))

    def predict_label(self, waveform):
        """waveform: 16,000 float samples at 16 kHz. Returns the command."""
        import torchaudio

        self.eval()
        x = torch.as_tensor(waveform, dtype=torch.float32)
        mel = torchaudio.transforms.MelSpectrogram(
            sample_rate=16000, n_fft=512, win_length=400, hop_length=160,
            n_mels=64,
        )(x.unsqueeze(0))
        logmel = torch.log10(mel + 1e-10).unsqueeze(1)
        with torch.no_grad():
            return LABELS[int(self.forward(logmel).argmax(1)[0])]


def register():
    AutoConfig.register(ToneConfig.model_type, ToneConfig, exist_ok=True)
    AutoModel.register(ToneConfig, ToneCNN, exist_ok=True)


register()