#!/usr/bin/env python3

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# Model takes 8 sets of input data per inference.
FRAMES = 8

# Image dimensions
IMG_WIDTH = 32
IMG_HEIGHT = 32

# Size of the input 1-d dense layer (per frame).
DENSE_IN = 64

# Ranges for the two output one-hot tensors.
OUT_A_SIZE = 9
OUT_B_SIZE = 2


def get_img_frontend(name):
    input = keras.Input(shape=[IMG_WIDTH, IMG_HEIGHT, FRAMES], name=name)
    img = layers.Conv2D(
        64 * FRAMES,
        8,
        activation="relu",
        kernel_initializer="random_normal",
        bias_initializer="zeros",
    )(input)
    img = layers.MaxPooling2D(2)(img)
    img = layers.Conv2D(
        16 * FRAMES,
        4,
        activation="relu",
        kernel_initializer="random_normal",
        bias_initializer="zeros",
    )(img)
    img = layers.MaxPooling2D(2)(img)
    img = layers.Reshape((2048,), name=f"{name}_flat")(img)
    return input, img


def get_model():
    img_a_in, img_a = get_img_frontend("img_a_in")
    img_b_in, img_b = get_img_frontend("img_b_in")

    foo_in = keras.Input(shape=[DENSE_IN * FRAMES], name="foo_in")
    foo = layers.Dense(
        DENSE_IN * FRAMES,
        kernel_initializer="random_normal",
        bias_initializer="zeros",
    )(foo_in)

    # Merge images and 'foo' togther, then run through a few dense layers.
    all = layers.concatenate([img_a, img_b, foo], name="all")
    all = layers.Dense(
        1024, kernel_initializer="random_normal", bias_initializer="zeros"
    )(all)
    all = layers.Dense(
        512, kernel_initializer="random_normal", bias_initializer="zeros"
    )(all)
    all = layers.Dense(
        256, kernel_initializer="random_normal", bias_initializer="zeros"
    )(all)

    # Want two outputs.
    out_a = layers.Dense(
        OUT_A_SIZE, kernel_initializer="random_normal", bias_initializer="zeros"
    )(all)
    out_a = keras.backend.argmax(out_a, name="out_a_one_hot")

    out_b = layers.Dense(
        OUT_B_SIZE, kernel_initializer="random_normal", bias_initializer="zeros"
    )(all)
    out_b = keras.backend.argmax(out_b, name="out_b_one_hot")

    # Create the model.
    model = keras.Model(
        inputs=[img_a_in, img_b_in, foo_in], outputs=[out_a, out_b], name="baz"
    )
    model.compile(optimizer="adam", loss="mean_squared_error")
    return model


def call_model_with_random_inputs(model):
    img_a = np.random.randn(IMG_WIDTH, IMG_HEIGHT, FRAMES)
    img_b = np.random.randn(IMG_WIDTH, IMG_HEIGHT, FRAMES)
    foo = np.random.randn(DENSE_IN * FRAMES)

    inputs = {"img_a_in": img_a, "img_b_in": img_b, "foo_in": foo}

    r = model.call(inputs)
    print(r)


def main():
    print("tf.version: ", tf.__version__)
    model = get_model()
    keras.utils.plot_model(model, "summary.png", show_shapes=True)
    call_model_with_random_inputs(model)


if __name__ == "__main__":
    main()