# Exporting Training Scan and Mask Data

import os 

import numpy as np

from PIL import Image

import matplotlib.pyplot as plt

import re

import numpy as np

import nibabel as nib

import matplotlib.pyplot as plt

from tqdm import tqdm

def WriteOutTextFile(data_array, save_dir, starting_slice):

    num_files = data_array.shape[0]

    # Create the save directory if it doesn't exist

    if not os.path.exists(save_dir):

        os.makedirs(save_dir)

    for i in range(num_files):

        file_name = f"data_{i + starting_slice}.txt"

        file_path = save_dir + '/' + file_name

        np.savetxt(file_path, data_array[i], delimiter=',', fmt='%.4f')

        print(f"Saved {file_path}")

def WriteOutImagePNGFiles(data_array, save_dir, starting_slice):

    num_files = data_array.shape[0]

    # Create the save directory if it doesn't exist

    if not os.path.exists(save_dir):

        os.makedirs(save_dir)

    for i in range(num_files):

        file_name = f"image_{i + starting_slice}.png"

        file_path = save_dir + '/' + file_name

        # Scale the pixel values to the range [0, 255] if necessary

        image_data = ((data_array[i] * 255).astype(np.uint8)).squeeze()

        # Create an Image object and save as PNG

        image = Image.fromarray(image_data)

        if True:

            image.save(file_path)

        print(f"Saved {file_path}")

def ReadInDatasets(folder_path):

    files = os.listdir(folder_path)

    num_files = len(files)

    result_array = np.empty((num_files, 512, 512))

    for i, file in enumerate(files):

        file_path = os.path.join(folder_path, file)

        with open(file_path, 'r') as f:

            # Read the contents of the text file

            content = f.read()

            # Convert the content into a 2D NumPy array of floats

            array_2d = np.array([list(map(float, line.split(','))) for line in content.splitlines()])

            # Store the 2D array in the result array

            result_array[i] = array_2d

    return result_array

def ReadInDatasetNPY(folder_path):

    # Get a list of all files in the folder

    file_list = os.listdir(folder_path)

    # Filter out only the .npy files

    npy_files = [file for file in file_list if file.endswith('.npy')]

    # Sort the .npy files based on the numeric order in their filenames

    npy_files.sort(key=lambda x: int(x.split('_')[0]))

    # Initialize an empty list to store arrays from each .npy file

    arrays_list = []

    # Read in each .npy file and append it to the list

    for file_name in npy_files:

        file_path = os.path.join(folder_path, file_name)

        array = np.load(file_path)

        arrays_list.append(array)

    # Concatenate the arrays into a single array

    concatenated_array = np.concatenate(arrays_list, axis=0)

    return concatenated_array

def Export2CompressedNifiti(imgs_train, scans_path, colab_fname, imgs_mask_train, orientation):

    nii_img_train = nib.Nifti1Image(imgs_train, affine=np.eye(4))

    scan_idx_fnamecolab = int((colab_fname[0].split('_'))[1])

    scan_idx_fnamecolab = '{:03d}'.format(scan_idx_fnamecolab)

    output_file_path = ('{}/nnUNet Data/unprocessed_scans/msk_{}.nii.gz').format(scans_path, scan_idx_fnamecolab)

    nib.save(nii_img_train, output_file_path)

    nii_img_mask = nib.Nifti1Image(imgs_mask_train, affine=np.eye(4))

    output_file_path = ('{}/nnUNet Data/masks/{}/{}_{}.nii.gz').format(scans_path, ((colab_fname[0]).split('_'))[0],colab_fname[0], orientation)

    nib.save(nii_img_mask, output_file_path)   

    print("Conversion to NIfTI complete with shape: ", imgs_mask_train.shape)

    print('\n')

    print('Successfully Exported Preprocessed Data!')