#!python -W ignore::DeprecationWarning
# coding: utf-8

"""
REQUIRES
> greedy
> HD-BET
> MNI atlas files
> Model files
"""
import os
import multiprocessing
import sys
import tempfile
import shutil
import argparse

from lst_ai.strip import run_hdbet, apply_mask
from lst_ai.register import mni_registration, rigid_reg, apply_warp_label, apply_warp_interp
from lst_ai.segment import unet_segmentation
from lst_ai.annotate import annotate_lesions
from lst_ai.stats import compute_stats
from lst_ai.utils import DATA_DIR, download_data, harmonize_affines

if __name__ == "__main__":

    print("###########################\n")
    print("Thank you for using LST-AI. If you publish your results, please cite our paper:")
    print("Wiltgen T, McGinnis J, Schlaeger S, Kofler F, Voon C, Berthele A, Bischl D, Grundl L, Will N, Metz M, Schinz D, "
          "Sepp D, Prucker P, Schmitz-Koep B, Zimmer C, Menze B, Rueckert D, Hemmer B, Kirschke J, Mühlau M, Wiestler B. "
          "LST-AI: A Deep Learning Ensemble for Accurate MS Lesion Segmentation. NeuroImage: Clinical, Volume 42, 2024. "
          "https://doi.org/10.1016/j.nicl.2024.103611.")
    print("###########################\n")

    parser = argparse.ArgumentParser(description='Segment / Label MS lesions according to McDonald criteria.')

    # Input Images
    parser.add_argument('--t1',
                        dest='t1',
                        help='Path to T1 image',
                        type=str,
                        required=True)
    parser.add_argument('--flair',
                        dest='flair',
                        help='Path to FLAIR image',
                        type=str,
                        required=True)
    # Output Images
    parser.add_argument('--output',
                        dest='output',
                        help='Path to output segmentation nifti path.',
                        type=str,
                        required=True)

    parser.add_argument('--existing_seg',
                        dest='existing_seg',
                        default='',
                        help='Path to output segmentation image (will be saved in FLAIR space)',
                        type=str)

    # Temporary directory
    parser.add_argument('--temp',
                        dest='temp',
                        default='',
                        help='Path to temp directory.',
                        type=str)

    # Mode (segment only, annotate only, segment+annotate (default))
    parser.add_argument('--segment_only',
                        action='store_true',
                        dest='segment_only',
                        help='Only perform the segmentation, and skip lesion annotation.')

    parser.add_argument('--annotate_only',
                        action='store_true',
                        dest='annotate_only',
                        help='Only annotate lesion files without segmentation of lesions.')

    parser.add_argument('--stripped',
                        action='store_true',
                        dest='stripped',
                        help='Images are already skull stripped. Skip skull-stripping.')

    # Model Settings
    parser.add_argument('--threshold',
                        dest='threshold',
                        help='Threshold for binarizing the joint segmentation (default: 0.5)',
                        type=float,
                        default=0.5)
    
    parser.add_argument('--lesion_threshold',
                        dest='lesion_threshold',
                        help='minimum lesion size',
                        type=int,
                        default=0)
        
    parser.add_argument('--clipping',
                        dest='clipping',
                        help='Clipping (min & max) for standardization of image intensities (default: 0.5 99.5).',
                        nargs='+',
                        type=str,
                        default=('0.5','99.5'))

    parser.add_argument('--fast-mode',
                        action='store_true',
                        dest='fast',
                        help='Only use one model for hd-bet.')
    
    parser.add_argument('--probability_map',
                        action='store_true',
                        dest='probability_map',
                        help='Additionally store the probability maps of the three models and of the ensemble network.',
                        default=False)

    # Computing Resources
    parser.add_argument('--device',
                        dest='device',
                        help='Either int for GPU ID or "cpu" for CPU (default: 0)',
                        type=str,
                        default='0')

    parser.add_argument('--threads',
                        dest='threads',
                        help='Number of threads to be used for registration (default: all available)',
                        type=int,
                        default=multiprocessing.cpu_count())

    args = parser.parse_args()

    print(f"Looking for model weights in {DATA_DIR}.")
    download_data(path=DATA_DIR)

    # Sanity Checks
    assert os.path.exists(args.t1), 'LST.AI aborted. T1w Image Path does not exist.'
    assert os.path.exists(args.flair), 'LST.AI aborted. Flair Image Path does not exist.'
    assert str(args.t1).endswith(".nii.gz"), 'Please provide T1w as a zipped nifti.'
    assert str(args.flair).endswith(".nii.gz"), 'Please provide FLAIR as a zipped nifti.'
    assert not os.path.isfile(args.output), 'Please provide an output path, not a filename.'
    assert len(args.clipping)==2, 'Please provide two values for clipping (min & max).'

    # convert input
    min_clip, max_clip = float(args.clipping[0]), float(args.clipping[1])

    if not args.temp:
        work_dir = tempfile.mkdtemp(prefix='lst_ai_')
    else:
        work_dir = os.path.abspath(args.temp)
        # make temp directory in case it does not exist
        if not os.path.exists(work_dir):
            os.makedirs(work_dir)

    # Harmonize the qform and sform of the inputs before anything reads them.
    # nibabel (used throughout this pipeline) prefers the sform, while greedy
    # reads the qform; if a prior co-registration updated only the sform, the two
    # disagree and the segmentation is mislocated. See CompImg/LST-AI#44. The
    # originals are left untouched; the pipeline continues on work-dir copies.
    harmonized_t1 = os.path.join(work_dir, 'input_T1w.nii.gz')
    harmonized_flair = os.path.join(work_dir, 'input_FLAIR.nii.gz')
    harmonize_affines(args.t1, harmonized_t1)
    harmonize_affines(args.flair, harmonized_flair)
    args.t1 = harmonized_t1
    args.flair = harmonized_flair
    # An existing segmentation is read in FLAIR space by the same tools, so it
    # has to be harmonized the same way or it would disagree with the harmonized
    # FLAIR exactly as the raw inputs disagreed with each other.
    if args.existing_seg and os.path.isfile(args.existing_seg):
        harmonized_seg = os.path.join(work_dir, 'input_existing_seg.nii.gz')
        harmonize_affines(args.existing_seg, harmonized_seg)
        args.existing_seg = harmonized_seg

    #  Define Image Paths (original space)
    path_orig_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-t1w_T1w.nii.gz')
    path_orig_flair = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_FLAIR.nii.gz')
    path_orig_stripped_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-t1w_desc-stripped_T1w.nii.gz')
    path_orig_stripped_flair = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_desc-stripped_FLAIR.nii.gz')

    #  Define Image Paths (FLAIR space)
    path_flair_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_T1w.nii.gz')
    path_flair_stripped_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_desc-stripped_T1w.nii.gz')

    #  Define Image Paths (MNI space)
    path_mni_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_T1w.nii.gz')
    path_mni_flair = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_FLAIR.nii.gz')
    path_mni_stripped_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_desc-stripped_T1w.nii.gz')
    path_mni_stripped_flair = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_desc-stripped_FLAIR.nii.gz')

    # Masks
    path_orig_brainmask_t1w = os.path.join(work_dir, 'sub-X_ses-Y_space-t1w_brainmask.nii.gz')
    path_orig_brainmask_flair = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_brainmask.nii.gz')
    path_mni_brainmask = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_brainmask.nii.gz')

    # Probability Maps MNI & FLAIR
    probmap_mni_segmentation = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-lst_prob.nii.gz')
    probmap_mni_model1 = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-lst_prob_1.nii.gz')
    probmap_mni_model2 = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-lst_prob_2.nii.gz')
    probmap_mni_model3 = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-lst_prob_3.nii.gz')
    probmap_FLAIR_segmentation = os.path.join(work_dir, 'sub-X_ses-Y_space-FLAIR_seg-lst_prob.nii.gz')
    probmap_FLAIR_model1 = os.path.join(work_dir, 'sub-X_ses-Y_space-FLAIR_seg-lst_prob_1.nii.gz')
    probmap_FLAIR_model2 = os.path.join(work_dir, 'sub-X_ses-Y_space-FLAIR_seg-lst_prob_2.nii.gz')
    probmap_FLAIR_model3 = os.path.join(work_dir, 'sub-X_ses-Y_space-FLAIR_seg-lst_prob_3.nii.gz')

    # Temp Segmentation results
    path_flair_segmentation =  os.path.join(work_dir, 'sub-X_ses-Y_space-flair_seg-lst.nii.gz')
    path_mni_segmentation = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-lst.nii.gz')
    path_flair_annotated_segmentation =  os.path.join(work_dir, 'sub-X_ses-Y_space-flair_desc-annotated_seg-lst.nii.gz')
    path_mni_annotated_segmentation = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_desc-annotated_seg-lst.nii.gz')

    # Output paths (in original space)
    filename_output_segmentation = "space-flair_seg-lst.nii.gz"
    filename_output_annotated_segmentation = "space-flair_desc-annotated_seg-lst.nii.gz"

    # FastSurfer segmentations
    path_flair_fastsurfer_seg = os.path.join(work_dir, 'sub-X_ses-Y_space-flair_seg-fastsurfer.nii.gz')
    path_mni_fastsurfer_seg = os.path.join(work_dir, 'sub-X_ses-Y_space-mni_seg-fastsurfer.nii.gz')

    # Stats
    filename_output_stats_segmentation = "lesion_stats.csv"
    filename_output_stats_annotated_segmentation = "annotated_lesion_stats.csv"

    # affines
    path_affine_mni_t1w = os.path.join(work_dir, 'affine_t1w_to_mni.mat')
    path_affine_mni_flair = os.path.join(work_dir, 'affine_flair_to_mni.mat')

    # atlas files
    t1w_atlas = os.path.join(DATA_DIR, "atlas", "sub-mni152_space-mni_t1.nii.gz")
    t1w_atlas_stripped = os.path.join(DATA_DIR, "atlas", "sub-mni152_space-mni_t1bet.nii.gz")

    # model weights
    model_directory = os.path.join(DATA_DIR, 'model')

    # make output path (in case it does not exist)
    if not os.path.exists(args.output):
        os.makedirs(args.output)

    if os.path.isfile(args.existing_seg) and not args.annotate_only:
        print("Existing segmentation may only be used with --annotate_only flag.")
        print("Aborted.")
        sys.exit(1)

    # Annotation only
    if args.annotate_only:
        print("LST-AI assumes existing segmentation to be in FLAIR space.")
        if os.path.isfile(args.existing_seg):
            shutil.copy(args.existing_seg, path_flair_segmentation)
        else:
            print("Existing segmentation does not exist.")
            print("Please provide a valid path via --existing_seg.")
            sys.exit(1)

        # Register T1w to FLAIR native
        t1w_native= os.path.join(work_dir,"T1w.nii.gz")
        native_affine = os.path.join(work_dir,"affine_native_t1_flair.mat")
        shutil.copy(args.t1, t1w_native)

        if not args.stripped:

            shutil.copy(args.flair, path_orig_flair)
            rigid_reg(moving=t1w_native,
                      fixed=path_orig_flair,
                      affine=native_affine,
                      destination=path_flair_t1w,
                      n_threads=args.threads)

            # strip only T1w (in FLAIR space)
            if args.fast:
                run_hdbet(input_image=path_flair_t1w, output_image=path_flair_stripped_t1w, device=args.device, mode="fast")
            else:
                run_hdbet(input_image=path_flair_t1w, output_image=path_flair_stripped_t1w, device=args.device, mode="accurate")

            # move processed mask to correct naming convention
            hdbet_mask = path_flair_stripped_t1w.replace(".nii.gz", "_bet.nii.gz")
            shutil.move(hdbet_mask, path_orig_brainmask_flair)

            # apply brain mask to FLAIR
            apply_mask(input_image=path_orig_flair,
                       mask=path_orig_brainmask_flair,
                       output_image=path_orig_stripped_flair)
        else:
            shutil.copy(args.flair, path_orig_stripped_flair)
            rigid_reg(moving=t1w_native,
                      fixed=path_orig_stripped_flair,
                      affine=native_affine,
                      destination=path_flair_stripped_t1w,
                      n_threads=args.threads)

        annotate_lesions(t1w_im=path_flair_stripped_t1w,
                         lesion_mask=path_flair_segmentation,
                         t1w_seg=path_flair_fastsurfer_seg,
                         lesion_mask_annotated=path_flair_annotated_segmentation,
                         device=args.device)

        shutil.copy(path_flair_annotated_segmentation,
                    os.path.join(args.output, filename_output_annotated_segmentation))


    # Segmentation only + (opt. Annotation)
    else:

        if args.stripped:
            print("Images have been provided as skull-stripped images. Skipping skull-stripping via HD-BET.")
            shutil.copy(args.t1, path_orig_stripped_t1w)
            shutil.copy(args.flair, path_orig_stripped_flair)

            mni_registration(t1w_atlas_stripped,
                             path_orig_stripped_t1w,
                             path_orig_stripped_flair,
                             path_mni_stripped_t1w,
                             path_mni_stripped_flair,
                             path_affine_mni_t1w,
                             path_affine_mni_flair,
                             n_threads=args.threads)

        else:
            print("Images need to be skull-stripped. Processing with HD-BET.")
            shutil.copy(args.t1, path_orig_t1w)
            shutil.copy(args.flair, path_orig_flair)

            # first register to MNI space
            mni_registration(t1w_atlas,
                             path_orig_t1w,
                             path_orig_flair,
                             path_mni_t1w,
                             path_mni_flair,
                             path_affine_mni_t1w,
                             path_affine_mni_flair,
                             n_threads=args.threads)

            # then skull strip T1w
            if args.fast:
                run_hdbet(input_image=path_mni_t1w, output_image=path_mni_stripped_t1w, device=args.device, mode="fast")
            else:
                run_hdbet(input_image=path_mni_t1w, output_image=path_mni_stripped_t1w, device=args.device, mode="accurate")

            # move processed mask to correct naming convention
            hdbet_mask = path_mni_stripped_t1w.replace(".nii.gz", "_bet.nii.gz")
            shutil.move(hdbet_mask, path_mni_brainmask)

            # then apply brain mask to FLAIR
            apply_mask(input_image=path_mni_flair,
                       mask=path_mni_brainmask,
                       output_image=path_mni_stripped_flair)

            # create masked images in original space
            apply_warp_label(image_org_space=path_orig_t1w,
                             affine=path_affine_mni_t1w,
                             origin=path_mni_brainmask,
                             target=path_orig_brainmask_t1w,
                             reverse=True,
                             n_threads=args.threads)

            apply_warp_label(image_org_space=path_orig_flair,
                             affine=path_affine_mni_flair,
                             origin=path_mni_brainmask,
                             target=path_orig_brainmask_flair,
                             reverse=True,
                             n_threads=args.threads)

            apply_mask(input_image=path_orig_t1w,
                       mask=path_orig_brainmask_t1w,
                       output_image=path_orig_stripped_t1w)

            apply_mask(input_image=path_orig_flair,
                       mask=path_orig_brainmask_flair,
                       output_image=path_orig_stripped_flair)


        # Segmentation
        print("Running LST Segmentation.")
        unet_segmentation(mni_t1=path_mni_stripped_t1w,
                          mni_flair=path_mni_stripped_flair,
                          model_path=model_directory,
                          output_segmentation_path=path_mni_segmentation,
                          device=args.device,
                          probmap=args.probability_map,
                          output_prob_path=probmap_mni_segmentation,
                          output_prob1_path=probmap_mni_model1,
                          output_prob2_path=probmap_mni_model2,
                          output_prob3_path=probmap_mni_model3,
                          input_shape=(192,192,192),
                          threshold=args.threshold,
                          clipping=(min_clip, max_clip),
                          lesion_thr=int(args.lesion_threshold))
        
        # warp probability maps to FLAIR space if flag is set
        if args.probability_map:
            apply_warp_interp(image_org_space=path_orig_stripped_flair,
                              affine=path_affine_mni_flair,
                              origin=probmap_mni_segmentation,
                              target=probmap_FLAIR_segmentation,
                              reverse=True,
                              n_threads=args.threads)
            apply_warp_interp(image_org_space=path_orig_stripped_flair,
                              affine=path_affine_mni_flair,
                              origin=probmap_mni_model1,
                              target=probmap_FLAIR_model1,
                              reverse=True,
                              n_threads=args.threads)
            apply_warp_interp(image_org_space=path_orig_stripped_flair,
                              affine=path_affine_mni_flair,
                              origin=probmap_mni_model2,
                              target=probmap_FLAIR_model2,
                              reverse=True,
                              n_threads=args.threads)
            apply_warp_interp(image_org_space=path_orig_stripped_flair,
                              affine=path_affine_mni_flair,
                              origin=probmap_mni_model3,
                              target=probmap_FLAIR_model3,
                              reverse=True,
                              n_threads=args.threads)

        # warp segmentation back to original FLAIR space
        # for now, we do not threshold the lesion sizes of the warped native image
        apply_warp_label(image_org_space=path_orig_stripped_flair,
                         affine=path_affine_mni_flair,
                         target=path_flair_segmentation,
                         origin=path_mni_segmentation,
                         reverse=True,
                         n_threads=args.threads)

        # store the segmentations
        shutil.copy(path_flair_segmentation, os.path.join(args.output, filename_output_segmentation))

        # Annotation
        if not args.segment_only:
            annotate_lesions(t1w_im=path_mni_stripped_t1w,
                             lesion_mask=path_mni_segmentation,
                             t1w_seg=path_mni_fastsurfer_seg,
                             lesion_mask_annotated=path_mni_annotated_segmentation,
                             device=args.device)

            # warp results back to original FLAIR space
            apply_warp_label(image_org_space=path_orig_stripped_flair,
                             affine=path_affine_mni_flair,
                             origin=path_mni_annotated_segmentation,
                             target=path_flair_annotated_segmentation,
                             reverse=True,
                             n_threads=args.threads)

            # store the segmentations
            shutil.copy(path_flair_annotated_segmentation, os.path.join(args.output, filename_output_annotated_segmentation))

    # Compute Stats of (annotated) segmentation if they exist
    if os.path.exists(path_flair_segmentation):
        compute_stats(mask_file=path_flair_segmentation,
                      output_file=os.path.join(args.output, filename_output_stats_segmentation),
                      multi_class=False)

    if os.path.exists(path_flair_annotated_segmentation):
        compute_stats(mask_file=path_flair_annotated_segmentation,
                      output_file=os.path.join(args.output, filename_output_stats_annotated_segmentation),
                      multi_class=True)

    print(f"Results in {work_dir}")
    if not args.temp:
        print("Delete temporary directory: {work_dir}")
        shutil.rmtree(work_dir)
    print("Done.")
