from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
from __future__ import print_function

import torch
from torch.utils.data import (SequentialSampler)
import numpy as np
import random
import os

from backup.modeling_meta import MTN, VNet
from modules.metrics import compute_metrics
import time
import argparse

from dataloaders.dataloader_meta_msrvtt import MSRVTT_1KA_Test, MSRVTT_7k_Test, MSRVTT_MIX_Train, MSRVTT_7k_Train, MSRVTT_1KA_Train
from dataloaders.dataloader_meta_youcook import Youcook_Train, Youcook_Retrieval_Test, Youcook_Caption_Test


from modules.optimization_pseudo import MetaSGD
from modules.tokenization import BertTokenizer
from modules.file_utils import PYTORCH_PRETRAINED_BERT_CACHE
from modules.modeling_MTN import UniVL
from modules.optimization_meta import BertAdam
from torch.utils.data import DataLoader

from modules.utils import str2list, AverageMeter, str2bool, parallel_apply, get_logger
from dataloaders.dataloader_youcook_retrieval import Youcook_DataLoader
from dataloaders.dataloader_meta_youcook import Youcook_Train
from dataloaders.dataloader_msrvtt_retrieval import MSRVTT_Retrieval_TestDataLoader
from dataloaders.dataloader_msrvtt_retrieval import MSRVTT_Retrieval_TrainDataLoader

import wandb
global logger
import socket
import torch.multiprocessing as mp

def get_args(description='UniVL on Retrieval Task'):
    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("--do_pretrain", action='store_true', help="Whether to run training.")
    parser.add_argument("--do_train", action='store_true', help="Whether to run training.")
    parser.add_argument("--do_eval", action='store_true', help="Whether to run eval on the dev set.")

    parser.add_argument('--youcook_train_csv', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_train.csv', help='')
    parser.add_argument('--youcook_val_csv', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_val.csv', help='')
    parser.add_argument('--youcook_features_path', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_videos_features.pickle', help='feature path')

    parser.add_argument('--youcook_data_path_0', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_data.no_transcript.pickle', help='data pickle file path')
    parser.add_argument('--youcook_data_path_1', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_data.transcript.pickle', help='data pickle file path')
    parser.add_argument('--youcook_data_path_2', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_data.transcript_v2.pickle', help='data pickle file path')
    parser.add_argument('--youcook_data_path_3', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_data.transcript_v3.pickle', help='data pickle file path')
    parser.add_argument('--youcook_data_path_4', type=str, default='../../../../../data/project/rw/youcookii_feature/youcookii_data.transcript_v4.pickle', help='data pickle file path')

    parser.add_argument('--VTT_train_csv', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_train.9k.csv', help='')
    parser.add_argument('--VTT_val_csv', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_JSFUSION_test.csv', help='')
    parser.add_argument('--VTT_data_path_0', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_data.json', help='data pickle file path')
    parser.add_argument('--VTT_data_path_1', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_transcript_data.json', help='data pickle file path')
    parser.add_argument('--VTT_data_path_2', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_transcript_data_v2.json', help='data pickle file path')
    parser.add_argument('--VTT_data_path_3', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_transcript_data_v3.json', help='data pickle file path')
    parser.add_argument('--VTT_features_path', type=str, default='../../../../../data/project/rw/msrvtt_feature/msrvtt_videos_features.pickle', help='feature path')
    parser.add_argument('--VTT_train_1KB_csv', type=str, default='../../../../../data/project/rw/msrvtt_feature/MSRVTT_train.1KB.csv', help='')
    parser.add_argument('--transcript', type=int, default=2, help='')

    parser.add_argument("--init_model", default="/data/project/rw/joonmyung/data/UniVL/univl.pretrained.bin", type=str, required=False, help="Initial model.")
    parser.add_argument('--num_thread_reader', type=int, default=1, help='')
    parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate')
    parser.add_argument('--epochs', type=int, default=20, help='upper epoch limit')
    parser.add_argument('--batch_size', type=int, default=256, help='batch size')
    parser.add_argument('--batch_size_val', type=int, default=3500, help='batch size eval')
    parser.add_argument('--lr_decay', type=float, default=0.9, help='Learning rate exp epoch decay')
    parser.add_argument('--n_display', type=int, default=10, help='Information display frequence')
    parser.add_argument('--video_dim', type=int, default=1024, help='video feature dimension')
    parser.add_argument('--seed', type=int, default=42, help='random seed')
    parser.add_argument('--max_words', type=int, default=20, help='')
    parser.add_argument('--max_frames', type=int, default=100, help='')
    parser.add_argument('--feature_framerate', type=int, default=1, help='')
    parser.add_argument('--margin', type=float, default=0.1, help='margin for loss')
    parser.add_argument('--hard_negative_rate', type=float, default=0.5, help='rate of intra negative sample')
    parser.add_argument('--negative_weighting', type=int, default=1, help='Weight the loss for intra negative')
    parser.add_argument('--n_pair', type=int, default=1, help='Num of pair to output from data loader')

    parser.add_argument("--output_dir", default=None, type=str, required=True,
                        help="The output directory where the model predictions and checkpoints will be written.")
    parser.add_argument("--bert_model", default="bert-base-uncased", type=str, required=True,
                        help="Bert pre-trained model")
    parser.add_argument("--visual_model", default="visual-base", type=str, required=False, help="Visual module")
    parser.add_argument("--cross_model", default="cross-base", type=str, required=False, help="Cross module")
    parser.add_argument("--decoder_model", default="decoder-base", type=str, required=False, help="Decoder module")
    parser.add_argument("--do_lower_case", action='store_true', help="Set this flag if you are using an uncased model.")
    parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.")
    parser.add_argument('--n_gpu', type=int, default=1, help="Changed in the execute process.")

    parser.add_argument("--cache_dir", default="", type=str,
                        help="Where do you want to store the pre-trained models downloaded from s3")

    parser.add_argument('--fp16', action='store_true',
                        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
    parser.add_argument('--fp16_opt_level', type=str, default='O1',
                        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
                             "See details at https://nvidia.github.io/apex/amp.html")

    parser.add_argument("--task_type", default="retrieval", type=str, help="Point the task `retrieval` to finetune.")
    parser.add_argument("--datatype", default="msrvttMIX", choices=['youcook', 'msrvtt1KA', 'msrvtt1KB', 'msrvtt7k', 'msrvttMIX'], type=str, help="Point the dataset `youcook` to finetune.")

    parser.add_argument("--world_size", default=0, type=int, help="distribted training")
    parser.add_argument("--local_rank", default=0, type=int, help="distribted training")
    parser.add_argument('--coef_lr', type=float, default=0.1, help='coefficient for bert branch.')
    parser.add_argument('--use_mil', action='store_true', help="Whether use MIL as Miech et. al. (2020).")
    parser.add_argument('--sampled_use_mil', action='store_true', help="Whether MIL, has a high priority than use_mil.")

    parser.add_argument('--text_num_hidden_layers', type=int, default=12, help="Layer NO. of text.")
    parser.add_argument('--visual_num_hidden_layers', type=int, default=6, help="Layer NO. of visual.")
    parser.add_argument('--cross_num_hidden_layers', type=int, default=2, help="Layer NO. of cross.")
    parser.add_argument('--decoder_num_hidden_layers', type=int, default=3, help="Layer NO. of decoder.")

    parser.add_argument('--train_sim_after_cross', action='store_true', help="Test retrieval after cross encoder.")
    parser.add_argument('--expand_msrvtt_sentences', action='store_true', help="")



    parser.add_argument('--use_wandb', default=True, type=str2bool)
    parser.add_argument('--wandb_project', default='ECCV2022_L2F', type=str)
    parser.add_argument('--wandb_table', default='Sample', type=str)
    parser.add_argument('--is_distributed', default=1, type=int, help='')
    parser.add_argument('--tasks', default=[1, 1, 1, 1, 1, 1, 1, 1, 0], type=str2list, help='task')
    parser.add_argument('--target_tasks', default=[0, 1, 0, 0, 0, 0, 0, 0, 0], type=str2list, help='target task')
    parser.add_argument('--start_eval', default=0, type=int, help='start evaluation')
    parser.add_argument('--lr_vnet', default=0.001, type=float, help='Vnet Learning Rate')
    parser.add_argument('--decay_vnet', default=0.0001, type=float, help='Vnet weight decay')
    parser.add_argument('--lr_meta', default=0.01, type=float, help='Meta Model Learning Rate(meta_lr)')
    parser.add_argument('--lr_meta_type', default=0, type=float, help='0:lr, 1:current lr')
    parser.add_argument('--activation', default='softmax', type=str)
    parser.add_argument('--activation_scaler', default='1.0', type=float)
    parser.add_argument('--t', default=1, help="temperature", type=float)
    parser.add_argument('--grad_clip', default=0, help="vnet gradient clip", type=float)
    parser.add_argument('--grad_clip_vnet', default=0, help="vnet gradient clip", type=float)
    parser.add_argument('--grad_clip_vnet_type', default=1, help="0:default, 1:ceil, 2:bottom", type=float)
    parser.add_argument('--ip_list', default=None, help="ipList", type=str2list)
    parser.add_argument('--loss_metric_align', default=2, help="", type=int)
    parser.add_argument('--loss_metric_joint', default=1, help="", type=int)
    parser.add_argument('--loss_type', default=0, help="0:all, 1:align + w*aux, 1:align_m_align + w*aux", type=int)
    parser.add_argument('--vnet_split', default=0, help="0: shared, 1: split", type=int)
    parser.add_argument('--vnet_eval', default=1, help="0: train, 1: eval", type=int)
    parser.add_argument('--vnet_kind', default=1, help="0: MLP, 1:Transformer", type=int)
    parser.add_argument('--vnet_S_type', default=0, help="[l, t, l*t, l+t, (l,t)]", type=int)
    parser.add_argument('--vnet_T_type', default=3, help="[s, t, s*t, s+t, (s,t)]", type=int)
    parser.add_argument('--task_embed', default=100, help="", type=int)
    parser.add_argument('--warmup', default='warmup_linear_down', help="warmup_cosine, warmup_constant, warmup_linear", type=str)
    parser.add_argument('--scheduler_epoch', default=4, help="", type=int)
    parser.add_argument('--milestones', default=None, help="", type=str2list)
    parser.add_argument('--milestones_proportion', default=0.1, help="", type=float)
    parser.add_argument("--warmup_proportion", default=0.1, type=float,
                        help="Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10%% of training.")
    parser.add_argument('--version', default="V1.0.0", type=str)
    parser.add_argument('--transformer_dim', default=[512,128,512], type=str2list)
    parser.add_argument('--debug', default=False, type=str2bool)

    parser.add_argument('--eval_task', default=[1, 0], type=str2list, help="[retrieval, caption] : Evaluation Task")
    parser.add_argument('--eval_start', default=[0, 0], type=str2list, help='start evaluation')
    parser.add_argument('--eval_term', default=[1, 1], type=str2list, help='start evaluation')

    parser.add_argument('--BMM', default=[120   , 48, 48], type=str2list, help="[BATCH, MAX_WORDS, MAX_FRAMES]")
    parser.add_argument('--mask_probability', default=0.15, type=float, help="")
    parser.add_argument('--val_mask_probability', default=0.15, type=float, help="")

    args = parser.parse_args()
    print("transcript : ", args.transcript)
    if args.transcript == 0:
        args.youcook_data_path = args.youcook_data_path_0
        args.VTT_data_path = args.VTT_data_path_0
    elif args.transcript == 1:
        args.youcook_data_path = args.youcook_data_path_1
        args.VTT_data_path = args.VTT_data_path_1
    elif args.transcript == 2:
        args.youcook_data_path = args.youcook_data_path_2
        args.VTT_data_path = args.VTT_data_path_2
    elif args.transcript == 3:
        args.youcook_data_path = args.youcook_data_path_3
        args.VTT_data_path = args.VTT_data_path_2
    elif args.transcript == 4:
        args.youcook_data_path = args.youcook_data_path_4
        args.VTT_data_path = args.VTT_data_path_2

    print("|- transcript : ", args.transcript)
    print("   |- youcook : ", args.youcook_data_path)
    print("   |- VTT     : ", args.VTT_data_path)

    args.port = random.randint(0, 20000)
    if args.datatype == "youcook":
        args.n_display = 20
    else:
        args.n_display = 100

    if args.eval_term[0] == 0:
        args.eval_term[0] = 1
    if args.eval_term[1] == 0:
        args.eval_term[1] = 1

    if args.eval_task[0] and args.eval_task[1]:
        raise ValueError("Please Select One Task for Training")
    elif args.eval_task[0] == 1:  # Retrieval Task
        args.downstream = "retrieval"
    elif args.eval_task[1] == 1:  # Captioning
        args.downstream = "caption"
    else:
        raise ValueError("Please Select One Task for Training")


    args.batch_size, args.max_words, args.max_frames = args.BMM[0], args.BMM[1], args.BMM[2]

    # Check paramenters
    if args.gradient_accumulation_steps < 1:
        raise ValueError("Invalid gradient_accumulation_steps parameter: {}, should be >= 1".format(
            args.gradient_accumulation_steps))
    if not args.do_train and not args.do_eval:
        raise ValueError("At least one of `do_train` or `do_eval` must be True.")

    args.batch_size = int(args.batch_size / args.gradient_accumulation_steps)

    return args

def set_seed_logger(args):
    global logger
    # predefining random initial seeds
    random.seed(args.seed)
    os.environ['PYTHONHASHSEED'] = str(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    if args.is_distributed:
        world_size = torch.distributed.get_world_size()
        torch.cuda.set_device(args.local_rank)
        args.world_size = world_size

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir, exist_ok=True)

    logger = get_logger(os.path.join(args.output_dir, "log.txt"))

    if args.local_rank == 0:
        logger.info("Effective parameters:")
        for key in sorted(args.__dict__):
            logger.info("  <<< {}: {}".format(key, args.__dict__[key]))


    return args

def init_device(args, local_rank):
    global logger

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu", local_rank)

    n_gpu = torch.cuda.device_count()
    logger.info("device: {} n_gpu: {}".format(device, n_gpu))
    args.n_gpu = n_gpu

    if args.batch_size % args.n_gpu != 0 or args.batch_size_val % args.n_gpu != 0:
        raise ValueError("Invalid batch_size/batch_size_val and n_gpu parameter: {}%{} and {}%{}, should be == 0".format(
            args.batch_size, args.n_gpu, args.batch_size_val, args.n_gpu))

    return device, n_gpu

def init_model(args, device, n_gpu, local_rank):

    if args.init_model:
        model_state_dict = torch.load(args.init_model, map_location='cpu')
    else:
        model_state_dict = None

    # Prepare model
    cache_dir = args.cache_dir if args.cache_dir else os.path.join(str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed')
    model = UniVL.from_pretrained(args.bert_model, args.visual_model, args.cross_model, args.decoder_model,
                                   cache_dir=cache_dir, state_dict=model_state_dict, task_config=args)

    if args.is_distributed:
        model.to(device)
    else:
        model.cuda()
    return model

def prep_optimizer(args, model, num_train_optimization_steps, num_train_warmup_steps, device, n_gpu, local_rank, coef_lr=1.):

    if hasattr(model, 'module'):
        model = model.module

    param_optimizer = list(model.named_parameters())
    no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']

    no_decay_param_tp = [(n, p) for n, p in param_optimizer if not any(nd in n for nd in no_decay)]
    decay_param_tp = [(n, p) for n, p in param_optimizer if any(nd in n for nd in no_decay)]

    no_decay_bert_param_tp = [(n, p) for n, p in no_decay_param_tp if "bert." in n]
    no_decay_nobert_param_tp = [(n, p) for n, p in no_decay_param_tp if "bert." not in n]

    decay_bert_param_tp = [(n, p) for n, p in decay_param_tp if "bert." in n]
    decay_nobert_param_tp = [(n, p) for n, p in decay_param_tp if "bert." not in n]

    optimizer_grouped_parameters = [
        {'params': [p for n, p in no_decay_bert_param_tp], 'weight_decay': 0.01, 'lr': args.lr * coef_lr},
        {'params': [p for n, p in no_decay_nobert_param_tp], 'weight_decay': 0.01},
        {'params': [p for n, p in decay_bert_param_tp], 'weight_decay': 0.0, 'lr': args.lr * coef_lr},
        {'params': [p for n, p in decay_nobert_param_tp], 'weight_decay': 0.0}
    ]

    scheduler = None
    optimizer = BertAdam(optimizer_grouped_parameters, lr=args.lr, warmup=args.warmup, warmup_proportion=args.warmup_proportion,
                         milestones=args.milestones, milestones_proportion=args.milestones_proportion,
                         t_total=num_train_optimization_steps, t_warmup=num_train_warmup_steps, weight_decay=0.01,
                         max_grad_norm=1.0)

    if args.is_distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank],
                                                          output_device=local_rank, find_unused_parameters=True)

    return optimizer, scheduler, model

def dataloader_youcook_train(args, tokenizer):
    youcook_dataset = Youcook_Train(
        csv=args.youcook_train_csv,
        data_path=args.youcook_data_path,
        features_path=args.youcook_features_path,
        max_words=args.max_words,
        feature_framerate=args.feature_framerate,
        tokenizer=tokenizer,
        max_frames=args.max_frames,
        downstream=args.downstream,
        p=args.mask_probability
    )

    if args.is_distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(youcook_dataset)
    else:
        train_sampler = None

    dataloader = DataLoader(
        youcook_dataset,
        batch_size=args.batch_size // args.world_size,
        num_workers=args.num_thread_reader,
        pin_memory=False,
        shuffle=(train_sampler is None),
        sampler=train_sampler,
        drop_last=True
    )

    return dataloader, len(youcook_dataset), train_sampler

def dataloader_youcook_test(args, tokenizer):
    youcook_retrieval = Youcook_Retrieval_Test(
        csv=args.youcook_val_csv,
        data_path=args.youcook_data_path,
        features_path=args.youcook_features_path,
        max_words=args.max_words,
        feature_framerate=args.feature_framerate,
        tokenizer=tokenizer,
        max_frames=args.max_frames,
    )


    test_retrieval_sampler = SequentialSampler(youcook_retrieval)
    dataloader_youcook_retrieval = DataLoader(
        youcook_retrieval,
        sampler=test_retrieval_sampler,
        batch_size=args.batch_size_val,
        num_workers=args.num_thread_reader,
        pin_memory=False,
    )

    youcook_caption = Youcook_Caption_Test(
        csv=args.youcook_val_csv,
        data_path=args.youcook_data_path,
        features_path=args.youcook_features_path,
        max_words=args.max_words,
        feature_framerate=args.feature_framerate,
        tokenizer=tokenizer,
        max_frames=args.max_frames,
    )

    test_caption_sampler = SequentialSampler(youcook_caption)
    dataloader_youcook_caption = DataLoader(
        youcook_caption,
        sampler=test_caption_sampler,
        batch_size=args.batch_size_val,
        num_workers=args.num_thread_reader,
        pin_memory=False,
    )

    if args.local_rank == 0:
        logger.info('YoucookII retrieval validation pairs: {}'.format(len(youcook_retrieval)))
        logger.info('YoucookII caption validation pairs: {}'.format(len(youcook_caption)))

    return dataloader_youcook_retrieval, len(youcook_retrieval), dataloader_youcook_caption, len(youcook_caption)

def dataloader_msrvtt_train(args, tokenizer):
    if args.datatype =="msrvtt1KA":
        msrvtt_dataset = MSRVTT_1KA_Train(
            csv_path=args.VTT_train_csv,
            json_path=args.VTT_data_path,
            features_path=args.VTT_features_path,
            max_words=args.max_words,
            feature_framerate=args.feature_framerate,
            tokenizer=tokenizer,
            max_frames=args.max_frames,
            unfold_sentences=args.expand_msrvtt_sentences,
            p = args.mask_probability
        )
    elif args.datatype =="msrvtt1KB":
        msrvtt_dataset = MSRVTT_1KA_Train(
            csv_path=args.VTT_train_1KB_csv,
            json_path=args.VTT_data_path,
            features_path=args.VTT_features_path,
            max_words=args.max_words,
            feature_framerate=args.feature_framerate,
            tokenizer=tokenizer,
            max_frames=args.max_frames,
            unfold_sentences=args.expand_msrvtt_sentences,
            p=args.mask_probability
        )
    elif args.datatype == "msrvtt7k":
        msrvtt_dataset = MSRVTT_7k_Train(
            csv_path=args.VTT_train_csv,
            json_path=args.VTT_data_path,
            features_path=args.VTT_features_path,
            max_words=args.max_words,
            feature_framerate=args.feature_framerate,
            tokenizer=tokenizer,
            max_frames=args.max_frames,
            split_type="train",
            p=args.mask_probability,
            downstream = args.downstream
        )
    elif args.datatype == "msrvttMIX": # From 7K
        msrvtt_dataset = MSRVTT_MIX_Train(
            csv_path=args.VTT_train_csv,
            json_path=args.VTT_data_path,
            features_path=args.VTT_features_path,
            max_words=args.max_words,
            feature_framerate=args.feature_framerate,
            tokenizer=tokenizer,
            max_frames=args.max_frames,
            split_type="train",
            p=args.mask_probability
        )
    else:
        raise argparse.ArgumentTypeError('Please Check Datatype')

    if args.is_distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(msrvtt_dataset)
    else:
        train_sampler = None

    dataloader = DataLoader(
        msrvtt_dataset,
        batch_size=args.batch_size // args.world_size,
        num_workers=args.num_thread_reader,
        pin_memory=False,
        shuffle=(train_sampler is None),
        sampler=train_sampler,
        drop_last=True,
    )

    return dataloader, len(msrvtt_dataset), train_sampler




def dataloader_msrvtt_test(args, tokenizer, split_type="test", ):
    if (args.datatype =="msrvtt1KA" or args.datatype == "msrvtt1KB") or args.datatype == "msrvttMIX":
        msrvtt_testset = MSRVTT_1KA_Test(
            csv_path=args.VTT_val_csv,
            features_path=args.VTT_features_path,
            max_words=args.max_words,
            feature_framerate=args.feature_framerate,
            tokenizer=tokenizer,
            max_frames=args.max_frames,
            )
    elif args.datatype == "msrvtt7k":
        msrvtt_testset = MSRVTT_7k_Test(
            csv_path=args.VTT_val_csv,
            json_path=args.VTT_data_path,
            features_path=args.VTT_features_path,
            max_words=args.max_words,
            feature_framerate=args.feature_framerate,
            tokenizer=tokenizer,
            max_frames=args.max_frames,
            split_type=split_type,
            )
    else:
        raise argparse.ArgumentTypeError('Please Check Datatype')

    test_msrvtt_sampler = SequentialSampler(msrvtt_testset)
    dataloader_msrvtt_test = DataLoader(
        msrvtt_testset,
        sampler=test_msrvtt_sampler,
        batch_size=args.batch_size_val,
        num_workers=args.num_thread_reader,
        pin_memory=False,
        drop_last=False,
    )

    return dataloader_msrvtt_test, len(msrvtt_testset), dataloader_msrvtt_test, len(msrvtt_testset)


def save_model(epoch, args, model, type_name=""):
    # Only save the model it-self
    model_to_save = model.module if hasattr(model, 'module') else model
    output_model_file = os.path.join(
        args.output_dir, "pytorch_model.bin.{}{}".format("" if type_name=="" else type_name+".", epoch))
    torch.save(model_to_save.state_dict(), output_model_file)
    logger.info("Model saved to %s", output_model_file)
    return output_model_file

def load_model(epoch, args, n_gpu, device, model_file=None):
    if model_file is None or len(model_file) == 0:
        model_file = os.path.join(args.output_dir, "pytorch_model.bin.{}".format(epoch))
    if os.path.exists(model_file):
        model_state_dict = torch.load(model_file, map_location='cpu')
        if args.local_rank == 0:
            logger.info("Model loaded from %s", model_file)
        # Prepare model
        cache_dir = args.cache_dir if args.cache_dir else os.path.join(str(PYTORCH_PRETRAINED_BERT_CACHE), 'distributed')
        model = UniVL.from_pretrained(args.bert_model, args.visual_model, args.cross_model, args.decoder_model,
                                       cache_dir=cache_dir, state_dict=model_state_dict, task_config=args)

        model.to(device)
    else:
        model = None
    return model
from tqdm import tqdm
import copy
def train_epoch(epoch, args, net, vnets, train_dataloader, device, n_gpu, optimizer, optimizer_vnet, scheduler, global_step, local_rank=0):
    global logger
    torch.cuda.empty_cache()
    net.train()
    log_step = args.n_display
    start_time = time.time()
    start_time_epoch = time.time()

    if args.activation == "sigmoid":
        activation = torch.nn.Sigmoid()
    elif args.activation == "softmax":
        activation = torch.nn.Softmax(dim=0)
    elif args.activation == "relu":
        activation = torch.nn.ReLU()
    elif args.activation == "softplus":
        activation = torch.nn.Softplus()
    total_losses, total_losses_aux, weights_prev, weights, ex_time = AverageMeter(), [AverageMeter() for _ in range(args.taskNum)], [AverageMeter() for _ in range(args.taskNum)], [AverageMeter() for _ in range(args.taskNum)], AverageMeter()
    ies, h1es, h2es, oes, w1es, b1es, w2es, b2es = AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter(), AverageMeter()
    grads_pseudos, grads_vnets = AverageMeter(), AverageMeter()
    for step, batch in enumerate(tqdm(train_dataloader)):
        if args.debug:
            break
        if n_gpu == 1 and args.is_distributed:
            batch = tuple(t.to(device=device, non_blocking=True) for t in batch)
        else:
            batch = tuple(t.cuda() for t in batch)

        input_ids, input_mask, segment_ids, video, video_mask, \
        pairs_masked_text, pairs_token_labels, masked_video, video_labels_index, \
        pairs_input_caption_ids, pairs_decoder_mask, pairs_output_caption_ids = batch

        pseudo_net = copy.deepcopy(net)
        if args.is_distributed:
            pseudo_net.to(device)
        else:
            pseudo_net.cuda()
        pseudo_net.train()


        losses = pseudo_net(input_ids, segment_ids, input_mask, video, video_mask,
                     pairs_masked_text=pairs_masked_text, pairs_token_labels=pairs_token_labels,
                     masked_video=masked_video, video_labels_index=video_labels_index,
                     input_caption_ids=pairs_input_caption_ids, decoder_mask=pairs_decoder_mask,
                     output_caption_ids=pairs_output_caption_ids, tasks= args.tasks)
        t_type = torch.eye(args.taskNum-args.loss_type, device=device) if args.is_distributed else torch.eye(args.taskNum-args.loss_type, device=torch.device("cuda"))
        l_type = torch.stack(losses).unsqueeze(1)
        pri_loss = 0
        if args.loss_type == 1:
            l_type = torch.cat((l_type[:1],l_type[2:]))
            pri_loss = losses[1]
        if args.loss_type == 2:
            l_type = torch.cat((l_type[:1], l_type[2:5], l_type[6:]))
            pri_loss = losses[1] + losses[5]
        # [l, t, l * t, (l, t)]
        if args.vnet_S_type == 0:
            vnet_inputs = l_type
        elif args.vnet_S_type == 1:
            vnet_inputs = t_type
        elif args.vnet_S_type == 2:
            vnet_inputs = t_type * l_type
        elif args.vnet_S_type == 3:
            vnet_inputs = torch.cat((l_type, t_type), dim=1)

        weight_prev, _, _, _, _, _, _, _, _ = vnets(vnet_inputs.data, args.vnet_T_type)
        weight = activation(weight_prev*args.t)

        pseudo_loss = (l_type * weight * args.activation_scaler).sum() / len(weight) + pri_loss
        if args.lr_meta_type == 0:
            pseudo_optimizer = MetaSGD(pseudo_net, pseudo_net.parameters(), lr=args.lr_meta)
        elif args.lr_meta_type == 1:
            pseudo_optimizer = MetaSGD(pseudo_net, pseudo_net.parameters(), lr=optimizer.param_groups[0]['lr'])
        elif args.lr_meta_type == 2:
            meta_lr = sorted(list(set(optimizer.get_lr(start=True))))[0]
            pseudo_optimizer = MetaSGD(pseudo_net, pseudo_net.parameters(), lr=meta_lr)

        if args.lr_vnet == 0:
            for param_group in optimizer_vnet.param_groups:
                param_group['lr'] = optimizer.param_groups[0]['lr'] if len(sorted(list(set(optimizer.get_lr())))) == 0 else sorted(list(set(optimizer.get_lr())))[0]


        pseudo_grads = torch.autograd.grad(pseudo_loss, pseudo_net.parameters(), create_graph=True, allow_unused=True)
        if args.grad_clip != 0:
            total_norm = torch.norm(torch.stack([torch.norm(p.detach()).to(device) for p in pseudo_grads if p is not None]), 2)
            if total_norm != 0:
                clip_coef = args.grad_clip / total_norm
                for g in pseudo_grads:
                    if g is not None:
                        g.mul_(clip_coef)


        pseudo_optimizer.meta_step(pseudo_grads)

        vnet_losses = pseudo_net(input_ids, segment_ids, input_mask, video, video_mask,
                     pairs_masked_text=pairs_masked_text, pairs_token_labels=pairs_token_labels,
                     masked_video=masked_video, video_labels_index=video_labels_index, tasks=args.target_tasks)

        vnet_loss = sum(vnet_losses) / len(vnet_losses)

        optimizer_vnet.zero_grad()
        vnet_loss.backward()

        if args.grad_clip_vnet != 0:
            total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach()).to(device) for p in vnets.parameters()]), 2)
            if total_norm != 0:
                clip_coef = args.grad_clip_vnet / total_norm
                if args.grad_clip_vnet_type == 0:
                    for p in vnets.parameters():
                        p.grad.mul_(clip_coef)
                elif args.grad_clip_vnet_type == 1:
                    if clip_coef <= 1:
                        for p in vnets.parameters():
                            p.grad.mul_(clip_coef)
                elif args.grad_clip_vnet_type == 2:
                    if clip_coef >= 1:
                        for p in vnets.parameters():
                            p.grad.mul_(clip_coef)

        optimizer_vnet.step()

        pseudo_net.zero_grad()
        pseudo_optimizer.zero_grad()
        total_losses.update(float(vnet_loss.detach().cpu()))
        del pseudo_net, pseudo_loss, vnet_loss, pseudo_optimizer

        losses = net(input_ids, segment_ids, input_mask, video, video_mask,
                     pairs_masked_text=pairs_masked_text, pairs_token_labels=pairs_token_labels,
                     masked_video=masked_video, video_labels_index=video_labels_index,
                     input_caption_ids=pairs_input_caption_ids, decoder_mask=pairs_decoder_mask,
                     output_caption_ids=pairs_output_caption_ids, tasks=args.tasks)

        t_type = torch.eye(args.taskNum-args.loss_type, device=device) if args.is_distributed else torch.eye(args.taskNum-args.loss_type, device=torch.device("cuda"))
        l_type = torch.stack(losses).unsqueeze(1)

        pri_loss = 0
        if args.loss_type == 1:
            l_type = torch.cat((l_type[:1], l_type[2:]))
            pri_loss = losses[1]
        if args.loss_type == 2:
            l_type = torch.cat((l_type[:1], l_type[2:5], l_type[6:]))
            pri_loss = losses[1] + losses[5]

        if args.vnet_S_type == 0:
            vnet_inputs = l_type
        elif args.vnet_S_type == 1:
            vnet_inputs = t_type
        elif args.vnet_S_type == 2:
            vnet_inputs = t_type * l_type
        elif args.vnet_S_type == 3:
            vnet_inputs = torch.cat((l_type, t_type), dim=1)

        with torch.no_grad():
            weight_prev, i, h1, h2, o, w1, b1, w2, b2 = vnets(vnet_inputs, args.vnet_T_type)
            weight = activation(weight_prev * args.t)

            ies.update(i); h1es.update(h1); h2es.update(h2); oes.update(o)
            w1es.update(w1); b1es.update(b1); w2es.update(w2); b2es.update(b2)
            c = 0
            for t in range(args.taskNum):
                if args.loss_type == 1 and t == 1:
                    weights_prev[t].update(0)
                    weights[t].update(1)
                    continue
                if args.loss_type == 2 and ((t == 1) or (t == 5)):
                    weights_prev[t].update(0)
                    weights[t].update(1)
                    continue

                weights_prev[t].update(float(weight_prev[c].detach().cpu().data))
                weights[t].update(weight[c].detach().cpu().data)
                c += 1

        loss = (l_type * weight * args.activation_scaler).sum() / len(weight) + pri_loss

        loss.backward()
        if args.is_distributed:
            torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0)

        if scheduler is not None:
            scheduler.step()  # Update learning rate schedule

        optimizer.step(epoch=epoch)
        optimizer.zero_grad()

        global_step += 1
        if global_step % log_step == 0 and local_rank == 0:
            logger.info("Epoch: %d/%s, Step: %d/%d, Lr: %s, Loss: %f, Time/step: %f", epoch + 1,
                        args.epochs, step + 1,
                        len(train_dataloader), "-".join([str('%.8f'%itm) for itm in sorted(list(set(optimizer.get_lr())))]),
                        total_losses.avg,
                        (time.time() - start_time) / (log_step * args.gradient_accumulation_steps))
            start_time = time.time()
        for t in range(args.taskNum):
            total_losses_aux[t].update(float(losses[t].detach().cpu().data))

    ex_time.update(time.time() - start_time_epoch)

    return total_losses, global_step, \
           [weights_prev, weights], \
           [grads_pseudos, grads_vnets], \
           total_losses_aux, \
           [ies, h1es, h2es, oes, w1es, b1es, w2es, b2es], \
            ex_time


def _run_on_single_gpu(net, batch_list_t, batch_list_v, batch_sequence_output_list, batch_visual_output_list):
    sim_matrix = []
    for idx1, b1 in enumerate(batch_list_t):
        input_ids, input_mask, segment_ids, _, _, _, _, _, _ = b1
        sequence_output = batch_sequence_output_list[idx1]
        each_row = []
        for idx2, b2 in enumerate(batch_list_v):
            _, _, _, video, video_mask, _, _, _, _ = b2
            visual_output = batch_visual_output_list[idx2]
            b1b2_logits = net.get_similarity_logits_align(sequence_output, visual_output, input_mask, video_mask)
            b1b2_logits = b1b2_logits.cpu().detach().numpy()
            each_row.append(b1b2_logits)
        each_row = np.concatenate(tuple(each_row), axis=-1)
        sim_matrix.append(each_row)
    return sim_matrix

def eval_epoch(args, model, test_dataloader, device, n_gpu):
    if hasattr(model, 'module'):
        model = model.module.to(device)
    else:
        model = model.to(device)
    model.eval()
    with torch.no_grad():
        batch_list = []
        batch_sequence_output_list, batch_visual_output_list = [], []
        for bid, batch in enumerate(tqdm(test_dataloader)):
            batch = tuple(t.to(device) for t in batch)

            input_ids, input_mask, segment_ids, video, video_mask, _, _, _, _ = batch

            sequence_output, visual_output = model.get_sequence_visual_output(input_ids, segment_ids, input_mask, video, video_mask)

            batch_sequence_output_list.append(sequence_output)
            batch_visual_output_list.append(visual_output)
            batch_list.append(batch)

            print("{}/{}\r".format(bid, len(test_dataloader)), end="")

        if n_gpu > 1:
            device_ids = list(range(n_gpu))
            batch_list_t_splits = []
            batch_list_v_splits = []
            batch_t_output_splits = []
            batch_v_output_splits = []
            bacth_len = len(batch_list)
            split_len = (bacth_len + n_gpu - 1) // n_gpu
            for dev_id in device_ids:
                s_, e_ = dev_id * split_len, (dev_id + 1) * split_len
                if dev_id == 0:
                    batch_list_t_splits.append(batch_list[s_:e_])
                    batch_list_v_splits.append(batch_list)

                    batch_t_output_splits.append(batch_sequence_output_list[s_:e_])
                    batch_v_output_splits.append(batch_visual_output_list)
                else:
                    devc = torch.device('cuda:{}'.format(str(dev_id)))
                    devc_batch_list = [tuple(t.to(devc) for t in b) for b in batch_list[s_:e_]]
                    batch_list_t_splits.append(devc_batch_list)
                    devc_batch_list = [tuple(t.to(devc) for t in b) for b in batch_list]
                    batch_list_v_splits.append(devc_batch_list)

                    devc_batch_list = [b.to(devc) for b in batch_sequence_output_list[s_:e_]]
                    batch_t_output_splits.append(devc_batch_list)
                    devc_batch_list = [b.to(devc) for b in batch_visual_output_list]
                    batch_v_output_splits.append(devc_batch_list)
            parameters_tuple_list = [(batch_list_t_splits[dev_id], batch_list_v_splits[dev_id],
                                      batch_t_output_splits[dev_id], batch_v_output_splits[dev_id]) for dev_id in device_ids]
            parallel_outputs = parallel_apply(_run_on_single_gpu, model, parameters_tuple_list, device_ids)
            sim_matrix = []
            for idx in range(len(parallel_outputs)):
                sim_matrix += parallel_outputs[idx]
            sim_matrix = np.concatenate(tuple(sim_matrix), axis=0)

        else:
            sim_matrix = _run_on_single_gpu(model, batch_list, batch_list, batch_sequence_output_list, batch_visual_output_list)

    metrics = compute_metrics(sim_matrix)
    logger.info('\t Length-T: {}, Length-V:{}'.format(len(sim_matrix), len(sim_matrix[0])))
    logger.info('\t>>>  R@1: {:.4f} - R@5: {:.4f} - R@10: {:.4f} - Median R: {}'.
                format(metrics['R1'], metrics['R5'], metrics['R10'], metrics['MR']))

    R1 = metrics['R1']
    return R1

DATALOADER_DICT = {}
DATALOADER_DICT["youcook"] = {"train":dataloader_youcook_train, "val":dataloader_youcook_test}
DATALOADER_DICT["msrvtt1KA"] = {"train":dataloader_msrvtt_train, "val":dataloader_msrvtt_test}
DATALOADER_DICT["msrvtt7k"] = {"train":dataloader_msrvtt_train, "val":dataloader_msrvtt_test}
DATALOADER_DICT["msrvttMIX"] = {"train":dataloader_msrvtt_train, "val":dataloader_msrvtt_test}

def main_worker(gpu, ngpus_per_node, args):
    global logger
    args.local_rank = gpu  # 0
    args.world_rank = args.local_rank + args.world_rank * ngpus_per_node
    if args.is_distributed:
        print(args.ip, args.local_rank, args.world_rank, args.world_size)
        torch.distributed.init_process_group(backend="nccl", init_method='tcp://{}:{}'.format(args.ip, args.port), world_size=args.world_size, rank=args.world_rank)  # 8, 4
    args = set_seed_logger(args)
    if args.activation == "softmax":
        args.vnet_split = 0


    device, n_gpu = init_device(args, args.local_rank)

    tokenizer = BertTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case)

    assert  args.task_type == "retrieval"
    model = init_model(args, device, n_gpu, args.local_rank)

    assert args.datatype in DATALOADER_DICT
    test_retrieval_dataloader, test_retrieval_length, test_caption_dataloader, test_caption_length = DATALOADER_DICT[args.datatype]["val"](args, tokenizer)

    if args.world_rank == 0 and args.use_wandb:
        wandb.init(project=args.wandb_project, name=args.wandb_table)
        args.dir = wandb.run.dir
        wandb.config.update(args)

    if args.local_rank == 0:
        logger.info("***** Running Retrieval test *****")
        logger.info("  Num examples = %d", test_retrieval_length)
        logger.info("  Batch size = %d", args.batch_size_val)
        logger.info("  Num steps = %d", len(test_retrieval_dataloader))
        logger.info("***** Running Captioning test *****")
        logger.info("  Num examples = %d", test_caption_length)
        logger.info("  Batch size = %d", args.batch_size_val)
        logger.info("  Num steps = %d", len(test_caption_dataloader))

    train_dataloader, train_length, train_sampler = DATALOADER_DICT[args.datatype]["train"](args, tokenizer)
    num_train_optimization_steps = (int(len(train_dataloader) + args.gradient_accumulation_steps - 1)
                                    / args.gradient_accumulation_steps) * args.epochs
    num_train_warmup_steps = (int(len(train_dataloader) + args.gradient_accumulation_steps - 1)
                              / args.gradient_accumulation_steps) * args.scheduler_epoch

    coef_lr = args.coef_lr
    if args.init_model:
        coef_lr = 1.0
    optimizer, scheduler, model = prep_optimizer(args, model, num_train_optimization_steps, num_train_warmup_steps, device, n_gpu, args.local_rank, coef_lr=coef_lr)

    vnets_params, task_embed, task_embed_shape = [], None, 0

    # vnet_S_type         : 0         : [l, t, l*t, (l,t)]
    if args.vnet_S_type == 0:
        vnet_input_dim = 1
    elif args.vnet_S_type == 1:
        vnet_input_dim = args.taskNum - args.loss_type
    elif args.vnet_S_type == 2:
        vnet_input_dim = args.taskNum - args.loss_type
    elif args.vnet_S_type == 3:
        vnet_input_dim = 1 + args.taskNum - args.loss_type
    else:
        raise ValueError("please check vnet Scale type")

    if args.vnet_kind == 0:
        if args.vnet_split == 0:
            vnets = VNet(vnet_input_dim, 100, 1, task_embed_shape=task_embed_shape).to(device) if args.is_distributed else VNet(vnet_input_dim, 100, 1, task_embed_shape=task_embed_shape).cuda()
            if args.vnet_eval:
                vnets.eval()
            vnets_params += vnets.parameters()
        elif args.vnet_split == 1:
            vnets = []
            for t in range(args.taskNum):
                vnets.append(VNet(vnet_input_dim, 100, 1, task_embed_shape=task_embed_shape).to(device) if args.is_distributed else VNet(vnet_input_dim, 100, 1, task_embed_shape=task_embed_shape).cuda())
                if args.vnet_eval:
                    vnets[t].eval()
                vnets_params += vnets[t].parameters()
    elif args.vnet_kind == 1:
        if args.vnet_split == 0:# t_dim, i_dim, h1_dim, h2_dim, o_dim
            vnets = MTN(args.taskNum - args.loss_type, args.transformer_dim[0], vnet_input_dim, args.transformer_dim[1], args.transformer_dim[2], 1, args.vnet_T_type).to(device) if args.is_distributed else MTN(args.taskNum - args.loss_type, args.transformer_dim[0], vnet_input_dim, args.transformer_dim[1], args.transformer_dim[2], 1, args.vnet_T_type).cuda()
            if args.vnet_eval:
                vnets.eval()
            vnets_params += vnets.parameters()
        elif args.vnet_split == 1:
            vnets = []
            for t in range(args.taskNum):
                vnets[t].append(MTN(args.taskNum - args.loss_type, args.transformer_dim[0], vnet_input_dim, args.transformer_dim[1], args.transformer_dim[2], 1, args.vnet_T_type).to(device) if args.is_distributed else MTN(args.taskNum - args.loss_type, args.transformer_dim[0], vnet_input_dim, args.transformer_dim[1], args.transformer_dim[2], 1, args.vnet_T_type).cuda())
                if args.vnet_eval:
                    vnets[t].eval()
                vnets_params += vnets[t].parameters()


    optimizer_vnet = torch.optim.Adam(vnets_params, lr=args.lr_vnet, weight_decay=args.decay_vnet)

    if args.local_rank == 0:
        logger.info("***** Running training *****")
        logger.info("  Num examples = %d", train_length)
        logger.info("  Batch size = %d", args.batch_size)
        logger.info("  Num steps = %d", num_train_optimization_steps * args.gradient_accumulation_steps)

    best_score, R1 = 0.00001, 0
    global_step = 0
    start_time_all = time.time()
    for epoch in range(args.epochs):
        if args.is_distributed:
            train_sampler.set_epoch(epoch)
        tr_loss, global_step, weight, grads, losses, iwes, ex_time = train_epoch(epoch, args, model, vnets, train_dataloader, device, n_gpu, optimizer, optimizer_vnet,
                                                                        scheduler, global_step, local_rank=args.local_rank)
        if args.world_rank == 0:
            ex_time_all = time.time() - start_time_all

            print("Epoch {}/{} Finished, Train Loss: {:0.4f}".format(epoch + 1, args.epochs, tr_loss.avg))
            if epoch >= args.start_eval:
                print("evaluation!!")
                R1 = eval_epoch(args, model, test_retrieval_dataloader, device, n_gpu)
                if best_score <= R1:
                    best_score = R1

            print("Best Accuracy : {:0.4f}, Test Accuracy: {:0.4f}".format(best_score, R1))

            lr = sorted(list(set(optimizer.get_lr())))[0] if len(set(optimizer.get_lr())) != 0 else -1
            if args.world_rank == 0:
                log_dict = {"train_loss": tr_loss.avg
                    , "curr_score": R1
                    , "best_score": best_score
                    , "lr":lr
                    , "vlr": optimizer_vnet.param_groups[0]['lr']
                    , "grads_pseudo": grads[0].avg
                    , "ex_time_epoch": ex_time.val
                    , "ex_time_all" : ex_time_all
                }
                # if args.vnet_split:
                #     for n, gv, vi, vh1, vh2, vo, w1, b1, w2, b2 in zip(args.taskName, grads[1], iwes[0], iwes[1], iwes[2], iwes[3], iwes[4], iwes[5], iwes[6], iwes[7]):
                #         log_dict["{}_grads_vnet".format(n)] = gv.avg
                #         log_dict["{}_vnet_input".format(n)] = vi.avg
                #         log_dict["{}_vnet_hidden_1".format(n)] = vh1.avg
                #         log_dict["{}_vnet_hidden_2".format(n)] = vh2.avg
                #         log_dict["{}_vnet_output".format(n)] = vo.avg
                #         log_dict["{}_vnet_l1_weight".format(n)] = w1.avg
                #         log_dict["{}_vnet_l1_bias".format(n)] = b1.avg
                #         log_dict["{}_vnet_l2_weight".format(n)] = w2.avg
                #         log_dict["{}_vnet_l2_bias".format(n)] = b2.avg
                # else:
                #     log_dict["grads_vnet"] = grads[1].avg
                #     log_dict["vnet_input"] = iwes[0].avg
                #     log_dict["vnet_hidden_1"] = iwes[1].avg
                #     log_dict["vnet_hidden_2"] = iwes[2].avg
                #     log_dict["vnet_output"] = iwes[3].avg
                #     log_dict["vnet_l1_weight"] = iwes[4].avg
                #     log_dict["vnet_l1_bias"] = iwes[5].avg
                #     log_dict["vnet_l2_weight"] = iwes[6].avg
                #     log_dict["vnet_l2_bias"] = iwes[7].avg

                for n, wp, w, l in zip(args.taskName, weight[0], weight[1], losses):
                    log_dict["{}_weight_prev".format(n)] = wp.avg
                    log_dict["{}_weight".format(n)] = w.avg
                    log_dict["{}_loss".format(n)] = l.avg

                if args.use_wandb:
                    wandb.log(log_dict)


def main():
    args = get_args()
    if args.is_distributed != 1:
        os.environ["CUDA_VISIBLE_DEVICES"] = "2"
    taskName = ["joint", "alignment", "mlm", "mfm", "m_joint", "m_align", "decoder", "m_decoder", "-alignment"]
    args.taskName = np.array(taskName)[np.array(args.tasks) == 1]
    args.taskNum = (np.array(args.tasks) == 1).sum()
    print("-"*10, args.taskName, "-"*10)

    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.connect(("8.8.8.8", 80))
    this_ip = s.getsockname()[0]
    s.close()

    if args.ip_list: # multi-node
        ip_list = args.ip_list
        args.ip = ip_list[0]
    else: # single-node
        args.ip = this_ip
        ip_list = [args.ip]

    args.world_size = len(ip_list)  # 8
    for i, ip in enumerate(ip_list):
        if ip == this_ip:
            args.world_rank = i

    ngpus_per_node = torch.cuda.device_count()  # 8
    args.world_size = ngpus_per_node * args.world_size  # 8 * 8
    if args.is_distributed:
        print(args.world_size)
        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args)) # args.world_size or ngpus_per_node
    else:
        main_worker(0, ngpus_per_node, args)


if __name__ == "__main__":
    main()