#!/usr/bin/env python
import argparse
import logging
import os
import time

from naps.utils.tracking import TrackingArray

logger = logging.getLogger("NAPS Logger")


def interactionParser():
    """
    Interaction Parser

    Assign the parameters for interaction assignment

    Parameters
    ----------
    sys.argv : list
            Parameters from command line

    Raises
    ------
    IOError
            If the specified files do not exist
    """

    def confirmFile():
        """Custom action to confirm file exists"""

        class customAction(argparse.Action):
            def __call__(self, parser, args, value, option_string=None):
                if not os.path.isfile(value):
                    raise IOError("%s not found" % value)
                setattr(args, self.dest, value)

        return customAction

    interaction_parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    # Input files
    interaction_parser.add_argument(
        "--sleap-file",
        help="Defines the input SLEAP file",
        type=str,
        action=confirmFile(),
        required=True,
    )
    interaction_parser.add_argument(
        "--yaml-model",
        help="Defines the YAML model file",
        type=str,
        action=confirmFile(),
        required=True,
    )

    # General options
    interaction_parser.add_argument(
        "--frame-start",
        help="Defines the frame to start interaction assignment",
        type=int,
        default=0,
    )
    interaction_parser.add_argument(
        "--frame-end",
        help="Defines the frame to end interaction assignment",
        type=int,
        default=-1,
    )
    interaction_parser.add_argument(
        "--threads", help="Number of threads to use", type=int, default=1
    )

    # interaction_parser.add_argument('--max-velocity', help = 'Defines the maximum velocity allowed in the dataset', type = int)
    # interaction_parser.add_argument('--remove-tags', help = 'Defines the tags to remove', type = str, nargs = '+', default = [])
    # interaction_parser.add_argument('--min-bin', help = 'Min interaction distance', type = int, default = 50)
    # interaction_parser.add_argument('--max-bin', help = 'Max interaction distance', type = int, default = 450)
    # interaction_parser.add_argument('--frame-buffer', help = 'Frame buffer - i.e. gap between frames allowed', type = int, default = 5)
    # interaction_parser.add_argument('--min-frames', help = 'Min frames for interaction', type = int, default = 5)

    # Output options
    interaction_parser.add_argument(
        "--out-file",
        help="Defines the output tsv file. Defaults to: out.tsv",
        type=str,
        default="out.tsv",
    )

    # Return the arguments
    return interaction_parser.parse_args()


def main():

    # Assign the interaction args
    interaction_args = interactionParser()

    # Load the tracking array
    logger.info("Loading the tracking array...")
    t0 = time.time()

    # Create the interactions object using a SLEAP file and the yaml model
    interactions = TrackingArray.fromSLEAP(
        interaction_args.sleap_file,
        model_yaml=interaction_args.yaml_model,
        threads=interaction_args.threads,
    )

    logger.info("Done loading in %s seconds.", time.time() - t0)

    # Assign interactions from the tracking array
    logger.info("Starting interaction assignment...")
    t0 = time.time()

    # Assign the interactions
    interactions.assign(
        interaction_args.out_file,
        frame_start=interaction_args.frame_start,
        frame_end=interaction_args.frame_end,
    )

    logger.info("Done with assignment in %s seconds.", time.time() - t0)


if __name__ == "__main__":
    main()
