import copy
import logging
import os
import os.path as osp
import io

try:
    import deepspeed
except Exception as e:
    print(e)
    print("deepspeed is not installed!!!")

from os.path import join


# try:
#     from petrel_client.client import Client
# except:
#     Client = None

import torch
from torch.utils.data import ConcatDataset, DataLoader

from ..dataset.resample_concat_dataset import ResampleConcatDataset
from ..models.backbones.internvideo2.pos_embed import interpolate_pos_embed_internvideo2_new
# from ..models.backbones.bert.tokenization_bert import BertTokenizer
from transformers import BertTokenizer
from ..utils.optimizer import create_optimizer
from ..utils.scheduler import create_scheduler
from ..utils.distributed import get_rank

logger = logging.getLogger(__name__)


def get_media_types(datasources):
    """get the media types for for all the dataloaders.

    Args:
        datasources (List): List of dataloaders or datasets.

    Returns: List. The media_types.

    """
    if isinstance(datasources[0], DataLoader):
        datasets = [dataloader.dataset for dataloader in datasources]
    else:
        datasets = datasources
    media_types = [
        dataset.datasets[0].media_type
        if isinstance(dataset, ConcatDataset) or isinstance(dataset, ResampleConcatDataset)
        else dataset.media_type
        for dataset in datasets
    ]

    return media_types


def setup_model(
    config, model_cls, add_decoder=False, pretrain=True, device="cuda"
):
    # logger.info("Creating model")
    config = copy.deepcopy(config)

    if "bert" in config.model.text_encoder.name:
        # logger.info(f"Using BertTokenizer: {config.model.text_encoder.pretrained}!")
        tokenizer = BertTokenizer.from_pretrained("google-bert/bert-large-uncased")#config.model.text_encoder.pretrained, local_files_only=True)
        model = model_cls(config=config, tokenizer=tokenizer, is_pretrain=pretrain)
    else:
        model = model_cls(config=config, is_pretrain=pretrain)
        tokenizer = model.tokenizer
        logger.info(f"Using model.tokenizer: {tokenizer}!")

    # if config.get('compile_model', False):
    #     torch.set_float32_matmul_precision('high')
    #     model = torch.compile(model)

    model = model.to(device)
    model_without_ddp = model


    # if hasattr(config, "deepspeed") and config.deepspeed.enable:
    #     # We move this to the back
    #     optimizer_params = create_optimizer(config.optimizer, model, return_group=True)
    #     scheduler = None
    #     scaler = None
    # else:
    #     if config.distributed:
    #         model = torch.nn.parallel.DistributedDataParallel(
    #             model,
    #             device_ids=[config.gpu],
    #             find_unused_parameters=find_unused_parameters,  # `False` for image-only task
    #         )

    #     optimizer = create_optimizer(config.optimizer, model)
    #     scaler = torch.cuda.amp.GradScaler(enabled=config.use_half_precision)  # This is never used actually if we fixed bf16
    #     scheduler = create_scheduler(config.scheduler, optimizer)


    # start_epoch = 0
    # global_step = 0

    # auto resume the latest checkpoint
    # if config.get("auto_resume", False):
        # logger.info("Auto resuming")
        # model_latest = join(config.output_dir, "ckpt_latest.pth")
        # model_best = join(config.output_dir, "ckpt_best.pth")
        # large_num = -1
        # for p in os.listdir(config.output_dir):
        #     if 'ckpt' in p:
        #         num = p.split('_')[1].split('.')[0]
        #         if str.isnumeric(num):
        #             if int(num) > large_num:
        #                 large_num = int(num)
        # if large_num != -1:
        #     model_latest = join(config.output_dir, f"ckpt_{large_num:02d}.pth")
        # if osp.isfile(model_latest):
        #     config.pretrained_path = model_latest
        #     config.resume = True
        # elif osp.isfile(model_best):
        #     config.pretrained_path = model_best
        #     config.resume = True
        # else:
        #     logger.info(f"Not found checkpoint in {config.output_dir}")


    if (config.pretrained_path.strip() and (osp.isfile(config.pretrained_path)) or "s3://" in config.pretrained_path):
        # if Client is not None:
        #     client = Client()
        #     with io.BytesIO(client.get(config.pretrained_path)) as buffer:
        #         checkpoint = torch.load(buffer, map_location="cpu")
        # else:
        checkpoint = torch.load(config.pretrained_path, map_location="cpu")
        # logger.info(f"Loading checkpoint from {config.pretrained_path}")
        try:
            if "model" in checkpoint.keys():
                state_dict = checkpoint["model"]
            else:
                state_dict = checkpoint["module"] # This is a deepspeed stage 1 model
        except:  
            state_dict = checkpoint

        if config.get('origin_num_frames', None) is not None:
            # logger.info(f"interpolate_pos_embed_internvideo2 (origin_num_frames={config.origin_num_frames})!!!")
            a = len(state_dict)
            interpolate_pos_embed_internvideo2_new(state_dict, model_without_ddp.vision_encoder, orig_t_size=config.origin_num_frames)
            assert a == len(state_dict), state_dict.keys()

        # if config.resume:
        #     assert not (hasattr(config, "deepspeed") and config.deepspeed.enable), "Deepspeed should run here!!!"
        #     optimizer.load_state_dict(checkpoint["optimizer"])
        #     scheduler.load_state_dict(checkpoint["scheduler"])
        #     scaler.load_state_dict(checkpoint["scaler"])
        #     if 'local_step' in checkpoint.keys():
        #         start_epoch = checkpoint['epoch']
        #     else:
        #         start_epoch = checkpoint["epoch"] + 1
        #     global_step = checkpoint["global_step"]

        elif not pretrain:  # downstream init from pretrained ckpt 


            if not config.evaluate or config.get("zero_shot", False):  # finetuning from a pretrained weights.
                if add_decoder:
                    logger.info("Init new decoder with encoder!!!")
                for key in list(state_dict.keys()):
                    if "text_encoder.bert" in key:
                        encoder_key = key.replace("bert.", "") 
                        state_dict[encoder_key] = state_dict[key]
                        if not add_decoder:
                            del state_dict[key]

                    # init text decoder as multimodal encoder (last 6 layers of model.text_encoder)
                    # only for generation tasks like VQA
                    if add_decoder and "text_encoder.bert" in key:
                        if "layer" in key:
                            encoder_keys = key.split(".")
                            layer_num = int(encoder_keys[4])
                            if layer_num < config.model.text_encoder.fusion_layer:
                                del state_dict[key]
                                continue
                            else:
                                decoder_layer_num = layer_num - config.model.text_encoder.fusion_layer
                                encoder_keys[4] = str(decoder_layer_num)
                                encoder_key = ".".join(encoder_keys)
                        else:
                            encoder_key = key
                        decoder_key = encoder_key.replace("text_encoder", "text_decoder")
                        state_dict[decoder_key] = state_dict[key]
                        del state_dict[key]


        msg = model.load_state_dict(state_dict, strict=True)
        # logger.info(msg)
        # logger.info(f"Loaded checkpoint from {config.pretrained_path}")
    else:
        raise ValueError("No pretrained checkpoint provided!!!")
        # if not config.resume:
        #     assert not config.evaluate, "No available pretrained checkpoint provided!!!"
        #     assert config.pretrained_path == "", config.pretrained_path
        #     logger.warning("No available pretrained checkpoint provided, training from scratch.")


    # if hasattr(config, "deepspeed") and config.deepspeed.enable:
    #     logger.info(f'Use deepspeed to initialize model (resume={config.resume}) !!!')
    #     model = model_without_ddp

    #     model, optimizer, _, _ = deepspeed.initialize(
    #         args=config, model=model, model_parameters=optimizer_params, dist_init_required=not config.distributed,
    #         lr_scheduler=lambda opt: create_scheduler(config.scheduler, opt)
    #     )


    #     if config.resume:
    #         logger.info(f'Resume deepspeed ckpt from {config.output_dir}, tag={config.pretrained_path}, load_module_strict={config.get("load_module_strict", True)}, load_lr_scheduler_states={config.get("load_lr_scheduler_states", True)}!!!')
    #         _, client_states = model.load_checkpoint(config.output_dir, tag=config.pretrained_path,  load_module_strict=config.get("load_module_strict", True), load_lr_scheduler_states=config.get("load_lr_scheduler_states", True))
    #         logger.info(client_states)
    #         if 'local_step' in client_states.keys():
    #             start_epoch = client_states['epoch']
    #         else:
    #             start_epoch = client_states['epoch'] + 1
    #         global_step = client_states['global_step']


    # logger.info(f"Cuda memory after create model: {torch.cuda.memory_allocated() // 1024**2}M, Max mem: {torch.cuda.max_memory_allocated() // 1024**2}M start_epoch={start_epoch}, global_step={global_step}")
    # print(f"\033[31m Cuda memory after create model: {torch.cuda.memory_allocated() // 1024**2}M, Max mem: {torch.cuda.max_memory_allocated() // 1024**2}M start_epoch={start_epoch}, global_step={global_step}\033[0m")

    return (
        model,
        # model_without_ddp,
        # optimizer,
        # scheduler,
        # scaler,
        tokenizer,
        # start_epoch,
        # global_step,
    )
