import argparse
from typing import Dict

import torch
import torch.nn as nn

from diffusers import SparseControlNetModel


KEYS_RENAME_MAPPING = {
    ".attention_blocks.0": ".attn1",
    ".attention_blocks.1": ".attn2",
    ".attn1.pos_encoder": ".pos_embed",
    ".ff_norm": ".norm3",
    ".norms.0": ".norm1",
    ".norms.1": ".norm2",
    ".temporal_transformer": "",
}


def convert(original_state_dict: Dict[str, nn.Module]) -> Dict[str, nn.Module]:
    converted_state_dict = {}

    for key in list(original_state_dict.keys()):
        renamed_key = key
        for new_name, old_name in KEYS_RENAME_MAPPING.items():
            renamed_key = renamed_key.replace(new_name, old_name)
        converted_state_dict[renamed_key] = original_state_dict.pop(key)

    return converted_state_dict


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt_path", type=str, required=True, help="Path to checkpoint")
    parser.add_argument("--output_path", type=str, required=True, help="Path to output directory")
    parser.add_argument(
        "--max_motion_seq_length",
        type=int,
        default=32,
        help="Max motion sequence length supported by the motion adapter",
    )
    parser.add_argument(
        "--conditioning_channels", type=int, default=4, help="Number of channels in conditioning input to controlnet"
    )
    parser.add_argument(
        "--use_simplified_condition_embedding",
        action="store_true",
        default=False,
        help="Whether or not to use simplified condition embedding. When `conditioning_channels==4` i.e. latent inputs, set this to `True`. When `conditioning_channels==3` i.e. image inputs, set this to `False`",
    )
    parser.add_argument(
        "--save_fp16",
        action="store_true",
        default=False,
        help="Whether or not to save model in fp16 precision along with fp32",
    )
    parser.add_argument(
        "--push_to_hub", action="store_true", default=False, help="Whether or not to push saved model to the HF hub"
    )
    return parser.parse_args()


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

    state_dict = torch.load(args.ckpt_path, map_location="cpu")
    if "state_dict" in state_dict.keys():
        state_dict: dict = state_dict["state_dict"]

    controlnet = SparseControlNetModel(
        conditioning_channels=args.conditioning_channels,
        motion_max_seq_length=args.max_motion_seq_length,
        use_simplified_condition_embedding=args.use_simplified_condition_embedding,
    )

    state_dict = convert(state_dict)
    controlnet.load_state_dict(state_dict, strict=True)

    controlnet.save_pretrained(args.output_path, push_to_hub=args.push_to_hub)
    if args.save_fp16:
        controlnet = controlnet.to(dtype=torch.float16)
        controlnet.save_pretrained(args.output_path, variant="fp16", push_to_hub=args.push_to_hub)