'''
    Author: Mahnoor Anjum
    Description:
        Autocolorization
        Model:
            neighboring pixels
            L + HAARS ----> A, B 
        Data preprocessed by:
            https://github.com/Abdullah230
'''



import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf
print(tf.__version__)

data = pd.read_csv('data/M15/m.csv')
cols = list(data.columns)
cols.remove('a')
cols.remove('b')
X_train = data.loc[:, cols]
y_train = data.loc[:, ['a', 'b']]

data_test = pd.read_csv('data/M15/test.csv')
X_test = data_test.loc[:, cols]
y_test = data_test.loc[:, ['a', 'b']]

print(X_train.shape)
print(y_train.shape)
print(X_test.shape)
print(y_test.shape)

miniL = X_train.min()
maxiL = X_train.max()
miniAB = y_train.min()
maxiAB = y_train.max()

from sklearn.preprocessing import StandardScaler
obj = StandardScaler() 
X_train = obj.fit_transform(X_train) 
X_test = obj.transform(X_test) 

objy = StandardScaler() 
y_train = objy.fit_transform(y_train) 
y_test = objy.transform(y_test) 
Y = y_train.shape[1]

N, D = X_train.shape


import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Dropout
from tensorflow.keras.models import Model
i_layer = Input(shape = (D,))
# h_layer = Dropout(0.4)(h_layer)
h_layer = Dense(16, activation='relu')(i_layer)
h_layer = Dropout(0.6)(h_layer)
h_layer = Dense(8, activation='relu')(h_layer)
#h_layer = Dropout(0.6)(h_layer)
# h_layer = Dense(256, activation='relu')(h_layer)
o_layer = Dense(Y)(h_layer)

model = Model(i_layer, o_layer)

model.summary()

optimizer = tf.keras.optimizers.RMSprop(0.001)

model.compile(loss='mse',
              optimizer='adam',
              metrics=['mae', 'mse'])

from tensorflow.keras.callbacks import EarlyStopping
callback = EarlyStopping(patience=3)

#report = model.fit(X_train, y_train,  epochs = 10)
report = model.fit(X_train, y_train, validation_data=(X_test, y_test), \
                   epochs = 100, batch_size = 64)


plt.plot(report.history['loss'], label="loss")
plt.plot(report.history['val_loss'], label="validation_loss")
plt.legend()

model.save('models/m15_landmarks') 


# print("Train eval: ", model.evaluate(X_train, y_train))
# print("Test eval: ", model.evaluate(X_test, y_test))

y_pred = model.predict(X_test)
y_pred = objy.inverse_transform(y_pred)
y_test = objy.inverse_transform(y_test)
X_test = obj.inverse_transform(X_test)
print(y_test.shape)
print(y_pred.shape)


shape = (64,64,1)
imageL = X_test[:,264].reshape(shape)
imagea = y_pred[:,0].reshape(shape)
imageb = y_pred[:,1].reshape(shape)

image = np.concatenate((imageL, imagea, imageb), axis=2)

import cv2
imageT = cv2.cvtColor(image.astype('float32'), cv2.COLOR_Lab2RGB)
cv2.imshow('colored',imageT)
cv2.waitKey(0)
cv2.destroyAllWindows()


imageL = X_test[:,264].reshape(shape)
imagea = y_test[:,0].reshape(shape)
imageb = y_test[:,1].reshape(shape)

image = np.concatenate((imageL, imagea, imageb), axis=2)
imageT = cv2.cvtColor(image.astype('float32'), cv2.COLOR_Lab2RGB)
cv2.imshow('original',imageT)
cv2.waitKey(0)
cv2.destroyAllWindows()


matplotlib.pyplot.legend pandas.read_csv tensorflow.keras.models.Model tensorflow.keras.layers.Dense tensorflow.keras.optimizers.RMSprop matplotlib.pyplot.plot numpy.concatenate tensorflow.keras.layers.Dropout sklearn.preprocessing.StandardScaler tensorflow.keras.callbacks.EarlyStopping tensorflow.keras.layers.Input
[(20, 'pandas.read_csv', 'pd.read_csv', (['"""data/M15/m.csv"""'], {}), True, 'import pandas as pd\n'), (27, 'pandas.read_csv', 'pd.read_csv', (['"""data/M15/test.csv"""'], {}), True, 'import pandas as pd\n'), (42, 'sklearn.preprocessing.StandardScaler', 'StandardScaler', ([], {}), False, 'from sklearn.preprocessing import StandardScaler\n'), (46, 'sklearn.preprocessing.StandardScaler', 'StandardScaler', ([], {}), False, 'from sklearn.preprocessing import StandardScaler\n'), (57, 'tensorflow.keras.layers.Input', 'Input', ([], {'shape': '(D,)'}), False, 'from tensorflow.keras.layers import Input, Dense, Dropout\n'), (66, 'tensorflow.keras.models.Model', 'Model', (['i_layer', 'o_layer'], {}), False, 'from tensorflow.keras.models import Model\n'), (70, 'tensorflow.keras.optimizers.RMSprop', 'tf.keras.optimizers.RMSprop', (['(0.001)'], {}), True, 'import tensorflow as tf\n'), (77, 'tensorflow.keras.callbacks.EarlyStopping', 'EarlyStopping', ([], {'patience': '(3)'}), False, 'from tensorflow.keras.callbacks import EarlyStopping\n'), (84, 'matplotlib.pyplot.plot', 'plt.plot', (["report.history['loss']"], {'label': '"""loss"""'}), True, 'import matplotlib.pyplot as plt\n'), (85, 'matplotlib.pyplot.plot', 'plt.plot', (["report.history['val_loss']"], {'label': '"""validation_loss"""'}), True, 'import matplotlib.pyplot as plt\n'), (86, 'matplotlib.pyplot.legend', 'plt.legend', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (107, 'numpy.concatenate', 'np.concatenate', (['(imageL, imagea, imageb)'], {'axis': '(2)'}), True, 'import numpy as np\n'), (111, 'cv2.imshow', 'cv2.imshow', (['"""colored"""', 'imageT'], {}), False, 'import cv2\n'), (112, 'cv2.waitKey', 'cv2.waitKey', (['(0)'], {}), False, 'import cv2\n'), (113, 'cv2.destroyAllWindows', 'cv2.destroyAllWindows', ([], {}), False, 'import cv2\n'), (120, 'numpy.concatenate', 'np.concatenate', (['(imageL, imagea, imageb)'], {'axis': '(2)'}), True, 'import numpy as np\n'), (122, 'cv2.imshow', 'cv2.imshow', (['"""original"""', 'imageT'], {}), False, 'import cv2\n'), (123, 'cv2.waitKey', 'cv2.waitKey', (['(0)'], {}), False, 'import cv2\n'), (124, 'cv2.destroyAllWindows', 'cv2.destroyAllWindows', ([], {}), False, 'import cv2\n'), (59, 'tensorflow.keras.layers.Dense', 'Dense', (['(16)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Input, Dense, Dropout\n'), (60, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.6)'], {}), False, 'from tensorflow.keras.layers import Input, Dense, Dropout\n'), (61, 'tensorflow.keras.layers.Dense', 'Dense', (['(8)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Input, Dense, Dropout\n'), (64, 'tensorflow.keras.layers.Dense', 'Dense', (['Y'], {}), False, 'from tensorflow.keras.layers import Input, Dense, Dropout\n')]