#!/usr/bin/python3

"""
Command line interface for gffutils.

Developer notes
---------------

To implement or modify subcommands (e.g., "gffutils-cli fetch"), simply edit
the correspondingly-named functions.

Since the `argh` module is used to wrap functions into argparse.ArgumentParser
objects (and sub-parsers), just add an argument to the appropriate function and
you'll get the cli for free.
"""

import os
import gffutils
import gffutils.helpers as helpers
from gffutils.exceptions import FeatureNotFoundError
import gffutils.gffwriter as gffwriter
import argh
from argcomplete.completers import EnvironCompleter
from argh import arg
import sys

# D.R.Y.
db_help = '''Database to use.  If a GFF or GTF file is provided instead for
this argument, a database will be created for you.  This can take some time
(several minutes), so it's best to create one ahead of time.'''


def handle_relations_args(ids, limit, exclude):
    """
    Handle args for children or parents subcommands
    """
    if limit:
        raise NotImplementedError('sorry, still working on --limit')
    if ids is None:
        ids = gffdb.features_of_type('gene')
    else:
        ids = ids.split(',')
    exclude = (exclude or "").split(',')
    return ids, limit, exclude


# fetch subcommand ------------------------------------------------------------
@arg('db', help=db_help)
@arg('ids', help='Comma-separated list of IDs to fetch')
def fetch(db, ids):
    """
    Fetch IDs.
    """
    if not os.path.isfile(db):
        raise Exception("Cannot fetch: %s does not exist." % (db))
    gff_db = None
    if not helpers.is_gff_db(db):
        gff_db = helpers.get_gff_db(db)
    else:
        gff_db = db
    for i in ids.split(','):
        try:
            yield gff_db[i]
        except FeatureNotFoundError:
            print ("%s not found" % i, file=sys.stderr)


# children subcommand ---------------------------------------------------------
@arg('db', help=db_help)
@arg('ids', nargs='?', help='''Comma-separated list of IDs. These IDs, along
     with their children, will be returned (subject to --limit and --exclude).
     If none provided, use all genes in the database''')
@arg('--limit', help='''Feature type (string), or level (integer).  No children
     below this level will be returned''')
@arg('--exclude', help='''Comma-separated list of featuretypes to filter out
     (like grep -v)''')
@arg('--exclude-self', help='''Use this to suppress reporting the IDs you've
     provided.''')
def children(db, ids, limit=None, exclude=None, exclude_self=False):
    """
    Fetch children from the database according to ID.
    """
    db = gffutils.FeatureDB(db)
    ids, limit, exclude = handle_relations_args(ids, limit, exclude)

    for i in ids:
        f = db[i]
        if f.featuretype in exclude:
            continue
        if not exclude_self:
            yield f
        for child1 in db.children(i):
            if child1.featuretype in exclude:
                continue
            yield child1
            for child2 in db.children(child1):
                if child2.featuretype in exclude:
                    continue
                yield child2


# parents subcommand ----------------------------------------------------------
@arg('db', help=db_help)
@arg('ids', nargs='?', help='''Comma-separated list of IDs. These IDs, along
     with their parents, will be returned (subject to --limit and --exclude).
     If none provided, use all exons in the database''')
@arg('--limit', help='''Feature type (string), or level (integer).  No parents
     below this level will be returned''')
@arg('--exclude', help='''Comma-separated list of featuretypes to filter out
     (like grep -v)''')
@arg('--exclude-self', help='''Use this to suppress reporting the IDs you've
     provided.''')
def parents(db, ids, limit=None, exclude=None, exclude_self=False):
    """
    Fetch parents from the database according to ID.
    """
    db = gffutils.FeatureDB(db)
    ids, limit, exclude = handle_relations_args(ids, limit, exclude)
    for i in ids:
        f = db[i]
        if f.featuretype in exclude:
            continue
        if not exclude_self:
            yield f
        for parent1 in db.parents(i):
            if parent1.featuretype in exclude:
                continue
            yield parent1
            for parent2 in db.parents(parent1):
                if parent2.featuretype in exclude:
                    continue
                yield parent2


# region subcommand -----------------------------------------------------------
@arg('db', help=db_help)
@arg('region', help='Genomic coordinates of the form "chrom:start-stop"')
def region(db, region):
    """
    Returns features within provided genomic coordinates.
    """
    raise NotImplementedError('sorry, still working on "region"')


# common subcommand------------------------------------------------------------
@arg('db', help=db_help)
def common(db):
    """
    Identify child features in common (e.g., common exons across multiple
    isoforms)
    """
    raise NotImplementedError('sorry, still working on "common"')


# create subcommand------------------------------------------------------------
@arg('filename', help='GFF or GTF file to use')
@arg('--output', help='''Database to create.  Default is to append ".db" to the
     end of the input filename''')
@arg('--force', help='''Overwrite an existing database''')
@arg('--quiet', help='''Suppress the reporting of timing information when
     creating the database''')
@arg('--merge', help='''Merge strategy to be used if if duplicate IDs are
     found.''')
@arg('--disable-infer-genes', help='''Disable inferring of gene
     extents for GTF files. Use this if your GTF file already has "gene"
     featuretypes''')
@arg('--disable-infer-transcripts', help='''Disable inferring of transcript
     extents for GTF files. Use this if your GTF file already has "transcript"
     featuretypes''')
def create(filename, output=None, force=False, quiet=False, merge="merge",
           disable_infer_genes=False, disable_infer_transcripts=False):
    """
    Create a database.
    """
    verbose = not quiet
    if output is None:
        output = filename + '.db'
    gffutils.create_db(filename, output,
                       force=force,
                       verbose=verbose,
                       merge_strategy=merge,
                       disable_infer_genes=disable_infer_genes,
                       disable_infer_transcripts=disable_infer_transcripts)


# clean subcommand ------------------------------------------------------------
@arg('filename', help='''GFF or GTF file to use''')
def clean(filename):
    """
    Perform various QC operations to clean a GFF or GTF file.
    """
    raise NotImplementedError('sorry, still working on "clean"')


# sanitize subcommand ---------------------------------------------------------
@arg('filename', help='''GFF or GTF file to use (or GFF database.)''')
@arg('--in-memory', help='''Load GFF into memory for processing.''')
@arg('--in-place',
     help='''Sanitize file in-place: overwrites current file with sanitized
     version.''')
def sanitize(filename,
             in_memory=True,
             in_place=False):
    """
    Sanitize a GFF file. Might get merged with clean feature later.

    Cleans and adds useful annotations to a GFF file:

      - Ensures that start > end in all entries
      - Adds an entry id (eid) to each entry to make files grep-able

    Outputs result to stdout unless asked to sanitize in place.
    """
    print("Sanitizing GFF %s" % filename, file=sys.stderr)
    if in_memory:
        print("  - Loading GFF in memory", file=sys.stderr)
    if in_place:
        print("  - Sanitizing file in place (overwriting current file)",
              file=sys.stderr)
    helpers.sanitize_gff_file(filename,
                              in_memory=in_memory,
                              in_place=in_place)


@arg('filename', help='''GFF or GTF file to use.''')
@arg('--in-place', help='''Remove duplicates in place (overwrite current
     file.)''')
def rmdups(filename, in_place=False):
    """
    Remove duplicates from a GFF file.
    """
    print("Removing duplicates from: %s" % filename)
    merge_strategy = "merge"
    if not in_place:
        # Write to stdout
        gff_out = gffwriter.GFFWriter(sys.stdout)
    else:
        # Write to file in place
        gff_out = gffwriter.GFFWriter(filename, in_place=True)
    # Create database with merge
    db = gffutils.create_db(filename, ":memory:",
                            verbose=False,
                            merge_strategy=merge_strategy)
    # Write out the records
    for rec in db.all_features():
        gff_out.write_rec(rec)
    gff_out.close()


# annotate subcommand ---------------------------------------------------------
@arg('filename', help='''GFF or GTF file to use''')
@arg('genes-table', help='''Genes table in GFF or GTF format to use.''')
def annotate(filename, genes_table):
    """
    Annotate a GFF file with useful information. For now, add annotation
    of gene IDs based on an input GFF annotation of genes.

    Computes the most inclusive transcription start/end coordinates
    for each gene, and then uses pybedtools to intersect (in strand-specific
    manner) with the input annotation.
    """
    raise NotImplementedError('implementation in progress')


# convert subcommand ----------------------------------------------------------
@arg('filename', help='''GFF or GTF file to convert''')
def convert(filename):
    """
    Convert a GTF file to GFF or vice versa.
    """
    raise NotImplementedError('implementation in progress')


# search subcommand -----------------------------------------------------------
@arg('db', help=db_help)
@arg('text', help='''Text to search for. Case-insensitive; use sql LIKE
     syntax''')
@arg('--featuretype', help='''Restrict to a particular featuretype.  This can
     be faster than doing a grep on the output, since it restricts the search
     space in the database''')
def search(db, text, featuretype=None):
    """
    Search the attributes.
    """
    db = gffutils.FeatureDB(db)
    for item in db.attribute_search(text, featuretype):
        yield item


if __name__ == "__main__":
    argh.dispatch_commands([
        fetch,
        children,
        parents,
        region,
        create,
        common,
        clean,
        search,
        sanitize,
        rmdups
    ])