import os
from typing import Optional

import numpy as np
import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras import backend as K
from tensorflow.keras import layers
from tensorflow.keras.applications import EfficientNetB0, ResNet50
from tensorflow.keras.optimizers import Adam

def l1_distance(vects) -> float:
    """
    Finds the L1 distance between two vectors.

    Args:
        vects: List containing two tensors of same length.

    Returns:
        Element-wise L1 distance.
    """
    x, y = vects

    return K.abs(x - y)

def create_model(
    target_shape = (224, 224, 3),
    path = None,
) -> Model:
    """
    Creates the siamese model.

    Args:
        target_shape: image dimensions.
        path: path to best weights.

    Returns:
        Siamese model.
    """
    input_1 = layers.Input(shape=target_shape, name="inp_1")
    input_2 = layers.Input(shape=target_shape, name="inp_2")
    # input_1aug = img_augmentation(input_1)
    # input_2aug = img_augmentation(input_2)

    input = layers.Input(shape=target_shape, name="input")
    lambda_1 = layers.Lambda(
        lambda image: tf.keras.applications.resnet.preprocess_input(image),
        name="pre_process",
    )(input)
    base_cnn = ResNet50(
        weights="imagenet",
        input_tensor=lambda_1,
        input_shape=target_shape,
        include_top=False,
    )
    # CONV/FC -> BatchNorm -> ReLu(or other activation) -> Dropout -> CONV/FC ->
    pool = layers.MaxPooling2D(pool_size=(2, 2))(base_cnn.output)
    flatten = layers.Flatten(name="base_output_flatten")(pool)
    dense1 = layers.BatchNormalization(name="dense1_norm")(flatten)
    dense1 = layers.Dense(512, activation="relu", name="dense1")(dense1)
    dense1 = layers.Dropout(0.3, name="dense1_dropout")(dense1)
    dense2 = layers.BatchNormalization(name="dense2_norm")(dense1)
    dense2 = layers.Dense(256, activation="relu", name="dense2")(dense2)
    dense2 = layers.Dropout(0.2, name="dense2_dropout")(dense2)
    output = layers.Dense(256, name="dense_output")(dense2)

    embedding = Model(input, output, name="Embedding")

    trainable = False
    for layer in base_cnn.layers:
        if layer.name == "conv5_block1_out":
            trainable = True
        layer.trainable = trainable

    tower_1 = embedding(input_1)
    tower_2 = embedding(input_2)

    merge_layer = layers.Lambda(l1_distance, name="l1")([tower_1, tower_2])
    normal_layer = tf.keras.layers.BatchNormalization(name="l1_norm")(merge_layer)
    comparison_layer = layers.Dense(
        1,
        activation="sigmoid",
        name="final_layer",
    )(normal_layer)
    siamese = Model(inputs=[input_1, input_2], outputs=comparison_layer)

    if path is not None:
        siamese.load_weights(path)

    return siamese

early_stopping_callback = tf.keras.callbacks.EarlyStopping(
        monitor="loss", patience=5
    )
tensorboard_callback = tf.keras.callbacks.TensorBoard(
        log_dir="/logs", histogram_freq=1
    )
model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
        filepath="/logs/weights{epoch:04d}.tf", save_weights_only=True, save_freq=1
    )

train_generator = get_train_generator(
        split_path, batch_size=batch_size, input_size=target_shape
    )
steps_per_epoch = len(train_generator)
clr = get_cyclical_lr(2 * steps_per_epoch)
optimizer = Adam(clr)

siamese = create_model(target_shape)

siamese.compile(
        loss=loss(margin=margin),
        optimizer=optimizer,
        metrics=[metrics.accuracy, metrics.precision, metrics.recall, metrics.f1],
    )
siamese.summary()

siamese.fit(
        train_generator,
        validation_data=get_valid_generator(
            split_path, batch_size=batch_size, input_size=target_shape
        ),
        epochs=epochs,
        callbacks=[
            early_stopping_callback,
            tensorboard_callback,
            model_checkpoint_callback,
        ],
        verbose=1,
    )

siamese = create_model(path="/content/weights00000012.h5")

ValueError                                Traceback (most recent call last)
<ipython-input-9-e1ec7f0dd441> in <module>()
----> 1 siamese = create_model(path="/content/weights00000012.h5")

3 frames
<__array_function__ internals> in transpose(*args, **kwargs)

/usr/local/lib/python3.7/dist-packages/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
     56 
     57     try:
---> 58         return bound(*args, **kwds)
     59     except TypeError:
     60         # A TypeError occurs if the object does have such a method in its

ValueError: axes don't match array