import sys
import rpy
from pylab import gca, figure, subplot,plot
from scipy import *
from scipy.linalg import norm
from lpls import correlation_loadings

import rpy_go
sys.path.append("../../fluents") # home of dataset
sys.path.append("../../fluents/lib") # home of cx_stats
sys.path.append("/home/flatberg/fluents/scripts/lpls")
import dataset
import cx_stats
from engines import nipals_lpls
from validation import lpls_val, lpls_jk
from plots_lpls import plot_corrloads, plot_dag
import plots_lpls

# Possible outliers
# http://www.pubmedcentral.nih.gov/articlerender.fcgi?tool=pubmed&pubmedid=16817967
sample_outliers = ['OV:NCI_ADR_RES', 'CNS:SF_295', 'CNS:SF_539', 'RE:SN12C', 'LC:NCI_H226', 'LC:NCI_H522', 'PR:PC_3', 'PR:DU_145']


####### OPTIONS ###########
# data
chip = "hgu133a"
use_data = 'uma'
subset = 'plsr'
small_test = False
use_sbg_subset = True # the sandberg nci-Ygroups subset
std_y = True
std_z = False
# go
ontology = "bp"
min_genes = 5
similarities = ("JiangConrath","Resnik","Lin","CoutoEnriched","CoutoJiangConrath","CoutoResnik","CoutoLin")
meth = similarities[2]
go_term_sim = "OA"
# lpls
a_max = 5
aopt = 2
xz_alpha = .4
w_alpha = .3
mean_ctr = [2, 0, 2]
nsets = None
qval_cutoff = 0.01
n_iter = 200

alpha_check = False
calc_rmsep = False


######## DATA ##########
if use_data=='smoker':
    # full smoker data
    DX = dataset.read_ftsv(open("../../data/smokers-full/Smokers.ftsv"))
    DY = dataset.read_ftsv(open("../../data/smokers-full/Yg.ftsv"))
    Y = DY.asarray().astype('d')
    gene_ids = DX.get_identifiers('gene_ids', sorted=True)
elif use_data=='scherf':
    DX = dataset.read_ftsv(open("../../data/scherf/scherfX.ftsv"))
    DY = dataset.read_ftsv(open("../../data/scherf/scherfY.ftsv"))
    Yg = DY.asarray().astype('d')
    gene_ids = DX.get_identifiers('gene_ids', sorted=True)
elif use_data=='staunton':
    pass
elif use_data=='uma':
    DX = dataset.read_ftsv(open("../../data/uma/X133.ftsv"))
    DYg = dataset.read_ftsv(open("../../data/uma/Yg133.ftsv"))
    DY = dataset.read_ftsv(open("../../data/uma/Yd.ftsv"))
    Y = DY.asarray().astype('d')
    Yg = DYg.asarray().astype('d')
    gene_ids = DX.get_identifiers('gene_ids', sorted=True)

# use subset with defined GO-terms
ic_all = 2026006.0 # sum of all ic in BP
max_ic = -log(1/ic_all)
ic_cutoff = -log(min_genes/ic_all)/max_ic
print "Information cutoff for min %d genes: %.2f" %(min_genes, ic_cutoff)

gene2goterms = rpy_go.goterms_from_gene(gene_ids, ic_cutoff=ic_cutoff)
all_terms = set()
for t in gene2goterms.values():
    all_terms.update(t)
terms = list(all_terms)
print "\nNumber of go-terms: %s" %len(terms)

# update genelist
gene_ids = gene2goterms.keys()
print "\nNumber of genes: %s" %len(gene_ids)

X = DX.asarray()
index = DX.get_indices('gene_ids', gene_ids)
X = X[:,index]

# Use only subset defined on GO
ontology = 'BP'
print "\n\nFiltering genes by Go terms "

# use subset based on SAM,PLSR or (IQR)

if subset=='sam':
    # select subset genes by SAM
    rpy.r.library("siggenes")
    rpy.r.library("qvalue")
    rpy.r.assign("data", X.T)
    cl = dot(DYg.asarray(), diag(arange(Yg.shape[1])+1)).sum(1)
    rpy.r.assign("cl", cl)
    rpy.r.assign("B", 20)
    # Perform a SAM analysis.
    print "Starting SAM"
    sam = rpy.r('sam.out<-sam(data=data,cl=cl,B=B,rand=123)')
    print "SAM done"
    # Compute the q-values of the genes.
    qq = rpy.r('qobj<-qvalue(sam.out@p.value)')
    qvals = asarray(qq['qvalues'])
    # cut off
    cutoff = 0.001
    index = where(qvals<cutoff)[0]
    if small_test:
        index = index[:20]
    
    # Subset data
    X = X[:,index]
    
    gene_ids = [gid for i, gid in enumerate(gene_ids) if i in index]
    print "\nNumber of genes: %s" %len(gene_ids)
    print "\nWorking on subset with %s genes " %len(gene_ids)

    # update valid go-terms
    gene2goterms = rpy_go.goterms_from_gene(gene_ids, ic_cutoff=ic_cutoff)
    all_terms = set()
    for t in gene2goterms.values():
        all_terms.update(t)
    terms = list(all_terms)
    print "\nNumber of go-terms: %s" %len(terms)
    # update genelist
    gene_ids = gene2goterms.keys()
    print "\nNumber of genes: %s" %len(gene_ids)
elif subset=='plsr':
    cx_stats.pls_qvals(X, Y)
else:
    # noimp (smoker data is prefiltered)
    pass


rpy.r.library("GOSim")
# Go-term similarity matrix

print "Term-term similarity matrix (method = %s)" %meth
print "\nCalculating term-term similarity matrix"

if meth=="CoutoEnriched":
    aa = 0
    ba = 0
    rpy.r.setEnrichmentFactors(alpha = aa, beta =ba) 
rpytmat = rpy.with_mode(rpy.NO_CONVERSION, rpy.r.getTermSim)(terms, method=meth,verbose=False)
tmat = rpy.r.assign("haha", rpytmat)
print "\n Calculating Z matrix"
Z = rpy_go.genego_sim(gene2goterms,gene_ids,terms,rpytmat,go_term_sim=go_term_sim,term_sim=meth)

# update data (X) matrix
newind = DX.get_indices('gene_ids', gene_ids)
Xr = DX.asarray()[:,newind]


######## LPLSR ######## 
print "LPLSR ..."
Y = Yg

if use_sbg_subset:
    Y_old = Y.copy()
    Xr_old = Xr.copy() 
    keep_samples = ['CN', 'ME', 'LE', 'CO', 'RE']
    sample_ids = DY.get_identifiers('cline', sorted=True)
    keep_ind = [i for i,name in enumerate(sample_ids) if name[:2] in keep_samples]
    Xr = Xr[keep_ind,:]
    Y = Y[keep_ind,:]
    Y = Y[:, where(Y.sum(0)>1)[0]]
    

# standardize Z?
sdtz = False
if sdtz:
    Z = Z/Z.std(0)

sdty = True
if sdty:
    Y = Y/Y.std(0)
lpls_result = nipals_lpls(Xr,Y,Z, a_max,alpha=xz_alpha,mean_ctr=mean_ctr)
globals().update(lpls_result)

# Correlation loadings
dx,Rx,rssx = correlation_loadings(Xr, T, P)
dx,Ry,rssy = correlation_loadings(Y, T, Q)
cadz,Rz,rssz = correlation_loadings(Z.T, W, L)

# Prediction error
if calc_rmsep:
    rmsep , yhat, class_error = lpls_val(Xr, Y, Z, a_max, alpha=xz_alpha,mean_ctr=mean_ctr)


if alpha_check:  
    Alpha = arange(0.01, 1, .1)
    Rmsep,Yhat, CE = [],[],[]
    for a in Alpha:
        print "alpha %f" %a
        rmsep , yhat, ce = lpls_val(Xr, Y, Z, a_max, alpha=a,mean_ctr=mean_ctr,nsets=nsets)
        Rmsep.append(rmsep.copy())
        #Yhat.append(yhat.copy())
        #CE.append(ce.copy())
    Rmsep = asarray(Rmsep)
    #Yhat = asarray(Yhat)
    #CE = asarray(CE)


# Significance Hotellings T
#Wx, Wz = lpls_jk(Xr, Y, Z, aopt, mean_ctr=mean_ctr, xz_alpha=xz_alpha, nsets=nsets)
#Ws = W*apply_along_axis(norm, 0, T)
#tsqx = cx_stats.hotelling(Wx, Ws[:,:aopt], alpha=w_alpha)
#tsqz = cx_stats.hotelling(Wz, L[:,:aopt], alpha=0)

# qvals
cal_tsq_z, pert_tsq_z, cal_tsq_x, pert_tsq_x = cx_stats.lpls_qvals(Xr, Y, Z, aopt=aopt, zx_alpha=xz_alpha, n_iter=n_iter)

qvalz = cx_stats.fdr(cal_tsq_z, pert_tsq_z, 'median')
qvalx = cx_stats.fdr(cal_tsq_x, pert_tsq_x, 'median')


# p-values, set-enrichment analysis
active_genes_ids = where(qvalx < qval_cutoff)[0]
active_genes = [name for i,name in enumerate(gene_ids) if i in active_genes_ids]
active_universe = gene_ids
gsea_result, gsea_params= rpy_go.gene_GO_hypergeo_test(genelist=active_genes,universe=active_universe,chip=chip,pval_cutoff=1.0,cond=False,test_direction="over")
active_goterms_ids = where(qvalz < qval_cutoff)[0]
active_goterms = [name for i,name in enumerate(terms) if i in active_goterms_ids]

gsea_t2p = dict(zip(gsea_result['GOBPID'], gsea_result['Pvalue']))




#### PLOTS ####

from pylab import *
from scipy import where
dg = plots_lpls.dag(terms, "bp")
pos = None

figure(300)
subplot(2,1,1)
pos = plots_lpls.plot_dag(dg, node_color=cal_tsq_z, pos=pos, nodelist=terms)
subplot(2,1,2)
pos = plot_dag(dg, node_color=qvalz, pos=pos, nodelist=terms)


if calc_rmsep:
    figure(1) #rmsep
    bar_w = .2
    bar_col = 'rgb'*5
    m = Y.shape[1]
    for a in range(m):
        bar(arange(a_max)+a*bar_w+.1, rmsep[:,a], width=bar_w, color=bar_col[a])
    ylim([rmsep.min()-.05, rmsep.max()+.05]) 
    title('RMSEP: Y(%s)' %Y.get_name())

#figure(2)
#for a in range(m):
#    bar(arange(a_max)+a*bar_w+.1, class_error[:,a], width=bar_w, color=bar_col[a])
#ylim([class_error.min()-.05, class_error.max()+.05])
#title('Classification accuracy')

figure(3) # Hyploid correlations
tsqz = cal_tsq_z
tsqx = cal_tsq_x
tsqz_s = 250*tsqz/tsqz.max()
td = rpy_go.goterm2desc(terms)
tlabels = [td[i] for i in terms]
keep = where(qvalz<0.01)[0]
k_Rz = Rz[keep,:]
k_tsqz_s = tsqz_s[keep]
k_tsq = tsqz[keep]
k_tlabels = [name for i,name in enumerate(tlabels) if i in keep]
plot_corrloads(Rz, pc1=0, pc2=1, s=tsqz_s, c=tsqz, zorder=5, expvar=evz, ax=None,alpha=.5,labels=None)
#plot_corrloads(k_Rz, pc1=0, pc2=1, s=k_tsqz_s, c=k_tsqz, zorder=5, expvar=evz, ax=None,alpha=.5,labels=None)
ax = gca()
yglabels = DYg.get_identifiers(DYg.get_dim_name()[1], sorted=True)
ylabels = DY.get_identifiers(DY.get_dim_name()[1], sorted=True)
blabels = yglabels[:]
blabels.append(ylabels[0])
plot_corrloads(Ry, pc1=0, pc2=1, s=150, c='g', marker='s', zorder=5, expvar=evy, ax=ax,labels=None,alpha=.9)
plot_corrloads(Rx, pc1=0, pc2=1, s=5, c='k', zorder=1, expvar=evx, ax=ax)


figure(4)
subplot(221)
ax = gca()
plot_corrloads(Rx, pc1=0, pc2=1, s=tsqx/2.0, c='b', zorder=5, expvar=evx, ax=ax)
# title('X correlation')
subplot(222)
ax = gca()
plot_corrloads(Ry, pc1=0, pc2=1, s=150, c='g', zorder=5, expvar=evy, ax=ax)
#title('Y correlation')
subplot(223)
ax = gca()
plot_corrloads(Rz, pc1=0, pc2=1, s=tsqz/10.0, c='r', zorder=5, expvar=evz, ax=ax)
#title('Z correlation')
subplot(224)
plot(arange(len(evx)), evx, 'b', label='X', linewidth=2)
plot(evy, 'g', label='Y', linewidth=2)
plot(evz, 'r', label='Z', linewidth=2)
legend(loc=2)
ylabel('Explained variance')
xlabel('Component')
show()