#!/usr/bin/env python
import sys
import getopt
import numpy as np
from emase.AlignmentPropertyMatrix import AlignmentPropertyMatrix as APM

__author__ = 'Kwangbom "KB" Choi, Ph. D.'


help_message = '''
    Usage:
        count-alignments -i <infile> -g <grpfile> -o <outbase>

    Input:
        <infile>  : Alignments stored in an EMASE format (PyTables HDF5)
        <grpfile> : Gene to isoform map if counting at gene level
        <outbase> : The name of output result file. (default: './emase')

    Parameters:
        -h, --help: shows this help message
'''


class Usage(Exception):
    def __init__(self, msg):
        self.msg = msg


def report_alignment_counts(alignments, filename):
    alignment_counts = alignments.count_alignments()
    allelic_unique_counts = alignments.count_unique_reads(ignore_haplotype=False)
    locus_unique_counts = alignments.count_unique_reads(ignore_haplotype=True)
    cntdata = np.vstack((alignment_counts, allelic_unique_counts))
    cntdata = np.vstack((cntdata, locus_unique_counts))
    fhout = open(filename, 'w')
    fhout.write("locus\t" + "\t".join([ 'aln_%s' % h for h in alignments.hname]) + "\t")
    fhout.write("\t".join([ 'uniq_%s' % h for h in alignments.hname ]) + "\t")
    fhout.write("locus_uniq" + "\n")
    for locus_id in xrange(alignments.num_loci):
        fhout.write("\t".join([alignments.lname[locus_id]] + map(str, cntdata[:, locus_id].ravel())) + "\n")
    fhout.close()


def main(argv=None):
    if argv is None:
        argv = sys.argv
    try:
        try:
            opts, args = getopt.getopt(argv[1:], "hi:g:o:", ["help"])
        except getopt.error, msg:
            raise Usage(msg)

        # Default values of vars
        infile = None
        grpfile = None
        outbase = './emase'

        # option processing (change this later with optparse)
        for option, value in opts:
            if option in ("-h", "--help"):
                raise Usage(help_message)
            if option == "-i":
                infile = value
            if option == "-g":
                grpfile = value
            if option == "-o":
                outbase = value

        # Check if the required options are given
        if infile is None or grpfile is None:  # If alignment file is not given
            raise Usage(help_message)

        #
        # Main body
        #

        alnmat = APM(h5file=infile, grpfile=grpfile)
        outfile1 = "%s.isoforms.alignment_counts" % outbase
        alnmat.report_alignment_counts(filename=outfile1)

        alnmat._bundle_inline(reset=True)
        outfile2 = "%s.genes.alignment_counts" % outbase
        alnmat.report_alignment_counts(filename=outfile2)

        #
        # End of main body
        #

    except Usage, err:
        print >> sys.stderr, sys.argv[0].split("/")[-1] + ": " + str(err.msg)
        return 2


if __name__ == "__main__":
    sys.exit(main())
