#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Script to resample a tractogram to a set number of streamlines.

Default behavior:
- If the number of requested streamlines is lower than the streamline count:
DOWNSAMPLE
- If the number of requested streamlines is higher than the streamline count:
UPSAMPLE  (to prevent upsampling when not desired, use --never_upsample).

Upsampling:
    Works by generating new streamlines, either by fitting a polynomial on
    a noisy version of the streamline (with option --point_wise_std) or by
    "parallel transport"; translating copies of existing streamlines by a
    random amount (with option --tube_radius). If both options are selected,
    streamlines are first translated, then noise is added.

    Includes smoothing (of value --gaussian) to compensate for the noisiness of
    new streamlines generated by the process.

    * Note that the new streamlines may not have the same step size as initial
    streamlines.

    Can be useful to build training sets for machine learning algorithms, to
    upsample under-represented bundles or downsample over-represented bundles.

Downsampling:
    Works by selecting a subset of streamlines.

    Includes the possibility of choosing randomly *per Quickbundle cluster* to
    ensure that all clusters are represented in the final tractogram.

Example usage:
$ scil_tractogram_resample.py input.trk 1000 output.trk \
      --point_wise_std 0.5 --gaussian 5 --keep_invalid_streamlines
$ scil_viz_bundle.py output.trk --local_coloring --width=0.1
"""

import argparse
import logging

from dipy.io.streamline import save_tractogram

from scilpy.io.streamlines import load_tractogram_with_reference
from scilpy.io.utils import (add_overwrite_arg, add_reference_arg,
                             add_verbose_arg,
                             assert_inputs_exist,
                             assert_outputs_exist, ranged_type,
                             add_compression_arg)
from scilpy.tractograms.tractogram_operations import (
    split_sft_randomly,
    split_sft_randomly_per_cluster,
    upsample_tractogram)
from scilpy.version import version_string


def _build_arg_parser():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawTextHelpFormatter,
                                epilog=version_string)

    p.add_argument('in_tractogram',
                   help='Input tractography file.')
    p.add_argument('nb_streamlines', type=int,
                   help='Number of streamlines to resample the tractogram to.')
    p.add_argument('out_tractogram',
                   help='Output tractography file.')

    p.add_argument('--never_upsample', action='store_true',
                   help='Make sure to never upsample a tractogram.\n'
                        'Useful when downsampling a batch of files at once.')

    # For upsampling:
    up_group = p.add_argument_group('Upsampling options')
    up_group.add_argument('--point_wise_std',
                          type=ranged_type(float, 0, None, min_excluded=True),
                          nargs='?', const=1.0,
                          help='Noise to add to existing streamlines '
                               'points to generate new ones.\n'
                               'Value is in mm. Default if set: %(const)s.')
    up_group.add_argument('--tube_radius',
                          type=ranged_type(float, 0, None, min_excluded=True),
                          nargs='?', const=1.0,
                          help='Maximum distance to generate streamlines '
                               'around the original ones. \n'
                               'Value is in mm. Default if set: %(const)s.')
    up_group.add_argument('--gaussian', metavar='SIGMA',
                          type=ranged_type(float, 0, None, min_excluded=True),
                          help='If set, smooth the upsampled streamlines. '
                               'Value is the gaussian sigma.\n Uses the '
                               'coordinates of surrounding points on the \n'
                               'streamline to blur the streamlines.\n'
                               'A good sigma choice would be around 5. '
                               'Default: no smoothing.')
    up_group.add_argument('--keep_invalid_streamlines',
                          action='store_true',
                          help='Keep invalid newly generated streamlines '
                               'that may go out of the \nbounding box.')
    add_compression_arg(up_group)

    # For downsampling:
    downsampling_group = p.add_argument_group('Downsampling options')
    downsampling_group.add_argument(
        '--downsample_per_cluster', action='store_true',
        help='If set, downsampling will be done per cluster (computed with \n'
             'Quickbundles) to ensure that at least some streamlines are \n'
             'kept per bundle. Else, random downsampling is performed '
             '(default).')
    downsampling_group.add_argument(
        '--qbx_thresholds', nargs='+', type=float, default=[40, 30, 20],
        help="If you chose option '--downsample_per_cluster', you may set \n"
             "the Quickbundles threshold value(s) here. Default: %(default)s")

    # General
    p.add_argument('--seed', default=None, type=int,
                   help='Use a specific random seed for the resampling.')

    add_reference_arg(p)
    add_verbose_arg(p)
    add_overwrite_arg(p)

    return p


def main():
    parser = _build_arg_parser()
    args = parser.parse_args()
    logging.getLogger().setLevel(logging.getLevelName(args.verbose))

    # Verifications
    assert_inputs_exist(parser, args.in_tractogram, args.reference)
    assert_outputs_exist(parser, args, args.out_tractogram)

    if args.point_wise_std and args.point_wise_std > 10:
        logging.warning("Careful. A value --point_wise_std of {} means that "
                        "each point of the streamlines could be moved up to "
                        "{}mm. This seems suspisouly high. Proceed with care.")
    if args.tube_radius and args.tube_radius > 10:
        logging.warning("Careful. A value --tube_radius of {} means that "
                        "the new streamlines could be translated of up to "
                        "{}mm. This seems suspisouly high. Proceed with care.")

    # Loading
    logging.info("Loading sft.")
    sft = load_tractogram_with_reference(parser, args, args.in_tractogram)
    original_number = len(sft.streamlines)

    # Processing
    if args.nb_streamlines > original_number:
        if args.never_upsample:
            logging.info(
                "Number of streamlines is higher than in the original "
                "tractogram, but upsampling was not allowed. Out_tractogram "
                "will be a copy of in_tractogram. To upsample, remove option "
                "--never_upsample.")
        else:
            # Waiting here to case the error, in case no tractogram requires
            # upsampling.
            if not args.point_wise_std and not args.tube_radius:
                parser.error("Tractogram requires upsampling, but no "
                             "upsampling option has been selected. Please "
                             "choose either --point_wise_std, --tube_radius, "
                             "or --never_upsample.")
            sft = upsample_tractogram(sft, args.nb_streamlines,
                                      args.point_wise_std, args.tube_radius,
                                      args.gaussian, args.compress_th,
                                      args.seed)
    elif args.nb_streamlines < original_number:
        if args.downsample_per_cluster:
            # output contains rejected streamlines, we don't use them.
            sft, _ = split_sft_randomly_per_cluster(
                sft, [args.nb_streamlines], args.seed, args.qbx_thresholds)
            logging.info("Kept {} out of {} expected streamlines."
                         .format(len(sft), args.nb_streamlines))
        else:
            # output is a list of two: kept and rejected.
            sft = split_sft_randomly(sft, args.nb_streamlines, args.seed)[0]
    else:
        logging.info("Number of streamlines in in_tractogram is already "
                     "correct! Out_tractogram will be a copy of "
                     "in_tractogram.")

    if not args.keep_invalid_streamlines:
        sft.remove_invalid_streamlines()
        logging.info("After removing invalid streamlines, the number of "
                     "remaining streamlines is: {}".format(len(sft)))
    save_tractogram(sft, args.out_tractogram,
                    bbox_valid_check=not args.keep_invalid_streamlines)


if __name__ == "__main__":
    main()
