import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
import os
from image_preprocessing import uniform_resizing

def CreateMasks4MulticlassMSK(input_scan_dir, scan_dir, directory, mask_index, TIBIA_encoding, FEMUR_encoding, FIBULA_encoding, PELVIS_encoding, output_dir, AOIThresholding, FriedLanderDataset):
    aoi_fnames = os.listdir(directory)
    resizing_size = 256
    print('Regions of Interest for Segmentation: ', aoi_fnames)

    fnames = []
    for aoi_fname in aoi_fnames:
        directory_aoi = (directory + ('/{}').format(aoi_fname))
        aoi_fname = os.listdir(directory_aoi)
        aoi_fname = sorted(aoi_fname)
        fnames.append(aoi_fname)

    n_segmentation_classes = len(fnames)
    print('Number of Segmentation Classes: ', n_segmentation_classes)



    suffix_index_dict = {}
    for idx, file_list in enumerate(fnames):
        for filename in file_list:
            # Extract the suffix index from the filename
            parts = filename.split('_')
            if len(parts) > 1:
                suffix_index = parts[1].split('.')[0]  # Extracting the index before '.'
                suffix_index = int(suffix_index)  # Convert index to integer
                # Add the filename to the corresponding list in the dictionary
                if suffix_index in suffix_index_dict:
                    suffix_index_dict[suffix_index].append(filename)
                else:
                    suffix_index_dict[suffix_index] = [filename]

    # for index, filenames in suffix_index_dict.items():
    #     print(f"Segmentation Mask {index}: {filenames}")

    combined_mask = None
    for fname_mask in suffix_index_dict[mask_index]:
        print(fname_mask)
        mask = nib.load(('{}/{}/{}').format(directory, (fname_mask.split('_'))[0], fname_mask))
        mask_data = mask.get_fdata().astype(np.uint8)
        aoi = (fname_mask.split('_'))[0]

        if (FriedLanderDataset == True):
            mask_data = np.expand_dims(mask_data, axis = -1)
            mask_data = mask_data.transpose(1, 0, 2, 3)
            print(('Unprocessed {} Mask Shape: {}').format(aoi, mask_data.shape))
            mask_data = uniform_resizing(mask_data, resizing_size)


        if (aoi == 'TIBIA'):
            mask_data = np.where(mask_data != 0, TIBIA_encoding, 0)
        if (aoi == 'FEMUR'):
            mask_data = np.where(mask_data != 0, FEMUR_encoding, 0)
        if (aoi == 'FIBULA'):
            mask_data = np.where(mask_data != 0, FIBULA_encoding, 0)
        if (aoi == 'PELVIS'):
            mask_data = np.where(mask_data != 0, PELVIS_encoding, 0)    
        if combined_mask is None:
            combined_mask = mask_data
        else:
            combined_mask += mask_data  # Combine the pixel arrays by adding them element-wise
            combined_mask = np.array(combined_mask)
            combined_mask[combined_mask > n_segmentation_classes] = 0

    if (AOIThresholding == True):
        threshold = 0

        combined_mask = np.array(combined_mask)
        indices = np.transpose(np.nonzero(combined_mask != 0))
        if indices.size > 0:
            first_non_zero_index_2d = tuple(indices[0])
        else:
            first_non_zero_index_2d = None

        if indices.size > 0:
            last_non_zero_index_2d = tuple(indices[-1])
        else:
            last_non_zero_index_2d = None


        print(first_non_zero_index_2d)
        print(last_non_zero_index_2d)
        first_slice_aoi = (first_non_zero_index_2d[0]) - int(threshold)
        last_slice_aoi = (last_non_zero_index_2d[0]) + int(threshold)
        print("AOI Slice Start (Final with Thresholding): ", first_slice_aoi)
        print("AOI Slice End (Final with Thresholding): ", last_slice_aoi)

        mask_index = '{:03d}'.format(mask_index)
        combined_mask = combined_mask[first_slice_aoi:last_slice_aoi, :, :, :]
        scan = nib.load(('{}/msk_{}.nii.gz').format(input_scan_dir, mask_index))
        scan_data = scan.get_fdata()
        # scan_og = nib.Nifti1Image(scan_data, scan.affine)
        # nib.save(scan_og, ('{}/msk_{}_og.nii.gz').format(output_dir, mask_index))

        if (FriedLanderDataset == True):
            scan_data = np.expand_dims(scan_data, axis = -1)
            scan_data = scan_data.transpose(1, 0, 2, 3) 
            print('Unprocessed Scan Shape: ', scan_data.shape)
            scan_data = uniform_resizing(scan_data, resizing_size)
            scan_data = scan_data.astype('float32')
            scan_data /= np.max(scan_data)  # scale scans to [0, 1]
            print('Max Pixel Value in Scan: ', np.max(scan_data))

        scan_data = scan_data[first_slice_aoi:last_slice_aoi, :, :, :]
        print('Final Training Scans Input Shape: ', scan_data.shape)
        print('Final Training Masks Input Shape: ', combined_mask.shape)

    # np.savetxt('output.txt', scan_data[100,:,:,0], fmt="%d", delimiter=",")
    combined_mask = combined_mask.astype(int)
    combined_img = nib.Nifti1Image(combined_mask, mask.affine)
    print('Multi-class Labels: ', np.unique(combined_mask))
    nib.save(combined_img, ('{}/msk_{}.nii.gz').format(output_dir, mask_index))

    combined_img = nib.Nifti1Image(scan_data, scan.affine)
    nib.save(combined_img, ('{}/msk_{}.nii.gz').format(scan_dir, mask_index))

    print(('MSK Multiclass Mask Made Using {}, Saved To {}, and Region of Interest Slice Thresholding = {} !').format(aoi_fnames, ('{}/msk_{}.nii.gz').format(output_dir, mask_index), AOIThresholding))
    print('\n')