import os
import numpy as np

import numpy as np
import pickle as pkl
from scipy import sparse

from time import time
import plotly.graph_objects as go

from custom_utils.supersmpl_io import writeOBJ, readOBJ


def makeGraph(F):
    G = {}
    for f in F:
        N = len(f)
        for i,v in enumerate(f):
            if not v in G: G[v] = {}
            G[v][f[(i+1) % N]] = 1
    return G

def extrapolateMissings(M, G):
    for i in range(6890, M.shape[0]):
        if np.isnan(M[i]).all():
            M[i] = neighAvg(M, G, i)
        elif np.isnan(M[i]).any():
            print("HOW?")
            print(M[i])
    if np.isnan(M).any(): extrapolateMissings(M, G)

def neighAvg(M, G, i):
    neigh = set(G[i].keys())
    avg = 0
    count = 0
    for n in neigh:
        if not np.isnan(M[n]).any():
            avg += M[n]
            count += 1
    if count: return avg / count
    return np.nan


def convert(smpl, G):
    supersmpl = np.empty((27578, 3), np.float32)
    supersmpl[:] = np.nan
    supersmpl[:6890] = smpl
    extrapolateMissings(supersmpl, G)
    return supersmpl


if __name__ == '__main__':

    g = 'f'
    # Read Super topology
    V, F, Vt, Ft = readOBJ('super_' + g + '_avg.obj')
    G = makeGraph(F)

    folder = '/home/workspace/data/dataset/cape/cape_release'
    test = np.load(os.path.join('/home/workspace/data/dataset/cape_zips/03383/shortlong_ATUsquat/shortlong_ATUsquat.000105.npz'))
    faces = np.load(os.path.join(folder, 'misc/smpl_tris.npy'))
    min_body = np.load(os.path.join(folder, 'minimal_body_shape/03383/03383_minimal.npy')) 
    offsets_cano = test['v_cano'] - min_body

    supersmpl = convert(test['v_posed'], G)

    x = supersmpl
    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()

    print()