from nnunetv2.dataset_conversion.generate_dataset_json import generate_dataset_json
from batchgenerators.utilities.file_and_folder_operations import *

import argparse

def main(args):
    
    labels = {
            "background": 0,
            'adrenal_gland_left':1,
            'adrenal_gland_right':2,
            'aorta':3,
            'bladder':4,
            'celiac_artery':5,
            'colon':6,
            'common_bile_duct':7,
            'duodenum':8,
            'femur_left':9,
            'femur_right':10,
            'gall_bladder':11,
            'kidney_left':12,
            'kidney_right':13,
            'liver':14,
            'lung_left':15,
            'lung_right':16,
            'pancreas':17,
            'pancreas_body':18,
            'pancreas_head':19,
            'pancreas_tail':20,
            'pancreatic_duct':21,
            'postcava':22,
            'prostate':23,
            'spleen':24,
            'stomach':25,
            'superior_mesenteric_artery':26,
            'veins':27,
            'pancreatic_tumor':28
        }
    generate_dataset_json(
        join(args.raw_dir, args.dataset_name),
        {0: 'CT'},  # this was a mistake we did at the beginning and we keep it like that here for consistency
        labels,
        3145,
        '.nii.gz',
        None,
        name=args.dataset_name,
        overwrite_image_reader_writer='NibabelIOWithReorient'
    )

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument("--raw_dir", type=str, default='/mnt/bodymaps/ePAI/nnUNet/raw', help='Path to the raw data directory')
    parser.add_argument("--dataset_name", type=str, default='Dataset1013_ePAI_3MM', help='Name of the target dataset')
    
    args = parser.parse_args()

    main(args)