import argparse
import os
import time
import traceback
from concurrent.futures import ThreadPoolExecutor, as_completed

from vae import get_vae

import tensorrt_llm
from tensorrt_llm._utils import release_gc
from tensorrt_llm.logger import logger
from tensorrt_llm.mapping import Mapping
from tensorrt_llm.models import STDiT3Model

PRETRAINED_STDIT_PATH = "hpcai-tech/OpenSora-STDiT-v3"


def pixel_size_to_latent_size(args):
    vae = get_vae(
        from_pretrained=args.vae_type,
        micro_frame_size=args.vae_micro_frame_size,
        micro_batch_size=args.vae_micro_batch_size,
    ).eval()
    spatial_patch_size = vae.spatial_vae.patch_size
    temporal_patch_size = vae.temporal_vae.patch_size
    vae_out_channels = vae.out_channels
    pixel_size = (args.num_frames, args.height, args.width)
    latent_size = vae.get_latent_size(pixel_size)
    return {
        'in_channels': vae_out_channels,
        'latent_size': latent_size,
        'spatial_patch_size': spatial_patch_size,
        'temporal_patch_size': temporal_patch_size,
    }


def size_str_to_list(repr):
    return [int(it) for it in repr.split('x')] if 'x' in repr else [int(repr)]


def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--timm_ckpt', type=str, default=None)
    parser.add_argument('--output_dir',
                        type=str,
                        default='tllm_checkpoint',
                        help='The path to save the TensorRT-LLM checkpoint')
    parser.add_argument('--caption_channels',
                        type=int,
                        default=4096,
                        help='The channel of input of caption embedder')
    parser.add_argument('--depth',
                        type=int,
                        default=28,
                        help='The number of STDiT blocks')
    parser.add_argument('--input_sq_size',
                        type=int,
                        default=512,
                        help='Base spatial position embedding size')
    parser.add_argument('--stdit_type',
                        type=str,
                        default="STDiT3",
                        choices=["STDiT3"])
    parser.add_argument('--stdit_patch_size',
                        type=str,
                        default='1x2x2',
                        help='The patch size of stdit for patchify')
    parser.add_argument('--width',
                        type=int,
                        default=640,
                        help='The width of image size')
    parser.add_argument('--height',
                        type=int,
                        default=360,
                        help='The height of image size')
    parser.add_argument('--num_frames',
                        type=int,
                        default=102,
                        help='The frames of generated video')
    parser.add_argument('--vae_type',
                        type=str,
                        default="hpcai-tech/OpenSora-VAE-v1.2",
                        choices=["hpcai-tech/OpenSora-VAE-v1.2"])
    parser.add_argument('--vae_micro_frame_size',
                        type=int,
                        default=17,
                        help='The micro_frame_size for vae')
    parser.add_argument('--vae_micro_batch_size',
                        type=int,
                        default=4,
                        help='The micro_batch_size for vae')
    parser.add_argument('--hidden_size',
                        type=int,
                        default=1152,
                        help='The hidden size of STDiT')
    parser.add_argument('--num_heads',
                        type=int,
                        default=16,
                        help='The number of heads of attention module')
    parser.add_argument(
        '--mlp_ratio',
        type=float,
        default=4.0,
        help=
        'The ratio of hidden size compared to input hidden size in MLP layer')
    parser.add_argument(
        '--class_dropout_prob',
        type=float,
        default=0.1,
        help='The probability to drop class token when training')
    parser.add_argument('--model_max_length',
                        type=int,
                        default=300,
                        help='The max number of tokens (default: 300)')
    parser.add_argument('--text_encoder_type',
                        type=str,
                        default="DeepFloyd/t5-v1_1-xxl",
                        choices=["DeepFloyd/t5-v1_1-xxl"])
    parser.add_argument('--learn_sigma',
                        type=bool,
                        default=True,
                        help='Whether the model learn sigma')
    parser.add_argument('--tp_size',
                        type=int,
                        default=1,
                        help='N-way tensor parallelism size')
    parser.add_argument('--cp_size',
                        type=int,
                        default=1,
                        help='Context parallelism size')
    parser.add_argument('--pp_size',
                        type=int,
                        default=1,
                        help='N-way pipeline parallelism size')
    parser.add_argument('--dtype',
                        type=str,
                        default='float16',
                        choices=['float32', 'bfloat16', 'float16'])
    parser.add_argument('--disable_qk_norm',
                        action='store_true',
                        help='Disable norm for qk in attention')
    parser.add_argument('--fp8',
                        action='store_true',
                        help='Whether use FP8 for layers')
    parser.add_argument(
        '--workers',
        type=int,
        default=1,
        help='The number of workers for converting checkpoint in parallel')
    parser.add_argument('--log_level', type=str, default='info')
    args = parser.parse_args()
    return args


def convert_and_save_model(args):
    # [NOTE] PP is not supported yet.
    world_size = args.tp_size * args.cp_size
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    def convert_and_save_rank(args, rank):
        mapping = Mapping(world_size=world_size,
                          rank=rank,
                          tp_size=args.tp_size,
                          cp_size=args.cp_size)
        # Process args
        runtime_config = {
            'architecture': "STDiT3",
            'checkpoint_path': os.path.abspath(args.timm_ckpt),
            'caption_channels': args.caption_channels,
            'num_hidden_layers': args.depth,
            'width': args.width,
            'height': args.height,
            'num_frames': args.num_frames,
            'hidden_size': args.hidden_size,
            'stdit_patch_size': size_str_to_list(args.stdit_patch_size),
            'input_sq_size': args.input_sq_size,
            'num_attention_heads': args.num_heads,
            'model_max_length': args.model_max_length,
            'mlp_ratio': args.mlp_ratio,
            'class_dropout_prob': args.class_dropout_prob,
            'learn_sigma': args.learn_sigma,
            'qk_norm': (not args.disable_qk_norm),
            'stdit_type': args.stdit_type,
            'vae_type': args.vae_type,
            'text_encoder_type': args.text_encoder_type,
        }
        runtime_config.update(pixel_size_to_latent_size(args))
        tik = time.time()
        stdit = STDiT3Model.from_pretrained(os.path.dirname(args.timm_ckpt),
                                            args.dtype,
                                            mapping=mapping,
                                            **runtime_config)
        stdit.save_checkpoint(args.output_dir, save_config=True)
        print(f'Total time of reading and converting: {time.time()-tik:.3f} s')
        tik = time.time()
        del stdit
        print(f'Total time of saving checkpoint: {time.time()-tik:.3f} s')

    execute(args.workers, [convert_and_save_rank] * world_size, args)
    release_gc()


def execute(workers, func, args):
    if workers == 1:
        for rank, f in enumerate(func):
            f(args, rank)
    else:
        with ThreadPoolExecutor(max_workers=workers) as p:
            futures = [p.submit(f, args, rank) for rank, f in enumerate(func)]
            exceptions = []
            for future in as_completed(futures):
                try:
                    future.result()
                except Exception as e:
                    traceback.print_exc()
                    exceptions.append(e)
            assert len(
                exceptions
            ) == 0, "Checkpoint conversion failed, please check error log."


def main():
    print(tensorrt_llm.__version__)
    args = parse_arguments()
    logger.set_level(args.log_level)

    assert args.pp_size == 1, "PP is not supported yet."

    tik = time.time()

    if args.timm_ckpt is None:
        print(
            f"No pretrained checkpoint provided, use default checkpoint from Huggingface instead."
        )
        args.timm_ckpt = "./pretrained_ckpt/model.safetensors"
        if not os.path.exists(args.timm_ckpt):
            from huggingface_hub import snapshot_download
            snapshot_download(repo_id=PRETRAINED_STDIT_PATH,
                              local_dir=os.path.dirname(args.timm_ckpt))

    convert_and_save_model(args)

    tok = time.time()
    t = time.strftime('%H:%M:%S', time.gmtime(tok - tik))
    print(f'Total time of converting checkpoints: {t}')


if __name__ == '__main__':
    main()
