from __future__ import annotations

import torch
import torch.nn as nn


class Decoder4FeatureExtractor(nn.Module):
    """MOSS audio-tokenizer codebook decode plus decoder blocks 0..4."""

    def __init__(
        self,
        audio_tokenizer: nn.Module,
        num_quantizers: int = 32,
        output_dtype: torch.dtype = torch.float16,
    ) -> None:
        super().__init__()
        quantizer = getattr(audio_tokenizer, "quantizer")
        self.quantizers = quantizer.quantizers
        self.output_proj = quantizer.output_proj
        self.decoder_prefix = nn.ModuleList(list(audio_tokenizer.decoder[:5]))
        self.rvq_dim = int(quantizer.rvq_dim)
        self.num_quantizers = int(num_quantizers)
        self.output_dtype = output_dtype

    def forward(self, codes: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        _, batch, frames = codes.shape
        emb = torch.zeros(batch, self.rvq_dim, frames, device=codes.device, dtype=self.output_dtype)
        for index, quantizer in enumerate(self.quantizers[: self.num_quantizers]):
            if self.output_dtype == torch.float16:
                z_q = quantizer.embed_code(codes[index]).transpose(1, 2).to(self.output_dtype)
                z_q = quantizer.out_proj(z_q)
            else:
                z_q = quantizer.decode_code(codes[index])
            emb = emb + z_q.to(self.output_dtype)
        features = self.output_proj(emb)
        feature_lengths = lengths
        for module in self.decoder_prefix:
            features, feature_lengths = module(features, feature_lengths)
        return features.to(self.output_dtype), feature_lengths