"""

# ===============Part_2: Dataset Production===============

To do: <Strategy>

    1. Read images/labels/labelsJSPN'path list

        Input:

            images[] (.png)

            labelsJSON[] (.json)

            labels[] (.txt)

    2. Create YOLO dataset: 

        [1] split set & copy files

        [2] create config file (.yaml) 

        Structure:

                dataset:(dir: sub_dir)

                        train/val/test: images

                                        labels -> YOLO TXT (annotation file)

                                        labelsJson

                dataset.yaml (config file)

                        -> format:

                                train: trainset_absolute_path

                                val: valset_absolute_path

                                # test: testset_absolute_path

                                nc: num(classes)

                                names: ['names_class', ...]

"""

import os

from sklearn.model_selection import train_test_split

from tqdm import tqdm  # progress bar

import shutil

MODELTYPES = ["YOLO", "UNET"]

def make_dataset(

    model, data_path, train_size=0.8, test_set=True

):

    # create dataset folder

    global MODELTYPES

    if model in MODELTYPES:

        dataset_path = os.path.join(os.path.dirname(__file__), f"{str(model)}Dataset")

        if model == "YOLO":

            yaml_path = os.path.join(dataset_path, "dataset.yaml")

    else:

        print(

            f"Missing creating {str(model)}'s Dataset solution. Recommended: {', '.join(MODELTYPES)}"

        )

        return

    ensure_dir(dataset_path)

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

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

    labels_path = os.path.join(data_path, "labels")

    labelsJSON_path = os.path.join(data_path, "labelsJSON")

    # 1. Read images/labels/labelsJSPN'path list

    images = []

    masks = []

    labels = []

    labelsJSON = []

    for image_name in os.listdir(images_path):

        image_path = os.path.join(images_path, image_name)

        mask_path = os.path.join(masks_path, image_name)

        name, ext = os.path.splitext(image_name)

        label_name = os.path.join(name + ".txt")

        label_path = os.path.join(labels_path, label_name)

        labelJSON_name = os.path.join(name + ".json")

        labelJSON_path = os.path.join(labelsJSON_path, labelJSON_name)

        if all(map(os.path.exists, [mask_path, label_path, labelJSON_path])):

            images.append(image_path)

            masks.append(mask_path)

            labels.append(label_path)

            labelsJSON.append(labelJSON_path)

        else:

            print("The initial data in images/masks/labels/labelsJSON do not correspond")

    # 2.1.1 Split data to train/val/test set

    data = list(zip(images, masks, labels, labelsJSON))

    if test_set:

        test_size = int(len(images) * (1 - train_size) / 2)

        train_val_data, test_data = train_test_split(

            data, test_size=test_size, random_state=42

        )  # split

        train_data, val_data = train_test_split(

            train_val_data, test_size=test_size, random_state=42

        )  # test_size: specific len(val_images)

    else:

        train_data, val_data = train_test_split(

            data, test_size=1 - train_size, random_state=42

        )  # split # test_size: radio

        test_data = list(), list(), list(), list()

    train_images, train_masks, train_labels, train_labelsJSON = zip(*train_data)  # unpack

    val_images, val_masks, val_labels, val_labelsJSON = zip(*val_data)

    test_images, test_masks, test_labels, test_labelsJSON = (zip(*test_data) if test_set else ([], [], [], []))

    # make CSV ???

    # dataset_structure

    dataset_structure = {

        "train": {

            "images": train_images,

            "masks": train_masks,

            "labels": train_labels,

            "labelsJson": train_labelsJSON,

        },

        "val": {

            "images": val_images,

            "masks": val_masks,

            "labels": val_labels,

            "labelsJson": val_labelsJSON,

        },

        "test": {

            "images": test_images,

            "masks": test_masks,

            "labels": test_labels,

            "labelsJson": test_labelsJSON,

        },

    }

    # 2.1.2 Copy data to train/val/test set

    dataset_paths = {}

    for base_dir, sub_dirs in dataset_structure.items():

        dataset_paths[base_dir] = {}

        for sub_dir, file_list in sub_dirs.items():

            # create each sub_dir & save paths

            sub_dir_path = os.path.join(dataset_path, base_dir, sub_dir)

            dataset_paths[base_dir][sub_dir] = sub_dir_path

            if base_dir == "test" and test_set is False:

                continue

            if sub_dir == "masks" and model == "YOLO":

                continue

            if (sub_dir == "labels" or sub_dir == "labelsJson") and model == "UNET":

                continue

            ensure_dir(sub_dir_path)

            # copy

            # print(f"\nlength: {base_dir + '_' + sub_dir}: {len(file_list)}")

            copy_files(sub_dir_path, file_list)

    # 2.2 Create config file (.yaml)

    if model == "YOLO":

        """ YAML

        format:

            train: trainset_absolute_path

            val: valset_absolute_path

            # test: testset_absolute_path

            nc: num(classes)

            names: ['names_class', ...]

        """

        class_mapping = {"pneumonia": 1}

        YAML_OUT = {

            "train": os.path.abspath(dataset_paths["train"]["images"]),

            "val": os.path.abspath(dataset_paths["val"]["images"]),

            'test': os.path.abspath(dataset_paths["test"]["images"]),

            "nc": len(class_mapping),

            "names": list(class_mapping.keys())

        }

        with open(yaml_path, "w") as yaml_output:

            idx = 0

            for key, value in YAML_OUT.items():

                idx += 1

                if key == 'test' and test_set is False:

                    continue

                row = f'{key}: {value}'

                if idx != len(YAML_OUT):

                    row = row + '\n'

                yaml_output.write(row)

    print(f"\n{str(model)}Dataset was successfully created.")

def ensure_dir(file_path):

    if not os.path.exists(file_path):

        os.makedirs(file_path)

def copy_files(sub_dir_path, file_list):

    # copy files to the dataset folders

    for file_path in tqdm(file_list, desc=sub_dir_path):

        shutil.copy(

            file_path,

            os.path.join(sub_dir_path, os.path.basename(file_path)),

        )

data_path = "./data/mosmed/data_preprocessing/"

# make_dataset("YOLO", data_path, train_size=0.8, test_set=False)

make_dataset("UNET", data_path, train_size=0.9, test_set=False)