from traceback import FrameSummary
import numpy as np
import SimpleITK as sitk
from os import listdir
from pip import main
from scipy.fftpack import fftshift, ifftshift,  ifftn
from umcglib.utils import apply_parallel
import time
import h5py
import matplotlib.pyplot as plt


mypath = f'/data/pca-rad/datasets/miccai_2022/K2S_MICCAI2022_GRP/train/data/TBrecon1/train/untarred'
files = [f for f in listdir(mypath)]

dict = {
     'background': [],
     'femoral_cartilage': [],
     'tibial_cartilage': [],
     'patellar_cartilage': [],
     'femur': [],
     'tibia': [],
     'patella': []
}

for file_idx in range(300):
     seg = []
     patient_id = files[file_idx]
     seg = np.load(f'{mypath}/{patient_id}/{patient_id}_seg.npy')
     # # 0: background; 1: femoral cartilage; 2: tibial cartilage; 3: patellar cartilage; 4: femur; 5: tibia; 6: patella. 

     dict['background'].append(np.sum(seg == 0))
     dict['femoral_cartilage'].append(np.sum(seg == 1))
     dict['tibial_cartilage'].append(np.sum(seg == 2))
     dict['patellar_cartilage'].append(np.sum(seg == 3))
     dict['femur'].append(np.sum(seg == 4)) 
     dict['tibia'].append(np.sum(seg == 5))
     dict['patella'].append(np.sum(seg == 6))

     print('done:',file_idx,'  ',f'{mypath}/{patient_id}/{patient_id}_seg.npy')

np.save('../dict.npy',dict)

for keys in dict:
     data = dict[f'keys']
     plt.hist(data)
     plt.title(f"historgram of {keys}")
     plt.savefig(f"../{keys}.png", dpi=300)
     print('done  ',keys)



# print(f.keys())
# print(np.shape(f['us_mask'][()]))
# print(type(f['us_mask'][()]))







# quit()
# seg = sitk.GetImageFromArray(np.squeeze(seg))
# sitk.WriteImage(seg, f"../test_seg_2.nii.gz")  
# img_s = sitk.GetImageFromArray(np.squeeze(image))
# # img_s.CopyInformation(seg)
# sitk.WriteImage(img_s, f"../test_image_3.nii.gz")
# img_s = sitk.GetImageFromArray(np.squeeze(coil_image))
# # img_s.CopyInformation(seg)
# sitk.WriteImage(img_s, f"../test_last_coil_image_2.nii.gz")