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

mypath = f'/data/pca-rad/datasets/miccai_2022/K2S_MICCAI2022_GRP/train/data/TBrecon1/train/untarred'
files = [f for f in listdir(mypath)]
KSPACE_U_PATH = f'/data/pca-rad/datasets/miccai_2022/K2S_MICCAI2022_GRP/train/data/TBrecon1/train/R_8_mask.h5'
f = h5py.File(KSPACE_U_PATH,'r')
kspace_mask = f['us_mask'][()]

def k2p_to_img(kspace_single_coil):
     kspace_padded = np.pad(np.squeeze(kspace_single_coil[:,:,:]),[(128,128),(128,128),(0,0)],'constant')
     coil_image = np.abs(fftshift(fftshift(ifftn(ifftshift(kspace_padded)),axes=0),axes=1))
     return coil_image

def img_from_ksp(
     kspace: np.array):
     kspace_masked = np.multiply(kspace, kspace_mask)
     kspace_masked = np.transpose(kspace_masked, (2, 0, 1, 3))
     coil_images = apply_parallel(kspace_masked, k2p_to_img, 12)
     image = np.round(np.sqrt(np.sum(np.power(coil_images,2),0)))
     image = image[:,:,2:-2]
     image = image[:,:,::-1]
     return image.astype(np.float32)

for file_idx in range(187,300):
     image_recon = []
     kspace = []
     save_file = []
     save_file = files[file_idx]
     kspace = np.load(f'{mypath}/{save_file}/{save_file}_ksp.npy')
     image_recon = img_from_ksp(kspace)
     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')

    #  img_s = sitk.GetImageFromArray(np.squeeze(image_recon))
    #  sitk.WriteImage(img_s, f"../job_test_undersampled_{save_file}.nii.gz")
     




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