#!/usr/bin/env python3
"""Generate OGU tables from read-to-reference maps.
"""

from os import listdir
from os.path import join, isfile, splitext
import gzip
import bz2
import lzma
import argparse


description = """example:
  ogu_from_maps.py bowtie2_result_dir output_name \\
    -m bowtie2 -e .sam.xz \\
    -s sample_ids_to_include.txt \\
    -t sequence_id_to_genome_id.txt
"""

epilog = """output OGU tables:
  - all: sum of all hits per genome.
  - norm: sum of hits per genome; if one hit is shared by k genomes, each
    genome receives 1/k hits.
  - uniq: sum of unique hits per genome; if one hit is shared by multiple
    genomes, it will be dropped.
"""


def main():
    # arguments
    parser = argparse.ArgumentParser(
        description=description,
        epilog=epilog,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('input_dir',
                        help='input directory hosting maps')
    parser.add_argument('output_name',
                        help='stem filename for output OGU tables')
    parser.add_argument('-m', dest='method',
                        help=('method that generated the maps, options: '
                              'Bowtie2, Centrifuge or leave empty, the last '
                              'of which applies to BLAST, DIAMOND, BURST and '
                              'other methods which generate a simple `query '
                              '<tab> subject <tab> ...` format)'))
    parser.add_argument('-t', dest='translation',
                        help='subject ID to genome ID translation table')
    parser.add_argument('-s', dest='samples',
                        help='sample ID list')
    parser.add_argument('-e', dest='extension',
                        help='filename extension following each sample ID')
    args = parser.parse_args()

    # sample Ids
    samples = []
    if args.samples is not None:
        if not isfile(args.samples):
            raise ValueError('Sample ID list file does not exist: %s.'
                             % args.samples)
        with open(args.samples, 'r') as f:
            samples = f.read().splitlines()
        print('Sample IDs to include: %d.' % len(samples))

    # sample Id to filename map
    sample2fname = {}
    fnames = listdir(args.input_dir)
    sample_set = set(samples)
    for fname in fnames:
        if args.extension is not None:
            if not fname.endswith(args.extension):
                continue
            sample = fname[:-len(args.extension)]
        else:
            sample = splitext(fname)[0]
        if sample_set:
            if sample in sample_set:
                sample2fname[sample] = fname
        else:
            samples.append(sample)
            sample2fname[sample] = fname
    if sample_set:
        samples = [x for x in samples if x in sample2fname]
    print('Samples to read: %d.' % len(samples))

    # subject ID to genome ID translation table
    sbj2g = {}
    if args.translation is not None:
        if not isfile(args.translation):
            raise ValueError('Translation table file does not exist: %s.'
                             % args.translation)
        with open(args.translation, 'r') as f:
            sbj2g = dict(x.split('\t') for x in f.read().splitlines())

    # parse maps and generate OGU tables
    data = {}
    for sample in samples:
        f = read(join(args.input_dir, sample2fname[sample]))
        data[sample] = parse_map(f, args.method, sbj2g)

    # write outputs
    for cat in ('all', 'norm', 'uniq'):
        gs = set()
        for sample in data:
            gs = gs.union(data[sample][cat])
        with open('%s.%s.tsv' % (args.output_name, cat), 'w') as f:
            f.write('#Genome ID\t%s\n' % '\t'.join(samples))
            for g in sorted(gs):
                out = [g]
                for sample in samples:
                    out.append(str(data[sample][cat][g])
                               if g in data[sample][cat] else '0')
                print('\t'.join(out), file=f)


def parse_map(f, method=None, sbj2g=None):
    """Parse a read-to-reference map generated by certain method.

    Parameters
    ----------
    f : file handle
        map file to parse
    method : str (optional)
        method for generating the map
    sbj2g : dict (optional)
        subject ID to genome ID map

    Returns
    -------
    dict of
        `all`, `norm`, `uniq`
    """
    m_ = method and method.lower()
    if m_ == 'centrifuge':
        return parse_centrifuge(f, sbj2g)
    elif m_ == 'bowtie2':
        return parse_simple(f, 2, sbj2g)
    elif m_ is None:
        return parse_simple(f, 1, sbj2g)
    else:
        raise ValueError('Unsupported method: %s.' % method)


def parse_simple(f, col=1, sbj2g=None):
    """Parse a read-to-reference map generated by certain method.

    Parameters
    ----------
    f : file handle
        map file to parse
    col : int (optional)
        index of column of subject IDs (default: 1)
    sbj2g : dict (optional)
        subject ID to genome ID map

    Returns
    -------
    dict of
        `all`, `norm`, `uniq`
    """
    seq2gs = {}
    for line in f:
        x = line.rstrip('\r\n').split('\t')
        g = x[col]  # subject Id
        if sbj2g:
            if g not in sbj2g:
                continue
            g = sbj2g[g]
        seq2gs.setdefault(x[0], []).append(g)
    res = {x: {} for x in ('all', 'norm', 'uniq')}
    for gs in seq2gs.values():
        g2n = {}
        for g in gs:
            g2n[g] = g2n.get(g, 0) + 1
        for g, n in g2n.items():
            res['all'][g] = res['all'].get(g, 0) + n
            res['norm'][g] = res['norm'].get(g, 0) + 1 / n
            if n == 1:  # unique match
                res['uniq'][g] = res['uniq'].get(g, 0) + 1
    for g in res['norm']:
        res['norm'][g] = int(res['norm'][g])
    return res


def parse_centrifuge(f, sbj2g=None):
    """Parse a read-to-reference map generated by Centrifuge.

    Parameters
    ----------
    f : file handle
        map file to parse
    sbj2g : dict (optional)
        subject ID to genome ID map

    Returns
    -------
    dict of
        `all`, `norm`, `uniq`

    Notes
    -----
    In a Centrifuge-generated map, columns are:
        readID, seqID, taxID, score, 2ndBestScore, hitLength, queryLength,
        numMatches
    """
    res = {x: {} for x in ('all', 'norm', 'uniq')}
    next(f)  # skip header
    for line in f:
        x = line.rstrip('\r\n').split('\t')
        g = x[1]  # subject Id
        if g == 'unclassified':
            continue
        if sbj2g:
            if g not in sbj2g:
                continue
            g = sbj2g[g]
        n = int(x[7])  # number of matches
        res['all'][g] = res['all'].get(g, 0) + 1
        res['norm'][g] = res['norm'].get(g, 0) + 1 / n
        if n == 1:  # unique match
            res['uniq'][g] = res['uniq'].get(g, 0) + 1
    for g in res['norm']:
        res['norm'][g] = int(res['norm'][g])
    return res


zipdict = {'.gz': gzip, '.bz2': bz2, '.xz': lzma, '.lz': lzma}


def read(fp):
    ext = splitext(fp)[1]
    zipfunc = getattr(zipdict[ext], 'open') if ext in zipdict else open
    return zipfunc(fp, 'rt')


if __name__ == "__main__":
    main()