#!/usr/bin/env python

"""
Copyright (c) 2015 Kwangbom Choi, The Jackson Laboratory
This software was developed by Kwangbom "KB" Choi in Gary Churchill's Lab.
This 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 software 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 software. If not, see <http://www.gnu.org/licenses/>.
"""

import sys
import getopt
import numpy as np
from emase.AlignmentPropertyMatrix import AlignmentPropertyMatrix as APM


help_message = '''
Usage:
    count-shared-multireads-pairwise -i <infile> -g <grpfile> -o <outbase>

Input:
    -i <infile>  : Alignments stored in an EMASE format (PyTables HDF5)
    -g <grpfile> : Gene to isoform map if counting at gene level
    -o <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 get_num_shared_multireads(alnmat):
    hapsum = alnmat.sum(axis=APM.Axis.HAPLOTYPE)
    hapsum.data = np.ones(hapsum.nnz)
    cntmat = hapsum.transpose() * hapsum
    #cntmat.setdiag(0)
    #cntmat.eliminate_zeros()
    return cntmat


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)
        cntmat = get_num_shared_multireads(alnmat)
        np.savez_compressed("%s.isoforms.shared_read_counts" % outbase, counts=cntmat)

        alnmat._bundle_inline(reset=True)
        cntmat = get_num_shared_multireads(alnmat)
        np.savez_compressed("%s.genes.shared_read_counts" % outbase, counts=cntmat)

        #
        # 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())
