import os
import sys
import random
import numpy as np
import plotly.graph_objects as go

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
from data.cloth3d_io import read_garment_mask, read_garment_template, read_posed_garment_frame, loadInfo, read_posed_body_frame, read_posed_body, read_posed_garment, read_curvatures, read_curvature_frame, writeOBJ
from custom_utils.custom_utils import _maskFaces


def get_frames(folder, frames_file):
    frames = []
    with open(os.path.join(folder, frames_file), 'r') as f:
        for line in f:
            frames.append(line.strip())
    return frames


def plot(mesh, faces):
    fig = go.Figure()

    fig.add_trace(go.Mesh3d(
                x=mesh[:,0],
                y=mesh[:,1],
                z=mesh[:,2],
                i = faces[:,0],
                j = faces[:,1],
                k = faces[:,2]
            ),)
    
    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
                    )
    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
                yaxis = dict(nticks=4, range=[-1,1],),
                zaxis = dict(nticks=4, range=[-1,1],),)
    fig.update_layout(scene = helper,
                        )

    fig.show()


def build_lookup_table(repeated):
    # iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element

    lookup = {}
    for i in range(len(repeated)):
        for j in range(len(repeated[i])):
            if len(repeated[i]) == 1:
                continue
            lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
    return lookup

def mask_faces(all_faces, mask, repeated, verts):
    # if any of the vertices of a face is masked, mask the face
    lookup = build_lookup_table(repeated)
    faces = []
    for face in all_faces:
        # if any two or all three vertices of the face have the same coordinates then continue
        if np.allclose(verts[face[0]], verts[face[1]]) or np.allclose(verts[face[0]], verts[face[2]]) or np.allclose(verts[face[1]], verts[face[2]]):
            continue

        if mask[face[0]] and mask[face[1]] and mask[face[2]]:
            faces.append(face)
        else:
            # create three arrays of vertices from the lookup table for each vertex of the face
            # but only for those vertices that are in the lookup table, otherwise createa a single element array with the vertex
            v1 = lookup[face[0]] if face[0] in lookup else [face[0]]
            v2 = lookup[face[1]] if face[1] in lookup else [face[1]]
            v3 = lookup[face[2]] if face[2] in lookup else [face[2]]
            # try each combination of vertices from the three arrays and check if the face is masked
            found_replacement = False
            for i in range(len(v1)):
                if not found_replacement:
                    for j in range(len(v2)):
                        if not found_replacement:
                            for k in range(len(v3)):
                                if not found_replacement:
                                    if mask[v1[i]] and mask[v2[j]] and mask[v3[k]]:
                                        # if the euclidean distance between any two of the vertices of the face is less than 0.05 then append the face
                                        # if np.linalg.norm(verts[v1[i]] - verts[v2[j]]) < 0.05 and np.linalg.norm(verts[v1[i]] - verts[v3[k]]) < 0.05 and np.linalg.norm(verts[v2[j]] - verts[v3[k]]) < 0.05:
                                        faces.append(np.array([v1[i], v2[j], v3[k]], dtype=np.int32))
                                        assert np.all(verts[face] == verts[faces[-1]])
                                        found_replacement = True
                                        break
                # for j in range(len(v2)):
                #     for k in range(len(v3)):
                #         if mask[v1[i]] and mask[v2[j]] and mask[v3[k]]:
                #             faces.append([v1[i], v2[j], v3[k]])
                #             break

    faces = np.array(faces)
    V_ind = np.zeros(mask.shape, dtype=np.int32)
    V_ind[mask] = np.arange(mask.sum())
    F_new = np.take(V_ind, faces.reshape(-1), axis=0).reshape((-1, 3))
    return F_new


def check(folder, frame):
    id, frame_nr, garment = frame.split()
    frame_nr = int(frame_nr)
    info = loadInfo(os.path.join(folder, id, 'info.mat'))
    mask = read_garment_mask(folder, id, garment)
    mesh_posed = read_posed_garment_frame(folder, id, garment, frame_nr, mask=mask, shape=info['shape'], pose=info['poses'][:,frame_nr], deorient=use_no_orient, unshape=no_shape)

    v = mesh_posed[mask,:]
    is_skirt = (garment == 'Dress' or garment == 'Skirt')
    skirt_mod = '_skirt' if is_skirt else ''
    F = np.load(f'custom_utils/smpl_faces{skirt_mod}.npy').astype(np.int32)
    _, faces = _maskFaces(np.zeros((14475,3)), F, mask)

    # plot(v, faces)

    unq, idx, count = np.unique(mesh_posed, axis=0, return_counts=True, return_index=True)
    repeated_groups = unq[count > 1]

    mask_2 = np.zeros((14475), dtype=bool)
    mask_2[idx] = True
    v_2 = mesh_posed[mask_2,:]
    _, faces_2 = _maskFaces(np.zeros((14475,3)), F, mask_2)
    # plot(v_2, faces_2)

    # writeOBJ('temp/investigate.obj', v, faces)

    repeated = []
    for repeated_group in repeated_groups:
        repeated_idx = np.argwhere(np.all(mesh_posed == repeated_group, axis=1))
        repeated.append(repeated_idx.ravel())

    faces_3 = mask_faces(F, mask_2, repeated, mesh_posed)
    writeOBJ('temp/investigate_3.obj', v_2, faces_3)
    plot(v_2, faces_3)


    pass
    


if __name__ == "__main__":
    folder = '/home/workspace/data/dataset/cloth3d/train'
    frames = get_frames(folder, 'frames_all.txt')
    root = '/'.join(folder.split('/')[:-1])
    type = folder.split('/')[-1]

    use_no_orient = False
    no_shape = False
    is_skirt = True

    skirt_modifier = '_skirt' if is_skirt else ''
    orient_mod = '_no_orient' if use_no_orient else ''
    shape_mod = '_no_shape' if no_shape else ''


    while True:
        frame = random.choice(frames)
        frame = '05936 0 Dress'

        check(folder, frame)

        pass