import argparse

import gc

import importlib

import os

import sys

import shutil

import numpy as np

import pandas as pd

import torch

from torch import nn

from monai.handlers.utils import from_engine

from monai.inferers import sliding_window_inference

from monai.data import decollate_batch

from torch.cuda.amp import GradScaler, autocast

from tqdm import tqdm

from utils import *

from monai.transforms import (

    Compose,

    Activations,

    AsDiscrete,

    Activationsd,

    AsDiscreted,

    KeepLargestConnectedComponentd,

    Invertd,

    LoadImage,

    Transposed,

)

import json

from metric import HausdorffScore

from monai.utils import set_determinism

from monai.losses import DiceLoss, DiceCELoss

from monai.networks.nets import UNet, SegResNet, DynUnet

from monai.optimizers import Novograd

from monai.metrics import DiceMetric

torch.backends.cudnn.enabled = True

torch.backends.cudnn.benchmark = True

def main(cfg):

    # data sequence

    if cfg.fold != -1:

        cfg.data_json_dir = cfg.data_dir + f"dataset_3d_fold_{cfg.fold}.json"

    else:

        cfg.data_json_dir = cfg.data_dir + f"dataset_3d_all.json"

    with open(cfg.data_json_dir, "r") as f:

        cfg.data_json = json.load(f)

    if cfg.fold != -1:

        fold_dir = f"fold{cfg.fold}"

    else:

        fold_dir = "all"

    os.makedirs(str(cfg.output_dir + f"/{fold_dir}/"), exist_ok=True)

    # # set random seed

    # set_determinism(cfg.seed)

    train_dataset = get_train_dataset(cfg)

    train_dataloader = get_train_dataloader(train_dataset, cfg)

    val_dataset = get_val_dataset(cfg)

    val_dataloader = get_val_dataloader(val_dataset, cfg)

    print(f"run fold {cfg.fold}, train len: {len(train_dataset)}")

    if cfg.model_type.startswith("segres"):

        model = SegResNet(

            spatial_dims = 3,

            in_channels = 1,

            out_channels = 3,

            init_filters = int(cfg.model_type.replace("segres", "")),

            norm = "BATCH",

            act = "PRELU"

        ).to(cfg.device)

    print(cfg.weights)

    if cfg.weights is not None:

        stt = torch.load(cfg.weights, map_location = "cpu")

        if "model" in stt:

            stt = stt["model"]

        if "state_dict" in stt:

            stt = stt["state_dict"]

            del stt["out.conv.conv.weight"], stt["out.conv.conv.bias"]

        model.load_state_dict(stt, strict = False)

        print(f"weights from: {cfg.weights} are loaded.")

    # set optimizer, lr scheduler

    total_steps = len(train_dataset)

    optimizer = get_optimizer(model, cfg)

    # optimizer = Novograd(model.parameters(), cfg.lr)

    scheduler = get_scheduler(cfg, optimizer, total_steps)

    seg_loss_func = DiceBceMultilabelLoss(w_dice=cfg.w_dice, w_bce=1-cfg.w_dice)

    # seg_loss_func = DiceLoss(sigmoid=True, smooth_nr=0.01, smooth_dr=0.01, include_background=True, batch=True)

    dice_metric = DiceMetric(reduction="mean")

    hausdorff_metric = HausdorffScore(reduction="mean")

    metric_function = [dice_metric, hausdorff_metric]

    post_pred = Compose([

        Activations(sigmoid=True),

        AsDiscrete(threshold=0.5),

    ])

    # train and val loop

    step = 0

    i = 0

    if cfg.eval is True:

        best_val_metric = run_eval(

            model=model,

            val_dataloader=val_dataloader,

            post_pred=post_pred,

            metric_function=metric_function,

            seg_loss_func=seg_loss_func,

            cfg=cfg,

            epoch=0,

        )

    else:

        best_val_metric = 0.0

    best_weights_name = "best_weights"

    for epoch in range(cfg.epochs):

        print("EPOCH:", epoch)

        gc.collect()

        if cfg.train is True:

            run_train(

                model=model,

                train_dataloader=train_dataloader,

                optimizer=optimizer,

                scheduler=scheduler,

                seg_loss_func=seg_loss_func,

                cfg=cfg,

                # writer=writer,

                epoch=epoch,

                step=step,

                iteration=i,

            )

        if (epoch + 1) % cfg.eval_epochs == 0 and cfg.eval is True and epoch > cfg.start_eval_epoch:

            val_metric = run_eval(

                model=model,

                val_dataloader=val_dataloader,

                post_pred=post_pred,

                metric_function=metric_function,

                seg_loss_func=seg_loss_func,

                cfg=cfg,

                epoch=epoch,

            )

            if val_metric > best_val_metric:

                print(f"Find better metric: val_metric {best_val_metric:.5} -> {val_metric:.5}")

                best_val_metric = val_metric

                checkpoint = create_checkpoint(

                    model,

                    optimizer,

                    epoch,

                    scheduler=scheduler,

                )

                torch.save(

                    checkpoint,

                    f"{cfg.output_dir}/{fold_dir}/{best_weights_name}.pth",

                )

            else:

                if cfg.load_best_weights is True:

                    try:

                        model.load_state_dict(torch.load(f"{cfg.output_dir}/{fold_dir}/{best_weights_name}.pth")["model"])

                        print(f"metric no improve, load the saved best weights with score: {best_val_metric}.")

                    except:

                        pass

        if (epoch + 1) == cfg.epochs:

            # save final best weights, with its distinct name in order to avoid mistakes.

            if os.path.exists(f"{cfg.output_dir}/{fold_dir}/{best_weights_name}.pth"):

                shutil.copyfile(

                    f"{cfg.output_dir}/{fold_dir}/{best_weights_name}.pth",

                    f"{cfg.output_dir}/{fold_dir}/{best_weights_name}_{best_val_metric:.4f}.pth",

                )

        torch.save(

            model.state_dict(),

            f"{cfg.output_dir}/{fold_dir}/last.pth",

        )

def run_train(

    model,

    train_dataloader,

    optimizer,

    scheduler,

    seg_loss_func,

    cfg,

    # writer,

    epoch,

    step,

    iteration,

):

    model.train()

    scaler = GradScaler()

    progress_bar = tqdm(range(len(train_dataloader)))

    tr_it = iter(train_dataloader)

    dataset_size = 0

    running_loss = 0.0

    for itr in progress_bar:

        iteration += 1

        batch = next(tr_it)

        inputs, masks = (

            batch["image"].to(cfg.device),

            batch["mask"].to(cfg.device),

        )

        step += cfg.batch_size

        if cfg.amp is True:

            with autocast():

                outputs = model(inputs)

                loss = seg_loss_func(outputs, masks)

        else:

            outputs = model(inputs)

            loss = seg_loss_func(outputs, masks)

        if cfg.amp is True:

            scaler.scale(loss).backward()

            torch.nn.utils.clip_grad_norm_(model.parameters(), 12)

            scaler.step(optimizer)

            scaler.update()

        else:

            loss.backward()

            optimizer.step()

        optimizer.zero_grad()

        scheduler.step()

        running_loss += (loss.item() * cfg.batch_size)

        dataset_size += cfg.batch_size

        losses = running_loss / dataset_size

        progress_bar.set_description(f"loss: {losses:.4f} lr: {optimizer.param_groups[0]['lr']:.6f}")

        del batch, inputs, masks, outputs, loss

    print(f"Train loss: {losses:.4f}")

    torch.cuda.empty_cache()

def run_eval(model, val_dataloader, post_pred, metric_function, seg_loss_func, cfg, epoch):

    model.eval()

    dice_metric, hausdorff_metric = metric_function

    progress_bar = tqdm(range(len(val_dataloader)))

    val_it = iter(val_dataloader)

    with torch.no_grad():

        for itr in progress_bar:

            batch = next(val_it)

            val_inputs, val_masks = (

                batch["image"].to(cfg.device),

                batch["mask"].to(cfg.device),

            )

            if cfg.val_amp is True:

                with autocast():

                    val_outputs = sliding_window_inference(val_inputs, cfg.roi_size, cfg.sw_batch_size, model)

            else:

                val_outputs = sliding_window_inference(val_inputs, cfg.roi_size, cfg.sw_batch_size, model)

            # cal metric

            if cfg.run_tta_val is True:

                tta_ct = 1

                for dims in [[2],[3],[2,3]]:

                    flip_val_outputs = sliding_window_inference(torch.flip(val_inputs, dims=dims), cfg.roi_size, cfg.sw_batch_size, model)

                    val_outputs += torch.flip(flip_val_outputs, dims=dims)

                    tta_ct += 1

                val_outputs /= tta_ct

            val_outputs = [post_pred(i) for i in val_outputs]

            val_outputs = torch.stack(val_outputs)

            # metric is slice level put (n, c, h, w, d) to (n, d, c, h, w) to (n*d, c, h, w)

            val_outputs = val_outputs.permute([0, 4, 1, 2, 3]).flatten(0, 1)

            val_masks = val_masks.permute([0, 4, 1, 2, 3]).flatten(0, 1)

            hausdorff_metric(y_pred=val_outputs, y=val_masks)

            dice_metric(y_pred=val_outputs, y=val_masks)

            del val_outputs, val_inputs, val_masks, batch

    dice_score = dice_metric.aggregate().item()

    hausdorff_score = hausdorff_metric.aggregate().item()

    dice_metric.reset()

    hausdorff_metric.reset()

    all_score = dice_score * 0.4 + hausdorff_score * 0.6

    print(f"dice_score: {dice_score} hausdorff_score: {hausdorff_score} all_score: {all_score}")

    torch.cuda.empty_cache()

    return all_score

if __name__ == "__main__":

    sys.path.append("configs")

    parser = argparse.ArgumentParser(description="")

    parser.add_argument("-c", "--config", default="cfg_unet_multilabel", help="config filename")

    parser.add_argument("-f", "--fold", type=int, default=0, help="fold")

    parser.add_argument("-s", "--seed", type=int, default=20220421, help="seed")

    parser.add_argument("-w", "--weights", default=None, help="the path of weights")

    parser_args, _ = parser.parse_known_args(sys.argv)

    cfg = importlib.import_module(parser_args.config).cfg

    cfg.fold = parser_args.fold

    cfg.seed = parser_args.seed

    cfg.weights = parser_args.weights

    main(cfg)