import math
import librosa
from typing import List
from typing import Union, Optional

import numpy as np
import torch
import torch.nn.functional as F
from audiotools import AudioSignal
from audiotools.ml import BaseModel
from torch import nn
from einops import rearrange

from .base import CodecMixin
from dac.nn.layers import Snake1d
from dac.nn.layers import WNConv1d
from dac.nn.layers import WNConvTranspose1d
from dac.nn.quantize import (
    ResidualVectorQuantize,
    PassthroughQuantize,
    GroupedResidualVectorQuantize,
)
from dac.nn.lfq import LFQ, ResidualLFQ, GroupedLFQ
from dac.nn.vae import VAEBottleneck


def init_weights(m):
    if isinstance(m, nn.Conv1d):
        nn.init.trunc_normal_(m.weight, std=0.02)
        nn.init.constant_(m.bias, 0)


class ResidualUnit(nn.Module):
    def __init__(self, dim: int = 16, dilation: int = 1):
        super().__init__()
        pad = ((7 - 1) * dilation) // 2
        self.block = nn.Sequential(
            Snake1d(dim),
            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
            Snake1d(dim),
            WNConv1d(dim, dim, kernel_size=1),
        )

    def forward(self, x):
        y = self.block(x)
        pad = (x.shape[-1] - y.shape[-1]) // 2
        if pad > 0:
            x = x[..., pad:-pad]
        return x + y


class SubbandProjection(nn.Module):
    def __init__(self, bandwidth, out_dim):
        super(SubbandProjection, self).__init__()
        self.layer_norm = nn.LayerNorm(bandwidth)
        self.fc = nn.Linear(bandwidth, out_dim)

    def forward(self, x):
        x = rearrange(x, "b f t -> b t f")
        x = self.layer_norm(x)
        x = self.fc(x)
        x = rearrange(x, "b t f -> b f t")
        return x


class NeuralFilterbank(nn.Module):
    def __init__(
        self, n_fft: int, n_filterbank: int, sample_rate: int = 48000, out_dim: int = 64
    ):
        super(NeuralFilterbank, self).__init__()

        self.bandwidth_indices = self.get_bandwidth_indices(
            sample_rate, n_fft, n_filterbank
        )
        self.projection_modules = self.get_projection_layers(out_dim)

    def get_bandwidth_indices(self, sample_rate, n_fft, n_filterbank):
        mel_basis = librosa.filters.mel(
            sr=sample_rate, n_fft=n_fft, n_mels=n_filterbank
        )
        bandwidth_indices = [np.where(row > 0)[0] for row in mel_basis]
        return bandwidth_indices

    def get_projection_layers(self, out_dim):
        projection_modules = nn.ModuleList([])
        for indices in self.bandwidth_indices:
            indices = indices[: min(64, len(indices))]
            projection_modules.append(SubbandProjection(len(indices), out_dim))
        return projection_modules

    def forward(self, spec):
        # spec: [batch, freq, time]
        emb = []
        for indices, layer in zip(self.bandwidth_indices, self.projection_modules):
            emb.append(layer(spec[:, indices[: min(64, len(indices))], :]))
        return rearrange(torch.stack(emb), "f b c t -> b c f t")


class Res2dModule(nn.Module):
    def __init__(self, idim, odim, stride=(2, 2), dilation=(1, 1), padding=(1, 1)):
        super(Res2dModule, self).__init__()
        self.conv1 = nn.Conv2d(
            idim, odim, 3, padding=padding, stride=stride, dilation=dilation
        )
        self.bn1 = nn.BatchNorm2d(odim)
        self.conv2 = nn.Conv2d(odim, odim, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(odim)
        self.relu = nn.ReLU()

        # residual
        self.diff = False
        if (idim != odim) or (stride[0] > 1):
            self.conv3 = nn.Conv2d(
                idim, odim, 3, padding=padding, stride=stride, dilation=dilation
            )
            self.bn3 = nn.BatchNorm2d(odim)
            self.diff = True

    def forward(self, x):
        out = self.bn2(self.conv2(self.relu(self.bn1(self.conv1(x)))))
        if self.diff:
            x = self.bn3(self.conv3(x))
        out = x + out
        out = self.relu(out)
        return out


class LayerNorm(nn.Module):
    r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
    with shape (batch_size, channels, height, width).
    """

    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise NotImplementedError
        self.normalized_shape = (normalized_shape,)

    def forward(self, x):
        if self.data_format == "channels_last":
            return F.layer_norm(
                x, self.normalized_shape, self.weight, self.bias, self.eps
            )
        elif self.data_format == "channels_first":
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x


class Block(nn.Module):
    r"""ConvNeXt Block. There are two equivalent implementations:
    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
    We use (2) as we find it slightly faster in PyTorch
    Args:
        dim (int): Number of input channels.
        drop_path (float): Stochastic depth rate. Default: 0.0
        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
    """

    def __init__(self, dim, drop_path=0.0, layer_scale_init_value=1e-6):
        super().__init__()
        self.dwconv = WNConv1d(
            dim, dim, kernel_size=7, padding=3, groups=dim
        )  # depthwise conv
        self.norm = LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(
            dim, 4 * dim
        )  # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * dim, dim)
        self.gamma = (
            nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
            if layer_scale_init_value > 0
            else None
        )
        self.drop_path = nn.Identity()

    def forward(self, x):
        input = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 1)  # (N, C, H) -> (N, H, C)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 2, 1)

        x = input + self.drop_path(x)
        return x


class ConvNeXtSimple(nn.Module):
    """No downsampling, just 4 blocks of ConvNeXt"""

    def __init__(self, dim, depth, drop_path=0.0, layer_scale_init_value=1e-6):
        super().__init__()
        self.blocks = nn.ModuleList(
            [Block(dim, drop_path, layer_scale_init_value) for _ in range(depth)]
        )

    def forward(self, x):
        for block in self.blocks:
            x = block(x)
        return x


class ConvNeXtBlock(nn.Module):
    """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.

    Args:
        dim (int): Number of input channels.
        intermediate_dim (int): Dimensionality of the intermediate layer.
        layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling.
            Defaults to None.
        adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm.
            None means non-conditional LayerNorm. Defaults to None.
    """

    def __init__(
        self,
        dim: int,
        intermediate_dim: int,
        layer_scale_init_value: float,
        adanorm_num_embeddings: Optional[int] = None,
    ):
        super().__init__()
        self.dwconv = nn.Conv1d(
            dim, dim, kernel_size=7, padding=3, groups=dim
        )  # depthwise conv
        self.adanorm = adanorm_num_embeddings is not None
        if adanorm_num_embeddings:
            self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6)
        else:
            self.norm = nn.LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(
            dim, intermediate_dim
        )  # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(intermediate_dim, dim)
        self.gamma = (
            nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True)
            if layer_scale_init_value > 0
            else None
        )

    def forward(
        self, x: torch.Tensor, cond_embedding_id: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        residual = x
        x = self.dwconv(x)
        x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
        if self.adanorm:
            assert cond_embedding_id is not None
            x = self.norm(x, cond_embedding_id)
        else:
            x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.transpose(1, 2)  # (B, T, C) -> (B, C, T)

        x = residual + x
        return x


class AdaLayerNorm(nn.Module):
    """
    Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes

    Args:
        num_embeddings (int): Number of embeddings.
        embedding_dim (int): Dimension of the embeddings.
    """

    def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.dim = embedding_dim
        self.scale = nn.Embedding(
            num_embeddings=num_embeddings, embedding_dim=embedding_dim
        )
        self.shift = nn.Embedding(
            num_embeddings=num_embeddings, embedding_dim=embedding_dim
        )
        torch.nn.init.ones_(self.scale.weight)
        torch.nn.init.zeros_(self.shift.weight)

    def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor:
        scale = self.scale(cond_embedding_id)
        shift = self.shift(cond_embedding_id)
        x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps)
        x = x * scale + shift
        return x


class Encoder1d(nn.Module):
    def __init__(
        self,
        d_model: int = 64,
        d_latent: int = 64,
        n_fft: int = 2048,
        hop_size: int = 480,
        n_bands: int = 128,
        sample_rate: int = 48000,
        nf_dim: int = 8,
        depth: int = 4,
    ):
        super().__init__()

        # stft parameters
        self.n_fft = n_fft
        self.hop_size = hop_size
        self.window = torch.hann_window(n_fft)

        # neural filterbank
        self.neural_filterbank = NeuralFilterbank(
            n_fft=n_fft,
            n_filterbank=n_bands,
            sample_rate=sample_rate,
            out_dim=nf_dim,
        )

        # First conv layer
        self.inp_layer = WNConv1d(
            nf_dim * n_bands * 4, d_model, kernel_size=3, padding=1
        )

        # Convnext
        self.conv_next = nn.ModuleList(
            [
                ConvNeXtBlock(
                    dim=d_model,
                    intermediate_dim=d_model * 3,
                    layer_scale_init_value=1 / depth,
                    adanorm_num_embeddings=None,
                )
                for _ in range(depth)
            ]
        )

        # last conv
        self.prj = [
            Snake1d(d_model),
            WNConv1d(d_model, d_latent, kernel_size=3, padding=1),
        ]
        self.prj = nn.Sequential(*self.prj)
        self.enc_dim = d_model

    def forward(self, x):
        """
        einops

        b: batch
        f: frequency
        d: feature dimension
        s: stereo
        c: complex
        """
        # reshape audio
        b, s, _ = x.shape
        x = rearrange(x, "b s t -> (b s) t")

        # short-time Fourier transform (b s) t -> (b s) f t c
        self.window = self.window.to(x.device)
        spec = torch.stft(
            x,
            n_fft=self.n_fft,
            hop_length=self.hop_size,
            win_length=self.n_fft,
            window=self.window,
            onesided=True,
            return_complex=True,
        )
        spec = torch.view_as_real(spec)
        t = spec.shape[-2]

        # neural filterbank (b s) f t c -> (b s) (c d) f' t)
        spec = rearrange(spec, "(b s) f t c -> (b s c) f t", b=b, s=s, t=t, c=2)
        nf_spec = self.neural_filterbank(spec)
        nf_spec = rearrange(
            nf_spec, "(b s c) d f t -> b (s c d f) t", b=b, s=s, t=t, c=2
        )

        # encoder b (s c d f) t -> b (s c d) t
        out = self.inp_layer(nf_spec)
        for conv_block in self.conv_next:
            out = conv_block(out)

        # projection b (s c d) t -> b c' t
        out = self.prj(out)

        return out


class Encoder2d(nn.Module):
    def __init__(
        self,
        d_model: int = 64,
        d_latent: int = 64,
        n_fft: int = 2048,
        hop_size: int = 480,
        n_bands: int = 128,
        sample_rate: int = 48000,
    ):
        super().__init__()

        # stft parameters
        self.n_fft = n_fft
        self.hop_size = hop_size
        self.window = torch.hann_window(n_fft)

        # neural filterbank
        self.neural_filterbank = NeuralFilterbank(
            n_fft=n_fft,
            n_filterbank=n_bands,
            sample_rate=sample_rate,
            out_dim=d_model // 2,
        )

        # encoder blocks
        self.res_block = []
        d_model *= 2
        for ix, stride in enumerate([2, 2, 2, 4, 4]):
            hop = 1 if ix == 0 else 2
            d_model *= 2
            self.res_block.append(
                Res2dModule(
                    idim=d_model // 2,
                    odim=d_model,
                    stride=(stride, 1),
                    dilation=(1, hop),
                    padding=(1, hop),
                )
            )
        self.res_block = nn.Sequential(*self.res_block)

        # last conv
        self.prj = [
            Snake1d(d_model),
            WNConv1d(d_model, d_latent, kernel_size=3, padding=1),
        ]
        self.prj = nn.Sequential(*self.prj)
        self.enc_dim = d_model

    def forward(self, x):
        """
        einops

        b: batch
        f: frequency
        d: feature dimension
        s: stereo
        c: complex
        """
        # reshape audio
        b, s, _ = x.shape
        x = rearrange(x, "b s t -> (b s) t")

        # short-time Fourier transform (b s) t -> (b s) f t c
        self.window = self.window.to(x.device)
        spec = torch.stft(
            x,
            n_fft=self.n_fft,
            hop_length=self.hop_size,
            win_length=self.n_fft,
            window=self.window,
            onesided=True,
            return_complex=True,
        )
        spec = torch.view_as_real(spec)
        t = spec.shape[-2]

        # neural filterbank (b s) f t c -> (b s) (c d) f' t)
        spec = rearrange(spec, "(b s) f t c -> (b s c) f t", b=b, s=s, t=t, c=2)
        nf_spec = self.neural_filterbank(spec)
        nf_spec = rearrange(
            nf_spec, "(b s c) d f t -> b  (s c d) f t", b=b, s=s, t=t, c=2
        )

        # encoder b (s c d) f' t -> b (s c d) t
        out = self.res_block(nf_spec)
        out = out.squeeze(2)

        # projection b (s c d) t -> b c' t
        out = self.prj(out)

        return out


class DecoderVocos(nn.Module):
    def __init__(
        self,
        input_channel,
        channels,
        d_out: int = 2,
        n_fft: int = 2048,
        hop_size: int = 480,
        depth: int = 4,
    ):
        super().__init__()

        # stft parameters
        self.n_fft = n_fft
        self.hop_size = hop_size
        self.window = torch.hann_window(n_fft)

        # Input BN
        # self.input_bn = nn.BatchNorm1d(input_channel)

        # First conv layer
        self.inp_layer = WNConv1d(input_channel, channels, kernel_size=7, padding=3)
        self.layer_norm = nn.LayerNorm(channels, eps=1e-6)

        # Convnext
        self.conv_next = nn.ModuleList(
            [
                ConvNeXtBlock(
                    dim=channels,
                    intermediate_dim=channels * 3,
                    layer_scale_init_value=1 / depth,
                    adanorm_num_embeddings=None,
                )
                for _ in range(depth)
            ]
        )

        # Final conv layer
        self.out_layer = nn.Linear(channels, 2 * (n_fft + 2))

    def forward(self, x):
        # decoding conv b (s c d) t -> b d' t
        # input layer
        # x = self.input_bn(x)
        x = self.inp_layer(x)

        # layer norm
        x = x.transpose(1, 2)
        x = self.layer_norm(x)
        x = x.transpose(1, 2)

        # convnext
        for conv_block in self.conv_next:
            x = conv_block(x)

        # layer norm
        x = x.transpose(1, 2)
        x = self.layer_norm(x)

        # projection
        x = self.out_layer(x)
        x = x.transpose(1, 2)

        # stereo
        x = rearrange(x, "b (s d) t -> (b s) d t", s=2)

        # stft params
        mag, p = x.chunk(2, dim=1)
        mag = torch.exp(mag)
        mag = torch.clip(
            mag, max=1e2
        )  # safeguard to prevent excessively large magnitudes
        # wrapping happens here. These two lines produce real and imaginary value
        x = torch.cos(p)
        y = torch.sin(p)
        # recalculating phase here does not produce anything new
        # only costs time
        # phase = torch.atan2(y, x)
        # S = mag * torch.exp(phase * 1j)
        # better directly produce the complex value
        S = mag * (x + 1j * y)

        # iSTFT
        self.window = self.window.to(S.device)
        audio = torch.istft(
            S,
            self.n_fft,
            self.hop_size,
            self.n_fft,
            self.window,
            center=True,
        )

        # reshape
        audio = rearrange(audio, "(b s) t -> b s t", s=2)

        return audio


class DAC(BaseModel, CodecMixin):
    def __init__(
        self,
        encoder_dim: int = 64,
        latent_dim: int = None,
        decoder_dim: int = 1536,
        n_codebooks: int = 8,
        n_codebooks_per_level: list[int] = [2, 2, 2, 2],
        codebook_size: int = 1024,
        codebook_dim: Union[int, list] = 8,
        quantizer_dropout: bool = False,
        sample_rate: int = 48000,
        hop_size: int = 480,
        quantizer_type: str = "rvq",
        use_vae: bool = False,
        residual_type: str = "subtract",
        enc_style: str = "2d",
        dec_style: str = "vocos",
        enc_depth: int = 4,
        dec_depth: int = 4,
    ):
        super().__init__()

        self.encoder_dim = encoder_dim
        self.decoder_dim = decoder_dim
        self.sample_rate = sample_rate
        self.quantizer_type = quantizer_type
        self.use_vae = use_vae
        self.residual_type = residual_type
        if use_vae and quantizer_type != "rvq":
            raise NotImplementedError
        self.enc_style = enc_style
        self.dec_style = dec_style
        self.enc_depth = enc_depth
        self.dec_depth = dec_depth

        if latent_dim is None:
            latent_dim = encoder_dim * (2 ** len(5))

        self.latent_dim = latent_dim

        self.hop_length = hop_size
        if enc_style == "1d":
            self.encoder = Encoder1d(encoder_dim, latent_dim, depth=enc_depth)
        elif enc_style == "2d":
            self.encoder = Encoder2d(encoder_dim, latent_dim)

        self.n_codebooks = n_codebooks
        self.n_codebooks_per_level = n_codebooks_per_level
        self.codebook_size = codebook_size
        self.codebook_dim = codebook_dim

        if quantizer_type == "rvq" or quantizer_type == "rlfq":
            self.quantizer = ResidualVectorQuantize(
                input_dim=latent_dim,
                n_codebooks=n_codebooks,
                codebook_size=codebook_size,
                codebook_dim=codebook_dim,
                quantizer_dropout=quantizer_dropout,
                use_vae=use_vae,
                residual_type=residual_type,
                use_lfq=quantizer_type == "rlfq",
            )
        elif quantizer_type == "grvq":
            self.quantizer = GroupedResidualVectorQuantize(
                input_dim=latent_dim,
                n_codebooks_per_level=n_codebooks_per_level,
                codebook_size=codebook_size,
                codebook_dim=codebook_dim,
                quantizer_dropout=quantizer_dropout,
            )
        elif quantizer_type == "passthrough":
            self.quantizer = PassthroughQuantize()
        elif quantizer_type == "lfq":
            assert isinstance(codebook_dim, int)
            self.quantizer = LFQ(
                codebook_size=codebook_size,
                num_codebooks=n_codebooks,
                dim=latent_dim,
                straight_through_activation=torch.tanh,
            )
        elif quantizer_type == "grouped_lfq":
            self.quantizer = GroupedLFQ(
                codebook_size=codebook_size,
                num_quantizers=n_codebooks,
                dim=latent_dim,
                quantize_dropout=quantizer_dropout,
            )
        elif quantizer_type == "vae":
            self.quantizer = VAEBottleneck(latent_dim=latent_dim, dim=codebook_dim)

        if dec_style == "vocos":
            self.decoder = DecoderVocos(
                latent_dim,
                decoder_dim,
                depth=dec_depth,
            )
        self.sample_rate = sample_rate
        self.apply(init_weights)

        self.delay = self.get_delay()

    def preprocess(self, audio_data, sample_rate):
        if sample_rate is None:
            sample_rate = self.sample_rate
        assert sample_rate == self.sample_rate

        length = audio_data.shape[-1]
        right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
        audio_data = nn.functional.pad(audio_data, (0, right_pad))

        return audio_data

    def encode(
        self,
        audio_data: torch.Tensor,
        n_quantizers: int = None,
    ):
        """Encode given audio data and return quantized latent codes

        Parameters
        ----------
        audio_data : Tensor[B x 1 x T]
            Audio data to encode
        n_quantizers : int, optional
            Number of quantizers to use, by default None
            If None, all quantizers are used.

        Returns
        -------
        dict
            A dictionary with the following keys:
            "z" : Tensor[B x D x T]
                Quantized continuous representation of input
            "codes" : Tensor[B x N x T]
                Codebook indices for each codebook
                (quantized discrete representation of input)
            "latents" : Tensor[B x N*D x T]
                Projected latents (continuous representation of input before quantization)
            "vq/commitment_loss" : Tensor[1]
                Commitment loss to train encoder to predict vectors closer to codebook
                entries
            "vq/codebook_loss" : Tensor[1]
                Codebook loss to update the codebook
            "length" : int
                Number of samples in input audio
        """
        z = self.encoder(audio_data)
        return self.quantizer(z, n_quantizers=n_quantizers)

    def decode(self, z: torch.Tensor):
        """Decode given latent codes and return audio data

        Parameters
        ----------
        z : Tensor[B x D x T]
            Quantized continuous representation of input
        length : int, optional
            Number of samples in output audio, by default None

        Returns
        -------
        dict
            A dictionary with the following keys:
            "audio" : Tensor[B x 1 x length]
                Decoded audio data.
        """
        return self.decoder(z)

    def forward(
        self,
        audio_data: torch.Tensor,
        sample_rate: int = None,
        n_quantizers: int = None,
    ):
        """Model forward pass

        Parameters
        ----------
        audio_data : Tensor[B x 1 x T]
            Audio data to encode
        sample_rate : int, optional
            Sample rate of audio data in Hz, by default None
            If None, defaults to `self.sample_rate`
        n_quantizers : int, optional
            Number of quantizers to use, by default None.
            If None, all quantizers are used.

        Returns
        -------
        dict
            A dictionary with the following keys:
            "z" : Tensor[B x D x T]
                Quantized continuous representation of input
            "codes" : Tensor[B x N x T]
                Codebook indices for each codebook
                (quantized discrete representation of input)
            "latents" : Tensor[B x N*D x T]
                Projected latents (continuous representation of input before quantization)
            "vq/commitment_loss" : Tensor[1]
                Commitment loss to train encoder to predict vectors closer to codebook
                entries
            "vq/codebook_loss" : Tensor[1]
                Codebook loss to update the codebook
            "length" : int
                Number of samples in input audio
            "audio" : Tensor[B x 1 x length]
                Decoded audio data.
        """

        audio_data = self.preprocess(audio_data, sample_rate)
        q_res = self.encode(audio_data, n_quantizers)

        x = self.decode(q_res["z"])

        return {
            "audio": x,
            **q_res,
        }


if __name__ == "__main__":
    import numpy as np
    from functools import partial

    model = DAC().to("cpu")

    for n, m in model.named_modules():
        o = m.extra_repr()
        p = sum([np.prod(p.size()) for p in m.parameters()])
        fn = lambda o, p: o + f" {p/1e6:<.3f}M params."
        setattr(m, "extra_repr", partial(fn, o=o, p=p))
    print(model)
    print("Total # of params: ", sum([np.prod(p.size()) for p in model.parameters()]))

    length = 88200 * 2
    x = torch.randn(1, 1, length).to(model.device)
    x.requires_grad_(True)
    x.retain_grad()

    # Make a forward pass
    out = model(x)["audio"]
    print("Input shape:", x.shape)
    print("Output shape:", out.shape)

    # Create gradient variable
    grad = torch.zeros_like(out)
    grad[:, :, grad.shape[-1] // 2] = 1

    # Make a backward pass
    out.backward(grad)

    # Check non-zero values
    gradmap = x.grad.squeeze(0)
    gradmap = (gradmap != 0).sum(0)  # sum across features
    rf = (gradmap != 0).sum()

    print(f"Receptive field: {rf.item()}")

    x = AudioSignal(torch.randn(1, 1, 44100 * 60), 44100)
    model.decompress(model.compress(x, verbose=True), verbose=True)