from typing import Any, Dict
import torch
import argparse
from diffusers.loaders.lora_base import LoraBaseMixin
from diffusers.models.modeling_utils import load_state_dict


def get_state_dict(saved_dict: Dict[str, Any]) -> Dict[str, Any]:
    state_dict = saved_dict
    if "model" in saved_dict.keys():
        state_dict = state_dict["model"]
    if "module" in saved_dict.keys():
        state_dict = state_dict["module"]
    if "state_dict" in saved_dict.keys():
        state_dict = state_dict["state_dict"]
    return state_dict


LORA_KEYS_RENAME = {
    'attention.query_key_value.matrix_A.0': 'attn1.to_q.lora_A.weight',
    'attention.query_key_value.matrix_A.1': 'attn1.to_k.lora_A.weight',
    'attention.query_key_value.matrix_A.2': 'attn1.to_v.lora_A.weight',
    'attention.query_key_value.matrix_B.0': 'attn1.to_q.lora_B.weight',
    'attention.query_key_value.matrix_B.1': 'attn1.to_k.lora_B.weight',
    'attention.query_key_value.matrix_B.2': 'attn1.to_v.lora_B.weight',
    'attention.dense.matrix_A.0': 'attn1.to_out.0.lora_A.weight',
    'attention.dense.matrix_B.0': 'attn1.to_out.0.lora_B.weight',
}


PREFIX_KEY = "model.diffusion_model."
SAT_UNIT_KEY = "layers"
LORA_PREFIX_KEY = "transformer_blocks"


def export_lora_weight(ckpt_path, lora_save_directory):
    merge_original_state_dict = get_state_dict(torch.load(ckpt_path, map_location="cpu", mmap=True))

    lora_state_dict = {}
    for key in list(merge_original_state_dict.keys()):
        new_key = key[len(PREFIX_KEY) :]
        for special_key, lora_keys in LORA_KEYS_RENAME.items():
            if new_key.endswith(special_key):
                new_key = new_key.replace(special_key, lora_keys)
                new_key = new_key.replace(SAT_UNIT_KEY, LORA_PREFIX_KEY)

                lora_state_dict[new_key] = merge_original_state_dict[key]

    # final length should be 240
    if len(lora_state_dict) != 240:
        raise ValueError("lora_state_dict length is not 240")

    lora_state_dict.keys()

    LoraBaseMixin.write_lora_layers(
        state_dict=lora_state_dict,
        save_directory=lora_save_directory,
        is_main_process=True,
        weight_name=None,
        save_function=None,
        safe_serialization=True,
    )


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--sat_pt_path", type=str, required=True, help="Path to original sat transformer checkpoint"
    )
    parser.add_argument(
        "--lora_save_directory",
        type=str,
        required=True,
        help="Path where converted lora should be saved",
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = get_args()

    export_lora_weight(args.sat_pt_path, args.lora_save_directory)