'''
File to handle motif/expression regression
'''

__author__='Anthony Soltis'
__email__='asoltis@mit.edu'

import sys,os,pickle,re
from optparse import OptionParser
import numpy as np
from scipy import stats
import fileinput
import matplotlib
matplotlib.use('pdf')
from matplotlib import pyplot as plt

def load_tgm(tgm_fn):
    '''
    Load tgm file and produce output matrix.
    Output is transposed numpy array object.
    '''
    print 'Loading tgm file...'
    tgm = []
    for line in fileinput.input(tgm_fn):
        l = line.strip('\n').split()
        tgm.append(l)

    # display results, return array
    s = np.asarray(tgm).T.shape
    print 'TGM file loaded with %d genes by %d motifs.'%(s[0],s[1])
    return np.asarray(tgm).T

def load_ids(ids_fn):
    '''
    Load ids filename and store as list.
    '''
    ids = []
    for line in fileinput.input(ids_fn):
        l = line.strip('\n')
        ids.append(l)

    return ids

def load_response(data_fn):
    '''
    Load ydata and return numpy vector.
    Input file should have one value per-row.
    '''
    r_data = []
    r_genes = []
    for line in fileinput.input(data_fn):
        row=line.strip('\n').split('\t')
        if len(row)>1:
            r_genes.append(row[0])
            r_data.append(float(row[1]))
        else:
            r_data.append(float(row[0]))
#        r_data.append(float(line.strip('\n')))
    
    print 'Response data file loaded with %d values.'%(len(r_data))
    return np.asarray(r_data),r_genes

def map_data(Xdata,Xnames,Ydata,Ynames):
    '''
    Map X (predictor) data to Y (response) data using X and Y data ids (i.e. gene names).
    '''
    # Intersect two gene lists
    Xinds = []
    Yinds = []
    #yn = []
    for i,Xgene in enumerate(Xnames):
        for j,Ygene in enumerate(Ynames):
            if Xgene == Ygene:
                Xinds.append(i)
                Yinds.append(j)
    #            yn.append(Ygene)
    Xdata_out = Xdata[Xinds,:]
    Ydata_out = Ydata[Yinds]
    print 'Found %d genes that have binding data and are in the expression output'%(len(Yinds))

    return Xdata_out,Ydata_out

def perform_regression(X,Y,motif_ids,norm,outdir,plot):
    '''
    Linear regression to assess whether X (predicted TF binding affinity)
    is predictive of Y (gene expression).  Optionally stores figures showing
    the regression results for each TF motif.
    '''
    reg_results = []
    for i in range(0,X.shape[1]):
        # Set up data
        x = np.array(X[:,i],dtype=float)
        if norm != None:
            if norm == 'log2':
                y = np.log2(Y+.1)
            elif norm == 'log10':
                y = np.log10(Y+.1)
        else: y = Y

        # Perform regression
        slope,intercept,r_val,p_val,std_err = stats.linregress(x,y)
        reg_results.append(([motif_ids[i],slope,p_val,i]))

	#regression plot
	if plot:
	    fig = plt.figure()
	    ax1 = fig.add_subplot(111)
	    ax1.plot(x,y,'bo',x,intercept+slope*x,'k')
	    ax1.set_title(motif_ids[i])
	    ax1.set_xlabel('Estimated transcription factor affinity')
	    ax1.set_ylabel('Expression log fold change')
	    #checking if a subdirectory is present to save plots
	    plotdir = os.path.join(os.path.split(outdir)[0],'regression_plots')
	    
	    if not os.path.isdir(plotdir):
			os.makedirs(plotdir)
	    #cleaning all motif ids to have all alphanumeric name
	    if not re.match(r'^[A-Za-z0-9.]*$', motif_ids[i]):
			motif_ids[i] = "".join(c for c in motif_ids[i] if c not in ('!','$','@','!','%','*','\\','/','_','-'))	    
	    #file name must be within max characters supported by os 
	    if len(motif_ids[i])>162:
			st = motif_ids[i]
			motif_ids[i] = st[0:160]
	    plotfile = os.path.join(plotdir,motif_ids[i]+'.pdf')
	    fig.savefig(open(plotfile,'w'),dpi=300)
	    plt.close()	    

    # First sort by p-value, then by magnitude of the regression slope,
    # then by motif id to break ties
    # Test case uses the values written to file by the str() conversion
    # so sort on the same truncated values
    return sorted(reg_results, key=lambda reg_res: (float(str(reg_res[2])), -float(str(abs(reg_res[1]))), reg_res[0]))

def fdr_correction(results):
    '''
    Compute FDR corrected p-values based on Benjamini-Hochberg procedure.
    '''
    new_results = []
    num_tests = len([r for r in results if str(r[1])!='nan'])
    print 'Correcting for '+str(num_tests)+' numeric values'
    for i in range(0,num_tests):
        tup = results[i]
        pval = tup[2]
        fdr = num_tests*pval/(i+1)
        if fdr > 1.0: fdr = 1.0
        tup+=(fdr,)
        new_results.append(tup)

    return new_results

def main():
    
    usage = "%prog [options] <scores.tgm or scores.tgm.pkl> <response_values.tab>"
    description = "Script that takes a predicted TF-Gene matrix and uses a linear regression to identify which TFs have binding scores correlated with gene expression changes."
    parser = OptionParser(usage=usage,description=description)
    
    # Options
    parser.add_option('--outdir','--out',dest="outdir",default='./test_out.txt',
                      help='Choose output file name. Default is %default.')
    parser.add_option('--motif-ids','--motif-ids',dest='motif_ids',default=None,
                      help='OPTIONAL: If input file is in text format (.tgm), provide motif ids corresponding to tgm file motifs.')
    parser.add_option('--tgm-genes',dest='tgm_genes',default=None,
                      help='OPTIONAL: If input file is in text format (.tgm), provide gene ids corresponding to tgm file genes.')
    parser.add_option('--response-genes',dest='response_genes',default=None,
                      help='OPTIONAL: If two-column file is not provided, add in gene ids corresponding to response values.')
    parser.add_option('--norm-type',dest='norm_type',default=None,
                      help='Choose normalization type for response data. Choices are: "log2", "log10".\
                            Default is %default.')    
    parser.add_option('--use-qval',dest='use_qval',action='store_true',default=False,help='If set this the Forest input file will contain -log(qval) instead of -log(pval) and threshold the output using qval. Default:%default')
    parser.add_option('--thresh',dest='thresh',type='string',default='0.9',help='P/Q-Value threshold to illustrate results. Default:%default')
    # Cannot set the default here because it depends on opts.outdir
    parser.add_option('--gifdir',dest='motifs',default=None,
                      help='Directory containing motif GIFs to illustrate results. Default is <path_to_this_script>/../data/matrix_files/gifs.')
    parser.add_option('--plot',dest='plot',action='store_true',default=False,help='Enable plot generation for regression results. Default:%default')

    # get options, arguments
    (opts,args) = parser.parse_args()

    # Handle arguments
    tgm_fn = args[0]
    response_data_fn = args[1]

    # Load in Y-vector data (gene expression, fold-changes, etc.)
    response_data,response_genes = load_response(response_data_fn)

    print 'Trying to get file type...'
    ext=tgm_fn.split('.')[-1]
    if ext.lower()=='pkl':
        print '...found PKL file'
        pkl=True
    else:
        print '...found text file, looking for additional data files in options'
        pkl=False

    # Handle options
    outdir = opts.outdir
    motif_ids = opts.motif_ids