import os

import argparse

import logging

import importlib

import torch

import torch.nn as nn

import torch.optim as optim

import torch.distributed as dist

from torch.utils.data import DataLoader

import torchvision

from tensorboardX import SummaryWriter

import _init_paths

from libs.configs.config_acdc import cfg

from libs.datasets import AcdcDataset

from libs.datasets import joint_augment as joint_augment

from libs.datasets import augment as standard_augment

from libs.datasets.collate_batch import BatchCollator

# from libs.losses.df_loss import EuclideanLossWithOHEM

# from libs.losses.surface_loss import SurfaceLoss

from libs.losses.create_losses import Total_loss

import train_utils.train_utils as train_utils

from train_utils.train_utils import load_checkpoint

from utils.init_net import init_weights

from utils.comm import get_rank, synchronize

parser = argparse.ArgumentParser(description="arg parser")

parser.add_argument("--local_rank", type=int, default=0, required=True, help="device_ids of DistributedDataParallel")

parser.add_argument("--batch_size", type=int, default=32, required=False, help="batch size for training")

parser.add_argument("--epochs", type=int, default=50, required=False, help="Number of epochs to train for")

parser.add_argument('--workers', type=int, default=4, help='number of workers for dataloader')

parser.add_argument("--ckpt_save_interval", type=int, default=5, help="number of training epochs")

parser.add_argument('--output_dir', type=str, default=None, help='specify an output directory if needed')

parser.add_argument('--mgpus', type=str, default=None, help='whether to use multiple gpu')

parser.add_argument("--ckpt", type=str, default=None, help="continue training from this checkpoint")

parser.add_argument('--train_with_eval', action='store_true', default=False, help='whether to train with evaluation')

args = parser.parse_args()

FILE_DIR = os.path.dirname(os.path.abspath(__file__))

if args.mgpus is not None:

    os.environ["CUDA_VISIBLE_DEVICES"] = args.mgpus

def create_logger(log_file, dist_rank):

    if dist_rank > 0:

        logger = logging.getLogger(__name__)

        logger.setLevel(logging.WARNING)

        return logger

    log_format = '%(asctime)s  %(levelname)5s  %(message)s'

    logging.basicConfig(level=logging.DEBUG, format=log_format, filename=log_file)

    console = logging.StreamHandler()

    console.setLevel(logging.DEBUG)

    console.setFormatter(logging.Formatter(log_format))

    logging.getLogger(__name__).addHandler(console)

    return logging.getLogger(__name__)

def create_dataloader(logger):

    train_joint_transform = joint_augment.Compose([

                            joint_augment.To_PIL_Image(),

                            joint_augment.RandomAffine(0,translate=(0.125, 0.125)),

                            joint_augment.RandomRotate((-180,180)),

                            joint_augment.FixResize(256)

                            ])

    transform = standard_augment.Compose([

                    standard_augment.to_Tensor(),

                    standard_augment.normalize([cfg.DATASET.MEAN], [cfg.DATASET.STD])])

    target_transform = standard_augment.Compose([

                        standard_augment.to_Tensor()])

    if cfg.DATASET.NAME == 'acdc':

        train_set = AcdcDataset(data_list=cfg.DATASET.TRAIN_LIST,

                                df_used=cfg.DATASET.DF_USED, df_norm=cfg.DATASET.DF_NORM,

                                boundary=cfg.DATASET.BOUNDARY,

                                joint_augment=train_joint_transform,

                                augment=transform, target_augment=target_transform)

    train_sampler = torch.utils.data.distributed.DistributedSampler(train_set,

                            num_replicas=dist.get_world_size(), rank=dist.get_rank())

    train_loader = DataLoader(train_set, batch_size=args.batch_size, pin_memory=True,

                              num_workers=args.workers, shuffle=False, sampler=train_sampler,

                              collate_fn=BatchCollator(size_divisible=32, df_used=cfg.DATASET.DF_USED,

                                                        boundary=cfg.DATASET.BOUNDARY))

    if args.train_with_eval:

        eval_transform = joint_augment.Compose([

                         joint_augment.To_PIL_Image(),

                         joint_augment.FixResize(256),

                         joint_augment.To_Tensor()])

        evalImg_transform = standard_augment.Compose([

                            standard_augment.normalize([cfg.DATASET.MEAN], [cfg.DATASET.STD])])

        if cfg.DATASET.NAME == 'acdc':

            test_set = AcdcDataset(data_list=cfg.DATASET.TEST_LIST,

                                df_used=cfg.DATASET.DF_USED, df_norm=cfg.DATASET.DF_NORM,

                                boundary=cfg.DATASET.BOUNDARY,

                                joint_augment=eval_transform,

                                augment=evalImg_transform)

        test_sampler = torch.utils.data.distributed.DistributedSampler(test_set,

                            num_replicas=dist.get_world_size(), rank=dist.get_rank())

        test_loader = DataLoader(test_set, batch_size=args.batch_size, pin_memory=True,

                                 num_workers=args.workers, shuffle=False, sampler=test_sampler,

                                 collate_fn=BatchCollator(size_divisible=32, df_used=cfg.DATASET.DF_USED,

                                                           boundary=cfg.DATASET.BOUNDARY))

    else:

        test_loader = None

    return train_loader, test_loader

def create_optimizer(model):

    if cfg.TRAIN.OPTIMIZER == "adam":

        optimizer = optim.Adam(model.parameters(), lr=cfg.TRAIN.LR, weight_decay=cfg.TRAIN.WEIGHT_DECAY)

    elif cfg.TRAIN.OPTIMIZER == "sgd":

        optimizer = optim.SGD(model.parameters(), lr=cfg.TRAIN.LR, weight_decay=cfg.TRAIN.WEIGHT_DECAY,

                              momentum=cfg.TRAIN.MOMENTUM)

    else:

        raise NotImplementedError

    return optimizer

def create_scheduler(model, optimizer, total_steps, last_epoch):

    def lr_lbmd(cur_epoch):

        cur_decay = 1

        for decay_step in cfg.TRAIN.DECAY_STEP_LIST:

            if cur_epoch >= decay_step:

                cur_decay = cur_decay * cfg.TRAIN.LR_DECAY

        return max(cur_decay, cfg.TRAIN.LR_CLIP / cfg.TRAIN.LR)

    lr_scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lbmd, last_epoch=last_epoch)

    return lr_scheduler

def create_model(cfg):

    network = cfg.TRAIN.NET

    module = 'libs.network.' + network[:network.rfind('.')] 

    model = network[network.rfind('.')+1:]

    mod = importlib.import_module(module)

    mod_func = importlib.import_module('libs.network.train_functions')

    net_func = getattr(mod, model)

    net = net_func(num_class=cfg.DATASET.NUM_CLASS)

    if network == 'unet.U_Net':

        train_func = getattr(mod_func, 'model_fn_decorator')

    elif network == 'unet_df.U_NetDF':

        net = net_func(selfeat=cfg.MODEL.SELFEATURE, num_class=cfg.DATASET.NUM_CLASS, shift_n=cfg.MODEL.SHIFT_N, auxseg=cfg.MODEL.AUXSEG)

        train_func = getattr(mod_func, 'model_DF_decorator')

    return net, train_func

def train():

    torch.cuda.set_device(args.local_rank)

    dist.init_process_group(backend="nccl", init_method="env://")

    synchronize()

    # create dataloader & network & optimizer

    model, model_fn_decorator = create_model(cfg)

    init_weights(model, init_type='kaiming')

    # model.to('cuda')

    model.cuda()

    model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)

    root_result_dir = args.output_dir

    os.makedirs(root_result_dir, exist_ok=True)

    log_file = os.path.join(root_result_dir, "log_train.txt")

    logger = create_logger(log_file, get_rank())

    logger.info("**********************Start logging**********************")

    # log to file

    gpu_list = os.environ['CUDA_VISIBLE_DEVICES'] if 'CUDA_VISIBLE_DEVICES' in os.environ.keys() else 'ALL'

    logger.info("CUDA_VISIBLE_DEVICES=%s" % gpu_list)

    for key, val in vars(args).items():

        logger.info("{:16} {}".format(key, val))

    logger.info("***********************config infos**********************")

    for key, val in vars(cfg).items():

        logger.info("{:16} {}".format(key, val))

    # log tensorboard

    if get_rank() == 0:

        tb_log = SummaryWriter(log_dir=os.path.join(root_result_dir, "tensorboard"))

    else:

        tb_log = None

    train_loader, test_loader = create_dataloader(logger)

    optimizer = create_optimizer(model)

    # load checkpoint if it is possible

    start_epoch = it = best_res = 0

    last_epoch = -1

    if args.ckpt is not None:

        pure_model = model.module if isinstance(model, (torch.nn.DataParallel, torch.nn.parallel.DistributedDataParallel)) else model

        it, start_epoch, best_res = load_checkpoint(pure_model, optimizer, args.ckpt, logger)

        last_epoch = start_epoch + 1

    lr_scheduler = create_scheduler(model, optimizer, total_steps=len(train_loader)*args.epochs,

                                    last_epoch=last_epoch)

    if cfg.DATASET.DF_USED:

        criterion = Total_loss(boundary=cfg.DATASET.BOUNDARY)

    else:

        criterion = nn.CrossEntropyLoss()

    # start training

    logger.info('**********************Start training**********************')

    ckpt_dir = os.path.join(root_result_dir, "ckpt")

    os.makedirs(ckpt_dir, exist_ok=True)

    trainer = train_utils.Trainer(model,

                                  model_fn=model_fn_decorator(),

                                  criterion=criterion,

                                  optimizer=optimizer,

                                  ckpt_dir=ckpt_dir,

                                  lr_scheduler=lr_scheduler,

                                  model_fn_eval=model_fn_decorator(),

                                  tb_log=tb_log,

                                  logger=logger,

                                  eval_frequency=1,

                                  grad_norm_clip=cfg.TRAIN.GRAD_NORM_CLIP,

                                  cfg=cfg)

    trainer.train(start_it=it,

                  start_epoch=start_epoch,

                  n_epochs=args.epochs,

                  train_loader=train_loader,

                  test_loader=test_loader,

                  ckpt_save_interval=args.ckpt_save_interval,

                  lr_scheduler_each_iter=False,

                  best_res=best_res)

    logger.info('**********************End training**********************')

# python -m torch.distributed.launch --nproc_per_node 2 --master_port $RANDOM tools/train.py --batch_size 20 --mgpus 2,3 --output_dir logs/... --train_with_eval

if __name__ == "__main__":

    train()