#!/usr/bin/python
###############################################################################
#
# amplishot.py - Pipeline for generating full-length 16S sequences and
#                performing community abundance measurements on the output
#
###############################################################################
#                                                                             #
#    This program is free software: you can redistribute it and/or modify     #
#    it under the terms of the GNU General Public License as published by     #
#    the Free Software Foundation, either version 3 of the License, or        #
#    (at your option) any later version.                                      #
#                                                                             #
#    This program is distributed in the hope that it will be useful,          #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of           #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the            #
#    GNU General Public License for more details.                             #
#                                                                             #
#    You should have received a copy of the GNU General Public License        #
#    along with this program. If not, see <http://www.gnu.org/licenses/>.     #
#                                                                             #
###############################################################################

__author__ = "Connor Skennerton"
__copyright__ = "Copyright 2012-2014"
__credits__ = ["Connor Skennerton"]
__license__ = "GPL3"
__version__ = "0.10.0"
__maintainer__ = "Connor Skennerton"
__email__ = "c.skennerton@gmail.com"
__status__ = "Development"

###############################################################################
import logging
import argparse
import sys
import os
import subprocess
import gzip
from collections import defaultdict
import amplishot.parse.fastx
from amplishot import taxon_segregator, otu_table, assign_taxonomy, assemble,\
        OrientationChecker
from amplishot.app import bowtie
from amplishot.config import AmplishotConfig, AmplishotConfigError
from amplishot.app.cd_hit import CD_HIT_EST
from amplishot.util import reverse_complement
from qiime.align_seqs import alignment_module_names,alignment_method_constructors,\
    pairwise_alignment_methods, CogentAligner, compute_min_alignment_length
from qiime.identify_chimeric_seqs import chimeraSlayer_identify_chimeras
###############################################################################

def align_seqs(config, input_fasta_fp):

    try:
        template_alignment = config.data['aligner_template']
        if template_alignment is None:
            raise KeyError()
    except KeyError:
        logging.warning('\'aligner_termplate\' not specified in the '\
                'configuration file or on the command line, but is required for' \
                ' the alignment step. PyNAST will not be executed and the' \
                ' reconstructed sequences will not be aligned or chimera checked')
        raise AmplishotConfigError()

    logging.info("Aligning rep set...")
    input_seqs_filepath = input_fasta_fp
    alignment_method = 'PyNAST'
    output_dir = os.path.join(config.data['output_directory'],
            'repset_aligned')
    if not os.path.exists(output_dir):
        os.mkdir(output_dir)

    # compute the minimum alignment length if a negative value was
    # provided (the default)
    min_length = config.data['minimum_reconstruction_length']


    result_path = os.path.join(output_dir, "repset_aligned.fasta")
    log_path = os.path.join(output_dir, "repset_aligned_log.txt")
    failure_path = os.path.join(output_dir, "repset_failures.fasta")

    # try/except was causing problems here, so replacing with
    # an explicit check
    # define the aligner params
    aligner_constructor =\
    alignment_method_constructors['pynast']
    aligner_type = alignment_method
    params = {'min_len': min_length,\
              'min_pct': 0.6,\
              'template_filepath': template_alignment,\
              'blast_db': None,
              'pairwise_alignment_method': 'blast'}
    # build the aligner object
    aligner = aligner_constructor(params)
    # apply the aligner
    aligner(input_seqs_filepath,result_path=result_path,\
     log_path=log_path,failure_path=failure_path)

    return result_path


def check_chimeras(config, input_filepath):
    #output_fp = os.path.join(config.data['output_directory'],
    #        'repset_aligned', 'chimeras.txt')

    return chimeraSlayer_identify_chimeras(input_filepath,
                                 output_fp=None,
                                 db_NAST_fp=config.data['aligner_template'],
                                 min_div_ratio=None,
                                 keep_intermediates=False)



def map_to_reference(params, sample_mapping_file):
    """Take overlapped reads and map them to reference sequences
       returns path to output file containing sam formated alignments
    """
    mapper_cmd = ['bowtie2']
    for k,v in params.items():
        mapper_cmd.extend([str(k),str(v)])

    samtools_cmd = ['samtools view -SubF 0x904 - | samtools sort -']
    samtools_cmd.append(sample_mapping_file)
    samtools_cmd.extend(['&& samtools index', sample_mapping_file + '.bam'])

    final_cmd = ' '.join(mapper_cmd) + ' | ' + ' '.join(samtools_cmd)
    logging.debug(final_cmd)

    with open(os.devnull, 'w') as devnull:
        retcode = subprocess.call(final_cmd, shell=True, stderr=devnull)
        if retcode != 0:
            raise RuntimeError("failed to create bam files")


def taxon_partition(samfile, taxonomyFile, rootdir='root'):
    """Call the taxon segregator and return a list of tuples of taxonomies that
       have adequate coverage for assembly
    """
    taxseg = amplishot.taxon_segregator.TaxonSegregator(taxonomyFile)
    taxseg.parse_sam(samfile)
    return taxseg.segregate(root=rootdir)


def overlap(forward, reverse, pairtig_dir, minPairtigLen=350, minOverlap=30,
        threads=1):
    pear_basename =  os.path.join(pairtig_dir,
           os.path.splitext(os.path.basename(forward))[0])
    # pear doesn't like gziped files therefore unzip them at teh start
    gzip_cmd = 'gzip -dc %s > %s'
    if forward[-2:] == 'gz':
        tmp_forward = os.path.join(pairtig_dir, '__ft.fastq')

        retcode = subprocess.call(gzip_cmd % (forward, tmp_forward), shell=True)
        if retcode != 0:
            raise RuntimeError("unzipping the forwad reads was unsuccessful")
        forward = tmp_forward

    if reverse[-2:] == 'gz':
        tmp_reverse = os.path.join(pairtig_dir, '__rt.fastq')
        retcode = subprocess.call(gzip_cmd % (reverse, tmp_reverse), shell=True)
        if retcode != 0:
            raise RuntimeError("unzipping the reverse reads was unsuccessful")
        reverse = tmp_reverse

    pear_cmd = "pear -f %s -r %s -o %s -n %d -v %d -j %d" % (forward, reverse,
           pear_basename, minPairtigLen, minOverlap, threads )
    with open(os.devnull, 'wb') as devnull:
        retcode = subprocess.call(pear_cmd, shell=True, stdout=devnull, stderr=devnull)
        if retcode != 0:
            raise RuntimeError("pear did not exit correctly: %d.\n%s\n" %(retcode, pear_cmd))

    # clean up the temp files
    if forward == os.path.join(pairtig_dir, '__ft.fastq'):
        os.remove(forward)
    if reverse == os.path.join(pairtig_dir, '__rt.fastq'):
        os.remove(reverse)

    return pear_basename + '.assembled.fastq'


def make_aliases(config):
    indata = []
    config.data['aliases'] = []
    if 'skip_pairtigs' not in config.data:
        for readfile in config.data['pairtig_read_files']:
            sample_file_name =\
            os.path.splitext(os.path.basename(readfile))[0]
            indata.append([readfile, sample_file_name])
            config.data['aliases'].append(sample_file_name)
    else:
        for r1, r2 in config.data['input_raw_reads']:
            sample_file_name =\
            os.path.splitext(os.path.basename(r1))[0]
            indata.append(([r1,r2], sample_file_name))
            config.data['aliases'].append(sample_file_name)

    return indata


def orientate(fullLengthSeqs):
    cs = OrientationChecker.OrientationChecker()
    for name in fullLengthSeqs.keys():
        if cs.isSeqReversed(fullLengthSeqs[name]):
            fullLengthSeqs[name] = reverse_complement(fullLengthSeqs[name])

    return fullLengthSeqs


def generate_pairtigs(config, root_dir):
    if len(config.data['pairtig_read_files']) != len(config.data['input_raw_reads']):
        pairtig_dir = os.path.join(root_dir, 'pairtigs')
        if not os.path.exists(pairtig_dir):
            os.mkdir(pairtig_dir)

        logging.info('generating pairtigs...')
        # lets not fork unnessessary processes
        panda_processes = len(config.data['input_raw_reads']) if\
        len(config.data['input_raw_reads']) < config.data['threads'] else\
        config.data['threads']

        panda_pool_results = []
        for f, r in config.data['input_raw_reads']:
            r = overlap(f, r, pairtig_dir, minPairtigLen=config.data['minimum_pairtig_length'],
                minOverlap=config.data['pair_overlap_length'],
                threads=config.data['threads'])
            panda_pool_results.append(r)
        for pairtig_file in panda_pool_results:
            config.data['pairtig_read_files'].append(os.path.abspath(pairtig_file))


def pair_reads(config, root_dir):
    # generate fragment consensus sequences - pandaseq
    if 'skip_pairtigs' in config.data:
        return

    for i in config.data['input_raw_reads']:
        if not isinstance(i, list):
            return

    generate_pairtigs(config, root_dir)


def fixup_aliases(config):
    if 'aliases' in config.data:
        if 'skip_pairtigs' not in config.data:
            if len(config.data['aliases']) != len(config.data['input_raw_reads']):
                logging.warn('There are a different number of aliases and input'\
                    'files. Aliases will not be used')
                indata = make_aliases(config)
            else:
                indata = zip(config.data['pairtig_read_files'], config.data['aliases'])
        else:
            indata = zip(config.data['input_raw_reads'], config.data['aliases'])
    else:
        indata = make_aliases(config)

    return indata


def segregate_and_assemble(config, indata, root_dir):
    # taxonomy assigner
    taxseg = amplishot.taxon_segregator.TaxonSegregator(config.data['taxonomy_file'],
            cutoffs=config.data['mapping_similarity_cutoffs'],
            neighboursfile=config.data['neighbours_file'])

    # assembly wrapper
    wrapper = assemble.AssemblyWrapper(config.data['assembly_method'], config)
    for readfile, samplename in indata:
        logging.info('processing %s...', readfile)
        sample_mapping_file = os.path.join(root_dir, samplename + '.initial_mapping')
        bowtie_params = {
            '-x': config.data['mapper_database'],
            '-p': config.data['threads'],
            #'-u': 10000,
            }

        if isinstance(readfile, list):
            bowtie_params['-1'] = readfile[0]
            bowtie_params['-2'] = readfile[1]
        else:
            bowtie_params['-U'] = readfile

        map_to_reference(bowtie_params, sample_mapping_file)

        # partition dataset based on mapped reads
        logging.info('partitioning...')
        taxon_root = os.path.join(root_dir, samplename, 'root')
        taxseg.parse_sam3(sample_mapping_file + '.bam')

        output_fastq = False
        output_qual = True
        output_fasta = True
        output_pairs = False
        sep12 = False
        if config.data['assembly_method'] != 'phrap':
            output_fastq = True
            output_qual = False
            output_fasta = False
            output_pairs = True
            sep12 = True

        good_taxons = taxseg.segregate2(root=taxon_root,
                minCount=config.data['taxon_coverage'][1],
                minCoverage=config.data['taxon_coverage'][0],
                sam=False, fastq=output_fastq, fasta=output_fasta,
                qual=output_qual, pairs=output_pairs, sep12=sep12)
        taxseg.clear()

        if len(good_taxons) < 1:
            logging.info("No taxons for assembly")
        else:
            logging.info('There are %i taxons suitable for assembly',
                    len(good_taxons))

            if good_taxons.has_key(()) and not config.data['assemble_unknowns']:
                del good_taxons[()]

            wrapper(good_taxons, samplename)

    return wrapper

def pipeline(config):
    """Run all the steps of amplishot
    """
    root_dir = config.data['output_directory']

    # generate fragment consensus sequences
    pair_reads(config, root_dir)

    # map overlapped fragments with bowtie
    # create the file for all the full-length seqs
    full_length_seqs =\
    open(os.path.join(root_dir,
        config.data.get('reconstruced_seq_file', 'full_length_sequences.fa')), 'w')

    indata = fixup_aliases(config)

    wrapper = segregate_and_assemble(config, indata, root_dir)

    # check to see if there are any full length sequences
    # if not we finish here
    if len(wrapper.fullLengthSeqs) < 1:
        logging.warn("No 16S reconstructions could be made")
    else:
        # fix up the orientation of the full-length sequences to they are all in
        # the forward direction then write them to file
        logging.info("Orientating full-length sequences...")
        wrapper.fullLengthSeqs = orientate(wrapper.fullLengthSeqs)
        wrapper.write(full_length_seqs)

        full_length_seqs.close()
        # pick OTUs and representative set
        out_otu_map = amplishot.otu_table.pick_otus(full_length_seqs.name,
                outputFileName=os.path.join(root_dir, 'full_length_otus.txt'),
                similarity=config.data['otu_clustering_similarity'],
                otuPickingMethod=config.data['otu_clustering_method'])

        rep_set = amplishot.otu_table.pick_rep_set(full_length_seqs.name, out_otu_map,
                outputFilePath=os.path.join(root_dir,
                    config.data.get('repset_output_file',
                        'full_length_sequences.repset.fa')))

        # align the repset and call chimeras
        chimera_ids = set()
        try:
            aligned_repset = align_seqs(config, rep_set)
            logging.info("Checking for chimeras in rep set...")
            chimeras = check_chimeras(config, aligned_repset)
            logging.info('there are %d chimeras found' % len(chimeras))
            if len(chimeras) > 0:
                logging.info('removing chimeras from rep set...')
                for _chim in chimeras:
                    #logging.info("\t%s" % str(_chim))
                    chimera_ids.add(_chim[0])
        except AmplishotConfigError:
            pass

        rep_set_seqs = dict()
        with open(rep_set) as fp:
            fxparser = amplishot.parse.fastx.FastxReader(fp)
            for name, seq, qual in fxparser.parse(
                    callback=amplishot.parse.fastx.not_include,
                    headers=chimera_ids):
                rep_set_seqs[name] = seq
        with open(rep_set, 'w') as fp:
            for name, seq in rep_set_seqs.items():
                fp.write('>%s\n%s\n' % (name, seq))

        logging.info('There are %i OTUs, assigning taxonomy...', len(rep_set_seqs))

        #assign taxonomy
        taxonomy_map = assign_taxonomy.assign_taxonomy(config.data['assign_taxonomy_method'],
                rep_set, config)

        observation_metadata = defaultdict(dict)
        score_field_name = 'confidence'

        if config.data['assign_taxonomy_method'] == 'blast':
            score_field_name = 'evalue'
        elif config.data['assign_taxonomy_method'] == 'bowtie':
            score_field_name = 'similarity'

        for otu, metadata in taxonomy_map.items():
            observation_metadata[otu]['taxonomy'] = metadata[0].split('; ')
            observation_metadata[otu][score_field_name] = metadata[1]

        logging.info('generating OTU table in biom format...')

        # create Bowtie database from representative set
        abundance = amplishot.otu_table.OTUTableGenerator(rep_set_seqs,
                root_dir)
        #generate abundance based on coverage and create OTU table
        for read_file, sample_name in indata:
            _params = {'-p': config.data['threads']}
            if isinstance(read_file, list):
                _params['-1'] = read_file[0]
                _params['-2'] = read_file[1]
            else:
                _params['-U'] = read_file
            abundance.generate_abundance(sample_name,
                    params=_params, tmpFile=False)

        abundance.generate_biom_table()
        abundance.add_metadata(observation_metadata=observation_metadata)
        with open(os.path.join(root_dir, 'abundances.biom'), 'w') as fp:
            fp.write(str(abundance))

        abundance.normalize()
        with open(os.path.join(root_dir, 'abundances.normalized.biom'), 'w') as fp:
            fp.write(str(abundance))

    logging.info('Done.')


###############################################################################
###############################################################################
###############################################################################
###############################################################################
class CommaStringToList(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        l = values.split(',')
        setattr(namespace, self.dest, map(float, l))


if __name__ == '__main__':
    config = AmplishotConfig()

    parser = argparse.ArgumentParser( formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument('--version', action='version', version='%(prog)s '+__version__)
    parser.add_argument('-c', '--config-file', dest='config',
            #default=os.path.join(os.path.expanduser('~'),'.amplishot_config.yml'),
            help='specify a custom location to the amplishot configuration\
            file. By default all values stored in the config file will\
            be used unless they are overwritten by command-line options')
    parser.add_argument('--quiet', dest='logger_quiet',
            action='store_const', const='CRITICAL', help='change the log verbosity so that only'\
                    'the most critical messages are reported')

    parser.add_argument('--log_level', choices=['DEBUG', 'INFO',
        'WARNING', 'ERROR', 'CRITICAL'], dest='log_level', default='INFO',
        help='Set the log level')

    parser.add_argument('-t', '--threads', type=int, dest='threads', default=1,
             help='The number of threads to use')

    parser.add_argument('-T','--taxonomy_file', #required=True,
    dest='taxonomy_file', help="File\
            containing the taxonomy mapping of the reference sequences in the\
                    bowtie database")

    parser.add_argument('-d', '--mapper_database', #required=True,
    dest='mapper_database', help="the\
            path to the bowtie database, without the suffixes; just like you\
            would write when using bowtie")

    parser.add_argument('-i', '--input_raw_reads', nargs=2, action='append',
            help='paths to first and pecond pairs of reads for \'pairtig\'\
            generation.  Can be specified multiple times for many samples',
            dest='input_raw_reads')

    parser.add_argument('-l', '--minimum_pairtig_length',
            dest='minimum_pairtig_length', default=350,
            help='The minimal acceptable length for\
            assembled pairs')
    parser.add_argument('-L', '--minimum_reconstruction_length', default=1000,
            help='minimum length for 16S assemblies to be included in final \
            set')

    parser.add_argument('-M', '--mapping_similarity_cutoffs',
            default=[0.88,0.92,0.94,0.98],
            action=CommaStringToList,
            help="A series of percentages (as a number between 0 and 1) that \
            are comma separated used to indicate a binning pattern for reads")

    parser.add_argument('-O', '--pair_overlap_length',
            dest='pair_overlap_length', type=int, default=30,
            help='the minimum length that pairs must overlap by')

    parser.add_argument('-o', '--output_directory', default='./',
             dest='output_directory',
            help='Root directoy under which all files and folders from the\
            pipeline will be generated from')

    parser.add_argument('-A', '--aligner_template',
            default=config.data.get('aligner_template', None),
            help='File path for prealigned 16S sequences from NAST')

    parser.add_argument('-s', '--otu_clustering_similarity', type=float,
            dest='otu_clustering_similarity', default=0.99,
            help='The similarity'\
                    ' required between sequences AFTER reconstruction'\
                    ' for OTU clustering')

    parser.add_argument('-p', '--otu_picking_method', choices=['cdhit',
        'uclust'], default='cdhit', help='The program used to cluster '\
            'sequences into OTUs')

    parser.add_argument('-R','--taxon_coverage',nargs=2, default=[2, 1000],
            metavar='INT',
            help='list of two numbers. The first is the minimum coverage, the \
            second is the number of bases that need to be covered')

    parser.add_argument('-a', '--assembly_method', default='ray',
            choices=['ray', 'phrap', 'velvet'], help='de novo assembly method \
            used for reconstruction')
    parser.add_argument('-b', '--assign_taxonomy_method', default='blast',
            choices=['blast'], help='method for assigning taxonomy to \
            full-length reconstructions')
    parser.add_argument('-n', '--neighbours_file', help='file containing OTUs \
            that are within X percentage, phylogenetic distance to each other')

    parser.add_argument('-B', '--taxonomy_blast_db', dest='blast_db',
            help='input file path for blast formatted database for assigning \
            taxonomy using blast')

    args = parser.parse_args()

    try:
        write_new_config = config.check_config_and_set_output(args)
    except AmplishotConfigError, e:
        print e.msg
    else:
        if args.logger_quiet is not None:
            config.data['log_level'] = args.logger_quiet

        numeric_level = getattr(logging, config.data['log_level'], None)
        if not isinstance(numeric_level, int):
            raise ValueError('Invalid log level: %s' % config.data['log_level'])

        logging.basicConfig(level=numeric_level,
                format='%(levelname)s\t%(asctime)s\t%(message)s')

        #if config.data['log_file'] is not None:
        #    logging.basicConfig(filename=config.data['log_file'])

        pipeline(config)
        if write_new_config:
            config.write_config()


###############################################################################
###############################################################################
###############################################################################
###############################################################################
