#!/usr/bin/env python
# Time-stamp: <2012-02-29 14:42:57 Tao Liu>

import os
import sys
import logging
from optparse import OptionParser

from MACS2.IO import cBedGraphIO
from MACS2.IO.cBedGraph import scoreTracktoBedGraph
from MACS2.cProb import poisson_cdf
from math import log as mlog

# ------------------------------------
# constants
# ------------------------------------
logging.basicConfig(level=20,
                    format='%(levelname)-5s @ %(asctime)s: %(message)s ',
                    datefmt='%a, %d %b %Y %H:%M:%S',
                    stream=sys.stderr,
                    filemode="w"
                    )

LOG10_E = 0.43429448190325176

# ------------------------------------
# Misc functions
# ------------------------------------
error   = logging.critical		# function alias
warn    = logging.warning
debug   = logging.debug
info    = logging.info
# ------------------------------------
# Main function
# ------------------------------------
pscore_dict = {}

def get_pscore ( observed, expectation ):
    key_value = (observed, expectation)
    if pscore_dict.has_key(key_value):
        return pscore_dict[key_value]
    else:
        score = -1*poisson_cdf(observed,expectation,False,True)
        pscore_dict[(observed,expectation)] = score
        return score

logLR_dict = {}

def logLR ( x, y ):
    """Calculate log10 Likelihood between H1 ( enriched ) and H0 (
    chromatin bias ). Then set minus sign for depletion.
    
    """
    key_value = ( x, y )
    if logLR_dict.has_key(key_value):
        return logLR_dict[key_value]
    else:
        if x > y:
            s = (x*(mlog(x+1)-mlog(y+1))+y-x)*LOG10_E
        elif x < y:
            s = (-1*x*(mlog(x+1)-mlog(y+1))-y+x)*LOG10_E
        else:
            s = 0
        logLR_dict[key_value] = s
        return s

def main():
    usage = "usage: %prog <-t TREATMENT.BEDGRAPH> <-c CONTROL.BEDGRAPH> <-o OUTPUT.BEDGRAPH> [-m METHOD] "
    description = "Calculate scores using certain method by comparing a bedGraph file from treatment and a file from control representing local bias."
    
    optparser = OptionParser(version="%prog 0.1",description=description,usage=usage,add_help_option=False)
    optparser.add_option("-h","--help",action="help",help="Show this help message and exit.")
    optparser.add_option("-t","--tfile",dest="tfile",type="string",
                         help="Required: Treatment bedGraph file, e.g. *_treat_pileup.bdg from MACSv2")
    optparser.add_option("-c","--cfile",dest="cfile",type="string",
                         help="Required: Control bedGraph file, e.g. *_control_lambda.bdg from MACSv2")
    optparser.add_option("-o","--output",dest="ofile",type="string",
                         help="Required: The output bedGraph file to write scores.")
    optparser.add_option("-m","--method",dest="method",type="string",
                         help="Method to use while calculating a score in any bin by comparing treatment value and control value. Available choices are: ppois, qpois, substract, divide, and logLR,. They represent Poisson Pvalue (-log10(pvalue) form) using control as lambda and treatment as observation, q-value through a BH process for poisson pvalues, substraction from treatment, fold change which may be problematic if there are zero in control and log10 likelihood between ChIP-enriched model and open chromatin model. Default option is ppois.",default="ppois")
    (options,args) = optparser.parse_args()

    if not options.tfile or not options.cfile or not options.ofile:
        optparser.print_help()
        sys.exit()

    available_methods = ['ppois','substract','divide','logLR','qpois']
    if options.method not in available_methods:
        sys.stderr.write("Method can only be %s" % ",".join(available_methods))
    else:
        method = options.method

    info("Read and build treatment bedGraph...")
    tbio = cBedGraphIO.bedGraphIO(options.tfile)
    tbtrack = tbio.build_bdgtrack()

    info("Read and build control bedGraph...")
    cbio = cBedGraphIO.bedGraphIO(options.cfile)
    cbtrack = cbio.build_bdgtrack()

    info("Calculate scores comparing treatment and control by %s..." % method)
    # build score track
    if method == 'ppois':
        sbtrack = tbtrack.make_scoreTrack_for_macs(cbtrack)
        sbtrack = scoreTracktoBedGraph(sbtrack,'-100logp')        
    elif method == 'qpois':
        sbtrack = tbtrack.make_scoreTrack_for_macs(cbtrack)
        pqtable = sbtrack.make_pq_table()
        sbtrack.assign_qvalue(pqtable)
        sbtrack = scoreTracktoBedGraph(sbtrack,'-100logq')
    elif method == 'substract':
        sbtrack = tbtrack.overlie(cbtrack,func=lambda x,y:x-y)
    elif method == 'divide':
        sbtrack = tbtrack.overlie(cbtrack,func=lambda x,y:float(x)/y)
    elif method == 'logLR':             # log likelihood
        sbtrack = tbtrack.overlie(cbtrack,func=logLR)
    else:
        raise Exception("Can't reach here!")

    info("Write to output bedGraph...")
    ofhd = open(options.ofile,"w")

    sbtrack.write_bedGraph(ofhd,name="%s_Scores" % (method.upper()),description="Scores calculated by %s" % (method.upper()))
    

if __name__ == '__main__':
    main()
