import sys
import rpy
from pylab import gca, figure, subplot
from scipy import *
from lpls import *
import rpy_go
sys.path.append("../../fluents") # home of dataset
sys.path.append("../../fluents/lib") # home of cx_stats
import dataset
import cx_stats
from plots_lpls import plot_corrloads

######## DATA ##########
use_data='uma'
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/Scherf.ftsv"))
    DY = dataset.read_ftsv(open("../../data/scherf/Yd.ftsv"))
    Y = 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"))
    DY = dataset.read_ftsv(open("../../data/uma/Yg133.ftsv"))
    Y = DY.asarray().astype('d')
    gene_ids = DX.get_identifiers('gene_ids', sorted=True)

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

gene2goterms = rpy_go.goterms_from_gene(gene_ids)
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)

# use subset based on SAM or IQR
subset = 'm'
if subset=='sam':
    # select subset genes by SAM
    rpy.r.library("siggenes")
    rpy.r.library("qvalue")
    data = DX.asarray().T
    # data = data[:100,:]
    rpy.r.assign("data", data)
    cl = dot(DY.asarray(), diag([1,2,3])).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]
    # Subset data
    X = DX.asarray()
    #Xr = X[:,index]
    gene_ids = DX.get_identifiers('gene_ids', index)
    print "\nWorking on subset with %s genes " %len(gene_ids)
else:
    # noimp (smoker data is prefiltered)
    pass


rpy.r.library("GOSim")
# Go-term similarity matrix
methods = ("JiangConrath","Resnik","Lin","CoutoEnriched","CoutoJiangConrath","CoutoResnik","CoutoLin")
meth = methods[2]
print "Term-term similarity matrix (method = %s)" %meth

print "\nCalculating term-term similarity matrix"

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="OA",term_sim=meth)

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


######## LPLSR ########
print "LPLSR ..."
a_max = 5
aopt = 3
xz_alpha = .5
w_alpha = .1
mean_ctr = [2, 0, 2]

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

T, W, P, Q, U, L, K, B, b0, evx, evy, evz = nipals_lpls(Xr,Y,Z, a_max,
                                                        alpha=xz_alpha,
                                                        mean_ctr=mean_ctr)

# 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
rmsep , yhat, class_error = cv_lpls(Xr, Y, Z, a_max, alpha=xz_alpha,mean_ctr=mean_ctr)
alpha_check=False
if alpha_check:
    Alpha = arange(0.01, 1, .1)
    Rmsep,Yhat, CE = [],[],[]
    for a in Alpha:
        rmsep , yhat, ce = cv_lpls(Xr, Y, Z, a_max, alpha=xz_alpha,mean_ctr=mean_ctr)
        Rmsep.append(rmsep)
        Yhat.append(yhat)
        CE.append(ce)
    Rmsep = asarray(Rmsep)
    Yhat = asarray(Yhat)
    CE = asarray(CE)

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


## plots ##
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')

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) # Hypoid correlations
tsqz_s = 250*tsqz/tsqz.max()
plot_corrloads(Rz, pc1=0, pc2=1, s=tsqz_s, c='b', zorder=5, expvar=evz, ax=None)
ax = gca()
ylabels = DY.get_identifiers('_status', sorted=True)
plot_corrloads(Ry, pc1=0, pc2=1, s=150, c='g', zorder=5, expvar=evy, ax=ax,labels=ylabels)

figure(3)
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()