import torch, torch.nn as nn, torch.nn.functional as F

HOP, STRIDE = 300, 8


class StyleEncoder(nn.Module):
    def __init__(self, style_dim=256, hidden=256, sr=24000):
        super().__init__()
        import torchaudio
        self.mel = torchaudio.transforms.MelSpectrogram(sample_rate=sr, n_fft=1024, hop_length=HOP, n_mels=80, power=1.0)
        self.convs = nn.ModuleList(nn.Conv1d(a, b, 5, stride=2, padding=2) for a, b in ((80, hidden), (hidden, hidden), (hidden, hidden)))
        self.out = nn.Linear(hidden, style_dim)
        nn.init.zeros_(self.out.weight)
        nn.init.zeros_(self.out.bias)

    def forward(self, audios):
        device = self.out.weight.device
        lengths = torch.tensor([len(a) for a in audios], device=device)
        x = (torch.log(self.mel(nn.utils.rnn.pad_sequence(audios, batch_first=True).to(device)) + 1e-5) + 4) / 4
        for conv in self.convs:
            x = F.silu(conv(x))
        frames = torch.div(lengths, HOP * STRIDE, rounding_mode="floor").clamp(min=1)
        mask = (torch.arange(x.shape[-1], device=device)[None] < frames[:, None]).float()
        return self.out((x * mask[:, None]).sum(-1) / mask.sum(-1, keepdim=True))