import tensorflow as tf

from tensorflow.keras import backend

from tensorflow.keras.layers import Layer

from . import build


@tf.function
def compute_u(h):

    h_as2d = tf.reshape(h, (-1, backend.shape(h)[-1]))

    u = tf.matmul(tf.transpose(h_as2d), h_as2d, name='u')

    return u


@tf.function
def partition_loss(h):

    u = compute_u(h)

    diag_eta = tf.linalg.diag_part(u) / backend.sum(u, axis=-1)
    diag_eta = backend.clip(diag_eta, min_value=backend.epsilon(), max_value=1.0)

    loss = backend.mean(-backend.log(diag_eta))

    return loss


@tf.function
def squared_l2_loss(y_true, y_pred):

    loss = 0.5 * (y_pred - y_true) ** 2

    return backend.sum(loss)


@tf.function
def binary_loss(y_true, y_pred):

    loss = backend.binary_crossentropy(y_true, y_pred)

    return backend.mean(loss)


@tf.function
def zero_loss(*_):

    return 0.0


def acrbm(inputs, units, visible: bool, h_name='h', v_name='v', *args, **kwargs):

    outputs = build.ACRBM(units, visible, *args, **kwargs)(inputs)

    h = Layer(name=h_name)(outputs[0])

    if visible:

        v = Layer(name=v_name)(outputs[1])

        return h, v

    else:

        return h,


def save_model(model, path='model.h5', **kwargs):

    model.save(path, **kwargs)
