from pathlib import Path

import torch
from diffusers.models import AutoencoderKL

from .autoencoder_kl_causal_3d import AutoencoderKLCausal3D
from .flux_vae import FluxAutoencoderKL
from ..constants import VAE_PATH
from ..utils.torch_utils import PRECISION_TO_TYPE


def transform_pytorch_ckpt_to_safetensors(vae_path):
    from safetensors.torch import save_file

    vae = AutoencoderKL.from_config(AutoencoderKL.load_config(vae_path))
    ckpt = torch.load(Path(vae_path) / "pytorch_model.pt", map_location=vae.device)
    if "state_dict" in ckpt:
        ckpt = ckpt["state_dict"]
    vae.load_state_dict(ckpt)

    save_file(vae.state_dict(), Path(vae_path) / "diffusion_pytorch_model.safetensors")
    return vae


def load_vae(vae_type,
             vae_precision=None,
             sample_size=None,
             vae_path=None,
             logger=None,
             device=None
             ):
    if vae_path is None:
        vae_path = VAE_PATH[vae_type]
    vae_compress_spec, vae_latent_channel, vae_name = vae_type.split("-")
    length = len(vae_compress_spec)

    if length == 2:
        if logger is not None:
            logger.info(f"Loading 2D VAE model ({vae_type}) from: {vae_path}")
        if vae_name == "flux":
            vae = FluxAutoencoderKL.from_pretrained(vae_path)
        else:
            vae = AutoencoderKL.from_pretrained(vae_path)

        # vae = transform_pytorch_ckpt_to_safetensors(vae_path)

        spatial_compression_ratio = 8
        time_compression_ratio = 1

    elif length == 3:
        if logger is not None:
            logger.info(f"Loading 3D VAE model ({vae_type}) from: {vae_path}")
        config = AutoencoderKLCausal3D.load_config(vae_path)
        if sample_size:
            vae = AutoencoderKLCausal3D.from_config(config, sample_size=sample_size)
        else:
            vae = AutoencoderKLCausal3D.from_config(config)
        ckpt = torch.load(Path(vae_path) / "pytorch_model.pt", map_location=vae.device)
        if "state_dict" in ckpt:
            ckpt = ckpt["state_dict"]
        # Internal checkpoints stored the VAE as a "vae." submodule; the public
        # HunyuanVideo VAE ships the keys unprefixed. Support both.
        if any(k.startswith("vae.") for k in ckpt):
            vae_ckpt = {k.replace("vae.", ""): v for k, v in ckpt.items() if k.startswith("vae.")}
        else:
            vae_ckpt = ckpt
        vae.load_state_dict(vae_ckpt)

        spatial_compression_ratio = vae.config.spatial_compression_ratio
        time_compression_ratio = vae.config.time_compression_ratio
    else:
        raise ValueError(f"Invalid VAE model: {vae_type}. Must be either 2D VAE in the format of '??-*' or "
                         f"3D VAE in the format of '???-*'.")

    if vae_precision is not None:
        vae = vae.to(dtype=PRECISION_TO_TYPE[vae_precision])

    vae.requires_grad_(False)

    if logger is not None:
        logger.info(f"VAE to dtype: {vae.dtype}")

    if device is not None:
        vae = vae.to(device)

    # Set vae to eval mode, even though it's dropout rate is 0.
    vae.eval()

    return vae, vae_path, spatial_compression_ratio, time_compression_ratio