# -*- coding: utf-8 -*-
"""
Created on Mon Nov 27 20:39:01 2017

@author: damodara
"""
from numpy.random import seed
seed(1)
from tensorflow import set_random_seed
set_random_seed(2)
import numpy as np
import matplotlib.pylab as plt
import dnn 
import os
import numpy as np
from sklearn.datasets import make_moons, make_blobs
from sklearn.decomposition import PCA

#from flipGradientTF import GradientReversal
#from revgrad import flip_gradient


source_traindata, source_trainlabel = make_blobs(1200, centers=[[0, 0], [0, 1]], cluster_std=0.2)
target_traindata, target_trainlabel = make_blobs(1200, centers=[[1, -1], [1, 0]], cluster_std=0.2)

plt.scatter(source_traindata[:,0], source_traindata[:,1], c=source_trainlabel, alpha=0.4)
plt.scatter(target_traindata[:,0], target_traindata[:,1], c=target_trainlabel, cmap='cool', alpha=0.4)


from keras.utils.np_utils import to_categorical
source_trainlabel_cat = to_categorical(source_trainlabel)
#source_testlabel_cat = to_categorical(source_testlabel)
#
target_trainlabel_cat = to_categorical(target_trainlabel)
#target_testlabel_cat = to_categorical(target_testlabel)
#target_label_cat =  to_categorical(target_label)
#%% MNIST
n_class = len(np.unique(source_trainlabel))
n_dim = np.shape(source_traindata)
optim = dnn.keras.optimizers.Adam(lr=0.001)
#%%
def feat_ext(main_input, l2_weight=0.0):
    '''
    main_input --> input shape tensor
    '''
    net = dnn.Dense(50, activation='relu')(main_input)  
    net = dnn.Dense(50, activation='relu')(net)            
#    self.domain_invariant_features = net
    return net

def classifier(model_input, nclass,l2_weight=0.0):
    net = dnn.Dense(nclass,activation='softmax', name='classifier_output')(model_input)
    return net
#%%
main_input = dnn.Input(shape=(n_dim[1],))
fe = feat_ext(main_input)
cl = classifier(fe, n_class)
source_model = dnn.Model(inputs=main_input, outputs=cl)
source_model.compile(loss='categorical_crossentropy', optimizer=optim, metrics=['accuracy'])
source_model.fit(source_traindata, source_trainlabel_cat, batch_size=32, epochs=10)
sonly_acc =source_model.evaluate(source_traindata, source_trainlabel_cat)
tacc = source_model.evaluate(target_traindata, target_trainlabel_cat)
print(sonly_acc, tacc)
feature_size = (fe.get_shape().as_list())[1]
#%%
main_input1=main_input = dnn.Input(shape=(n_dim[1],))
fe1 = feat_ext(main_input1)
cl1 = classifier(fe1,n_class)
con_cat = dnn.concatenate([cl1,fe1], axis=1)
coralmodel = dnn.Model(inputs=main_input1,outputs=con_cat)
coralmodel.compile(loss='categorical_crossentropy', optimizer=optim, metrics=['accuracy'])
#%%

#%%
#dnn.K.clear_session()
from SoA_methods import DeepCoral
optim = dnn.keras.optimizers.Adam(lr=0.001)
coral_model = DeepCoral(coralmodel, nclass=n_class,
                        batch_size=32,feature_size=feature_size, optimizer=optim, coral_param=0.001) 
#%%
coral_model.fit(source_traindata, source_trainlabel_cat, target_traindata,
                epochs=50)  
#%%
coral_acc = coral_model.evaluate(target_traindata, target_trainlabel_cat)
coral_sacc = coral_model.evaluate(source_traindata, source_trainlabel_cat)
print("coral model on source data",coral_sacc)
print("coral model on target data",coral_acc)

#%% MMD
from SoA_methods import DeepMMD
optim = dnn.keras.optimizers.Adam(lr=0.001)
mmd_model = DeepMMD(coralmodel, nclass=n_class,
                        batch_size=128,feature_size=feature_size, optimizer=optim, mmd_param=10) 
#%%
mmd_model.fit(source_traindata, source_trainlabel_cat, target_traindata,
                epochs=10)  
#%%
mmd_acc = mmd_model.evaluate(target_traindata, target_trainlabel_cat)
mmd_sacc = mmd_model.evaluate(source_traindata, source_trainlabel_cat)
print("MMD model on source data",mmd_sacc)
print("MMD model on target data",mmd_acc)


#%%DANN model
from SoA_methods import DANN
domain_lambda=2.0
loss_mode ='auto'
optim = dnn.keras.optimizers.SGD(lr=0.01, momentum=0.9)
dann_model = DANN(train_mode='dann', n_dim = n_dim, nclass=n_class,
                        batch_size=64,optimizer=optim, dann_param=domain_lambda, loss_mode=loss_mode) 

dann_model.fit(source_traindata, source_trainlabel_cat, target_traindata,
                epochs=20)  
apred = dann_model.predict(target_traindata)
dann_acc = dann_model.evaluate(target_traindata, target_trainlabel_cat)
print("The accuracy using auto mode\n",dann_acc) 

#%% 
from keras.models import load_model
import deep_JDOT as dtl
 

model = source_model
bcd_model = model
#model.set_weights(model1.get_weights())
lr=0.001
sgd= dnn.keras.optimizers.adam(lr=lr)#, momentum=0.9, decay=0.000, nesterov=True)
model.compile(optimizer=sgd,loss='categorical_crossentropy', metrics =['accuracy'])
bcd_model.compile(optimizer=sgd,loss='categorical_crossentropy', metrics =['accuracy'])


alpha=10.0
model_jdot = dtl.DeepJDOT(model, alpha, bcd_model)

batch_size=64
n_epoch=10

model_jdot.fit(source_traindata, source_trainlabel_cat, target_traindata, 
               target_trainlabel_cat, batch_size, n_epoch)

jdot_s_predict_label = model_jdot.predict(source_traindata)
jdot_s_acc = np.mean(np.argmax(jdot_s_predict_label,1) == np.argmax(source_trainlabel_cat,1))

jdot_predict_label = model_jdot.predict(target_traindata)
jdot_acc = np.mean(np.argmax(jdot_predict_label,1) == np.argmax(target_trainlabel_cat,1))

print('Learning with TLOT:')
print('TBR_app={:0.3f}'.format(jdot_s_acc))
print('TBR_test={:0.3f}'.format(jdot_acc)) 