# -*- coding: utf-8 -*-
import keras
import numpy as np
#from keras.layers.core import Dense, Dropout, Activation 
from keras.optimizers import SGD,Adam
from keras.models import Sequential
from keras.layers import Input, Dense, Dropout
from keras.models import Model
from keras.callbacks import ModelCheckpoint, EarlyStopping
from keras.utils.layer_utils import layer_from_config
from itertools import izip_longest
import tempfile
import shutil
import pandas 
from collections import namedtuple
from sklearn.metrics import accuracy_score as perf
save_tuple = namedtuple("save_tuple",["pred_train","pred_dev","pred_test"])


def ft_dsae(train,dev,test,
        y_train=None,y_dev=None,y_test=None,
        ae_hidden=[20],transfer_hidden=[20],
        start_weights=None,transfer_weights=None,end_weights=None,
        input_activation="tanh", output_activation="tanh",
        init="glorot_uniform",
        ae_dropouts=[None], transfer_do=[None],
        sgd="sgd", loss="mse", patience=5, verbose=0, epochs=5, batch_size=8):

    if not start_weights :
        start_weights = [ None ] * len(ae_hidden)
    if not transfer_weights :
        transfer_weights = [None ] * len(transfer_hidden)
    if not end_weights :
        end_weights = [ None ] * len(end_weights)
    if not transfer_do :
        transfer_do = [0] * len(transfer_hidden) 
    predict_y = True
    if  y_train is None or y_dev is None or y_test is None :
        y_train = train
        y_dev = dev
        y_test = test
        predict_y = False
    param_predict = [ train, dev, test ]
    if predict_y :
        param_predict += [ y_train, y_dev ,y_test ]

    pred_by_level = [] # Contient les prediction par niveaux de transfert 
    layers = [Input(shape=(train.shape[1],))]
    #for w in transfer_weights:
        #print "TW",[ [ y.shape for y in x ]  for x in w] 
    #print "SW",[ [ y.shape for y in x] for x in start_weights]
    #print "EW",[ [ y.shape for y in x ]  for x in end_weights] 
    for cpt in range(1,len(ae_hidden)):
        #print ae_hidden,cpt
        #print cpt, "before" 
        #print "before2", [ [ x.shape for x in y] for y in start_weights[:cpt] ]
        #print "before3", [ [ x.shape for x in y] for y in transfer_weights[cpt]]
        #print "before4", [ [ x.shape for x in y] for y in end_weights[cpt:]]
        sizes = ae_hidden[:cpt] + transfer_hidden + ae_hidden[cpt:]
        weights =  start_weights[:cpt] + transfer_weights[(cpt-1)] + end_weights[cpt:]
        #print "SIZES", sizes
        #print "AW",[ [ y.shape for y in x ]  for x in weights] 
        #print "WEI", len(weights) , [ len(x) for x in weights ]
        if len(ae_dropouts) == len(ae_hidden):
                do = ae_dropouts[:cpt] + transfer_do + ae_dropouts[cpt:]
        else : 
                do = [ 0 ] * (len(ae_hidden) + len(transfer_hidden))
        for w in weights[:-1]:
            #print "STEP", size
            layers.append(Dense(w[1].shape[0],activation=input_activation,init=init,weights=w)(layers[-1]))
            if do :
                d = do.pop(0)
                if d > 0 : 
                    layers.append(Dropout(d)(layers[-1]))
               
        layers.append(Dense(y_train.shape[1],activation=output_activation)(layers[-1]))
        models = [Model(input=layers[0] , output=x) for x in layers[1:]]
        models[-1].compile(optimizer=sgd,loss=loss)
        models[-1].fit(train,y_train,nb_epoch=epochs,batch_size=batch_size,callbacks=[EarlyStopping(monitor='val_loss', patience=patience, verbose=0)],validation_data=(dev,dev),verbose=verbose)
        predictions = [ [x.predict(y) for y in param_predict  ] for x in models ]
        pred_by_level.append(predictions)
  
    return pred_by_level

def train_mlp_proj(x_train,y_train,x_dev,y_dev,x_test,y_test,hidden_size,input_activation="relu",hidden_activation="relu",output_activation="softmax",loss="mse",init="glorot_uniform",dropouts=None,sgd=None,epochs=1200,batch_size=16,fit_verbose=1,patience=20,test_verbose=0):

    #model_tempfile=tempfile.mkstemp()
    tempfold = tempfile.mkdtemp()
    model_tempfile= tempfold+"/model.hdf"
    
    layers = [Input(shape=(x_train.shape[1],))]

    for h in hidden_size:
        print h
        if dropouts:   
            d = dropouts.pop(0)
            if d > 0 :
                ldo = Dropout(d)(layers[-1])
                print 'append'
                layers.append(Dense(h,init=init,activation=input_activation)(ldo))
        else :
            print " append" 
            layers.append(Dense(h,init=init,activation=input_activation)(layers[-1]))


    if dropouts:   
        d = dropouts.pop(0)
        if d > 0 :
            ldo =Dropout(d)(layers[-1])
            print "end"
            layers.append(Dense( y_train.shape[1],activation=output_activation,init=init)(ldo))
    else: 
        print "end"
        layers.append(Dense( y_train.shape[1],activation=output_activation,init=init)(layers[-1]))
    
    models = []
    for l in layers[1:] :
        models.append(Model(layers[0] , l))
    print "nb models : ", len(models), "h :",hidden_size , "layer", len(layers)
    if not sgd:
        sgd = SGD(lr=0.01, decay=0, momentum=0.9)

    models[-1].compile(loss=loss, optimizer=sgd,metrics=['accuracy'])
    callbacks = [ModelCheckpoint(model_tempfile, monitor='val_acc', verbose=test_verbose, save_best_only=True, save_weights_only=True, mode='auto'),
                 EarlyStopping(monitor='val_acc', patience=patience, verbose=test_verbose) ] # On pourrai essayer avec la loss aussi
    print models[-1].summary()
    hist=models[-1].fit(x_train, y_train, nb_epoch=epochs, batch_size=batch_size,verbose=fit_verbose,validation_data=(x_dev,y_dev),callbacks=callbacks)
    models[-1].load_weights(model_tempfile, by_name=False)
    proj = []
    for layer,model in enumerate(models):
        proj.append((model.predict(x_train),model.predict(x_dev),model.predict(x_test)))

    shutil.rmtree(tempfold)
    return models[-1].summary(),proj





def train_mlp_pred(x_train,y_train,x_dev,y_dev,x_test,y_test,hidden_size,input_activation="relu",hidden_activation="relu",output_activation="softmax",loss="mse",init="glorot_uniform",dropouts=None,sgd=None,epochs=1200,batch_size=16,fit_verbose=1,patience=20,test_verbose=0):

    #model_tempfile=tempfile.mkstemp()
    tempfold = tempfile.mkdtemp()
    model_tempfile= tempfold+"/model.hdf"
    
    layers = [Input(shape=(x_train.shape[1],))]

    for h in hidden_size:
        if dropouts:   
            d = dropouts.pop(0)
            if d > 0 :
                ldo = Dropout(d)(layers[-1])
                layers.append(Dense(h,init=init,activation=input_activation)(ldo))
        else :
            layers.append(Dense(h,init=init,activation=input_activation)(layers[-1]))


    if dropouts:   
        d = dropouts.pop(0)
        if d > 0 :
            ldo =Dropout(d)(layers[-1])
            layers.append(Dense( y_train.shape[1],activation=output_activation,init=init)(ldo))
    else: 
        layers.append(Dense( y_train.shape[1],activation=output_activation,init=init)(layers[-1]))
    
    model=Model(layers[0] , layers[-1])
    if not sgd:
        sgd = SGD(lr=0.01, decay=0, momentum=0.9)

    model.compile(loss=loss, optimizer=sgd,metrics=['accuracy'])
    callbacks = [ModelCheckpoint(model_tempfile, monitor='val_acc', verbose=test_verbose, save_best_only=True, save_weights_only=True, mode='auto'),
                 EarlyStopping(monitor='val_acc', patience=patience, verbose=test_verbose) ] # On pourrai essayer avec la loss aussi
    print model.summary()
    hist=model.fit(x_train, y_train, nb_epoch=epochs, batch_size=batch_size,verbose=fit_verbose,validation_data=(x_dev,y_dev),callbacks=callbacks)
    model.load_weights(model_tempfile, by_name=False)
    pred=(model.predict(x_train),model.predict(x_dev),model.predict(x_test))

    shutil.rmtree(tempfold)
    return pred,hist







def train_mlp(x_train,y_train,x_dev,y_dev,x_test,y_test,hidden_size,input_activation="relu",hidden_activation="relu",output_activation="softmax",loss="mse",init="glorot_uniform",dropouts=None,sgd=None,epochs=1200,batch_size=16,fit_verbose=1,test_verbose=0,save_pred=False,keep_histo=False):

    layers = [Input(shape=(x_train.shape[1],))]

    for h in hidden_size:
        if dropouts:   
            d = dropouts.pop(0)
            if d > 0 :
                layers.append(Dropout(d)(layers[-1]))

        layers.append(Dense(h,init=init,activation=input_activation)(layers[-1]))
            #if dropouts:
            #    drop_prob=dropouts.pop(0)
            #    if drop_prob > 0:
            #        model.add(Dropout(drop_prob))

        #if dropouts:
        #    drop_prob=dropouts.pop(0)
        #    if drop_prob > 0:
        #        model.add(Dropout(drop_prob))

        #if dropouts:
        #    model.add(Dropout(dropouts.pop(0)))
    if dropouts:   
        d = dropouts.pop(0)
        if d > 0 :
            layers.append(Dropout(d)(layers[-1]))
    print y_train[2:10]
    layers.append(Dense( y_train.shape[1],activation=output_activation,init=init)(layers[-1]))

    model =  Model(layers[0] , layers[-1])
    if not sgd:
        sgd = SGD(lr=0.01, decay=0, momentum=0.9)

    model.compile(loss=loss, optimizer=sgd,metrics=['accuracy'])

    scores_dev=[]
    scores_test=[]
    scores_train=[]
    save=None
    for i in range(epochs):
        hist=model.fit(x_train, y_train, nb_epoch=1, batch_size=batch_size,verbose=fit_verbose,validation_data=(x_dev,y_dev))
        pred_train=model.predict(x_train)
        pred_dev=model.predict(x_dev)
        pred_test=model.predict(x_test)

        scores_train.append(perf(np.argmax(y_train,axis=1),np.argmax(pred_train,axis=1)))
        scores_dev.append(perf(np.argmax(y_dev,axis=1),np.argmax(pred_dev,axis=1)))
        scores_test.append(perf(np.argmax(y_test,axis=1),np.argmax(pred_test,axis=1)))
        if fit_verbose :
            print "{} {} {} {}".format(i,scores_train[-1],scores_dev[-1],scores_test[-1])
        if save is None or (len(scores_dev)>2 and scores_dev[-1] > scores_dev[-2]):
            save=save_tuple(pred_train,pred_dev,pred_test)
    arg_dev = np.argmax(scores_dev)
    best_dev=scores_dev[arg_dev]
    best_test=scores_test[arg_dev]
    max_test=np.max(scores_test)
    if fit_verbose:
        print " res : {} {} {}".format(best_dev,best_test,max_test)

    res=[scores_train,scores_dev,scores_test]
    if save_pred:
        res.append(save)
    if keep_histo:
        res.append(hist)
    return res

def train_ae(train,dev,test,hidden_sizes,y_train=None,y_dev=None,y_test=None,dropouts=None,input_activation="tanh",output_activation="tanh",loss="mse",sgd=None,epochs=500,batch_size=8,test_verbose=0,verbose=1,patience=20,get_weights=False,set_weights=[],best_mod=False):
     
    input_vect = Input(shape=(train.shape[1],))

    previous = [input_vect]

    if dropouts is None:
        dropouts = [ 0 ] * (len(hidden_sizes) +1)
    if sgd is None : 
        sgd = SGD(lr=0.01, decay=0, momentum=0.9)
    did_do = False
    if dropouts :
        d = dropouts.pop(0)
        if d :
            previous.append(Dropout(d)(previous[-1]))
            did_do = True

    for h_layer,weight_layer in izip_longest(hidden_sizes,set_weights,fillvalue=None) :
        # ,weights=w
        if weight_layer :
            w = weight_layer[0] 
        else :
            w = None
        #print "ADD SIZE" , h_layer
        if did_do : 
            p = previous.pop()
            did_do = False
        else :
            p = previous[-1]
        previous.append(Dense(h_layer,activation=input_activation,weights=w)(previous[-1]))
        if dropouts:
            d = dropouts.pop(0)
            if d :
                previous.append(Dropout(d)(previous[-1]))
                did_do = True

    predict_y = True
    if y_train is None or  y_dev is None or y_test is None :
        y_train = train
        y_dev = dev
        y_test = test
        predict_y = False
    previous.append(Dense(y_train.shape[1],activation=output_activation)(previous[-1]))
    models = [Model(input=previous[0] , output=x) for x in previous[1:]]
    print "MLP", sgd, loss
    models[-1].compile(optimizer=sgd,loss=loss)
    cb = [EarlyStopping(monitor='val_loss', patience=patience, verbose=0)]
    if best_mod:
        tempfold = tempfile.mkdtemp()
        model_tempfile= tempfold+"/model.hdf"
        cb.append( ModelCheckpoint(model_tempfile, monitor='val_loss', verbose=test_verbose, save_best_only=True, save_weights_only=True, mode='auto') )

    models[-1].summary()
    models[-1].fit(train,y_train,nb_epoch=epochs,batch_size=batch_size,callbacks=cb,validation_data=(dev,dev),verbose=verbose)
    if best_mod:
        models[-1].load_weights(model_tempfile)
        shutil.rmtree(tempfold)
    param_predict = [ train, dev, test ]
    if predict_y :
        param_predict += [ y_train, y_dev ,y_test ]
    predictions = [ [x.predict(y) for y in param_predict  ] for x in models ]
    if get_weights : 
        weights = [ x.get_weights()  for x in models[-1].layers if x.get_weights() ]
        return ( predictions , weights )
    else :
        return predictions

def train_sae(train,dev,test,hidden_sizes,y_train=None,y_dev=None,y_test=None,dropouts=None,input_activation="tanh",output_activation="tanh",loss="mse",sgd=None,epochs=500,batch_size=8,verbose=1,patience=20):

    weights = []
    predictions = [[(train,dev,test),()]]
    ft_pred = []
    past_sizes = []


    for size in hidden_sizes :
        #print "DO size " , size , "FROM" , hidden_sizes
        res_pred, res_wght = train_ae(predictions[-1][-2][0], predictions[-1][-2][1],predictions[-1][-2][2],[size],
                                      dropouts=dropouts, input_activation=input_activation,
                                      output_activation=output_activation, loss=loss, sgd=sgd,
                                      epochs=epochs, batch_size=batch_size, verbose=verbose,
                                      patience=patience,get_weights=True)
        past_sizes.append(size)
        weights.append(res_wght)
        predictions.append(res_pred)
        #print "FINE TUNE "
        res_ftpred = train_ae(train,dev,test,past_sizes,y_train=y_train,y_dev=y_dev,y_test=y_test,
                              dropouts=dropouts,
                              input_activation=input_activation,
                              output_activation=output_activation,
                              loss=loss,sgd=sgd,epochs=epochs,
                              batch_size=batch_size,verbose=verbose,patience=patience,
                              set_weights=weights)
        ft_pred.append(res_ftpred)

    return ( predictions[1:] , ft_pred)