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 ##########
# full smoker data
DX = dataset.read_ftsv(open("../../data/smokers-full/Xfull.ftsv"))
DY = dataset.read_ftsv(open("../../data/smokers-full/Yg.ftsv"))
Y = DY.asarray()
# 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", 100)
# 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
co = 0.001
index = where(qvals<0.01)[0]

# Subset data
X = DX.asarray()
Xr = X[:,index]
gene_ids = DX.get_identifiers('gene_ids', index)
print "\nWorkiing on subset with %s genes " %len(gene_ids)
### Build GO data ####

print "Go terms ..."
goterms = rpy_go.goterms_from_gene(gene_ids)
terms = set()
for t in goterms.values():
    terms.update(t)
terms = list(terms)
print "Number of go-terms: %s" %len(terms)
rpy.r.library("GOSim")
# Go-term similarity matrix
methods = ("JiangConrath","Resnik","Lin","CoutoEnriched","CoutoJiangConrath","CoutoResnik","CoutoLin")
meth = methods[0]
print "Term-term similarity matrix (method = %s)" %meth
if meth=="CoutoEnriched":
    rpy.r('setEnrichmentFactors(alpha=0.1,beta=0.5)')
print "Calculating term-term similarity matrix"
tmat = rpy.r.getTermSim(terms, verbose=False, method=meth)
# check if all terms where found
nanindex = where(isnan(tmat[:,0]))[0]
keep=[]
has_miss = False
if len(nanindex)>0:
    has_miss = True
    print "Some terms missing in similarity matrix"
    keep = where(isnan(tmat[:,0])!=True)[0]
    print "Number of nans: %d" %len(nanindex)
    tmat_new = tmat[:,keep][keep,:]
    new_terms = [i for ind,i in enumerate(terms) if ind in keep]
    bad_terms = [i for ind,i in enumerate(terms) if ind not in keep]
    # update go-term dict
    for gene,trm in goterms.items():
        for t in trm:
            if t in bad_terms:
                trm.remove(t)
                if len(trm)==0:
                    print "Removing gene: %s" %gene
                goterms[gene]=trm
    terms = new_terms
    tmat = tmat_new
# Z-matrix
# func (min, max, median, mean, etc),
# func decides on the representation of gene-> goterm when multiple
# goterms exist for one gene
Z, newind = rpy_go.genego_matrix(goterms, tmat, gene_ids, terms,func=mean)
Z = Z.T
# update X matrix (no go-terms available)
Xr = Xr[:,newind]
gene_ids = asarray(gene_ids)[newind]


######## LPLSR ########
print "LPLSR ..."
a_max = 5
aopt = 2
alpha=.6
T, W, P, Q, U, L, K, B, b0, evx, evy, evz = nipals_lpls(Xr,Y,Z, a_max, alpha)

# 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=alpha)

# Significance Hotellings T
Wx, Wz, Wy, = jk_lpls(Xr, Y, Z, aopt)
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])

figure(2) # Hypoid correlations
title('RMSEP')
plot_corrloads(Rz, pc1=0, pc2=1, s=tsqz/10.0, c='b', zorder=5, expvar=evz, ax=None)
ax = gca()
ylabels = DY.get_identifiers('_cat', 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()