import sys, argparse, math, random
import timeit
import numpy as np
from pyearth import Earth
from sets import Set
from scipy.stats import nbinom
import scipy, scipy.stats
from rpy2.robjects import r
import rpy2.robjects as robjects
import math
import os
import re

try:
        import HTSeq
except ImportError:
        sys.stderr.write( "Error: Could not import HTSeq. Please install the HTSeq Python framework\n" )
        sys.stderr.write( "available from http://www-huber.embl.de/users/anders/HTSeq\n" )
        sys.exit(1)

parser = argparse.ArgumentParser(
        description = "This generates negative-binomial distributed isoform-level read counts for in-silico AS studies based on real experimental data.",
        epilog =  "Written for Windows. Copyrighted by Ming Jie Wong, NTU"    
        )

parser.add_argument("gene_model", type=str, help='Reference annotation in GFF3 format.')
parser.add_argument("-g1", "--group1", nargs='+', type=argparse.FileType('r'), required=True, help='First group of sam files to be separated by a spacing.')
parser.add_argument("-g2", "--group2", nargs='+', type=argparse.FileType('r'), required=True, help='Second group of sam files to be separated by a spacing.')
parser.add_argument("-n", "--num-reps", type=int, default = 3, dest='nreps',
        help = "Number of replicates. Default is 3." )
parser.add_argument("-l", "--num-target-gene", type=int, default = 1000, dest='ntarg',
        help = "Number of AS genes. Default is 1000. ")
parser.add_argument("-m", "--mode", choices = ['all', 'chrom'], default = 'chrom',
        help = "Reading from DNA or from chromosomal DNA without pseudogenes. Default is chrom.")
parser.add_argument("-f", "--fit", choices = ['loess', 'mars', 'auto'], default = 'auto',
        help = "Fitting. loess, mars or auto. Default is auto.")
parser.add_argument("-c", "--cutoff", choices = ['h', 'l'], default = 'h',
        help = "Cutoff for genes chosen for AS. high (10 read counts) or low (0 read counts). Default is h.")

# Options for gene mode: all simulates annotated AS genes only while chrom simulates all genes in annotation file derived from chromosomal DNA.

if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

GENE = Set(["gene", "transposable_element_gene","pseudogene"])
EXON = Set(["exon", "pseudogenic_exon"])
## All the parsed inputs
args = parser.parse_args()
name_gtf = args.gene_model
NREPS=args.nreps
MODE=args.mode
NTARG = args.ntarg
FITTED = args.fit
CUTOFF = args.cutoff

def loessGCV(non0_df):
    ''' n X 2-fold Generalized Cross Validation '''
    #first fold
    population = r.nrow(non0_df)
    #initialize first sample for learning
    sample1mean = r.sample(non0_df[1],population[0]/2)
    index1 = np.array(r.match(sample1mean,non0_df[1])).tolist()
    #apply minus to all indices due to python index start from 0
    index1[:] = [x - 1 for x in index1]
    sample1var = robjects.FloatVector([non0_df[0][i] for i in index1])
    #initialize second sample for testing
    sample2mean = robjects.FloatVector([x for x in np.array(non0_df[1]).tolist() if x not in np.array(sample1mean).tolist()])
    index2 = np.array(r.match(sample2mean,non0_df[1])).tolist()
    #apply minus to all due to python index start from 0
    index2[:] = [x - 1 for x in index2]
    sample2var = robjects.FloatVector([non0_df[0][i] for i in index2])

    #applying learning to sample 1
    df1 = robjects.DataFrame({"mean": sample1mean, "var": sample1var})
    loess_fitlearn1 = r.loess("var ~ mean", data=df1, degree=2)
    span1 = loess_fitlearn1.rx2('pars').rx2('span')

    #applying testing to sample 2
    var_pred = r.predict(loess_fitlearn1, sample2mean)
    #calculate first root-mean-square error
    SEP1 = [(x - y)**2 for x, y in zip(var_pred, sample2var)]
    RMSEP1 = math.sqrt(sum(SEP1)/len(SEP1))

    #second fold
    #applying learning to sample 2
    df2 = robjects.DataFrame({"mean": sample2mean, "var": sample2var})
    loess_fitlearn2 = r.loess("var ~ mean", data=df2, degree=2)
    span2 = loess_fitlearn2.rx2('pars').rx2('span')

    #applying testing to sample 1
    var_pred = r.predict(loess_fitlearn2, sample1mean)
    #calculate first root-mean-square error
    SEP2 = [(x - y)**2 for x, y in zip(var_pred, sample1var)]
    RMSEP2 = math.sqrt(sum(SEP2)/len(SEP2))

    if RMSEP1 < RMSEP2:
        loess_fitprior = loess_fitlearn1
        var_pred = r.predict(loess_fitprior, non0_df[1])
        SEprior = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
        RMSEPprior = math.sqrt(sum(SEprior)/len(SEprior))

    else:
        loess_fitprior = loess_fitlearn2
        var_pred = r.predict(loess_fitprior, non0_df[1])
        SEprior = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
        RMSEPprior = math.sqrt(sum(SEprior)/len(SEprior))   
    
    counter = True
    while(counter):
        population = r.nrow(non0_df)
        #initialize first sample for learning
        sample1mean = r.sample(non0_df[1],population[0]/2)
        index1 = np.array(r.match(sample1mean,non0_df[1])).tolist()
        #apply minus to all indices due to python index start from 0
        index1[:] = [x - 1 for x in index1]
        sample1var = robjects.FloatVector([non0_df[0][i] for i in index1])
        #initialize second sample for testing
        sample2mean = robjects.FloatVector([x for x in np.array(non0_df[1]).tolist() if x not in np.array(sample1mean).tolist()])
        index2 = np.array(r.match(sample2mean,non0_df[1])).tolist()
        #apply minus to all due to python index start from 0
        index2[:] = [x - 1 for x in index2]
        sample2var = robjects.FloatVector([non0_df[0][i] for i in index2])

        #applying learning to sample 1
        df1 = robjects.DataFrame({"mean": sample1mean, "var": sample1var})
        loess_fitlearn1 = r.loess("var ~ mean", data=df1, degree=2)
        span1 = loess_fitlearn1.rx2('pars').rx2('span')

        #applying testing to sample 2
        var_pred = r.predict(loess_fitlearn1, sample2mean)
        #calculate first root-mean-square error
        SEP1 = [(x - y)**2 for x, y in zip(var_pred, sample2var)]
        RMSEP1 = math.sqrt(sum(SEP1)/len(SEP1))

        #second fold
        #applying learning to sample 2
        df2 = robjects.DataFrame({"mean": sample2mean, "var": sample2var})
        loess_fitlearn2 = r.loess("var ~ mean", data=df2, degree=2)
        span2 = loess_fitlearn2.rx2('pars').rx2('span')

        #applying testing to sample 1
        var_pred = r.predict(loess_fitlearn2, sample1mean)
        #calculate first root-mean-square error
        SEP2 = [(x - y)**2 for x, y in zip(var_pred, sample1var)]
        RMSEP2 = math.sqrt(sum(SEP2)/len(SEP2))

        if RMSEP1 < RMSEP2:
            var_pred = r.predict(loess_fitlearn1, non0_df[1])
            SE = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
            RMSEP = math.sqrt(sum(SE)/len(SE))
            loess_fit = loess_fitlearn1

        else:
            var_pred = r.predict(loess_fitlearn2, non0_df[1])
            SE = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
            RMSEP = math.sqrt(sum(SE)/len(SE))
            loess_fit = loess_fitlearn2

        if RMSEP >= RMSEPprior:
            counter = False

        else:
            RMSEPprior =  RMSEP
            loess_fitprior = loess_fit

    return loess_fitprior

def fivetimestwofold(non0_df,option):
    RMSEPsum = 0
    population = r.nrow(non0_df)
    #initialize first sample for learning
    for i in range(5):
        sample1mean = r.sample(non0_df[1],population[0]/2)
        index1 = np.array(r.match(sample1mean,non0_df[1])).tolist()
        #apply minus to all indices due to python index start from 0
        index1[:] = [x - 1 for x in index1]
        sample1var = robjects.FloatVector([non0_df[0][i] for i in index1])
        #initialize second sample for testing
        sample2mean = robjects.FloatVector([x for x in np.array(non0_df[1]).tolist() if x not in np.array(sample1mean).tolist()])
        index2 = np.array(r.match(sample2mean,non0_df[1])).tolist()
        #apply minus to all due to python index start from 0
        index2[:] = [x - 1 for x in index2]
        sample2var = robjects.FloatVector([non0_df[0][i] for i in index2])

        #applying learning to sample 1
        df1 = robjects.DataFrame({"mean": sample1mean, "var": sample1var})
        if option == "loess":
            fitlearn1 = r.loess("var ~ mean", data=df1, degree=2)
            var_pred = r.predict(fitlearn1, sample2mean)
            SEP1 = [(x - y)**2 for x, y in zip(var_pred, sample2var)]
            RMSEP1 = math.sqrt(sum(SEP1)/len(SEP1))

        if option == "lm":
            fitlearn1 = r.lm("var ~ mean",data=df1)
            var_pred = r.predict(fitlearn1, robjects.DataFrame({"mean":sample2mean}))
            SEP1 = [(x - y)**2 for x, y in zip(var_pred, sample2var)]
            RMSEP1 = math.sqrt(sum(SEP1)/len(SEP1))

        if option == "mars":
            model = Earth()
            a_earth = np.array(df1.rx2("mean"))
            b_earth = np.array(df1.rx2("var"))
            model.fit(a_earth,b_earth)
            var_pred = model.predict(np.array(sample2mean)).tolist()
            SEP1 = [(x - y)**2 for x, y in zip(var_pred, sample2var)]
            RMSEP1 = math.sqrt(sum(SEP1)/len(SEP1))

        #second fold
        #applying learning to sample 2
        df2 = robjects.DataFrame({"mean": sample2mean, "var": sample2var})
        if option == "loess":
            fitlearn2 = r.loess("var ~ mean", data=df2, degree=2)
            var_pred = r.predict(fitlearn2, sample1mean)
            SEP2 = [(x - y)**2 for x, y in zip(var_pred, sample1var)]
            RMSEP2 = math.sqrt(sum(SEP2)/len(SEP2))

        if option == "lm":
            fitlearn2 = r.lm("var ~ mean",data=df2)
            var_pred = r.predict(fitlearn2, robjects.DataFrame({"mean":sample1mean}))
            SEP2 = [(x - y)**2 for x, y in zip(var_pred, sample1var)]
            RMSEP2 = math.sqrt(sum(SEP2)/len(SEP2))

        if option == "mars":
            model = Earth()
            a_earth = np.array(df2.rx2("mean"))
            b_earth = np.array(df2.rx2("var"))
            model.fit(a_earth,b_earth)
            var_pred = model.predict(np.array(sample1mean)).tolist()
            SEP2 = [(x - y)**2 for x, y in zip(var_pred, sample1var)]
            RMSEP2 = math.sqrt(sum(SEP2)/len(SEP2))

        RMSEPsum = (RMSEP1 + RMSEP2)/2 + RMSEPsum
    return RMSEPsum/5

def modelselection(non0_df):
    print "Running Model Selection..."
    RMSEPloess = fivetimestwofold(non0_df,"loess")
    RMSEPmars = fivetimestwofold(non0_df,"mars")
    RMSEPlm = fivetimestwofold(non0_df,"lm")
    RMSEP = [RMSEPloess,RMSEPmars,RMSEPlm]
    print "Root Mean Squared Error Prediction (RMSEP) Summary: "
    print "LOESS: ", RMSEPloess
    print "MARS: ", RMSEPmars
    print "LM: ", RMSEPlm
    if (min(RMSEP) == RMSEPloess):
        print "LOESS is selected."
        return "loess"
    elif (min(RMSEP) == RMSEPmars):
        print "MARS is selected."
        return "mars"
    elif (min(RMSEP) == RMSEPlm):
        print "lm is selected."
        return "lm"

def fullbagging(non0_df,a,option):
    population = r.nrow(non0_df)
    var_predsum = robjects.FloatVector(np.zeros(len(a)).tolist())
    for i in range(1000):
        sample1mean = r.sample(non0_df[1],population[0]*1/2,replace=True)
        index1 = np.array(r.match(sample1mean,non0_df[1])).tolist()
        #apply minus to all indices due to python index start from 0
        index1[:] = [x - 1 for x in index1]
        sample1var = robjects.FloatVector([non0_df[0][i] for i in index1])

        df1 = robjects.DataFrame({"mean": sample1mean, "var": sample1var})
        if option == "loess":
            loess_fitlearn1 = r.loess("var ~ mean", data=df1, degree=2)
            span1 = loess_fitlearn1.rx2('pars').rx2('span')
            var_pred1 = r.predict(loess_fitlearn1, a)
            var_predsum = robjects.FloatVector([x + y for x,y in zip(var_predsum,var_pred1)])

        elif option == "mars":
            model = Earth()
            a_earth = np.array(df1.rx2("mean"))
            b_earth = np.array(df1.rx2("var"))
            model.fit(a_earth,b_earth)
            var_pred1 = model.predict(np.array(a)).tolist()
            var_predsum = robjects.FloatVector([x + y for x,y in zip(var_predsum,var_pred1)])

        elif option == "lm":
            fitlearn1 = r.lm("var ~ mean",data=df1)
            var_pred1 = r.predict(fitlearn1, robjects.DataFrame({"mean":a}))
            var_predsum = robjects.FloatVector([x + y for x,y in zip(var_predsum,var_pred1)])

    return robjects.FloatVector([x/1000 for x in var_predsum])

def scedasticity_check(non0_df):
    fitlearn = r.lm("var ~ mean",data=non0_df)
    var_pred = r.predict(fitlearn, robjects.DataFrame({"mean":non0_df[1]}))
    SE = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
    sigma2 = sum(SE)/len(SE)
    #w = [(x - sigma2)**2 for x in SE]
    w = [(x - sigma2) for x in SE]
    neww = robjects.FloatVector(w)
    df1 = robjects.DataFrame({"mean": non0_df[1], "var": neww})
    fitlearn1 = r.lm("var ~ mean",data=df1)
    fv = r.predict(fitlearn1, robjects.DataFrame({"mean":non0_df[1]}))
    #fv = [(x - y)**2 for x, y in zip(w, non0_df[1])]
    fv2 = [(x - y)**2 for x, y in zip(fv, w)]
    sstot = [(x - sum(w)/len(w))**2 for x in w]
    #fv2 = [x**2 for x in fv]
    w2 = [y**2 for y in w]
    bp = len(SE) * (1-sum(fv2)/sum(sstot))
    #bp = len(SE) * sum(fv2)/sum(w2)

    pvalue = scipy.stats.chisqprob(bp, 1)
    chisquare = bp

    return chisquare, pvalue

def stabilitycheck(non0_df,option):
    RMSEPsum = 0
    population = r.nrow(non0_df)
    RMSEP_array = []
    #initialize first sample for learning
    for i in range(100):
        sample1mean = r.sample(non0_df[1],population[0]/2)
        index1 = np.array(r.match(sample1mean,non0_df[1])).tolist()
        #apply minus to all indices due to python index start from 0
        index1[:] = [x - 1 for x in index1]
        sample1var = robjects.FloatVector([non0_df[0][i] for i in index1])
        #initialize second sample for testing
        sample2mean = robjects.FloatVector([x for x in np.array(non0_df[1]).tolist() if x not in np.array(sample1mean).tolist()])
        index2 = np.array(r.match(sample2mean,non0_df[1])).tolist()
        #apply minus to all due to python index start from 0
        index2[:] = [x - 1 for x in index2]
        sample2var = robjects.FloatVector([non0_df[0][i] for i in index2])

        #applying learning to sample 1
        df1 = robjects.DataFrame({"mean": sample1mean, "var": sample1var})
        if option == "loess":
            fitlearn1 = r.loess("var ~ mean", data=df1, degree=2)
            var_pred = r.predict(fitlearn1, sample2mean)
            SE1 = [(x - y)**2 for x, y in zip(var_pred, sample2var)]
            RMSEP1 = math.sqrt(sum(SE1)/len(SE1))

        if option == "lm":
            fitlearn1 = r.lm("var ~ mean",data=df1)
            var_pred = r.predict(fitlearn1, robjects.DataFrame({"mean":sample2mean}))
            SE1 = [(x - y)**2 for x, y in zip(var_pred, sample2var)]
            RMSEP1 = math.sqrt(sum(SE1)/len(SE1))

        if option == "mars":
            model = Earth()
            a_earth = np.array(df1.rx2("mean"))
            b_earth = np.array(df1.rx2("var"))
            model.fit(a_earth,b_earth)
            var_pred = model.predict(np.array(sample2mean)).tolist()
            SE1 = [(x - y)**2 for x, y in zip(var_pred, sample2var)]
            RMSEP1 = math.sqrt(sum(SE1)/len(SE1))

        #second fold
        #applying learning to sample 2
        df2 = robjects.DataFrame({"mean": sample2mean, "var": sample2var})
        if option == "loess":
            fitlearn2 = r.loess("var ~ mean", data=df2, degree=2)
            var_pred = r.predict(fitlearn2, sample1mean)
            SE2 = [(x - y)**2 for x, y in zip(var_pred, sample1var)]
            RMSEP2 = math.sqrt(sum(SE2)/len(SE2))

        if option == "lm":
            fitlearn2 = r.lm("var ~ mean",data=df2)
            var_pred = r.predict(fitlearn2, robjects.DataFrame({"mean":sample1mean}))
            SE2 = [(x - y)**2 for x, y in zip(var_pred, sample1var)]
            RMSEP2 = math.sqrt(sum(SE2)/len(SE2))

        if option == "mars":
            model = Earth()
            a_earth = np.array(df2.rx2("mean"))
            b_earth = np.array(df2.rx2("var"))
            model.fit(a_earth,b_earth)
            var_pred = model.predict(np.array(sample1mean)).tolist()
            SE2 = [(x - y)**2 for x, y in zip(var_pred, sample1var)]
            RMSEP2 = math.sqrt(sum(SE2)/len(SE2))

        if RMSEP1 < RMSEP2:
            if option == "loess":
                    var_pred = r.predict(fitlearn1, non0_df[1])
                    SE_final = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
                    RMSEP_final = math.sqrt(sum(SE_final)/len(SE_final))
                    
            if option == "lm":
                    var_pred = r.predict(fitlearn1, robjects.DataFrame({"mean":non0_df[1]}))
                    SE_final = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
                    RMSEP_final = math.sqrt(sum(SE_final)/len(SE_final))

            if option == "mars":
                    model = Earth()
                    a_earth = np.array(df1.rx2("mean"))
                    b_earth = np.array(df1.rx2("var"))
                    model.fit(a_earth,b_earth)
                    var_pred = model.predict(np.array(non0_df[1])).tolist()
                    SE_final = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
                    RMSEP_final = math.sqrt(sum(SE_final)/len(SE_final))

        else:
            if option == "loess":
                    var_pred = r.predict(fitlearn2, non0_df[1])
                    SE_final = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
                    RMSEP_final = math.sqrt(sum(SE_final)/len(SE_final))
                    
            if option == "lm":
                    var_pred = r.predict(fitlearn2, robjects.DataFrame({"mean":non0_df[1]}))
                    SE_final = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
                    RMSEP_final = math.sqrt(sum(SE_final)/len(SE_final))

            if option == "mars":
                    model = Earth()
                    a_earth = np.array(df2.rx2("mean"))
                    b_earth = np.array(df2.rx2("var"))
                    model.fit(a_earth,b_earth)
                    var_pred = model.predict(np.array(non0_df[1])).tolist()
                    SE_final = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
                    RMSEP_final = math.sqrt(sum(SE_final)/len(SE_final))

        RMSEP_array.append(RMSEP_final)

    if option == "loess":
        fitlearn = r.loess("var ~ mean", data=non0_df, degree=2)
        var_pred = r.predict(fitlearn, non0_df[1])
        SE = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
        RMSEP = math.sqrt(sum(SE)/len(SE))

    if option == "lm":
        fitlearn = r.lm("var ~ mean",data=non0_df)
        var_pred = r.predict(fitlearn, robjects.DataFrame({"mean":non0_df[1]}))
        SE = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
        RMSEP = math.sqrt(sum(SE)/len(SE))

    if option == "mars":
        model = Earth()
        a_earth = np.array(non0_df.rx2("mean"))
        b_earth = np.array(non0_df.rx2("var"))
        model.fit(a_earth,b_earth)
        var_pred = model.predict(np.array(non0_df[1])).tolist()
        SE = [(x - y)**2 for x, y in zip(var_pred, non0_df[0])]
        RMSEP = math.sqrt(sum(SE)/len(SE))

    success = sum([x<RMSEP for x in RMSEP_array])
    pvalue = scipy.stats.binom.cdf(success, len(RMSEP_array), 0.5)

    return pvalue

def countSam(_sam_file, _genes, _dic, _idx):
        '''
        This function process a single sam file and count the reads sitting on each gene
        Results are saved in the form of global variable.
        '''
        num_reads = 0
        for alnmt in _sam_file:
                if alnmt.aligned: 
                        intersection_set = None
                        #print list(_genes[alnmt.iv].steps())
                        for iv1, step_set in _genes[alnmt.iv].steps():
                                #looks at all the considered genes with alnmt.iv intervals
                                #print iv1, step_set separately
                                #iv1: the chromosome information and intervals
                                #step_set: the gene 
                                if intersection_set is None:
                                        intersection_set = step_set.copy()
                                else:
                                        #https://docs.python.org/2.4/lib/types-set.html
                                        intersection_set.intersection_update(step_set)
                        if len(intersection_set) == 1:
                                _dic[ list(intersection_set)[0] ][_idx] += 1
                num_reads += 1
                if num_reads % 100000 == 0:
                        sys.stderr.write("%d reads have been processed.\n" % num_reads)
        return _dic

def meanVar(_files, _gff_file , _output):


        NFILE=len(_files)
        if NFILE == 1:
                sys.stderr.write("Error: At least two samples for each group are required.\n")
                sys.exit(1)
        #####

        _dict_counts = dict() ## dictionary of gene counts
        _genes = HTSeq.GenomicArrayOfSets("auto",stranded=False)
        idx=0
        count = 0
        count2 = 0
        count3 = 0
        rest = 'none'
        if MODE == "all" or MODE == "chrom":
                ## Bug: Does not report last gene in gff if it has at least two transcript
                transcript= set()
                cur_line = None
                for feature in _gff_file:
                        if feature.type in GENE:
                                if len(transcript) >1:
                                        _dict_counts[ cur_line.name ] = [0]*NFILE #cur_line.name is the coordinates
                                        _genes[cur_line.iv] += cur_line.name #gene name/ id
                                        count +=1
                                cur_line = feature
                                transcript.clear()
                        else:
                                if count ==  0 and feature.type in Set(["mRNA"]):
                                        if len(transcript) >1:
                                                _dict_counts[ cur_line.attr["geneID"] ] = [0]*NFILE
                                                # if len(transcript) >= 1, then count3 is total amount (in Arabidopsis is 35385)
                                                # GenomicFeature .name .type (mRNA etc) .iv (interval)
                                                #sep = '.'
                                                #rest = cur_line.name.split(sep,1)[0] #mRNA IDs consist of .1, .2,... as they maybe isoforms of the same gene
                                                _genes[cur_line.iv] += cur_line.attr["geneID"]
                                                count3 +=1
                                                #print cur_line.attr["geneID"]
                                        cur_line = feature
                                        transcript.clear()                                      
                        if feature.type in EXON:
                                count2 +=1
                                #sep = '.'
                                #featparent = feature.attr["Parent"].split(sep,1)[0]
                                featparent = feature.attr["Parent"]
                                #adding same elements will not result in repetitions e.g.
                                #hey = set() >>> hey.add("me") >>> hey >>> set(['me']) >>> hey.add("wo") >>> hey >>> set(['me', 'wo']) >>> hey.add("me") >>> hey >>> set(['me', 'wo'])
                                transcript.add(featparent)
        print "Summary of GFF3 File... "
        print "number of genes: ", count
        print "number of exons: ", count2
        print "number of mRNA: ", count3
        print "Running HTSeq... "
        _file_raw_count = open(_output+'.rawcounts','w')
        ## Reads through the input list of sam files and then call countSam for each inputed sam file  
        for f in _files:
                sam_file=HTSeq.SAM_Reader(f)
                _dict_counts=countSam(sam_file, _genes,_dict_counts, idx)
                f.close()
                idx += 1
                sys.stderr.write("library %d has been generated.\n" % idx)
        ## Print raw counts in file specified by <out>
        if MODE == "all":
                for key, value in sorted(_dict_counts.iteritems()):
                        if "gene" in key or "transcript" in key:
                                sep = ':'
                                rest = key.split(sep,1)[1]
                                if '.' in rest:
                                        sep2 = '.'
                                        rest2 = rest.split(sep2,1)[0]
                                        _file_raw_count.write(rest2+"\t"+"\t".join(map(str,value))+"\n")
                                else:
                                        _file_raw_count.write(rest+"\t"+"\t".join(map(str,value))+"\n")

                        else:
                                _file_raw_count.write(key+"\t"+"\t".join(map(str,value))+"\n")
        elif MODE == "chrom":
                for key, value in sorted(_dict_counts.iteritems()):
                        if "gene" in key:
                                sep = ':'
                                rest = key.split(sep,1)[1]
                                _file_raw_count.write(rest+"\t"+"\t".join(map(str,value))+"\n")
                        ## remove pseudogenes
                        elif "transcript" not in key:
                                _file_raw_count.write(key+"\t"+"\t".join(map(str,value))+"\n")
        _file_raw_count.close()

        RMSEP = 1
        RMSEP2 = 0
        _dict_counts_org = _dict_counts
        while RMSEP > RMSEP2:
                
                ## calculate group mean and variance
                list_mean = list()
                list_var = list()
                # loop <for all genes and their values> 
                for key, value in sorted(_dict_counts_org.iteritems()):
                                # print value
                                # print key
                                # print np.mean(np.array(value))
                                # print np.var(np.array(value))
                                list_mean.append(np.mean(np.array(value)))
                                list_var.append(np.var(np.array(value)))
                                
                ## The following code uses rpy2 module
                ## loess is a non-parametric locally weighted regression using a nearest neighbor approach to identify pattern
                ## in the relationship between between the mean and variance for all the genes considered without a
                ## priori knowledge of the "model" and the requirements of the parameters to formulate such a "model".
                ## https://en.wikipedia.org/wiki/Local_regression
                ## http://stackoverflow.com/questions/7550582/when-to-choose-nls-over-loess

                a = robjects.FloatVector(list_mean)
                b = robjects.FloatVector(list_var)
                df = robjects.DataFrame({"mean": a, "var": b})
                non0_df=df.rx(df.rx2("mean").ro > 0, True) ## subsetting if mean > 0
                        

                if FITTED == "loess":
                        ## http://stats.stackexchange.com/questions/25092/what-is-the-difference-between-simple-linear-model-and-loess-model
                        ## http://blog.minitab.com/blog/adventures-in-statistics/why-is-there-no-r-squared-for-nonlinear-regression
                        print "Formulating LOESS..."
                        loess_fit = r.loess("var ~ mean", data=non0_df, degree=2)
                        print "Formulating Prediction..."
                        var_pred = r.predict(loess_fit, a)

                if FITTED == "mars":
                        print "Formulating MARS..."
                        model = Earth()
                        a_earth = np.array(non0_df.rx2("mean"))
                        b_earth = np.array(non0_df.rx2("var"))
                        model.fit(a_earth,b_earth)
                        print model.trace()
                        print model.summary()
                        var_pred = model.predict(np.array(a)).tolist()

                if FITTED == "auto":
                        print "Running Dataset Check for Appropriate Model Assessment..."
                        option = modelselection(non0_df)
                        print "Running Breusch Pagan Test for heteroskedasticity..."
                        chisq, pvalue = scedasticity_check(non0_df)
                        baggingcounter = 0
                        if (pvalue <= 0.05):
                                        #print chisq
                                        #print pvalue
                                        print "Rejection of homoskedasticity..."
                                        if option == "loess":
                                                        print "Model Approved..."
                                        else:
                                                        option = "loess"
                                                        print "LOESS is now selected instead..."
                                        print "Running one-sided Sign Test..."
                                        pvalue2 = stabilitycheck(non0_df,option)
                                        #pvalue2 = 0.01
                                        #print pvalue2
                                        if (pvalue2 <= 0.05):
                                                        print "Formulating LOESS for 2 fold GCV..."
                                                        loess_fit = loessGCV(non0_df)
                                                        var_pred = r.predict(loess_fit, a)
                                                        SEnew = [(x - y)**2 for x, y in zip(var_pred, b)]
                                                        RMSEPnew = math.sqrt(sum(SEnew)/len(SEnew))
                                        else:
                                                        print "Formulating LOESS for BootStrap Aggregating..."
                                                        var_pred = fullbagging(non0_df,a,option)
                                                        baggingcounter = 1
                                        
                        else:
                                        print "Model Approved for Bootstrap Aggregating..."
                                        print "Formulating Bootstrap Aggregating..."
                                        var_pred = fullbagging(non0_df,a,option)
                                        baggingcounter = 1

                        if option == "loess":
                                fitlearn = r.loess("var ~ mean", data=non0_df, degree=2)
                                var_preddef = r.predict(fitlearn, a)
                                SEdef = [(x - y)**2 for x, y in zip(var_preddef, b)]
                                RMSEPdef = math.sqrt(sum(SEdef)/len(SEdef))

                        if option == "lm":
                                fitlearn = r.lm("var ~ mean",data=non0_df)
                                var_preddef = r.predict(fitlearn, robjects.DataFrame({"mean":a}))
                                SEdef = [(x - y)**2 for x, y in zip(var_preddef, b)]
                                RMSEPdef = math.sqrt(sum(SEdef)/len(SEdef))

                        if option == "mars":
                                model = Earth()
                                a_earth = np.array(non0_df.rx2("mean"))
                                b_earth = np.array(non0_df.rx2("var"))
                                model.fit(a_earth,b_earth)
                                var_preddef = model.predict(np.array(a)).tolist()
                                SE = [(x - y)**2 for x, y in zip(var_preddef, b)]
                                RMSEPdef = math.sqrt(sum(SEdef)/len(SEdef))

                        if baggingcounter == 0:
                                print "Default Root Mean Squared Error Predicition (RMSEP):", RMSEPdef
                                print "New Root Mean Squared Error Predicition (RMSEP):", RMSEPnew

                print "Proceeding with simulation of negative-binomial distributed counts... "
                # This loop overwrite global variable dict_counts for recording new count data
                count_idx = 0
                for key, value in sorted(_dict_counts.iteritems()):
                                if var_pred[count_idx] > list_mean[count_idx]:
                                                n = math.pow(list_mean[count_idx],2)/(var_pred[count_idx]-list_mean[count_idx])
                                                if math.isnan(n) == True:
                                                                n = 0
                                                n = int(n) # n: number of failures
                                                if n<=0:
                                                                _dict_counts[key] = [0]*NREPS
                                                else:
                                                                p = n/float(n+list_mean[count_idx]) # p: prob of failure
                                                                _dict_counts[key] = nbinom.rvs(n, p, size=NREPS).tolist()
                                else:
                                                if list_mean[count_idx] == 0:
                                                                _dict_counts[key] = [0]*NREPS
                                                else:
                                                                _dict_counts[key] = scipy.stats.poisson.rvs(list_mean[count_idx], size=NREPS).tolist()
                                count_idx += 1
                #var_pred = r.predict(loess_fit, a)
                _file_nb_count = open(_output+'.nbcounts','w')
                if MODE == "all":
                                for key, value in sorted(_dict_counts.iteritems()):
                                                ## include pseudogenes
                                                if "gene" in key or "transcript" in key:
                                                                sep = ':'
                                                                rest = key.split(sep,1)[1]
                                                                if '.' in rest:
                                                                                sep2 = '.'
                                                                                rest2 = rest.split(sep2,1)[0]
                                                                                _file_nb_count.write(rest2+"\t"+"\t".join(map(str,value))+"\n")
                                                                else:
                                                                                _file_nb_count.write(rest+"\t"+"\t".join(map(str,value))+"\n")

                                                else:
                                                                _file_nb_count.write(key+"\t"+"\t".join(map(str,value))+"\n")
                elif MODE == "chrom":
                                for key, value in sorted(_dict_counts.iteritems()):
                                                if "gene" in key:
                                                                sep = ':'
                                                                rest = key.split(sep,1)[1]
                                                                _file_nb_count.write(rest+"\t"+"\t".join(map(str,value))+"\n")
                                                ## remove pseudogenes
                                                elif "transcript" not in key:
                                                                _file_nb_count.write(key+"\t"+"\t".join(map(str,value))+"\n")
                                                
                _file_nb_count.close()
                _file_raw_count.close()

                if FITTED == "loess" or FITTED == "mars":
                    return _dict_counts

                else:

                    print "Running mean read counts validation..."
                    count_idx = 0

                    if option == "loess":
                            ## http://stats.stackexchange.com/questions/25092/what-is-the-difference-between-simple-linear-model-and-loess-model
                            ## http://blog.minitab.com/blog/adventures-in-statistics/why-is-there-no-r-squared-for-nonlinear-regression
                            loess_fit2 = r.loess("var ~ mean", data=non0_df, degree=2)
                            print "Formulating Prediction..."
                            var_pred2 = r.predict(loess_fit2, a)

                    if option == "mars":
                            model = Earth()
                            a_earth = np.array(non0_df.rx2("mean"))
                            b_earth = np.array(non0_df.rx2("var"))
                            model.fit(a_earth,b_earth)
                            print model.trace()
                            print model.summary()
                            var_pred2 = model.predict(np.array(a)).tolist()

                    if option == "lm":
                            fitlearn2 = r.lm("var ~ mean",data=non0_df)
                            var_pred2 = r.predict(fitlearn2, robjects.DataFrame({"mean":a}))

                    checkRMSEList = []

                    for key, value in sorted(_dict_counts.iteritems()):
                                    if var_pred2[count_idx] > list_mean[count_idx]:
                                                    n = math.pow(list_mean[count_idx],2)/(var_pred2[count_idx]-list_mean[count_idx])
                                                    if math.isnan(n) == True:
                                                                    n = 0
                                                    n = int(n) # n: number of failures
                                                    if n<=0:
                                                                    checkRMSEList.append(np.mean([0]*NREPS))
                                                    else:
                                                                    p = n/float(n+list_mean[count_idx]) # p: prob of failure
                                                                    checkRMSEList.append(np.mean(nbinom.rvs(n, p, size=NREPS).tolist()))
                                    else:
                                                    if list_mean[count_idx] == 0:
                                                                    checkRMSEList.append(np.mean([0]*NREPS))
                                                    else:
                                                                    checkRMSEList.append(np.mean(scipy.stats.poisson.rvs(list_mean[count_idx], size=NREPS).tolist()))
                                    count_idx += 1

                    rawcountsList = []
                    nbcountsList = []
                    if os.path.isfile('group2.nbcounts'):
                            with open('group2.rawcounts') as oldfile, open('group2.nbcounts') as newfile:
                                    for line in oldfile:
                                            line1 = re.split('\t+',line)
                                            line1[-1] = line1[-1].strip()
                                            rawcountsList.append((float(line1[1])+float(line1[2]))/2)

                                    for line in newfile:
                                            line1 = re.split('\t+',line)
                                            line1[-1] = line1[-1].strip()
                                            nbcountsList.append((float(line1[1])+float(line1[2]))/2)

                    elif os.path.isfile('group1.nbcounts'):
                            with open('group1.rawcounts') as oldfile, open('group1.nbcounts') as newfile:
                                    for line in oldfile:
                                            line1 = re.split('\t+',line)
                                            line1[-1] = line1[-1].strip()
                                            rawcountsList.append((float(line1[1])+float(line1[2]))/2)

                                    for line in newfile:
                                            line1 = re.split('\t+',line)
                                            line1[-1] = line1[-1].strip()
                                            nbcountsList.append((float(line1[1])+float(line1[2]))/2)

                    oldfile.close()
                    newfile.close()

                    #print nbcountsList
                    #print rawcountsList
                    #print checkRMSEList

                    SEP = [(x - y)**2 for x, y in zip(rawcountsList, nbcountsList)]
                    RMSEP = math.sqrt(sum(SEP)/len(SEP))
                    #print RMSEP

                    SEP2 = [(x - y)**2 for x, y in zip(rawcountsList, checkRMSEList)]
                    RMSEP2 = math.sqrt(sum(SEP2)/len(SEP2))
                    #print RMSEP2

                    if RMSEP <= RMSEP2:
                           print "Approved..."
                    else:
                           print "Rejected. Return to model building..."


        return _dict_counts

def main():
        ## first group
        start_time = timeit.default_timer()
        group1_f = args.group1
        file_gtf=open(name_gtf,'r')
        gff_file=HTSeq.GFF_Reader(file_gtf)
        
        ##### Sanity Check
        samfile=HTSeq.SAM_Reader(group1_f[0])
        is_chr_sam = None
        set_chr_gff = set()
        for almnt in samfile:
                is_chr_sam = almnt.iv.chrom
                break
        for feature in gff_file:
                #if not "chr" in feature.iv.chrom and "chr" in is_chr_sam:
                #       feature.iv.chrom = 'chr'+feature.iv.chrom
                set_chr_gff.add(feature.iv.chrom)
        
        if  is_chr_sam not in set_chr_gff:
                sys.stderr.write("Error: Chromosome ID in SAM files and GFF3 file do not agree.\n")
                sys.exit(1)
        #####

        if MODE == "chrom":
                print "Writing new GFF3 file..."
                bad_words = ['chrC','chrM','Pt','Mt','chrMt', 'chrPt']
                with open(name_gtf) as oldfile, open(name_gtf.split(".",1)[0]+'_splicesim.gff3', 'w') as newfile:
                        for line in oldfile:
                                if not any(bad_word in line.split("\t",1)[0] for bad_word in bad_words):
                                        newfile.write(line)
                newfile.close()
                name_gtfnew = name_gtf.split(".",1)[0]+'_splicesim.gff3'
        
        #####

        if MODE == "chrom":
                file_gtf=open(name_gtfnew,'r')
        elif MODE == "all":
                file_gtf=open(name_gtf,'r')
        gff_file=HTSeq.GFF_Reader(file_gtf)
        counts1=meanVar(group1_f, gff_file,'group1')
        file_gtf.close()
        ## second group
        
        group2_f = args.group2
        if MODE == "chrom":
                file_gtf=open(name_gtfnew,'r')
        elif MODE == "all":
                file_gtf=open(name_gtf,'r')
        gff_file=HTSeq.GFF_Reader(file_gtf)
        counts2=meanVar(group2_f, gff_file,'group2')
        
        file_gtf.close()
        merged_count = {k1: v1+v2 for (k1,v1) in counts1.iteritems() for (k2,v2) in counts2.iteritems() if k1 == k2}
        non0_exp_list = list()
        for k,v in sorted(merged_count.iteritems()):
                #print k,v
                value=np.array(v)
                #ensures all counts must not be 0 for both counts1 and counts2 and at least 10
                #value[value ==0] produces 1s for any zeroes and 0s for any nonzeroes
                if CUTOFF == 'h':
                        if len(value[value <10]) == 0:
                                        #remove pseudogenes and pt, mt genes from consideration
                                        if "transcript" not in k and "ATCG" not in k and "ATMG" not in k:
                                                non0_exp_list.append(k)
                elif CUTOFF == 'l':
                        if len(value[value==0]) == 0:
                                        #remove pseudogenes and pt, mt genes from consideration
                                        if "transcript" not in k and "ATCG" not in k and "ATMG" not in k:
                                                non0_exp_list.append(k)        
        num_non0 = len(non0_exp_list)
        sys.stderr.write("Randomly choose %d out of %d genes to be assumed to undergo AS\n" % (NTARG,num_non0))
        if NTARG > num_non0 and num_non0 >= 100:
                print "Simulated genes are not enough for the required number of AS genes."
                newtarg = int(np.floor(num_non0/100)*100)
                print "Running %d number of AS genes instead..." %newtarg
                l=random.sample(xrange(num_non0), newtarg)
        elif NTARG > num_non0 and num_non0 < 100:
                print "Simulated genes are not enough for the required number of AS genes."
                newtarg = int(np.floor(num_non0/10)*10)
                print "Running %d number of AS genes instead..." %newtarg
                l=random.sample(xrange(num_non0), newtarg)
        else:
                l=random.sample(xrange(num_non0), NTARG)
                
        out1 = open('All_genes_list.txt','w')
        for i in non0_exp_list:
            if "gene" in i or "transcript" in i:
                sep = ':'
                temp_gene = i
                rest = temp_gene.split(sep,1)[1]
                if '.' in rest:
                        sep2 = '.'
                        rest2 = rest.split(sep2,1)[0]
                        out1.write(rest2+"\n")
                else:
                        out1.write(rest+"\n")
            else:
                out1.write(i+"\n")
        out1.close()
        
        out = open('AS_genes_list.txt','w')
        for i in l:
                #if "gene" in non0_exp_list[i]:
                if "gene" in non0_exp_list[i] or "transcript" in non0_exp_list[i]:
                        sep = ':'
                        temp_gene = non0_exp_list[i]
                        rest = temp_gene.split(sep,1)[1]
                        if '.' in rest:
                                sep2 = '.'
                                rest2 = rest.split(sep2,1)[0]
                                out.write(rest2+"\n")
                        else:
                                out.write(rest+"\n")
                else:
                        out.write(non0_exp_list[i]+"\n")
        out.close()
        stop_time = timeit.default_timer()
        total_time = stop_time - start_time
        print "Completed GenCount. Total time lapsed: ", total_time
if __name__ == '__main__': main()
