import argparse
import os

import torch
from huggingface_hub import create_repo, upload_folder
from safetensors.torch import load_file, save_file


def convert_motion_module(original_state_dict):
    converted_state_dict = {}
    for k, v in original_state_dict.items():
        if "pos_encoder" in k:
            continue

        else:
            converted_state_dict[
                k.replace(".norms.0", ".norm1")
                .replace(".norms.1", ".norm2")
                .replace(".ff_norm", ".norm3")
                .replace(".attention_blocks.0", ".attn1")
                .replace(".attention_blocks.1", ".attn2")
                .replace(".temporal_transformer", "")
            ] = v

    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(
        "--push_to_hub",
        action="store_true",
        default=False,
        help="Whether to push the converted model to the HF or not",
    )

    return parser.parse_args()


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

    if args.ckpt_path.endswith(".safetensors"):
        state_dict = load_file(args.ckpt_path)
    else:
        state_dict = torch.load(args.ckpt_path, map_location="cpu")

    if "state_dict" in state_dict.keys():
        state_dict = state_dict["state_dict"]

    conv_state_dict = convert_motion_module(state_dict)

    # convert to new format
    output_dict = {}
    for module_name, params in conv_state_dict.items():
        if type(params) is not torch.Tensor:
            continue
        output_dict.update({f"unet.{module_name}": params})

    os.makedirs(args.output_path, exist_ok=True)

    filepath = os.path.join(args.output_path, "diffusion_pytorch_model.safetensors")
    save_file(output_dict, filepath)

    if args.push_to_hub:
        repo_id = create_repo(args.output_path, exist_ok=True).repo_id
        upload_folder(repo_id=repo_id, folder_path=args.output_path, repo_type="model")