# coding: utf8

import argparse
from distutils.util import strtobool

from colorama import Fore

TRAIN_CATEGORIES = {
    # General parent group
    'POSITIONAL': '%sPositional arguments%s' % (Fore.BLUE, Fore.RESET),
    'COMPUTATIONAL': '%sComputational resources%s' % (Fore.BLUE, Fore.RESET),
    'DATA': '%sData management%s' % (Fore.BLUE, Fore.RESET),
    'CROSS-VALIDATION': '%sCross-validation arguments%s' % (Fore.BLUE, Fore.RESET),
    'OPTIMIZATION': '%sOptimization parameters%s' % (Fore.BLUE, Fore.RESET),
    # Other parent groups
    'TRANSFER LEARNING': '%sTransfer learning%s' % (Fore.BLUE, Fore.RESET),
    'AUTOENCODER': '%sAutoencoder specific%s' % (Fore.BLUE, Fore.RESET),
    # Slice-level
    'SLICE': '%sSlice-level parameters%s' % (Fore.BLUE, Fore.RESET),
    'SLICE CNN': '%sSlice-level CNN parameters%s' % (Fore.BLUE, Fore.RESET),
    # Patch arguments
    'PATCH': '%sPatch-level parameters%s' % (Fore.BLUE, Fore.RESET),
    'PATCH CNN': '%sPatch-level CNN parameters%s' % (Fore.BLUE, Fore.RESET),
    # ROI-based arguments
    'ROI': '%sROI-based parameters%s' % (Fore.BLUE, Fore.RESET),
    'ROI CNN': '%sROI-based CNN parameters%s' % (Fore.BLUE, Fore.RESET),
    # Other optional arguments
    'OPTIONAL': '%sOther options%s' % (Fore.BLUE, Fore.RESET),
    # Model selection
    'MODEL': '%sModel selection%s' % (Fore.BLUE, Fore.RESET),
    # Display
    'DISPLAY': '%sResults display%s' % (Fore.BLUE, Fore.RESET),
}

def extract_tensors(args):
    import sys
    from clinica.utils.stream import FilterOut
    from clinica.pipelines.deeplearning_prepare_data.deeplearning_prepare_data_cli import DeepLearningPrepareDataCLI

    sys.stdout = FilterOut(sys.stdout)

    dl_prepare_data_cli = DeepLearningPrepareDataCLI()
    dl_prepare_data_cli.run_command(args)

def qc_func(args):
    from .quality_check.t1_linear.quality_check import quality_check as linear_qc
    from .quality_check.t1_volume.quality_check import quality_check as volume_qc

    if args.preprocessing == "t1-linear":
        linear_qc(
            args.caps_dir,
            args.output_path,
            tsv_path=args.subjects_sessions_tsv,
            threshold=args.threshold,
            batch_size=args.batch_size,
            num_workers=args.nproc,
            gpu=not args.use_cpu
        )
    elif args.preprocessing == "t1-volume":
        volume_qc(
            args.caps_dir,
            args.output_dir,
            args.group_label
        )

def generate_data_func(args):
    from .tools.data.generate_data import (generate_random_dataset,
                                           generate_trivial_dataset,
                                           generate_shepplogan_dataset)

    if args.mode == "random":
        generate_random_dataset(
            caps_dir=args.caps_dir,
            tsv_path=args.subjects_sessions_tsv,
            output_dir=args.output_dir,
            n_subjects=args.n_subjects,
            mean=args.mean,
            sigma=args.sigma,
            preprocessing=args.preprocessing)
    elif args.mode == "trivial":
        generate_trivial_dataset(
            caps_dir=args.caps_dir,
            tsv_path=args.subjects_sessions_tsv,
            output_dir=args.output_dir,
            n_subjects=args.n_subjects,
            preprocessing=args.preprocessing,
            mask_path=args.mask_path,
            atrophy_percent=args.atrophy_percent,
        )
    else:
        labels_distribution = {"AD": args.AD_subtypes_distribution, "CN": args.CN_subtypes_distribution}
        generate_shepplogan_dataset(
            output_dir=args.output_dir,
            img_size=args.image_size,
            labels_distribution=labels_distribution,
            samples=args.n_subjects,
            smoothing=args.smoothing
        )

def rs_func(args):
    from .train.random_search import launch_search

    launch_search(args)

# Function to dispatch training to corresponding function
def train_func(args):
    from .train import train_autoencoder, train_multi_cnn, train_single_cnn

    if args.network_type == "autoencoder":
        args.transfer_learning_selection = "best_loss"
        train_autoencoder(args)
    elif args.network_type == "cnn":
        train_single_cnn(args)
    elif args.network_type == "multicnn":
        train_multi_cnn(args)
    else:
        raise NotImplementedError('Framework %s not implemented in clinicadl' % args.network_type)

# Function to dispatch command line options from classify to corresponding
# function
def classify_func(args):
    from .classify.inference import classify

    classify(
        args.caps_directory,
        args.tsv_path,
        args.model_path,
        args.prefix_output,
        labels=not args.no_labels,
        gpu=not args.use_cpu,
        prepare_dl=args.use_extracted_features,
        selection_metrics=args.selection_metrics,
        diagnoses=args.diagnoses,
        verbose=args.verbose
    )

# Functions to dispatch command line options from tsvtool to corresponding
# function
def tsv_restrict_func(args):
    from .tools.tsv.restriction import aibl_restriction, oasis_restriction

    if args.dataset == "AIBL":
        aibl_restriction(args.merged_tsv, args.results_path)
    elif args.dataset == "OASIS":
        oasis_restriction(args.merged_tsv, args.results_path)

def tsv_getlabels_func(args):
    from .tools.tsv.data_formatting import get_labels

    get_labels(
        args.merged_tsv,
        args.missing_mods,
        args.results_path,
        diagnoses=args.diagnoses,
        modality=args.modality,
        restriction_path=args.restriction_path,
        time_horizon=args.time_horizon,
        variables_of_interest=args.variables_of_interest,
        remove_smc=not args.keep_smc,
        verbose=args.verbose
    )

def tsv_split_func(args):
    from .tools.tsv.data_split import split_diagnoses

    split_diagnoses(
        args.formatted_data_path,
        n_test=args.n_test,
        subset_name=args.subset_name,
        MCI_sub_categories=args.MCI_sub_categories,
        t_val_threshold=args.t_val_threshold,
        p_val_threshold=args.p_val_threshold,
        ignore_demographics=args.ignore_demographics,
        verbose=args.verbose
    )

def tsv_kfold_func(args):
    from .tools.tsv.kfold_split import split_diagnoses

    split_diagnoses(
        args.formatted_data_path,
        n_splits=args.n_splits,
        subset_name=args.subset_name,
        MCI_sub_categories=args.MCI_sub_categories,
        stratification=args.stratification,
        verbose=args.verbose
    )

def tsv_analysis_func(args):
    from .tools.tsv.demographics_analysis import demographics_analysis

    demographics_analysis(
        args.merged_tsv,
        args.formatted_data_path,
        args.results_path,
        diagnoses=args.diagnoses
    )

def interpret_func(args):
    from .interpret.group_backprop import group_backprop
    from .interpret.individual_backprop import individual_backprop

    if args.task == "group":
        group_backprop(args)
    elif args.task == "individual":
        individual_backprop(args)
    else:
        raise ValueError("Unknown task %s for interpretation" % args.task)

def parse_command_line():
    parser = argparse.ArgumentParser(
        prog='clinicadl',
        description='Deep learning software for neuroimaging datasets')

    parser.add_argument('-l', '--logname',
                        dest='logname',
                        default="clinicaDL.log",
                        metavar=('file.log'),
                        help='Define the log file name (default: clinicaDL.log)')
    parser.add_argument("-V", "--version",
                        dest='version',
                        action='store_true', default=False,
                        help="ClinicaDL's installed version")

    parent_parser = argparse.ArgumentParser(add_help=False)
    parent_parser.add_argument('--verbose', '-v', action='count', default=0)

    subparser = parser.add_subparsers(
        title='''Task to execute with clinicadl:''',
        description='''What kind of task do you want to use with clinicadl?
            (tsvtool, preprocessing, generate, train, classify).''',
        dest='task',
        help='''****** Tasks proposed by clinicadl ******''')

    # subparser.required = True

    # Generate synthetic data
    generate_parser = subparser.add_parser(
        'generate',
        help='Generate synthetic data for functional tests.'
    )

    generate_subparser = generate_parser.add_subparsers(
        title='''Type of synthetic data generated''',
        description='''What type of synthetic data do you want to generate?
                (random, trivial).''',
        dest='mode',
        help='''****** Synthetic datasets proposed by clinicadl ******''')

    # generate_subparser.required = True

    # Positional arguments
    generate_rs_parent_parser = argparse.ArgumentParser(add_help=False)
    generate_rs_parent_parser.add_argument(
        'caps_dir',
        help='Data using CAPS structure.',
        default=None
    )
    generate_rs_parent_parser.add_argument(
        'preprocessing',
        type=str,
        choices=['t1-linear', 't1-extensive'],
        help="Preprocessing used to generate synthetic data."
    )
    generate_rs_parent_parser.add_argument(
        'output_dir',
        help='Folder containing the synthetic dataset.',
    )
    generate_rs_parent_parser.add_argument(
        "--subjects_sessions_tsv", "-tsv",
        help='TSV file containing a list of subjects with their sessions.',
        type=str, default=None
    )
    generate_rs_parent_parser.add_argument(
        '--n_subjects',
        type=int,
        default=300,
        help="Number of subjects in each class of the synthetic dataset."
    )

    generate_random_parser = generate_subparser.add_parser(
        "random",
        parents=[
            generate_rs_parent_parser],
        help="Generate a random dataset in which gaussian noise is added to brain images.")
    generate_random_parser.add_argument(
        '--mean',
        type=float,
        default=0,
        help="Mean value of the noise added for the random dataset."
    )
    generate_random_parser.add_argument(
        '--sigma',
        type=float,
        default=0.5,
        help="Standard deviation of the noise added for the random dataset."
    )
    generate_random_parser.set_defaults(func=generate_data_func)

    generate_trivial_parser = generate_subparser.add_parser(
        "trivial",
        parents=[
            generate_rs_parent_parser],
        help="Generate a trivial dataset in which gaussian half of the brain is atrophied.")
    generate_trivial_parser.add_argument(
        '--mask_path',
        type=str,
        help='path to the extracted masks to generate the two labels.',
        default=None
    )
    generate_trivial_parser.add_argument(
        '--atrophy_percent',
        type=float,
        default=60,
        help='percentage of atrophy applied'
    )
    generate_trivial_parser.set_defaults(func=generate_data_func)

    # generate_shepplogan_parser = generate_subparser.add_parser(
    #     "shepplogan",
    #     help="Generate a dataset of 2D images including 3 subtypes based on Shepp Logan phantom."
    # )
    # generate_shepplogan_parser.add_argument(
    #     'output_dir',
    #     help='Folder containing the synthetic dataset.',
    # )
    # generate_shepplogan_parser.add_argument(
    #     '--n_subjects',
    #     type=int,
    #     default=300,
    #     help="Number of subjects in each class of the synthetic dataset."
    # )
    # generate_shepplogan_parser.add_argument(
    #     '--image_size',
    #     type=int,
    #     default=128,
    #     help="Size in pixels of the squared images."
    # )
    # generate_shepplogan_parser.add_argument(
    #     '--CN_subtypes_distribution', '-Csd',
    #     type=float, nargs='+',
    #     default=[1.0, 0.0, 0.0],
    #     help="Probability of each subtype to be drawn in CN label."
    # )
    # generate_shepplogan_parser.add_argument(
    #     '--AD_subtypes_distribution', '-Asd',
    #     type=float, nargs='+',
    #     default=[0.05, 0.85, 0.10],
    #     help="Probability of each subtype to be drawn in AD label."
    # )
    # generate_shepplogan_parser.add_argument(
    #     '--smoothing',
    #     action='store_true',
    #     default=False,
    #     help='Adds random smoothing to generated data.'
    # )
    # generate_shepplogan_parser.set_defaults(func=generate_data_func)

    # Preprocessing
    from clinica.pipelines.t1_linear.t1_linear_cli import T1LinearCLI
    from clinica.engine.cmdparser import init_cmdparser_objects
    from .preprocessing.t1_extensive.t1_extensive_cli import T1ExtensiveCli
    preprocessing_parser = subparser.add_parser(
        'preprocessing',
        help='Preprocess T1w-weighted images with t1-linear or t1-extensive pipelines.'
    )

    preprocessing_subparser = preprocessing_parser.add_subparsers(
        title='''Preprocessing task to execute with clinicadl:''',
        description='''What kind of task do you want to perform with clinicadl?
                (run, quality-check, extract-tensor).''',
        dest='preprocessing_task',
        help='''****** Tasks proposed by clinicadl ******''')
    preprocessing_subparser.required = True

    run_parser = preprocessing_subparser.add_parser(
        'run',
        help='Preprocess T1w-weighted images with t1-linear or t1-extensive pipelines.'
    )
    run_parser._positionals.title = ('%sclinicadl preprocessing expects one of the following pipelines%s'
                                     % (Fore.GREEN, Fore.RESET))

    def preprocessing_help(args):
        print('%sNo pipeline was specified. Type clinica preprocessing -h for details%s' %
              (Fore.RED, Fore.RESET))

    run_parser.set_defaults(func=preprocessing_help)

    init_cmdparser_objects(
        parser,
        run_parser.add_subparsers(dest='preprocessing'),
        [
            T1LinearCLI(),
            T1ExtensiveCli(),
        ]
    )

    extract_parser = preprocessing_subparser.add_parser(
        'extract-tensor',
        help='Prepare data generated Clinica for PyTorch with Tensor extraction (image, patches or slices).'
    )
    clinica_comp = extract_parser.add_argument_group('%sClinica mandatory arguments%s' % (Fore.BLUE, Fore.RESET))
    clinica_comp.add_argument(
        "caps_directory",
        help='Path to the CAPS directory.'
    )
    clinica_comp.add_argument("modality",
                              help='''For which modality the tensor will be extracted.
            't1-linear': images prepocessed with t1-linear pipeline.
            't1-extensive': images preprocessed with t1-extensive pipeline.
            'custom': find images with a custom suffix in their filename and
            transform them to tensor format.''',
                              choices=['t1-linear', 't1-extensive', 'custom'], default='t1-linear'
                              )

    clinica_comp.add_argument(
        "extract_method",
        help='''Format of the extracted features. Three options:
        'image' to convert to PyTorch tensor the complete 3D image,
        'patch' to extract 3D volumetric patches and
        'slice' to extract 2D slices from the image.
        By default the features are extracted from the cropped image.''',
        choices=['image', 'slice', 'patch'],
        default='image'
    )

    optional = extract_parser.add_argument_group('%sPipeline options%s' % (Fore.BLUE, Fore.RESET))
    optional.add_argument(
        '-uui', '--use_uncropped_image',
        help='''Use the uncropped image instead of the
        cropped image generated by t1-linear.''',
        default=False, action="store_true"
    )

    optional_patch = extract_parser.add_argument_group(
        "%sPipeline options if you chose ‘patch’ extraction%s" % (Fore.BLUE, Fore.RESET)
    )
    optional_patch.add_argument(
        '-ps', '--patch_size',
        help='''Patch size (default: --patch_size 50).''',
        type=int, default=50
    )
    optional_patch.add_argument(
        '-ss', '--stride_size',
        help='''Stride size (default: --stride_size 50).''',
        type=int, default=50
    )
    optional_slice = extract_parser.add_argument_group(
        "%sPipeline options if you chose ‘slice’ extraction%s" % (Fore.BLUE, Fore.RESET)
    )
    optional_slice.add_argument(
        '-sd', '--slice_direction',
        help='''Slice direction. Three options:
            '0' -> Sagittal plane,
            '1' -> Coronal plane or
            '2' -> Axial plane
            (default: sagittal plane i.e. --slice_direction 0)''',
        type=int, default=0
    )
    optional_slice.add_argument(
        '-sm', '--slice_mode',
        help='''Slice mode. Two options: 'rgb' to save the slice in
            three identical channels, ‘single’ to save the slice in a
            single channel (default: --slice_mode rgb).''',
        choices=['rgb', 'single'], default='rgb'
    )
    optional_custom = extract_parser.add_argument_group(
        "%sPipeline options if you chose ‘custom’ modality%s" % (Fore.BLUE, Fore.RESET)
    )
    optional_custom.add_argument(
        '-cn', '--custom_suffix',
        help='''Custom suffix filename, e.g.:
            'graymatter_space-Ixi549Space_modulated-off_probability.nii.gz', or
            'segm-whitematter_probability.nii.gz'
            ''',
        type=str, default=''
    )

    # Clinica standard arguments (e.g. --n_procs)
    clinica_standard_options = extract_parser.add_argument_group('%sClinica standard options%s' % (Fore.BLUE, Fore.RESET))
    clinica_standard_options.add_argument(
        "-tsv", "--subjects_sessions_tsv",
        help='TSV file containing a list of subjects with their sessions.'
    )
    clinica_standard_options.add_argument(
        "-wd", "--working_directory",
        help='Temporary directory to store pipelines intermediate results.'
    )
    clinica_standard_options.add_argument(
        "-np", "--n_procs",
        metavar='N', type=int,
        help='Number of cores used to run in parallel.'
    )

    extract_parser.set_defaults(func=extract_tensors)

    qc_parser = preprocessing_subparser.add_parser(
        'quality-check',
        help='Performs quality check procedure for t1-linear pipeline.'
             'Original code can be found at https://github.com/vfonov/deep-qc'
    )

    qc_subparsers = qc_parser.add_subparsers(
        title='''Preprocessing pipelines available''',
        description='''Which preprocessing pipeline do you want to check?''',
        dest='preprocessing',
        help='''****** Preprocessing pipelines ******''')

    qc_subparsers.required = True
    qc_linear_parser = qc_subparsers.add_parser(
        't1-linear',
        help='Performs quality check on t1-linear pipeline.'
    )
    qc_linear_parser.add_argument("caps_dir",
                                  help='Data using CAPS structure.',
                                  type=str)
    qc_linear_parser.add_argument("output_path",
                                  help="Path to the output tsv file (filename included).",
                                  type=str)

    qc_linear_parser.add_argument("--subjects_sessions_tsv", "-tsv",
                                  help='TSV file containing a list of subjects with their sessions.',
                                  type=str, default=None)
    qc_linear_parser.add_argument("--threshold",
                                  help='The threshold on the output probability to decide if the image '
                                       'passed or failed. (default=0.5)',
                                  type=float, default=0.5)
    qc_linear_parser.add_argument('--batch_size',
                                  help='Batch size used in DataLoader (default=1).',
                                  default=1, type=int)
    qc_linear_parser.add_argument("-np", "--nproc",
                                  help='Number of cores used the quality check. (default=2)',
                                  type=int, default=2)
    qc_linear_parser.add_argument('-cpu', '--use_cpu', action='store_true',
                                  help='If provided, will use CPU instead of GPU.',
                                  default=False)
    qc_linear_parser.set_defaults(func=qc_func)

    qc_volume_parser = qc_subparsers.add_parser(
        't1-volume',
        help='Performs quality check on t1-volume pipeline.'
    )
    qc_volume_parser.add_argument("caps_dir",
                                  help='Data using CAPS structure.',
                                  type=str)
    qc_volume_parser.add_argument("output_dir",
                                  help="Path to the output directory containing TSV files.",
                                  type=str)
    qc_volume_parser.add_argument("group_label",
                                  help="Identifier for the group of subjects used to create the DARTEL template.",
                                  type=str)
    qc_volume_parser.set_defaults(func=qc_func)

    # random search parsers
    rs_generate_parser = subparser.add_parser(
        'random-search',
        parents=[parent_parser],
        help='Generate random networks to explore hyper parameters space.',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    # rs_subparsers = rs_parser.add_subparsers(
    #     title='''Possibilities for random network training''',
    #     description='''You can generate and train a new random network,
    #     or relaunch a previous random job with some alterations.''',
    #     dest='random_task',
    #     help='''****** Possible tasks ******'''
    # )
    # rs_subparsers.required = True

    # rs_generate_parser = rs_subparsers.add_parser(
    #     'generate',
    #     parents=[parent_parser],
    #     help='Sample a new network and train it.',
    #     formatter_class=argparse.ArgumentDefaultsHelpFormatter
    # )
    rs_pos_group = rs_generate_parser.add_argument_group(
        TRAIN_CATEGORIES["POSITIONAL"]
    )
    rs_pos_group.add_argument("launch_dir", type=str,
                              help="Directory containing the random_search.json file.")
    rs_pos_group.add_argument("name", type=str,
                              help="Name of the job.")

    # Data management
    rs_data_group = rs_generate_parser.add_argument_group(
        TRAIN_CATEGORIES["CROSS-VALIDATION"]
    )
    rs_data_group.add_argument("--n_splits", type=int, default=None,
                               help="If a value is given will load data of a k-fold CV")
    rs_data_group.add_argument("--split", type=int, default=None, nargs="+",
                               help="Will load the specific split wanted.")

    rs_comp_group = rs_generate_parser.add_argument_group(
        TRAIN_CATEGORIES["COMPUTATIONAL"]
    )
    rs_comp_group.add_argument(
        '-cpu', '--use_cpu', action='store_true',
        help='If provided, will use CPU instead of GPU.',
        default=False)
    rs_comp_group.add_argument(
        '-np', '--nproc',
        help='Number of cores used during the training.',
        type=int, default=2)
    rs_comp_group.add_argument(
        '--batch_size',
        default=2, type=int,
        help='Batch size for training.')
    rs_comp_group.add_argument(
        '--evaluation_steps', '-esteps',
        default=0, type=int,
        help='Fix the number of iterations to perform before computing an evaluations. Default will only '
             'perform one evaluation at the end of each epoch.')

    rs_generate_parser.set_defaults(func=rs_func)

    # retrain_parent_parser = return_train_parent_parser(retrain=True)
    # rs_retrain_parser = rs_subparsers.add_parser(
    #     'retrain',
    #     parents=[parent_parser, retrain_parent_parser],
    #     help='Train a network previously created by generate.',
    #     formatter_class=argparse.ArgumentDefaultsHelpFormatter
    # )

    train_parser = subparser.add_parser(
        'train',
        help='Train with your data and create a model.')

    train_subparser = train_parser.add_subparsers(
        title='''Inputs types implemented in clinicadl''',
        description='''What type of input do you want to use?
                (image, patch, roi, slice).''',
        dest='mode',
        help='''****** Input types proposed by clinicadl ******''')

    train_subparser.required = True

    # Transfer learning
    transfer_learning_parent = argparse.ArgumentParser(add_help=False)
    transfer_learning_group = transfer_learning_parent.add_argument_group(
        TRAIN_CATEGORIES["TRANSFER LEARNING"])
    transfer_learning_group.add_argument(
        '--transfer_learning_path',
        help="If an existing path is given, a pretrained model is used.",
        type=str, default=None)

    # Autoencoder
    autoencoder_parent = argparse.ArgumentParser(add_help=False)
    autoencoder_group = autoencoder_parent.add_argument_group(
        TRAIN_CATEGORIES["AUTOENCODER"])
    autoencoder_group.add_argument(
        '--visualization',
        help='Save examples of image reconstructions.',
        action="store_true",
        default=False)

    ######################
    # IMAGE
    ######################
    train_image_parser = train_subparser.add_parser(
        "image",
        help="Train a 3D image-level network.")

    train_image_subparser = train_image_parser.add_subparsers(
        title='''Task to be performed''',
        description='''Autoencoder reconstruction or cnn classification ?''',
        dest='network_type',
        help='''****** Choose a type of network ******''')

    train_parent_parser = return_train_parent_parser(retrain=False)
    train_image_ae_parser = train_image_subparser.add_parser(
        "autoencoder",
        parents=[
            parent_parser,
            train_parent_parser,
            autoencoder_parent,
            transfer_learning_parent],
        help="Train an image-level autoencoder.")

    train_image_ae_parser.set_defaults(func=train_func)

    train_image_cnn_parser = train_image_subparser.add_parser(
        "cnn",
        parents=[
            parent_parser,
            train_parent_parser,
            transfer_learning_parent],
        help="Train an image-level CNN.")
    # /!\ If parents list is changed the arguments won't be in the right group anymore !
    train_image_cnn_parser._action_groups[-1].add_argument(
        '--transfer_learning_selection',
        help="If transfer_learning from CNN, chooses which best transfer model is selected.",
        type=str, default="best_balanced_accuracy", choices=["best_loss", "best_balanced_accuracy"])

    train_image_cnn_parser.set_defaults(func=train_func)

    #########################
    # PATCH
    #########################
    train_patch_parser = train_subparser.add_parser(
        "patch",
        help="Train a 3D patch-level network.")

    train_patch_parent = argparse.ArgumentParser(add_help=False)
    train_patch_group = train_patch_parent.add_argument_group(
        TRAIN_CATEGORIES["PATCH"])
    train_patch_group.add_argument(
        '-ps', '--patch_size',
        help='Patch size',
        type=int, default=50)
    train_patch_group.add_argument(
        '-ss', '--stride_size',
        help='Stride size',
        type=int, default=50)
    train_patch_group.add_argument(
        '--use_extracted_patches',
        help='''If provided the outputs of extract preprocessing are used, else the whole
             MRI is loaded.''',
        default=False, action="store_true")

    train_patch_subparser = train_patch_parser.add_subparsers(
        title='''Task to be performed''',
        description='''Autoencoder reconstruction or (multi)cnn classification ?''',
        dest='network_type',
        help='''****** Choose a type of network ******''')
    train_patch_subparser.required = True

    train_patch_ae_parser = train_patch_subparser.add_parser(
        "autoencoder",
        parents=[parent_parser, train_parent_parser, train_patch_parent, autoencoder_parent, transfer_learning_parent],
        help="Train a 3D patch-level autoencoder.")

    train_patch_ae_parser.set_defaults(func=train_func)

    train_patch_cnn_parser = train_patch_subparser.add_parser(
        "cnn",
        parents=[
            parent_parser,
            train_parent_parser,
            train_patch_parent,
            transfer_learning_parent],
        help="Train a 3D patch-level CNN.")
    # /!\ If parents list is changed the arguments won't be in the right group anymore !
    train_patch_cnn_parser._action_groups[-1].add_argument(
        '--transfer_learning_selection',
        help="If transfer_learning from CNN, chooses which best transfer model is selected.",
        type=str, default="best_balanced_accuracy", choices=["best_loss", "best_balanced_accuracy"])

    train_patch_cnn_group = train_patch_cnn_parser.add_argument_group(
        TRAIN_CATEGORIES["PATCH CNN"])
    train_patch_cnn_group.add_argument(
        '--selection_threshold',
        help='''Threshold on the balanced accuracies to compute the
             subject-level performance. Patches are selected if their balanced
             accuracy > threshold. Default corresponds to no selection.''',
        type=float, default=0.0)

    train_patch_cnn_parser.set_defaults(func=train_func)

    train_patch_multicnn_parser = train_patch_subparser.add_parser(
        "multicnn",
        parents=[
            parent_parser,
            train_parent_parser,
            train_patch_parent,
            transfer_learning_parent],
        help="Train a 3D patch-level multi-CNN (one CNN is trained per patch location).")
    # /!\ If parents list is changed the arguments won't be in the right group anymore !
    train_patch_multicnn_parser._action_groups[-1].add_argument(
        '--transfer_learning_selection',
        help="If transfer_learning from CNN, chooses which best transfer model is selected.",
        type=str, default="best_balanced_accuracy", choices=["best_loss", "best_balanced_accuracy"])

    train_patch_multicnn_group = train_patch_multicnn_parser.add_argument_group(
        TRAIN_CATEGORIES["PATCH CNN"])
    train_patch_multicnn_group.add_argument(
        '--selection_threshold',
        help='''Threshold on the balanced accuracies to compute the
                 subject-level performance. Patches are selected if their balanced
                 accuracy > threshold. Default corresponds to no selection.''',
        type=float, default=0.0)

    train_patch_multicnn_parser.set_defaults(func=train_func)

    #########################
    # ROI
    #########################
    train_roi_parser = train_subparser.add_parser(
        "roi",
        help="Train a ROI-based level network.")

    train_roi_subparser = train_roi_parser.add_subparsers(
        title='''Task to be performed''',
        description='''Autoencoder reconstruction or cnn classification ?''',
        dest='network_type',
        help='''****** Choose a type of network ******''')
    train_roi_subparser.required = True

    train_roi_ae_parser = train_roi_subparser.add_parser(
        "autoencoder",
        parents=[
            parent_parser,
            train_parent_parser,
            autoencoder_parent,
            transfer_learning_parent
        ],
        help="Train a ROI-based autoencoder.")

    train_roi_ae_parser.set_defaults(func=train_func)

    train_roi_cnn_parser = train_roi_subparser.add_parser(
        "cnn",
        parents=[
            parent_parser,
            train_parent_parser,
            transfer_learning_parent],
        help="Train a ROI-based CNN.")
    # /!\ If parents list is changed the arguments won't be in the right group anymore !
    train_roi_cnn_parser._action_groups[-1].add_argument(
        '--transfer_learning_selection',
        help="If transfer_learning from CNN, chooses which best transfer model is selected.",
        type=str, default="best_balanced_accuracy", choices=["best_loss", "best_balanced_accuracy"])

    train_roi_cnn_group = train_roi_cnn_parser.add_argument_group(
        TRAIN_CATEGORIES["ROI CNN"])
    train_roi_cnn_group.add_argument(
        '--selection_threshold',
        help='''Threshold on the balanced accuracies to compute the
             subject-level performance. ROIs are selected if their balanced
             accuracy > threshold. Default corresponds to no selection.''',
        type=float, default=0.0)

    train_roi_cnn_parser.set_defaults(func=train_func)

    train_roi_multicnn_parser = train_roi_subparser.add_parser(
        "multicnn",
        parents=[
            parent_parser,
            train_parent_parser,
            transfer_learning_parent],
        help="Train a ROI-based multi-CNN (one CNN is trained per patch location).")
    # /!\ If parents list is changed the arguments won't be in the right group anymore !
    train_roi_multicnn_parser._action_groups[-1].add_argument(
        '--transfer_learning_selection',
        help="If transfer_learning from CNN, chooses which best transfer model is selected.",
        type=str, default="best_balanced_accuracy", choices=["best_loss", "best_balanced_accuracy"])

    train_roi_multicnn_group = train_roi_multicnn_parser.add_argument_group(
        TRAIN_CATEGORIES["ROI CNN"])
    train_roi_multicnn_group.add_argument(
        '--selection_threshold',
        help='''Threshold on the balanced accuracies to compute the
                         subject-level performance. Patches are selected if their balanced
                         accuracy > threshold. Default corresponds to no selection.''',
        type=float, default=0.0)

    train_roi_multicnn_parser.set_defaults(func=train_func)

    #########################
    # SLICE
    #########################
    train_slice_parser = train_subparser.add_parser(
        "slice",
        help="Train a 2D slice-level CNN.")

    train_slice_subparser = train_slice_parser.add_subparsers(
        title='''Task to be performed''',
        description='''Autoencoder reconstruction or cnn classification ?''',
        dest='network_type',
        help='''****** Choose a type of network ******''')
    train_slice_subparser.required = True

    train_slice_parent = argparse.ArgumentParser(add_help=False)
    train_slice_group = train_slice_parent.add_argument_group(
        TRAIN_CATEGORIES["SLICE"])
    train_slice_group.add_argument(
        '--slice_direction', '-sd',
        help='''Which coordinate axis to take for slicing the MRI.
                 0 for sagittal
                 1 for coronal
                 2 for axial direction.''',
        default=0, type=int)
    train_slice_group.add_argument(
        '--discarded_slices',
        help='''Number of slices discarded from respectively the beginning and
        the end of the MRI volume.  If only one argument is given, it will be
        used for both sides.''',
        default=20, type=int, nargs='+'
    )
    train_slice_group.add_argument(
        '--use_extracted_slices',
        help='''If provided the outputs of extract preprocessing are used, else the whole
                 MRI is loaded.''',
        default=False, action="store_true")

    train_slice_ae_parser = train_slice_subparser.add_parser(
        "autoencoder",
        parents=[parent_parser, train_parent_parser, train_slice_parent, transfer_learning_parent],
        help="Train a 2D slice-level autoencoder.")

    train_slice_ae_parser.set_defaults(func=train_func)

    train_slice_cnn_parser = train_slice_subparser.add_parser(
        "cnn",
        parents=[parent_parser, train_parent_parser, train_slice_parent, transfer_learning_parent],
        help="Train a 2D slice-level CNN.")
    # /!\ If parents list is changed the arguments won't be in the right group anymore !
    train_slice_cnn_parser._action_groups[-1].add_argument(
        '--transfer_learning_selection',
        help="If transfer_learning from CNN, chooses which best transfer model is selected.",
        type=str, default="best_balanced_accuracy", choices=["best_loss", "best_balanced_accuracy"])

    train_slice_cnn_group = train_slice_cnn_parser.add_argument_group(
        TRAIN_CATEGORIES["SLICE CNN"])
    train_slice_cnn_group.add_argument(
        '--selection_threshold',
        help='''Threshold on the balanced accuracies to compute the
                 subject-level performance. Slices are selected if their balanced
                 accuracy > threshold. Default corresponds to no selection.''',
        type=float, default=0.0)

    train_slice_cnn_parser.set_defaults(func=train_func)

    train_slice_multicnn_parser = train_slice_subparser.add_parser(
        "multicnn",
        parents=[parent_parser, train_parent_parser, train_slice_parent, transfer_learning_parent],
        help="Train a 2D slice-level multi-CNN.")
    # /!\ If parents list is changed the arguments won't be in the right group anymore !
    train_slice_multicnn_parser._action_groups[-1].add_argument(
        '--transfer_learning_selection',
        help="If transfer_learning from CNN, chooses which best transfer model is selected.",
        type=str, default="best_balanced_accuracy", choices=["best_loss", "best_balanced_accuracy"])

    train_slice_multicnn_group = train_slice_multicnn_parser.add_argument_group(
        TRAIN_CATEGORIES["SLICE CNN"])
    train_slice_multicnn_group.add_argument(
        '--selection_threshold',
        help='''Threshold on the balanced accuracies to compute the
                 subject-level performance. Slices are selected if their balanced
                 accuracy > threshold. Default corresponds to no selection.''',
        type=float, default=0.0)

    train_slice_multicnn_parser.set_defaults(func=train_func)

    # Classify - Classify a subject or a list of tsv files with the CNN
    # provided as argument.
    # classify_parser: get command line arguments and options

    classify_parser = subparser.add_parser(
        'classify',
        parents=[parent_parser],
        help='''Classify one image or a list of images with your previously
                 trained model.''')
    classify_pos_group = classify_parser.add_argument_group(
        TRAIN_CATEGORIES["POSITIONAL"])
    classify_pos_group.add_argument(
        'caps_directory',
        help='Data using CAPS structure.',
        default=None)
    classify_pos_group.add_argument(
        'tsv_path',
        help='''Path to the file with subjects/sessions to process.
        If it includes the filename will load the tsv file directly.
        Else will load the baseline tsv files of wanted diagnoses produced by tsvtool.''',
        default=None)
    classify_pos_group.add_argument(
        'model_path',
        help='''Path to the folder where the model is stored. Folder structure
                should be the same obtained during the training.''',
        default=None)
    classify_pos_group.add_argument(
        'prefix_output',
        help='Prefix to name the files resulting from the classify task.',
        type=str)

    # Computational resources
    classify_comput_group = classify_parser.add_argument_group(
        TRAIN_CATEGORIES["COMPUTATIONAL"])
    classify_comput_group.add_argument(
        '-cpu', '--use_cpu', action='store_true',
        help='Uses CPU instead of GPU.',
        default=False)
    classify_comput_group.add_argument(
        '-np', '--nproc',
        help='Number of cores used during the task.',
        type=int, default=2)
    classify_comput_group.add_argument(
        '--batch_size',
        default=2, type=int,
        help='Batch size for data loading. (default=2)')

    # Specific classification arguments
    classify_specific_group = classify_parser.add_argument_group(
        TRAIN_CATEGORIES["OPTIONAL"]
    )
    classify_specific_group.add_argument(
        '-nl', '--no_labels', action='store_true',
        help='Add this flag if your dataset does not contain a ground truth.',
        default=False)
    classify_specific_group.add_argument(
        '--use_extracted_features',
        help='''If True the extract slices or patche are used, otherwise the they
                will be extracted on the fly (if necessary).''',
        default=False, action="store_true")
    classify_specific_group.add_argument(
        '--selection_metrics',
        help='''List of metrics to find the best models to evaluate. Default will
        classify best model based on balanced accuracy.''',
        choices=['loss', 'balanced_accuracy'],
        default=['balanced_accuracy'],
        nargs='+'
    )
    classify_specific_group.add_argument(
        "--diagnoses",
        help="List of participants that will be classified.",
        nargs="+", type=str, choices=['AD', 'CN', 'MCI', 'sMCI', 'pMCI'], default=None)

    classify_parser.set_defaults(func=classify_func)

    tsv_parser = subparser.add_parser(
        'tsvtool',
        help='''Handle tsv files for metadata processing and data splits.''')

    tsv_subparser = tsv_parser.add_subparsers(
        title='''Task to execute with tsv tool:''',
        description='''What kind of task do you want to use with tsv tool?
                (restrict, getlabels, split, kfold, analysis).''',
        dest='tsv_task',
        help='''****** Tasks proposed by clinicadl tsv tool ******''')

    tsv_subparser.required = True

    tsv_restrict_subparser = tsv_subparser.add_parser(
        'restrict',
        help='Reproduce restrictions applied to AIBL and OASIS datasets')

    tsv_restrict_subparser.add_argument(
        "dataset",
        help="dataset on which the restriction is performed.",
        choices=["AIBL", "OASIS"],
        type=str)

    tsv_restrict_subparser.add_argument(
        "merged_tsv",
        help="Path to the file obtained by the command clinica iotools merge-tsv.",
        type=str)
    tsv_restrict_subparser.add_argument(
        "results_path",
        help="Path to the output tsv file (filename included).",
        type=str)

    tsv_restrict_subparser.set_defaults(func=tsv_restrict_func)

    tsv_getlabels_subparser = tsv_subparser.add_parser(
        'getlabels',
        parents=[parent_parser],
        help='Get labels in separate tsv files.')

    tsv_getlabels_subparser.add_argument(
        "merged_tsv",
        help="Path to the file obtained by the command clinica iotools merge-tsv.",
        type=str)
    tsv_getlabels_subparser.add_argument(
        "missing_mods",
        help="Path to the folder where the outputs of clinica iotools missing-mods are.",
        type=str)
    tsv_getlabels_subparser.add_argument(
        "results_path",
        type=str,
        help="Path to the folder where tsv files are extracted.")

    # Optional arguments
    tsv_getlabels_subparser.add_argument(
        "--modality", "-mod",
        help="Modality to select sessions. Sessions which do not include the modality will be excluded.",
        default="t1w", type=str)
    tsv_getlabels_subparser.add_argument(
        "--diagnoses",
        help="Labels that must be extracted from merged_tsv.",
        nargs="+", type=str, choices=['AD', 'CN', 'MCI', 'sMCI', 'pMCI'], default=['AD', 'CN'])
    tsv_getlabels_subparser.add_argument(
        "--time_horizon",
        help="Time horizon to analyse stability of MCI subjects.",
        default=36, type=int)
    tsv_getlabels_subparser.add_argument(
        "--restriction_path",
        help="Path to a tsv containing the sessions that can be included.",
        type=str, default=None)
    tsv_getlabels_subparser.add_argument(
        "--variables_of_interest",
        help="Variables of interest that will be kept in the final lists."
             "Default will keep the diagnosis, age and the sex needed for the split procedure.",
        type=str, nargs="+", default=None)
    tsv_getlabels_subparser.add_argument(
        "--keep_smc",
        help="This flag allows to keep SMC participants, else they are removed.",
        default=False, action="store_true"
    )

    tsv_getlabels_subparser.set_defaults(func=tsv_getlabels_func)

    tsv_split_subparser = tsv_subparser.add_parser(
        'split',
        parents=[parent_parser],
        help='Performs one stratified shuffle split on participant level.')
    tsv_split_subparser.add_argument(
        "formatted_data_path",
        help="Path to the folder containing data extracted by clinicadl tsvtool getlabels.",
        type=str)

    # Optional arguments
    tsv_split_subparser.add_argument(
        "--n_test",
        help="If >= 1, number of subjects to put in set with name 'subset_name'. "
             "If < 1, proportion of subjects to put set with name 'subset_name'. "
             "If 0, no training set is created and the whole dataset is considered as one set with name 'subset_name.",
        type=float, default=100.)
    tsv_split_subparser.add_argument(
        "--MCI_sub_categories",
        help="Deactivate default managing of MCI sub-categories to avoid data leakage.",
        action="store_false", default=True)
    tsv_split_subparser.add_argument(
        "--t_val_threshold", "-t",
        help="The threshold used for the chi2 test on sex distributions.",
        default=0.0642, type=float)
    tsv_split_subparser.add_argument(
        "--p_val_threshold", "-p",
        help="The threshold used for the T-test on age distributions.",
        default=0.80, type=float)
    tsv_split_subparser.add_argument(
        "--subset_name",
        help="Name of the subset that is complementary to train.",
        type=str, default="test")
    tsv_split_subparser.add_argument(
        "--ignore_demographics",
        help="If True do not use age and sex to create the splits.",
        default=False, action="store_true"
    )

    tsv_split_subparser.set_defaults(func=tsv_split_func)

    tsv_kfold_subparser = tsv_subparser.add_parser(
        'kfold',
        parents=[parent_parser],
        help='Performs a k-fold split on participant level.')

    tsv_kfold_subparser.add_argument(
        "formatted_data_path",
        help="Path to the folder containing data extracted by clinicadl tsvtool getlabels.",
        type=str)

    # Optional arguments
    tsv_kfold_subparser.add_argument(
        "--n_splits",
        help="Number of folds in the k-fold split."
             "If 0, there is no training set and the whole dataset is considered as a test set.",
        type=int, default=5)
    tsv_kfold_subparser.add_argument(
        "--MCI_sub_categories",
        help="Deactivate default managing of MCI sub-categories to avoid data leakage.",
        action="store_false", default=True)
    tsv_kfold_subparser.add_argument(
        "--subset_name",
        help="Name of the subset that is complementary to train.",
        type=str, default="validation")
    tsv_kfold_subparser.add_argument(
        "--stratification",
        help="Name of a variable used to stratify the k-fold split.",
        type=str, default=None)

    tsv_kfold_subparser.set_defaults(func=tsv_kfold_func)

    tsv_analysis_subparser = tsv_subparser.add_parser(
        'analysis',
        help='Produces a demographic analysis of the extracted labels.')

    tsv_analysis_subparser.add_argument(
        "merged_tsv",
        help="Path to the file obtained by the command clinica iotools merge-tsv.",
        type=str)
    tsv_analysis_subparser.add_argument(
        "formatted_data_path",
        help="Path to the folder containing data extracted by clinicadl tsvtool getlabels.",
        type=str)
    tsv_analysis_subparser.add_argument(
        "results_path",
        help="Path to the output tsv file (filename included).",
        type=str)

    # Modality selection
    tsv_analysis_subparser.add_argument(
        "--diagnoses",
        help="Labels selected for the demographic analysis.",
        default=['AD', 'CN'], nargs="+", type=str, choices=['AD', 'CN', 'MCI', 'sMCI', 'pMCI'])

    tsv_analysis_subparser.set_defaults(func=tsv_analysis_func)

    interpret_parser = subparser.add_parser(
        'interpret',
        help='''Interpret classification performed by a CNN with saliency maps.''')

    interpret_subparser = interpret_parser.add_subparsers(
        title='''Type of saliency map to perform:''',
        description='''Do you want to perform a group saliency map or individual ones?''',
        dest='task',
        help='''****** Saliency maps proposed by clinicadl ******'''
    )
    interpret_subparser.required = True

    interpret_parent_parser = argparse.ArgumentParser(add_help=False)

    interpret_pos_group = interpret_parent_parser.add_argument_group(TRAIN_CATEGORIES["POSITIONAL"])
    interpret_pos_group.add_argument(
        "model_path", type=str,
        help="Path to the model output directory.")
    interpret_pos_group.add_argument(
        "name", type=str,
        help="Name of the interpretation map.")

    interpret_comput_group = interpret_parent_parser.add_argument_group(TRAIN_CATEGORIES["COMPUTATIONAL"])
    interpret_comput_group.add_argument(
        "--batch_size", default=1, type=int,
        help="Batch size for selection of images (keep_true).")
    interpret_comput_group.add_argument(
        '-cpu', '--use_cpu',
        action='store_true', default=False,
        help='Uses gpu instead of cpu if cuda is available.')
    interpret_comput_group.add_argument(
        '-np', '--nproc',
        default=2, type=int,
        help='the number of batches being loaded in parallel.')

    interpret_model_group = interpret_parent_parser.add_argument_group(TRAIN_CATEGORIES["MODEL"])
    interpret_model_group.add_argument(
        "--selection", default=['best_loss'], type=str, nargs="+",
        choices=['best_loss', 'best_balanced_accuracy'],
        help="Loads the model selected on minimal loss or maximum accuracy on validation.")

    interpret_data_group = interpret_parent_parser.add_argument_group(TRAIN_CATEGORIES["DATA"])
    interpret_data_group.add_argument(
        "--tsv_path", type=str, default=None,
        help="TSV path with subjects/sessions to process, if different from classification task.")
    interpret_data_group.add_argument(
        "--caps_dir", type=str, default=None,
        help="Path to input dir of the MRI (preprocessed CAPS_dir), if different from classification task")
    interpret_data_group.add_argument(
        "--diagnosis", "-d", default='AD', type=str,
        help="The images corresponding to this diagnosis only will be loaded.")
    interpret_data_group.add_argument(
        "--target_diagnosis", default=None, type=str,
        help="Which class the gradients explain. If None is given will be equal to diagnosis.")
    interpret_data_group.add_argument(
        "--baseline", action="store_true", default=False,
        help="If provided, only the baseline sessions are used for training.")
    interpret_data_group.add_argument(
        "--keep_true", type=lambda x: bool(strtobool(x)), default=None,
        help="Chooses false or true positive values of the classification. No selection by default")
    interpret_data_group.add_argument(
        "--nifti_template_path", type=str, default=None,
        help="Path to a nifti template to retrieve affine values.")

    interpret_display_group = interpret_parent_parser.add_argument_group(TRAIN_CATEGORIES["DISPLAY"])
    interpret_display_group.add_argument(
        "--vmax", type=float, default=0.5,
        help="Maximum value used in 2D image display.")

    interpret_group_parser = interpret_subparser.add_parser(
        "group",
        parents=[parent_parser, interpret_parent_parser],
        help="Mean saliency map over a list of sessions"
    )

    interpret_group_parser.set_defaults(func=interpret_func)

    interpret_individual_parser = interpret_subparser.add_parser(
        "individual",
        parents=[parent_parser, interpret_parent_parser],
        help="Individual saliency maps for each session in the input TSV file."
    )

    interpret_individual_parser.set_defaults(func=interpret_func)

    return parser

def return_train_parent_parser(retrain=False):
    # Main train parent parser common to train and random search
    train_parent_parser = argparse.ArgumentParser(add_help=False)
    train_pos_group = train_parent_parser.add_argument_group(
        TRAIN_CATEGORIES["POSITIONAL"])
    if retrain:
        train_pos_group.add_argument(
            "model_path", type=str,
            help="Path to the trained model folder.")
        train_pos_group.add_argument(
            "output_dir", type=str,
            help="Directory in which the new job is stored.")
    else:
        train_pos_group.add_argument(
            'caps_dir',
            help='Data using CAPS structure.',
            default=None)
        train_pos_group.add_argument(
            'preprocessing',
            help='Defines the type of preprocessing of CAPS data.',
            choices=['t1-linear', 't1-extensive'], type=str)
        train_pos_group.add_argument(
            'tsv_path',
            help='TSV path with subjects/sessions to process.',
            default=None)
        train_pos_group.add_argument(
            'output_dir',
            help='Folder containing results of the training.',
            default=None)
        train_pos_group.add_argument(
            'model',
            help='CNN Model to be used during the training.',
            default='Conv5_FC3')

    train_comput_group = train_parent_parser.add_argument_group(
        TRAIN_CATEGORIES["COMPUTATIONAL"])
    train_comput_group.add_argument(
        '-cpu', '--use_cpu', action='store_true',
        help='If provided, will use CPU instead of GPU.',
        default=False)
    train_comput_group.add_argument(
        '-np', '--nproc',
        help='Number of cores used during the training. (default=2)',
        type=int, default=2)
    train_comput_group.add_argument(
        '--batch_size',
        default=2, type=int,
        help='Batch size for training. (default=2)')
    train_comput_group.add_argument(
        '--evaluation_steps', '-esteps',
        default=0, type=int,
        help='Fix the number of iterations to perform before computing an evaluations. Default will only '
             'perform one evaluation at the end of each epoch.')

    train_data_group = train_parent_parser.add_argument_group(
        TRAIN_CATEGORIES["DATA"])

    if retrain:
        train_data_group.add_argument(
            "--caps_dir", type=str, default=None,
            help="Data using CAPS structure.")
        train_data_group.add_argument(
            "--tsv_path", type=str, default=None,
            help="TSV path with subjects/sessions to process.")

    train_data_group.add_argument(
        '--diagnoses', '-d',
        help='List of diagnoses that will be selected for training.',
        default=None if retrain else ['AD', 'CN'], nargs='+', type=str, choices=['AD', 'CN', 'MCI', 'sMCI', 'pMCI'])
    train_data_group.add_argument(
        '--baseline',
        help='If provided, only the baseline sessions are used for training.',
        action="store_true",
        default=None if retrain else False)
    train_data_group.add_argument(
        '--unnormalize', '-un',
        help='Disable default MinMaxNormalization.',
        action="store_true",
        default=None if retrain else False)
    train_data_group.add_argument(
        "--data_augmentation", nargs="+", default=None if retrain else False,
        choices=["None", "Noise", "Erasing", "CropPad", "Smoothing"],
        help="Randomly applies transforms on the training set.")
    train_data_group.add_argument(
        '--sampler', '-s',
        help="Sampler choice (random, or weighted for imbalanced datasets)",
        default="random", type=str, choices=["random", "weighted"])

    train_cv_group = train_parent_parser.add_argument_group(
        TRAIN_CATEGORIES["CROSS-VALIDATION"])
    train_cv_group.add_argument(
        '--n_splits',
        help='If a value is given will load data of a k-fold CV. Else will load a single split.',
        type=int, default=None)
    train_cv_group.add_argument(
        '--split',
        help='Train the list of given folds. By default train all folds.',
        type=int, default=None, nargs='+')

    train_optim_group = train_parent_parser.add_argument_group(
        TRAIN_CATEGORIES["OPTIMIZATION"])
    train_optim_group.add_argument(
        '--epochs',
        help='Maximum number of epochs.',
        default=None if retrain else 20, type=int)
    train_optim_group.add_argument(
        '--learning_rate', '-lr',
        help='Learning rate of the optimization.',
        default=None if retrain else 1e-4, type=float)
    train_optim_group.add_argument(
        '--weight_decay', '-wd',
        help='Weight decay value used in optimization.',
        default=None if retrain else 1e-4, type=float)
    train_optim_group.add_argument(
        '--dropout',
        help='rate of dropout that will be applied to dropout layers in CNN.',
        default=None if retrain else 0, type=float)
    train_optim_group.add_argument(
        '--patience',
        help='Number of epochs for early stopping patience.',
        type=int, default=None if retrain else 0)
    train_optim_group.add_argument(
        '--tolerance',
        help='Value for the early stopping tolerance.',
        type=float, default=None if retrain else 0.0)
    train_optim_group.add_argument(
        '--accumulation_steps', '-asteps',
        help='Accumulates gradients during the given number of iterations before performing the weight update '
             'in order to virtually increase the size of the batch.',
        default=None if retrain else 1, type=int)
    # train_optim_group.add_argument(
    #     "--loss",
    #     help="Replaces default losses: cross-entropy for CNN and MSE for autoencoders.",
    #     type=str, default=None if retrain else "default",
    #     choices=["default", "L1", "L1Norm", "SmoothL1", "SmoothL1Norm"])

    return train_parent_parser


import matplotlib.pyplot as plt
import numpy as np

Fs = 300  # sample rate
Ts = 1 / Fs  # sample period
N = 2048  # number of samples to simulate

t = Ts * np.arange(N)
x = np.exp(1j * 2 * np.pi * 50 * t)  # simulates sinusoid at +50 Hz

n = (np.random.randn(N) + 1j * np.random.randn(N)) / np.sqrt(2)  # AWGN with unity power
noise_power = 2
r = x + n * np.sqrt(noise_power)

PSD = (np.abs(np.fft.fft(r)) / N) ** 2
PSD_log = 10.0 * np.log10(PSD)
PSD_shifted = np.fft.fftshift(PSD_log)

f = np.linspace(Fs / -2.0, Fs / 2.0, N)  # lazy method

plt.plot(f, PSD_shifted)
plt.grid(True)
plt.xlabel("Frequency [Hz]")
plt.ylabel("Magnitude [dB]")

# plt.show()

plt.savefig('../_static/fft_example1.svg', bbox_inches='tight')


import logging
from unittest import TestCase

from tests.st.utils.utils import (get_ip, wipe_etcdv2, wipe_etcdv3,
                                  set_version_etcdv2, set_ready_etcdv2, get_value_etcdv2)

HOST_IPV4 = get_ip()

logging.basicConfig(level=logging.DEBUG, format="%(message)s")
logger = logging.getLogger(__name__)

class TestBase(TestCase):
    """
    Base class for test-wide methods.
    """

    def setUp(self):
        """
        Clean up before every test.
        """
        self.ip = HOST_IPV4
        self.wipe_etcdv2()
        self.wipe_etcdv3()
        minimum_starting_state()

        # Log a newline to ensure that the first log appears on its own line.
        logger.info("")

    def wipe_etcdv2(self):
        wipe_etcdv2(self.ip)

    def wipe_etcdv3(self):
        wipe_etcdv3(self.ip)

def minimum_starting_state():
    logger.debug("INFO: Setting minimum starting state config for calicoctl v1")
    set_version_etcdv2("v2.6.6-2-g0f0e2184")
    set_ready_etcdv2("true")
    ready_output = get_value_etcdv2("/calico/v1/Ready")
    assert ready_output == "true"


import json
from .exec import exec_and_capture_output
from .view import View
from .json_utils import JsonDeltaCalculator
from .openai_utils import openai_streaming_response

functions = [{
    "name": "run_code",
    "description": "Executes given code in an IPython shell, capturing prints, return values, terminal outputs, and tracebacks.",
    "parameters": {
        "type": "object",
        "properties": {
            "code": {
                "type": "string",
                "description": "The code to execute as a JSON decodable string. Can include standard Python and IPython commands."
            }
        },
        "required": ["code"],
    },
    "function": exec_and_capture_output
}]

system_message = """
Only use the function you have been provided with, which has one keyword argument: code.
Remember to write your messages to the user in Markdown.
You can access the internet. Run whatever code you'd like to achieve the goal, and if at first you don't succeed, try again and again.
If you run_code and it fails, write a normal message explaining what happened, theorizing why, and planning a new way forward.
You can install new packages with !pip.
While you can generate and display static plots (like those from Matplotlib), interactive and dynamic visualizations (like those from Plotly) won't be displayed correctly, so if you need to do something like that, save it as an image and display it.
""".strip()

class Interpreter:
    def __init__(self):
        self.messages = []
        self._logs = []
        self.system_message = system_message
        self.openai_api_key = None

    def reset(self):
        self.messages = []
        self._logs = []

    def load(self, messages):
        self.messages = messages

    def chat(self, message=None, return_chat=False):
        if message:
            self.messages.append({"role": "user", "content": message})
            self.respond()

        else:
            print("Type 'exit' to leave the chat.\n")

            while True:
                user_input = input("> ")

                if user_input == 'exit':
                    break

                self.messages.append({"role": "user", "content": user_input})
                self.respond()

        if return_chat:
            return self.messages

    def display(self, delta):

        old_delta = delta

        if delta == None:
          return

        if "content" in delta and delta["content"] != None:
          delta = {"type": "message", "content": delta["content"]}
        elif "function_call" in delta:
          delta = {"type": "function", "content": delta["function_call"]}

        self._logs.append(["old delta:", old_delta, "new delta:", delta])
        self.view.process_delta(delta)

    def respond(self):

        # You always need a new view.
        self.view = View()

        try:

            # make openai call
            gpt_functions = [{k: v for k, v in d.items() if k != 'function'} for d in functions]
            response = openai_streaming_response(
                self.messages,
                gpt_functions,
                self.system_message,
                "gpt-4-0613",
                self.openai_api_key
            )

            base_event = {
                "role": "assistant",
                "content": ""
            }
            event = base_event

            func_call = {
                "name": None,
                "arguments": "",
            }

            for chunk in response:

                delta = chunk.choices[0].delta

                if "function_call" in delta:
                    if "name" in delta.function_call:

                        # New event!
                        if event != base_event:
                          self.messages.append(event)
                        event = {
                            "role": "assistant",
                            "content": None
                        }

                        func_call["name"] = delta.function_call["name"]
                        self.display(delta)

                        delta_calculator = JsonDeltaCalculator()

                    if "arguments" in delta.function_call:
                        func_call["arguments"] += delta.function_call["arguments"]

                        argument_delta = delta_calculator.receive_chunk(delta.function_call["arguments"])

                        # Reassemble it as though OpenAI did this properly

                        if argument_delta != None:
                          self.display({"content": None, "function_call": argument_delta})

                if chunk.choices[0].finish_reason == "function_call":

                    event["function_call"] = func_call
                    self.messages.append(event)

                    # For interpreter
                    if func_call["name"] != "run_code":
                      func_call["name"] = "run_code"

                    function = [f for f in functions if f["name"] == func_call["name"]][0]["function"]
                    self._logs.append(func_call["arguments"])

                    # For interpreter. Sometimes it just sends the code??
                    try:
                      function_args = json.loads(func_call["arguments"])
                    except:
                      function_args = {"code": func_call["arguments"]}

                    # The output might use a rich Live display so we need to finalize ours.
                    self.view.finalize()

                    output = function(**function_args)

                    event = {
                        "role": "function",
                        "name": func_call["name"],
                        "content": output
                    }
                    self.messages.append(event)

                    # Go around again
                    self.respond()

                if "content" in delta and delta.content != None:
                    event["content"] += delta.content
                    self.display(delta)

                if chunk.choices[0].finish_reason and chunk.choices[0].finish_reason != "function_call":
                    self.messages.append(event)

        finally:
            self.view.finalize()


"""
The Demo Ballot module contains the document specifications,
page templates, and specific pages
"""
from datetime import datetime
from functools import partial
from pathlib import Path

from electos.ballotmaker.ballots.contest_layout import (
    BallotMeasureData,
    BallotMeasureLayout,
    CandidateContestData,
    CandidateContestLayout,
)
from electos.ballotmaker.ballots.instructions import Instructions
from electos.ballotmaker.ballots.page_layout import PageLayout
from electos.ballotmaker.demo_data import spacetown_data
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import inch
from reportlab.platypus import (
    BaseDocTemplate,
    Frame,
    NextPageTemplate,
    PageBreak,
    PageTemplate,
    Paragraph,
)
from reportlab.platypus.flowables import CondPageBreak

# set up frames
# 1 = True, 0 = FALSE
SHOW_BOUNDARY = 0
# get page layout settings
margin = PageLayout.margin
c_width = PageLayout.col_width
c_height = PageLayout.col_height
c_space = PageLayout.col_space
# define 3-column layout with header
left_frame = Frame(
    margin * inch,
    margin * inch,
    width=c_width * inch,
    height=c_height * inch,
    topPadding=0,
    showBoundary=SHOW_BOUNDARY,
)
mid_frame = Frame(
    (margin + c_width + c_space) * inch,
    margin * inch,
    width=c_width * inch,
    height=c_height * inch,
    topPadding=0,
    showBoundary=SHOW_BOUNDARY,
)
right_frame = Frame(
    (margin + (2 * (c_width + c_space))) * inch,
    margin * inch,
    width=c_width * inch,
    height=c_height * inch,
    topPadding=0,
    showBoundary=SHOW_BOUNDARY,
)

one_frame = Frame(
    margin * inch,
    margin * inch,
    width=7 * inch,
    height=c_height * inch,
    topPadding=0,
    showBoundary=SHOW_BOUNDARY,
)

def get_election_header() -> dict:
    return {
        "Name": "General Election",
        "StartDate": "2024-11-05",
        "EndDate": "2024-11-05",
        "Type": "general",
        "ElectionScope": "Spacetown Precinct, Orbit City",
    }

def add_header_line(font_size, line_text, new_line=False):
    line_end = "
" if new_line else ""
    return f"{line_text}{line_end}"

def build_header_text():
    elect_dict = get_election_header()
    font_size = 12
    formatted_header = add_header_line(
        font_size, f"Sample Ballot for {elect_dict['Name']}", new_line=True
    )
    formatted_header += add_header_line(
        font_size, elect_dict["ElectionScope"], new_line=True
    )
    end_date = datetime.fromisoformat(elect_dict["EndDate"])
    formatted_date = end_date.strftime("%B %m, %Y")
    formatted_header += add_header_line(font_size, formatted_date)

    return formatted_header

def header(canvas, doc, content):
    canvas.saveState()
    # these variables are used elsewhere by ReportLab
    width, height = content.wrap(doc.width, doc.topMargin)
    content.drawOn(canvas, PageLayout.margin * inch, 10.75 * inch)
    canvas.restoreState()

def build_ballot() -> str:
    # create PDF filename; include
    # datestamp string for PDF
    now = datetime.now()
    date_time = now.strftime("%Y_%m_%dT%H%M%S")
    home_dir = Path.home()
    ballot_name = f"{home_dir}/ballot_demo_{date_time}.pdf"

    doc = BaseDocTemplate(ballot_name)

    styles = getSampleStyleSheet()
    normal = styles["Normal"]
    head_text = build_header_text()
    header_content = Paragraph(head_text, normal)
    three_column_template = PageTemplate(
        id="3col",
        frames=[left_frame, mid_frame, right_frame],
        onPage=partial(header, content=header_content),
    )
    one_column_template = PageTemplate(
        id="1col",
        frames=[one_frame],
        onPage=partial(
            header,
            content=header_content,
        ),
    )
    doc.addPageTemplates(three_column_template)
    doc.addPageTemplates(one_column_template)
    # add a ballot contest to the second frame (colomn)
    layout_1 = CandidateContestLayout(
        CandidateContestData(spacetown_data.can_con_1)
    )
    layout_2 = CandidateContestLayout(
        CandidateContestData(spacetown_data.can_con_2)
    )
    layout_3 = CandidateContestLayout(
        CandidateContestData(spacetown_data.can_con_3)
    )
    layout_4 = CandidateContestLayout(
        CandidateContestData(spacetown_data.can_con_4)
    )
    layout_5 = BallotMeasureLayout(
        BallotMeasureData(spacetown_data.ballot_measure_1)
    )
    layout_6 = BallotMeasureLayout(
        BallotMeasureData(spacetown_data.ballot_measure_2)
    )
    elements = []
    # add voting instructions
    inst = Instructions()
    elements = inst.instruction_list
    elements.append(NextPageTemplate("3col"))
    elements.append(layout_1.contest_table)
    elements.append(layout_2.contest_table)
    elements.append(CondPageBreak(c_height * inch))
    elements.append(layout_4.contest_table)
    elements.append(layout_3.contest_table)
    elements.append(NextPageTemplate("1col"))
    elements.append(PageBreak())
    elements.append(layout_5.contest_table)
    elements.append(layout_6.contest_table)
    doc.build(elements)
    return str(ballot_name)

if __name__ == "__main__":  # pragma: no cover
    build_ballot()


import torch
import torch.nn as nn
import torch.optim as optim
import config
import os
from dataset import TrainDataset, TestDataset
from generator_model import Generator
from discriminator_model import Discriminator
from VGGNet import VGGNet
from torch.utils.data import DataLoader
from tqdm import tqdm
from torchvision.utils import save_image
from utils import save_test_examples, load_checkpoint, save_checkpoint, save_training_images

def initialization_phase(gen, loader, opt_gen, l1_loss, VGG, pretrain_epochs):
    for epoch in range(pretrain_epochs):
        loop = tqdm(loader, leave=True)
        losses = []

        for idx, (sample_photo, _, _) in enumerate(loop):
            sample_photo = sample_photo.to(config.DEVICE)

            # train generator G
            #with torch.cuda.amp.autocast():
            reconstructed = gen(sample_photo)

            sample_photo_feature = VGG(sample_photo)
            reconstructed_feature = VGG(reconstructed)
            reconstruction_loss = l1_loss(reconstructed_feature, sample_photo_feature.detach())

            losses.append(reconstruction_loss.item())

            opt_gen.zero_grad()

            reconstruction_loss.backward()
            opt_gen.step()

            #opt_gen.zero_grad()
            #g_scaler.scale(reconstruction_loss).backward()
            #g_scaler.step(opt_gen)
            #g_scaler.update()
            loop.set_postfix(epoch=epoch + 1)

        print('[%d/%d] - Recon loss: %.8f' % ((epoch + 1), pretrain_epochs, torch.mean(torch.FloatTensor(losses))))

        save_image(sample_photo*0.5+0.5, os.path.join(config.RESULT_TRAIN_DIR, str(epoch + 1) + "_initialization_phase_photo.png"))
        save_image(reconstructed*0.5+0.5, os.path.join(config.RESULT_TRAIN_DIR, str(epoch + 1) + "_initialization_phase_reconstructed.png"))

def train_fn(total_epoch, disc, gen, train_loader, val_loader, opt_disc, opt_gen, l1_loss, mse, VGG):
    step = 0

    for epoch in range(total_epoch):
        loop = tqdm(train_loader, leave=True)

        # Training
        for idx, (sample_photo, sample_cartoon, sample_edge) in enumerate(loop):
            sample_photo = sample_photo.to(config.DEVICE)
            sample_cartoon = sample_cartoon.to(config.DEVICE)
            sample_edge = sample_edge.to(config.DEVICE)

            # Train Discriminator
            #Pass samples into Discriminator: Fake Cartoon, real Cartoon and Edge
            fake_cartoon = gen(sample_photo)

            D_real = disc(sample_cartoon)
            D_fake = disc(fake_cartoon.detach())
            D_edge = disc(sample_edge)

            #Compute loss (edge-promoting adversarial loss)
            D_real_loss = mse(D_real, torch.ones_like(D_real))
            D_fake_loss = mse(D_fake, torch.zeros_like(D_fake))
            D_edge_loss = mse(D_edge, torch.zeros_like(D_edge))

            # Author's code divided it by 3.0, I believe it has similar thoughts to CycleGAN (divided by 2 with only 2 loss)
            D_loss = (D_real_loss + D_fake_loss + D_edge_loss) / 3.0

            opt_disc.zero_grad() # clears old gradients from the last step

            D_loss.backward()
            opt_disc.step()

            # Train Generator
            D_real = disc(sample_cartoon)
            D_fake = disc(fake_cartoon.detach())

            G_fake_loss = mse(D_fake, torch.ones_like(D_fake))

            # Content loss
            sample_photo_feature = VGG(sample_photo)
            fake_cartoon_feature = VGG(fake_cartoon)
            content_loss = l1_loss(fake_cartoon_feature, sample_photo_feature.detach())

            # Compute loss (adversarial loss + lambda*content loss)
            G_loss = config.LAMBDA_ADV * G_fake_loss + config.LAMBDA_CONTENT * content_loss

            opt_gen.zero_grad()

            G_loss.backward()
            opt_gen.step()

            if step % config.SAVE_IMG_PER_STEP == 0:
                save_image(sample_photo*0.5+0.5, os.path.join(config.RESULT_TRAIN_DIR, "epoch_" + str(epoch + 1) + "step_" + str(step + 1) + "_photo.png"))
                save_image(fake_cartoon*0.5+0.5, os.path.join(config.RESULT_TRAIN_DIR, "epoch_" + str(epoch + 1) + "step_" + str(step + 1) + "_fakecartoon.png"))
                print('[Epoch: %d| Step: %d] - D loss: %.12f' % ((epoch + 1), (step+1), D_loss.item()))
                print('[Epoch: %d| Step: %d] - G Content loss (w/o lambda): %.12f' % ((epoch + 1), (step+1), content_loss.item()))
                print('[Epoch: %d| Step: %d] - G Content loss (w/ lambda): %.12f' % ((epoch + 1), (step+1), config.LAMBDA_CONTENT * content_loss.item()))
                print('[Epoch: %d| Step: %d] - G ADV loss (w/o lambda): %.12f' % ((epoch + 1), (step+1), G_fake_loss.item()))
                print('[Epoch: %d| Step: %d] - G ADV loss (w/ lambda): %.12f' % ((epoch + 1), (step+1), config.LAMBDA_ADV * G_fake_loss.item()))
                print('[Epoch: %d| Step: %d] - G loss: %.12f' % ((epoch + 1), (step+1), G_loss.item()))

            step+= 1

            loop.set_postfix(step=step, epoch=epoch+1)

        if config.SAVE_MODEL and epoch % config.SAVE_MODEL_FREQ == 0:
            save_checkpoint(gen, opt_gen, epoch, folder=config.CHECKPOINT_FOLDER, filename=config.CHECKPOINT_GEN)
            save_checkpoint(disc, opt_disc, epoch, folder=config.CHECKPOINT_FOLDER, filename=config.CHECKPOINT_DISC)

        # Test Some data
        save_test_examples(gen, val_loader, epoch, folder=config.RESULT_TEST_DIR)

def main():
    print(config.DEVICE)
    disc = Discriminator(in_channels=3).to(config.DEVICE)
    gen = Generator(in_channels=3).to(config.DEVICE)
    opt_disc = optim.Adam(disc.parameters(), lr=config.LEARNING_RATE, betas=(0.5, 0.999))
    opt_gen = optim.Adam(gen.parameters(), lr=config.LEARNING_RATE, betas=(0.5, 0.999))

    VGG19 = VGGNet(in_channels=3, VGGtype="VGG19", init_weights=config.VGG_WEIGHTS, batch_norm=False, feature_mode=True)
    VGG19 = VGG19.to(config.DEVICE)
    VGG19.eval()

    if config.LOAD_MODEL:
        is_gen_loaded = load_checkpoint(
            gen, opt_gen, config.LEARNING_RATE, folder=config.CHECKPOINT_FOLDER, checkpoint_file=config.LOAD_CHECKPOINT_GEN
        )
        is_disc_loaded = load_checkpoint(
            disc, opt_disc, config.LEARNING_RATE, folder=config.CHECKPOINT_FOLDER, checkpoint_file=config.LOAD_CHECKPOINT_DISC
        )

    #BCE_Loss = nn.BCELoss()
    L1_Loss = nn.L1Loss()
    MSE_Loss = nn.MSELoss() # went through the author's code and found him using LSGAN, LSGAN should gives better training

    train_dataset = TrainDataset(config.TRAIN_PHOTO_DIR, config.TRAIN_CARTOON_EDGE_DIR)
    train_loader = DataLoader(train_dataset, batch_size=config.BATCH_SIZE, shuffle=True, num_workers=config.NUM_WORKERS)

    val_dataset = TestDataset(config.VAL_PHOTO_DIR)
    val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=config.NUM_WORKERS)

    # Initialization Phase
    if not(is_gen_loaded):
        print("="*80)
        print("=> Initialization Phase")
        initialization_phase(gen, train_loader, opt_gen, L1_Loss, VGG19, pretrain_epochs=config.PRETRAIN_EPOCHS)
        print("Finished Initialization Phase")
        print("="*80)

        if config.SAVE_MODEL:
            save_checkpoint(gen, opt_gen, 'i', folder=config.CHECKPOINT_FOLDER, filename=config.CHECKPOINT_GEN)

    # Do the training
    print("=> Start Training")
    train_fn(config.NUM_EPOCHS, disc, gen, train_loader, val_loader, opt_disc, opt_gen, L1_Loss, MSE_Loss, VGG19)
    print("Finished Training")

if __name__ == "__main__":
    main()

#TODO
#https://github.com/SystemErrorWang/CartoonGAN/blob/master/old_code/main.py


import os
import numpy as np
from scipy import stats
import pandas as pd
import statsmodels.api as sm
import statsmodels.formula.api as smf
import matplotlib.pyplot as plt
import sys

path = str(sys.path[0]).replace('\\run','')

churn = pd.read_csv(path + '/data/telecom_churn.csv')
# print(churn.head())
# print(churn.info())
# 数据没有缺失情况

# 相关关系分析
# 列联表
cross_table = pd.crosstab(churn.posTrend, churn.churn, margins=True)
# print(cross_table)
def percConvert(ser):
    return ser/float(ser[-1])
cross_table.apply(percConvert, axis=1)
# 列联表的卡方独立性检验：若显著，则说明二者存在相关关系
chi2, p, dof, expected_freq = stats.chi2_contingency(cross_table.iloc[:2, :2])
print(f"chisq = {chi2:.4f}\np-value = {p:.4f}\ndof = {dof}\nexpected_frep = {expected_freq}")
# 随机抽样 建立训练集与测试集
train = churn.sample(frac=0.7, random_state=666)
test = churn[~ churn.index.isin(train.index)]

print(' 训练集样本量: %i \n 测试样本量: %i' %(len(train), len(test)))

# 建立一元逻辑回归模型
lg = smf.logit('churn ~ duration', train).fit()
# 使用summary函数查看模型的一些信息
print(lg.summary())

# 变量选择 向前回归法
def forward_select(data, response):
    remaining = set(data.columns)
    remaining.remove(response)
    selected = []
    current_score, best_new_score = float('inf'), float('inf')
    while remaining:
        aic_with_cadidates = []
        for candidate in remaining:
            formula = "{} ~ {}".format(response, ' + '.join(selected + [candidate]))
            aic = smf.logit(formula=formula, data=data).fit().aic
            aic_with_cadidates.append((aic, candidate))

        aic_with_cadidates.sort(reverse=True)
        best_new_score, best_candidate = aic_with_cadidates.pop()
        if current_score > best_new_score:
            remaining.remove(best_candidate)
            selected.append(best_candidate)
            current_score = best_new_score
            print('aic is {},continuing!'.format(current_score))
        else:
            print('forward selection over!')
            break

        formula = "{} ~ {} ".format(response,' + '.join(selected))
        print('final formula is {}'.format(formula))
        model = smf.logit(formula=formula, data=data).fit()
        return model
# 对连续型变量进行筛选
candidates = ['churn', 'duration', 'AGE', 'incomeCode', 'nrProm', 'peakMinAv', 'peakMinDiff', 'call_10086']
data_for_select = train[candidates]
lg_ml = forward_select(data=data_for_select, response='churn')
lg_ml.summary()
print(f'原来有 {len(candidates)-1} 个变量')
print(f'筛选剩下 {len(lg_ml.params.index.tolist())} 个（包含 intercept 截距项）。')

# 对分类型变量进行筛选 判断个分类变量的显著性
class_col = ['posTrend', 'prom', 'edu_class', 'feton', 'curPlan', 'avgplan', 'gender', 'negTrend', 'planChange', 'posPlanChange', 'negPlanChange']
for i in class_col:
    tab = pd.crosstab(churn[i], churn.churn)
    print(i,''' p-value = %6.4f''' %stats.chi2_contingency(tab)[1])

# 建立逻辑回归模型
formula = '''churn ~ duration + AGE + incomeCode + peakMinDiff + call_10086 + C(edu_class) + C(feton) + C(posTrend) +
C(negTrend) + C(curPlan) + C(avgplan) + C(gender)'''
lg_m = smf.logit(formula=formula, data=train).fit()
print(lg_m.summary())

# 方差膨胀因子检测
def vif(df, col_i):
    from statsmodels.formula.api import ols
    cols = list(df.columns)
    cols.remove(col_i)
    cols_noti = cols
    formula = col_i + '~' + '+'.join(cols_noti)
    r2 = ols(formula, df).fit().rsquared
    return 1. / (1. - r2)

exog = train[candidates].drop(['churn'], axis=1)
for i in exog.columns:
    print(i, '\t', vif(df=exog, col_i=i))

train['proba'] = lg_m.predict(train)
test['proba'] = lg_m.predict(test)
print(test.info())

# ROC曲线
# 模型的准确率
# acc_m = sum(test['prediction'] == test['churn']) / np.float(len(test))
# print('The accurancy is %.2f' %acc_m)

# 混淆矩阵
test['prediction'] = (test['proba'] > 0.5).astype('int')
print(pd.crosstab(test.churn, test.prediction, margins=True))

# ROC曲线
for i in np.arange(0.1, 0.9, 0.1):
    prediction = (test['proba'] > i).astype('int')
    confusion_matrix = pd.crosstab(prediction, test.churn, margins=True)
    precision = confusion_matrix.loc[0, 0] / confusion_matrix.loc['All', 0]
    recall = confusion_matrix.loc[0, 0] / confusion_matrix.loc[0, 'All']
    Specificity = confusion_matrix.loc[1, 1] / confusion_matrix.loc[1, 'All']
    f1_score = 2 * (precision * recall) / (precision + recall)
    print('threshold:%s, precision:%.2f, recall:%.2f, Specificity:%.2f, f1_score:%.2f' %(i, precision, recall, Specificity, f1_score))

import sklearn.metrics as metrics
fpr_test, tpr_test, th_test = metrics.roc_curve(test.churn, test.proba)
fpr_train, tpr_train, th_train = metrics.roc_curve(train.churn, train.proba)

print('AUC = %.4f' %metrics.auc(fpr_test, tpr_test))
# 计算最优阈值
print(th_test[(tpr_test - fpr_test).argmax()])

plt.figure(figsize=[3, 3])
plt.plot(fpr_test, tpr_test, 'b--')
plt.plot(fpr_train, tpr_train, 'r-')
plt.title('ROC curve')
plt.show()


from django.contrib.auth.decorators import login_required
from django.contrib.auth.mixins import LoginRequiredMixin
from django.shortcuts import render
from django.views.generic import ListView, DetailView

from expenses.forms import NewQuickExpenseForm
from expenses.models import Expense

@login_required
def home_page(request):
    new_quick_expense_form = NewQuickExpenseForm(data=request.POST if request.POST else None,
                                                 initial={'paid_by': request.user, 'added_by': request.user})

    if request.POST and new_quick_expense_form.is_valid():
        new_quick_expense_form.save()
        new_quick_expense_form = NewQuickExpenseForm()

    recent_expenses = Expense.objects.order_by('-last_edited')[:10]
    return render(request, 'home.html', context={'form': new_quick_expense_form, 'recent_expenses': recent_expenses})

class ExpenseListView(LoginRequiredMixin, ListView):
    queryset = Expense.objects.order_by("-date")
    context_object_name = 'expenses'

class ExpenseDetailView(LoginRequiredMixin, DetailView):
    queryset = Expense.objects.all()


from rest_framework import serializers
from users.models import User

from .models import Review

class CriticSerializer(serializers.ModelSerializer):
    class Meta:
        model = User
        fields = ['id', 'first_name', 'last_name']

class ReviewSerializer(serializers.ModelSerializer):
    critic = CriticSerializer(read_only=True)

    class Meta:
        model = Review
        fields = ['id', 'stars', 'review', 'spoilers', 'recommendation','critic', 'movie_id']
        read_only_fields = ['id', 'movie_id', 'critic']
        extra_kwargs = {'stars': {'max_value':10, 'min_value':1}}

    def create(self, validated_data):
        review = Review.objects.create(**validated_data)

        return review


from flask import Flask
from flask import render_template, jsonify, request, g, abort, redirect, url_for

import json

#from xinui.rest import Rest
from rest import Rest

app = Flask(__name__)

global_flow_table = {}

# Summary page
@app.route("/")
def summary():
    sw_desc = Rest.get_switch_desc()
    if sw_desc is False:
        abort(404)
    return render_template("summary.html", sw_desc=sw_desc)

# Policy page
@app.route("/policy/")
def policy():
    sw_dpid_list = Rest.get_switch_list()

    if sw_dpid_list is False:
        abort(404)

    return render_template("policy.html", sw_dpid_list=sw_dpid_list)

@app.route("/policy/<dpid>")
def dpid_policy(dpid):
    sw_dpid_list = Rest.get_switch_list()

    if sw_dpid_list is False or int(dpid) not in sw_dpid_list:
        abort(404)

    flow_table = None
    flow_table = Rest.get_flow_table(dpid)
    global global_flow_table
    global_flow_table = flow_table
    port = Rest.get_switch_port(dpid)

    return render_template("policy.html", sw_dpid_list=sw_dpid_list, dpid=dpid, port=port, flow_table=flow_table)

# Topology page
@app.route("/topology")
def topology():
    topo = Rest.get_topology()
    return render_template("topology.html", topo=topo)

@app.route("/_query_flow/")
def query_flow(dpid):
    #dpid = request.args.get("dpid")
    flow_table = Rest.get_flow_table(dpid)

    #print flow_table

    return render_template("policy.html", flow_table=flow_table)
    #return jsonify(flow_table)

@app.route("/_query_port")
def query_port():
    dpid = request.args.get("dpid")
    port = Rest.get_switch_port(dpid)

    return jsonify(port)

@app.route("/_add_flow", methods=["POST"])
def add_flow():
    req = json.loads(request.form["flow_cmd"])

    flow_cmd = {}
    match_dict = {}
    action_list = []

    flow_cmd["dpid"] = req["common"]["dpid"]
    flow_cmd["priority"] = req["common"]["priority"]
    flow_cmd["idle_timeout"] = req["common"]["idle"]
    flow_cmd["hard_timeout"] = req["common"]["hard"]

    if req["match"]["input"] != "Any":
        match_dict["in_port"] = int(req["match"]["input"])

    if req["match"]["dl_saddr"] != "None":
        match_dict["dl_src"] = req["match"]["dl_saddr"]

    if req["match"]["dl_daddr"] != "None":
        match_dict["dl_dst"] = req["match"]["dl_daddr"]

    if req["match"]["nw_saddr"] != "None":
        match_dict["nw_src"] = req["match"]["nw_saddr"]

    if req["match"]["nw_daddr"] != "None":
        match_dict["nw_dst"] = req["match"]["nw_daddr"]

    if req["match"]["nw_saddr"] != "None" or req["match"]["nw_daddr"] != "None" or req["match"]["l4_proto"]:
        match_dict["dl_type"] = 2048

    if req["match"]["l4_proto"] != "None":
        if req["match"]["l4_proto"] == "TCP":
            match_dict["nw_proto"] = 6
        elif req["match"]["l4_proto"] == "UDP":
            match_dict["nw_proto"] = 17

    if req["match"]["sport"] != "None":
        match_dict["tp_src"] = int(req["match"]["sport"])

    if req["match"]["dport"] != "None":
        match_dict["tp_dst"] = int(req["match"]["dport"])

    if req["match"]["vlan_id"] != "None":
        match_dict["dl_vlan"] = int(req["match"]["vlan_id"])

    if req["action"]["output"] != "Drop":
        for value in req["action"]["output"].split(" "):
            action_dict = {}
            action_dict["port"] = int(value)
            action_dict["type"] = "OUTPUT"
            action_list.append(action_dict)

    if req["action"]["vlan_action"] != "None":
        if "Strip" in req["action"]["vlan_action"]:
            action_dict = {}
            action_dict["type"] = "POP_VLAN"
            action_list.append(action_dict)
        elif "Swap" in req["action"]["vlan_action"]:
            action_dict = {}
            action_dict["field"] = "vlan_vid"
            action_dict["value"] = req["action"]["vlan_action"].split(" ")[-1]
            action_dict["type"] = "SET_FIELD"
            action_list.append(action_dict)
        elif "New" in req["action"]["vlan_action"]:
            push_vlan_dict = {}
            push_vlan_dict["ethertype"] = 33024
            push_vlan_dict["type"] = "PUSH_VLAN"
            action_list.append(push_vlan_dict)

            action_dict = {}
            action_dict["field"] = "vlan_vid"
            action_dict["value"] = req["action"]["vlan_action"].split(" ")[-1]
            action_dict["type"] = "SET_FIELD"
            action_list.append(action_dict)

    flow_cmd["match"] = match_dict
    flow_cmd["actions"] = action_list

    #print json.dumps(flow_cmd)
    Rest.add_flow(json.dumps(flow_cmd))

    return "foo"

@app.route("/_del_flow", methods=["POST"])
def del_flow():
    index = request.form["index"]

    flow_cmd = {}
    match_dict = {}

    for key, list in global_flow_table.iteritems():
        flow_cmd["dpid"] = key
        match_dict = list[int(index)]["match"]

    flow_cmd["match"] = match_dict
    #print json.dumps(flow_cmd)

    Rest.del_flow(json.dumps(flow_cmd))
    return "foo"

@app.errorhandler(404)
def page_not_found(error):
    return render_template("page_not_found.html"), 404

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=5566, debug=False)


import copy
import json

import httpretty
from datahub.access.tests.fixture import conftest
from datahub.access.tests.fixture.db import *  # noqa
from datahub.access.tests.mock_api import cc
from datahub.access.tests.modules.collector.conftest import (
    common_inner_error,
    common_param_error,
    common_success,
)
from datahub.access.tests.utils import get, post
from datahub.databus.tests.fixture.channel_fixture import add_channel  # noqa
from datahub.databus.tests.mock_api import config_server, meta
from datahub.databus.tests.mock_api.meta import get_inland_tag_ok
from datahub.storekit.tests.utils.conftest import mock_meta_content_language_configs

from datahub.databus import model_manager

param = {
    "bk_app_secret": "xxx",
    "bkdata_authentication_method": "user",
    "bk_app_code": "bk_dataweb",
    "bk_username": "admin",
    "data_scenario": "db",
    "bk_biz_id": 2,
    "description": "xx",
    "access_raw_data": {
        "raw_data_name": "db_new_01",
        "maintainer": "xxxx",
        "raw_data_alias": "asdfsaf",
        "data_source": "svr",
        "data_encoding": "UTF-8",
        "sensitivity": "private",
        "description": "xxxxx",
    },
    "access_conf_info": {
        "collection_model": {
            "collection_type": "time",
            "start_at": 0,
            "period": 1,
            "time_format": "yyyy-MM-dd HH:mm:ss",
            "increment_field": "create_at",
            "before_time": 100,
        },
        "filters": [
            {
                "op": "=",
                "logic_op": "and",
                "key": "key",
                "value": "val",
            }
        ],
        "resource": {
            "scope": [
                {
                    "db_host": "test.db.com",
                    "db_port": 10000,
                    "db_user": "user",
                    "db_pass": "pwd",
                    "db_name": "db_name",
                    "table_name": "table_name",
                    "db_type_id": 1,
                }
            ]
        },
    },
}

@pytest.mark.usefixtures("init_access_manager_config")
@pytest.mark.usefixtures("init_host_config_type")
@pytest.mark.usefixtures("add_channel")
@pytest.mark.django_db
def test_ok(mocker):
    """
    :return:
    """
    __mock_deploy_obj(mocker)

    res = post("/v3/access/deploy_plan/", param)

    want = {
        "code": "1500200",
        "message": "ok",
        "errors": None,
        "data": {"raw_data_id": 123},
        "result": True,
    }
    assert json.dumps(res, sort_keys=True) == json.dumps(want, sort_keys=True)

@pytest.mark.usefixtures("init_access_manager_config")
@pytest.mark.usefixtures("init_host_config_type")
@pytest.mark.django_db
def test_get_db_type_error():
    """
    :return:
    """
    httpretty.enable()
    httpretty.reset()
    conftest.mock_user_perm("admin")
    conftest.mock_app_perm("bk_dataweb")
    conftest.mock_create_data_id()
    conftest.mock_get_data_id("db")
    conftest.mock_get_mysql_type_error()

    res = post("/v3/access/deploy_plan/", param)

    common_inner_error(res)

@pytest.mark.usefixtures("init_access_manager_config")
@pytest.mark.usefixtures("init_host_config_type")
@pytest.mark.django_db
def test_access_db_type_url_down():
    """
    :return:
    """
    httpretty.enable()
    httpretty.reset()
    conftest.mock_user_perm("admin")
    conftest.mock_app_perm("bk_dataweb")
    conftest.mock_create_data_id()
    conftest.mock_get_data_id("db")

    res = post("/v3/access/deploy_plan/", param)

    assert not res["result"]

@pytest.mark.usefixtures("init_access_manager_config")
@pytest.mark.usefixtures("init_host_config_type")
@pytest.mark.usefixtures("add_channel")
@pytest.mark.django_db
def test_wrong_db_host(mocker):
    """
    :return:
    """
    __mock_deploy_obj(mocker)

    wrong_param = copy.deepcopy(param)
    wrong_param["db_host"] = "127.0.0.1:3306"
    wrong_param["access_conf_info"]["resource"]["scope"][0]["db_host"] = "127.0.0.1:3306"
    res = post("/v3/access/deploy_plan/", wrong_param)

    common_param_error(res, "db_host")

@pytest.mark.usefixtures("init_access_manager_config")
@pytest.mark.usefixtures("init_host_config_type")
@pytest.mark.usefixtures("add_channel")
@pytest.mark.django_db
def test_wrong_poll_type(mocker):
    __mock_deploy_obj(mocker)

    wrong_param = copy.deepcopy(param)
    wrong_param["access_conf_info"]["collection_model"]["collection_type"] = "unknown"
    res = post("/v3/access/deploy_plan/", wrong_param)
    common_param_error(res, "collection_type")

@pytest.mark.usefixtures("init_access_manager_config")
@pytest.mark.usefixtures("init_host_config_type")
@pytest.mark.usefixtures("add_channel")
@pytest.mark.django_db
def test_wrong_poll_type_param(mocker):
    __mock_deploy_obj(mocker)

    wrong_param = copy.deepcopy(param)
    del wrong_param["access_conf_info"]["collection_model"]["increment_field"]
    wrong_param["access_conf_info"]["collection_model"]["collection_type"] = "pri"
    res = post("/v3/access/deploy_plan/", wrong_param)
    common_param_error(res, "increment_field")

    wrong_param["access_conf_info"]["collection_model"]["period"] = 0
    res = post("/v3/access/deploy_plan/", wrong_param)
    common_param_error(res, "period")

    wrong_param = copy.deepcopy(param)
    del wrong_param["access_conf_info"]["collection_model"]["increment_field"]
    res = post("/v3/access/deploy_plan/", wrong_param)
    common_param_error(res, "increment_field")

    wrong_param = copy.deepcopy(param)
    wrong_param["access_conf_info"]["filters"][0]["op"] = "xx"
    res = post("/v3/access/deploy_plan/", wrong_param)
    common_param_error(res, "op")

@pytest.mark.usefixtures("init_access_manager_config")
@pytest.mark.usefixtures("init_host_config_type")
@pytest.mark.usefixtures("add_channel")
@pytest.mark.django_db
def test_wrong_time_format(mocker):
    __mock_deploy_obj(mocker)

    conftest.mock_db_data_get_task_ip_log('[{"create_at": "111:11 111"}]')

    res = post("/v3/access/deploy_plan/", param)

    common_param_error(res, "time_format")

@pytest.mark.usefixtures("init_access_manager_config")
@pytest.mark.usefixtures("init_host_config_type")
@pytest.mark.usefixtures("add_channel")
@pytest.mark.usefixtures("init_access_time_db_resource")
@pytest.mark.django_db
def test_get_resource(mocker):
    """
    :return:
    """
    __mock_deploy_obj(mocker)

    res = get("/v3/access/deploy_plan/123/?bk_username=admin")
    common_success(res)

def __mock_deploy_obj(mocker):
    httpretty.enable()
    httpretty.reset()
    conftest.mock_user_perm("admin")
    conftest.mock_app_perm("bk_dataweb")
    conftest.mock_create_data_id()
    conftest.mock_create_data_auth()
    conftest.mock_get_data_auth()
    conftest.mock_create_alert()
    conftest.mock_get_data_id("db")
    conftest.mock_fast_execute_script()
    conftest.mock_db_get_task_ip_log()
    conftest.mock_get_mysql_type()
    conftest.mock_databus_task_puller()
    cc.get_biz_location_ok()
    cc.get_tag_target_ok()
    get_inland_tag_ok()
    mock_meta_content_language_configs()
    meta.post_sync_hook_ok()
    meta.get_dm_category_ok()
    config_server.post_add_route_ok()

    mocker.patch(
        "datahub.access.collectors.base_collector.BaseAccess.check_perm_by_scenario",
        return_value=True,
    )
    mocker.patch("datahub.access.tags.create_tags", return_value=True)
    mocker.patch("datahub.access.handlers.puller.DatabusPullerHandler.create", return_value=True)
    outer_channel = model_manager.get_channel_by_id(1002)
    assert outer_channel
    mocker.patch(
        "datahub.access.raw_data.rawdata.generate_raw_data_channel_id",
        return_value=(1002, outer_channel),
    )
    mocker.patch(
        "common.auth.check_perm",
        return_value=True,
    )
    mocker.patch(
        "datahub.access.handlers.v3.paas.PaasHandler.get_app_name_by_code",
        return_value=None,
    )
    mocker.patch(
        "datahub.databus.task.puller_task.process_puller_task",
        return_value=None,
    )