import os

import numpy as np

from src.dataset.utils.nifi_volume import load_nifi_volume_return_nib, save_segmask_as_nifi_volume, load_nifi_volume

from src.uncertainty.uncertainty import brats_normalize

from tqdm import tqdm

if __name__ == "__main__":

    setx = "valid"

    task = "uncertainty_task_tta/"

    model_id = "model_1598640035"

    gen_path = "/mnt/gpid07/users/laura.mora/"

    model_path = f"{gen_path}/results/checkpoints/{model_id}/"

    ground_truth_path = f"{gen_path}/datasets/2020/train/no_patch"

    input_dir = os.path.join(model_path, f"{task}/{setx}")

    output_dir = os.path.join(model_path, f"{task}/{setx}/normalized")

    if not os.path.exists(output_dir):

        os.makedirs(output_dir)

    file_list = sorted([file for file in os.listdir(input_dir) if "unc" in file and "nii.gz"])

    file_list_all = sorted([file for file in os.listdir(input_dir) if "nii.gz" in file])

    max_uncertainty = 0

    min_uncertainty = 10000

    for uncertainty_map in tqdm(file_list, total=len(file_list), desc="Getting min and max"):

        # Load Uncertainty maps

        patient_name = uncertainty_map.split(".")[0].split("_unc")[0]

        path_gt = os.path.join(ground_truth_path, patient_name, f"{patient_name}_flair.nii.gz")

        flair = load_nifi_volume(path_gt, normalize=False)

        brain_mask = np.zeros(flair.shape, np.float)

        brain_mask[flair > 0] = 1

        path = os.path.join(input_dir, uncertainty_map)

        unc_map, _ = load_nifi_volume_return_nib(path, normalize=False)

        tmp_max = np.max(unc_map[brain_mask == 1])

        tmp_min = np.min(unc_map[brain_mask == 1])

        if tmp_max > max_uncertainty:

            max_uncertainty = tmp_max

        if tmp_min < min_uncertainty:

            min_uncertainty = tmp_min

    for uncertainty_map_path in tqdm(file_list_all, total=len(file_list_all), desc="Normalizing.."):

        path = os.path.join(input_dir, uncertainty_map_path)

        output_path = os.path.join(output_dir, uncertainty_map_path)

        unc_map, nib_data = load_nifi_volume_return_nib(path, normalize=False)

        if "unc" in uncertainty_map_path:

            uncertainty_map_normalized = brats_normalize(unc_map, max_unc=max_uncertainty, min_unc=min_uncertainty)

            print(f"Saving to: {output_path}")

            save_segmask_as_nifi_volume(uncertainty_map_normalized,  nib_data.affine, output_path)

        else:

            save_segmask_as_nifi_volume(unc_map, nib_data.affine, output_path)