"""
DroneClassifier — CRNN for binary drone audio detection.

Architecture: CNN feature pyramid (3 blocks) → BiGRU bottleneck → FC head.
Input : log-mel spectrogram (B, 1, 64, T)
Output: raw logit (B, 1)  — pass through sigmoid for probability

Parameter count: 1 486 113  (~5.94 MB FP32 / ~1.49 MB INT8)
"""

import torch
import torch.nn as nn


def _conv_block(in_ch: int, out_ch: int) -> nn.Sequential:
    return nn.Sequential(
        nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),
        nn.BatchNorm2d(out_ch),
        nn.ReLU(inplace=True),
        nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),
        nn.BatchNorm2d(out_ch),
        nn.ReLU(inplace=True),
    )


class SharedEncoder(nn.Module):
    """
    Input : (B, 1, 64, T)   log-mel spectrogram, T ≈ 101 frames for 1 s
    Output: (B, T//2, 256)  BiGRU hidden states

    Pooling stages:
      pool1: MaxPool2d(2,2)       freq 64→32,  time T→T//2
      pool2: MaxPool2d(2,2)       freq 32→16,  time unchanged
      pool3: MaxPool2d((2,1))     freq 16→8,   time unchanged
    Bottleneck shape: (B, 128, 8, T//2) → reshape → (B, T//2, 1024) → BiGRU
    """

    GRU_INPUT  = 128 * 8   # channels × freq bins at bottleneck
    GRU_HIDDEN = 128

    def __init__(self):
        super().__init__()
        self.enc1  = _conv_block(1,   32)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.enc2  = _conv_block(32,  64)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.enc3  = _conv_block(64, 128)
        self.pool3 = nn.MaxPool2d(kernel_size=(2, 1), stride=(2, 1))

        self.gru = nn.GRU(
            self.GRU_INPUT,
            self.GRU_HIDDEN,
            num_layers=2,
            batch_first=True,
            bidirectional=True,
            dropout=0.2,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.pool1(self.enc1(x))    # (B, 32,  32, T//2)
        x = self.pool2(self.enc2(x))    # (B, 64,  16, T//2)
        x = self.pool3(self.enc3(x))    # (B, 128,  8, T//2)

        B, C, F, Tp = x.shape
        rnn_in  = x.permute(0, 3, 1, 2).reshape(B, Tp, C * F)  # (B, T//2, 1024)
        rnn_out, _ = self.gru(rnn_in)                            # (B, T//2, 256)
        return rnn_out


class ClassifierHead(nn.Module):
    """Global-average-pool over GRU time steps → binary logit."""

    def __init__(self, in_features: int = 256):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(in_features, 64),
            nn.ReLU(inplace=True),
            nn.Dropout(0.3),
            nn.Linear(64, 1),
        )

    def forward(self, rnn_out: torch.Tensor) -> torch.Tensor:
        feat = rnn_out.mean(dim=1)   # (B, 256)
        return self.fc(feat)         # (B, 1)  raw logit


class DroneClassifier(nn.Module):
    """Full model: encoder + classifier head."""

    def __init__(self):
        super().__init__()
        self.encoder    = SharedEncoder()
        self.classifier = ClassifierHead()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Returns raw logit (B, 1). Apply sigmoid for probability."""
        return self.classifier(self.encoder(x))


def load_classifier(checkpoint: str, device: torch.device | str = "cpu") -> "DroneClassifier":
    """
    Load a DroneClassifier from a checkpoint.
    Accepts both new (classifier-only) and old (joint model with separator) state dicts —
    separator keys are silently ignored.
    """
    model = DroneClassifier()
    state = torch.load(checkpoint, map_location=device, weights_only=False)
    if "model_state_dict" in state:
        state = state["model_state_dict"]
    model.load_state_dict(state, strict=False)
    return model.to(device).eval()


# Back-compat alias so code that imports JointModel still works
JointModel = DroneClassifier