#!/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" #####the path you put the pre_processed data.

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="./pre-trained_model/best_metric_model_unet.pth"      ######path to the pretrained UNet model

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 only the parameters from ENN are optimized

optimizer = torch.optim.Adam(params, 1e-2)    ####optimize only the ENN

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_step1.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()