"""Decoder half of the YuE2 48 kHz stereo waveform VAE."""

import math

import torch
from torch import nn

SAMPLE_RATE = 48000
HOP = 1920


def WNConv1d(*args, **kwargs):
    return nn.utils.parametrizations.weight_norm(nn.Conv1d(*args, **kwargs))


def WNConvT1d(*args, **kwargs):
    return nn.utils.parametrizations.weight_norm(nn.ConvTranspose1d(*args, **kwargs))


class SnakeBeta(nn.Module):
    def __init__(self, channels: int):
        super().__init__()
        self.alpha = nn.Parameter(torch.zeros(channels))
        self.beta = nn.Parameter(torch.zeros(channels))

    def forward(self, x):
        alpha = self.alpha.unsqueeze(0).unsqueeze(-1).exp()
        beta = self.beta.unsqueeze(0).unsqueeze(-1).exp()
        return x + (1.0 / (beta + 1e-9)) * torch.sin(x * alpha).pow(2)


class ResUnit(nn.Module):
    def __init__(self, channels: int, dilation: int):
        super().__init__()
        self.layers = nn.Sequential(
            SnakeBeta(channels),
            WNConv1d(
                channels,
                channels,
                7,
                dilation=dilation,
                padding=(dilation * 6) // 2,
            ),
            SnakeBeta(channels),
            WNConv1d(channels, channels, 1),
        )

    def forward(self, x):
        return x + self.layers(x)


class DecBlock(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, stride: int):
        super().__init__()
        self.layers = nn.Sequential(
            SnakeBeta(in_channels),
            WNConvT1d(
                in_channels,
                out_channels,
                2 * stride,
                stride=stride,
                padding=math.ceil(stride / 2),
                output_padding=stride % 2,
            ),
            ResUnit(out_channels, 1),
            ResUnit(out_channels, 3),
            ResUnit(out_channels, 9),
        )

    def forward(self, x):
        return self.layers(x)


class SequenceWrapper(nn.Module):
    def __init__(self, *modules):
        super().__init__()
        self.layers = nn.Sequential(*modules)

    def forward(self, x):
        return self.layers(x)


def convert_weight_norm_keys(state_dict: dict) -> dict:
    converted = {}
    for key, value in state_dict.items():
        if key.endswith(".weight_g"):
            key = key[: -len(".weight_g")] + ".parametrizations.weight.original0"
        elif key.endswith(".weight_v"):
            key = key[: -len(".weight_v")] + ".parametrizations.weight.original1"
        converted[key] = value
    return converted


class YuE2VAEDecoder(nn.Module):
    def __init__(
        self,
        in_channels: int = 2,
        channels: int = 64,
        latent_dim: int = 64,
        channel_multipliers=(1, 2, 4, 8, 16, 32),
        strides=(2, 2, 4, 4, 5, 6),
    ):
        super().__init__()
        multipliers = [1] + list(channel_multipliers)
        decoder = [WNConv1d(latent_dim, multipliers[-1] * channels, 7, padding=3)]
        for index in range(len(multipliers) - 1, 0, -1):
            decoder.append(
                DecBlock(
                    multipliers[index] * channels,
                    multipliers[index - 1] * channels,
                    strides[index - 1],
                )
            )
        decoder += [
            SnakeBeta(multipliers[0] * channels),
            WNConv1d(multipliers[0] * channels, in_channels, 7, padding=3, bias=False),
        ]
        self.decoder = SequenceWrapper(*decoder)

    @classmethod
    def load_from_state_dict(cls, state_dict: dict, dtype=torch.float32):
        model = cls()
        state_dict = convert_weight_norm_keys(state_dict)
        model.load_state_dict({key: value.to(dtype) for key, value in state_dict.items()}, strict=True)
        return model.to(dtype)

    @property
    def device(self):
        return next(self.parameters()).device

    @property
    def dtype(self):
        return next(self.parameters()).dtype

    def decode(self, latents):
        return self.decoder(latents)

    def tiled_decode(self, latents, core_frames: int = 750, halo_frames: int = 16):
        total = latents.shape[-1]
        if total <= core_frames + 2 * halo_frames:
            return self.decoder(latents)
        pieces = []
        for start in range(0, total, core_frames):
            end = min(start + core_frames, total)
            low = max(0, start - halo_frames)
            high = min(total, end + halo_frames)
            audio = self.decoder(latents[..., low:high])
            pieces.append(
                audio[..., (start - low) * HOP : (start - low + end - start) * HOP]
            )
        return torch.cat(pieces, dim=-1)