from typing import Dict, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers import AutoencoderKL
from diffusers.configuration_utils import register_to_config
from diffusers.models.autoencoders.vae import DecoderOutput, DiagonalGaussianDistribution
from diffusers.models.modeling_outputs import AutoencoderKLOutput
from diffusers.utils.accelerate_utils import apply_forward_hook


class AutoencoderKLNextStep(AutoencoderKL):
    @register_to_config
    def __init__(
        self,
        in_channels: int = 3,
        out_channels: int = 3,
        down_block_types: Tuple[str] = ("DownEncoderBlock2D",),
        up_block_types: Tuple[str] = ("UpDecoderBlock2D",),
        block_out_channels: Tuple[int] = (64,),
        layers_per_block: int = 1,
        act_fn: str = "silu",
        latent_channels: int = 4,
        norm_num_groups: int = 32,
        sample_size: int = 32,
        scaling_factor: float = 0.18215,
        shift_factor: Optional[float] = None,
        latents_mean: Optional[Tuple[float]] = None,
        latents_std: Optional[Tuple[float]] = None,
        force_upcast: bool = True,
        use_quant_conv: bool = True,
        use_post_quant_conv: bool = True,
        mid_block_add_attention: bool = True,
        deterministic: bool = False,
        normalize_latents: bool = False,
        patch_size: Optional[int] = None,
    ):
        super().__init__(
            in_channels=in_channels,
            out_channels=out_channels,
            down_block_types=down_block_types,
            up_block_types=up_block_types,
            block_out_channels=block_out_channels,
            layers_per_block=layers_per_block,
            act_fn=act_fn,
            latent_channels=latent_channels,
            norm_num_groups=norm_num_groups,
            sample_size=sample_size,
            scaling_factor=scaling_factor,
            shift_factor=shift_factor,
            latents_mean=latents_mean,
            latents_std=latents_std,
            force_upcast=force_upcast,
            use_quant_conv=use_quant_conv,
            use_post_quant_conv=use_post_quant_conv,
            mid_block_add_attention=mid_block_add_attention,
        )
        self.deterministic = deterministic
        self.normalize_latents = normalize_latents
        self.patch_size = patch_size

    def patchify(self, x: torch.Tensor) -> torch.Tensor:
        b, c, h, w = x.shape
        p = self.patch_size
        h_, w_ = h // p, w // p

        x = x.reshape(b, c, h_, p, w_, p)
        x = torch.einsum("bchpwq->bcpqhw", x)
        x = x.reshape(b, c * p ** 2, h_,  w_)
        return x

    def unpatchify(self, x: torch.Tensor) -> torch.Tensor:
        b, _, h_, w_ = x.shape
        p = self.patch_size
        c = x.shape[1] // (p ** 2)

        x = x.reshape(b, c, p, p, h_, w_)
        x = torch.einsum("bcpqhw->bchpwq", x)
        x = x.reshape(b, c, h_ * p, w_ * p)
        return x

    @apply_forward_hook
    def encode(
        self, x: torch.Tensor, return_dict: bool = True
    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
        if self.use_slicing and x.shape[0] > 1:
            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
            h = torch.cat(encoded_slices)
        else:
            h = self._encode(x)

        mean, logvar = torch.chunk(h, 2, dim=1)
        if self.patch_size is not None:
            mean = self.patchify(mean)
        if self.normalize_latents:
            mean = mean.permute(0, 2, 3, 1)
            mean = F.layer_norm(mean, mean.shape[-1:], eps=1e-6)
            mean = mean.permute(0, 3, 1, 2)
        if self.patch_size is not None:
            mean = self.unpatchify(mean)
        h = torch.cat([mean, logvar], dim=1).contiguous()
        posterior = DiagonalGaussianDistribution(h, deterministic=self.deterministic)

        if not return_dict:
            return (posterior,)

        return AutoencoderKLOutput(latent_dist=posterior)

    def forward(
        self,
        sample: torch.Tensor,
        sample_posterior: bool = False,
        return_dict: bool = True,
        generator: Optional[torch.Generator] = None,
        noise_strength: float = 0.0,
    ) -> Union[DecoderOutput, torch.Tensor]:
        x = sample
        posterior = self.encode(x).latent_dist
        if sample_posterior:
            z = posterior.sample(generator=generator)
        else:
            z = posterior.mode()
        if noise_strength > 0.0:
            p = torch.distributions.Uniform(0, noise_strength)
            z = z + p.sample((z.shape[0],)).reshape(-1, 1, 1, 1).to(z.device) * randn_tensor(
                z.shape, device=z.device, dtype=z.dtype
            )
        dec = self.decode(z).sample

        if not return_dict:
            return (dec,)

        return DecoderOutput(sample=dec)