# from preprocessing import  ReadIn_MRIScans_Masks
# from writeout_dataset import Export2CompressedNifiti
import numpy as np
import os
import pydicom
import pandas as pd
import nibabel as nib

# cutoffSlice = 400
# seg_scan_dir = 'D:/MRI - Tairawhiti (User POV)'
# folders = ['10_AutoBindWATER_450_13A']
# output_seg_scan_dir = ('{}/Patient Segmentation Tasks (Google Colab)').format(seg_scan_dir)

def preprocessTestScans(cutoffSlice, seg_scan_dir, folders):
    output_seg_scan_dir = ('{}/Patient Segmentation Tasks (Google Colab)').format(seg_scan_dir)
    def read_dicom_files(directory):
        dicom_files = []
        files = os.listdir(directory)
        sorted_files = sorted(files)
        for filename in sorted_files:
            filepath = os.path.join(directory, filename)
            if os.path.isfile(filepath) and filename.endswith('.dcm'):
                try:
                    dicom_file = pydicom.dcmread(filepath)
                    dicom_files.append(dicom_file)
                except pydicom.errors.InvalidDicomError:
                    print(f"Skipping file: {filename}. It is not a valid DICOM file.")
        return dicom_files

    def flatten_2d_array(arr):
        flattened = []
        for row in arr:
            flattened.extend(row)
        return flattened

    def ReadIn_MRIScans_Masks(scans_path, folders):
        print('Patient Scan Data: ', folders)
        scan_pixel_data = []
        scan_coordinate_data = []
        single_scan_pixel_data = []
        scan_coordinate_data = []
        scan_orientation_data = []
        scan_pixelspacing_data = []


        single_paitent_scans_path =  scans_path + ('/Raw DICOM MRI Scans/{}').format(folders)
        dicom_files = read_dicom_files(single_paitent_scans_path)

        # Extracting pixel data
        for i in range (len(dicom_files)):
            single_scan_pixel_data.append(dicom_files[i].pixel_array)
        scan_pixel_data.append(single_scan_pixel_data)

        training_scans = flatten_2d_array(scan_pixel_data)
        training_scans = np.array(training_scans)
    
        # Coordinate Data
        single_paitent_scans_path =  scans_path + '/{}'.format(folders)
        for i in range (len(dicom_files)):
            scan_coordinate_data.append(dicom_files[i].ImagePositionPatient)
            scan_orientation_data.append(dicom_files[i].ImageOrientationPatient)
            scan_pixelspacing_data.append(dicom_files[i].PixelSpacing)
            
        coord_data = pd.DataFrame(scan_coordinate_data, columns=["x", "y", "z"])
        return training_scans

    def Export2CompressedNifiti(imgs_train, scans_path, colab_fname):
        nii_img_train = nib.Nifti1Image(imgs_train, affine=np.eye(4))
        output_file_path = ('{}/msk_{}.nii.gz').format(scans_path, colab_fname)
        nib.save(nii_img_train, output_file_path)

        print("Conversion to NIfTI complete with shape: ", nii_img_train.shape)
        print(('Successfully Exported Preprocessed Data! -> {}').format(output_file_path))
        print('\n')


    for patient in folders:
        patient_scan_data = ReadIn_MRIScans_Masks(seg_scan_dir, patient)

        patient_scan_data = patient_scan_data[(int(cutoffSlice)):-1]

        patient_scan_data = np.array(patient_scan_data)
        patient_scan_data = patient_scan_data.astype('float32')
        patient_scan_data /= 255.  # scale scans to [0, 1]

        patient_scan_data = np.expand_dims(patient_scan_data, axis=-1)
        print('Prediction Input Scan Shape: ', patient_scan_data.shape)

        
        scan_idx_fnamecolab = (patient.split('_'))[0]
        scan_idx_fnamecolab = '{:03d}'.format(int(scan_idx_fnamecolab))
        Export2CompressedNifiti(patient_scan_data, output_seg_scan_dir, scan_idx_fnamecolab)