import numpy as np
import tensorflow as tf


def m2uv_array(m2uv):
    result = np.zeros((14475, 2), dtype=np.int32)
    for k, v in m2uv.items():
        if len(v) == 2:
            result[k] = list(v)
        elif len(v) == 1:
            x = v.pop()
            result[k] = [x, x]
        else:
            assert False

    return result


class UV2Mesh():
    def __init__(self, res=256):
        super(UV2Mesh, self).__init__()
        self.res = res
        uv_map = np.load('custom_utils/stretched_uv_map_3.npy')
        self.uv_map_discreet = np.floor(uv_map * self.res).astype(np.int32)
        m2uv = np.load('custom_utils/m2uv.npy', allow_pickle=True).item()
        m2uv = m2uv_array(m2uv)
        self.uv_pixels = self.uv_map_discreet[m2uv]
        # self.batch_size = batch_size

        # self.x1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.zeros((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
        # self.y1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
        # self.z1 = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
        # self.x2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.zeros((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
        # self.y2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
        # self.z2 = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [self.batch_size, 1, 1]), tf.ones((self.batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)

    # input = channels_first
    @tf.function
    def uv_to_mesh(self, uv):
        uv = tf.transpose(uv, [0,2,3,1])
        batch_size = uv.shape[0]

        x1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
        y1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
        z1_ind = tf.concat((tf.tile(self.uv_pixels[:, 0, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)
        x2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.zeros((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
        y2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32)), axis=-1)
        z2_ind = tf.concat((tf.tile(self.uv_pixels[:, 1, :][tf.newaxis, :, :], [batch_size, 1, 1]), tf.ones((batch_size, 14475, 1), dtype=tf.int32) * 2), axis=-1)


        x1 = tf.gather_nd(uv, x1_ind, batch_dims=1)
        y1 = tf.gather_nd(uv, y1_ind, batch_dims=1)
        z1 = tf.gather_nd(uv, z1_ind, batch_dims=1)
        x2 = tf.gather_nd(uv, x2_ind, batch_dims=1)
        y2 = tf.gather_nd(uv, y2_ind, batch_dims=1)
        z2 = tf.gather_nd(uv, z2_ind, batch_dims=1)

        p1 = tf.stack((x1, y1, z1), axis=-1)
        p2 = tf.stack((x2, y2, z2), axis=-1)

        mesh = (p1 + p2) / 2
        return mesh

    def __call__(self, uv):
        return self.uv_to_mesh(uv)
        


def _maskFaces(V, F, mask):
    V_mask = mask
    F_mask = np.take(V_mask, F.reshape(-1), axis=0)
    F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3

    V_ind = np.zeros(V_mask.shape)
    V_ind[V_mask] = np.arange(V_mask.sum())
    F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
    V_new = V[V_mask, :]

    return V_new, F_new

def maskFacesBatch(V, F, mask):
    Vs = []
    Fs = []
    for i in range(mask.shape[0]):
        v, f = _maskFaces(V[i], F, mask[i])
        Vs.append(v)
        Fs.append(f)
    return Vs, Fs