from glob import glob
from os.path import normpath, basename
import SimpleITK as sitk



def get_paths(main_dir):
    all_niftis = glob(main_dir, recursive=True)

    dwis_b800 = [i for i in all_niftis if ("diff" in i.lower() or "dwi" in i.lower()) and ("b-800" in i.lower() or "b800" in i.lower())] 
    
    dwis_b400 = [i for i in all_niftis if ("diff" in i.lower() or "dwi" in i.lower()) and ("b-400" in i.lower() or "b400" in i.lower())] 

    return dwis_b800, dwis_b400


pat_numbers = ['pat0132','pat0091','pat0352','pat0844','pat1006','pat0406','pat0128','pat0153','pat0062','pat0758','pat0932','pat0248','pat0129','pat0429','pat0181','pat0063','pat0674','pat0176','pat0366','pat0082']
load_path = '../../datasets/radboud_new/{pat_number}/2016/**/*.nii.gz'

for idx, pat_number in enumerate(pat_numbers):
    dwis_b800,dwis_b400 = get_paths(f'../../datasets/radboud_new/{pat_number}/2016/**/*.nii.gz')

    # load 
    dwi_b800 = sitk.ReadImage(dwis_b800, sitk.sitkFloat32)
    dwi_b400 = sitk.ReadImage(dwis_b400, sitk.sitkFloat32)
    # write 
    output_path_b800 = f'../temp/check_by_derya/{idx}_{pat_number}_b800.nii.gz'
    output_path_b400 = f'../temp/check_by_derya/{idx}_{pat_number}_b400.nii.gz'
    sitk.WriteImage(dwi_b800, output_path_b800)
    sitk.WriteImage(dwi_b400, output_path_b400)