#!/usr/bin/env python
# Time-stamp: <2009-10-16 14:46:37 Tao Liu>

"""Description: MACS main executable

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

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

@status:  beta
@version: $Revision$
@author:  Yong Zhang, Tao Liu
@contact: taoliu@jimmy.harvard.edu
"""

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

import os
import sys
import logging
import math
import re
from optparse import OptionParser

import gzip

# ------------------------------------
# own python modules
# ------------------------------------
from MACS.OptValidator import opt_validate
from MACS.OutputWriter import *
from MACS.Prob import binomial_cdf_inv
from MACS.PeakModel import PeakModel
from MACS.PeakDetect import PeakDetect
from MACS.Constants import *
# ------------------------------------
# Main function
# ------------------------------------
def main():
    """The Main function/pipeline for MACS.
    
    """
    # Parse options...
    options = opt_validate(prepare_optparser())
    # end of parsing commandline options
    info = options.info
    debug = options.debug
    #0 output arguments
    info("\n"+options.argtxt)
    
    #1 Read tag files
    info("#1 read tag files...")
    (treat, control) = load_tag_files_options (options)
    debug("#1 calculate max duplicate tags in single position based on binomal distribution...")
    options.max_dup_tags = cal_max_dup_tags(options.gsize,treat.total)
    debug("#1  max_dup_tags based on binomal = %d" % (options.max_dup_tags))
    debug("#1  unique tags in treatment: %d" % (treat.total_unique))
    debug("#1  total tags in treatment: %d" % (treat.total))
    tagsinfo = "# unique tags in treatment: %d\n" % (treat.total_unique)
    tagsinfo += "# total tags in treatment: %d\n" % (treat.total)
    if control:
        debug("#1  unique tags in control: %d" % (control.total_unique))
        debug("#1  total tags in control: %d" % (control.total))
        tagsinfo += "# unique tags in control: %d\n" % (control.total_unique)
        tagsinfo += "# total tags in control: %d\n" % (control.total)
    
    options.bg_redundant = bg_r (treat,options.max_dup_tags) # average
                                                             # number
                                                             # of
                                                             # duplicate
                                                             # tags
                                                             # per tag
                                                             # position
    info("#1  Background Redundant rate: %.2f" % (options.bg_redundant))
    info("#1 finished!")

    #2 Build Model
    info("#2 Build Peak Model...")
    if options.nomodel:
        info("#2 Skipped...")
        options.d=options.shiftsize*2
        options.scanwindow=max(2*options.d,2*options.bw) # scan window is no less than 2* bandwidth
    else:
        peakmodel = PeakModel(treatment = treat,
                              max_pairnum = MAX_PAIRNUM,
                              opt = options
                              )
        info("#2 finished!")
        debug("#2  Summary Model:")
        debug("#2   min_tags: %d" % (peakmodel.min_tags))
        debug("#2   d: %d" % (peakmodel.d))
        debug("#2   scan_window: %d" % (peakmodel.scan_window))
        info("#2.2 Generate R script for model : %s" % (options.modelR))
        model2r_script(peakmodel,options.modelR,options.name)
        options.d = peakmodel.d
        options.scanwindow= peakmodel.scan_window
        
    #3 Call Peaks
    options.info("#3 Call peaks...")
    if options.nolambda:
        options.info("# local lambda is disabled!")

    peakdetect = PeakDetect(treat = treat,
                            control = control,
                            opt = options
                            )
    peakdetect.call_peaks()    
    diag_result = peakdetect.diag_result()
    #4 output
    #4.1 peaks in XLS
    options.info("#4 Write output xls file... %s" % (options.peakxls))
    ofhd_xls = open(options.peakxls,"w")
    ofhd_xls.write("# This file is generated by MACS\n")
    ofhd_xls.write(options.argtxt+"\n")

    ofhd_xls.write(tagsinfo)

    ofhd_xls.write("# d = %d\n" % (options.d))
    if options.nolambda:
        ofhd_xls.write("# local lambda is disabled!\n")
    ofhd_xls.write(peakdetect.toxls())
    ofhd_xls.close()
    #4.2 peaks in BED
    options.info("#4 Write output bed file... %s" % (options.peakbed))
    ofhd_bed = open(options.peakbed,"w")
    ofhd_bed.write("track name=\"MACS peaks for %s\"\n" % (options.name))
    ofhd_bed.write(peakdetect.tobed())
    ofhd_bed.close()
    #4.3 negative peaks in XLS
    if options.cfile:
        options.info("#4 Write output xls file for negative peaks... %s" % (options.negxls))
        ofhd_xls = open(options.negxls,"w")
        ofhd_xls.write(peakdetect.neg_toxls())
        ofhd_xls.close()

    #4.4 diag result
    if diag_result:
        options.info("#4 Write diagnosis result ... %s" % (options.name+"_diag.xls"))
        diag_write (options.diagxls, diag_result)

    options.info("#5 Done! Check the output files!\n")

def prepare_optparser ():
    """Prepare optparser object. New options will be added in this
    function first.
    
    """
    usage = """usage: %prog <-t tfile> [options]
Example: %prog -t ChIP.bed -c Control.bed --name test
"""
    description = "%prog -- Model-based Analysis for ChIP-Sequencing"

    optparser = OptionParser(version="%prog 1.3.7.1 (Oktoberfest, bug fixed #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","--treatment",dest="tfile",type="string",
                         help="ChIP-seq treatment files. REQUIRED. When ELANDMULTIPET is selected, you must provide two files separated by comma, e.g. s_1_1_eland_multi.txt,s_1_2_eland_multi.txt")
    optparser.add_option("-c","--control",dest="cfile",type="string",
                         help="Control files. When ELANDMULTIPET is selected, you must provide two files separated by comma, e.g. s_2_1_eland_multi.txt,s_2_2_eland_multi.txt")
    optparser.add_option("--name",dest="name",type="string",
                         help="Experiment name, which will be used to generate output file names. DEFAULT: \"NA\"",
                         default="NA")
    optparser.add_option("--format",dest="format",type="string",
                         help="Format of tag file, \"BED\" or \"ELAND\" or \"ELANDMULTI\" or \"ELANDMULTIPET\" or \"SAM\" or \"BAM\" or \"BOWTIE\". Please check the definition in 00README file before you choose ELAND/ELANDMULTI/ELANDMULTIPET/SAM/BAM/BOWTIE. DEFAULT: \"BED\"",
                         default="BED")
    optparser.add_option("--gsize",dest="gsize",type="int",default=2700000000,
                         help="Effective genome size, default:2.7e+9")
    optparser.add_option("--tsize",dest="tsize",type="int",default=25,
                         help="Tag size. DEFAULT: 25")
    optparser.add_option("--bw",dest="bw",type="int",default=300,
                         help="Band width. This value is used while building the shifting model. If --nomodel is set, 2 time of this value will be used as a scanwindow width. DEFAULT: 300")
    optparser.add_option("--pvalue",dest="pvalue",type="float",default=1e-5,
                         help="Pvalue cutoff for peak detection. DEFAULT: 1e-5")
    optparser.add_option("--mfold",dest="mfold",type="int",default=32,
                         help="Select the regions with MFOLD high-confidence enrichment ratio against background to build model. DEFAULT:32")
    optparser.add_option("--wig",dest="store_wig",action="store_true",
                          help="Whether or not to save shifted raw tag count at every bp into a wiggle file. WARNING: this process is time/space consuming!!",
                          default=False)
    optparser.add_option("--wigextend",dest="wigextend",type="int",
                         help="If set as an integer, when MACS saves wiggle files, it will extend tag from its middle point to a wigextend size fragment. By default it is modeled d. Use this option if you want to increase the resolution in wiggle file. It doesn't affect peak calling.")
    optparser.add_option("--space",dest="space",type="int",
                          help="The resoluation for saving wiggle files, by default, MACS will save the raw tag count every 10 bps. Usable only with '--wig' option.",
                          default=10)
    optparser.add_option("--nolambda",dest="nolambda",action="store_true",
                         help="If True, MACS will use fixed background lambda as local lambda for every peak region. Normally, MACS calculates a dynamic local lambda to reflect the local bias due to potential chromatin structure. ",
                         default=False)
    optparser.add_option("--lambdaset",dest="lambdaset",type="string",default="1000,5000,10000",
                         help="Three levels of nearby region in basepairs to calculate dynamic lambda, DEFAULT: \"1000,5000,10000\" ")
    optparser.add_option("--nomodel",dest="nomodel",action="store_true",
                         help="Whether or not to build the shifting model. If True, MACS will not build model. by default it means shifting size = 100, try to set shiftsize to change it. DEFAULT: False",
                         default=False)
    optparser.add_option("--shiftsize",dest="shiftsize",type="int",default=100,
                         help="The arbitrary shift size in bp. When nomodel is true, MACS will regard this value as 'modeled' d. DEFAULT: 100 ")
    optparser.add_option("--diag",dest="diag",action="store_true",
                         help="Whether or not to produce a diagnosis report. It's up to 9X time consuming. Please check 00README file for detail. DEFAULT: False",default=False)
    optparser.add_option("--futurefdr",dest="futurefdr",action="store_true",
                         help="Whether or not to perform the new peak detection method which is though to be more suitable for sharp peaks. The default method only consider the peak location, 1k, 5k, and 10k regions in the control data; whereas the new future method also consider the 5k, 10k regions in treatment data to calculate local bias. DEFAULT: False",default=False)
    optparser.add_option("--petdist",dest="petdist",type="int",default=200,
                         help="Best distance between Pair-End Tags. Only available when format is 'ELANDMULTIPET'. DEFAULT: 200 ")
    optparser.add_option("--verbose",dest="verbose",type="int",default=2,
                         help="Set verbose level. 0: only show critical message, 1: show additional warning message, 2: show process information, 3: show debug messages. DEFAULT:2")
    optparser.add_option("--fe-min",dest="femin",type="int",default=0,
                         help="For diagnostics, min fold enrichment to consider. DEFAULT: 0")
    optparser.add_option("--fe-max",dest="femax",type="int",
                         help="For diagnostics, max fold enrichment to consider. DEFAULT: maximum fold enrichment")
    optparser.add_option("--fe-step",dest="festep",type="int",default=FESTEP,
                         help="For diagnostics, fold enrichment step.  DEFAULT: 20")

    return optparser
   
def bg_r (trackI, max_dup_tags):
    total_tags  = trackI.total	# number for all tags
    total_duplicates = 0		# total duplicated number of tags over
                                # maximum dup tag number
    chrs = trackI.get_chr_names()
    for chrom in chrs:
        countlist = trackI.get_comments_by_chr (chrom)
        for c in countlist[0]:
            total_duplicates += max(c-max_dup_tags,0) 
        for c in countlist[1]:
            total_duplicates += max(c-max_dup_tags,0)
    return float(total_duplicates)/(total_tags)

def cal_max_dup_tags ( genome_size, tags_number, p=1e-5 ):
    """Calculate the maximum duplicated tag number based on genome
    size, total tag number and a p-value based on binomial
    distribution. Brute force algorithm to calculate reverse CDF no
    more than MAX_LAMBDA(100000).
    
    """
    return binomial_cdf_inv(1-p,tags_number,1.0/genome_size)

def load_tag_files_options ( options ):
    """From the options, load treatment tags and control tags (if available).

    """
    if options.format == "ELANDMULTIPET":
        options.info("#1 read PET treatment tags...")
        treat = options.build(open2(options.tfile[0]),open2(options.tfile[1]),options.petdist)
        treat.merge_overlap()
        if options.cfile:
            options.info("#1.2 read input tags...")
            control = options.build(open2(options.cfile[0]),open2(options.cfile[1]),options.petdist)
            control.merge_overlap()
        else:
            control = None
    else:
        options.info("#1 read treatment tags...")
        treat = options.build(open2(options.tfile, gzip_flag=options.gzip_flag))
        treat.merge_overlap()
        if options.cfile:
            options.info("#1.2 read input tags...")
            control = options.build(open2(options.cfile, gzip_flag=options.gzip_flag))
            control.merge_overlap()
        else:
            control = None
    return (treat, control)

def open2(path, mode='r', bufsize=-1, gzip_flag=False):
    if path.endswith('.gz') or gzip_flag:
        f = gzip.open(path, mode)
    else:
        f = open(path, mode, bufsize)
    return f



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