import torch
from collections import namedtuple, OrderedDict
from safetensors import safe_open
from .attention_processor import init_attn_proc
from .ip_adapter import MultiIPAdapterImageProjection
from .resampler import Resampler
from transformers import (
    AutoModel, AutoImageProcessor,
    CLIPVisionModelWithProjection, CLIPImageProcessor)


def init_adapter_in_unet(
        unet,
        image_proj_model=None,
        pretrained_model_path_or_dict=None,
        adapter_tokens=64,
        embedding_dim=None,
        use_lcm=False,
        use_adaln=True,
    ):
        device = unet.device
        dtype = unet.dtype
        if image_proj_model is None:
            assert embedding_dim is not None, "embedding_dim must be provided if image_proj_model is None."
            image_proj_model = Resampler(
                embedding_dim=embedding_dim,
                output_dim=unet.config.cross_attention_dim,
                num_queries=adapter_tokens,
            )
        if pretrained_model_path_or_dict is not None:
            if not isinstance(pretrained_model_path_or_dict, dict):
                if pretrained_model_path_or_dict.endswith(".safetensors"):
                    state_dict = {"image_proj": {}, "ip_adapter": {}}
                    with safe_open(pretrained_model_path_or_dict, framework="pt", device=unet.device) as f:
                        for key in f.keys():
                            if key.startswith("image_proj."):
                                state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key)
                            elif key.startswith("ip_adapter."):
                                state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key)
                else:
                    state_dict = torch.load(pretrained_model_path_or_dict, map_location=unet.device)
            else:
                state_dict = pretrained_model_path_or_dict
            keys = list(state_dict.keys())
            if "image_proj" not in keys and "ip_adapter" not in keys:
                state_dict = revise_state_dict(state_dict)

        # Creat IP cross-attention in unet.
        attn_procs = init_attn_proc(unet, adapter_tokens, use_lcm, use_adaln)
        unet.set_attn_processor(attn_procs)

        # Load pretrinaed model if needed.
        if pretrained_model_path_or_dict is not None:
            if "ip_adapter" in state_dict.keys():
                adapter_modules = torch.nn.ModuleList(unet.attn_processors.values())
                missing, unexpected = adapter_modules.load_state_dict(state_dict["ip_adapter"], strict=False)
                for mk in missing:
                    if "ln" not in mk:
                        raise ValueError(f"Missing keys in adapter_modules: {missing}")
            if "image_proj" in state_dict.keys():
                image_proj_model.load_state_dict(state_dict["image_proj"])

        # Load image projectors into iterable ModuleList.
        image_projection_layers = []
        image_projection_layers.append(image_proj_model)
        unet.encoder_hid_proj = MultiIPAdapterImageProjection(image_projection_layers)

        # Adjust unet config to handle addtional ip hidden states.
        unet.config.encoder_hid_dim_type = "ip_image_proj"
        unet.to(dtype=dtype, device=device)


def load_adapter_to_pipe(
        pipe,
        pretrained_model_path_or_dict,
        image_encoder_or_path=None,
        feature_extractor_or_path=None,
        use_clip_encoder=False,
        adapter_tokens=64,
        use_lcm=False,
        use_adaln=True,
    ):

        if not isinstance(pretrained_model_path_or_dict, dict):
            if pretrained_model_path_or_dict.endswith(".safetensors"):
                state_dict = {"image_proj": {}, "ip_adapter": {}}
                with safe_open(pretrained_model_path_or_dict, framework="pt", device=pipe.device) as f:
                    for key in f.keys():
                        if key.startswith("image_proj."):
                            state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key)
                        elif key.startswith("ip_adapter."):
                            state_dict["ip_adapter"][key.replace("ip_adapter.", "")] = f.get_tensor(key)
            else:
                state_dict = torch.load(pretrained_model_path_or_dict, map_location=pipe.device)
        else:
            state_dict = pretrained_model_path_or_dict
        keys = list(state_dict.keys())
        if "image_proj" not in keys and "ip_adapter" not in keys:
            state_dict = revise_state_dict(state_dict)

        # load CLIP image encoder here if it has not been registered to the pipeline yet
        if image_encoder_or_path is not None:
            if isinstance(image_encoder_or_path, str):
                feature_extractor_or_path = image_encoder_or_path if feature_extractor_or_path is None else feature_extractor_or_path

                image_encoder_or_path = (
                    CLIPVisionModelWithProjection.from_pretrained(
                        image_encoder_or_path
                    ) if use_clip_encoder else
                    AutoModel.from_pretrained(image_encoder_or_path)
                )

        if feature_extractor_or_path is not None:
            if isinstance(feature_extractor_or_path, str):
                feature_extractor_or_path = (
                    CLIPImageProcessor() if use_clip_encoder else
                    AutoImageProcessor.from_pretrained(feature_extractor_or_path)
                )

        # create image encoder if it has not been registered to the pipeline yet
        if hasattr(pipe, "image_encoder") and getattr(pipe, "image_encoder", None) is None:
            image_encoder = image_encoder_or_path.to(pipe.device, dtype=pipe.dtype)
            pipe.register_modules(image_encoder=image_encoder)
        else:
            image_encoder = pipe.image_encoder

        # create feature extractor if it has not been registered to the pipeline yet
        if hasattr(pipe, "feature_extractor") and getattr(pipe, "feature_extractor", None) is None:
            feature_extractor = feature_extractor_or_path
            pipe.register_modules(feature_extractor=feature_extractor)
        else:
            feature_extractor = pipe.feature_extractor

        # load adapter into unet
        unet = getattr(pipe, pipe.unet_name) if not hasattr(pipe, "unet") else pipe.unet
        attn_procs = init_attn_proc(unet, adapter_tokens, use_lcm, use_adaln)
        unet.set_attn_processor(attn_procs)
        image_proj_model = Resampler(
            embedding_dim=image_encoder.config.hidden_size,
            output_dim=unet.config.cross_attention_dim,
            num_queries=adapter_tokens,
        )

        # Load pretrinaed model if needed.
        if "ip_adapter" in state_dict.keys():
            adapter_modules = torch.nn.ModuleList(unet.attn_processors.values())
            missing, unexpected = adapter_modules.load_state_dict(state_dict["ip_adapter"], strict=False)
            for mk in missing:
                if "ln" not in mk:
                    raise ValueError(f"Missing keys in adapter_modules: {missing}")
        if "image_proj" in state_dict.keys():
            image_proj_model.load_state_dict(state_dict["image_proj"])

        # convert IP-Adapter Image Projection layers to diffusers
        image_projection_layers = []
        image_projection_layers.append(image_proj_model)
        unet.encoder_hid_proj = MultiIPAdapterImageProjection(image_projection_layers)

        # Adjust unet config to handle addtional ip hidden states.
        unet.config.encoder_hid_dim_type = "ip_image_proj"
        unet.to(dtype=pipe.dtype, device=pipe.device)


def revise_state_dict(old_state_dict_or_path, map_location="cpu"):
    new_state_dict = OrderedDict()
    new_state_dict["image_proj"] = OrderedDict()
    new_state_dict["ip_adapter"] = OrderedDict()
    if isinstance(old_state_dict_or_path, str):
        old_state_dict = torch.load(old_state_dict_or_path, map_location=map_location)
    else:
        old_state_dict = old_state_dict_or_path
    for name, weight in old_state_dict.items():
        if name.startswith("image_proj_model."):
            new_state_dict["image_proj"][name[len("image_proj_model."):]] = weight
        elif name.startswith("adapter_modules."):
            new_state_dict["ip_adapter"][name[len("adapter_modules."):]] = weight
    return new_state_dict


# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_image
def encode_image(image_encoder, feature_extractor, image, device, num_images_per_prompt, output_hidden_states=None):
    dtype = next(image_encoder.parameters()).dtype

    if not isinstance(image, torch.Tensor):
        image = feature_extractor(image, return_tensors="pt").pixel_values

    image = image.to(device=device, dtype=dtype)
    if output_hidden_states:
        image_enc_hidden_states = image_encoder(image, output_hidden_states=True).hidden_states[-2]
        image_enc_hidden_states = image_enc_hidden_states.repeat_interleave(num_images_per_prompt, dim=0)
        return image_enc_hidden_states
    else:
        if isinstance(image_encoder, CLIPVisionModelWithProjection):
            # CLIP image encoder.
            image_embeds = image_encoder(image).image_embeds
        else:
            # DINO image encoder.
            image_embeds = image_encoder(image).last_hidden_state
        image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
        return image_embeds


def prepare_training_image_embeds(
    image_encoder, feature_extractor,
    ip_adapter_image, ip_adapter_image_embeds,
    device, drop_rate, output_hidden_state, idx_to_replace=None
):
    if ip_adapter_image_embeds is None:
        if not isinstance(ip_adapter_image, list):
            ip_adapter_image = [ip_adapter_image]

        # if len(ip_adapter_image) != len(unet.encoder_hid_proj.image_projection_layers):
        #     raise ValueError(
        #         f"`ip_adapter_image` must have same length as the number of IP Adapters. Got {len(ip_adapter_image)} images and {len(unet.encoder_hid_proj.image_projection_layers)} IP Adapters."
        #     )

        image_embeds = []
        for single_ip_adapter_image in ip_adapter_image:
            if idx_to_replace is None:
                idx_to_replace = torch.rand(len(single_ip_adapter_image)) < drop_rate
            zero_ip_adapter_image = torch.zeros_like(single_ip_adapter_image)
            single_ip_adapter_image[idx_to_replace] = zero_ip_adapter_image[idx_to_replace]
            single_image_embeds = encode_image(
                image_encoder, feature_extractor, single_ip_adapter_image, device, 1, output_hidden_state
            )
            single_image_embeds = torch.stack([single_image_embeds], dim=1) # FIXME

            image_embeds.append(single_image_embeds)
    else:
        repeat_dims = [1]
        image_embeds = []
        for single_image_embeds in ip_adapter_image_embeds:
            if do_classifier_free_guidance:
                single_negative_image_embeds, single_image_embeds = single_image_embeds.chunk(2)
                single_image_embeds = single_image_embeds.repeat(
                    num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
                )
                single_negative_image_embeds = single_negative_image_embeds.repeat(
                    num_images_per_prompt, *(repeat_dims * len(single_negative_image_embeds.shape[1:]))
                )
                single_image_embeds = torch.cat([single_negative_image_embeds, single_image_embeds])
            else:
                single_image_embeds = single_image_embeds.repeat(
                    num_images_per_prompt, *(repeat_dims * len(single_image_embeds.shape[1:]))
                )
            image_embeds.append(single_image_embeds)

    return image_embeds