import os
import torch
import logging
import pytorch_lightning as pl

from ldm.modules.bigvgan.bigvgan import VocoderBigVGAN
from ldm.modules.latent_encoder.model import Encoder, Decoder

from ldm.modules.diffusionmodules.distributions import (
    DiagonalGaussianDistribution,
)


import soundfile as sf


class AutoencoderKL(pl.LightningModule):
    def __init__(
        self,
        ddconfig=None,
        embed_dim=None,
        reload_from_ckpt=None,
        ignore_keys=[],
        vocoder_path=None,
    ):
        super().__init__()

        self.encoder = Encoder(**ddconfig)
        self.decoder = Decoder(**ddconfig)
        self.vocoder_path = vocoder_path

        self.quant_conv = torch.nn.Conv2d(2 * ddconfig["z_channels"], 2 * embed_dim, 1)
        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)

        self.vocoder = VocoderBigVGAN(self.vocoder_path, device="cuda")
        self.reload_from_ckpt = reload_from_ckpt

        if self.reload_from_ckpt is not None:
            checkpoint = torch.load(self.reload_from_ckpt)

            load_todo_keys = {}

            pretrained_state_dict = checkpoint
            current_state_dict = self.state_dict()
            for key in current_state_dict:
                if (
                    key in pretrained_state_dict.keys()
                    and pretrained_state_dict[key].size()
                    == current_state_dict[key].size()
                ):
                    load_todo_keys[key] = pretrained_state_dict[key]
                else:
                    logging.info("Key %s mismatch during loading, seems fine" % key)

            self.load_state_dict(load_todo_keys, strict=False)
            self.reloaded = True
        else:
            logging.info("Train from scratch")

    def encode(self, x):
        h = self.encoder(x)
        moments = self.quant_conv(h)
        posterior = DiagonalGaussianDistribution(moments)

        return posterior

    def decode(self, z):
        z = self.post_quant_conv(z)
        fbk = self.decoder(z)

        return fbk

    def forward(self, input, sample_posterior=True):
        posterior = self.encode(input)
        z = posterior.sample()
        dec = self.decode(z)

        return dec, posterior