#!/usr/bin/env python
import collections
import subprocess
import itertools
import argparse
import tempfile
import logging
import os.path
import sqlite3
import atexit
import shutil
import errno
import shlex
import csv

from taxtastic.refpkg import Refpkg
from Bio import SeqIO

log = logging.getLogger(__name__)

def logging_check_call(cmd, *a, **kw):
    log.info(' '.join(cmd))
    return subprocess.check_call(cmd, *a, **kw)

def silently_unlink(path):
    try:
        os.unlink(path)
    except OSError, e:
        if e.errno != errno.ENOENT:
            raise

input_suffixes = {'align-each': 'fasta', 'merge-each': 'sto', 'none': 'sto'}

def main():
    logging.basicConfig(
        level=logging.INFO, format="%(levelname)s: %(message)s")

    parser = argparse.ArgumentParser(
        description="Classify query sequences using a hrefpkg.")
    parser.add_argument('hrefpkg', help="hrefpkg to classify using")
    parser.add_argument('query_seqs', help="input query sequences")
    parser.add_argument('classification_db', help="output sqlite database")
    parser.add_argument('-j', '--ncores', default=1, type=int, metavar='CORES',
                        help="number of cores to tell commands to use")
    parser.add_argument('-r', '--classification-rank', metavar='RANK',
                        help="rank to perform the initial NBC classification at")
    parser.add_argument('--classifier', default='pplacer',
                        help="which classifier to use with guppy classify")
    parser.add_argument('--pplacer-args', default=[], type=shlex.split,
                        help="additional arguments for pplacer")
    parser.add_argument('--classification-args', default=[], type=shlex.split,
                        help="additional arguments for guppy classification")
    parser.add_argument('--post-prob', default=False, action='store_true',
                        help="place with posterior probabilities")
    parser.add_argument('--workdir', metavar='DIR',
                        help=("directory to write intermediate files to "
                              "(default: a temporary directory)"))
    parser.add_argument('--disable-cleanup', default=False, action='store_true',
                        help="don't remove the work directory as the final step")
    parser.add_argument('--use-mpi', default=False, action='store_true',
                        help="run refpkg_align with MPI")
    parser.add_argument(
        '--alignment', choices=['align-each', 'merge-each', 'none'], default='align-each',
        help=("respectively: align each input sequence; subset an input stockholm "
              "alignment and merge each sequence to a reference alignment; only subset "
              "an input stockholm alignment (default: align-each)"))
    parser.add_argument(
        '--cmscores', type=argparse.FileType('w'), metavar='FILE',
        help="in align-each mode, write out a file containing the cmalign scores")

    programs = parser.add_argument_group("external binaries")
    programs.add_argument(
        '--pplacer', default='pplacer', metavar='PROG', help="pplacer binary to call")
    programs.add_argument(
        '--guppy', default='guppy', metavar='PROG', help="guppy binary to call")
    programs.add_argument(
        '--rppr', default='rppr', metavar='PROG', help="rppr binary to call")
    programs.add_argument(
        '--refpkg-align', default='refpkg_align.py', metavar='PROG',
        help="refpkg_align binary to call")
    programs.add_argument(
        '--cmalign', default='cmalign', metavar='PROG', help="cmalign binary to call")


    args = parser.parse_args()

    if args.ncores < 0:
        args.error('ncores must be >= 0')

    mpi_args = []
    if args.use_mpi and args.ncores > 0:
        mpi_args = ['--use-mpi', '--mpi-arguments', '-np %d' % (args.ncores,)]

    if args.workdir is None:
        workdir = tempfile.mkdtemp()
    else:
        workdir = args.workdir
        try:
            os.makedirs(workdir)
        except OSError, e:
            if e.errno != errno.EEXIST:
                raise

    if not args.disable_cleanup:
        @atexit.register
        def cleanup_workdir():
            shutil.rmtree(workdir, ignore_errors=True)

    classif_db = os.path.join(workdir, 'classifications.sqlite')
    index_refpkg = os.path.join(args.hrefpkg, 'index.refpkg')
    index = Refpkg(index_refpkg)
    index_rank = index.metadata('index_rank')
    classif_rank = args.classification_rank or index_rank
    index_counts = os.path.join(args.hrefpkg, 'index-%s.counts' % (classif_rank,))
    log.info('performing initial classification at %s', classif_rank)
    silently_unlink(classif_db)
    logging_check_call(
        [args.rppr, 'prep_db', '--sqlite', classif_db, '-c', index_refpkg])
    logging_check_call(
        [args.guppy, 'classify', '--sqlite', classif_db, '-c', index_refpkg,
         '--classifier', 'nbc', '--nbc-rank', classif_rank, '--no-pre-mask',
         '--nbc-sequences', args.query_seqs, '--nbc-counts', index_counts,
         '-j', str(args.ncores)])

    with open(os.path.join(args.hrefpkg, 'index.csv'), 'rU') as fobj:
        refpkg_map = dict(csv.reader(fobj))

    log.info('determining bins at %s', index_rank)
    conn = sqlite3.connect(classif_db)
    curs = conn.cursor()
    curs.execute("""
        SELECT name,
               tax_id
        FROM   multiclass
        WHERE  want_rank = ?
               AND rank = want_rank
    """, (index_rank,))
    seq_bins = dict(curs)
    all_bins = set(seq_bins.itervalues()) & set(refpkg_map)
    log.info('classified into %d bins', len(all_bins))
    log.debug('bins: %r', all_bins)
    bin_outputs = {}
    bin_counts = collections.Counter()

    input_suffix = input_suffixes[args.alignment]
    def binfile(bin):
        return os.path.join(workdir, '%s.%s' % (bin, input_suffix))

    for bin in all_bins:
        bin_outputs[bin] = open(binfile(bin), 'w')

    if args.alignment == 'align-each':
        for seq in SeqIO.parse(args.query_seqs, 'fasta'):
            bin = seq_bins.get(seq.id)
            if bin is None or bin not in bin_outputs:
                continue
            SeqIO.write([seq], bin_outputs[bin], 'fasta')
            bin_counts[bin] += 1

    # Gross ad-hoc stockholm parser. This is only necessary because biopython
    # can't handle the metadata that cmalign adds to the file.
    else:
        def write_all(lines):
            for line in lines:
                for fobj in bin_outputs.itervalues():
                    fobj.write(line)

        with open(args.query_seqs, 'r') as fobj:
            # Copy the stockholm header,
            write_all(itertools.islice(fobj, 3))

            # then partition the sequences,
            for line in fobj:
                # (stopping if we reach the metadata at the end (hopefully it's at
                # the end))
                if line.startswith('#=GC'):
                    break
                id, seq = line.split()
                bin = seq_bins.get(id)
                if bin is None or bin not in bin_outputs:
                    continue
                bin_outputs[bin].write(line)
                bin_counts[bin] += 1

            # and then copy the footer.
            write_all([line])
            write_all(fobj)

    for fobj in bin_outputs.itervalues():
        fobj.close()

    # This is a bit nasty but I think it's Good Enough. We can change it if we
    # encounter problems or need to be more portable.
    if args.cmscores is not None:
        cmscores_proc = subprocess.Popen(
            ['romp', 'cmscores', '/dev/stdin'],
            stdin=subprocess.PIPE, stdout=args.cmscores)
        cmscores_args = [
            '--alignment-method', 'INFERNAL',
            '--stdout', '/dev/fd/%d' % (cmscores_proc.stdin.fileno(),)]
        cmscores_stdout = cmscores_proc.stdout
    else:
        cmscores_proc = None
        cmscores_args = []
        cmscores_stdout = open(os.devnull, 'w')

    for e, bin in enumerate(all_bins):
        log.info('classifying bin %s (%d/%d; %d seqs)',
                 bin, e + 1, len(all_bins), bin_counts[bin])
        input = binfile(bin)
        aligned = os.path.join(workdir, bin + '-aligned.sto')
        placed = os.path.join(workdir, bin + '.jplace')
        refpkg = os.path.join(args.hrefpkg, refpkg_map[bin])

        if args.alignment == 'align-each':
            logging_check_call(
                [args.refpkg_align, 'align', '--output-format', 'stockholm',
                 refpkg, input, aligned] + mpi_args + cmscores_args)
        elif args.alignment == 'merge-each':
            refpkg_obj = Refpkg(refpkg)
            logging_check_call(
                [args.cmalign, '--merge', '-o', aligned, '-1', '--hbanded',
                 '--sub', '--dna', refpkg_obj.resource_path('profile'),
                 refpkg_obj.resource_path('aln_sto'), input],
                stdout=cmscores_stdout)
        elif args.alignment == 'none':
            raise NotImplementedError('none')

        classifier_args, pplacer_args = [], []
        if args.classifier.startswith('nbc') or args.classifier.startswith('hybrid'):
            if '--no-pre-mask' not in args.classification_args:
                classifier_args.append('--no-pre-mask')
            classifier_args.extend(['--nbc-sequences', input])
        if args.post_prob:
            classifier_args.append('--pp')
            pplacer_args.extend(
                ['-p', '--inform-prior', '--prior-lower', '0.01'])

        logging_check_call(
            [args.pplacer, '--discard-nonoverlapped', '-c', refpkg,
             '-j', str(args.ncores), aligned, '-o', placed]
            + args.pplacer_args
            + pplacer_args)

        logging_check_call(
            [args.guppy, 'classify', '--sqlite', classif_db, '-c', refpkg, placed,
             '-j', str(args.ncores), '--classifier', args.classifier]
            + args.classification_args
            + classifier_args)

    if cmscores_proc is not None:
        cmscores_proc.communicate()

    log.info('cleaning up `multiclass` table')
    conn.rollback()
    curs = conn.cursor()
    curs.execute("""
        CREATE TEMPORARY TABLE classified AS
          SELECT name,
                 COUNT(DISTINCT run_id) n_runs
            FROM multiclass
                 JOIN placements USING (placement_id)
           GROUP BY name
          HAVING n_runs > 1
    """)
    curs.execute("SELECT name FROM classified WHERE n_runs > 2")
    too_many_classifications = [name for name, in curs]
    if too_many_classifications:
        raise ValueError("some sequences got classified more than twice",
                         too_many_classifications)
    curs.execute("""
        DELETE FROM multiclass
         WHERE (SELECT run_id
                  FROM placements p
                 WHERE p.placement_id = multiclass.placement_id) = 1
           AND name IN (SELECT name
                          FROM classified)
    """)
    conn.commit()

    shutil.copy(classif_db, args.classification_db)

main()