#!/usr/bin/env python
from __future__ import print_function
import sys,os
import re
import gzip
import logging
import argparse
from pprint import pprint

desc = 'Annotate a genome based on mapping to UniRef via diamond'
epi = """DESCRIPTION:
Adding UniRef IDs and taxonomy to genes from a particular genome.
"""
parser = argparse.ArgumentParser(description=desc,
                                 epilog=epi,
                                 formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument('diamond_hits', metavar='diamond_hits', type=str,
                    help='tab-delim table of diamond hits')
parser.add_argument('genes_fasta_nuc', metavar='genes_fasta_nuc', type=str,
                    help='Genes in nucletide fasta format')
parser.add_argument('genes_fasta_AA', metavar='genes_fasta_AA', type=str,
                    help='Genes in amino acid fasta format')
parser.add_argument('taxonomy', metavar='taxonomy', type=str,
                    help='Taxonomy of the genome')
parser.add_argument('taxID', metavar='taxID', type=str,
                    help='NCBI TaxID of the genome')
parser.add_argument('--columns', type=str, default='qseqid,sseqid,pident,length,qstart,qend,qlen,sstart,send,slen,evalue',
                        help='Diamond output columns (default:  %(default)s)')                    
parser.add_argument('--outdir', type=str, default='genes_annotated',
                        help='Output directory (default:  %(default)s)')
parser.add_argument('--dmnd-db', type=str, default='/ebio/abt3_projects2/databases_no-backup/humann2/uniref50/uniref50_annotated.1.1.dmnd',
                    help='UniRef dmnd db for annotating genes (default: %(default)s)')
parser.add_argument('--percid', type=float, default=50.0,
                        help='Percent sequence ID cutoff for calling a hit (default:  %(default)s)')
parser.add_argument('--overlap', type=float, default=80.0,
                        help='Perc. overlap cutoff (longest sequence) for calling a hit (default:  %(default)s)')
parser.add_argument('--skip', action='store_true', default=False,
                        help='Skip diamond-based annotation if the diamond hits file exists (default:  %(default)s)')
parser.add_argument('--gzip', action='store_true', default=False,
                        help='gzip output (default:  %(default)s)')
parser.add_argument('--threads', type=int, default=1,
                       help='Threads used for diamond (default:  %(default)s)')
parser.add_argument('--version', action='version', version='0.0.1')

logging.basicConfig(format='%(asctime)s - %(message)s', level=logging.DEBUG)


def make_best_hit_index(dmnd_hit_file, outfmt_cols):
    """Making index of best diamond hits for each query
    """
    # column index
    column_idx = {x:i for i,x in enumerate(outfmt_cols)}
    # streaming hits
    hits = {}
    with open(dmnd_hit_file) as inF:
        longest_seq_len = None
        for line in inF:
            line = line.rstrip().split('\t')
            # Perc ID >= cutoff?
            try:
                pident = float(line[column_idx['pident']])
            except KeyError:
                raise KeyError('Cannot find "pident" column')
            if pident < args.percid:
                continue
            # overlap >= cutoff?
            ## longest sequence?
            try:
                query_len = float(line[column_idx['qlen']])
            except KeyError:
                raise KeyError('Cannot find "qlen" column')
            try:
                subject_len = float(line[column_idx['slen']])
            except KeyError:
                raise KeyError('Cannot find "slen" column')
            if query_len >= subject_len:
                longest_seq_len = query_len
            else:
                longest_seq_len = subject_len
            ## overlap vs longest sequence
            try:
                aln_len = float(line[column_idx['length']])
            except KeyError:
                raise KeyError('Cannot find "length" column')
            perc_overlap = aln_len / longest_seq_len * 100.0
            if perc_overlap < args.overlap:
                continue
            # better than current best hit for query?
            try:
                qseqid = line[column_idx['qseqid']]
            except KeyError:
                raise KeyError('Cannot find "qseqid" column')
            try:
                sseqid = line[column_idx['sseqid']]
            except KeyError:
                raise KeyError('Cannot find "sseqid" column')
            try:
                best_hit = hits[qseqid]
            except KeyError:
                best_hit = None
            if best_hit is None or (pident >= best_hit[1] and perc_overlap >= best_hit[2]):
                hits[qseqid] = [sseqid, pident, perc_overlap]
    return hits

def rename_seqs(best_hits, fasta_file, taxonomy, outfile, gzip_output=False):
    """Renaming sequences based on uniref hits.
    Using naming format: `gene_family|gene_length|taxonomy`
    Taxonomy format: `g__{genus};s__{species}_taxID{taxID}`
    (see https://bitbucket.org/biobakery/humann2/wiki/Home).
    """
    seq_name = None
    seq = ''
    annot_cnt = 0
    annot_skip_cnt = 0
    if gzip_output == True:
        _open = lambda x: gzip.open(x, 'ab')
        outfile += '.gz'
    else:
        _open = lambda x: open(x, 'a')
    
    with open(fasta_file) as inF, _open(outfile) as outF:
        for line in inF:
            if line.startswith('>'):
                # previous sequence
                if seq_name is not None and seq != '':
                    seq = seq.rstrip().strip('*')
                    x = '\n'.join(['>' + seq_name, seq]) + '\n'
                    if gzip_output == True:
                        x = x.encode()   
                    outF.write(x)
                    annot_cnt += 1
                else:
                    annot_skip_cnt += 1
                seq = ''
                # hit for sequence?
                query = line.rstrip().lstrip('>').split(' ')[0]
                try:
                    best_hit = best_hits[query]
                except KeyError:
                    best_hit = None
                # renaming
                if best_hit is None:
                    seq_name = None
                else:
                    seq_name = '|'.join([best_hit[0], taxonomy])
            else:
                seq += line.rstrip()
        # final sequence
        if seq_name is not None:
            seq = seq.rstrip().strip('*')
            x = '\n'.join(['>' + seq_name, seq]) + '\n'
            if gzip_output == True:
                x = x.encode()
            outF.write(x)
            annot_cnt += 1
        else:
            annot_skip_cnt += 1

    logging.info('Number of genes with an annotation: {}'.format(annot_cnt))            
    logging.info('Number of genes skipped due to no annotation: {}'.format(annot_skip_cnt))  
    logging.info('File written: {}'.format(outfile))
     
def format_taxonomy(tax, taxID):
    """Formatting taxonomy string
    """
    logging.info('Taxonomy string provided {}'.format(tax))
    logging.info('TaxID provided {}'.format(taxID))
    
    try:
        taxID = int(float(taxID.strip()))
    except ValueError:
        msg = 'ERROR: taxID "{}" is not an integer!'
        raise ValueError(msg)
    tax = re.sub('[^A-Za-z0-9-_;]+', '_', tax).split(';')
    
    if not len(tax) == 7:
        species = 's__unclassified'
    else:
        species = tax[6]
    if not len(tax) >= 6:
        genus = 'g__unclassified'
    else:
        genus = tax[5]

    if genus.startswith('G__'):
        genus = genus[3:]
    if not genus.startswith('g__'):
        genus = 'g__' + genus

    if species.startswith('S__'):
        species = genus[3:]
    if not species.startswith('s__'):
        species = 'g__' + species

    if genus == 'g__':
        genus = 'g__unclassified'
    if species == 's__':
        species = 's__unclassified'
        
    tax = '.'.join([genus, species])
    tax += '__taxID{}'.format(taxID)
    logging.info('Converted taxonomy string to {}'.format(tax))
    return tax

def main(args):
    # formatting taxonomy
    args.taxonomy = format_taxonomy(args.taxonomy, args.taxID)
    
    # filtering diamond hits
    logging.info('Finding best hit for each gene')
    outfmt_cols = args.columns.split(',')
    best_hits = make_best_hit_index(args.diamond_hits, outfmt_cols)
    #pprint(hits)

    logging.info('Renaming genes')
    # nuc
    outfile = os.path.join(args.outdir, 'annot.fna')
    rename_seqs(best_hits, args.genes_fasta_nuc,
                args.taxonomy, outfile=outfile,
                gzip_output=args.gzip)
    # AA
    outfile = os.path.join(args.outdir, 'annot.faa')
    rename_seqs(best_hits, args.genes_fasta_AA,
                args.taxonomy, outfile=outfile,
                gzip_output=args.gzip)

    
if __name__ == '__main__':
    args = parser.parse_args()
    main(args)
