"""
Student model for LoRaVoiceLink: a compact spectrogram U-Net that maps a
Codec2-degraded magnitude spectrogram to an enhanced one.

Kept deliberately small (edge-deployment target: real-time inference on a
Jetson) — unlike the teacher models used only offline during data
preparation, this is the model that actually ships.
"""

import torch
import torch.nn as nn


def conv_block(in_ch, out_ch):
    return nn.Sequential(
        nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
        nn.BatchNorm2d(out_ch),
        nn.LeakyReLU(0.2, inplace=True),
        nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
        nn.BatchNorm2d(out_ch),
        nn.LeakyReLU(0.2, inplace=True),
    )


class SpectrogramUNet(nn.Module):
    """
    Input:  (B, 1, F, T) degraded magnitude spectrogram — F and T must be
            multiples of 8 (three pooling stages); see dataset.py's
            pad_to_multiple().
    Output: (B, 1, F, T) enhanced magnitude spectrogram, produced by
            predicting a [0, 1] mask and multiplying it against the input
            (standard, more stable to train than predicting raw magnitude
            directly).
    """

    def __init__(self, base_channels: int = 32):
        super().__init__()
        c = base_channels

        self.enc1 = conv_block(1, c)
        self.enc2 = conv_block(c, c * 2)
        self.enc3 = conv_block(c * 2, c * 4)
        self.bottleneck = conv_block(c * 4, c * 8)

        self.pool = nn.MaxPool2d(2)

        self.up3 = nn.ConvTranspose2d(c * 8, c * 4, kernel_size=2, stride=2)
        self.dec3 = conv_block(c * 8, c * 4)
        self.up2 = nn.ConvTranspose2d(c * 4, c * 2, kernel_size=2, stride=2)
        self.dec2 = conv_block(c * 4, c * 2)
        self.up1 = nn.ConvTranspose2d(c * 2, c, kernel_size=2, stride=2)
        self.dec1 = conv_block(c * 2, c)

        self.out_conv = nn.Conv2d(c, 1, kernel_size=1)

    def forward(self, x):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        b = self.bottleneck(self.pool(e3))

        d3 = self.up3(b)
        d3 = self.dec3(torch.cat([d3, e3], dim=1))
        d2 = self.up2(d3)
        d2 = self.dec2(torch.cat([d2, e2], dim=1))
        d1 = self.up1(d2)
        d1 = self.dec1(torch.cat([d1, e1], dim=1))

        mask = torch.sigmoid(self.out_conv(d1))
        return mask * x


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


if __name__ == "__main__":
    model = SpectrogramUNet()
    dummy = torch.randn(2, 1, 256, 128)
    out = model(dummy)
    print("output shape:", out.shape)
    print("parameters:", f"{count_parameters(model):,}")