import numpy as np
import tensorflow as tf
from tensorflow import keras, int32

from snake.constants.constants import GAME_SHAPE
from snake.snake_game import game_state

input_dims = (GAME_SHAPE[0], GAME_SHAPE[1], 3)

NUM_CLASSES = 3
DECODER_UNITS = 128


def residual_depthwise_block(x, num_repeats):
    residual = x
    for i in range(num_repeats):
        x = keras.layers.DepthwiseConv2D(3, padding='same', strides=1, activation='relu',
                                         kernel_regularizer=tf.keras.regularizers.l2(0.01))(x)

    return keras.layers.Add()([x, residual])


def residual_cov(x, num_channels, filt, strides, num_repeats):
    x = keras.layers.Conv2D(num_channels, filt, strides, padding='valid')(x)
    x = keras.layers.BatchNormalization(axis=-1, momentum=0.99, epsilon=0.001, center=True, scale=True)(x)
    x = keras.layers.ReLU()(x)

    ident = x
    for _ in range(num_repeats - 1):
        x = keras.layers.Conv2D(num_channels, filt, strides=1, padding='same')(x)
        x = keras.layers.BatchNormalization(axis=-1, momentum=0.99, epsilon=0.001, center=True, scale=True)(x)
        x = keras.layers.ReLU()(x)

    return keras.layers.Add()([x, ident])


def create_model(name='q_learning_model'):
    model_input = keras.Input(shape=input_dims, name='input_frames')
    selection_input = keras.Input(shape=tf.shape(0), dtype=int32, name="selector_input")

    x = model_input
    # x = x - 10
    # x = tf.keras.layers.ZeroPadding2D(2)(x)
    # x = x + 10
    # x = residual_cov(x, num_channels=16, filt=3, strides=1, num_repeats=3)
    # x = residual_cov(x, num_channels=32, filt=3, strides=1, num_repeats=2)
    #
    encoding = keras.layers.Flatten()(x)
    init = tf.keras.initializers.HeUniform()

    y = keras.layers.Dense(DECODER_UNITS, activation='relu', kernel_initializer=init)(encoding)
    y = keras.layers.Dense(DECODER_UNITS, activation='relu', kernel_initializer=init)(encoding)
    y = keras.layers.Dense(DECODER_UNITS//2, activation='relu', kernel_initializer=init)(y)
    y = keras.layers.Dense(NUM_CLASSES, activation='linear', kernel_initializer=init)(y)

    # y = encoding
    # y = keras.layers.Dense(DECODER_UNITS, activation='relu', use_bias=True,
    #                        kernel_regularizer=tf.keras.regularizers.l2(0.01))(y)
    # y = keras.layers.Dropout(0.2)(y)
    # y = keras.layers.Dense(DECODER_UNITS/2, activation='relu', use_bias=True,
    #                        kernel_regularizer=tf.keras.regularizers.l2(0.01))(y)
    # y = keras.layers.Dense(NUM_CLASSES, activation=None, name='output',
    #                        kernel_regularizer=tf.keras.regularizers.l2(0.01))(y)

    output = y
    selected_output = tf.gather(y, selection_input, axis=1, batch_dims=1, name='single_output')

    model = keras.Model(inputs=[model_input, selection_input], outputs=[y, selected_output], name=name)

    model.compile(loss=[None, tf.keras.losses.Huber()], optimizer=tf.keras.optimizers.Adam(lr=0.00001), metrics=['mae'])

    # print("Creating new model:")
    model.summary()
    return model


def compute_target(target_model):
    pass


def board_preprocess(board):
    board = np.copy(board)

    mask = board == game_state.FOOD_SYMBOL
    board[mask] = 1

    mask = board == game_state.BODY_SYMBOL
    board[mask] = -1

    mask = board == game_state.HEAD_SYMBOL
    board[mask] = -1

    return board  # , mask

import numpy as np
import tensorflow as tf
from tensorflow import keras, int32

from snake.constants.constants import GAME_SHAPE
from snake.snake_game import game_state

input_dims = (GAME_SHAPE[0], GAME_SHAPE[1], 3)

NUM_CLASSES = 3
DECODER_UNITS = 128


def residual_depthwise_block(x, num_repeats):
    residual = x
    for i in range(num_repeats):
        x = keras.layers.DepthwiseConv2D(3, padding='same', strides=1, activation='relu',
                                         kernel_regularizer=tf.keras.regularizers.l2(0.01))(x)

    return keras.layers.Add()([x, residual])


def residual_cov(x, num_channels, filt, strides, num_repeats):
    x = keras.layers.Conv2D(num_channels, filt, strides, padding='valid')(x)
    x = keras.layers.BatchNormalization(axis=-1, momentum=0.99, epsilon=0.001, center=True, scale=True)(x)
    x = keras.layers.ReLU()(x)

    ident = x
    for _ in range(num_repeats - 1):
        x = keras.layers.Conv2D(num_channels, filt, strides=1, padding='same')(x)
        x = keras.layers.BatchNormalization(axis=-1, momentum=0.99, epsilon=0.001, center=True, scale=True)(x)
        x = keras.layers.ReLU()(x)

    return keras.layers.Add()([x, ident])


def create_model(name='q_learning_model'):
    model_input = keras.Input(shape=input_dims, name='input_frames')
    selection_input = keras.Input(shape=tf.shape(0), dtype=int32, name="selector_input")

    x = model_input
    # x = x - 10
    # x = tf.keras.layers.ZeroPadding2D(2)(x)
    # x = x + 10
    # x = residual_cov(x, num_channels=16, filt=3, strides=1, num_repeats=3)
    # x = residual_cov(x, num_channels=32, filt=3, strides=1, num_repeats=2)
    #
    encoding = keras.layers.Flatten()(x)
    init = tf.keras.initializers.HeUniform()

    y = keras.layers.Dense(DECODER_UNITS, activation='relu', kernel_initializer=init)(encoding)
    y = keras.layers.Dense(DECODER_UNITS, activation='relu', kernel_initializer=init)(encoding)
    y = keras.layers.Dense(DECODER_UNITS//2, activation='relu', kernel_initializer=init)(y)
    y = keras.layers.Dense(NUM_CLASSES, activation='linear', kernel_initializer=init)(y)

    # y = encoding
    # y = keras.layers.Dense(DECODER_UNITS, activation='relu', use_bias=True,
    #                        kernel_regularizer=tf.keras.regularizers.l2(0.01))(y)
    # y = keras.layers.Dropout(0.2)(y)
    # y = keras.layers.Dense(DECODER_UNITS/2, activation='relu', use_bias=True,
    #                        kernel_regularizer=tf.keras.regularizers.l2(0.01))(y)
    # y = keras.layers.Dense(NUM_CLASSES, activation=None, name='output',
    #                        kernel_regularizer=tf.keras.regularizers.l2(0.01))(y)

    output = y
    selected_output = tf.gather(y, selection_input, axis=1, batch_dims=1, name='single_output')

    model = keras.Model(inputs=[model_input, selection_input], outputs=[y, selected_output], name=name)

    model.compile(loss=[None, tf.keras.losses.Huber()], optimizer=tf.keras.optimizers.Adam(lr=0.00001), metrics=['mae'])

    # print("Creating new model:")
    model.summary()
    return model


def compute_target(target_model):
    pass


def board_preprocess(board):
    board = np.copy(board)

    mask = board == game_state.FOOD_SYMBOL
    board[mask] = 1

    mask = board == game_state.BODY_SYMBOL
    board[mask] = -1

    mask = board == game_state.HEAD_SYMBOL
    board[mask] = -1

    return board  # , mask

# x = x - 10 # x = tf.keras.layers.ZeroPadding2D(2)(x) # x = x + 10 # x = residual_cov(x, num_channels=16, filt=3, strides=1, num_repeats=3) # x = residual_cov(x, num_channels=32, filt=3, strides=1, num_repeats=2) # # y = encoding # y = keras.layers.Dense(DECODER_UNITS, activation='relu', use_bias=True, #                        kernel_regularizer=tf.keras.regularizers.l2(0.01))(y) # y = keras.layers.Dropout(0.2)(y) # y = keras.layers.Dense(DECODER_UNITS/2, activation='relu', use_bias=True, #                        kernel_regularizer=tf.keras.regularizers.l2(0.01))(y) # y = keras.layers.Dense(NUM_CLASSES, activation=None, name='output', #                        kernel_regularizer=tf.keras.regularizers.l2(0.01))(y) # print("Creating new model:") # , mask