from math import pi
import tensorflow as tf

squares = [chr(i) for i in range(0x1F7E5, 0x1F7EC)]
tfsquares = tf.constant(squares)
colors = tf.constant([10, 219, 38, 48, 80,  282, 20] *
                     tf.constant(pi/180.0), dtype=tf.float32)
color_vectors = tf.transpose(
    tf.stack([
        tf.math.cos(colors),
        tf.math.sin(colors)]
    )
)


def vector_to_index(tensor):
    return tf.argmax((tf.einsum("...ijk,lk->...ijl", tensor, color_vectors)), axis=-1)


def index_to_string(tensor):
    tfstring = tf.strings.join(
        tf.map_fn(
            lambda v:
            tf.strings.join(
                tf.gather(tfsquares, tf.cast(v, tf.int64))), tensor, fn_output_signature=tf.string), "\n")

    return tfstring.numpy().decode()


def vector_to_string(tensor):
    return index_to_string(vector_to_index(tensor))