"""

data:

    CT         :     used

    mask       :     used

    labels(txt): not used

    labelsJson : not used

"""

# U-Net Structure:

# 1->64->64...............................................................->(*/)128->64->64=>2      # 1: input image, 2: output segmentation map

#    (-+)64->128->128...........................................->(*/)256->128->128

#             (-+)128->256->256.......................->(*/)512->256->256

#                       (-+)256->512->512..->(*/)1024->512->512     # 1024: 512+512

#                                 (-+)512->1024->1024

# ->  : conv 3x3, RELU

# ..->: copy & crop

# (-+): max pool 2x2

# (*/): up-conv 2x2

# =>  : conv 1x1

'''

Issues:

1. 要不直接把preprocessing_tmp1 经过"data_generation"分割90% 出来用作训练集, 并存为dataset放在外面目录下

 ----> 看下之后test部分的图片怎么预处理, 如果处理方式一样那就dataset拿出来放到外面去

'''

from keras._tf_keras import keras  # CPU - keras > 3.*

from keras._tf_keras.keras.layers import *  # CPU - keras > 3.*

from keras._tf_keras.keras.preprocessing.image import (

    ImageDataGenerator,

)  # CPU - keras > 3.*

# from keras.layers import *  # GPU - keras > 2.*

# from keras.callbacks import ModelCheckpoint  # GPU - keras > 2.*

# from keras.preprocessing.image import ImageDataGenerator  # GPU - keras > 2.*

from keras import Model

from keras import backend as K

import os

import numpy as np

# from data_preparation import draw_image

import matplotlib.pyplot as plt

import cv2

# img & mask

# 3. data augmentation: (import tensorflow.keras.preprocessing.Image)

#         [1] define an image_generator -> ImageDataGenerator()

#         [2] image data augmentation -> flow_from_directory()

#         [3] image normalization

#     问题:

#         [1] 先进行.nii -> png/json/txt, 后进一步keras数据增强

#             有个问题: images/mask增强后随之的json/txt是否也要发生改变 ---> ???

#         [2] tensorflow和torch一起用 ---> 可以

#             model -> Yolo 使用的是pytorch

#             data augmentation 使用的是 tensorflow->keras

#         [3] gene后一定要跟fit()均值化，否则会提示：

#               F:\AI_Outils\Anaconda\1\envs\opencv_CPU\Lib\site-packages\keras\src\legacy\preprocessing\image.py:1263: UserWarning: This ImageDataGenerator specifies `featurewise_center`, but it hasn't been fit on any training data. Fit it first by calling `.fit(numpy_data)`.

# data augmentation for train

def train_generator(dataset_path, type):

    data_path = os.path.join(dataset_path, type)

    data_pre_path = os.path.join(dataset_path, f"{type}_generator")

    img_png_path = os.path.join(data_pre_path, "images")

    mask_png_path = os.path.join(data_pre_path, "masks")

    PATH = {

        data_pre_path,

        img_png_path,

        mask_png_path,

    }

    for path in PATH:

        os.makedirs(path, exist_ok=True)

    # 3.1 define an image_generator: to perform various transformations on object

    generator_args = dict(

        rotation_range=0.1,

        width_shift_range=0.05,

        height_shift_range=0.05,

        shear_range=0.05,

        zoom_range=0.05,

        horizontal_flip=False,

        vertical_flip=False,

    )

    generator_image = ImageDataGenerator(generator_args)

    generator_mask = ImageDataGenerator(generator_args)

    # 3.2 implement further data augmentation for image & mask

    generation_image = generator_image.flow_from_directory(

        directory=data_path,

        classes=["images"],

        class_mode=None,

        color_mode="grayscale",

        target_size=(512, 512),

        batch_size=2,

        save_to_dir=os.path.join(data_pre_path, "images"),

        # save_prefix='ct_',

        seed=123,

    )

    generation_mask = generator_mask.flow_from_directory(

        directory=data_path,

        classes=["masks"],

        class_mode=None,

        color_mode="grayscale",

        target_size=(512, 512),

        batch_size=2,

        save_to_dir=os.path.join(data_pre_path, "masks"),

        # save_prefix='mask_',

        seed=123,

    )

    generation = zip(generation_image, generation_mask)

    print("2--------------------------")

    i = 0

    # 3.3 image normalization (image -> not normalized yet, mask -> binary)

    for image, mask in generation:

        '''

        i = i + 1

        # output image data to TXT

        arr = np.array(image[0][:, :, 0])

        np.savetxt("array_0.txt", arr)

        print(f"image: min: {np.nanmin(arr)}, max: {np.nanmax(arr)}.")

        # output image data to TXT

        arr = np.array(normalization(image)[0][:, :, 0])

        np.savetxt("array_1.txt", arr)

        print(

            f"normalization_image: min: {np.nanmin(arr)}, max: {np.nanmax(arr)}."

        )

        print(image.shape, mask.shape)  # (2, 256, 256, 1) (2, 256, 256, 1) --> batch_size=2

        # visualization

        data_img_slices = [

            image[0][:, :, 0],

            normalization(image)[0][:, :, 0],

            mask[0][:, :, 0],

            image[1][:, :, 0],

            normalization(image)[1][:, :, 0],

            mask[1][:, :, 0],

        ]

        draw_image(data_img_slices, 1, 6, None)

        if i == 1:

            break

        '''

        yield (normalization(image), mask)

        # image[0][:, :, 0] = normalization(image[0][:, :, 0])

        # image[1][:, :, 0] = normalization(image[1][:, :, 0])

    print("Further data augmentation was completed successfully.")

def binarization(data):

    """

    Binarization: Converts data to only two values, e.g. 0 & 1

    To do: To highlight certain features in the image

    Processing: x'[x/255.0 > 0.5] = 1.0

                x'[x/255.0 <= 0.5] = 0.0

    """

    data_binary = data / 255.0

    data_binary[data_binary > 0.5] = 1.0

    data_binary[data_binary <= 0.5] = 0.0

    return data_binary

def standardization(data):

    """

    Standardization: Converts the data into a new distribution with a mean of 0 and a standard deviation of 1

    To do: to have comparability between different features (if data feature value range/unit is quite different --> perform standardization).

        --> Standardization does not change the distribution of feature data

    Processing: x' = (x - mean) / std

    """

    mean = np.mean(data)

    std = np.std(data)

    data_std = (data - mean) / std

    return data_std

def normalization(data):

    """

    Normalization: Scale the data to a specific range, e.g. 0-1

    To do: To make the influence of each feature on the target variable consistent.

        --> Data normalization changes the distribution of feature data

    Processing: x' = (x - min)/(max - min)

    Note: In the medical field, normalization is generally performed --> to accelerate the convergence of the network and make the model more stable

    """

    min = np.nanmin(data)

    max = np.nanmax(data)

    data_nor = (data - min) / (max - min)

    return data_nor

def u_net(input_size = (512, 512, 1), path=None):

    # layer 1-1

    inputs_L1_1 = Input(input_size)

    conv1_L1_1 = Conv2D(64, 3, activation="relu", padding="same", kernel_initializer="he_normal")(inputs_L1_1)  # filters: 64, kernel_size:3x3, kernel_initializer: use normal distribution to initializer Weights of kernel

    conv2_L1_1 = Conv2D(64, 3, activation="relu", padding="same", kernel_initializer="he_normal")(conv1_L1_1)

    pool1_L1_1 = MaxPool2D(pool_size=(2,2))(conv2_L1_1)

    # layer 2-1

    conv3_L2_1 = Conv2D(128, 3, activation="relu", padding="same", kernel_initializer="he_normal")(pool1_L1_1)

    conv4_L2_1 = Conv2D(128, 3, activation="relu", padding="same", kernel_initializer="he_normal")(conv3_L2_1)

    pool2_L2_1 = MaxPool2D(pool_size=(2, 2))(conv4_L2_1)

    # layer 3-1

    conv5_L3_1 = Conv2D(256, 3, activation="relu", padding="same", kernel_initializer="he_normal")(pool2_L2_1)

    conv6_L3_1 = Conv2D(256, 3, activation="relu", padding="same", kernel_initializer="he_normal")(conv5_L3_1)

    pool3_L3_1 = MaxPool2D(pool_size=(2, 2))(conv6_L3_1)

    # layer 4-1

    conv7_L4_1 = Conv2D(512, 3, activation="relu", padding="same", kernel_initializer="he_normal")(pool3_L3_1)

    conv8_L4_1 = Conv2D(512, 3, activation="relu", padding="same", kernel_initializer="he_normal")(conv7_L4_1)

    pool4_L4_1 = MaxPool2D(pool_size=(2, 2))(conv8_L4_1)

    # layer 5

    conv9_L5 = Conv2D(1024, 3, activation="relu", padding="same", kernel_initializer="he_normal")(pool4_L4_1)

    conv10_L5 = Conv2D(1024, 3, activation="relu", padding="same", kernel_initializer="he_normal")(conv9_L5)

    up1_L5 = UpSampling2D(size=(2, 2))(conv10_L5)  # deconvolution

    # layer 4-2

    conv11_L4_2 = Conv2D(512, 3, activation="relu", padding="same", kernel_initializer="he_normal")(concatenate([up1_L5, conv8_L4_1], axis=3))  # concatenation

    conv12_L4_2 = Conv2D(512, 3, activation="relu", padding="same", kernel_initializer="he_normal")(conv11_L4_2)

    up2_L4_2 = UpSampling2D(size=(2, 2))(conv12_L4_2)

    # layer 3-2

    conv13_L3_2 = Conv2D(256, 3, activation="relu", padding="same", kernel_initializer="he_normal")(concatenate([up2_L4_2, conv6_L3_1], axis=3))

    conv14_L3_2 = Conv2D(256, 3, activation="relu", padding="same", kernel_initializer="he_normal")(conv13_L3_2)

    up3_L3_2 = UpSampling2D(size=(2, 2))(conv14_L3_2)

    # layer 2-2

    conv15_L2_2 = Conv2D(128, 3, activation="relu", padding="same", kernel_initializer="he_normal")(concatenate([up3_L3_2, conv4_L2_1], axis=3))

    conv16_L2_2 = Conv2D(128, 3, activation="relu", padding="same", kernel_initializer="he_normal")(conv15_L2_2)

    up4_L2_2 = UpSampling2D(size=(2, 2))(conv16_L2_2)

    # layer 1-2

    conv17_L1_2 = Conv2D(64, 3, activation="relu", padding="same", kernel_initializer="he_normal")(concatenate([up4_L2_2, conv2_L1_1], axis=3))

    conv18_L1_2 = Conv2D(64, 3, activation="relu", padding="same", kernel_initializer="he_normal")(conv17_L1_2)

    outputs_L1_2 = Conv2D(1, 1, activation="sigmoid")(conv18_L1_2)

    # build model

    model = Model(inputs = inputs_L1_1, outputs = outputs_L1_2)

    # compile model

    '''

    Loss     : 0-1 binary cross-entropy (binary_crossentropy)

    Optimizer: Adaptive Descent (Adam)

    Callback : After each epoch is trained, autosave a best pre-trained model(optimal weights). (keras.callbacks.ModelCheckpoint)

    '''

    model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])

    return model

class ShowMask(keras.callbacks.Callback):

    def __init__(self):

        super().__init__()

    def on_epoch_end(self, epoch, logs=None):

        print()

        idx = 0

        for img, mask in gene:

            compare_list = [img[0], mask[0], model.predict(img[0].reshape(1, 512, 512, 1))[0]]

            for i in range(0, len(compare_list)):

                plt.subplot(1, 3, i+1)

                plt.imshow(compare_list[i], cmap="gray")

                plt.axis(False)

            # plt.show()

            plt.savefig(f"compare_{idx}.png")

            idx = idx + 1

            break

        # return super().on_epoch_end(epoch, logs)

# Test

# dataset: data augmentation for train data

UNETDataset_path = "./UNETDataset"

gene = train_generator(UNETDataset_path, "train")  # could be used as input to the model and directly as training

# model params

steps_per_epoch = 50

epochs = 100

model_name = f"u_net-512-512-1-pneumonia_{epochs}_{steps_per_epoch}.keras"

models_path = "./models/"

model_path = os.path.join(models_path, model_name)

model_ckpt = keras.callbacks.ModelCheckpoint(model_path, save_best_only=False, verbose=1)

# train

K.clear_session() # keras

model = u_net(path=model_path) # structure

# print(model.summary())

# model.fit(gene, steps_per_epoch=steps_per_epoch, epochs=epochs, callbacks=[model_ckpt, ShowMask()]) # train

model.fit(

    gene,

    steps_per_epoch=steps_per_epoch,

    epochs=epochs,

    callbacks=[model_ckpt],

)  # train

# 思路：

# 1. Dataset_mini -> CPU OK, GPU KO

# 2. cudnn v8 -> GPU KO

# 3. 缩减 input_size -> 512->256

# 4. 缩减 unet structure

# evalution

data_val_generator_path = os.path.join(UNETDataset_path, "val_generator")

compare_path = os.path.join(data_val_generator_path, "compare")

PATH = {data_val_generator_path, compare_path}

for path in PATH:

    os.makedirs(path, exist_ok=True)

# model_test= keras.models.load_model(model_path)

'''

gene = train_generator(UNETDataset_path, "val")

idx = 0

for img, mask in gene:

    predict_mask = model_test.predict(img)[0]

    # predict_mask_np = (predict_mask * 255).numpy()

    _, real_mask = cv2.threshold(mask[0], 127, 255, 0)

    real_mask = (real_mask).astype('uint8')

    real_contours, _ = cv2.findContours(real_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    real_overlap_img = cv2.drawContours(img[0].copy(), real_contours, -1, (0, 255, 0), 2)

    _, pred_mask = cv2.threshold((predict_mask * 255).astype("uint8"), 127, 255, 0)

    pred_mask = (pred_mask).astype('uint8')

    pred_contours, _ = cv2.findContours(pred_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    pred_overlap_img = cv2.drawContours(img[0].copy(), pred_contours, -1, (255, 0, 0), 2)

    compare_list = [img[0], pred_mask, real_overlap_img, pred_overlap_img]

    for i in range(0, len(compare_list)):

        plt.subplot(1, 4, i+1)

        plt.imshow(compare_list[i], cmap="gray")

        plt.axis(False)

    # plt.show()

    save_path = os.path.join(compare_path, f"compare_{idx}.png")

    plt.savefig(save_path)

    idx = idx + 1

'''

'''

# test save compare_png

images_path = os.path.join(data_val_generator_path, "images")

masks_path = os.path.join(data_val_generator_path, "masks")

idx = 0

for png_name in os.listdir(images_path):

    # predict_mask = model_test.predict(img)[0]

    img = cv2.imread(os.path.join(images_path, png_name))

    mask = cv2.imread(os.path.join(masks_path, png_name), cv2.IMREAD_GRAYSCALE)

    _, real_mask = cv2.threshold(mask, 127, 255, 0)

    real_mask = (real_mask).astype("uint8")

    real_contours, _ = cv2.findContours(

        real_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE

    )

    real_overlap_img = cv2.drawContours(

        img.copy(), real_contours, -1, (0, 255, 0), 2

    )

    compare_list = [img, real_overlap_img]

    for i in range(0, len(compare_list)):

        plt.subplot(1, 2, i + 1)

        plt.imshow(compare_list[i], cmap="gray")

        plt.axis(False)

    # plt.show()

    save_path = os.path.join(compare_path, f"compare_{idx}.png")

    plt.savefig(save_path)

    idx = idx + 1

'''