from .svd_image_encoder import SVDImageEncoder
from transformers import CLIPImageProcessor
import torch


class IpAdapterXLCLIPImageEmbedder(SVDImageEncoder):
    def __init__(self):
        super().__init__(embed_dim=1664, encoder_intermediate_size=8192, projection_dim=1280, num_encoder_layers=48, num_heads=16, head_dim=104)
        self.image_processor = CLIPImageProcessor()

    def forward(self, image):
        pixel_values = self.image_processor(images=image, return_tensors="pt").pixel_values
        pixel_values = pixel_values.to(device=self.embeddings.class_embedding.device, dtype=self.embeddings.class_embedding.dtype)
        return super().forward(pixel_values)


class IpAdapterImageProjModel(torch.nn.Module):
    def __init__(self, cross_attention_dim=2048, clip_embeddings_dim=1280, clip_extra_context_tokens=4):
        super().__init__()
        self.cross_attention_dim = cross_attention_dim
        self.clip_extra_context_tokens = clip_extra_context_tokens
        self.proj = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
        self.norm = torch.nn.LayerNorm(cross_attention_dim)

    def forward(self, image_embeds):
        clip_extra_context_tokens = self.proj(image_embeds).reshape(-1, self.clip_extra_context_tokens, self.cross_attention_dim)
        clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
        return clip_extra_context_tokens


class IpAdapterModule(torch.nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.to_k_ip = torch.nn.Linear(input_dim, output_dim, bias=False)
        self.to_v_ip = torch.nn.Linear(input_dim, output_dim, bias=False)

    def forward(self, hidden_states):
        ip_k = self.to_k_ip(hidden_states)
        ip_v = self.to_v_ip(hidden_states)
        return ip_k, ip_v


class SDXLIpAdapter(torch.nn.Module):
    def __init__(self):
        super().__init__()
        shape_list = [(2048, 640)] * 4 + [(2048, 1280)] * 50 + [(2048, 640)] * 6 + [(2048, 1280)] * 10
        self.ipadapter_modules = torch.nn.ModuleList([IpAdapterModule(*shape) for shape in shape_list])
        self.image_proj = IpAdapterImageProjModel()
        self.set_full_adapter()

    def set_full_adapter(self):
        map_list = sum([
            [(7, i) for i in range(2)],
            [(10, i) for i in range(2)],
            [(15, i) for i in range(10)],
            [(18, i) for i in range(10)],
            [(25, i) for i in range(10)],
            [(28, i) for i in range(10)],
            [(31, i) for i in range(10)],
            [(35, i) for i in range(2)],
            [(38, i) for i in range(2)],
            [(41, i) for i in range(2)],
            [(21, i) for i in range(10)],
        ], [])
        self.call_block_id = {i: j for j, i in enumerate(map_list)}

    def set_less_adapter(self):
        map_list = sum([
            [(7, i) for i in range(2)],
            [(10, i) for i in range(2)],
            [(15, i) for i in range(10)],
            [(18, i) for i in range(10)],
            [(25, i) for i in range(10)],
            [(28, i) for i in range(10)],
            [(31, i) for i in range(10)],
            [(35, i) for i in range(2)],
            [(38, i) for i in range(2)],
            [(41, i) for i in range(2)],
            [(21, i) for i in range(10)],
        ], [])
        self.call_block_id = {i: j for j, i in enumerate(map_list) if j>=34 and j<44}

    def forward(self, hidden_states, scale=1.0):
        hidden_states = self.image_proj(hidden_states)
        hidden_states = hidden_states.view(1, -1, hidden_states.shape[-1])
        ip_kv_dict = {}
        for (block_id, transformer_id) in self.call_block_id:
            ipadapter_id = self.call_block_id[(block_id, transformer_id)]
            ip_k, ip_v = self.ipadapter_modules[ipadapter_id](hidden_states)
            if block_id not in ip_kv_dict:
                ip_kv_dict[block_id] = {}
            ip_kv_dict[block_id][transformer_id] = {
                "ip_k": ip_k,
                "ip_v": ip_v,
                "scale": scale
            }
        return ip_kv_dict

    @staticmethod
    def state_dict_converter():
        return SDXLIpAdapterStateDictConverter()


class SDXLIpAdapterStateDictConverter:
    def __init__(self):
        pass

    def from_diffusers(self, state_dict):
        state_dict_ = {}
        for name in state_dict["ip_adapter"]:
            names = name.split(".")
            layer_id = str(int(names[0]) // 2)
            name_ = ".".join(["ipadapter_modules"] + [layer_id] + names[1:])
            state_dict_[name_] = state_dict["ip_adapter"][name]
        for name in state_dict["image_proj"]:
            name_ = "image_proj." + name
            state_dict_[name_] = state_dict["image_proj"][name]
        return state_dict_
    
    def from_civitai(self, state_dict):
        return self.from_diffusers(state_dict)