"""
Usage example: 
python split_labels.py --input_dir /projects/bodymaps/Pedro/data/JHH_liver_segments/ --output_dir /projects/bodymaps/Pedro/data/JHH_liver_segments_sep_labels/ --JHH --parts 3 --part 0
python split_labels.py --input_dir /projects/bodymaps/Data/UFO_OrganSubSegments_nnUNet/subSegments_output --output_dir /projects/bodymaps/Data/UCSF_pseudo_masks_subsegments --parts 3 --part 0

"""

import os
import glob
import argparse
import numpy as np
import nibabel as nib
from concurrent.futures import ProcessPoolExecutor
import concurrent
from tqdm import tqdm


labels = ['background',
            'kidney_right',
            'kidney_left',
            'kidney_lesion',
            'kidney_lesion_kidney_right',
            'kidney_lesion_kidney_left',
            'pancreas',
            'pancreatic_lesion',
            'liver',
            'liver_lesion',
            'spleen',
            'spleen_lesion',
            'colon',
            'colon_lesion',
            'stomach',
            'stomach_lesion',
            'duodenum',
            'duodenum_lesion',
            'common_bile_duct',
            'intestine',
            'aorta',
            'postcava',
            'adrenal_gland_left',
            'adrenal_gland_right',
            'adrenal_lesion',
            'adrenal_lesion_adrenal_gland_left',
            'adrenal_lesion_adrenal_gland_right',
            'gall_bladder',
            'gallbladder_lesion',
            'gallbladder_lesion_gall_bladder',
            'bladder',
            'bladder_lesion',
            'esophagus',
            'esophagus_lesion',
            'hepatic_vessel',
            'portal_vein_and_splenic_vein',
            'lung_left',
            'lung_right',
            'lung_lesion',
            'lung_lesion_lung_left',
            'lung_lesion_lung_right',
            'prostate',
            'prostate_lesion',
            'uterus_lesion',
            'uterus_lesion_prostate',#this is a bug, our prostate annotation should actually be called prostate/uterus
            'femur_left',
            'femur_right',
            'superior_mesenteric_artery',
            'veins',
            'bone_lesion',
            'breast_lesion',]

labels = {name: i for i, name in enumerate(labels, 0)}


# Classes we want to output explicitly
#out_labels = [
#    "liver",
#    "pancreas",
#    "kidney_right",
#    "kidney_left",
#    "liver_lesion",
#    "kidney_lesion",
#    "pancreatic_lesion",
#    "pancreas_head",
#    "pancreas_body",
#    "pancreas_tail"
#]
# Add "liver_segment_1" through "liver_segment_8"
#for i in range(1, 9):
#    out_labels.append("liver_segment_" + str(i))
out_labels = list(labels.keys())
#out_labels+=['liver','pancreas']

def parse_args():
    parser = argparse.ArgumentParser(
        description="Split segmentation labels into multiple .nii.gz files."
    )
    parser.add_argument(
        "--input_dir",
        type=str,
        required=True,
        help="Directory containing input .nii.gz files to process.",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        required=True,
        help="Directory where the splitted .nii.gz files will be saved.",
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=20,
        help="Number of parallel processes to use for splitting.",
    )
    parser.add_argument(
        "--restart",
        action='store_true',
        help="Overwrites files already saved.",
    )

    parser.add_argument(
        "--parts",
        type=int,
        default=1,
        help="Total number of parts to split the files into.",
    )
    parser.add_argument(
        "--part",
        type=int,
        default=0,
        help="Index of the part to process (0-indexed).",
    )
    parser.add_argument(
        "--JHH",
        action='store_true',
        help="Processes only JHH files.",
    )
    return parser.parse_args()


def build_mapping():
    """
    Build the 'mapping' dictionary:
    Each key is an output class name,
    each value is a list of label IDs that should be merged into that class.
    """

    # Map from lesion label to the organ sub-parts it can overlap with
    lesion_to_organs = {
        'kidney_lesion': ['kidney_right', 'kidney_left'],
        'adrenal_lesion': ['adrenal_gland_left', 'adrenal_gland_right'],
        'gallbladder_lesion': ['gall_bladder'],
        'lung_lesion': ['lung_left', 'lung_right'],
    }

    mapping = {}
    for key in out_labels:
        ids = []
        for name, val in labels.items():
            if name == key or name.startswith(key + "_"):
                ids.append(val)

        # Include overlap labels where this organ is the sub-segment
        # e.g., kidney_lesion_kidney_right should be included for kidney_right
        for lesion_name, organ_keys in lesion_to_organs.items():
            if key in organ_keys:
                overlap_label = lesion_name + "_" + key
                if overlap_label in labels:
                    ids.append(labels[overlap_label])

        # Special case: "pancreas" also includes "pancreatic"
        if key == "pancreas":
            for name, val in labels.items():
                if name == "pancreatic_lesion" or name.startswith("pancreatic_lesion_"):
                    ids.append(val)

        mapping[key] = sorted(set(ids))

    return mapping



def split_segmentation(input_nifti_path, output_folder, classes_dict):
    """
    Reads a NIfTI segmentation, then saves separate .nii.gz masks
    for each class (merging labels if necessary).
    """
    print('Splitting', input_nifti_path)
    seg_img = nib.load(input_nifti_path)
    # *** Fix: load as float, then cast to int16 (or directly from dataobj) ***
    seg_data = seg_img.get_fdata().astype(np.int16)
    # Or use: seg_data = np.asanyarray(seg_img.dataobj).astype(np.int16)

    # Ensure output_folder exists
    os.makedirs(output_folder, exist_ok=True)

    # We'll name output files using:
    #   <base of input file>_<class_name>.nii.gz
    base_name = os.path.basename(input_nifti_path)
    if base_name.endswith(".nii.gz"):
        base_name = base_name[:-7]
    else:
        base_name = os.path.splitext(base_name)[0]

    # Optionally create a subfolder for each input file
    this_out_folder = os.path.join(output_folder, base_name, 'predictions')
    os.makedirs(this_out_folder, exist_ok=True)

    for class_name, label_list in classes_dict.items():
        # Create binary mask: 1 if label in label_list, else 0
        mask = np.isin(seg_data, label_list).astype(np.uint8)

        mask_nifti = nib.Nifti1Image(mask, seg_img.affine, seg_img.header)
        out_file = os.path.join(this_out_folder, f"{class_name}.nii.gz")
        nib.save(mask_nifti, out_file)

    print(f"Finished splitting {input_nifti_path}!")

def process_file(input_path, output_dir, classes_dict):
    """Wrapper to process a single file."""
    split_segmentation(input_path, output_dir, classes_dict)

def main():
    args = parse_args()
    os.makedirs(args.output_dir, exist_ok=True)

    # Build the label mapping once
    classes_dict = build_mapping()
    print(classes_dict)

    # Gather all .nii.gz files in the input directory
    nifti_files = glob.glob(os.path.join(args.input_dir, "*.nii.gz"))
    if len(nifti_files) == 0:
        print(f"No NIfTI files found in {args.input_dir}!")
        return

    if args.JHH:
        nifti_files = [f for f in nifti_files if (('BDMAP_A' in f) or ('BDMAP_V' in f))]

    # remove from the list all files aleady saved
    if not args.restart:
        new_nifti_files = []
        for fpath in nifti_files:
            base_name = os.path.basename(fpath)
            if base_name.endswith(".nii.gz"):
                base_name = base_name[:-7]
            else:
                base_name = os.path.splitext(base_name)[0]
            
            # The same logic used in split_segmentation to build the output folder
            this_out_folder = os.path.join(args.output_dir, base_name, 'predictions')
            
            # Check if every class_name has a corresponding <class_name>.nii.gz
            all_exist = True
            for class_name in classes_dict.keys():
                out_file = os.path.join(this_out_folder, f"{class_name}.nii.gz")
                if not os.path.exists(out_file):
                    all_exist = False
                    break
            
            # Keep this file if not all masks are present
            if not all_exist:
                new_nifti_files.append(fpath)
        
        nifti_files = new_nifti_files

    # Split the nifti_files list into parts and process only the selected part.
    total_files = len(nifti_files)
    if args.parts > 1:
        files_per_part = (total_files + args.parts - 1) // args.parts  # round up
        start_idx = args.part * files_per_part
        end_idx = min(start_idx + files_per_part, total_files)
        nifti_files = nifti_files[start_idx:end_idx]
        print(f"Processing part {args.part+1}/{args.parts}: files {start_idx} to {end_idx-1}")
    else:
        print("Processing all files (only one part).")
    print(f"Total files to process: {len(nifti_files)}")

    # Use parallel processing
    with ProcessPoolExecutor(max_workers=args.num_workers) as executor:
        futures = []
        for fpath in nifti_files:
            futures.append(
                executor.submit(process_file, fpath, args.output_dir, classes_dict)
            )
        #for future in futures:
        #    future.result()
        for future in tqdm(concurrent.futures.as_completed(futures), total=len(futures), desc="Processing files"):
            future.result()

    print(f"All files processed. Output saved to: {args.output_dir}")

if __name__ == "__main__":
    main()

# version 2 - fixed label dict inversion, fixed substring matching in build_mapping (bladder/gall_bladder, prostate/uterus ambiguities)