import torch 
from torch import Tensor
from transformers import PreTrainedModel
from audio_encoders_pytorch import MelE1d, TanhBottleneck
from audio_diffusion_pytorch import DiffusionAE, UNetV0, LTPlugin, VDiffusion, VSampler
from .config import DMAE1dConfig 


class DMAE1d(PreTrainedModel):

    config_class = DMAE1dConfig

    def __init__(self, config: DMAE1dConfig):
        super().__init__(config)

        UNet = LTPlugin(
            UNetV0,
            num_filters=128,
            window_length=64,
            stride=32,
        )

        self.model = DiffusionAE(
            net_t=UNet, 
            dim=1, 
            in_channels=2,
            channels=[256, 512, 512, 512, 512],
            factors=[1, 2, 2, 2, 2],
            items=[1, 2, 2, 2, 4],
            inject_depth=4,
            encoder=MelE1d(
                in_channels=2,
                channels=512,
                multipliers=[1, 1],
                factors=[2],
                num_blocks=[12],
                out_channels=32,
                mel_channels=80,
                mel_sample_rate=48000,
                mel_normalize_log=True,
                bottleneck=TanhBottleneck()
            ),
            diffusion_t=VDiffusion, 
            sampler_t=VSampler
        )
    
    def forward(self, *args, **kwargs):
        return self.model(*args, **kwargs)

    def encode(self, *args, **kwargs):
        return self.model.encode(*args, **kwargs)
    
    @torch.no_grad()
    def decode(self, *args, **kwargs):
        return self.model.decode(*args, **kwargs)