import tensorflow as tf
from agent_code.big_bertha_v1.parameters import (ACTIVATION_FUNCTION,
                                                 DENSE_LAYER_DIMS,
                                                 LEARNING_RATE, LOSS_FUNCTION,
                                                 NUMBER_OF_ACTIONS,
                                                 STATE_SHAPE)
from tensorflow.keras.layers import Dense
from tensorflow.keras.optimizers import Adam
from tensorflow.python.keras import Input, Model
from tensorflow.python.keras.layers import Add, Flatten, Lambda, Subtract
from tensorflow.python.ops.init_ops_v2 import variance_scaling_initializer


def compile_dueling_deep_q_network():
    he = variance_scaling_initializer()

    model_input = Input(shape=(STATE_SHAPE,))
    x = Dense(DENSE_LAYER_DIMS, input_shape=(STATE_SHAPE,), kernel_initializer=he,
              activation=ACTIVATION_FUNCTION)(model_input)
    # x = Dense(DENSE_LAYER_DIMS, kernel_initializer=he, activation=ACTIVATION_FUNCTION)(x)

    # val_stream, adv_stream = Lambda(lambda w: tf.split(w, 2, 3))(x)
    val_stream, adv_stream = Lambda(lambda w: tf.split(w, 2, 1))(x)

    val_stream = Flatten()(val_stream)
    val = Dense(1, kernel_initializer=he)(val_stream)

    adv_stream = Flatten()(adv_stream)
    adv = Dense(NUMBER_OF_ACTIONS, kernel_initializer=he)(adv_stream)

    # Combine streams into Q-Values
    reduce_mean = Lambda(lambda w: tf.reduce_mean(w, axis=1, keepdims=True))  # custom layer for reduce mean

    q_vals = Add()([val, Subtract()([adv, reduce_mean(adv)])])

    # Build model
    model = Model(model_input, q_vals)
    model.compile(Adam(LEARNING_RATE), loss=LOSS_FUNCTION)

    return model
