import torch
import numpy as np
import torch.nn.functional as F
import torch.nn as nn

from torch import Tensor, nn, no_grad
from .autoencoders import OobleckDecoder, OobleckEncoder

from .transformer import ContinuousTransformer
LRELU_SLOPE = 0.1
padding_mode = "zeros"
sample_eps = 1e-6

def vae_sample(mean, scale):
    stdev = nn.functional.softplus(scale)
    var = stdev * stdev + sample_eps
    logvar = torch.log(var)
    latents = torch.randn_like(mean) * stdev + mean

    kl = (mean * mean + var - logvar - 1).sum(1).mean()
    
    return latents, kl


class EAR_VAE(nn.Module):

    def __init__(self, model_config: dict = None):
        super().__init__()

        if model_config is None:
            model_config = {
                "encoder": {
                    "config": {
                        "in_channels": 2,
                        "channels": 128,
                        "c_mults": [1, 2, 4, 8, 16],
                        "strides": [2, 4, 4, 4, 8],
                        "latent_dim": 128,
                        "use_snake": True
                    }
                },
                "decoder": {
                    "config": {
                        "out_channels": 2,
                        "channels": 128,
                        "c_mults": [1, 2, 4, 8, 16],
                        "strides": [2, 4, 4, 4, 8],
                        "latent_dim": 64,
                        "use_nearest_upsample": False,
                        "use_snake": True,
                        "final_tanh": False,
                    },
                },
                "latent_dim": 64,
                "downsampling_ratio": 1024,
                "io_channels": 2,
            }
        else:
            model_config = model_config

        if model_config.get("transformer") is not None:
            self.transformers = ContinuousTransformer(
                dim=model_config["decoder"]["config"]["latent_dim"],
                depth=model_config["transformer"]["depth"],
                **model_config["transformer"].get("config", {}),
            )
        else:
            self.transformers = None

        self.encoder = OobleckEncoder(**model_config["encoder"]["config"])
        self.decoder = OobleckDecoder(**model_config["decoder"]["config"])

    def forward(self, audio) -> Tensor:
        """
        audio: Input audio tensor [B,C,T]
        """
        status = self.encoder(audio)
        mean, scale = status.chunk(2, dim=1)
        z, kl = vae_sample(mean, scale)
        
        if self.transformers is not None:
            z = z.permute(0, 2, 1)
            z = self.transformers(z)
            z = z.permute(0, 2, 1)

        x = self.decoder(z)
        return x, kl

    def encode(self, audio, use_sample=True):
        x = self.encoder(audio)
        mean, scale = x.chunk(2, dim=1)
        if use_sample:
            z, _ = vae_sample(mean, scale)
        else:
            z = mean
        return z

    def decode(self, z):
        
        if self.transformers is not None:
            z = z.permute(0, 2, 1)
            z = self.transformers(z)
            z = z.permute(0, 2, 1)
            
        x = self.decoder(z)
        return x

    @no_grad()
    def inference(self, audio):
        z = self.encode(audio)
        recon_audio = self.decode(z)
        return recon_audio