#!/usr/bin/env python3
from accessoryFunctions.accessoryFunctions import SetupLogging
from geneseekr.parser import modify_usage_error, objector
from geneseekr.blast import BLAST
from time import time
import click
import sys

__author__ = 'adamkoziol'

SetupLogging()
start = time()

# https://stackoverflow.com/a/40195800
click_options = [
    click.version_option(version='1.0.0'),
    click.option('-s', '--sequencepath',
                 required=True,
                 help='Specify input fasta folder'),
    click.option('-t', '--targetpath',
                 required=True,
                 help='Specify folder of targets'),
    click.option('-r', '--reportpath',
                 required=True,
                 help='Specify output folder for csv'),
    click.option('-c', '--cutoff',
                 type=int,
                 default=70,
                 help='Minumum sequence identity threshold for a match to be reported. Default is 70%'),
    click.option('-e', '--evalue',
                 default='1E-5',
                 help='Minimum evalue to use for BLAST analyses. Default is 1E-5'),
    click.option('-n', '--numthreads',
                 type=int,
                 help='Specify number of threads. Defaults to the number of cores in the system'),
    click.option('-a', '--align',
                 is_flag=True,
                 help='Optionally output alignments of genes with less than 100% identity to reference '
                      'genes. This alignment will use amino acid sequences for both query and reference'),
    click.option('-u', '--unique',
                 is_flag=True,
                 help='Do not report multiple hits at the same location in a contig. Instead, store the '
                      'best hit, and ignore the rest')
]

click_nt_options = [
    click.option('-R', '--resfinder',
                 is_flag=True,
                 help='Perform ResFinder-like analyses'),
    click.option('-V', '--virulencefinder',
                 is_flag=True,
                 help='Perform VirulenceFinder-like analyses'),
    click.option('-M', '--MLST',
                 is_flag=True,
                 help='Perform allele typing (e.g. rMLST, cgMLST, wgMLST, MLST, etc.')
]


def add_options(options):
    def _add_options(func):
        for option in reversed(options):
            func = option(func)
        return func
    return _add_options


@click.group(context_settings=dict(help_option_names=['-h', '--help']))
def group():
    pass


@group.command()
@add_options(click_options)
@add_options(click_nt_options)
def blastn(**kwargs):
    """
    nt query: nt db
    """
    metadata, pipeline = objector(kwargs, start)
    metadata.program = 'blastn'
    geneseekr = BLAST(args=metadata,
                      pipeline=pipeline)
    geneseekr.seekr()


@group.command()
@add_options(click_options)
def blastp(**kwargs):
    """
    protein query: protein db
    """
    metadata, pipeline = objector(kwargs, start)
    metadata.program = 'blastp'
    geneseekr = BLAST(args=metadata,
                      pipeline=pipeline)
    geneseekr.seekr()


@group.command()
@add_options(click_options)
def blastx(**kwargs):
    """
    translated nt query: protein db
    """
    metadata, pipeline = objector(kwargs, start)
    metadata.program = 'blastx'
    geneseekr = BLAST(args=metadata,
                      pipeline=pipeline)
    geneseekr.seekr()


@group.command()
@add_options(click_options)
@add_options(click_nt_options)
def tblastn(**kwargs):
    """
    protein query: translated nt db
    """
    metadata, pipeline = objector(kwargs, start)
    metadata.program = 'tblastn'
    geneseekr = BLAST(args=metadata,
                      pipeline=pipeline)
    geneseekr.seekr()


@group.command()
@add_options(click_options)
@add_options(click_nt_options)
def tblastx(**kwargs):
    """
    translated nt query: translated nt db
    """
    metadata, pipeline = objector(kwargs, start)
    metadata.program = 'tblastx'
    geneseekr = BLAST(args=metadata,
                      pipeline=pipeline)
    geneseekr.seekr()


def find_subcommand(subcommand_string):
    """
    Associate the appropriate function with the supplied BLAST program to use
    :param subcommand_string:
    :return:
    """
    if subcommand_string == 'blastn':
        subcommand = blastn
    elif subcommand_string == 'blastp':
        subcommand = blastp
    elif subcommand_string == 'blastx':
        subcommand = blastn
    elif subcommand_string == 'tblastn':
        subcommand = blastn
    else:
        subcommand = tblastx
    return subcommand


# Extract the BLAST command to use from the command line arguments
try:
    program = sys.argv[1] if sys.argv[1] in ['blastn', 'blastp', 'blastx', 'tblastn', 'tblastx'] else str()
except IndexError:
    program = str()

# Convert the program string to the appropriate subscommand to use when modifying the usage error - ResFinder
# and VirulenceFinder analyses are only available for BLAST programs that use a nt database
sub_command = find_subcommand(program)
# Change the default behaviour of click to print the help menu when a subcommand is specified, but is missing arguments
modify_usage_error(sub_command)

if __name__ == '__main__':
    group()
