#!/usr/bin/env python3
import argparse
from contextlib import contextmanager


@contextmanager
def subparser(parser, name, description):
    yield parser.add_parser(name, conflict_handler='resolve', help=description)

@contextmanager
def mutex_group(parser, required):
    group = parser.add_argument_group(f'Mutually exclusive '
                                      f'{"required" if required else "optional"} '
                                      f'arguments')
    yield group.add_mutually_exclusive_group(required=required)

@contextmanager
def arg_group(parser, name):
    yield parser.add_argument_group(name)

def __all_plots(group, required=False):
    group.add_argument('--all_plots', action='store_true', required=required,
                       help='Generate a combined plot for CUGO, AA sequence length and TMHMM')

def __all_proteins(group, required=False):
    group.add_argument('-a', '--all_proteins', type=str, required=required,
                       help='Path to FASTA file containing all GlobDB protein sequences')

def __all_metadata(group, required=False):
    group.add_argument('--all_metadata', action='store_true', required=required,
                       help='Include all metadata in output file')

def __annotation(group, required=False):
    group.add_argument('--annotation', type=str, default='COG_ID', required=required,
                       help='Annotation: COG_ID, KEGG_ID, Pfam_ID (default: COG_ID)')

def __bin_width(group, required=False):
    group.add_argument('-b', '--bin_width', type=int, default=50, required=required,
                       help='Bin width for amino acid sequence size plotting (default: 50)')

def __blast_output(group, required=False):
    group.add_argument('-b', '--blast_output', type=str, default=None, required=required,
                       help='Path to BLAST/DIAMOND output file')

def __block(group, required=False):
    group.add_argument('-b', '--block', type=int, default=6, required=required,
                       help='blastp sequence block size in billions of letters (default: 6)')

def __bsr(group, required=False):
    group.add_argument('-b', '--bsr', type=str, required=required,
                       help='Path to blast score ratio input file.')

def __bsr_cutoff(group, required=False):
    group.add_argument('-b', '--bsr_cutoff', type=float, default=None, required=required,
                       help="Blast Score Ratio cutoff for inclusion in updated dataset")

def __chunk(group, required=False):
    group.add_argument('-c', '--chunk', type=int, default=2, required=required,
                       help='Number of chunks for diamond blastp index processing (default: 2)')

def __cluster_path(group, required=False):
    group.add_argument('-c', '--cluster_path', type=str, default=None, required=required,
                       help='Path to cluster TSV file')

def __create_yaml(group, required=False):
    group.add_argument('--create_yaml', action='store_true', required=required,
                       help='Create metadata yaml file from command line parameters')

def __cog_gff(group, required=False):
    group.add_argument('-c', '--cog_gff', type=str, required=required,
                       help='Path to (.tar.gz) GFF directory containing COG annotations')

def __column_info_path(group, required=False):
    group.add_argument('-c', '--column_info_path', type=str, default=None, required=required,
                       help='Path to json file containing columns the blast output file generated by aastk search')

def __context_path(group, required=False):
    group.add_argument('-c', '--context_path', type=str, required=required,
                       help='Path to CUGO context file')

def __cugo(group, required=False):
    group.add_argument('--cugo', action='store_true', required=required,
                       help='Generate CUGO plot')

def __cugo_range(group, required=False):
    group.add_argument('-r', '--cugo_range', type=int, required=required,
                       help='CUGO range for genomic context analysis')

def __culture_collection_path(group, required=False):
    group.add_argument('-c', '--culture_collection_path', type=str, required=required,
                       help='Path to culture collection data TSV file')

def __dataset(group, required=False):
    group.add_argument('-d', '--dataset', type=str, required=required,
                       help='Dataset name')

def __diamond_db(group, required=False):
    group.add_argument('-d', '--diamond_db', type=str, required=required,
                       help='Path to DIAMOND database')

def __dbmin(group, required=False):
    group.add_argument('-d', '--dbmin', type=int, default=None, required=required,
                       help='Lower database score cutoff for inclusion in updated dataset')

def __db_path(group, required=False):
    group.add_argument('-d', '--db_path', type=str, default='temp_genome_data.db', required=required,
                       help='Path to SQLite database')

def __exaggeration(group, required=False):
    group.add_argument('-e', '--exaggeration', type=int, default=6, required=required,
                       help='Exaggeration value for tSNE embedding (default: 6)')

def __extracted(group, required=False):
    group.add_argument('-e', '--extracted', type=str, required=required,
                       help='Path to FASTA file containing extracted matching sequences')

def __fasta(group, required=False):
    group.add_argument('-f', '--fasta', type=str, required=required,
                       help='Path to FASTA file')

def __filter_seqs(group, required=False):
    group.add_argument('--filter_seqs', action='store_true', required=required,
                       help='Filter datasets in FASTA format for improved homogeneity')

def __force(group, required=False):
    group.add_argument('--force', action='store_true', required=required,
                       help='Overwrite existing files in specified path')

def __flank_lower(group, required=False):
    group.add_argument('-l', '--flank_lower', type=int, required=required,
                       help='Start of flanking window (inclusive)')

def __flank_upper(group, required=False):
    group.add_argument('-u', '--flank_upper', type=int, required=required,
                       help='End of flanking window (inclusive)')

def __globdb_version(group, required=False):
    group.add_argument('-g', '--globdb_version', type=str, required=required,
                       help='GlobDB version (example: r226)')

def __help(group, required=False):
    group.add_argument('-h', '--help', action='help',
                       help='Display help text')

def __high_level_environment_path(group, required=False):
    group.add_argument('-h', '--high_level_environment_path', type=str, required=required,
                       help='Path to high level environment data TSV file')

def __id_list(group, required=False):
    group.add_argument('-i', '--id_list', type=str, required=required,
                       help='Path to list of GlobDB protein IDs')

def __include_annotation(group, required=False):
    group.add_argument('--include_annotation', action='store_true', required=required,
                       help='Include annotation metadata in output')

def __include_culture_collection(group, required=False):
    group.add_argument('--include_culture_collection', action='store_true', required=required,
                       help='Include culture collection metadata in output')

def __include_high_level_environment(group, required=False):
    group.add_argument('--include_high_level_environment', action='store_true', required=required,
                       help='Include high level environment metadata in output')

def __include_low_level_environment(group, required=False):
    group.add_argument('--include_low_level_environment', action='store_true', required=required,
                       help='Include low level environment metadata in output')

def __include_taxonomy(group, required=False):
    group.add_argument('--include_taxonomy', action='store_true', required=required,
                       help='Include taxonomy metadata in output')

def __iterations(group, required=False):
    group.add_argument('-i', '--iterations', type=str, default=500, required=required,
                       help='Number of tSNE embedding iterations (default: 500)')

def __keep(group, required=False):
    group.add_argument('--keep', action='store_true', required=required,
                       help='Keep intermediate files in specified paths')

def __kegg_gff(group, required=False):
    group.add_argument('-k', '--kegg_gff', type=str, required=required,
                       help='Path to (.tar.gz) GFF directory containing KEGG annotations')

def __key_column(group, required=False):
    group.add_argument('-k', '--key_column', type=int, default=0, required=required,
                       help='Column index in the BLAST tab file to pull unique IDs from (default: 0)')

def __low_level_environment_path(group, required=False):
    group.add_argument('-l', '--low_level_environment_path', type=str, required=required,
                       help='Path to high level environment data TSV file')

def __matched(group, required=False):
    group.add_argument('-m', '--matched', type=str, required=required,
                       help='Path to matched sequences FASTA file from previous PASR run')

def __matrix(group, required=False):
    group.add_argument('-m', '--matrix', type=str, default='BLOSUM45', required=required,
                       help='BLOSUM substitution matrix: BLOSUM45, BLOSUM62 (default: BLOSUM45)')

def __matrix_path(group, required=False):
    group.add_argument('-m', '--matrix_path', type=str, required=required,
                       help='Path to alignment matrix .npy file')

def __max_score_max(group, required=False):
    group.add_argument('-u', '--max_score_max', type=int, required=required,
                       help='Upper cutoff for max score range')

def __max_score_min(group, required=False):
    group.add_argument('-l', '--max_score_min', type=int, required=required,
                       help='Lower cutoff for max score range')

def __max_scores(group, required=False):
    group.add_argument('-m', '--max_scores', type=str, required=required,
                       help='Path to file containing max scores')

def __metadata_protein(group, required=False):
    group.add_argument('--metadata_protein', type=str, required=required,
                       help='Select metadata for plotting; run "aastk metadata_categories" to view available categories"')

def __metadata_matrix(group, required=False):
    group.add_argument('--metadata_matrix', type=str, required=required,
                       help='Path to alignment matrix metadata file')

def __no_cluster(group, required=False):
    group.add_argument('-n', '--no_cluster', type=int, required=required,
                       help='Number of cluster of choice in TSV file')

def __output(group, required=False):
    group.add_argument('-o', '--output', type=str, required=required,
                       help='Output directory (default: current working directory)')

def __params(group, required=False):
    group.add_argument('--params', action='store_true', required=required,
                       help='Select subset of matched sequences based on command line parameters')

def __perplexity(group, required=False):
    group.add_argument('-p', '--perplexity', type=int, default=50, required=required,
                       help='Perplexity value for tSNE embedding (default: 50)')

def __pfam_gff(group, required=False):
    group.add_argument('-p', '--pfam_gff', type=str, required=required,
                       help='Path to (.tar.gz) GFF directory containing Pfam annotations')

def __position(group, required=False):
    group.add_argument('-p', '--position', type=int, required=required,
                       help='CUGO position to retrieve protein IDs for')

def __protein_ids(group, required=False):
    group.add_argument('-p', '--protein_ids', type=str, required=required,
                       help='Path to file containing list of protein IDs')

def __query(group, required=False):
    group.add_argument('-q', '--query', type=str, default=None, required=required,
                       help='Path to DIAMOND query FASTA file')

def __score_column(group, required=False):
    group.add_argument('-s', '--score_column', type=int, default=None, required=required,
                       help='Index of the raw score column of BLAST output file')

def __seed(group, required=False):
    group.add_argument('-s', '--seed', type=str, default=None, required=required,
                       help='Path to seed FASTA file for DIAMOND database creation')

def __sensitivity(group, required=False):
    group.add_argument('--sensitivity', required=required,
                       help='DIAMOND search sensitivity level: fast, sensitive, mid-sensitive, very-sensitive, '
                            'ultra-sensitive, or faster (default: fast)')

def __show(group, required=False):
    group.add_argument('-s', '--show', action='store_true', required=required,
                       help='Show cluster number at cluster centers in output plots')

def __size(group, required=False):
    group.add_argument('--size', action='store_true', required=required,
                       help='Generate AA sequence length plot')

def __sql(group, required=False):
    group.add_argument('--sql', action='store_true', required=required,
                       help='Retrieve sequences from AASTK SQLite database')

def __subset(group, required=False):
    group.add_argument('-s', '--subset', type=str, required=required,
                       help='Path to subset FASTA')

def __subset_size(group, required=False):
    group.add_argument('--subset_size', type=int, required=required,
                       help='Number of sequences in subset FASTA file')

def __svg(group, required=False):
    group.add_argument('--svg', action='store_true', required=required,
                       help='Generate plot in SVG format')

def __taxonomy_path(group, required=False):
    group.add_argument('-t', '--taxonomy_path', type=str, default=None, required=required,
                       help='Path to gzipped taxonomy TSV file')

def __threads(group, required=False):
    group.add_argument('-n', '--threads', type=int, default=1, required=required,
                       help='Number of threads (default: 1)')

def __tmhmm_dir(group, required=False):
    group.add_argument('-t', '--tmhmm_dir', type=str, required=required,
                       help='Directory containing tmhmm files')

def __tmh_y_range(group, required=False):
    group.add_argument('-t', '--tmh_y_range', type=int, required=required,
                       help='Upper limit for TMH plot y-axis (default: maximum number of transmembrane helices)')

def __top_n(group, required=False):
    group.add_argument('-t', '--top_n', type=int, default=3, required=required,
                       help='Number of top annotations to plot per position (default: 3)')

def __tsv(group, required=False):
    group.add_argument('-t', '--tsv', type=str, required=required,
                       help='Path to ASM_clust tab separated output file')

def __update(group, required=False):
    group.add_argument('--update', action='store_true', required=required,
                       help='Update subset of data using metadata yaml file')

def __yaml(group, required=False):
    group.add_argument('-y', '--yaml', type=str, required=required,
                       help='Path to metadata yaml file')

def __y_range(group, required=False):
    group.add_argument('-y', '--y_range', type=int, required=required,
                       help='Upper limit for sequence length plot y-axis (default: maximum length)')

def get_main_parser():
    main_parser = argparse.ArgumentParser(
        prog='aastk', add_help=False, conflict_handler='resolve')

    main_parser.add_argument('--silent', action='store_true',
                             help='Suppress all console output except errors')
    sub_parsers = main_parser.add_subparsers(help="--", dest='subparser_name')


    ### PARSER FOR PASR FUNCTIONALITIES AND WORKFLOW ###
    with subparser(sub_parsers, 'build', 'Build DIAMOND database from seed sequence(s)') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __seed(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __threads(grp)
            __force(grp)

    with subparser(sub_parsers, 'search', 'Search DIAMOND reference database for homologous sequences') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __diamond_db(grp, required=True),
            __query(grp, required=True),
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __threads(grp)
            __block(grp)
            __chunk(grp)
            __sensitivity(grp)
            __force(grp)

    with subparser(sub_parsers, 'get_hit_seqs', 'Extract reads that have DIAMOND hits against custom database') as parser:
        with mutex_group(parser, required=True) as grp:
            __sql(grp)
            __query(grp)
        with arg_group(parser, 'Required arguments') as grp:
            __blast_output(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __db_path(grp)
            __threads(grp)
            __key_column(grp)
            __force(grp)

    with subparser(sub_parsers, 'max_score', 'Calculate max scores for extracted sequences using BLOSUM matrix') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __extracted(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __matrix(grp)
            __force(grp)

    with subparser(sub_parsers, 'bsr', 'Compute BSR (Blast Score Ratio) using a BLAST tab file and max scores from a TSV.') as parser:
        with mutex_group(parser, required=True) as grp:
            __column_info_path(grp)
            __score_column(grp)
        with arg_group(parser, 'Required arguments') as grp:
            __blast_output(grp, required=True)
            __max_scores(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __key_column(grp)
            __force(grp)

    with subparser(sub_parsers, 'pasr_plot', 'Plot the Blast Score Ratio of query sequences against the DIAMOND database') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __bsr(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __svg(grp)
            __force(grp)
            __update(grp)
            __yaml(grp)


    with subparser(sub_parsers, 'pasr_select', 'Select target sequences in accordance with metadata cutoffs') as parser:
        with mutex_group(parser, required=True) as grp:
            __yaml(grp)
            __params(grp)
        with arg_group(parser, 'Required arguments') as grp:
            __matched(grp, required=True)
            __bsr(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __create_yaml(grp)
            __max_score_min(grp)
            __max_score_max(grp)
            __dbmin(grp)
            __bsr_cutoff(grp)
            __force(grp)


    with subparser(sub_parsers, 'pasr', 'PASR: protein alignment score ratio') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __query(grp, required=True)
            __seed(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __matrix(grp)
            __threads(grp)
            __db_path(grp)
            __key_column(grp)
            __block(grp)
            __chunk(grp)
            __sensitivity(grp)
            __update(grp)
            __yaml(grp)
            __sql(grp)
            __keep(grp)
            __svg(grp)
            __force(grp)

    ### PARSER FOR CUGO FUNCTIONALITIES AND WORKFLOW ###
    with subparser(sub_parsers, 'database', 'Create AASTK SQLite database') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __all_proteins(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __cog_gff(grp)
            __kegg_gff(grp)
            __pfam_gff(grp)
            __taxonomy_path(grp)
            __culture_collection_path(grp)
            __high_level_environment_path(grp)
            __low_level_environment_path(grp)
            __globdb_version(grp)
            __output(grp)
            __tmhmm_dir(grp)

    with subparser(sub_parsers, 'database_check', 'Checks for missing data in AASTK SQLite database') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __db_path(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __force(grp)

    with subparser(sub_parsers, 'meta', 'Retrieve metadata from AASTK SQLite database') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __db_path(grp, required=True)
            __fasta(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __threads(grp)
            __include_annotation(grp)
            __include_culture_collection(grp)
            __include_high_level_environment(grp)
            __include_low_level_environment(grp)
            __include_taxonomy(grp)
            __all_metadata(grp)
            __force(grp)




    with subparser(sub_parsers, 'context', 'Parse context information from CUGO input file') as parser:
        with mutex_group(parser, required=True) as grp:
            __fasta(grp)
            __id_list(grp)
        with arg_group(parser, 'Required arguments') as grp:
            __db_path(grp, required=True)
            __cugo_range(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __annotation(grp)
            __output(grp)
            __threads(grp)
            __force(grp)

    with subparser(sub_parsers, 'cugo_plot', 'Plot CUGO context') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __context_path(grp, required=True)
            __flank_lower(grp, required=True)
            __flank_upper(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __annotation(grp)
            __top_n(grp)
            __cugo(grp)
            __size(grp)
            __all_plots(grp)
            __bin_width(grp)
            __y_range(grp)
            __tmh_y_range(grp)
            __output(grp)
            __force(grp)
            __svg(grp)



    with subparser(sub_parsers, 'cugo', 'CUGO: Co-localized gene organization') as parser:
        with mutex_group(parser, required=True) as grp:
            __fasta(grp)
            __id_list(grp)
        with arg_group(parser, 'Required arguments') as grp:
            __db_path(grp, required=True)
            __cugo_range(grp, required=True)
            __flank_lower(grp, required=True)
            __flank_upper(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __annotation(grp)
            __output(grp)
            __top_n(grp)
            __bin_width(grp)
            __y_range(grp)
            __tmh_y_range(grp)
            __output(grp)
            __threads(grp)
            __svg(grp)
            __force(grp)


    with subparser(sub_parsers, 'cugo_select', 'Retrieve protein IDs for select CUGO position') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __context_path(grp, required=True)
            __position(grp, required=True)
            __db_path(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __threads(grp)
            __filter_seqs(grp)
            __force(grp)

    with subparser(sub_parsers, 'filter', 'Filter datasets in FASTA format for improved homogeneity') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __fasta(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __threads(grp)
            __sql(grp)
            __db_path(grp)
            __force(grp)

    with subparser(sub_parsers, 'matrix', 'Create alignment matrix for tSNE embedding and DBSCAN clustering') as parser:
        with mutex_group(parser, required=True) as grp:
            __subset(grp)
            __subset_size(grp)
        with arg_group(parser, 'Required arguments') as grp:
            __fasta(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __threads(grp)
            __force(grp)

    with subparser(sub_parsers, 'cluster', 'Run tSNE embedding and DBSCAN clustering on input matrix and matrix metadata') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __matrix_path(grp, required=True)
            __metadata_matrix(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __perplexity(grp)
            __iterations(grp)
            __exaggeration(grp)
            __threads(grp)
            __force(grp)


    with subparser(sub_parsers, 'casm_plot', 'Plot CASM .tsv output files') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __cluster_path(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __db_path(grp)
            __metadata_protein(grp)
            __show(grp)
            __svg(grp)
            __force(grp)

    with subparser(sub_parsers, 'casm', 'CASM: protein clustering using alignment score matrices') as parser:
        with mutex_group(parser, required=True) as grp:
            __subset(grp)
            __subset_size(grp)
        with arg_group(parser, 'Required arguments') as grp:
            __fasta(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __threads(grp)
            __output(grp)
            __db_path(grp)
            __perplexity(grp)
            __iterations(grp)
            __exaggeration(grp)
            __metadata_protein(grp)
            __force(grp)
            __keep(grp)
            __show(grp)
            __svg(grp)

    with subparser(sub_parsers, 'casm_select', 'Pick CASM clusters to generate .faa file for further analysis') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __cluster_path(grp, required=True)
            __fasta(grp, required=True)
            __no_cluster(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __force(grp)

    with subparser(sub_parsers, 'list_metadata', 'Display avaiilable metadata categories') as parser:
        pass

    with subparser(sub_parsers, 'export_fasta', 'Create FASTA file containing all GlobDB protein sequences') as parser:
        with arg_group(parser, 'Required arguments') as grp:
            __db_path(grp, required=True)
        with arg_group(parser, 'Optional') as grp:
            __output(grp)
            __threads(grp)
            __force(grp)


    return main_parser
