import os
import sys
os.environ["CUDA_VISIBLE_DEVICES"] = '2'
import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print(e)

import numpy as np
import plotly.graph_objects as go
from custom_utils.smpl import SMPL
from custom_utils.supersmpl_io import writeOBJ
from custom_utils.custom_utils import _maskFaces


def getAdjMatrix(faces):
    adj = np.zeros((np.max(faces[:])+1, np.max(faces[:])+1))

    for face in faces:
        adj[face[0], face[1]] = 1
        adj[face[1], face[0]] = 1
        adj[face[2], face[1]] = 1
        adj[face[0], face[2]] = 1
        adj[face[1], face[2]] = 1
        adj[face[2], face[0]] = 1

    return adj


if __name__ == '__main__':

    smpl = SMPL()

    pose = np.zeros((1, 72), dtype=np.float32)
    shape = np.zeros((1, 11), dtype=np.float32)
    body = smpl.run(smpl.run_rodrigues(pose), shape, remove=False)[0].numpy()[0]

    faces = np.load(os.path.join('/home/workspace/data/dataset/cape/cape_release', 'misc/smpl_tris.npy')).astype(np.int32)
    mask = np.zeros((27578), dtype=bool)
    mask[:6890] = True
    verts, faces = _maskFaces(body, faces, mask)
    adj = getAdjMatrix(faces.astype(np.int32))
    adj = np.triu(adj)
    edges = np.stack(np.nonzero(adj)).transpose()

    x = body[:6890]
    fig = go.Figure(data=[
            go.Mesh3d(
                # 8 vertices of a cube
                x=x[:,0],
                y=x[:,1],
                z=x[:,2],
                # i, j and k give the vertices of triangles
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2]
            ),
        ])

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

    fig.show()

    writeOBJ('temp/smpl_orig_f.obj', x, faces)

    print()