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 rotary_embedding_torch import RotaryEmbedding

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
from dac.nn.attend import Attend


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")


# Frequency attention module
def exists(val):
    return val is not None


# norm
def l2norm(t):
    return F.normalize(t, dim=-1, p=2)


class RMSNorm(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.scale = dim**0.5
        self.gamma = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return F.normalize(x, dim=-1) * self.scale * self.gamma


# attention


class FeedForward(nn.Module):
    def __init__(self, dim, mult=4, dropout=0.0):
        super().__init__()
        dim_inner = int(dim * mult)
        self.net = nn.Sequential(
            RMSNorm(dim),
            nn.Linear(dim, dim_inner),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(dim_inner, dim),
            nn.Dropout(dropout),
        )

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


class Attention(nn.Module):
    def __init__(
        self, dim, heads=8, dim_head=64, dropout=0.0, rotary_embed=None, flash=True
    ):
        super().__init__()
        self.heads = heads
        self.scale = dim_head**-0.5
        dim_inner = heads * dim_head

        self.rotary_embed = rotary_embed

        self.attend = Attend(flash=flash, dropout=dropout)

        self.norm = RMSNorm(dim)
        self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False)

        self.to_gates = nn.Linear(dim, heads)

        self.to_out = nn.Sequential(
            nn.Linear(dim_inner, dim, bias=False), nn.Dropout(dropout)
        )

    def forward(self, x):
        x = self.norm(x)

        q, k, v = rearrange(
            self.to_qkv(x), "b n (qkv h d) -> qkv b h n d", qkv=3, h=self.heads
        )

        if exists(self.rotary_embed):
            q = self.rotary_embed.rotate_queries_or_keys(q)
            k = self.rotary_embed.rotate_queries_or_keys(k)

        out = self.attend(q, k, v)

        gates = self.to_gates(x)
        out = out * rearrange(gates, "b n h -> b h n 1").sigmoid()

        out = rearrange(out, "b h n d -> b n (h d)")
        return self.to_out(out)


class Transformer(nn.Module):
    def __init__(
        self,
        *,
        dim,
        depth,
        dim_head=64,
        heads=8,
        attn_dropout=0.0,
        ff_dropout=0.0,
        ff_mult=4,
        norm_output=True,
        rotary_embed=None,
        flash_attn=True,
    ):
        super().__init__()
        self.layers = nn.ModuleList([])

        for _ in range(depth):
            attn = Attention(
                dim=dim,
                dim_head=dim_head,
                heads=heads,
                dropout=attn_dropout,
                rotary_embed=rotary_embed,
                flash=flash_attn,
            )

            self.layers.append(
                nn.ModuleList(
                    [attn, FeedForward(dim=dim, mult=ff_mult, dropout=ff_dropout)]
                )
            )

        self.norm = RMSNorm(dim) if norm_output else nn.Identity()

    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x

        return self.norm(x)


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 Encoder(nn.Module):
    def __init__(
        self,
        d_latent: int = 64,
        n_fft: int = 2048,
        hop_size: int = 480,
        n_bands: int = 64,
        sample_rate: int = 48000,
        nf_dim: int = 32,
        depth: int = 4,
        n_heads: int = 8,
    ):
        super().__init__()

        d_model = nf_dim * 2 * 2

        # 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,
        )

        # frequency attention
        self.depth = depth
        transformer_kwargs = dict(
            dim=d_model,
            heads=n_heads,
            dim_head=d_model // n_heads,
            attn_dropout=0.0,
            ff_dropout=0.0,
            flash_attn=True,
            norm_output=False,
        )
        rotary_emb = RotaryEmbedding(dim=d_model // n_heads)
        freq_attn = []
        for _ in range(depth):
            freq_attn.append(
                Transformer(
                    depth=1, rotary_embed=rotary_emb, **transformer_kwargs
                ).bfloat16()
            )
        self.freq_attn = nn.ModuleList(freq_attn)

        # 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 * n_bands),
            WNConv1d(d_model * n_bands, 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)
        out = rearrange(nf_spec, "(b s c) d f t -> (b t) f (s c d)", b=b, s=s, t=t, c=2)

        # frequency attention + convnext
        for i in range(self.depth):
            out = self.freq_attn[i](out.bfloat16())
            out = rearrange(
                out.float(), "(b t) f (s c d) -> (b f) (s c d) t", b=b, s=s, t=t, c=2
            ).float()
            out = self.conv_next[i](out)
            out = rearrange(
                out, "(b f) (s c d) t -> (b t) f (s c d)", b=b, s=s, t=t, c=2
            )

        # rearrange
        out = rearrange(out, "(b t) f (s c d) -> b (s f c d) t", b=b, s=s, t=t, c=2)

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

        return out


class Decoder(nn.Module):
    def __init__(
        self,
        input_channel,
        channels,
        n_bands: int = 1025,
        n_fft: int = 2048,
        hop_size: int = 480,
        depth: int = 4,
        n_heads: int = 8,
    ):
        super().__init__()

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

        # First conv layer
        d_model = channels * 2 * 2
        self.inp_layer = WNConv1d(
            input_channel, d_model * n_bands, kernel_size=7, padding=3
        )

        # 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)
            ]
        )

        # frequency attention
        self.depth = depth
        transformer_kwargs = dict(
            dim=d_model,
            heads=n_heads,
            dim_head=d_model // n_heads,
            attn_dropout=0.0,
            ff_dropout=0.0,
            flash_attn=True,
            norm_output=False,
        )
        rotary_emb = RotaryEmbedding(dim=d_model // n_heads)
        freq_attn = []
        for _ in range(depth):
            freq_attn.append(
                Transformer(
                    depth=1, rotary_embed=rotary_emb, **transformer_kwargs
                ).bfloat16()
            )
        self.freq_attn = nn.ModuleList(freq_attn)

        # Final conv layer
        self.prj = [Snake1d(d_model), WNConv1d(d_model, 4, kernel_size=7, padding=3)]
        self.prj = nn.Sequential(*self.prj)

    def forward(self, x):
        # projection b c' t -> (b f) (s c d) t
        b, _, t = x.shape
        out = self.inp_layer(x)
        out = rearrange(
            out,
            "b (s f c d) t -> (b f) (s c d) t",
            s=2,
            f=self.n_bands,
            c=2,
            d=self.n_dim,
        )

        # convnext + frequency attention
        for i in range(self.depth):
            out = self.conv_next[i](out)
            out = rearrange(
                out, "(b f) (s c d) t -> (b t) f (s c d)", b=b, s=2, t=t, c=2
            )
            out = self.freq_attn[i](out.bfloat16())
            out = rearrange(
                out.float(), "(b t) f (s c d) -> (b f) (s c d) t", b=b, s=2, t=t, c=2
            ).float()

        # projection
        out = rearrange(out, "(b f) (s c d) t -> b (s c d) (f t)", b=b, s=2, t=t, c=2)
        out = self.prj(out)
        out = rearrange(out, "b (s c) (f t) -> (b s) c f t", b=b, s=2, t=t, c=2)

        # stft params
        mag = out[:, 0, :]
        p = out[:, 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,
        encoder_heads: int = 8,
        latent_dim: int = None,
        n_bands: int = 64,
        decoder_dim: int = 64,
        decoder_heads: int = 8,
        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_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_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
        self.encoder = Encoder(
            d_latent=latent_dim,
            n_bands=n_bands,
            nf_dim=encoder_dim,
            depth=enc_depth,
            n_heads=encoder_heads,
        )

        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)

        self.decoder = Decoder(
            latent_dim,
            decoder_dim,
            depth=dec_depth,
            n_heads=decoder_heads,
        )
        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)