import tensorflow as tf

from colors import vector_to_index


@tf.function
def aesthetic_loss(y_pics):

    flipx = tf.reverse(y_pics, axis=(-2,))
    flipy = tf.reverse(y_pics, axis=(-3,))
    center = tf.reverse(y_pics, axis=(-2, -3,))
    trans = tf.transpose(y_pics, perm=(0, 2, 1, 3))
    anti_trans = tf.reverse(tf.transpose(tf.reverse(
        y_pics, axis=(-2,)), perm=(0, 2, 1, 3)), axis=(-2,))

    flipx_score = compute_score(y_pics, flipx)
    flipy_score = compute_score(y_pics, flipy)
    trans_score = compute_score(y_pics, trans)
    anti_trans_score = compute_score(y_pics, anti_trans)
    center_score = compute_score(y_pics, center)

    # Extract patches for neighborhood check
    patches = tf.image.extract_patches(
        y_pics, (1, 3, 3, 1), (1, 1, 1, 1), (1, 1, 1, 1), padding="SAME")
    # Dot product with central pixel
    # neigh = tf.reduce_sum(
    #     tf.maximum(0.0, -tf.einsum("...k,...k->...", tf.repeat(y_pics, 9, -1), patches)))
    neigh = tf.reduce_sum(
        tf.abs(tf.einsum("...k,...k->...", tf.repeat(y_pics, 9, -1), patches)))

    divers = tf.reduce_sum(
        tf.square(tf.einsum("ijkl,mjkl->imjk", y_pics, y_pics)))

    # print(flipx_score)
    # print( tf.reduce_max(tf.stack([flipx_score, flipy_score, trans_score, anti_trans_score, center_score]), axis=(0,)))

    # idxs = vector_to_index(y_pics)
    # ncol = tf.map_fn(lambda p: tf.size(
    #     tf.unique(tf.reshape(p, [-1]),)[0], out_type=tf.int64), idxs)
    # print(ncol)

    return tf.reduce_sum([
        -1.0/(64) * tf.reduce_sum(tf.reduce_max(tf.stack([flipx_score, flipy_score,
                                                          trans_score, anti_trans_score, center_score]), axis=(0,))),
        -1.0/(64*9) * 0.1 * neigh,
        #   -1.0/7 * 10 * tf.cast(ncol, tf.float32),
        # 1.0/64 * 0.01 * tf.reduce_sum(tf.square(y_pics)),
        1.0/64 * 10 * divers
    ], )


@tf.function
def compute_score(pic1, pic2):
    return tf.reduce_sum(
        tf.abs(
            tf.einsum(
                "...k,...k->...",
                pic1,
                pic2)), axis=(-2, -1))
    # return tf.reduce_sum(
    #     -tf.math.maximum(0.0,
    #                     -tf.einsum(
    #                         "...k,...k->...",
    #                         pic1,
    #                         pic2)), axis=(-2, -1))