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

mypath = f'/data/pca-rad/datasets/miccai_2022/K2S_MICCAI2022_GRP/train/data/TBrecon1/train/untarred'
files = [f for f in listdir(mypath)]
# save_file = files[1]
# seg = np.load(f'{mypath}/{save_file}/{save_file}_seg.npy')
# kspace = np.load(f'{mypath}/{save_file}/{save_file}_ksp.npy')

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 = np.transpose(kspace, (2, 0, 1, 3))
     coil_images = apply_parallel(kspace, 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

for file_idx in range(167,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.npy',image_recon)
     print('done:',file_idx,'  ',f'{mypath}/{save_file}/{save_file}_rss_recon.npy')


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