#!/usr/bin/env python

# coding: utf-8

# In[]:

###########################  IMPORTS   ############################################# 

import torch

import torch.nn as nn

import torch.nn.functional as F

import monai

from monai.networks.nets import UNet,VNet,DynUNet,UNet_ENN_KMEANS

from monai.networks.utils import one_hot

from monai.transforms import (

    AsDiscrete,

    AddChanneld,

    AsChannelFirstd,

    Compose,

    LoadNiftid,

    RandCropByPosNegLabeld,

    RandRotate90d,

    ScaleIntensityd,

    ToTensord,

)

from monai.visualize import plot_2d_or_3d_image

from monai.data.utils import list_data_collate, worker_init_fn

from monai.inferers import sliding_window_inference

from monai.metrics import DiceMetric

from monai.metrics import compute_meandice

from torch.autograd import Variable

import torch.optim as optim

from torch.optim import lr_scheduler

from torch.utils.data import DataLoader

from torch.utils.data import Dataset

import torch.backends.cudnn as cudnn

import torchvision

from torchvision import datasets, models, transforms

import csv

import time

import SimpleITK as sitk

from os.path import splitext,basename

import random

from glob import glob

import matplotlib.pyplot as plt

from matplotlib.backends.backend_pdf import PdfPages

from copy import copy

import os

import numpy as np

from torch.utils.tensorboard import SummaryWriter

#from global_tools.tools import display_loading_bar

from class_modalities.transforms import LoadNifti, Roi2Mask, ResampleReshapeAlign, Sitk2Numpy, ConcatModality

from monai.utils import first, set_determinism

train_transforms = Compose(

    [  # read img + meta info

        LoadNifti(keys=["pet_img", "ct_img", "mask_img"]),

        Sitk2Numpy(keys=['pet_img', 'ct_img', 'mask_img']),

        ConcatModality(keys=['pet_img', 'ct_img']),

        AddChanneld(keys=["mask_img"]),  # Add channel to the first axis

        ToTensord(keys=["image", "mask_img"]),

    ])

# without data augmentation for validation

val_transforms = Compose(

    [  # read img + meta info

        LoadNifti(keys=["pet_img", "ct_img", "mask_img"]),

        Sitk2Numpy(keys=['pet_img', 'ct_img', 'mask_img']),

        ConcatModality(keys=['pet_img', 'ct_img']),

        AddChanneld(keys=["mask_img"]),  # Add channel to the first axis

        ToTensord(keys=["image", "mask_img"]),

    ])

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

base_path="/home/lab/hualing/2.5_SUV_dilation"

pet_path = base_path + '/' + 'pet_test'

ct_path = base_path + '/' + 'ct_test'

mask_path = base_path + '/' + 'pet_test_mask'

PET_ids = sorted(glob(os.path.join(pet_path, '*pet.nii')))

CT_ids = sorted(glob(os.path.join(ct_path, '*ct.nii')))

MASK_ids = sorted(glob(os.path.join(mask_path, '*mask.nii')))

data_dicts= zip(PET_ids, CT_ids, MASK_ids)

files=list(data_dicts)

train_files = [{"pet_img": PET, "ct_img": CT, 'mask_img': MASK} for  PET, CT, MASK in files[:138]]

val_files = [{"pet_img": PET, "ct_img": CT, 'mask_img': MASK} for  PET, CT, MASK in files[138:156]]

test_files = [{"pet_img": PET, "ct_img": CT, 'mask_img': MASK} for  PET, CT, MASK in files[156:]]

train_ds = monai.data.Dataset(data=train_files,transform=train_transforms)

val_ds = monai.data.Dataset(data=val_files,transform=val_transforms)

test_ds = monai.data.Dataset(data=test_files,transform=val_transforms)

train_loader = DataLoader(

        train_ds,

        batch_size=6,

        shuffle=True,

        num_workers=4,

        collate_fn=list_data_collate,

        pin_memory=torch.cuda.is_available(),)

val_loader = DataLoader(val_ds, batch_size=1, num_workers=4, collate_fn=list_data_collate)

test_loader = DataLoader(test_ds, batch_size=1, num_workers=4, collate_fn=list_data_collate)

trained_model_path="./ENN_best_metric_model_segmentation3d_dict_step1.pth" #####path to the ENN model after step 1

model = UNet_ENN_KMEANS(

        dimensions=3,  # 3D

        in_channels=2,

        out_channels=2,

        kernel_size=5,

        channels=(8,16, 32, 64,128),

        strides=(2, 2, 2, 2),

        num_res_units=2,).to(device)

model_dict = model.state_dict()

pre_dict = torch.load(trained_model_path)

pre_dict = {k: v for k, v in pre_dict.items() if k in model_dict}

model_dict.update(pre_dict)

model.load_state_dict(model_dict)

params = filter(lambda p: p.requires_grad, model.parameters())

for name, param in model.named_parameters():

    if param.requires_grad==True:

        print(name)  ####code to make sure the parameters from the whole model are optimized

optimizer = torch.optim.Adam(params, 1e-4)      ####finetune the whole models

dice_metric = monai.metrics.DiceMetric( include_background=False,reduction="mean")

scheduler=torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,'min',patience=10)

loss_function = monai.losses.DiceLoss(include_background=False,softmax=False,squared_pred=True,to_onehot_y=True)

# TODO : generate a learning rate scheduler

val_interval = 1

best_metric = -1

best_metric_epoch = -1

epoch_loss_values = list()

metric_values = list()

writer = SummaryWriter()

post_pred = AsDiscrete(argmax=True, to_onehot=True, n_classes=2)

post_label = AsDiscrete(to_onehot=True, n_classes=2)

for epoch in range(100):

    print("-" * 10)

    print(f"epoch {epoch + 1}/{100}")

    model.train()

    epoch_loss = 0

    step = 0

    for batch_data in train_loader:

        step += 1

        inputs, labels = batch_data["image"].to(device), batch_data["mask_img"].to(device)

        optimizer.zero_grad()

        outputs = model(inputs)

        output=outputs[:, :2, :, :, :]+0.5*outputs[:, 2, :, :, :].unsqueeze(1)

        dice_loss=loss_function(output, labels)

        loss = dice_loss

        loss.backward()

        optimizer.step()

        epoch_loss += loss.item()

        epoch_len = len(train_ds) // train_loader.batch_size

        print(f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")

        writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step)

    epoch_loss /= step

    epoch_loss_values.append(epoch_loss)

    print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

    scheduler.step(epoch_loss)

    if (epoch + 1) % val_interval == 0:

        model.eval()

        with torch.no_grad():

            metric_sum = 0.0

            metric_count = 0

            val_images = None

            val_labels = None

            val_outputs = None

            for val_data in val_loader:

                val_images, val_labels = val_data["image"].to(device), val_data["mask_img"].to(device)

                output = model(val_images)

                val_outputs = output[:, :2, :, :, :]+0.5*output[:, 2, :, :, :].unsqueeze(1)

                value = dice_metric(y_pred=val_outputs, y=val_labels)

                metric_count += len(value)

                metric_sum += value.item() * len(value)

            metric = metric_sum / metric_count

            metric_values.append(metric)

            if metric > best_metric:

                best_metric = metric

                best_metric_epoch = epoch + 1

                torch.save(model.state_dict(), "ENN_best_metric_model_segmentation3d_dict_step2.pth")

                print("saved new best metric model")

            print(

                "current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}".format(

                    epoch + 1, metric, best_metric, best_metric_epoch

                )

            )

print(f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}")

writer.close()

# LEARNING PROCEDURE

model.load_state_dict(torch.load("ENN_best_metric_model_segmentation3d_dict_step2.pth"))

model.eval()

###########################  RULES   #############################################

PREDICTION_VALIDATION_SET = True

#path_results = '/home/tongxue/huangling/code-hl/(IJAR)hl_medical-segmentation-master/result'

path_results='/home/lab/hualing/(IJAR_new)hl_medical-segmentation-master/result_enn_kmeans'

# generates folders

if not os.path.exists(path_results):

    os.makedirs(path_results)

##################

def PREDICT_MASK(data_set_ids, path_predictions, model):

    # generates folder

    if not os.path.exists(path_predictions):

        os.makedirs(path_predictions)

    filenames_predicted_masks = []

    n_patients = len(val_ds)

    val_loader = DataLoader(data_set_ids, batch_size=1, num_workers=4, collate_fn=list_data_collate)

    metric_sum = 0.0

    metric_sum_sen = 0.0

    metric_sum_spe = 0.0

    metric_sum_pre = 0.0

    metric_count =0

    os.chdir(r'/home/lab/hualing/(IJAR_new)hl_medical-segmentation-master/result_enn_kmeans')

    # for i,data_set_id in enumerate(data_set_ids):

    for i,val_data in enumerate(val_loader):

        val_images, val_labels = val_data["image"].to(device), val_data["mask_img"].to(device)

        prediction = model(val_images)

        pm=prediction

        #####save mass to .npy##########

        mass_out=prediction.data.cpu().numpy()

        name=splitext(basename(test_files[i]["mask_img"]))[0]

        val_outputs = prediction[:, :2, :, :, :]+0.5*prediction[:, 2, :, :, :].unsqueeze(1)

        np.save(name, mass_out)

        #####save results to .nii##########

        prediction = torch.argmax(prediction, axis=1)

        prediction=prediction.permute(0,3,1,2)# output from a multiclass softmax

        prediction = prediction.squeeze().cpu().numpy()

        # conversion in unsigned int 8 to store mask with less memory requirement

        mask = np.asarray(prediction, dtype=np.uint8)

        new_filename = path_predictions + "/pred_" + splitext(basename(test_files[i]["mask_img"]))[0] + '.nii'

        filenames_predicted_masks.append(new_filename)

        sitk.WriteImage(sitk.GetImageFromArray(mask), new_filename)

        ########calculate sen,spe,pre,acc,f1#########

        value = dice_metric(y_pred=val_outputs, y=val_labels)

        val_outputs=torch.argmax(pm, axis=1)

        val_outputs=val_outputs.unsqueeze(1)

        sensitivity = monai.metrics.compute_confusion_metric(y_pred=val_outputs, y=val_labels, to_onehot_y=False,

                                                             metric_name='sensitivity')

        specificity = monai.metrics.compute_confusion_metric(y_pred=val_outputs, y=val_labels, to_onehot_y=False,

                                                             metric_name='specificity')

        precision = monai.metrics.compute_confusion_metric(y_pred=val_outputs, y=val_labels, to_onehot_y=False,

                                                           metric_name='precision')

        print(len(value))

        metric_count += len(value)

        metric_sum += value.item() * len(value)

        metric_sum_sen += sensitivity.item() * len(value)

        metric_sum_spe += specificity.item() * len(value)

        metric_sum_pre += precision.item() * len(value)

    metric_dice = metric_sum / metric_count

    metric_sen = metric_sum_sen / metric_count

    metric_spe = metric_sum_spe / metric_count

    metric_pre = metric_sum_pre / metric_count

    print("dice:", metric_dice)

    print("sen:", metric_sen)

    print("spe",metric_spe)

    print("pre",metric_pre)

    return filenames_predicted_masks

    ####################################################################################################

###########################  Testing and visulation   #############################################

if PREDICTION_VALIDATION_SET:

    print("Prediction on validation set :")

    # use to fine tune and evaluate model performances

    print("Generating predictions :")

    valid_prediction_ids = PREDICT_MASK(data_set_ids=test_ds,

                                                 path_predictions=path_results + '/valid_predictions',

                                                 model=model)

    print("fini")