# import os
# from batchgenerators.utilities.file_and_folder_operations import join, load_json, isfile, save_json, maybe_mkdir_p
# import nibabel as nib

# val_liver = load_json('/storage/chenqi/cvpr25/annotated_exp/nnUNet_preprocessed/Dataset016_AbdomenAtlas2.0_base_liver/splits_final.json')[0]['val']
# val_pancreas = load_json('/storage/chenqi/cvpr25/annotated_exp/nnUNet_preprocessed/Dataset017_AbdomenAtlas2.0_base_pancreas/splits_final.json')[0]['val']
# val_kidney = load_json('/storage/chenqi/cvpr25/annotated_exp/nnUNet_preprocessed/Dataset018_AbdomenAtlas2.0_base_kidney/splits_final.json')[0]['val']
# val_colon = load_json('/storage/chenqi/code/GEM-3D-ct/data_split/splits_final_colon_bdmap.json')[0]['val']
# val_eso = load_json('/storage/chenqi/code/GEM-3D-ct/data_split/splits_final_eso_bdmap.json')[0]['val']
# val_ute = load_json('/storage/chenqi/code/GEM-3D-ct/data_split/splits_final_uterus_bdmap.json')[0]['val']

# all_files = os.listdir('/storage/chenqi/data/our_data/nnUNet_preprocessed/Dataset104_AbdomenAtlas2.0_1102/gt_segmentations')
# all_files = [i.split('.')[0] for i in all_files]
# train_liver = list(set(all_files ) - set( val_liver) - set(val_pancreas) - set(val_kidney) - set(val_colon) - set(val_eso) - set(val_ute) )
# # breakpoint()
# write_content = []

# valid_train_liver = []
# for name in train_liver:
#     data = nib.load(os.path.join('/storage/chenqi/data/our_data/nnUNet_preprocessed/Dataset104_AbdomenAtlas2.0_1102/gt_segmentations', name)).get_fdata()
#     if (data.shape[0]==512 and data.shape[1]==512):
#         valid_train_liver.append(name)

# write_content.append({"train":train_liver, "val":val_liver})

# # breakpoint()

# splits_file = 'splits_final.json'
# save_json(write_content, splits_file)

import os
from batchgenerators.utilities.file_and_folder_operations import join, load_json, isfile, save_json, maybe_mkdir_p
import nibabel as nib
from multiprocessing import Pool, cpu_count

def is_valid_file(name):
    data = nib.load(os.path.join('/storage/chenqi/data/our_data/nnUNet_preprocessed/Dataset105_AbdomenAtlas2.0_1106/gt_segmentations', name+'.nii.gz')).get_fdata()
    return name if (data.shape[0] == 512 and data.shape[1] == 512) else None

def main():
    val_liver = load_json('/storage/chenqi/cvpr25/annotated_exp/nnUNet_preprocessed/Dataset016_AbdomenAtlas2.0_base_liver/splits_final.json')[0]['val']
    val_pancreas = load_json('/storage/chenqi/cvpr25/annotated_exp/nnUNet_preprocessed/Dataset017_AbdomenAtlas2.0_base_pancreas/splits_final.json')[0]['val']
    val_kidney = load_json('/storage/chenqi/cvpr25/annotated_exp/nnUNet_preprocessed/Dataset018_AbdomenAtlas2.0_base_kidney/splits_final.json')[0]['val']
    val_colon = load_json('/storage/chenqi/code/GEM-3D-ct/data_split/splits_final_colon_bdmap.json')[0]['val']
    val_eso = load_json('/storage/chenqi/code/GEM-3D-ct/data_split/splits_final_eso_bdmap.json')[0]['val']
    val_ute = load_json('/storage/chenqi/code/GEM-3D-ct/data_split/splits_final_uterus_bdmap.json')[0]['val']

    all_files = os.listdir('/storage/chenqi/data/our_data/nnUNet_preprocessed/Dataset105_AbdomenAtlas2.0_1106/gt_segmentations')
    all_files = [i.split('.')[0] for i in all_files]
    train_liver = list(set(all_files) - set(val_liver) - set(val_pancreas) - set(val_kidney) - set(val_colon) - set(val_eso) - set(val_ute))

    with Pool(cpu_count()) as pool:
        valid_train_liver = pool.map(is_valid_file, train_liver)
    
    valid_train_liver = [name for name in valid_train_liver if name is not None]

    write_content = [{"train": valid_train_liver, "val": val_liver}]

    splits_file = 'splits_final.json'
    save_json(write_content, splits_file)

if __name__ == "__main__":
    main()