# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# SPDX-License-Identifier: Apache-2.0

import os
import tempfile

import numpy as np
import torch
import torch.nn.functional as F
from diffusers import AutoencoderDC
from diffusers.models import AutoencoderKL
from diffusers.models.autoencoders import AutoencoderKLLTX2Video
from mmcv import Registry
from termcolor import colored
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    CLIPVisionModel,
    SiglipImageProcessor,
    SiglipVisionModel,
    T5EncoderModel,
    T5Tokenizer,
)
from transformers import logging as transformers_logging

from diffusion.data.datasets.video.sana_video_data import SanaZipDataset
from diffusion.model.dc_ae.efficientvit.ae_model_zoo import DCAE_HF, DCAEWithTemporal_HF
from diffusion.model.qwen.qwen_vl import QwenVLEmbedder
from diffusion.model.utils import set_fp32_attention, set_grad_checkpoint
from diffusion.model.wan2_2.vae import Wan2_2_VAE
from diffusion.model.wan.clip import CLIPModel
from diffusion.model.wan.vae import WanVAE

MODELS = Registry("models")

transformers_logging.set_verbosity_error()


def build_model(cfg, use_grad_checkpoint=False, use_fp32_attention=False, gc_step=1, **kwargs):
    if isinstance(cfg, str):
        cfg = dict(type=cfg)
    model = MODELS.build(cfg, default_args=kwargs)

    if use_grad_checkpoint:
        set_grad_checkpoint(model, gc_step=gc_step)
    if use_fp32_attention:
        set_fp32_attention(model)
    return model


def get_tokenizer_and_text_encoder(name="T5", device="cuda"):
    text_encoder_dict = {
        "T5": "DeepFloyd/t5-v1_1-xxl",
        "T5-small": "google/t5-v1_1-small",
        "T5-base": "google/t5-v1_1-base",
        "T5-large": "google/t5-v1_1-large",
        "T5-xl": "google/t5-v1_1-xl",
        "T5-xxl": "google/t5-v1_1-xxl",
        "gemma-2b": "google/gemma-2b",
        "gemma-2b-it": "google/gemma-2b-it",
        "gemma-2-2b": "google/gemma-2-2b",
        "gemma-2-2b-it": "Efficient-Large-Model/gemma-2-2b-it",
        "gemma-2-9b": "google/gemma-2-9b",
        "gemma-2-9b-it": "google/gemma-2-9b-it",
        "Qwen2-5-VL-3B-Instruct": "Qwen/Qwen2.5-VL-3B-Instruct",
        "Qwen2-5-VL-7B-Instruct": "Qwen/Qwen2.5-VL-7B-Instruct",
    }
    assert name in list(text_encoder_dict.keys()), f"not support this text encoder: {name}"
    if "T5" in name:
        tokenizer = T5Tokenizer.from_pretrained(text_encoder_dict[name])
        text_encoder = T5EncoderModel.from_pretrained(text_encoder_dict[name], torch_dtype=torch.float16).to(device)
    elif "gemma" in name:
        tokenizer = AutoTokenizer.from_pretrained(text_encoder_dict[name])
        tokenizer.padding_side = "right"
        text_encoder = (
            AutoModelForCausalLM.from_pretrained(text_encoder_dict[name], torch_dtype=torch.bfloat16)
            .get_decoder()
            .to(device)
        )
    elif "Qwen" in name:
        text_handler = QwenVLEmbedder(model_id=text_encoder_dict[name], device=device)
        return None, text_handler
    else:
        print("error load text encoder")
        exit()

    return tokenizer, text_encoder


def get_image_encoder(name, model_path, tokenizer_path=None, device="cuda", dtype=None, config=None):
    if name == "CLIP":
        image_encoder = CLIPModel(dtype, device, model_path, tokenizer_path)
    elif name == "flux-siglip":
        image_encoder = SiglipVisionModel.from_pretrained(model_path, subfolder="image_encoder", torch_dtype=dtype).to(
            device
        )
        image_processor = SiglipImageProcessor.from_pretrained(model_path, subfolder="feature_extractor")
        return image_encoder.eval().requires_grad_(False), image_processor
    else:
        raise ValueError(f"Unsupported image encoder: {name}")

    return image_encoder


@torch.no_grad()
def encode_image(name, image_encoder, images, device="cuda", image_processor=None, dtype=None):
    if image_encoder is None:
        return None
    if name == "CLIP":
        image_embeds = image_encoder.visual(images.to(image_encoder.device))
        return image_embeds.to(device, images.dtype)
    elif name == "flux-siglip":
        dtype = dtype or image_encoder.dtype
        images = (images + 1) / 2.0  # [-1, 1] -> [0, 1]
        images = image_processor(images=images.clamp(0, 1), return_tensors="pt", do_rescale=False).to(
            device=device, dtype=image_encoder.dtype
        )
        image_embeds = image_encoder(**images).last_hidden_state
        return image_embeds.to(dtype=dtype)
    else:
        raise ValueError(f"Unsupported image encoder: {name}")


def get_vae(name, model_path, device="cuda", dtype=None, config=None):
    if name == "sdxl" or name == "sd3":
        vae = AutoencoderKL.from_pretrained(model_path).to(device).to(torch.float16)
        if name == "sdxl":
            vae.config.shift_factor = 0
        return vae.to(dtype)
    elif ("dc-ae" in name and not "st-dc-ae" in name) or "dc-vae" in name:
        print(colored(f"[DC-AE] Loading model from {model_path}", attrs=["bold"]))
        dc_ae = DCAE_HF.from_pretrained(model_path).to(device).eval()
        return dc_ae.to(dtype)
    elif "st-dc-ae" in name:
        print(colored(f"[ST-DC-AE] Loading model from {model_path}", attrs=["bold"]))
        dc_ae = DCAEWithTemporal_HF.from_pretrained(model_path, model_name=name).to(device).eval()
        if config.scaling_factor is not None:
            dc_ae.cfg.scaling_factor = torch.tensor(config.scaling_factor).to(dtype).to(device)
        return dc_ae.to(dtype)

    elif "AutoencoderDC" in name:
        print(colored(f"[AutoencoderDC] Loading model from {model_path}", attrs=["bold"]))
        dc_ae = AutoencoderDC.from_pretrained(model_path).to(device).eval()
        return dc_ae.to(dtype)
    elif "WanVAE" in name:
        assert config is not None, "config.vae is required for WanVAE"
        print(colored(f"[WanVAE] Loading model from {model_path}", attrs=["bold"]))
        vae = WanVAE(
            z_dim=config.vae_latent_dim,
            vae_pth=config.vae_pretrained,
            dtype=dtype,
            device=device,
        )
        return vae
    elif "Wan2_2_VAE" in name:
        assert config is not None, "config.vae is required for Wan2_2_VAE"
        print(colored(f"[Wan2_2_VAE] Loading model from {model_path}", attrs=["bold"]))
        vae = Wan2_2_VAE(
            z_dim=config.vae_latent_dim,
            vae_pth=config.vae_pretrained,
            dtype=dtype,
            device=device,
        )
        return vae
    elif "LTX2VAE_diffusers_causal" in name:
        # Causal LTX-2 VAE (AutoencoderKLCausalLTX2Video) — encoder is causal (same
        # latent contract as the bidirectional sibling) and the decoder is also
        # causal, enabling chunk-by-chunk streaming decode with a persistent
        # per-layer feature cache. vae_pretrained should point at a directory with
        # config.json + diffusion_pytorch_model.safetensors (no "vae" subfolder).
        from diffusion.model.ltx2.causal_vae import AutoencoderKLCausalLTX2Video

        assert config is not None, "config.vae is required for LTX2VAE_diffusers_causal"
        print(colored(f"[LTX2VAE_diffusers_causal] Loading model from {config.vae_pretrained}", attrs=["bold"]))
        vae = AutoencoderKLCausalLTX2Video.from_pretrained(config.vae_pretrained, torch_dtype=dtype).to(device)
        vae.eval()
        return vae
    elif "LTX2VAE_chunk_tile" in name:
        # Public LTX-2 VAE loaded through the local causal wrapper so long V2V
        # inference can decode with temporal-only chunk tiling.
        from diffusion.model.ltx2.causal_vae import AutoencoderKLCausalLTX2Video

        assert config is not None, "config.vae is required for LTX2VAE_chunk_tile"
        print(colored(f"[LTX2VAE_chunk_tile] Loading model from {config.vae_pretrained}", attrs=["bold"]))
        vae = (
            AutoencoderKLCausalLTX2Video.from_pretrained(config.vae_pretrained, subfolder="vae", torch_dtype=dtype)
            .to(device)
            .eval()
        )
        vae.enable_tiling(tile_sample_min_num_frames=24, tile_sample_stride_num_frames=8)
        return vae
    elif "LTX2VAE_diffusers" in name:
        # Use diffusers AutoencoderKLLTX2Video for LTX2
        assert config is not None, "config.vae is required for LTX2VAE_diffusers"
        print(colored(f"[LTX2VAE_diffusers] Loading model from {config.vae_pretrained}", attrs=["bold"]))
        vae = (
            AutoencoderKLLTX2Video.from_pretrained(config.vae_pretrained, subfolder="vae", torch_dtype=dtype)
            .to(device)
            .eval()
        )
        return vae
    else:
        print("error load vae")
        exit()


@torch.no_grad()
def vae_encode(name, vae, images, sample_posterior=True, device="cuda", cache_key=None, if_cache=False, data_info=None):
    dtype = images.dtype
    if name == "sdxl" or name == "sd3":
        posterior = vae.encode(images.to(device)).latent_dist
        if sample_posterior:
            z = posterior.sample()
        else:
            z = posterior.mode()
        z = (z - vae.config.shift_factor) * vae.config.scaling_factor
    elif "dc-ae" in name and not "st-dc-ae" in name:
        ae = vae
        scaling_factor = ae.cfg.scaling_factor if ae.cfg.scaling_factor is not None else 0.41407
        z = ae.encode(images.to(device))
        z = z * scaling_factor
    elif "dc-vae" in name or "st-dc-ae" in name:
        ae = vae
        scaling_factor = ae.cfg.scaling_factor if ae.cfg.scaling_factor is not None else 0.493
        if isinstance(cache_key, list) and ae.cfg.cache_dir is not None:
            cache_file = [os.path.join(ae.cfg.cache_dir, f"{key}.npz") for key in cache_key]
        else:
            cache_file = None

        z = None
        try:
            if data_info is None:
                z = torch.stack([torch.from_numpy(np.load(cf)["z"]).to(device) for cf in cache_file], dim=0)
            elif data_info is not None and data_info.get("zip_file", None) is not None:
                z = []
                for zip_file, key, dataset_name in zip(
                    data_info["zip_file"], data_info["key"], data_info["dataset_name"]
                ):
                    vae_zip_file = os.path.join(ae.cfg.cache_dir, dataset_name, os.path.basename(zip_file))
                    if os.path.exists(vae_zip_file):
                        z_vae_cache = SanaZipDataset.open_zip_file(vae_zip_file)
                        with z_vae_cache.open(key + ".npz", "r") as f:
                            z.append(np.load(f)["z"] if "z" in np.load(f) else np.load(f))
                z = torch.from_numpy(np.stack(z)).to(device)
        except:
            z = None
        if z is None or len(z) == 0:
            z = ae.encode(images.to(device))
            if isinstance(scaling_factor, float):
                z = z * scaling_factor
            else:
                z = z * scaling_factor[None].view(1, -1, 1, 1, 1)
            if cache_file is not None and if_cache:
                tempdir = os.path.join(ae.cfg.cache_dir, ".tmp")
                os.makedirs(tempdir, exist_ok=True, mode=0o777)
                for i, cf in enumerate(cache_file):
                    if os.path.exists(cf):
                        continue
                    os.makedirs(os.path.dirname(cf), exist_ok=True)
                    with tempfile.NamedTemporaryFile(dir=tempdir) as f:
                        np.savez_compressed(f, z=z[i].float().cpu().numpy())  # bf16 not support for cpu
                        try:
                            os.link(f.name, cf)
                        except:
                            pass
    elif "AutoencoderDC" in name:
        ae = vae
        scaling_factor = ae.config.scaling_factor if ae.config.scaling_factor else 0.41407
        z = ae.encode(images.to(device))[0]
        z = z * scaling_factor
    elif "WanVAE" in name:
        ae = vae

        if isinstance(cache_key, list) and ae.cfg.cache_dir is not None:
            cache_file = [os.path.join(ae.cfg.cache_dir, f"{key}.npz") for key in cache_key]
        else:
            cache_file = None

        z = None
        try:
            if data_info is None:
                z = torch.stack([torch.from_numpy(np.load(cf)["z"]).to(device) for cf in cache_file], dim=0)
            elif data_info is not None and data_info.get("zip_file", None) is not None:
                z = []
                for zip_file, key, dataset_name in zip(
                    data_info["zip_file"], data_info["key"], data_info["dataset_name"]
                ):
                    vae_zip_file = os.path.join(ae.cfg.cache_dir, dataset_name, os.path.basename(zip_file))

                    if os.path.exists(vae_zip_file):
                        z_vae_cache = SanaZipDataset.open_zip_file(vae_zip_file)
                        with z_vae_cache.open(key + ".npz", "r") as f:
                            z.append(np.load(f)["z"] if "z" in np.load(f) else np.load(f))
                z = [torch.from_numpy(_z).to(device) for _z in z]
        except:
            z = None

        if z is None or len(z) == 0:
            z = ae.encode(images.to(device))
            if cache_file is not None and if_cache:
                tempdir = os.path.join(ae.cfg.cache_dir, ".tmp")
                os.makedirs(tempdir, exist_ok=True, mode=0o777)
                for i, cf in enumerate(cache_file):
                    if os.path.exists(cf):
                        continue
                    os.makedirs(os.path.dirname(cf), exist_ok=True)
                    with tempfile.NamedTemporaryFile(dir=tempdir) as f:
                        np.savez_compressed(f, z=z[i].float().cpu().numpy())  # bf16 not support for cpu
                        try:
                            os.link(f.name, cf)
                        except:
                            pass

        z = torch.stack(z, dim=0)
    elif "Wan2_2_VAE" in name:
        ae = vae
        z = ae.encode(images.to(device))
        z = torch.stack(z, dim=0)
    elif "LTX2VAE_chunk_tile" in name:
        posterior = vae.encode(images.to(device=vae.device, dtype=vae.dtype), causal=True).latent_dist
        z = posterior.mode()
        latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to(z.device, z.dtype)
        latents_std = vae.latents_std.view(1, -1, 1, 1, 1).to(z.device, z.dtype)
        z = (z - latents_mean) * vae.config.scaling_factor / latents_std
    elif "LTX2VAE_diffusers" in name:
        # Diffusers LTX-2 VAE uses full-video encode.
        posterior = vae.encode(images.to(device=vae.device, dtype=vae.dtype)).latent_dist
        z = posterior.mode()
        latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to(z.device, z.dtype)
        latents_std = vae.latents_std.view(1, -1, 1, 1, 1).to(z.device, z.dtype)
        z = (z - latents_mean) * vae.config.scaling_factor / latents_std
    else:
        print(f"{name} encode error")
        exit()
    return z.to(dtype)


def vae_decode(name, vae, latent):
    if name == "sdxl" or name == "sd3":
        latent = (latent.detach() / vae.config.scaling_factor) + vae.config.shift_factor
        samples = vae.decode(latent).sample
    elif "dc-ae" in name and not "st-dc-ae" in name:
        ae = vae
        vae_scale_factor = (
            2 ** (len(ae.config.encoder_block_out_channels) - 1)
            if hasattr(ae, "config") and ae.config is not None
            else 32
        )
        scaling_factor = ae.cfg.scaling_factor if ae.cfg.scaling_factor else 0.41407
        if latent.shape[-1] * vae_scale_factor > 4000 or latent.shape[-2] * vae_scale_factor > 4000:
            from patch_conv import convert_model

            ae = convert_model(ae, splits=4)
        samples = ae.decode(latent.detach() / scaling_factor)
    elif "dc-vae" in name or "st-dc-ae" in name:
        ae = vae
        scaling_factor = ae.cfg.scaling_factor if ae.cfg.scaling_factor is not None else 0.493
        if isinstance(scaling_factor, float):
            latent = latent.detach() / scaling_factor
        else:
            latent = latent.detach() / scaling_factor[None].view(1, -1, 1, 1, 1)
        samples = ae.decode(latent)
    elif "AutoencoderDC" in name:
        ae = vae
        scaling_factor = ae.config.scaling_factor if ae.config.scaling_factor else 0.41407
        try:
            samples = ae.decode(latent / scaling_factor, return_dict=False)[0]
        except torch.cuda.OutOfMemoryError as e:
            print("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
            ae.enable_tiling(tile_sample_min_height=1024, tile_sample_min_width=1024)
            samples = ae.decode(latent / scaling_factor, return_dict=False)[0]
    elif "WanVAE" in name:
        samples = vae.decode(latent)
    elif "Wan2_2_VAE" in name:
        samples = vae.decode(latent)
    elif "LTX2VAE_chunk_tile" in name:
        latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to(latent.device, latent.dtype)
        latents_std = vae.latents_std.view(1, -1, 1, 1, 1).to(latent.device, latent.dtype)
        latent = latent * latents_std / vae.config.scaling_factor + latents_mean
        latent = latent.to(vae.dtype)
        samples = vae.decode_chunk_tile(latent, temb=None, causal=False, return_dict=False)[0]
    elif "LTX2VAE_diffusers" in name:
        # Covers both bidirectional ("LTX2VAE_diffusers") and causal
        # ("LTX2VAE_diffusers_causal") variants — they share the same
        # latents_mean/std/scaling_factor and identical .decode() signature.
        # For the causal variant, .decode() internally chunks via
        # chunk_num_latent_frames (default 3) using the per-layer feature cache.
        # The chunked streaming path (CausalVaeStreamingDecoder) bypasses
        # this function and drives decode_with_cache directly.
        latents_mean = vae.latents_mean.view(1, -1, 1, 1, 1).to(latent.device, latent.dtype)
        latents_std = vae.latents_std.view(1, -1, 1, 1, 1).to(latent.device, latent.dtype)
        latent = latent * latents_std / vae.config.scaling_factor + latents_mean
        latent = latent.to(vae.dtype)
        samples = vae.decode(latent, temb=None, return_dict=False)[0]
    else:
        print(f"{name} decode error")
        exit()
    return samples