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, print_stats_np
import time
import h5py
import matplotlib.pyplot as plt
from multiprocessing import set_start_method
import gzip
import os


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

for file_idx in range(188):
     image_recon = []
     save_file = []
     save_file = files[file_idx]
     image_recon = np.load(f'{mypath}/{save_file}/{save_file}_rss_recon_usampled.npy').astype(np.float32)
     np.save(f'{mypath}/{save_file}/{save_file}_rss_recon_usampled.npy',image_recon)
     print('done:',file_idx,'  ',f'{mypath}/{save_file}/{save_file}_rss_recon_usampled.npy')









quit()
mypath = f'/data/pca-rad/datasets/miccai_2022/K2S_MICCAI2022_GRP/train/data/TBrecon1/train/untarred'
patient_id = 'pat00335'
seg = np.load(f'{mypath}/{patient_id}/{patient_id}_seg.npy')
print_stats_np(seg,'seg_')
image = np.load(f'{mypath}/{patient_id}/{patient_id}_rss_recon.npy')
print_stats_np(image,'image_')
image_u = np.load(f'{mypath}/{patient_id}/{patient_id}_rss_recon_usampled.npy')
print_stats_np(image_u,'image_u_')

seg = sitk.GetImageFromArray(np.squeeze(seg))
sitk.WriteImage(seg, f"../test_seg_335.nii.gz")  
image = sitk.GetImageFromArray(np.squeeze(image))
sitk.WriteImage(image, f"../test_image_335.nii.gz")
image_u = sitk.GetImageFromArray(np.squeeze(image_u))
sitk.WriteImage(image_u, f"../test_image_u_335.nii.gz")
print(np.shape(seg))




# ROOT_H5_GZ  = r'../../../../../scratch/p290820/MICCAI/TBrecon-01-02-00335.h5.gz'
# SAVE_TO_DIR = r'/data/pca-rad/datasets/miccai_2022/K2S_MICCAI2022_GRP/train/data/TBrecon1/train/untarred/pat00335'
# unzipped = gzip.open(ROOT_H5_GZ,'rb')
# h5_file = h5py.File(unzipped,'r') # Keys: kspace, kspace_info, seg, seg_info

# # ksp_n = h5_file['kspace'][()]                   #np.complex64
# seg_n = h5_file['seg'][()].astype(np.float32)   #np.float32

# # fname_ksp = os.path.join(patient_path, f"pat{pat_num}_ksp.npy")
# fname_seg = os.path.join(SAVE_TO_DIR, f"pat00335_seg_test.npy")
# # np.save(fname_ksp, ksp_n)
# np.save(fname_seg, seg_n)



quit()

def inladen(patient_id,mypath):
     seg = np.load(f'{mypath}/{patient_id}/{patient_id}_seg.npy').astype(int)
     return seg


def load_seg(file,mypath):          
     patient_id = file
     print(patient_id)
     seg = inladen(patient_id,mypath)
     return seg

def calculate_hist(segs):
     background = np.sum(segs == 0)
     femoral_cartilage = np.sum(segs == 1)
     tibial_cartilage = np.sum(segs == 2)
     patellar_cartilage = np.sum(segs == 3)
     femur = np.sum(segs == 4)
     tibia = np.sum(segs == 5)
     patella = np.sum(segs == 6)
     output = [background,femoral_cartilage,tibial_cartilage,patellar_cartilage,femur,tibia,patella]
     # print('background',background)
     return output
     
if __name__ == '__main__':
     path = f'/data/pca-rad/datasets/miccai_2022/K2S_MICCAI2022_GRP/train/data/TBrecon1/train/untarred'
     files = [f for f in listdir(path)]
     print('files length',len(files))
     # For multiprocessing
     set_start_method("spawn")
     print('spawn done')
     segs_list = apply_parallel(files[:100], load_seg, 4, mypath = path)
     print('done loading 100')

     segs_list_200 = apply_parallel(files[100:200], load_seg, 4, mypath = path)
     print('done loading 200')

     segs_list_300 = apply_parallel(files[200:300], load_seg, 4, mypath = path)
     print('done loading 300')

     segs_list.extend(segs_list_200)
     segs_list.extend(segs_list_300)

     segs_n = np.stack(segs_list, axis=0)
     print('load done',np.shape(segs_n))

     output = apply_parallel(segs_list,calculate_hist,4)
     print('hist calc done',np.shape(np.stack(output)))
     # print(np.stack(output)[:,0])

     dict = {
          'background': np.stack(output)[:,0],
          'femoral_cartilage': np.stack(output)[:,0],
          'tibial_cartilage': np.stack(output)[:,0],
          'patellar_cartilage': np.stack(output)[:,0],
          'femur': np.stack(output)[:,0],
          'tibia': np.stack(output)[:,0],
          'patella': np.stack(output)[:,0]
     }

     # 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")