# coding: utf-8

# In[2]:

# Import
import pandas
# Alignement
import nltk
import codecs
import gensim
from scipy import sparse
import itertools
from sklearn.feature_extraction.text import CountVectorizer
import scipy.sparse
import scipy.io
from sklearn import preprocessing
from keras.models import Sequential
from keras.layers.core import Dense, Dropout, Activation,AutoEncoder
from keras.optimizers import SGD,Adam
from keras.layers import containers
from mlp import *
import mlp
import sklearn.metrics
import shelve
import pickle
from utils import *
import sys
import json
# In[4]:

db=shelve.open("{}.shelve".format(sys.argv[2]),writeback=True)

sparse_model=shelve.open("{}.shelve".format(sys.argv[1]))
#['ASR', 'TRS', 'LABEL']
# In[6]:

do_do=False
try:
    do_do = True if sys.argv[3] == 1 else False
    hidden_size = int(sys.argv[4]) if sys.argv[4] else 100
    w1_size = int(sys.argv[5]) if sys.argv[5] else 300

except IndexError :
    do_do = False
    hidden_size = 100
    w1_size = 300


ASR=sparse_model["ASR"]
TRS=sparse_model["TRS"]
LABEL=sparse_model["LABEL"]

db["ASR_SPARSE"]=ASR
db["TRS_SPARSE"]=TRS
db["LABEL"]=LABEL
print "todo label"
def select(elm):
    return int(elm.split("_")[-1])
#z.apply(select)
label_bin={}
lb = preprocessing.LabelBinarizer(neg_label=0)
lb.fit(LABEL["TRAIN"].apply(select))
for i in ASR.keys():
    label_bin=lb.transform(LABEL[i].apply(select))

input_activation="tanh"
out_activation="tanh"
loss="mse"
epochs=500
batch=1
patience=60
sgd = Adam(lr=0.0001)#SGD(lr=0.00001,nesterov=False) #'rmsprop'# Adam(lr=0.00001)#SGD(lr=0.001, momentum=0.9, nesterov=True)
try :
    sgd_repr=sgd.get_config()
except AttributeError :
    sgd_repr=sgd
json.dump({ "h1" : hidden_size,
	"inside_activation" : input_activation,
	"out_activation" : out_activation,
        "hidden_size" : hidden_size,
        "w1_size" : w1_size,
        "do dropout": do_do,
	"loss" : loss,
	"epochs" : epochs ,
	"batch_size" : batch,
	"patience" : patience,
        "sgd" : sgd_repr},
	open("{}.json".format(sys.argv[2]),"w"),
	indent=4)
print "gogo autoencoder ASR"
autoencode=Sequential()
autoencode.add(Dense(hidden_size,input_dim=ASR["TRAIN"].shape[1],init='glorot_uniform',activation=input_activation))
if do_do :
    autoencode.add(Dropout(0.5))
autoencode.add(Dense(ASR["DEV"].todense().shape[1],input_dim=hidden_size,init="glorot_uniform",activation=out_activation))

#autoencode.compile(optimizer=sgd,loss=loss)

autoencode.compile(optimizer=sgd,loss=loss)


# In[ ]:

autoencode.fit(ASR["TRAIN"].todense(),ASR["TRAIN"].todense(),nb_epoch=epochs,batch_size=batch,
               callbacks=[keras.callbacks.EarlyStopping(monitor='val_loss',
                   patience=patience, verbose=0)],validation_data=(ASR["DEV"].todense(),ASR["DEV"].todense()),verbose=1)


# In[ ]:

auto_decoder=Sequential()
auto_decoder.add(Dense(hidden_size,input_dim=ASR["DEV"].todense().shape[1],init='uniform',activation=input_activation,weights=autoencode.get_weights()[:2]))
auto_decoder.compile(optimizer=sgd,loss=loss)


# In[77]:

#autoencode.predict(ASR["DEV"].todense())


# In[ ]:

print "auto encoder et auto decoder asr okay"

ASR_AE_H1={}
for i in ASR.keys():
    ASR_AE_H1[i]=auto_decoder.predict(ASR[i].todense())
    #TRS[i]=dico.transform(TRS[i][2])

db["ASR_AE_H1"]=ASR_AE_H1

print "auto encoder trs learning"
# In[68]:/
autoencode_trs=Sequential()
autoencode_trs.add(Dense(hidden_size,input_dim=TRS["DEV"].todense().shape[1],init='glorot_uniform',activation=input_activation))
autoencode_trs.add(Dense(TRS["DEV"].todense().shape[1],input_dim=hidden_size,init="glorot_uniform",activation=out_activation))

#autoencode_trs.compile(optimizer=sgd_trs,loss=loss)

autoencode_trs.compile(optimizer=sgd,loss=loss)


# In[69]:

autoencode_trs.fit(TRS["TRAIN"].todense(),TRS["TRAIN"].todense(),nb_epoch=epochs,batch_size=batch,
               callbacks=[keras.callbacks.EarlyStopping(monitor='val_loss',
                                                        patience=patience, verbose=0)],
               validation_data=(TRS["DEV"].todense(),TRS["DEV"].todense()),verbose=1)


# In[87]:



auto_decoder_trs=Sequential()
auto_decoder_trs.add(Dense(hidden_size,input_dim=ASR["DEV"].todense().shape[1],activation=input_activation,weights=autoencode_trs.get_weights()[:2]))
auto_decoder_trs.compile(optimizer=sgd,loss=loss)


# In[88]:
print "auto encoder trs okay"
TRS_AE_H1={}

for i in TRS.keys():
    TRS_AE_H1[i]=auto_decoder_trs.predict(TRS[i].todense())
    #TRS[i]=dico.transform(TRS[i][2])

db["TRS_AE_H1"]=TRS_AE_H1


db.sync()




# In[261]:

#pred_dev= model_TRS_AE.predict(TRS_AE["DEV"],batch_size=1)

TRS_AE={}
ASR_AE={}
for i in TRS.keys():
    TRS_AE[i]=autoencode_trs.predict(TRS[i].todense())
    ASR_AE[i]=autoencode.predict(ASR[i].todense())


db["TRS_AE_OUT"]=TRS_AE
db["ASR_AE_OUT"]=ASR_AE

db.sync()
# # Transfert de couche
# ICI
# In[138]:
print "learn transform ae H1({})".format(hidden_size)
model_TRANS = Sequential()
model_TRANS.add(Dense(hidden_size,input_dim=ASR["TRAIN"].shape[1],init='glorot_uniform',activation=input_activation,weights=autoencode.get_weights()[:2]))
model_TRANS.add(Dense( w1_size,input_dim=hidden_size, init='glorot_uniform', activation=input_activation))
model_TRANS.add(Dense( hidden_size,input_dim=w1_size, init='glorot_uniform', activation=input_activation))
model_TRANS.add(Dense(TRS["TRAIN"].shape[1],input_dim=hidden_size,init='glorot_uniform',activation=out_activation,weights=autoencode_trs.get_weights()[2:4]))
sgd_TRANS = SGD(lr=0.01, decay=1e-4, momentum=0.9, nesterov=True)
#model_TRANS.compile(loss='mse', optimizer=sgd_TRANS)

model_TRANS.compile(loss='mse', optimizer=sgd)


# In[146]:

model_TRANS.fit(ASR["TRAIN"].todense(),TRS["TRAIN"].todense(),nb_epoch=epochs,batch_size=batch,
               callbacks=[keras.callbacks.EarlyStopping(monitor='val_loss',
                                                        patience=patience, verbose=0)],
               validation_data=(ASR["DEV"].todense(),TRS["DEV"].todense()),verbose=1)


# In[140]:
print "make trans projection H1"
asr_transformer={}
for i in ASR_AE.keys():
    asr_transformer[i]=model_TRANS.predict(ASR[i].todense())

db["ASR_TRANFORMED_OUT"]=asr_transformer
# In[ ]:

db.sync()


auto_decoder_trans=Sequential()
auto_decoder_trans.add(Dense(hidden_size,input_dim=ASR["TRAIN"].shape[1],activation=input_activation,weights=model_TRANS.get_weights()[:2]))
auto_decoder_trans.compile(optimizer=sgd,loss=loss)

asr_trans_h1={}
for i in ASR_AE.keys():
    asr_trans_h1[i]=auto_decoder_trans.predict(ASR[i].todense())
db["ASR_H1_TRANSFORMED"]=asr_trans_h1
print "shape",ASR_AE["TRAIN"].shape[1]




auto_decoder_trans=Sequential()
auto_decoder_trans.add(Dense(hidden_size,input_dim=ASR["TRAIN"].shape[1],activation=input_activation,weights=model_TRANS.get_weights()[:2]))
auto_decoder_trans.add(Dense(w1_size,input_dim=hidden_size,activation=input_activation,weights=model_TRANS.get_weights()[2:4]))
auto_decoder_trans.compile(optimizer=sgd,loss=loss)

asr_trans_w1={}
for i in ASR_AE.keys():
    asr_trans_w1[i]=auto_decoder_trans.predict(ASR[i].todense())
db["ASR_W1_TRANSFORMED"]=asr_trans_w1
print "shape",ASR_AE["TRAIN"].shape[1]

auto_decoder_trans=Sequential()
auto_decoder_trans.add(Dense(hidden_size,input_dim=ASR["TRAIN"].shape[1],activation=input_activation,weights=model_TRANS.get_weights()[:2]))
auto_decoder_trans.add(Dense(w1_size,input_dim=hidden_size,activation=input_activation,weights=model_TRANS.get_weights()[2:4]))
auto_decoder_trans.add(Dense(hidden_size,input_dim=w1_size,activation=input_activation,weights=model_TRANS.get_weights()[4:6]))
auto_decoder_trans.compile(optimizer=sgd,loss=loss)

asr_trans_h2={}
for i in ASR_AE.keys():
    asr_trans_h2[i]=auto_decoder_trans.predict(ASR[i].todense())
db["ASR_H2_TRANSFORMED"]=asr_trans_h2
print "shape",ASR_AE["TRAIN"].shape[1]



db.sync()
db.close()