import os
import sys
import cv2
import numpy as np
import trimesh as tm
import matplotlib.pyplot as plt
import plotly.graph_objects as go


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


def uv_to_mesh(uv):
    uv_res = 256
    # m2uv = np.load('/workspace/Code/featurelines/MeshUtils/m2uv.npy', allow_pickle=True).item()
    m2uv = np.load('custom_utils\m2uv.npy', allow_pickle=True).item()
    m2uv = m2uv_array(m2uv)
    # uv_map = np.load('/workspace/Code/featurelines/Meshutils/stretched_uv_map_3.npy')
    uv_map = np.load('custom_utils\stretched_uv_map_3.npy')
    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)

    mesh = np.zeros((14475, 3))
    uv_pixels = uv_map_discreet[m2uv]
    # counter = np.ones((14475))
    # counter[m2uv[:, 0] == m2uv[:, 1]] = 2
    mesh = np.mean(uv[uv_pixels[:, :, 0], uv_pixels[:, :, 1], :], axis=1) # / counter[:, np.newaxis]
    return mesh


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

def plot(uv, mesh, faces, body, F, template, uv_gt, mesh_gt):
    # fig = plt.figure(figsize=(15, 15))
    # ax = fig.add_subplot(111, projection='3d')
    # ax.set_xlim3d(-1, 1)
    # ax.set_ylim3d(-1, 1)
    # ax.set_zlim3d(-1, 1)
    # ax.set_xlabel("X")
    # ax.set_ylabel("Y")
    # ax.set_zlabel("Z")
    # ax.scatter(mesh[:, 0], mesh[:, 1], mesh[:, 2], s=5)
    # # plt.show()

    fig = plt.figure(figsize=(15, 15))
    ax = fig.add_subplot(131)
    ax.imshow(np.round((uv + 1) * 128).astype(np.int32))
    ax = fig.add_subplot(132)
    ax.imshow(np.round((template + 1) * 128).astype(np.int32))
    ax = fig.add_subplot(133)
    ax.imshow(np.round((uv_gt + 1) * 128).astype(np.int32))
    plt.show()

    collision = np.concatenate((body, mesh), axis=0)
    collision_faces = np.concatenate((F, faces + 14475), axis=0)
    collision_color = np.concatenate((np.ones_like(body, dtype=np.int32)*100, np.ones_like(mesh, dtype=np.int32)*230), axis=0)


    plot_3d(mesh_gt, mesh, faces, collision, collision_faces, collision_color)
    tm.Trimesh(mesh, faces).show()
    tm.Trimesh(mesh_gt, faces).show()
    # tm.Trimesh(collision, collision_faces, vertex_colors=collision_color).show()


def plot_3d(mesh_gt, mesh_pred, faces, collision, collision_faces, collision_color):
    mesh_pred[:,1] += 1
    collision[:,1] += 2
    fig = go.Figure(data=[
        go.Mesh3d(
            # 8 vertices of a cube
            x=mesh_gt[:,0],
            y=mesh_gt[:,1],
            z=mesh_gt[:,2],
            # i, j and k give the vertices of triangles
            i = faces[:,0],
            j = faces[:,1],
            k = faces[:,2]
        ),
        go.Mesh3d(
            # 8 vertices of a cube
            x=mesh_pred[:,0],
            y=mesh_pred[:,1],
            z=mesh_pred[:,2],
            # i, j and k give the vertices of triangles
            i = faces[:,0],
            j = faces[:,1],
            k = faces[:,2],
        ),
        go.Mesh3d(
            # 8 vertices of a cube
            x=collision[:,0],
            y=collision[:,1],
            z=collision[:,2],
            # i, j and k give the vertices of triangles
            i = collision_faces[:,0],
            j = collision_faces[:,1],
            k = collision_faces[:,2],
            vertexcolor=collision_color
        )
    ])

    fig.update_layout(scene_aspectmode='manual',
                    scene_aspectratio=dict(x=1, y=2, z=1))
    fig.update_layout(
        scene = dict(
            xaxis = dict(nticks=4, range=[-1,1],),
                        yaxis = dict(nticks=4, range=[-1,3],),
                        zaxis = dict(nticks=4, range=[-1,1],),)
                        )

    fig.show()


if __name__ == '__main__':
    # file = sys.argv[1]

    for i in range(50):
        file = r'Z:\Work\Huni\Barcelona\Output\stylegan\test-15-115000\\' + str(i) + '.npz'

        F = np.load('custom_utils\smpl_faces.npy')
        data = np.load(file)
        uv = np.transpose(data['uv'], [1,2,0])
        uv_gt = np.transpose(data['gt'], [1,2,0])
        mask = data['mask'].astype(bool)
        template = np.transpose(data['template'], [1,2,0])


        mesh = uv_to_mesh(uv)
        mesh_gt = uv_to_mesh(uv_gt)

        template_mesh = uv_to_mesh(template)
        filter = np.abs(np.sum(mesh, axis=-1)) < 1e-2
        overlap = np.sum((((np.sum(template_mesh, axis=-1) * mask) * (np.sum(mesh, axis=-1) * filter * mask))) > 0)

        mesh, faces = _maskFaces(mesh, F, mask)
        plot(uv, mesh, faces, data['body'], F, template, uv_gt, mesh_gt[mask, :])

        pass

    pass