# coding: utf-8

# In[2]:

# Import
import gensim
from scipy import sparse
import itertools
from sklearn import preprocessing
from keras.models import Sequential
from keras.optimizers import SGD,Adam
from mlp import *
import mlp
import sklearn.metrics
import shelve
import pickle

from utils import *
import sys
import os
import json
# In[4]:

infer_model=shelve.open("{}".format(sys.argv[2]))
in_dir = sys.argv[1]
#['ASR', 'TRS', 'LABEL']
# In[6]:
if len(sys.argv) > 4 :
    features_key = sys.argv[4]
else :
    features_key = "LDA"

json_conf =json.load(open(sys.argv[3]))

dsae_conf = json_conf["dsae"]

hidden_size= dsae_conf["hidden_size"]
input_activation=dsae_conf["input_activation"]
output_activation=dsae_conf["output_activation"]
loss=dsae_conf["loss"]
epochs=dsae_conf["epochs"]
batch_size=dsae_conf["batch"]
patience=dsae_conf["patience"]
do_do=dsae_conf["do"]
try:
    k = dsae_conf["sgd"]
    if dsae_conf["sgd"]["name"] == "adam":
        sgd = Adam(lr=dsae_conf["sgd"]["lr"])#SGD(lr=0.00001,nesterov=False) #'rmsprop'# Adam(lr=0.00001)#SGD(lr=0.001, momentum=0.9, nesterov=True)
    elif dsae_conf["sgd"]["name"] == "sgd":
        sgd = SGD(lr=dsae_conf["sgd"]["lr"])
except: 
    sgd = dsae_conf["sgd"]

trans_conf = json_conf["dsae"]["transform"]
trans_hidden_size=trans_conf["hidden_size"]
trans_input_activation=trans_conf["input_activation"]
trans_output_activation=trans_conf["output_activation"]
trans_loss=trans_conf["loss"]
trans_epochs=trans_conf["epochs"]
trans_batch_size=trans_conf["batch"]
trans_patience=trans_conf["patience"]
trans_do=trans_conf["do"]
try:
    k = trans_conf["sgd"]
    if trans_conf["sgd"]["name"] == "adam":
        trans_sgd = Adam(lr=trans_conf["sgd"]["lr"])#SGD(lr=0.00001,nesterov=False) #'rmsprop'# Adam(lr=0.00001)#SGD(lr=0.001, momentum=0.9, nesterov=True)
    elif trans_conf["sgd"]["name"] == "sgd":
        trans_sgd = SGD(lr=trans_conf["sgd"]["lr"])
except e : 
    trans_sgd = trans_conf["sgd"]


mlp_conf = json_conf["mlp"]
mlp_h = mlp_conf["hidden_size"]
mlp_loss = mlp_conf["loss"]
mlp_dropouts = mlp_conf["do"]
mlp_epochs = mlp_conf["epochs"]
mlp_batch_size = mlp_conf["batch"]
mlp_input_activation=mlp_conf["input_activation"]
mlp_output_activation=mlp_conf["output_activation"]
try:
    k = mlp_conf["sgd"]
    if mlp_conf["sgd"]["name"] == "adam":
        mlp_sgd = Adam(lr=mlp_conf["sgd"]["lr"])#SGD(lr=0.00001,nesterov=False) #'rmsprop'# Adam(lr=0.00001)#SGD(lr=0.001, momentum=0.9, nesterov=True)
    elif mlp_conf["sgd"]["name"] == "sgd":
        mlp_sgd = SGD(lr=mlp_conf["sgd"]["lr"])
except: 
    mlp_sgd = mlp_conf["sgd"]


name = json_conf["name"]
try:
    os.mkdir("{}/{}".format(in_dir,name))
except:
    pass
db = shelve.open("{}/{}/ae_model.shelve".format(in_dir,name),writeback=True)
#

keys = ["ASR","TRS"]



db["DSAE"] = {}

db["DSAEFT"] = {}
mod = "ASR"
res_tuple_ASR = train_ae(infer_model[features_key][mod]["TRAIN"],
                         infer_model[features_key][mod]["DEV"],
                         infer_model[features_key][mod]["TEST"],
                         hidden_size,dropouts=do_do,
                         patience = patience,sgd=sgd,
                         input_activation=input_activation,
                         output_activation=output_activation,loss=loss,epochs=epochs,
                         batch_size=batch_size,verbose=0,get_weights=True)
mlp_res_list = []
for layer in res_tuple_ASR[0]: 
    mlp_res_list.append(train_mlp(layer[0],infer_model['LABEL'][mod]["TRAIN"],
                                  layer[1],infer_model["LABEL"][mod]["DEV"],
                                  layer[2],infer_model["LABEL"][mod]["TEST"],
                                  mlp_h,loss=mlp_loss,dropouts=mlp_dropouts,
                                  sgd=mlp_sgd,epochs=mlp_epochs,
                                  output_activation=mlp_output_activation,
                                  input_activation=mlp_input_activation,
                                  batch_size=mlp_batch_size,fit_verbose=0))

db["DSAE"][mod] = mlp_res_list
mod = "TRS"
res_tuple_TRS = train_ae(infer_model[features_key][mod]["TRAIN"],
                         infer_model[features_key][mod]["DEV"],
                         infer_model[features_key][mod]["TEST"],
                         hidden_size,dropouts=do_do,
                         sgd=sgd,input_activation=input_activation,
                         output_activation=output_activation,loss=loss,epochs=epochs,
                         batch_size=batch_size,patience=patience,
                         verbose=0,get_weights=True)

mlp_res_list = []
for layer in res_tuple_TRS[0]: 
    mlp_res_list.append(train_mlp(layer[0],infer_model['LABEL'][mod]["TRAIN"],
                                  layer[1],infer_model["LABEL"][mod]["DEV"],
                                  layer[2],infer_model["LABEL"][mod]["TEST"],
                                  mlp_h,loss=mlp_loss,dropouts=mlp_dropouts,
                                  sgd=mlp_sgd,epochs=mlp_epochs,
                                  output_activation=mlp_output_activation,
                                  input_activation=mlp_input_activation,
                                  batch_size=mlp_batch_size,fit_verbose=0))

db["DSAE"][mod] = mlp_res_list



transfert = []

print " get weight trans" 

#for asr_pred, trs_pred in zip(res_tuple_ASR[0], res_tuple_TRS[0]):
 #   print "ASR", [ x.shape for x in asr_pred]

  #  print "TRS", [ x.shape for x in trs_pred]

for asr_pred, trs_pred in zip(res_tuple_ASR[0], res_tuple_TRS[0]):
 #   print "ASR", [ x.shape for x in asr_pred]

  #  print "TRS", [ x.shape for x in trs_pred]
  #  print " TRANS SGD", trans_sgd
    transfert.append( train_ae(asr_pred[0],
                               asr_pred[1],
                               asr_pred[2],
                               trans_hidden_size,
                               dropouts=trans_do,
                               y_train = trs_pred[0],
                               y_dev=trs_pred[1],
                               y_test = trs_pred[2],
                               patience = trans_patience,sgd=trans_sgd,
                               input_activation=trans_input_activation,
                               output_activation=trans_output_activation,
                               loss=trans_loss,
                               epochs=trans_epochs,
                               batch_size=trans_batch_size,verbose=0,get_weights=True) )
mod = "ASR"
mlp_res_bylvl = []
print " MLP on transfert "
for level, w  in transfert  :
    mlp_res_list = []
    for layer in level :
        mlp_res_list.append(train_mlp(layer[0],infer_model['LABEL'][mod]["TRAIN"],
                                      layer[1],infer_model["LABEL"][mod]["DEV"],
                                      layer[2],infer_model["LABEL"][mod]["TEST"],
                                      mlp_h,loss=mlp_loss,dropouts=mlp_dropouts,
                                      sgd=mlp_sgd,epochs=mlp_epochs,
                                      output_activation=mlp_output_activation,
                                      input_activation=mlp_input_activation,
                                      batch_size=mlp_batch_size,fit_verbose=0))
    mlp_res_bylvl.append(mlp_res_list)
db["DSAE"]["transfert"] = mlp_res_bylvl


print " FT " 
WA = res_tuple_ASR[1]
#print "WA", len(WA), [ len(x) for x in WA]
WT = res_tuple_TRS[1]

#print "WT", len(WT), [ len(x) for x in WT]
Wtr = [ x[1] for x in transfert]

#print "Wtr", len(Wtr), [ len(x) for x in Wtr],[ len(x[1]) for x in Wtr]

ft_res = ft_dsae(infer_model[features_key]["ASR"]["TRAIN"],
                 infer_model[features_key]["ASR"]["DEV"],
                 infer_model[features_key]["ASR"]["TEST"],
                 y_train=infer_model[features_key]["TRS"]["TRAIN"],
                 y_dev=infer_model[features_key]["TRS"]["DEV"],
                 y_test=infer_model[features_key]["TRS"]["TEST"],
                 ae_hidden = hidden_size,
                 transfer_hidden = trans_hidden_size,
                 start_weights = WA,
                 transfer_weights = Wtr,
                 end_weights = WT,
                 input_activation = input_activation,
                 output_activation = output_activation,
                 ae_dropouts= do_do,
                 transfer_do = trans_do,
                 sgd =  sgd,
                 loss = loss ,
                 patience = patience,
                 batch_size = batch_size,
                 epochs= epochs)
mlps_by_lvls= []
for level  in ft_res  :
    mlp_res_list = []
    for layer in level :
        mlp_res_list.append(train_mlp(layer[0],infer_model['LABEL'][mod]["TRAIN"],
                                      layer[1],infer_model["LABEL"][mod]["DEV"],
                                      layer[2],infer_model["LABEL"][mod]["TEST"],
                                      mlp_h,loss=mlp_loss,dropouts=mlp_dropouts,
                                      sgd=mlp_sgd,epochs=mlp_epochs,
                                      output_activation=mlp_output_activation,
                                      input_activation=mlp_input_activation,
                                      batch_size=mlp_batch_size,fit_verbose=0))
    mlps_by_lvls.append(mlp_res_list)


db["DSAEFT"]["transfert"] = mlps_by_lvls

db.close()