#!/usr/bin/env python
# Time-stamp: <2012-02-27 12:20:36 Tao Liu>

"""Module Description

Copyright (c) 2008 Tao Liu <taoliu@jimmy.harvard.edu>

This code is free software; you can redistribute it and/or modify it
under the terms of the BSD License (see the file COPYING included with
the distribution).

@status:  experimental
@version: $Revision$
@author:  Tao Liu
@contact: taoliu@jimmy.harvard.edu
"""

# ------------------------------------
# python modules
# ------------------------------------

import os
import sys
from array import array as pyarray
from optparse import OptionParser
import pymc
from pymc import deterministic
from math import log

# ------------------------------------
# constants
# ------------------------------------
LOG2E = log(2.718281828459045,2)
# BUGFIX FOR PYMC
PROGRESS_BAR_ENABLED = 'progress_bar' in getargspec(pymc.MCMC.sample)[0]
# ------------------------------------
# Misc functions
# ------------------------------------

# ------------------------------------
# Classes
# ------------------------------------

# ------------------------------------
# Main function
# ------------------------------------
def main():
    usage = "usage: %prog [options]"
    description = "MCMC to calculate ratios between two Posterior Poisson with a given probability cutoff c. Left tail, right tail according to cutoff 0.01 and 0.05, and mean value will be stored in a table, then compress the table in gzip file."
    
    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("-m","--min",dest="mincount",type="int",
                         help="Minimum of pileup height. >=0. Default: 0", default=0)
    optparser.add_option("-M","--max",dest="maxcount",type="int",
                         help="Maximum of pileup height. <=10000. Default: 100", default=100)
    optparser.add_option("-N","--samplesize",dest="samplesize",type="int",
                         help="Sampling size for MCMC. Default: 15000", default=15000)
    optparser.add_option("-B","--burn",dest="burnsize",type="int",
                         help="Burn step for MCMC. First this number of samples will be trashed. Default: 5000", default=5000)
    optparser.add_option("-c","--cutoff",dest="cutoff",type="float",
                         help="Cutoff, must be <1. Default: 0.01", default=0.01)
    optparser.add_option("-o","--ofile",dest="ofile",
                         help="output file, a gzip file.") 
    (options,args) = optparser.parse_args()
    if options.mincount < 0 or options.maxcount > 10000 or options.samplesize <= 500 or options.cutoff>=1 or not options.ofile:
        optparser.print_help()
        sys.exit(1)

    ofhd = file(options.ofile,"w")

    sample_number = options.samplesize
    cutoff = options.cutoff

    gfold = pyarray('f',[])

    for i in xrange(options.mincount,options.maxcount+1):
        for j in xrange(options.mincount,options.maxcount+1):        
            P_X = MCMCPoissonPosteriorRatio(sample_number,1000,i,j)
            c = int(sample_number * cutoff)
            #print i,j,P_X[c],P_X[-1*c]
            if i > j:
                # X >= 0
                ret = max(0,P_X[c])
            else:
                # X < 0
                ret = min(0,P_X[-1*c])
            print i,j,ret
            gfold.append(ret)

    gfold.tofile(ofhd)
    ofhd.close()

def MCMCPoissonPosteriorRatio (sample_number, burn, count1, count2):
    """MCMC method to calculate ratio distribution of two Posterior Poisson distributions.

    sample_number: number of sampling. It must be greater than burn, however there is no check.
    burn: number of samples being burned.
    count1: observed counts of condition 1
    count2: observed counts of condition 2

    return: list of log2-ratios
    """
    lam1 = pymc.Uniform('U1',0,10000)   # prior of lambda is uniform distribution
    lam2 = pymc.Uniform('U2',0,10000)   # prior of lambda is uniform distribution    
    poi1 = pymc.Poisson('P1',lam1,value=count1,observed=True) # Poisson with observed value count1
    poi2 = pymc.Poisson('P2',lam2,value=count2,observed=True) # Poisson with observed value count2
    @deterministic
    def ratio (l1=lam1,l2=lam2):
        return log(l1) - log(l2)
    mcmcmodel  = pymc.MCMC([ratio,lam1,lam2,poi1,poi2])
    if PROGRESS_BAR_ENABLED:
        mcmcmodel.sample(iter=sample_number, progress_bar=False, burn=burn)
    else:
        mcmcmodel.sample(iter=sample_number, burn=burn)
    #print mcmcmodel.stats()
    #sys.exit(1)
    return map(lambda x:x*LOG2E, ratio.trace())

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        sys.stderr.write("User interrupt me! ;-) See you!\n")
        sys.exit(0)