import json
import os
import nibabel as nib
import csv

from operator import itemgetter

# PATH TO PREPROCESSED DATA
raw_data_path = '/home/lab/nnUNet_data/nnUNet_raw_data_base/nnUNet_raw_data/Task500_BrainMets'
pixdim_ind = [1,2,3] # Indexes at which the voxel size [x,y,z] is stored


# PATH TO JSON FILE
with open('/home/lab/nnUNet_data/RESULTS_FOLDER/nnUNet/3d_fullres/Task500_BrainMets/nnUNetTrainerV2__nnUNetPlansv2.1/fold_4/validation_raw/summary.json') as file:
  data = json.load(file)

with open('json_parsed.csv', mode='w') as csv_file:
    csv_writer = csv.writer(csv_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)
    csv_writer.writerow(['Case Number', 'Dice Score', 'Voxel Size-X', 'Voxel Size-Y', 'Voxel Size-Z'])

    for img in data['results']['all']: 
        # Get dice score on image
        dice = img['1']['Dice']

        # Get nifti data on image
        img_filename = (os.path.basename(img['reference']).split('.'))[0]
        img_ni = nib.load(raw_data_path + '/imagesTr/' + img_filename + '_0000.nii.gz')
        label_ni = nib.load(raw_data_path + '/labelsTr/' + img_filename + '.nii.gz')

        voxel_size = itemgetter(*pixdim_ind)(img_ni.header["pixdim"])

        # Get tumor dimensions
        # tumor_size = 

        # Get case number corresponding to image
        case_number = img_filename.split('_')[1]

        # Write to csv file
        csv_writer = csv.writer(csv_file, delimiter=',', quotechar='"', quoting=csv.QUOTE_MINIMAL)

        csv_writer.writerow([case_number, dice, voxel_size[0], voxel_size[1], voxel_size[2]])


[((15, 9, 15, 24), 'json.load', 'json.load', ({(15, 19, 15, 23): 'file'}, {}), '(file)', False, 'import json\n'), ((18, 17, 18, 94), 'csv.writer', 'csv.writer', (), '', False, 'import csv\n'), ((27, 17, 27, 87), 'nibabel.load', 'nib.load', ({(27, 26, 27, 86): "raw_data_path + '/imagesTr/' + img_filename + '_0000.nii.gz'"}, {}), "(raw_data_path + '/imagesTr/' + img_filename + '_0000.nii.gz')", True, 'import nibabel as nib\n'), ((28, 19, 28, 84), 'nibabel.load', 'nib.load', ({(28, 28, 28, 83): "raw_data_path + '/labelsTr/' + img_filename + '.nii.gz'"}, {}), "(raw_data_path + '/labelsTr/' + img_filename + '.nii.gz')", True, 'import nibabel as nib\n'), ((39, 21, 39, 98), 'csv.writer', 'csv.writer', (), '', False, 'import csv\n'), ((30, 21, 30, 44), 'operator.itemgetter', 'itemgetter', ({(30, 32, 30, 43): '*pixdim_ind'}, {}), '(*pixdim_ind)', False, 'from operator import itemgetter\n'), ((26, 24, 26, 58), 'os.path.basename', 'os.path.basename', ({(26, 41, 26, 57): "img['reference']"}, {}), "(img['reference'])", False, 'import os\n')]