import os
import inspect
import imageio
import numpy as np
from typing import Union
from omegaconf import OmegaConf

import torch
import torchvision
from safetensors import safe_open

import transformers
transformers.logging.set_verbosity_error()
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers import AutoencoderKL, DDIMScheduler
from diffusers.utils.import_utils import is_xformers_available
from animatediff.models.unet import UNet3DConditionModel
from animatediff.pipelines.pipeline_animation import AnimationPipeline

from tqdm import tqdm
from einops import rearrange
from collections import defaultdict
from .convert_from_ckpt import convert_ldm_unet_checkpoint, convert_ldm_clip_checkpoint, convert_ldm_vae_checkpoint


def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8):
    videos = rearrange(videos, "b c t h w -> t b c h w")
    outputs = []
    for x in videos:
        x = torchvision.utils.make_grid(x, nrow=n_rows)
        x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
        if rescale:
            x = (x + 1.0) / 2.0  # -1,1 -> 0,1
        x = (x * 255).numpy().astype(np.uint8)
        outputs.append(x)

    os.makedirs(os.path.dirname(path), exist_ok=True)
    imageio.mimsave(path, outputs, fps=fps)


# DDIM Inversion
@torch.no_grad()
def init_prompt(prompt, pipeline):
    uncond_input = pipeline.tokenizer([""], padding="max_length", max_length=pipeline.tokenizer.model_max_length, return_tensors="pt")
    uncond_embeddings = pipeline.text_encoder(uncond_input.input_ids.to(pipeline.device))[0]
    text_input = pipeline.tokenizer(
        [prompt],
        padding="max_length",
        max_length=pipeline.tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt",
    )
    text_embeddings = pipeline.text_encoder(text_input.input_ids.to(pipeline.device))[0]
    context = torch.cat([uncond_embeddings, text_embeddings])

    return context


def next_step(model_output: Union[torch.FloatTensor, np.ndarray], timestep: int, sample: Union[torch.FloatTensor, np.ndarray], ddim_scheduler):
    timestep, next_timestep = min(timestep - ddim_scheduler.config.num_train_timesteps // ddim_scheduler.num_inference_steps, 999), timestep
    alpha_prod_t = ddim_scheduler.alphas_cumprod[timestep] if timestep >= 0 else ddim_scheduler.final_alpha_cumprod
    alpha_prod_t_next = ddim_scheduler.alphas_cumprod[next_timestep]
    beta_prod_t = 1 - alpha_prod_t
    next_original_sample = (sample - beta_prod_t**0.5 * model_output) / alpha_prod_t**0.5
    next_sample_direction = (1 - alpha_prod_t_next) ** 0.5 * model_output
    next_sample = alpha_prod_t_next**0.5 * next_original_sample + next_sample_direction
    return next_sample


def get_noise_pred_single(latents, t, context, unet):
    noise_pred = unet(latents, t, encoder_hidden_states=context)["sample"]
    return noise_pred


@torch.no_grad()
def ddim_loop(pipeline, ddim_scheduler, latent, num_inv_steps, prompt):
    context = init_prompt(prompt, pipeline)
    uncond_embeddings, cond_embeddings = context.chunk(2)
    all_latent = [latent]
    latent = latent.clone().detach()
    for i in tqdm(range(num_inv_steps)):
        t = ddim_scheduler.timesteps[len(ddim_scheduler.timesteps) - i - 1]
        noise_pred = get_noise_pred_single(latent, t, cond_embeddings, pipeline.unet)
        latent = next_step(noise_pred, t, latent, ddim_scheduler)
        all_latent.append(latent)
    return all_latent


@torch.no_grad()
def ddim_inversion(pipeline, ddim_scheduler, video_latent, num_inv_steps, prompt=""):
    ddim_latents = ddim_loop(pipeline, ddim_scheduler, video_latent, num_inv_steps, prompt)
    return ddim_latents


def load_ckpt(model_path, device):
    if model_path.endswith(".ckpt"):
        state_dict = torch.load(model_path, map_location=device)
        state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
    elif model_path.endswith(".safetensors"):
        state_dict = {}
        with safe_open(model_path, framework="pt", device=device) as f:
            for key in f.keys():
                state_dict[key] = f.get_tensor(key)
    else:
        raise ValueError(f"Unknown model type {model_path}, must be either .ckpt or .safetensors, otherwise convert it first")
    return state_dict


def init_pipeline(pretrained_model_path, inference_config, device, dtype):
    tokenizer    = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer", torch_dtype=dtype)
    text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder", torch_dtype=dtype)
    vae          = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae", torch_dtype=dtype)
    unet         = UNet3DConditionModel.from_pretrained_2d(
        pretrained_model_path,
        subfolder="unet",
        unet_additional_kwargs=OmegaConf.to_container(inference_config.Model.unet_additional_kwargs),
        device=device,
        torch_dtype=dtype,
    )

    if is_xformers_available():
        unet.enable_xformers_memory_efficient_attention()
    else:
        assert False, "xformers is not available"

    pipeline = AnimationPipeline(
        vae=vae,
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        unet=unet,
        scheduler=DDIMScheduler(
            **OmegaConf.to_container(inference_config.Model.noise_scheduler_kwargs), steps_offset=1, clip_sample=False
        ),
    )


    return pipeline

def reload_motion_module(pipeline, motion_module, device):
    *_, func_args = inspect.getargvalues(inspect.currentframe())
    func_args = dict(func_args)

    motion_module_state_dict = {}
    with safe_open(motion_module, framework="pt", device=device) as f:
        for key in f.keys():
            motion_module_state_dict[key] = f.get_tensor(key)
    if "global_step" in motion_module_state_dict:
        func_args.update({"global_step": motion_module_state_dict["global_step"]})
    missing, unexpected = pipeline.unet.load_state_dict(motion_module_state_dict, strict=False)
    assert len(unexpected) == 0, f"Unexpected keys in motion module: {unexpected}"
    del motion_module_state_dict
    return pipeline

def load_base_model(pipeline, model_dir, base_model, device, dtype):
    # load base model
    base_model = os.path.join(model_dir, base_model)
    base_state_dict = load_ckpt(base_model, device)
    # vae
    converted_vae_checkpoint = convert_ldm_vae_checkpoint(base_state_dict, pipeline.vae.config)
    try:
        pipeline.vae.load_state_dict(converted_vae_checkpoint)
    except:
        converted_vae_checkpoint = {
            key.replace(".query.", ".to_q.").replace(".key.", ".to_k.").replace(".value.", ".to_v.").replace(".proj_attn.", ".to_out.0."): value
            for key, value in converted_vae_checkpoint.items()
        }
        pipeline.vae.load_state_dict(converted_vae_checkpoint)

    # unet
    converted_unet_checkpoint = convert_ldm_unet_checkpoint(base_state_dict, pipeline.unet.config)
    pipeline.unet.load_state_dict(converted_unet_checkpoint, strict=False)
    # text_model
    pipeline.text_encoder = convert_ldm_clip_checkpoint(base_state_dict)
    del base_state_dict

    return pipeline


def apply_motion_lora(pipeline, model_dir, motion_lora_config, device, dtype):
    motion_lora_path, motion_lora_scale = motion_lora_config
    print(f"Loading motion lora from {motion_lora_path} with scale {motion_lora_scale}")
    motion_lora_path = os.path.join(model_dir, motion_lora_path)
    state_dict = load_ckpt(motion_lora_path, device)
    # directly update weight in diffusers model
    for key in state_dict:
        # only process lora down key
        if "up." in key:
            continue

        up_key = key.replace(".down.", ".up.")
        model_key = key.replace("processor.", "").replace("_lora", "").replace("down.", "").replace("up.", "")
        model_key = model_key.replace("to_out.", "to_out.0.")
        layer_infos = model_key.split(".")[:-1]

        curr_layer = pipeline.unet
        while len(layer_infos) > 0:
            temp_name = layer_infos.pop(0)
            curr_layer = curr_layer.__getattr__(temp_name)

        weight_down = state_dict[key].to(dtype)
        weight_up = state_dict[up_key].to(dtype)
        curr_layer.weight.data += motion_lora_scale * torch.mm(weight_up, weight_down).to(curr_layer.weight.data.device)
    del state_dict
    return pipeline


def apply_lora(pipeline, model_dir, lora_configs, device, dtype):
    for lora_name in lora_configs:
        lora_path, lora_scale = lora_configs[lora_name]
        print(f"Loading {lora_name} from {lora_path} with scale {lora_scale}")
        lora_path = os.path.join(model_dir, lora_path)
        state_dict = load_ckpt(lora_path, device)

        LORA_PREFIX_UNET = "lora_unet"
        LORA_PREFIX_TEXT_ENCODER = "lora_te"

        updates = defaultdict(dict)
        for key, value in state_dict.items():
            # it is suggested to print out the key, it usually will be something like below
            # "lora_te_text_model_encoder_layers_0_self_attn_k_proj.lora_down.weight"

            layer, elem = key.split(".", 1)
            updates[layer][elem] = value

        # directly update weight in diffusers model
        for layer, elems in updates.items():
            if "text" in layer:
                layer_infos = layer.split(LORA_PREFIX_TEXT_ENCODER + "_")[-1].split("_")
                curr_layer = pipeline.text_encoder
            else:
                layer_infos = layer.split(LORA_PREFIX_UNET + "_")[-1].split("_")
                curr_layer = pipeline.unet

            # find the target layer
            temp_name = layer_infos.pop(0)
            while len(layer_infos) > -1:
                try:
                    curr_layer = curr_layer.__getattr__(temp_name)
                    if len(layer_infos) > 0:
                        temp_name = layer_infos.pop(0)
                    elif len(layer_infos) == 0:
                        break
                except Exception:
                    if len(temp_name) > 0:
                        temp_name += "_" + layer_infos.pop(0)
                    else:
                        temp_name = layer_infos.pop(0)

            # get elements for this layer
            weight_up = elems["lora_up.weight"].to(dtype)
            weight_down = elems["lora_down.weight"].to(dtype)
            alpha = elems["alpha"]
            if alpha:
                alpha = alpha.item() / weight_up.shape[1]
            else:
                alpha = 1.0

            # update weight
            if len(weight_up.shape) == 4:
                curr_layer.weight.data += (
                    lora_scale * alpha * torch.mm(weight_up.squeeze(3).squeeze(2), weight_down.squeeze(3).squeeze(2)).unsqueeze(2).unsqueeze(3)
                )
            else:
                curr_layer.weight.data += lora_scale * alpha * torch.mm(weight_up, weight_down)
        del state_dict, updates

    return pipeline
