import tensorflow as tf
import tensorflow.keras.backend as K


def dense_depth_loss_function(y, y_pred):
    """
    Implementation of the loss from the dense depth paper https://arxiv.org/pdf/1812.11941.pdf
    """
    # Point-wise L1 loss
    l_depth = K.mean(K.abs(y_pred - y), axis=-1)

    # L1 loss over image gradients
    dy, dx = tf.image.image_gradients(y)
    dy_pred, dx_pred = tf.image.image_gradients(y_pred)
    l_grad = K.mean(K.abs(dy_pred - dy) + K.abs(dx_pred - dx), axis=-1)

    #  Structural Similarity (SSIM)
    l_ssim = (1 - tf.image.ssim(y, y_pred, 500)) / 2

    return 0.1 * K.mean(l_depth) + l_grad + l_ssim