import torch
import torch.nn as nn

from .inference_hubert import InferenceHubertBase
from .vae_memory_bank import VAEMemoryBank


def create_padding_mask(waveforms_lengths: torch.Tensor = None):
    if waveforms_lengths is None:
        return None
    batch = waveforms_lengths.shape[0]
    max_len = waveforms_lengths.max()
    device = waveforms_lengths.device
    padding_mask = torch.ones(batch, max_len, dtype=torch.bool, device=device)
    for idx, length in enumerate(waveforms_lengths):
        padding_mask[idx, :length] = 0
    return padding_mask


def unfreeze_layers(model: nn.Module, root_name: str):
    for name, param in model.named_parameters():
        if root_name in name[: len(root_name)]:
            param.requires_grad = True


class PosteriorHubert(nn.Module):
    def __init__(
        self, out_channels, feature_channels, downsample_channels, output_layer=11
    ) -> None:
        super().__init__()
        self.out_channels = out_channels
        self.feature_channels = feature_channels
        self.downsample_channels = downsample_channels
        self.output_layer = output_layer

        self.hubert = InferenceHubertBase()
        self.memory_bank = VAEMemoryBank(
            n_hidden_dims=feature_channels,
            bank_size=1000,
            output_channels=downsample_channels,
        )

        self.proj = nn.Conv1d(downsample_channels, out_channels * 2, 1)

    def forward(self, waveforms: torch.Tensor, waveforms_lengths: torch.Tensor, g=None):
        features, features_mask = self.hubert.extract_features(
            source=waveforms,
            padding_mask=create_padding_mask(waveforms_lengths),
            output_layer=self.output_layer,
        )
        x = self.memory_bank(features.transpose(1, 2))
        x_mask = (~features_mask).unsqueeze(1).to(torch.float32)
        x = x[:, :, : x_mask.shape[-1]]

        stats = self.proj(x) * x_mask
        m, logs = torch.split(stats, self.out_channels, dim=1)

        z = (m + torch.randn_like(m) * torch.exp(logs)) * x_mask
        return z, m, logs, x_mask