from typing import List, Optional

import click

from clinica.pipelines import cli_param
from clinica.utils.pet import LIST_SUVR_REFERENCE_REGIONS

pipeline_name = "deeplearning-prepare-data"


@click.command(name=pipeline_name)
@cli_param.argument.caps_directory
@click.argument(
    "modality",
    type=click.Choice(["t1-linear", "t1-extensive", "pet-linear", "custom"]),
    default="t1-linear",
)
@click.argument(
    "extract-method",
    type=click.Choice(["image", "slice", "patch", "roi"]),
    default="image",
)
@cli_param.option_group.pipeline_specific_options
@cli_param.option_group.option(
    "-uui",
    "--use_uncropped_image",
    is_flag=True,
    help="Use the uncropped image instead of the cropped image generated by t1-linear or pet-linear.",
)
@cli_param.option_group.custom_pipeline_options(
    "Pipeline options if you chose `patch` extraction"
)
@cli_param.option_group.option(
    "-ps",
    "--patch_size",
    default=50,
    show_default=True,
    help="Patch size if EXTRACT_METHOD is `patch`.",
)
@cli_param.option_group.option(
    "-ss",
    "--stride_size",
    default=50,
    show_default=True,
    help="Stride size if EXTRACT_METHOD is ``patch.",
)
@cli_param.option_group.custom_pipeline_options(
    "Pipeline options if you chose `slice` extraction"
)
@cli_param.option_group.option(
    "-sd",
    "--slice_direction",
    type=click.IntRange(0, 2),
    default=0,
    show_default=True,
    help="Slice direction if EXTRACT_METHOD is `slice`. 0: Sagittal plane, 1: Coronal plane, 2: Axial plane.",
)
@cli_param.option_group.option(
    "-sm",
    "--slice_mode",
    type=click.Choice(["rgb", "single"]),
    default="rgb",
    show_default=True,
    help=(
        "Slice mode if EXTRACT_METHOD is `slice`. rgb: Save the slice in three identical channels, "
        "single: Save the slice in a single channel."
    ),
)
@cli_param.option_group.custom_pipeline_options(
    "Pipeline options if you chose `roi` extraction"
)
@cli_param.option_group.option(
    "--roi_list",
    multiple=True,
    help="List of regions to be extracted.",
)
@cli_param.option_group.option(
    "--roi_uncrop_output",
    is_flag=True,
    help="Disable cropping option so the output tensors have the same size than the whole image.",
)
@cli_param.option_group.option(
    "-ct",
    "--custom_template",
    help="Name of the template used when modality is set to `custom`.",
)
@cli_param.option_group.option(
    "-cmp",
    "--custom_mask_pattern",
    default="",
    help=(
        "Select only the masks containing this pattern. The mask with the shortest name is taken. "
        "This argument is taken into account only when the modality is set to `custom`."
    ),
)
@cli_param.option_group.custom_pipeline_options(
    "Pipeline options if you chose `pet-linear` pipeline"
)
@cli_param.option_group.option(
    "--acq_label",
    type=click.Choice(["av45", "fdg"]),
    help=(
        "Name of the label given to the PET acquisition, specifying  the tracer used (acq-<acq_label>). "
        "For instance it can be 'fdg' for fluorodeoxyglucose or 'av45' for florbetapir. This option only applies to "
        "the `pet-linear` pipeline."
    ),
)
@cli_param.option_group.option(
    "-suvr",
    "--suvr_reference_region",
    type=click.Choice(LIST_SUVR_REFERENCE_REGIONS),
    help=(
        "Intensity normalization using the average PET uptake in reference regions resulting in a standardized uptake "
        "value ratio (SUVR) map. It can be cerebellumPons or cerebellumPon2 (used for amyloid tracers) or pons or "
        "pons2 (used for 18F-FDG tracers). This option only applies to `pet-linear` pipeline."
    ),
)
@cli_param.option_group.custom_pipeline_options(
    "Pipeline options if you chose `custom` modality"
)
@cli_param.option_group.option(
    "-cn",
    "--custom_suffix",
    default="",
    help=(
        "Suffix to append to filenames for a custom modality, for instance "
        "`graymatter_space-Ixi549Space_modulated-off_probability.nii.gz`, or "
        "`segm-whitematter_probability.nii.gz`"
    ),
)
@cli_param.option_group.common_pipelines_options
@cli_param.option.subjects_sessions_tsv
@cli_param.option.working_directory
@cli_param.option.n_procs
def cli(
    caps_directory: str,
    modality: str,
    extract_method: str,
    use_uncropped_image: bool = False,
    patch_size: int = 50,
    stride_size: int = 50,
    slice_direction: int = 0,
    slice_mode: str = "rgb",
    roi_list: Optional[List[str]] = None,
    roi_uncrop_output: bool = False,
    custom_template: Optional[str] = None,
    custom_mask_pattern: str = "",
    acq_label: Optional[str] = None,
    suvr_reference_region: Optional[str] = None,
    custom_suffix: str = "",
    subjects_sessions_tsv: Optional[str] = None,
    working_directory: Optional[str] = None,
    n_procs: Optional[int] = None,
) -> None:
    """Prepare data generated by Clinica for PyTorch with tensor extraction.

       It currently supports three preprocessed modalities (t1-linear, t1-extensive and pet-linear).

       Extraction of custom modalities is possible using the right argument and the option `--custom_suffix`.

       Extracted tensors can be obtained in three forms: the whole 3D image, 3D patches or 2D slices.

    See https://aramislab.paris.inria.fr/clinica/docs/public/latest/Pipelines/DeepLearning_PrepareData/
    """
    from networkx import Graph

    from clinica.utils.ux import print_end_pipeline

    from .deeplearning_prepare_data_pipeline import DeepLearningPrepareData

    parameters = {
        "modality": modality,
        "extract_method": extract_method,
        "patch_size": patch_size,
        "stride_size": stride_size,
        "slice_direction": slice_direction,
        "slice_mode": slice_mode,
        "roi_list": roi_list,
        "roi_uncrop_output": roi_uncrop_output,
        "custom_template": custom_template,
        "custom_mask_pattern": custom_mask_pattern,
        "use_uncropped_image": use_uncropped_image,
        "custom_suffix": custom_suffix,
        "acq_label": acq_label,
        "suvr_reference_region": suvr_reference_region,
    }

    pipeline = DeepLearningPrepareData(
        caps_directory=caps_directory,
        tsv_file=subjects_sessions_tsv,
        base_dir=working_directory,
        parameters=parameters,
        name=pipeline_name,
    )

    exec_pipeline = (
        pipeline.run(plugin="MultiProc", plugin_args={"n_procs": n_procs})
        if n_procs
        else pipeline.run()
    )

    if isinstance(exec_pipeline, Graph):
        print_end_pipeline(
            pipeline_name, pipeline.base_dir, pipeline.base_dir_was_specified
        )


if __name__ == "__main__":
    cli()
