"""

# ===============Part_1: Original Data Format Conversion===============

To do: <Strategy>

    1. init data [.nii(Nifit file) -> .png]: 

            CT data: 3D -> 2D

            save: [1] .png (dir: data_preprocessing/images)

    2. mask data [.nii -> .json / .txt]: 

            CT mask data: 3D -> 2D 

                          Filter the masks(contains 'pneumonia') -> JSON & YOLO_TXT 

            save: [1] .png (dir: data_preprocessing/masks)

                  [2] .json (dir: data_preprocessing/labelJSON) 

                            (recognizable by Labelme.exe)

                  [3] .txt(dir: data_preprocessing/labels)

                            (for YOLO input: 

                                YOLO format annotation data file -> describe label's classID & relative position) 

    3. data augmentation:

            [1] CT (images) - optimize:

                [1.1] reduce the CT range and to increase the contrast (pkg: windowing -> transform window width & center)

                [1.2] Filter all CT & masks which contains 'pneumonia'

            [2] Masks - optimize: (for YOLOv8)

                [2.1] binarize masks (0 & 1)

                [2.2] padded & filled mask_tmp's connected area

                [2.3] close all contours in mask

                [2.4] contours -> polygons

                [2.5] discard the polygons(contours) if it contains few points

            [3] CT & Masks - optimize: -> create a train_data_generator to perform various transformations (for UNET)

                [3.1] define an image_generator -> keras.preprocessing.Image.ImageDataGenerator()

                [3.2] image data augmentation -> flow_from_directory()

                [3.3] image normalization (??? false -> binarize CT in real)

                    recommended: 

                        in U-Net : CT -> normalization is better than binarization

"""

import os

import SimpleITK as sitk

import matplotlib.pyplot as plt

import math

import numpy as np

import io

import base64

from skimage import measure

from scipy.ndimage import (

    binary_dilation,

    binary_fill_holes,

)  # dilation & filling for mask

from PIL import Image  # Pillow Imaging Library

import json

import tensorflow as tf

import cv2

import numpy as np

import pandas as pd

import matplotlib.pyplot as plt

# ========================1. Read Data========================

# Read .nii file

# def read_nifit(nii_path):

#     path_parts = nii_path.split("\\")

#     nii_name = path_parts[-1]

#     nii_type = path_parts[-2]

#     # read nii

#     img_slices = sitk.ReadImage(nii_path)

#     wincenter = img_slices

#     data_img_slices = sitk.GetArrayFromImage(img_slices)  # convert sitk image to numpy array

#     # write in txt

#     with open(f"{nii_type}_{nii_name.split('.')[0]}.txt", "w", encoding="utf-8") as f:

#         f.write(f"{nii_path}: \n {img_slices}\n")

#         f.write("===================================\n")

#         f.write(f"Shape of the data: {data_img_slices.shape}\n")

#         f.write(f"data_img_slices: \n")

#         for i in data_img_slices:

#             for j in i:

#                 f.write(f"{j}\n")

#             f.write("===================================\n")

#     return data_img_slices, nii_type, nii_name

# Make plt_savefig Name

def make_fig_name(nii_type, nii_name):

    return nii_type + "_" + nii_name.split(".")[0]

# Draw Nifit CT

def draw_image(data_img_slices, idx1, idx2, savefig=None):

    row = math.floor(math.sqrt(idx2 + 1 - idx1))

    col = math.ceil(math.sqrt(idx2 + 1 - idx1))

    if row * col < (idx2 + 1 - idx1):

        row += 1

    j = 1

    for i in range(idx1 - 1, idx2):

        plt.subplot(row, col, j)

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

        plt.axis(False)

        j += 1

    plt.show()

    if savefig and isinstance(savefig, str):

        plt.savefig(f"CT_{savefig}.png")

# Draw CT histogram

def draw_histogram(data_img_slices, idx, savefig=None):

    ct_values = data_img_slices[idx - 1].flatten()  # flatten 3D -> 1D

    plt.figure(figsize=(10, 6))

    plt.hist(ct_values, bins=100, color="blue", alpha=0.7)

    plt.title("Histogram of CT Values")

    plt.xlabel("CT Value")

    plt.ylabel("Frequency")

    plt.grid(True)

    plt.show()

    if savefig and isinstance(savefig, str):

        plt.savefig(f"CT_histogram_{savefig}.png")

# Output nii's metadata

def create_metadata_dir(data_path):

    # create metadata folder & sub_dir

    meta_path = os.path.join(data_path, "metadata_tmp")

    img_meta_path = os.path.join(meta_path, "images")

    mask_meta_path = os.path.join(meta_path, "masks")

    PATH = {meta_path, img_meta_path, mask_meta_path}

    for path in PATH:

        os.makedirs(path, exist_ok=True)

    return img_meta_path, mask_meta_path

# Output nii's metadata

def write_metadata_in_txt(nii_path, sitk_nii, metadata, save_dir_path):

    meta_path = os.path.join(save_dir_path, f"{nii_name.split('.')[0]}_metadata.txt")

    with open(meta_path, "w", encoding="utf-8") as f:

        f.write(f"{nii_path}: \n {sitk_nii}\n")

        f.write("===================================\n")

        f.write(f"Shape of the data: {metadata.shape}\n")

        f.write(f"data_img_slices: \n")

        for i in metadata:

            f.write("\n".join(map(str, i)))

            f.write("===================================\n")

# ========================2. Data Augmentation========================

# img

def window_transform(sitkImage, winwidth=250, wincenter=80):

    """

    To do: transform window width & center

           to reduce the CT range and to increase the contrast (pkg: windowing)

    :param sitkImage(input) -> SimpleITK object

    :param winwidth -> CT values range --> winwidth: smaller, contrast: stronger

    :param wincenter -> all displayed CT values' center position --> wincenter: smaller, brightness: brighter

    :return sitkImage(output) -> sitkImage after transform

    """

    # define borders

    min = int(wincenter - winwidth / 2.0)

    max = int(wincenter + winwidth / 2.0)

    # define SimpleITK object -> transform window width & center

    intensityWindow = sitk.IntensityWindowingImageFilter()

    intensityWindow.SetWindowMinimum(min)

    intensityWindow.SetWindowMaximum(max)

    # print(

    #     f"new window width: {intensityWindow.GetWindowMinimum()} ~ {intensityWindow.GetWindowMaximum()}"

    # )

    sitkImage = intensityWindow.Execute(sitkImage)

    return sitkImage

# img

def get_winwidth_wincenter(img_nii_path, mask_nii_path):

    """

    To do: According to CT's Mask flexibly winwidth && wincenter

    :param img_nii_path -> nii_path

    :param mask_nii_path -> nii_path

    :return winwidth, wincenter

    Note: for Lung CT: default winwidth: 1000 ~ 1600HU, default wincenter:-600 ~ -800HU

    """

    all_min_max_hu = []

    for nii_name in os.listdir(mask_nii_path):

        if nii_name.endswith(".nii.gz") or nii_name.endswith(".nii"):

            # read nii file

            imgs_nii = sitk.ReadImage(os.path.join(img_nii_path, nii_name))

            masks_nii = sitk.ReadImage(os.path.join(mask_nii_path, nii_name))

            # nii -> array[]

            imgs_init = sitk.GetArrayFromImage(imgs_nii)

            masks = sitk.GetArrayFromImage(masks_nii)

            masks_hu_init = np.where(masks == 1, imgs_init, np.nan)

            for mask in masks_hu_init:

                arr_mask = mask.flatten()

                if not np.all(np.isnan(arr_mask)):

                    all_min_max_hu.extend([np.nanmin(arr_mask), np.nanmax(arr_mask)])

    # print(f"all_min_max_hu: {all_min_max_hu}")

    lower_bound_min = (

        math.floor(np.nanmin(all_min_max_hu) / 100) * 100

    )  # find nearest bounds

    upper_bound_max = math.ceil(np.nanmax(all_min_max_hu) / 100) * 100

    winwidth = upper_bound_max - lower_bound_min

    wincenter = np.around(1 / 2 * (upper_bound_max + lower_bound_min))

    # print(f"min_bound: {min}, max_bound: {max}")

    # print(f"min: {lower_bound_min}, max: {upper_bound_max}")

    print(f"winwidth: {winwidth}, wincenter: {wincenter}")

    return winwidth, wincenter

# img

def img_to_base64str(img_pil):

    """

    To do:  Image convertion: PIL Image -> Base64 string

    :img_pil    : PIL Image Object(gray)

    :return     : (Base64)      str(utf-8)

    PIL: Python Imaging Library

    (Binary) byte str ==> use in Image/Video/...

    (Base64) byte str (ASCII character) ==> use in Print/Mail/URL/...

    (Base64)      str (utf-8: Unicode character (text str)) ==> use in JSON/XML/...

    This allows image data to be transmitted over the network as text strings or stored in text files without being affected by the compatibility or transmission issues that binary data can bring.

    """

    ENCODING = "utf-8"

    img_byte = io.BytesIO()  # Create a binary data stream in memory

    img_pil.save(img_byte, format="PNG")  # save img_pil in img_byte

    binary_byte_str = img_byte.getvalue()  # img -> Binary byte data

    base64_byte_str = base64.b64encode(

        binary_byte_str

    )  # (Binary) byte str -> (Base64) byte str

    base64_str = base64_byte_str.decode(ENCODING)  # (Base64) byte str -> (Base64) str

    return base64_str

# mask

def close_contour(contour):

    """

    To do: close contour [traite image shapes and borders]

           sure that a 2D contour is connected end-to-end to form a closed contour

    """

    if not np.array_equal(contour[0], contour[-1]):  # contours: not closed

        contour = np.vstack((contour, contour[0]))  # add contours[0] at the end

    return contour

# mask

def binary_mask_to_polygon(binary_mask, tolerance=0):

    """

    To do: binary_mask -> polygon(COCO dataset format)

    From project "pycococreator-master"

    :binary_mask : array[](uint8)

    :tolerance=0 : control the fineness of the polygon approximation process

    :return      : polygons (list)

    contours: {[c_1][c_2]...[c_N]} -> N contours in mask

    contour : [[y1,x1]...[yn,xn]]  -> n points in this contour

              -> shape: (n,2)

              -> y: row coordinate, x: col coordinate

    polygons: [[segmentation_1], ..., [segmentation_N]] -> N polygons in mask

    """

    # treat mask

    padded_binary_mask = np.pad(

        binary_mask, pad_width=1, mode="constant", constant_values=0

    )  # padding: to avoid losing the boundary info when check the contour later

    # check all contours in mask

    contours = measure.find_contours(

        padded_binary_mask, 0.5

    )  # threshold: 0.5 (for binary img) # default contour's precision = 1 -> decimals=1

    # treat contours

    polygons = []

    for contour in contours:

        contour = np.subtract(contour, 1)  # restore: original = contour - padding(1)

        contour = close_contour(contour)  # make sure contour closed

        contour = measure.approximate_polygon(contour, tolerance)

        if len(contour) < 3:  # check points num # filter the 'contour' whose points < 3

            continue

        contour = np.flip(contour, axis=1)  # flip the axis: img(y,x) -> COCOdata(x,y)

        segmentation = contour.ravel().tolist()  # contour[points] -> 1D list

        segmentation = [

            0 if i < 0 else i for i in segmentation

        ]  # make sure segmentation >= 0

        polygons.append(segmentation)

    return polygons

# img & mask

def img_mask_to_json_txt(img, mask, class_names, class_mapping):

    """

    To do: img + mask -> json

    json: for img annotation tool (eg : Labelme)

    :img         : PIL Image Object (gray)

    :mask        : PIL Image Object (gray)

    :class_names : ['_background_', 'pneumonia']

    """

    # init JSON

    JSON_OUTPUT = {

        "version": "5.4.1",  # Labelme version

        "flags": {},  # Additional flag info

        "shapes": [],  # [{label, points, group_id, shape_type, flags}...{}]

        "imagePath": {},

        "imageData": {},  # Image data (Base64 encoded)

        "imageHeight": {},

        "imageWidth": {},

    }

    # init TXT -> YOLO .txt format: one polygon(object) <-> one row

    TXT_OUTPUT = []

    """

    TXT_OUTPUT = [

        # if bboxes is rectangle:

        [classID, x_center_nor, y_center_nor, width_nor, height_nor]

        ...  # here: nor-> normalized

        # for polygons in this case:

        [classID, [x1_nor, y1_nor], ...]

        ... # one polygon(object) <-> one row

    ] # list

    """

    # JSON - "imagePath" -> set externally

    # JSON - "imageData"

    imageData = img_to_base64str(img)  # Image -> base64_str ('utf-8')

    JSON_OUTPUT["imageData"] = imageData

    # JSON - "imageHeight" & "imageWidth"

    binary_mask = np.asarray(mask).astype(

        np.uint8

    )  # Image -> array ('uint-8') # binary_mask just contains 0 & 1

    JSON_OUTPUT["imageHeight"] = binary_mask.shape[0]  # height

    JSON_OUTPUT["imageWidth"] = binary_mask.shape[1]  # width

    # JSON - "shapes": label, points, group_id,  shape_type, flags

    success = False

    for i in np.unique(binary_mask):  # loop through each img's mask(frame)

        # Filter the masks(contains 'pneumonia')

        if i != 0:  # 0 : typically background

            # original mask_patch

            mask_tmp = np.where(

                binary_mask == i, 1, 0

            )  # check the connected area (set: 1, rest: 0 -> binarization)

            """

            Description for "np.where(condition, x(True), y(False))":

                mask_tmp = binary_mask

                for j in mask_tmp:

                    if j == i:

                        True

                        j = 1

                    else:

                        False

                        j = 0

                return mask_tmp: a new binary mask (contain only 0 & 1)

            """

            # Draw mask -> plot its boundary points

            # plt.subplot(1, 2, 1) # position: left

            # plt.imshow(mask_tmp, cmap='gray')

            # padded & filled mask_tmp's connected area

            mask_tmp = binary_dilation(

                mask_tmp, structure=np.ones((3, 3)), iterations=1

            )  # padding

            mask_tmp = binary_fill_holes(mask_tmp, structure=np.ones((5, 5)))  # filling

            # plt.subplot(1, 2, 2) # position: right

            # plt.imshow(mask_tmp, cmap='gray')

            # plt.show()

            # polygon extraction

            polygons = binary_mask_to_polygon(mask_tmp, tolerance=2)

            for polygon in polygons:

                if len(polygon) > 10:  # discard the contour if it contains few points

                    # "shapes" -> labels

                    label = class_names[i]  # class_names[1] : 'pneumonia'

                    # "shapes" -> points

                    # YOLO TXT -> one row, one contour(polygon)

                    points = []

                    yolotxt_row = str(class_mapping[label])  # YOLO txt_row -> classID

                    for j in range(0, len(polygon), 2):  # step_range: 2

                        x = polygon[j]  # default decimals = 1

                        y = polygon[j + 1]

                        points.append([x, y])  # point[x, y]

                        # YOLO txt_row -> point[x, y] -> normalization

                        x = round(x / JSON_OUTPUT["imageWidth"], 6)  # normalized

                        y = round(

                            y / JSON_OUTPUT["imageHeight"], 6

                        )  # normalized, 6: decimals = 6

                        yolotxt_row += " " + str(x) + " " + str(y)

                    # JSON - "shapes"

                    shape = {

                        "label": label,

                        "points": points,

                        "group_id": None,

                        # 'description': "",

                        "shape_type": "polygon",

                        "flags": {},

                        # 'mask' : None

                    }

                    JSON_OUTPUT["shapes"].append(shape)

                    # TXT - "row <-> contour"

                    TXT_OUTPUT.append(yolotxt_row)

                    # width_normalized, height_normalized ???

                    success = True

    return success, JSON_OUTPUT, TXT_OUTPUT

# Test Code

data_path = "./data/mosmed"

img_nii_path = os.path.join(data_path, "dataNii")

mask_nii_path = os.path.join(data_path, "maskNii")

# create data_preprocessing folder & sub_dir

data_pre_path = os.path.join(data_path, "data_preprocessing")

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

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

json_path = os.path.join(data_pre_path, "labelsJSON")

txt_path = os.path.join(data_pre_path, "labels")

PATH = {

    data_pre_path,  # preprocessing_tmp1 --> data augmentation for mask

    img_png_path,

    mask_png_path,

    json_path,

    txt_path

}

for path in PATH:

    os.makedirs(path, exist_ok=True)

class_names = ["_background_", "pneumonia"]

class_mapping = {"pneumonia": 0}

print("Starting: Nii -> PNG/JSON/TXT ...")

winwidth, wincenter = get_winwidth_wincenter(img_nii_path, mask_nii_path)

for nii_name in os.listdir(mask_nii_path):

    if nii_name.endswith('.nii.gz') or nii_name.endswith('.nii'):

        # read nii file

        imgs_nii = sitk.ReadImage(os.path.join(img_nii_path, nii_name))

        masks_nii = sitk.ReadImage(os.path.join(mask_nii_path, nii_name))

        # imgs_init = sitk.GetArrayFromImage(imgs_nii)

        # windowing

        imgs_nii = window_transform(imgs_nii, winwidth=winwidth, wincenter=wincenter)

        # nii -> array[]

        imgs = sitk.GetArrayFromImage(imgs_nii)

        masks = sitk.GetArrayFromImage(masks_nii).astype(np.uint8) # uint8 ???

        """

        # write nii's metadata in txt && Draw

        img_meta_path, mask_meta_path = create_metadata_dir(data_path)

        write_metadata_in_txt(f"{os.path.join(img_nii_path, nii_name)}", imgs_nii, imgs_init, img_meta_path)

        write_metadata_in_txt(f"{os.path.join(mask_nii_path, nii_name)}", masks_nii, masks, mask_meta_path)

        draw_image(imgs, 1, len(imgs), make_fig_name("image", nii_name))

        draw_histogram(imgs_init, 17, make_fig_name("image_init", nii_name, 17))

        """

        for idx in range(masks.shape[0]): # masks.shape: (depth, height, width)

            # get img/mask whose mask value > 0 -> which has pneumonia in CT

            if np.sum(masks[idx, :, :]) > 0: 

                img_png = Image.fromarray(imgs[idx, :, :]).convert('L') # array -> PIL Image object # L: RGB -> gray

                mask_png = Image.fromarray(masks[idx, :, :]).convert('L') # L: RGB -> gray

                # mask -> json

                sub_name = str(nii_name.split('.')[0] + '_' + str(idx))

                img_mask_json_path = os.path.join(json_path, sub_name + '.json' )

                img_mask_txt_path = os.path.join(txt_path, sub_name + '.txt' )

                success, JSON_OUTPUT, TXT_OUTPUT = img_mask_to_json_txt(img_png, mask_png, class_names, class_mapping)

                if success:

                    # 1.1 & 2.1 save PNG

                    img_png.save(os.path.join(img_png_path, sub_name + '.png'))

                    mask_png.save(os.path.join(mask_png_path, sub_name + '.png'))  # Save the original mask, not the mask after padding & filling & close_contours

                    # 2.2 write in JSON

                    # JSON - "imagePath"

                    JSON_OUTPUT['imagePath'] = sub_name + '.png'

                    with open(img_mask_json_path, 'w') as json_output:

                        json.dump(JSON_OUTPUT, json_output, indent=4) # Serialize python_data as a JSON file

                    # 2.3 write in TXT

                    with open(img_mask_txt_path, 'w') as txt_output:

                        TXT_OUTPUT = np.array(TXT_OUTPUT)

                        for idx, row in enumerate(TXT_OUTPUT):

                            if idx != len(TXT_OUTPUT) - 1:

                                row = row + '\n'

                            txt_output.write(row)

print("The conversion was completed successfully.")