import os

import time

import argparse

import pathlib

import tqdm

import cv2 as cv

import numpy as np

import pandas as pd

import tensorflow as tf

from tensorflow import keras

import matplotlib.pyplot as plt

from tensorflow.compat.v1 import ConfigProto

from tensorflow.compat.v1 import InteractiveSession

# Custom package

import data

import loss

import utils

import module

import performance

from model import dr_unet

config = ConfigProto()

config.gpu_options.allow_growth = True

# 1. Parameter settings

parser = argparse.ArgumentParser(description="Segment Use Args")

parser.add_argument('--model-name', default='DR_UNet', type=str)

parser.add_argument('--dims', default=32, type=int)

parser.add_argument('--epochs', default=50, type=int)

parser.add_argument('--batch-size', default=16, type=int)

parser.add_argument('--lr', default=2e-4, type=float)

# Training data, testing, verification parameter settings

parser.add_argument('--height', default=256, type=int)

parser.add_argument('--width', default=256, type=int)

parser.add_argument('--channel', default=1, type=int)

parser.add_argument('--pred-height', default=4 * 256, type=int)

parser.add_argument('--pred-width', default=4 * 256, type=int)

parser.add_argument('--total-samples', default=5000, type=int)

parser.add_argument('--invalid-samples', default=1000, type=int)

parser.add_argument('--regularize', default=False, type=bool)

parser.add_argument('--record-dir', default=r'', type=str, help='the save dir of tfrecord')

parser.add_argument('--train-record-name', type=str, default=r'train_data', help='the train record save name')

parser.add_argument('--test-image-dir', default=r'', type=str, help='the path of test images dir')

parser.add_argument('--invalid-record-name', type=str, default=r'test_data', help='the invalid record save name')

parser.add_argument('--gt-mask-dir', default=r'', type=str, help='the ground truth dir of validation set')

parser.add_argument('--invalid-volume-dir', default=r'', type=str, help='estimation bleeding volume')

args = parser.parse_args()

class Segmentation:

    def __init__(self, params):

        self.params = params

        self.input_shape = [params.height, params.width, params.channel]

        self.mask_shape = [params.height, params.width, 1]

        self.model_name = params.model_name

        self.crop_height = params.pred_height

        self.crop_width = params.pred_width

        self.regularize = params.regularize

        # Obtain a segmentation model

        self.seg_model = dr_unet.dr_unet(input_shape=self.input_shape, dims=params.dims)

        self.seg_model.summary()

        # Optimization function

        self.optimizer = tf.keras.optimizers.Adam(lr=params.lr)

        # Every epoch, predict invalid-images to test the segmentation performance of the model

        self.save_dir = str(params.model_name).upper()

        self.weight_save_dir = os.path.join(self.save_dir, 'checkpoint')

        self.pred_invalid_save_dir = os.path.join(self.save_dir, 'invalid_pred')

        self.invalid_crop_save_dir = os.path.join(self.save_dir, 'invalid_pred_crop')

        self.pred_test_save_dir = os.path.join(self.save_dir, 'test_pred')

        utils.check_file([

            self.save_dir, self.weight_save_dir, self.pred_invalid_save_dir,

            self.pred_test_save_dir, self.invalid_crop_save_dir]

        )

        # Save model parameters

        train_steps = tf.Variable(0, tf.int32)

        self.save_ckpt = tf.train.Checkpoint(

            train_steps=train_steps, seg_model=self.seg_model, model_optimizer=self.optimizer)

        self.save_manger = tf.train.CheckpointManager(

            self.save_ckpt, directory=self.weight_save_dir, max_to_keep=1)

        # Set the loss function

        self.loss_fun = loss.bce_dice_loss

    def load_model(self):

        if self.save_manger.latest_checkpoint:

            self.save_ckpt.restore(self.save_manger.latest_checkpoint)

            print('Loading model: {}'.format(self.save_manger.latest_checkpoint))

        else:

            print('Retrain the model！')

        return

    @tf.function

    def train_step(self, inputs, target):

        tf.keras.backend.set_learning_phase(True)

        with tf.GradientTape() as tape:

            pred_mask = self.seg_model(inputs)

            loss = self.loss_fun(target, pred_mask)

            if self.regularize:

                loss = tf.reduce_sum(loss) + tf.reduce_sum(self.seg_model.losses)

        gradient = tape.gradient(loss, self.seg_model.trainable_variables)

        self.optimizer.apply_gradients(zip(gradient, self.seg_model.trainable_variables))

        return tf.reduce_mean(loss)

    @tf.function

    def inference(self, inputs):

        tf.keras.backend.set_learning_phase(True)

        pred = self.seg_model(inputs)

        return pred

    @staticmethod

    def calculate_volume_by_mask(mask_dir, save_dir, model_name, dpi=96, thickness=0.45):

        all_mask_file_paths = utils.list_file(mask_dir)

        pd_record = pd.DataFrame(columns=['file_name', 'Volume'])

        for file_dir in tqdm.tqdm(all_mask_file_paths):

            file_name = pathlib.Path(file_dir).stem

            each_blood_volume = module.calculate_volume(file_dir, thickness=thickness, dpi=dpi)

            pd_record = pd_record.append({'file_name': file_name, 'Volume': each_blood_volume}, ignore_index=True)

            pd_record.to_csv(

                os.path.join(save_dir, '{}_{}.csv'.format(model_name, file_name)), index=True, header=True)

        return

    def predict_blood_volume(self, input_dir, save_dir, calc_nums=-1, dpi=96, thickness=0.45):

        """

         :param input_dir: The directory for testing bleeding volume images,

                           there are multiple folders under the directory, each folder represents a CT image of a patient

         :param save_dir: The predicted segmented image save directory

         :param calc_nums: predict how many images in the folder

         :param dpi: image parameters

         :param thickness: slice thickness

        """

        # Loading weights of model

        self.load_model()

        save_pred_images_dir = os.path.join(save_dir, 'pred_images')

        save_pred_csv_dir = os.path.join(save_dir, 'pred_csv')

        utils.check_file([save_pred_images_dir, save_pred_csv_dir])

        all_file_dirs = utils.list_file(input_dir)

        cost_time_list = []

        total_images = 0

        for file_dir in tqdm.tqdm(all_file_dirs):

            file_name = pathlib.Path(file_dir).stem

            image_names, ori_images, normed_images = data.get_test_data(

                test_data_path=file_dir, image_shape=self.input_shape, image_nums=calc_nums)

            total_images += len(image_names)

            start_time = time.time()

            pred_mask = self.inference(normed_images)

            end_time = time.time()

            print('FPS: {}'.format(pred_mask.shape[0] / (end_time - start_time)),

                  pred_mask.shape, end_time - start_time)

            denorm_pred_mask = module.reverse_pred_image(pred_mask.numpy())  # (image_nums, 256, 256, 1)

            if denorm_pred_mask.ndim == 2 and self.input_shape[-1] == 1:

                denorm_pred_mask = np.expand_dims(denorm_pred_mask, 0)

            drawed_images = []

            blood_areas = []

            pd_record = pd.DataFrame(columns=['image_name', 'Square Centimeter', 'Volume'])

            for index in range(denorm_pred_mask.shape[0]):

                drawed_image, blood_area = module.draw_contours(ori_images[index], denorm_pred_mask[index], dpi=dpi)

                drawed_images.append(drawed_image)

                blood_areas.append(blood_area)

                pd_record = pd_record.append({'image_name': image_names[index], 'Square Centimeter': blood_area},

                                             ignore_index=True)

            one_pred_save_dir = os.path.join(save_pred_images_dir, file_name)

            module.save_invalid_data(ori_images, drawed_images, denorm_pred_mask,

                                     image_names, reshape=True, save_dir=one_pred_save_dir)

            # Calculate the amount of bleeding based on the area of each layer of hematoma

            blood_volume = module.count_volume(blood_areas, thickness=thickness)

            pd_record = pd_record.append({'Volume': blood_volume}, ignore_index=True)

            pd_record.to_csv(os.path.join(save_pred_csv_dir, '{}_{}.csv'.format(self.model_name, file_name)),

                             index=True, header=True)

            cost_time_list.append(end_time - start_time)

            print('FileName: {} time: {}'.format(file_name, end_time - start_time))

        print('total_time: {:.2f}, mean_time: {:.2f}, total_images: {}'.format(

            np.sum(cost_time_list), np.mean(cost_time_list), total_images))

        return

    def predict_and_save(self, input_dir, save_dir, calc_nums=-1, batch_size=16):

        """ predict bleeding image and save

         :param input_dir: There are several images waiting to be tested under the input_dir folder

         :param save_dir: The file directory where the segmented image predicted by the model is saved

         :param calc_nums: How many images are taken from the directory to participate in the calculation

         :param batch_size: how many images to test each time

         :return:

        """

        mask_save_dir = os.path.join(save_dir, 'pred_mask')

        drawed_save_dir = os.path.join(save_dir, 'drawed_image')

        utils.check_file([mask_save_dir, drawed_save_dir])

        self.load_model()

        test_image_list = utils.list_file(input_dir)

        for index in range(len(test_image_list) // 128 + 1):

            input_test_list = test_image_list[index * 128:(index + 1) * 128]

            image_names, ori_images, normed_images = data.get_test_data(

                test_data_path=input_test_list, image_shape=self.input_shape, image_nums=-1,

            )

            if calc_nums != -1:

                ori_images = ori_images[:calc_nums]

                normed_images = normed_images[:calc_nums]

                image_names = image_names[:calc_nums]

            inference_times = normed_images.shape[0] // batch_size + 1

            for inference_time in range(inference_times):

                this_normed_images = normed_images[

                                     inference_time * batch_size:(inference_time + 1) * batch_size, ...]

                this_ori_images = ori_images[

                                  inference_time * batch_size:(inference_time + 1) * batch_size, ...]

                this_image_names = image_names[

                                   inference_time * batch_size:(inference_time + 1) * batch_size]

                this_pred_mask = self.inference(this_normed_images)

                this_denorm_pred_mask = module.reverse_pred_image(this_pred_mask.numpy())

                if ori_images.shape[0] == 1:

                    this_denorm_pred_mask = np.expand_dims(this_denorm_pred_mask, 0)

                for i in range(this_denorm_pred_mask.shape[0]):

                    bin_denorm_pred_mask = this_denorm_pred_mask[i]

                    this_drawed_image, this_blood_area = module.draw_contours(

                        this_ori_images[i], bin_denorm_pred_mask, dpi=96

                    )

                    cv.imwrite(os.path.join(

                        mask_save_dir, '{}'.format(this_image_names[i])), bin_denorm_pred_mask

                    )

                    cv.imwrite(os.path.join(

                        drawed_save_dir, '{}'.format(this_image_names[i])), this_drawed_image

                    )

        return

    def train(self, start_epoch=1):

        # get training dataset

        train_data = data.get_tfrecord_data(

            self.params.record_dir, self.params.train_record_name,

            self.input_shape, batch_size=self.params.batch_size)

        self.load_model()

        pd_record = pd.DataFrame(columns=['Epoch', 'Iteration', 'Loss', 'Time'])

        data_name, original_test_image, norm_test_image = data.get_test_data(

            test_data_path=self.params.test_image_dir, image_shape=self.input_shape, image_nums=-1

        )

        start_time = time.time()

        best_dice = 0.0

        for epoch in range(start_epoch, self.params.epochs):

            for train_image, gt_mask in tqdm.tqdm(

                    train_data, total=self.params.total_samples // self.params.batch_size):

                self.save_ckpt.train_steps.assign_add(1)

                iteration = self.save_ckpt.train_steps.numpy()

                # training step

                train_loss = self.train_step(train_image, gt_mask)

                if iteration % 100 == 0:

                    print('Epoch: {}, Iteration: {}, Loss: {:.2f}, Time: {:.2f} s'.format(

                        epoch, iteration, train_loss, time.time() - start_time))

                    # test step

                    test_pred = self.inference(norm_test_image)

                    module.save_images(

                        image_shape=self.mask_shape, pred=test_pred,

                        save_path=self.pred_test_save_dir, index=iteration, split=False

                    )

                    pd_record = pd_record.append({

                        'Epoch': epoch, 'Iteration': iteration, 'Loss': train_loss.numpy(),

                        'Time': time.time() - start_time}, ignore_index=True

                    )

                    pd_record.to_csv(os.path.join(

                        self.save_dir, '{}_record.csv'.format(self.params.model_name)), index=True, header=True

                    )

            m_dice = self.invalid(epoch)

            if m_dice > best_dice:

                best_dice = m_dice

                print('Best Dice:{}'.format(best_dice))

                self.save_manger.save(checkpoint_number=epoch)

        return

    def invalid(self, epoch):

        invalid_data = data.get_tfrecord_data(

            self.params.record_dir, self.params.invalid_record_name,

            self.input_shape, batch_size=self.params.batch_size, shuffle=False)

        epoch_pred_save_dir = None

        for index, (invalid_image, invalid_mask) in enumerate(

                tqdm.tqdm(invalid_data, total=self.params.invalid_samples // self.params.batch_size + 1)):

            invalid_pred = self.inference(invalid_image)

            epoch_pred_save_dir = os.path.join(self.pred_invalid_save_dir, f'epoch_{epoch}')

            module.save_images(

                image_shape=self.mask_shape, pred=invalid_pred,

                save_path=epoch_pred_save_dir, index=f'{index}', split=False

            )

        # Test model performance

        epoch_cropped_save_dir = os.path.join(

            self.invalid_crop_save_dir, f'epoch_{epoch}'

        )

        utils.crop_image(epoch_pred_save_dir, epoch_cropped_save_dir,

                         self.crop_width, self.crop_height,

                         self.input_shape[0], self.input_shape[1]

                         )

        m_dice, m_iou, m_precision, m_recall = performance.save_performace_to_csv(

            pred_dir=epoch_cropped_save_dir, gt_dir=self.params.gt_mask_dir,

            img_resize=(self.params.height, self.params.width),

            csv_save_name=f'{self.model_name}_epoch_{epoch}',

            csv_save_path=epoch_cropped_save_dir

        )

        return m_dice