from keras.initializers import truncated_normal,Constant

from keras.layers import Convolution2D, Input
from keras.layers.core import Dense, Flatten, Lambda, Activation, Dropout
from keras.models import Model, Sequential
from keras.optimizers import SGD,Adam
from keras.regularizers import l2
from keras.losses import MSE
import tensorflow as tf
import keras.backend as K

# from keras import backend as K
# K.set_image_dim_ordering('tf')

def atan_layer(x):
    print(x, tf.multiply(tf.atan(x), 2))
    return tf.multiply(tf.atan(x), 2)

def atan_layer_shape(input_shape):
    return input_shape

def atan(x):
    return tf.atan(x)

def customLoss(yTrue, yPred):
    return tf.reduce_mean(tf.square(tf.subtract(yTrue, yPred)))

def NVIDA():
    REG = 0.001
    from keras.initializers import truncated_normal, Constant

    inputs = Input(shape=(66, 200, 3))
    conv_1 = Convolution2D(24, (5, 5), activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1),   name='conv_1',padding='valid', strides=(2, 2))(inputs)
    conv_2 = Convolution2D(36, (5, 5), activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1),   name='conv_2',padding='valid', strides=(2, 2))(conv_1)
    conv_3 = Convolution2D(48, (5, 5), activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1),   name='conv_3',padding='valid', strides=(2, 2))(conv_2)
    conv_4 = Convolution2D(64, (3, 3), activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1),   name='conv_4',padding='valid', strides=(1, 1))(conv_3)
    conv_5 = Convolution2D(64, (3, 3), activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1),   name='conv_5',padding='valid', strides=(1, 1))(conv_4)

    flat = Flatten()(conv_5)

    dense_1 = Dense(1164)(flat)
    dense_1 = Dropout(.8)(dense_1)
    dense_2 = Dense(100, activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1)  )(dense_1)
    dense_2 = Dropout(.8)(dense_2)
    dense_3 = Dense(50, activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1)  )(dense_2)
    dense_3 = Dropout(.8)(dense_3)
    dense_4 = Dense(10, activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1)  )(dense_3)
    dense_4 = Dropout(.8)(dense_4)

    final = Dense(1)(dense_4)
    final = Lambda(lambda x: tf.multiply(tf.atan(x), 2))(final)

    model = Model(input=inputs, outputs=final)
    model.compile(optimizer=Adam(lr=1e-4), loss= customLoss)
    model.summary()
    return model

def NVIDA1():
    REG = 0.001
    from keras.initializers import truncated_normal, Constant
    model= Sequential()
    model.add(Input(shape=(66, 200, 3)))
    model.add(Convolution2D(24, (5, 5), activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1),   name='conv_1',padding='valid', strides=(2, 2)))
    model.add(Convolution2D(36, (5, 5), activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1),   name='conv_2',padding='valid', strides=(2, 2)))
    model.add(Convolution2D(48, (5, 5), activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1),   name='conv_3',padding='valid', strides=(2, 2)))
    model.add(Convolution2D(64, (3, 3), activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1),   name='conv_4',padding='valid', strides=(1, 1)))
    model.add(Convolution2D(64, (3, 3), activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1),   name='conv_5',padding='valid', strides=(1, 1)))

    model.add(Flatten())

    model.add(Dense(1164))
    model.add(Dropout(.8))
    model.add(Dense(100, activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1)  ))
    model.add(Dropout(.8))
    model.add(Dense(50, activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1)  ))
    model.add(Dropout(.8))
    model.add(Dense(10, activation='relu',bias_initializer=Constant(0.1), kernel_initializer=truncated_normal(0.1)  ))
    model.add(Dropout(.8))
    model.add(Dense(1))
    model.add( Lambda(lambda x: tf.multiply(tf.atan(x), 2)))

    model.compile(optimizer=Adam(lr=1e-4), loss= customLoss)
    model.summary()
    return model