"""A small CNN+BiGRU classifier (CRNN), trained end-to-end from scratch.

The third architecture tried on this task, after the frozen-WavLM head
(model.py) and the from-scratch transformer (transformer_model.py). It exists
to test a specific hypothesis about why the transformer plateaued: measured
on the test split, that model reached 80.1% while its train loss kept falling
to 0.108 -- it had capacity to spare and was spending it memorising, not
generalising. A transformer has almost no built-in assumptions about its
input, so with 40k clips it has to *learn* that a spectrogram has local
time-frequency structure. A CRNN is told that up front:

  - 2D convolutions over (frequency x time) share weights across both axes,
    so a harmonic stack or an onset is recognised the same wherever it lands
    in the clip. The transformer has to learn that invariance from data.
  - A bidirectional GRU then reads the conv output along time, which is a
    cheaper way to model sequence structure than all-pairs attention when
    the useful context is mostly local.

Whether that trade is worth it here is an empirical question -- hence the
identical training loop, seed and data as the transformer, so the only thing
that differs is the architecture.
"""

import torch
from torch import nn

# Imported from the training project when it is present, with literal
# fallbacks so this file also works on its own inside the inference package,
# where the rest of the project is not shipped. The fallbacks are asserted
# against the real values by tests/test_crnn_model.py.
try:
    import smad_config as config

    NUM_CLASSES = config.NUM_CLASSES
except ImportError:  # standalone inference
    NUM_CLASSES = 4

try:
    from mel_features import N_MELS
except ImportError:  # standalone inference
    N_MELS = 80


class ConvBlock(nn.Module):
    """Conv -> BatchNorm -> ReLU -> MaxPool, the standard CRNN unit.

    BatchNorm rather than the transformer's LayerNorm: convolutions share
    weights across positions, so normalising per-channel over the batch is
    the matching choice, and it also removes any dependence on the input's
    absolute dB offset.
    """

    def __init__(self, in_ch, out_ch, pool):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(pool),
        )

    def forward(self, x):
        return self.block(x)


class TinyAudioCRNN(nn.Module):
    """(B, T, N_MELS) log-mel spectrogram -> 4 class logits."""

    def __init__(self, n_mels=N_MELS, channels=(32, 64, 128, 128), rnn_hidden=128,
                 dropout=0.2, num_classes=NUM_CLASSES, rnn_type="gru"):
        super().__init__()
        # Same standardisation contract as TinyAudioTransformer: statistics
        # live in state_dict so predict.py cannot feed a different scale than
        # the model trained on. Defaults are the identity.
        self.register_buffer("feat_mean", torch.zeros(n_mels))
        self.register_buffer("feat_std", torch.ones(n_mels))

        # MaxPool2d((a, b)) pools height by a and width by b, and the input is
        # laid out (freq, time) -- so every block halves frequency, while only
        # the earlier ones halve time. Frequency detail is being folded into
        # the channel dimension anyway, whereas the GRU downstream needs
        # whatever temporal resolution it can get.
        # Derived from len(channels) rather than fixed, so the depth stays a
        # free parameter: a hardcoded 4-tuple silently mismatched the block
        # count when a test built a 2-block model.
        pools = [(2, 2)] * (len(channels) - 1) + [(2, 1)]
        blocks, in_ch = [], 1
        for out_ch, pool in zip(channels, pools):
            blocks.append(ConvBlock(in_ch, out_ch, pool))
            in_ch = out_ch
        self.conv = nn.Sequential(*blocks)

        freq_out = n_mels
        for freq_pool, _ in pools:
            freq_out //= freq_pool
        if freq_out < 1:
            raise ValueError(
                f"{len(channels)} conv blocks pool {n_mels} mel bins down to nothing")
        rnn_in = channels[-1] * freq_out

        self.dropout = nn.Dropout(dropout)
        # GRU by default, LSTM available for comparison. The two differ in
        # whether a separate cell state carries information alongside the
        # hidden state; on 50 time steps per clip that extra memory has little
        # room to pay for its third of extra parameters, but that is a claim
        # worth measuring rather than assuming.
        rnn_cls = {"gru": nn.GRU, "lstm": nn.LSTM}[rnn_type.lower()]
        self.rnn = rnn_cls(rnn_in, rnn_hidden, num_layers=1, batch_first=True,
                           bidirectional=True)
        # mean AND max over time, concatenated: mean describes the clip as a
        # whole, max fires on the single best moment. For "is there a voice
        # anywhere in these 4 seconds" a brief clear phrase should count even
        # when the rest of the clip is only background, which mean-pooling
        # alone would dilute.
        self.classifier = nn.Linear(rnn_hidden * 2 * 2, num_classes)

    def set_normalization(self, mean, std):
        with torch.no_grad():
            self.feat_mean.copy_(torch.as_tensor(mean, dtype=self.feat_mean.dtype))
            self.feat_std.copy_(torch.as_tensor(std, dtype=self.feat_std.dtype).clamp_min(1e-5))

    def forward(self, x):
        x = (x - self.feat_mean) / self.feat_std
        # (B, T, F) -> (B, 1, F, T): convolutions want frequency as the
        # "height" axis so a kernel spans neighbouring bins and frames alike.
        x = x.transpose(1, 2).unsqueeze(1)
        x = self.conv(x)                       # (B, C, F', T')
        b, c, f, t = x.shape
        x = x.permute(0, 3, 1, 2).reshape(b, t, c * f)   # (B, T', C*F')
        x = self.dropout(x)
        x, _ = self.rnn(x)                     # (B, T', 2*hidden)
        pooled = torch.cat([x.mean(dim=1), x.max(dim=1).values], dim=-1)
        return self.classifier(self.dropout(pooled))


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)