#coding=utf-8
import argparse
import os
import time
import logging
import numpy as np
import wandb
import torch
import torch.optim
import sys
# from tensorboardX import SummaryWriter
from utils.random_seed import setup_seed
from IMFuse_no1skip import Model
from IMFuse import IMFuse
from data.transforms import *
from data.datasets_nii import Brats_loadall_nii, Brats_loadall_test_nii, Brats_loadall_val_nii
from data.data_utils import init_fn
from utils import Parser,criterions
from utils.parser import setup 
from utils.lr_scheduler import LR_Scheduler, record_loss, MultiEpochsDataLoader 
from torch.optim.lr_scheduler import CosineAnnealingLR
from predict import AverageMeter, test_softmax

DEBUG_ITER = 1

parser = argparse.ArgumentParser()
parser.add_argument('--batch_size', default=1, type=int, help='Batch size')
parser.add_argument('--datapath', default=None, type=str)
parser.add_argument('--dataname', default='BRATS2018', type=str)
parser.add_argument('--savepath', default=None, type=str)
parser.add_argument('--resume', default=None, type=str)
parser.add_argument('--pretrain', default=None, type=str)
parser.add_argument('--lr', default=2e-4, type=float)
parser.add_argument('--weight_decay', default=3e-5, type=float)
parser.add_argument('--num_epochs', default=1000, type=int)
parser.add_argument('--iter_per_epoch', default=150, type=int)
parser.add_argument('--region_fusion_start_epoch', default=0, type=int)
parser.add_argument('--seed', default=999, type=int)
parser.add_argument('--debug', action='store_true', default=False)
parser.add_argument('--interleaved_tokenization', action='store_true', default=False)
parser.add_argument('--mamba_skip', action='store_true', default=False)
parser.add_argument('--first_skip', action='store_true', default=False)
path = os.path.dirname(__file__)

## parse arguments
args = parser.parse_args()
setup(args, 'training')
args.train_transforms = 'Compose([RandCrop3D((128,128,128)), RandomRotion(10), RandomIntensityChange((0.1,0.1)), RandomFlip(0), NumpyType((np.float32, np.int64)),])'
args.test_transforms = 'Compose([NumpyType((np.float32, np.int64)),])'

ckpts = args.savepath
os.makedirs(ckpts, exist_ok=True)

###tensorboard writer
# writer = SummaryWriter(os.path.join(args.savepath, 'summary'))

###modality missing mask
masks = [[False, False, False, True], [False, True, False, False], [False, False, True, False], [True, False, False, False],
         [False, True, False, True], [False, True, True, False], [True, False, True, False], [False, False, True, True], [True, False, False, True], [True, True, False, False],
         [True, True, True, False], [True, False, True, True], [True, True, False, True], [False, True, True, True],
         [True, True, True, True]]
masks_torch = torch.from_numpy(np.array(masks))
mask_name = ['t2', 't1c', 't1', 'flair', 
            't1cet2', 't1cet1', 'flairt1', 't1t2', 'flairt2', 'flairt1ce',
            'flairt1cet1', 'flairt1t2', 'flairt1cet2', 't1cet1t2',
            'flairt1cet1t2']
print (masks_torch.int())

val_check = [50, 100, 150, 200, 300, 400, 500, 600, 700, 800, 850, 900, 910, 920, 930, 940, 950, 955, 960, 965, 970, 975, 980, 985, 990, 995, 1000] 
print(f"Validation checks: {val_check}")

def main():
    ##########setting seed
    setup_seed(args.seed)
    
    ##########print args
    for k, v in args._get_kwargs():
        pad = ' '.join(['' for _ in range(25-len(k))])
        print(f"{k}:{pad} {v}", flush=True)

    ##########init wandb
    slurm_job_id = os.getenv("SLURM_JOB_ID") 
    wandb_name_and_id = f'{args.dataname}_IMFuse{"no_1_skip" if not args.first_skip else ""}_{"Interleaved" if args.interleaved_tokenization else ""}{"Skip" if args.mamba_skip else ""}_jobid{slurm_job_id}'
    wandb_mode = 'online'
    # if args.debug:
    #     wandb_mode = 'disabled'
    wandb.init(
        project="SegmentationMM",
        name=wandb_name_and_id,
        # entity="NeuroTumor",
        id=wandb_name_and_id,
        mode=wandb_mode,
        resume="allow",
        config={
            "architecture": "IMFuse",
            "learning_rate": args.lr, 
            "batch_size": args.batch_size,
            "iter_per_epoch": args.iter_per_epoch,
            "num_epochs": args.num_epochs,
            "datapath": args.datapath,
            "region_fusion_start_epoch": args.region_fusion_start_epoch,
            "interleaved_tokenization": args.interleaved_tokenization,
            "mamba_skip": args.mamba_skip,
        }
    )
    
    ##########setting models
    if args.dataname in ['BRATS2023', 'BRATS2021', 'BRATS2020', 'BRATS2018']:
        num_cls = 4
    elif args.dataname == 'BRATS2015':
        num_cls = 5
    else:
        print ('dataset is error')
        exit(0)
    if args.first_skip:
        model = IMFuse(
                    num_cls=num_cls, 
                    interleaved_tokenization=args.interleaved_tokenization,
                    mamba_skip=args.mamba_skip,
            )
    else:
        model = Model(
                        num_cls=num_cls, 
                        interleaved_tokenization=args.interleaved_tokenization,
                        mamba_skip=args.mamba_skip,
                )
    print (model)
    model = torch.nn.DataParallel(model).cuda()

    ########## Setting learning scheduler and optimizer    
    train_params = [{'params': model.parameters(), 'lr': args.lr, 'weight_decay':args.weight_decay}]
    optimizer = torch.optim.RAdam(train_params)
    lr_schedule = LR_Scheduler(args.lr, args.num_epochs)

    ########## Setting data
    if args.dataname in ['BRATS2023', 'BRATS2020', 'BRATS2015']:
        train_file = 'datalist/train.txt'
        test_file = 'datalist/test15splits2.csv'
        val_file = 'datalist/val15splits.csv'
        #test_file = 'datalist/test.txt'
        #val_file = 'datalist/val.txt'
    elif args.dataname == 'BRATS2018':
        #### BRATS2018 contains three splits (1,2,3)
        test_file = 'datalist/Brats18_test15splits.csv'
        val_file = 'datalist/Brats18_val15splits.csv'
        train_file = 'datalist/train3.txt'

    logging.info(str(args))
    train_set = Brats_loadall_nii(transforms=args.train_transforms, 
                                    root=args.datapath, 
                                    num_cls=num_cls, 
                                    train_file=train_file)
    test_set = Brats_loadall_test_nii(transforms=args.test_transforms, 
                                    root=args.datapath, 
                                    test_file=test_file)
    val_set = Brats_loadall_val_nii(transforms=args.test_transforms, 
                                    root=args.datapath, 
                                    num_cls=num_cls, 
                                    val_file=val_file)
    train_loader = MultiEpochsDataLoader(
        dataset=train_set,
        batch_size=args.batch_size,
        num_workers=8,
        pin_memory=True,
        shuffle=True,
        worker_init_fn=init_fn)
    test_loader = MultiEpochsDataLoader(
        dataset=test_set,
        batch_size=1,
        shuffle=False,
        num_workers=0,
        pin_memory=True)
    val_loader = MultiEpochsDataLoader(
        dataset=val_set,
        batch_size=1,
        shuffle=False,
        num_workers=0,
        pin_memory=True)

    ##########Training
    start = time.time()
    torch.set_grad_enabled(True)
    logging.info('#############training############')
    # iter_per_epoch = args.iter_per_epoch
    iter_per_epoch = len(train_loader) #number of batches
    train_iter = iter(train_loader)
    val_Dice_best = -999999
    start_epoch = 0

    ##########Resume Training
    if args.resume is not None:
        checkpoint = torch.load(args.resume)
        logging.info('best epoch: {}'.format(checkpoint['epoch']))
        model.load_state_dict(checkpoint['state_dict'])
        val_Dice_best = checkpoint['val_Dice_best']
        optimizer.load_state_dict(checkpoint['optim_dict'])
        start_epoch = checkpoint['epoch'] + 1

    for epoch in range(start_epoch, args.num_epochs):
        step_lr = lr_schedule(optimizer, epoch)
        # writer.add_scalar('lr', step_lr, global_step=(epoch+1))
        b = time.time()
        model.train()
        model.module.is_training = True

        prm_cross_loss_epoch = 0.0
        prm_dice_loss_epoch = 0.0
        fuse_cross_loss_epoch = 0.0
        fuse_dice_loss_epoch = 0.0
        sep_cross_loss_epoch = 0.0
        sep_dice_loss_epoch = 0.0
        loss_epoch = 0.0

        ########## training epoch
        for i in range(iter_per_epoch):
            step = (i+1) + epoch*iter_per_epoch
            ###Data load
            try:
                data = next(train_iter)
            except:
                train_iter = iter(train_loader)
                data = next(train_iter)
            x, target, mask = data[:3] #x=(B, M=4, 128, 128, 128), target = (B, C, 128, 128, 128), mask = (B, 4)
            x = x.cuda(non_blocking=True)
            target = target.cuda(non_blocking=True)
            mask = mask.cuda(non_blocking=True)

            fuse_pred, sep_preds, prm_preds = model(x, mask)

            ###Loss compute
            ### fuse modality segmentation loss
            fuse_cross_loss = criterions.softmax_weighted_loss(fuse_pred, target, num_cls=num_cls)
            fuse_dice_loss = criterions.dice_loss(fuse_pred, target, num_cls=num_cls)
            fuse_loss = fuse_cross_loss + fuse_dice_loss

            fuse_cross_loss_epoch += fuse_cross_loss
            fuse_dice_loss_epoch += fuse_dice_loss

            ### separated modality segmentation loss
            sep_cross_loss = torch.zeros(1).cuda().float()
            sep_dice_loss = torch.zeros(1).cuda().float()
            for sep_pred in sep_preds:
                sep_cross_loss += criterions.softmax_weighted_loss(sep_pred, target, num_cls=num_cls)
                sep_dice_loss += criterions.dice_loss(sep_pred, target, num_cls=num_cls)
            sep_loss = sep_cross_loss + sep_dice_loss

            sep_cross_loss_epoch += sep_cross_loss
            sep_dice_loss_epoch += sep_dice_loss

            ### pyramid segmentation loss
            prm_cross_loss = torch.zeros(1).cuda().float()
            prm_dice_loss = torch.zeros(1).cuda().float()
            for prm_pred in prm_preds:
                prm_cross_loss += criterions.softmax_weighted_loss(prm_pred, target, num_cls=num_cls)
                prm_dice_loss += criterions.dice_loss(prm_pred, target, num_cls=num_cls)
            prm_loss = prm_cross_loss + prm_dice_loss

            prm_cross_loss_epoch += prm_cross_loss
            prm_dice_loss_epoch += prm_dice_loss

            ### total segmentation loss
            if epoch < args.region_fusion_start_epoch:
                loss = fuse_loss * 0.0 + sep_loss + prm_loss
            else:
                loss = fuse_loss + sep_loss + prm_loss

            loss_epoch += loss

            ### backpropagation
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            ###log
            msg = 'Epoch {}/{}, Iter {}/{}, Loss {:.4f}, '.format((epoch+1), args.num_epochs, (i+1), iter_per_epoch, loss.item())
            msg += 'fusecross:{:.4f}, fusedice:{:.4f},'.format(fuse_cross_loss.item(), fuse_dice_loss.item())
            msg += 'sepcross:{:.4f}, sepdice:{:.4f},'.format(sep_cross_loss.item(), sep_dice_loss.item())
            msg += 'prmcross:{:.4f}, prmdice:{:.4f},'.format(prm_cross_loss.item(), prm_dice_loss.item())
            logging.info(msg)

            """
            if args.debug:
                break
            """

        logging.info('train time per epoch: {}'.format(time.time() - b))

        ########## log current epoch metrics and save current model 
        wandb.log({
            "train/epoch": epoch,
            "train/loss": loss_epoch.cpu().detach().item() / iter_per_epoch,
            "train/fusecross": fuse_cross_loss_epoch.cpu().detach().item() / iter_per_epoch,
            "train/fusedice": fuse_dice_loss_epoch.cpu().detach().item() / iter_per_epoch,
            "train/sepcross": sep_cross_loss_epoch.cpu().detach().item() / iter_per_epoch,
            "train/sepdice": sep_dice_loss_epoch.cpu().detach().item() / iter_per_epoch,
            "train/prmcross": prm_cross_loss_epoch.cpu().detach().item() / iter_per_epoch,
            "train/prmdice": prm_dice_loss_epoch.cpu().detach().item() / iter_per_epoch,
            "train/learning_rate": step_lr,
        })

        file_name = os.path.join(ckpts, 'model_last.pth')
        torch.save({
            'epoch': epoch,
            'state_dict': model.state_dict(),
            'optim_dict': optimizer.state_dict(),
            'val_Dice_best': val_Dice_best,
            },
            file_name)
        
        ########## validation and test
        if epoch+1 in val_check or args.debug:
            print('validate ...')
            with torch.no_grad():
                dice_score, seg_loss = test_softmax(
                    val_loader,
                    model,
                    dataname = args.dataname)
        
            val_WT, val_TC, val_ET, val_ETpp = dice_score #validate(model, val_loader)
            logging.info('Validate epoch = {}, WT = {:.2}, TC = {:.2}, ET = {:.2}, ETpp = {:.2}, loss = {:.2}'.format(epoch, val_WT.item(), val_TC.item(), val_ET.item(), val_ETpp.item(), seg_loss.cpu().item()))
            val_dice = (val_ET + val_WT + val_TC)/3
            wandb.log({
                "val/epoch":epoch,
                "val/val_ET_Dice": val_ET.item(),
                "val/val_ETpp_Dice": val_ETpp.item(),
                "val/val_WT_Dice": val_WT.item(),
                "val/val_TC_Dice": val_TC.item(),
                "val/val_Dice": val_dice.item(), 
                "val/seg_loss": seg_loss.cpu().item(),       
            })
            
            if val_dice > val_Dice_best:
                val_Dice_best = val_dice.item()
                print('save best model ...')
                file_name = os.path.join(ckpts, 'best.pth')
                torch.save({
                    'epoch': epoch,
                    'state_dict': model.state_dict(),
                    'optim_dict': optimizer.state_dict(),
                    'val_Dice_best': val_Dice_best,
                    },
                    file_name)
                
            print('testing ...')
            with torch.no_grad():
                dice_score, seg_loss = test_softmax(
                    test_loader,
                    model,
                    dataname = args.dataname)
            test_WT, test_TC, test_ET, test_ETpp = dice_score   
            logging.info('Testing epoch = {}, WT = {:.2}, TC = {:.2}, ET = {:.2}, ET_postpro = {:.2}'.format(epoch, test_WT.item(), test_TC.item(), test_ET.item(), test_ETpp.item()))
            test_dice = (test_ET + test_WT + test_TC)/3
            wandb.log({
                "test/epoch":epoch,
                "test/test_WT_Dice": test_WT.item(),
                "test/test_TC_Dice": test_TC.item(),
                "test/test_ET_Dice": test_ET.item(),
                "test/test_ETpp": test_ETpp.item(),
                "test/test_Dice": test_dice.item(),  
                "test/seg_loss": seg_loss.cpu().item(),   
            })

            model.train()
            model.module.is_training=True


    msg = 'total time: {:.4f} hours'.format((time.time() - start)/3600)
    logging.info(msg)

    ##########Evaluate the last epoch model
    """
    test_score = AverageMeter()
    with torch.no_grad():
        logging.info('###########test set wi/wo postprocess###########')
        for i, mask in enumerate(masks):
            logging.info('{}'.format(mask_name[i]))
            dice_score = test_softmax(
                            test_loader,
                            model,
                            dataname = args.dataname,
                            feature_mask = mask)
            test_score.update(dice_score)
        logging.info('Avg scores: {}'.format(test_score.avg))
    """

if __name__ == '__main__':
    main()
