import vdemsinversioncode as vdemi
import numpy as np
from keras.models import Sequential
from keras.layers import Dense
from keras.layers import Dropout
from keras.layers import BatchNormalization
from math import fabs


def prepareVdem(vdm, path):
    vdm.load_all_pkl(path)
    vdm.set_deltax(deltax=27, offsetx=0)
    vdm.truevdemraster(iw=[0, 1, 2, 3])
    vdm.set_deltax(deltax=27, offsetx=0)
    vdm.rastersyn(iw=[0, 1, 2, 3])

vdemr1=vdemi.vdemsinversioncode()
# vdemr2=vdemi.vdemsinversioncode()
vdemr3=vdemi.vdemsinversioncode()
vdemr4=vdemi.vdemsinversioncode()

prepareVdem(vdemr1, "data/vdem_cuda_revvel_raster_incl193_tg=4.7-7.5_0.2_vel=400_10.0_it=2_woNNature.opy")
# prepareVdem(vdemr2, "data/vdeminv_seppix79.6000_incl193_1he_400_15_temp=4.7-7.5-0.2_woNViggohcv2.opy")
prepareVdem(vdemr3, "data/vdeminv_seppix79.6000_incl193_1he_400_15_temp=4.7-7.5-0.2_woNViggoen024.opy")
prepareVdem(vdemr4, "data/vdem_cuda_revvel_denser_raster_incl193_tg=4.7-7.5_0.2_vel=400_20.0_it=2_woN.opy")

###-----------Edit this-------------###
###---------------------------------###
test_number = 6
###---------------------------------###
###---------------------------------###

#%% Format the data

print(np.max(vdemr1.rassyn_arr[1][0, :, :]))


#%% Format the data
from sklearn import preprocessing
from sklearn.decomposition import PCA
from sklearn.externals import joblib


def get_x_train(vdm):
    for i in range(0, vdm.rassyn_arr[0].shape[0]):
        if (i == 0):
            tmp_x_train = np.vstack((np.vstack((vdm.rassyn_arr[0][0, :, :], vdm.rassyn_arr[1][0, :, :])),
                                     vdm.rassyn_arr[3][0, :, :]))  # Not the index 2 for rassyn_arr
            print("tmp_x_train = {}".format(tmp_x_train.shape))
        else:
            to_be_stacked = np.vstack((np.vstack((vdm.rassyn_arr[0][i, :, :], vdm.rassyn_arr[1][i, :, :])),
                                       vdm.rassyn_arr[3][i, :, :]))  # Not the index 2 for rassyn_arr
            tmp_x_train = np.hstack((tmp_x_train, to_be_stacked))
    return np.swapaxes(tmp_x_train, 0, 1)

def get_y_train(vdm):
    factor = int(vdm.vdemraster_arr.shape[
                     2] / 35)  # This is the number of steps taken by the raster. For example the raster can take 7 steps, then every 7th value belong together.
                               # So we should look at every xth value where x is the factor.
    for i in range(0, factor):
        if (i == 0):
            tmp_y_train = reshapeYArray(vdm.vdemraster_arr[:, :, 0::factor, :]).reshape(-1, vdm.vdemraster_arr[:, :,0::factor, :].shape[-1])
        else:
            reshaped_y_train = reshapeYArray(vdm.vdemraster_arr[:, :, i::factor, :]).reshape(-1, vdm.vdemraster_arr[:, :, i::factor, :].shape[-1])
            tmp_y_train = np.hstack((tmp_y_train, reshaped_y_train))
    return np.swapaxes(tmp_y_train,0,1)

def reshapeYArray(yArr):
    shape = yArr.shape
    factor = int(shape[1]/20)  # The number of bins to combine into one. For example with shape[1] = 41, factor = 2. With shape[2] = 81, factor = 4
    reformatedArr = np.zeros((shape[0], int(shape[1]/factor), shape[2], shape[3]))
    for y in range(0, int(np.shape(yArr)[1] / factor)):
        combinedYarrValue = yArr[:, y * factor, :, :]
        for l in range (1, factor):
            combinedYarrValue += yArr[:, y * factor + l, :, :]
        reformatedArr[:, y, :, :] = combinedYarrValue
    return reformatedArr



x_unfitted1 = get_x_train(vdemr1)
x_unfitted3 = get_x_train(vdemr3)
x_unfitted4 = get_x_train(vdemr4)
x_unfitted = np.vstack((np.vstack((x_unfitted1, x_unfitted3)),x_unfitted4))

print("x_unfitted.shape = {} x_unfitted1 = {}, x_unfitted3 = {}, x_uniftted4 = {} ".format(x_unfitted.shape, x_unfitted1.shape, x_unfitted3.shape, x_unfitted4.shape))
inputScaler = preprocessing.StandardScaler()
x = inputScaler.fit(x_unfitted).transform(x_unfitted)

x_train = x[0:int(x.shape[0]*0.8),:]
x_test = x[int(x.shape[0]*0.8):-1,:]


y_unfitted1 = get_y_train(vdemr1)
y_unfitted3 = get_y_train(vdemr3)
y_unfitted4 = get_y_train(vdemr4)
y_unfitted = np.vstack((np.vstack((y_unfitted1,y_unfitted3)),y_unfitted4))
print("y_unfitted.shape = {} y_unfitted1 = {}, y_uniftted3 = {}, y_unfitted4 = {}".format(y_unfitted.shape, y_unfitted1.shape, y_unfitted3.shape, y_unfitted4.shape))

#Reduce the number of dimensions with PCA:

# pca = joblib.load("/Users/lars/Python/muse/tests/test6/pca.save")
# pca = PCA(0.99)
# pca.fit(y_unfitted)
# y_dim_reduced = pca.transform(y_unfitted)

y_dim_reduced = y_unfitted

outputScaler = preprocessing.StandardScaler()
y = outputScaler.fit(y_dim_reduced).transform(y_dim_reduced)

joblib.dump(outputScaler, "tests/test{}/outputScaler.save".format(test_number))
# joblib.dump(pca, "tests/test{}/pca.save".format(test_number))


y_train = y[0:int(x.shape[0]*0.8),:]
y_test = y[int(y.shape[0]*0.8):-1,:]

#%% Get correct dimentions for array


print("The shape is: {}".format(get_y_train(vdemr1).shape))

#%% Train the network

clf = Sequential()
x_train_dim = x_train.shape[1]
y_train_dim = y_train.shape[1]

print("xtraindim = {}, ytraindim = {}".format(x_train_dim, y_train_dim))

clf.add(Dense(x_train_dim, input_dim=x_train_dim, activation='relu'))
clf.add(Dense(int((x_train_dim+y_train_dim) / 1.3 ), activation='relu'))
clf.add(Dropout(0.4))
clf.add(Dense(int((x_train_dim+y_train_dim) / 1.5), activation='relu'))
clf.add(BatchNormalization())
# clf.add(Dropout(0.2))
clf.add(Dense(int((x_train_dim+y_train_dim) / 1.8), activation='tanh'))
clf.add(BatchNormalization())
# clf.add(Dropout(0.4))
clf.add(Dense(int((x_train_dim+y_train_dim) / 2), activation='tanh'))
# clf.add(BatchNormalization())
# clf.add(Dropout(0.4))
clf.add(Dense(y_train_dim))

# Compile model
clf.compile(loss='mean_squared_error', optimizer='adam', metrics=['accuracy'])

clf.fit(x_train, y_train, epochs=300, batch_size=100, shuffle=True, validation_data=(x_test,y_test))
clf.save('tests/test{}/graph.h5'.format(test_number))

#%% Predict using the model
import matplotlib.pyplot as plt
from matplotlib import colors



def mean_absolute_percentage_error(actual, predict):
    tmp, n = 0.0, 0
    for i in range(0, len(actual)):
        for l in range(len(actual[i])):
            if actual[i][l] != 0:
                diff = fabs(fabs(actual[i][l]-predict[i][l])/actual[i][l])
                tmp += diff
                n += 1

    return (tmp/n)


from keras.models import load_model
savedModel = load_model('/Users/lars/Python/muse/tests/test5/graphWith20veltest5')
predictedYs = savedModel.predict(x_test)
print("PredictedYs = {}".format(predictedYs.shape))
# y_test_scaled = outputScaler.inverse_transform(y_test)
# scaledpredictedYs = outputScaler.inverse_transform(predictedYs)
# y_test_scaled = outputScaler.inverse_transform(y_test)
# y_predict_scaled = outputScaler.inverse_transform(predictedYs)

# y_test_original = pca.inverse_transform(y_test_scaled)
# y_predict_original = pca.inverse_transform(y_predict_scaled)

y_test_original = y_test
y_predict_original = predictedYs

y_test_scaled_real_dim = np.swapaxes(y_test_original,0,1)
y_test_scaled_real_dim = y_test_scaled_real_dim.reshape((15,20,35,1705))


y_predict_scaled_real_dim = np.swapaxes(y_predict_original,0,1)
y_predict_scaled_real_dim = y_predict_scaled_real_dim.reshape((15,20,35,1705))

print("The dim = {}".format(y_predict_scaled_real_dim.shape))


plt.title("Test {}".format(5))
plt.xlabel("Real values")
plt.ylabel("Predicted values")



plt.plot(y_test_original, y_predict_original, 'ro', markersize=1)
plt.axis([np.min(y_test_original), np.max(y_test_original), np.min(y_predict_original), np.max(y_predict_original)])
# print("Keras score = {}".format(savedModel.evaluate(x_test, y_test)))
plt.show()
 #%% show 2D plot

print(y_test_scaled_real_dim[0][15][15][1700])


#%% show 2D plot
plt.axis([np.min(y_test_scaled), 1e31, np.min([y_predict_scaled]), 1e31])
plt.show()

#%% show 2D plot
# print(predictedYs[0].shape)
plt.hist2d(y_test.ravel(),predictedYs.ravel(), bins=60,cmap='gnuplot2',
           range=[[y_test.min(),y_test.max()],[predictedYs.min(),predictedYs.max()]],normed=True,norm = colors.LogNorm())
plt.title("Test 6")
plt.xlabel('Real y')
plt.ylabel('predicted y')
plt.colorbar(label='Probability Density')
plt.show()
# babs_weighted_error = np.average(np.abs(Y_test-Y_testread_synlines_pred)*babs_test.reshape(shape_test[0]*shape_test[1]))/np.average(babs_test.reshape(shape_test[0]*shape_test[1]))
# error = np.average(np.abs(Y_test-Y_test_pred))
# print("Mean error (L1):", error, ". |B|-weighted mean error (L1): ",babs_weighted_error)

#%% Print the data


print(y_predict_scaled_real_dim[:, :, 15, 15])



#%% Plot observations as images to compare
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from matplotlib.mlab import bivariate_normal
import matplotlib.animation as animation
from matplotlib.pyplot import savefig


# for sample in range (0, 10):

sample = 15

# for sample in range (0, y_unfitted.shape[1]):
print("Going to {}".format(y_unfitted.shape[1]))
print(y_unfitted.shape)
print(y_test_scaled_real_dim.shape)
y_test_scaled_for_plot = np.swapaxes(y_unfitted[int(y_unfitted.shape[0]*0.8):-1,:],0,1).reshape((15,20,35,1705))

y_test_to_plot = y_test_scaled_for_plot[:, :, 15, sample]
y_predict_to_plot = y_predict_scaled_real_dim[:, :, 15, sample]


if (np.min(y_test_to_plot) < np.min(y_predict_to_plot)):
    plt_min = np.min(y_test_to_plot)
else:
    plt_min = np.min(y_predict_to_plot)

if(np.max(y_test_to_plot) > np.max(y_predict_to_plot)):
    plt_max = np.max(y_test_to_plot)
else:
    plt_max = np.max(y_predict_to_plot)


fig, ax = plt.subplots(2, 1)
pcm = ax[0].pcolormesh(y_test_to_plot,
                       norm=colors.SymLogNorm(linthresh=0.03, linscale=0.03,
                                              vmin=plt_min, vmax=plt_max), cmap='RdBu_r')
pcm2 = ax[1].pcolormesh(y_predict_to_plot,
                        norm=colors.SymLogNorm(linthresh=0.03, linscale=0.03,
                                               vmin=plt_min, vmax=plt_max), cmap='RdBu_r')

fig.colorbar(pcm, ax=ax[0], extend='both')
fig.colorbar(pcm2, ax=ax[1], extend='both')

# savefig("tests/test{}/moviePngs/pn{}.png".format(test_number, sample))

# cax = plt.axes([0.95, 0.1, 0.075, 0.8])
# plt.colorbar(cax=cax)
plt.show()

# savedModel.evaluate(X_test, Y_predict)
#%% Look at correlation with intensity
from keras.models import load_model
from matplotlib import colors
import matplotlib.pyplot as plt

def doPlot(y_test, y_predict, title, xMax = None, yMax = None):
    plt.plot(y_test, y_predict, 'ro', markersize=1)
    if (xMax is None): xMax = np.max(y_test)
    if (yMax is None): yMax = np.max(y_predict)
    plt.axis([np.min(y_test), 1e33, np.min(y_predict), 1e33])
    plt.xlabel("Real intensity")
    plt.ylabel("Predicted intensity")
    plt.title(title)
    plt.show()
    plt.savefig('figures/'+title+".png")

def do2dPlot(y_test, y_predict, title, xMax = None, yMax = None):

    plt.hist2d(y_test.ravel(), y_predict.ravel(), bins=60, cmap='gnuplot2',
               range=[[y_test.min(), 1e33], [y_predict.min(), 1e33]], normed=True,
               norm=colors.LogNorm())
    plt.title(title)
    plt.xlabel("Real intensity")
    plt.ylabel("Predicted intensity")
    plt.colorbar(label='Probability Density')
    plt.show()
    plt.savefig('figures/' + title + ".png")

print("xtest = {}, ytest = {}".format(x_test.shape, y_test.shape))
vdemr1.read_synlines()

savedModel = load_model('/Users/lars/Python/muse/graphs/graphWith20veltest6')
# x_test = np.swapaxes(vdemr1.rassyn_arr[0][0], 0,1)
y_predict = outputScaler.inverse_transform(savedModel.predict(x_test))
y_predict_scaled = np.swapaxes(y_predict,0,1)
y_predict_scaled = y_predict.reshape((15,20,35,1705))

print("Evaluation = {}".format(savedModel.evaluate(x_test, y_test)))

print("y_test shape = {}".format(y_test.shape))
print("y_predict shape = {}".format(y_predict.shape))

print("predict max = {}, test max = {}".format(np.max(y_predict), np.max(y_test)))
scaled_y_test = np.swapaxes(outputScaler.inverse_transform(y_test),0,1).reshape(15,20,35,1705)
vdemr1.calc_moments_vdemsol_arr(scaled_y_test)
y_intens_test1 = vdemr1.synprofsol[0]['intens']
y_intens_test2 = vdemr1.synprofsol[1]['intens']
y_intens_test3 = vdemr1.synprofsol[2]['intens']

# scaled_y_predict = outputScaler.inverse_transform(y_predict_scaled)
vdemr1.calc_moments_vdemsol_arr(y_predict_scaled)
y_intens_predict1 = vdemr1.synprofsol[0]['intens']
y_intens_predict2 = vdemr1.synprofsol[1]['intens']
y_intens_predict3 = vdemr1.synprofsol[2]['intens']

print("y_intens_test1 {}".format(y_intens_test1.shape))

doPlot(y_intens_test1, y_intens_predict1, "Predicted vs real intensity (fe_19_108.355), test5")
doPlot(y_intens_test2, y_intens_predict2, "Predicted vs real intensity (fe_15_284.163), test5")
doPlot(y_intens_test3, y_intens_predict3, "Predicted vs real intensity (fe_9_171.073), test5")

do2dPlot(y_intens_test1, y_intens_predict1, "Predicted vs real intensity (fe_19_108.355), test5")
do2dPlot(y_intens_test2, y_intens_predict2, "Predicted vs real intensity (fe_15_284.163), test5")
do2dPlot(y_intens_test3, y_intens_predict3, "Predicted vs real intensity (fe_9_171.073), test5")

# synteticVdm = vdemr.vdemsol_arr[str([0, 1, 2, 3])]
#
# vdemr.calc_moments_vdemsol_arr(synteticVdm)
# print(vdemr.synprofsol[0]['intens'].shape)