#!/usr/bin/env python
#model like in https://www.youtube.com/watch?v=tcpmucSLKo8

import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Flatten
from tensorflow.keras.layers import Conv2D
from tensorflow.keras import optimizers
from tensorflow.keras import backend as K
import sklearn
from utils import Sample

# Global variable
OUT_SHAPE = 1
INPUT_SHAPE = (Sample.IMG_H, Sample.IMG_W, Sample.IMG_D)


def customized_loss(y_true, y_pred, loss='euclidean'):
    # Simply a mean squared error that penalizes large joystick summed values
    if loss == 'L2':
        L2_norm_cost = 0.001
        val = K.mean(K.square((y_pred - y_true)), axis=-1) \
                    + K.sum(K.square(y_pred), axis=-1)/2 * L2_norm_cost
    # euclidean distance loss
    elif loss == 'euclidean':
        val = K.sqrt(K.sum(K.square(y_pred-y_true), axis=-1))
    return val


def create_model_1(keep_prob = 0.8):
    model = Sequential()

    # NVIDIA's model
    model.add(Conv2D(24, kernel_size=(5, 5), strides=(2, 2), activation='relu', input_shape= INPUT_SHAPE))
    model.add(Conv2D(36, kernel_size=(5, 5), strides=(2, 2), activation='relu'))
    model.add(Conv2D(48, kernel_size=(5, 5), strides=(2, 2), activation='relu'))
    model.add(Conv2D(64, kernel_size=(3, 3), activation='relu'))
    model.add(Conv2D(64, kernel_size=(3, 3), activation='relu'))
    model.add(Flatten())
    model.add(Dense(1164, activation='relu'))
    drop_out = 1 - keep_prob
    model.add(Dropout(drop_out))
    model.add(Dense(100, activation='relu'))
    model.add(Dropout(drop_out))
    model.add(Dense(50, activation='relu'))
    model.add(Dropout(drop_out))
    model.add(Dense(10, activation='relu'))
    model.add(Dropout(drop_out))
    
    model.add(Dense(OUT_SHAPE, activation='softsign'))
    model.summary()

    return model

def create_model_2(keep_prob = 1.0):
    model = Sequential()

    # NVIDIA's model
    model.add(Conv2D(24, kernel_size=(5, 5), strides=(2, 2), activation='elu', input_shape= INPUT_SHAPE))
    model.add(Conv2D(36, kernel_size=(5, 5), strides=(2, 2), activation='elu'))
    model.add(Conv2D(48, kernel_size=(5, 5), strides=(2, 2), activation='elu'))
    model.add(Conv2D(64, kernel_size=(3, 3), activation='elu'))
    model.add(Conv2D(64, kernel_size=(3, 3), activation='elu'))
    model.add(Flatten())
    drop_out = 1 - keep_prob
    model.add(Dropout(drop_out))
    model.add(Dense(100, activation='elu'))
    model.add(Dropout(drop_out))
    model.add(Dense(50, activation='elu'))
    model.add(Dropout(drop_out))
    model.add(Dense(10, activation='elu'))
    model.add(Dropout(drop_out))
    model.add(Dense(OUT_SHAPE))
    model.summary()

    return model

import sys
import os
if __name__ == '__main__':

    # Load Training Data
    x_train = np.load(os.path.join(sys.argv[2], "X.npy"))
    y_train = np.load(os.path.join(sys.argv[2], "y.npy"))

    x_train, y_train = sklearn.utils.shuffle(x_train, y_train)
    # print(y_train)

    print(x_train.shape[0], 'train samples')


    # Training loop variables
    epochs = 120
    batch_size = 70
    print("1")
    if(sys.argv[1] == "1"):
        model = create_model_1(0.8)
    elif(sys.argv[1] == "2"):
        model = create_model_2(0.8)

    print("2")
    #learning_rate=0.0001
    model.compile(loss='mse', optimizer=optimizers.Adam(learning_rate=0.0001))
    print("3")
    model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, shuffle=True, validation_split=0.2)
    print("4")
    model.save_weights('model_weights.h5')

tensorflow.keras.layers.Dense sklearn.utils.shuffle tensorflow.keras.layers.Conv2D tensorflow.keras.backend.square tensorflow.keras.optimizers.Adam tensorflow.keras.layers.Dropout tensorflow.keras.models.Sequential tensorflow.keras.layers.Flatten
[(31, 'tensorflow.keras.models.Sequential', 'Sequential', ([], {}), False, 'from tensorflow.keras.models import Sequential\n'), (56, 'tensorflow.keras.models.Sequential', 'Sequential', ([], {}), False, 'from tensorflow.keras.models import Sequential\n'), (86, 'sklearn.utils.shuffle', 'sklearn.utils.shuffle', (['x_train', 'y_train'], {}), False, 'import sklearn\n'), (34, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(24)'], {'kernel_size': '(5, 5)', 'strides': '(2, 2)', 'activation': '"""relu"""', 'input_shape': 'INPUT_SHAPE'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (35, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(36)'], {'kernel_size': '(5, 5)', 'strides': '(2, 2)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (36, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(48)'], {'kernel_size': '(5, 5)', 'strides': '(2, 2)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (37, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(64)'], {'kernel_size': '(3, 3)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (38, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(64)'], {'kernel_size': '(3, 3)', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (39, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (40, 'tensorflow.keras.layers.Dense', 'Dense', (['(1164)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (42, 'tensorflow.keras.layers.Dropout', 'Dropout', (['drop_out'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (43, 'tensorflow.keras.layers.Dense', 'Dense', (['(100)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (44, 'tensorflow.keras.layers.Dropout', 'Dropout', (['drop_out'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (45, 'tensorflow.keras.layers.Dense', 'Dense', (['(50)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (46, 'tensorflow.keras.layers.Dropout', 'Dropout', (['drop_out'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (47, 'tensorflow.keras.layers.Dense', 'Dense', (['(10)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (48, 'tensorflow.keras.layers.Dropout', 'Dropout', (['drop_out'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (50, 'tensorflow.keras.layers.Dense', 'Dense', (['OUT_SHAPE'], {'activation': '"""softsign"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (59, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(24)'], {'kernel_size': '(5, 5)', 'strides': '(2, 2)', 'activation': '"""elu"""', 'input_shape': 'INPUT_SHAPE'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (60, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(36)'], {'kernel_size': '(5, 5)', 'strides': '(2, 2)', 'activation': '"""elu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (61, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(48)'], {'kernel_size': '(5, 5)', 'strides': '(2, 2)', 'activation': '"""elu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (62, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(64)'], {'kernel_size': '(3, 3)', 'activation': '"""elu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (63, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(64)'], {'kernel_size': '(3, 3)', 'activation': '"""elu"""'}), False, 'from tensorflow.keras.layers import Conv2D\n'), (64, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (66, 'tensorflow.keras.layers.Dropout', 'Dropout', (['drop_out'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (67, 'tensorflow.keras.layers.Dense', 'Dense', (['(100)'], {'activation': '"""elu"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (68, 'tensorflow.keras.layers.Dropout', 'Dropout', (['drop_out'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (69, 'tensorflow.keras.layers.Dense', 'Dense', (['(50)'], {'activation': '"""elu"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (70, 'tensorflow.keras.layers.Dropout', 'Dropout', (['drop_out'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (71, 'tensorflow.keras.layers.Dense', 'Dense', (['(10)'], {'activation': '"""elu"""'}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (72, 'tensorflow.keras.layers.Dropout', 'Dropout', (['drop_out'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (73, 'tensorflow.keras.layers.Dense', 'Dense', (['OUT_SHAPE'], {}), False, 'from tensorflow.keras.layers import Dense, Dropout, Flatten\n'), (83, 'os.path.join', 'os.path.join', (['sys.argv[2]', '"""X.npy"""'], {}), False, 'import os\n'), (84, 'os.path.join', 'os.path.join', (['sys.argv[2]', '"""y.npy"""'], {}), False, 'import os\n'), (103, 'tensorflow.keras.optimizers.Adam', 'optimizers.Adam', ([], {'learning_rate': '(0.0001)'}), False, 'from tensorflow.keras import optimizers\n'), (22, 'tensorflow.keras.backend.square', 'K.square', (['(y_pred - y_true)'], {}), True, 'from tensorflow.keras import backend as K\n'), (26, 'tensorflow.keras.backend.square', 'K.square', (['(y_pred - y_true)'], {}), True, 'from tensorflow.keras import backend as K\n'), (23, 'tensorflow.keras.backend.square', 'K.square', (['y_pred'], {}), True, 'from tensorflow.keras import backend as K\n')]