import modules.util.convert.convert_diffusers_to_ckpt_util as util

import torch
from torch import Tensor


def __swap_chunks(tensor: Tensor) -> Tensor:
    chunk_0, chunk_1 = tensor.chunk(2, dim=0)
    return torch.cat([chunk_1, chunk_0], dim=0)

def __map_double_transformer_block(in_states: dict, out_prefix: str, in_prefix: str) -> dict:
    out_states = {}

    out_states[util.combine(out_prefix, "img_attn.qkv.weight")] = torch.cat([
        in_states[util.combine(in_prefix, "attn.to_q.weight")],
        in_states[util.combine(in_prefix, "attn.to_k.weight")],
        in_states[util.combine(in_prefix, "attn.to_v.weight")],
    ], 0)

    out_states[util.combine(out_prefix, "img_attn.qkv.bias")] = torch.cat([
        in_states[util.combine(in_prefix, "attn.to_q.bias")],
        in_states[util.combine(in_prefix, "attn.to_k.bias")],
        in_states[util.combine(in_prefix, "attn.to_v.bias")],
    ], 0)

    out_states[util.combine(out_prefix, "txt_attn.qkv.weight")] = torch.cat([
        in_states[util.combine(in_prefix, "attn.add_q_proj.weight")],
        in_states[util.combine(in_prefix, "attn.add_k_proj.weight")],
        in_states[util.combine(in_prefix, "attn.add_v_proj.weight")],
    ], 0)

    out_states[util.combine(out_prefix, "txt_attn.qkv.bias")] = torch.cat([
        in_states[util.combine(in_prefix, "attn.add_q_proj.bias")],
        in_states[util.combine(in_prefix, "attn.add_k_proj.bias")],
        in_states[util.combine(in_prefix, "attn.add_v_proj.bias")],
    ], 0)

    out_states[util.combine(out_prefix, "img_attn.norm.key_norm.scale")] = in_states[util.combine(in_prefix, "attn.norm_k.weight")]
    out_states[util.combine(out_prefix, "img_attn.norm.query_norm.scale")] = in_states[util.combine(in_prefix, "attn.norm_q.weight")]
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "img_attn.proj"), util.combine(in_prefix, "attn.to_out.0"))
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "img_mlp.0"), util.combine(in_prefix, "ff.net.0.proj"))
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "img_mlp.2"), util.combine(in_prefix, "ff.net.2"))

    out_states[util.combine(out_prefix, "txt_attn.norm.key_norm.scale")] = in_states[util.combine(in_prefix, "attn.norm_added_k.weight")]
    out_states[util.combine(out_prefix, "txt_attn.norm.query_norm.scale")] = in_states[util.combine(in_prefix, "attn.norm_added_q.weight")]
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "txt_attn.proj"), util.combine(in_prefix, "attn.to_add_out"))
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "txt_mlp.0"), util.combine(in_prefix, "ff_context.net.0.proj"))
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "txt_mlp.2"), util.combine(in_prefix, "ff_context.net.2"))

    return out_states


def __map_single_transformer_block(in_states: dict, out_prefix: str, in_prefix: str) -> dict:
    out_states = {}

    out_states[util.combine(out_prefix, "linear1.weight")] = torch.cat([
        in_states[util.combine(in_prefix, "attn.to_q.weight")],
        in_states[util.combine(in_prefix, "attn.to_k.weight")],
        in_states[util.combine(in_prefix, "attn.to_v.weight")],
        in_states[util.combine(in_prefix, "proj_mlp.weight")],
    ], 0)

    out_states[util.combine(out_prefix, "linear1.bias")] = torch.cat([
        in_states[util.combine(in_prefix, "attn.to_q.bias")],
        in_states[util.combine(in_prefix, "attn.to_k.bias")],
        in_states[util.combine(in_prefix, "attn.to_v.bias")],
        in_states[util.combine(in_prefix, "proj_mlp.bias")],
    ], 0)

    out_states[util.combine(out_prefix, "norm.key_norm.scale")] = in_states[util.combine(in_prefix, "attn.norm_k.weight")]
    out_states[util.combine(out_prefix, "norm.query_norm.scale")] = in_states[util.combine(in_prefix, "attn.norm_q.weight")]
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "linear2"), util.combine(in_prefix, "proj_out"))

    return out_states

def __map_distilled_guidance_layer_layers(in_states: dict, out_prefix: str, in_prefix: str) -> dict:
    out_states = {}

    out_states |= util.map_wb(in_states, util.combine(out_prefix, "in_layer"), util.combine(in_prefix, "linear_1"))
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "out_layer"), util.combine(in_prefix, "linear_2"))

    return out_states

def __map_distilled_guidance_layer_norms(in_states: dict, out_prefix: str, in_prefix: str) -> dict:
    out_states = {}

    out_states[util.combine(out_prefix, "scale")] = in_states[util.combine(in_prefix, "weight")]

    return out_states


def __map_transformer(in_states: dict, out_prefix: str, in_prefix: str) -> dict:
    out_states = {}

    out_states |= util.map_wb(in_states, util.combine(out_prefix, "txt_in"), util.combine(in_prefix, "context_embedder"))
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "final_layer.linear"), util.combine(in_prefix, "proj_out"))
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "img_in"), util.combine(in_prefix, "x_embedder"))
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "distilled_guidance_layer.in_proj"), util.combine(in_prefix, "distilled_guidance_layer.in_proj"))
    out_states |= util.map_wb(in_states, util.combine(out_prefix, "distilled_guidance_layer.out_proj"), util.combine(in_prefix, "distilled_guidance_layer.out_proj"))

    i = 0
    while any(key.startswith(util.combine(in_prefix, f"transformer_blocks.{i}")) for key in in_states):
        out_states |= __map_double_transformer_block(in_states, util.combine(out_prefix, f"double_blocks.{i}"), util.combine(in_prefix, f"transformer_blocks.{i}"))
        i += 1

    i = 0
    while any(key.startswith(util.combine(in_prefix, f"single_transformer_blocks.{i}")) for key in in_states):
        out_states |= __map_single_transformer_block(in_states, util.combine(out_prefix, f"single_blocks.{i}"), util.combine(in_prefix, f"single_transformer_blocks.{i}"))
        i += 1

    i = 0
    while any(key.startswith(util.combine(in_prefix, f"distilled_guidance_layer.layers.{i}")) for key in in_states):
        out_states |= __map_distilled_guidance_layer_layers(in_states, util.combine(out_prefix, f"distilled_guidance_layer.layers.{i}"), util.combine(in_prefix, f"distilled_guidance_layer.layers.{i}"))
        i += 1

    i = 0
    while any(key.startswith(util.combine(in_prefix, f"distilled_guidance_layer.norms.{i}")) for key in in_states):
        out_states |= __map_distilled_guidance_layer_norms(in_states, util.combine(out_prefix, f"distilled_guidance_layer.norms.{i}"), util.combine(in_prefix, f"distilled_guidance_layer.norms.{i}"))
        i += 1

    return out_states


def convert_chroma_diffusers_to_ckpt(
        transformer_state_dict: dict,
) -> dict:
    state_dict = {}

    state_dict |= __map_transformer(transformer_state_dict, "", "")

    return state_dict
