import torch
from torch import nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput

from .configuration_smad import SmadConfig


class ConvBlock(nn.Module):
    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):
    def __init__(
        self,
        n_mels=80,
        channels=(32, 64, 128, 128),
        rnn_hidden=128,
        dropout=0.2,
        num_classes=4,
        rnn_type="gru",
    ):
        super().__init__()
        self.register_buffer("feat_mean", torch.zeros(n_mels))
        self.register_buffer("feat_std", torch.ones(n_mels))

        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)
        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,
        )
        self.classifier = nn.Linear(rnn_hidden * 2 * 2, num_classes)

    def forward(self, x):
        x = (x - self.feat_mean) / self.feat_std
        x = x.transpose(1, 2).unsqueeze(1)
        x = self.conv(x)
        b, c, f, t = x.shape
        x = x.permute(0, 3, 1, 2).reshape(b, t, c * f)
        x = self.dropout(x)
        x, _ = self.rnn(x)
        pooled = torch.cat([x.mean(dim=1), x.max(dim=1).values], dim=-1)
        return self.classifier(self.dropout(pooled))


class SmadForAudioClassification(PreTrainedModel):
    config_class = SmadConfig
    base_model_prefix = "smad"
    main_input_name = "input_features"
    all_tied_weights_keys = {}

    def __init__(self, config):
        super().__init__(config)
        self.smad = TinyAudioCRNN(
            n_mels=config.num_mels,
            channels=tuple(config.channels),
            rnn_hidden=config.rnn_hidden,
            dropout=config.dropout,
            num_classes=config.num_labels,
            rnn_type=config.rnn_type,
        )

    def forward(self, input_features=None, labels=None, return_dict=None, **kwargs):
        if input_features is None:
            input_features = kwargs.pop("inputs", None)
        if input_features is None:
            raise ValueError("Pass log-mel features as `input_features`.")

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        logits = self.smad(input_features)
        loss = None
        if labels is not None:
            loss = nn.functional.cross_entropy(logits, labels)
        if not return_dict:
            output = (logits,)
            return ((loss,) + output) if loss is not None else output
        return SequenceClassifierOutput(loss=loss, logits=logits)

    @torch.no_grad()
    def predict_proba(self, input_features):
        logits = self(input_features=input_features).logits
        return torch.softmax(logits / float(self.config.temperature), dim=-1)