import os
import numpy as np
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 plotly.graph_objects as go
import plotly.express as ex
from custom_utils.smpl import SMPL
from custom_utils.custom_utils import m2uv_array, _maskFaces

is_skirt = False

smpl = SMPL(orig_smpl=False)
smpl_orig = SMPL(orig_smpl=True)



def get_faces_orig_smpl():
    faces = np.load('custom_utils/smpl_faces_orig.npy')
    return faces


def filter_uv_map():
    uv_map = np.load('custom_utils/stretched_uv_map_' + ("skirt" if is_skirt == 1 else "3") + '.npy')
    m2uv = np.load('custom_utils/m2uv' + ("_skirt" if is_skirt == 1 else "") + '.npy', allow_pickle=True).item()
    # m2uv = m2uv_array(m2uv)
    m2uv_orig = { key: m2uv[key] for key in m2uv.keys() if key < 3638 }

    all_ids = []
    for k in m2uv_orig.keys():
        all_ids.extend(list(m2uv_orig[k]))
    all_ids = np.asarray(all_ids)

    uv_map_orig = np.asarray([uv_map[i] for i in range(uv_map.shape[0]) if i in all_ids])
    mapping = [i for i in range(uv_map.shape[0]) if i in all_ids]

    new_m2uv = dict()
    for key in m2uv_orig.keys():
        elems = list(m2uv_orig[key])
        new_elem = set()
        for elem in elems:
            new_elem.add(mapping.index(elem))
        new_m2uv[key] = new_elem

    np.save('custom_utils/stretched_uv_map_3_orig.npy', uv_map_orig)
    np.save('custom_utils/m2uv_orig.npy', new_m2uv)
    print()


def test():
    from custom_utils.custom_utils import _maskFaces

    pose = np.zeros((1, 72), dtype=np.float32)
    shape = np.zeros((1, 11), dtype=np.float32)
    body = smpl_orig.run(smpl_orig.run_rodrigues(pose), shape, remove=True)[0].numpy()[0]
    x = body
    faces = get_faces_orig_smpl()
    y = x
    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]
            ),
            go.Mesh3d(
                # 8 vertices of a cube
                x=y[:,0],
                y=y[:,1],
                z=y[:,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.5,1.5],),
            yaxis = dict(nticks=4, range=[-1.5,1.5],),
            zaxis = dict(nticks=4, range=[-1.5,1.5],),)
            )

    fig.show()

if __name__ == '__main__':
    # test()
    filter_uv_map()