import torch
import trimesh
import numpy as np
import skimage.io as io
from PIL import Image
from smplx import SMPL
from matplotlib import cm as mpl_cm, colors as mpl_colors
from trimesh.visual.color import face_to_vertex_color, vertex_to_face_color, to_rgba

from common import constants
from .colorwheel import make_color_wheel_image


def get_smpl_uv():
    uv_obj = 'data/body_models/smpl_uv_20200910/smpl_uv.obj'

    uv_map = []
    with open(uv_obj) as f:
        for line in f.readlines():
            if line.startswith('vt'):
                coords = [float(x) for x in line.split(' ')[1:]]
                uv_map.append(coords)

    uv_map = np.array(uv_map)

    return uv_map


def show_uv_texture():
    # image = io.imread('data/body_models/smpl_uv_20200910/smpl_uv_20200910.png')
    image = make_color_wheel_image(1024, 1024)
    image = Image.fromarray(image)

    uv = np.load('data/body_models/smpl_uv_20200910/uv_table.npy') # get_smpl_uv()
    material = trimesh.visual.texture.SimpleMaterial(image=image)
    tex_visuals = trimesh.visual.TextureVisuals(uv=uv, image=image, material=material)

    smpl = SMPL(constants.SMPL_MODEL_DIR)

    faces = smpl.faces
    verts = smpl().vertices[0].detach().numpy()

    # assert(len(uv) == len(verts))
    print(uv.shape)
    vc = tex_visuals.to_color().vertex_colors
    fc = trimesh.visual.color.vertex_to_face_color(vc, faces)
    face_colors = fc.copy()
    fc = fc.astype(float)
    vc = vc.astype(float)
    fc[:,:3] = fc[:,:3] / 255.
    vc[:,:3] = vc[:,:3] / 255.
    print(fc[:,:3].max(), fc[:,:3].min(), fc[:,:3].mean())
    print(vc[:, :3].max(), vc[:, :3].min(), vc[:, :3].mean())
    np.save('data/body_models/smpl/color_wheel_face_colors.npy', fc)
    np.save('data/body_models/smpl/color_wheel_vertex_colors.npy', vc)
    print(fc.shape)
    mesh = trimesh.Trimesh(verts, faces, validate=True, process=False, face_colors=face_colors)
    # mesh = trimesh.load('data/body_models/smpl_uv_20200910/smpl_uv.obj', process=False)
    # mesh.visual = tex_visuals

    # import ipdb; ipdb.set_trace()
    # print(vc.shape)
    mesh.show()


def show_colored_mesh():
    cm = mpl_cm.get_cmap('jet')
    norm_gt = mpl_colors.Normalize()

    smpl = SMPL(constants.SMPL_MODEL_DIR)

    faces = smpl.faces
    verts = smpl().vertices[0].detach().numpy()

    m = trimesh.Trimesh(verts, faces, process=False)

    mode = 1
    if mode == 0:
        # mano_segm_labels = m.triangles_center
        face_labels = m.triangles_center
        face_colors = (face_labels - face_labels.min()) / np.ptp(face_labels)

    elif mode == 1:
        # print(face_labels.shape)
        face_labels = m.triangles_center
        face_labels = np.argsort(np.linalg.norm(face_labels, axis=-1))
        face_colors = np.ones((13776, 4))
        face_colors[:, 3] = 1.0
        face_colors[:, :3] = cm(norm_gt(face_labels))[:, :3]
    elif mode == 2:
        # breakpoint()
        fc = np.load('data/body_models/smpl_uv_20200910/data/vertex_texture.npy')[0, :, 0, 0, 0, :]
        face_colors = np.ones((13776, 4))
        face_colors[:, :3] = fc
    mesh = trimesh.Trimesh(verts, faces, process=False, face_colors=face_colors)
    mesh.show()


def get_tenet_texture(mode='smplpix'):
    # mode = 'smplpix', 'decomr'

    smpl = SMPL(constants.SMPL_MODEL_DIR)

    faces = smpl.faces
    verts = smpl().vertices[0].detach().numpy()

    m = trimesh.Trimesh(verts, faces, process=False)
    if mode == 'smplpix':
        # mano_segm_labels = m.triangles_center
        face_labels = m.triangles_center
        face_colors = (face_labels - face_labels.min()) / np.ptp(face_labels)
        texture = np.zeros((1, faces.shape[0], 1, 1, 1, 3), dtype=np.float32)
        texture[0, :, 0, 0, 0, :] = face_colors[:, :3]
        texture = torch.from_numpy(texture).float()
    elif mode == 'decomr':
        texture = np.load('data/body_models/smpl_uv_20200910/data/vertex_texture.npy')
        texture = torch.from_numpy(texture).float()
    elif mode == 'colorwheel':
        face_colors = np.load('data/body_models/smpl/color_wheel_face_colors.npy')
        texture = np.zeros((1, faces.shape[0], 1, 1, 1, 3), dtype=np.float32)
        texture[0, :, 0, 0, 0, :] = face_colors[:, :3]
        texture = torch.from_numpy(texture).float()
    else:
        raise ValueError(f'{mode} is not defined!')

    return texture


def save_tenet_textures(mode='smplpix'):
    # mode = 'smplpix', 'decomr'

    smpl = SMPL(constants.SMPL_MODEL_DIR)

    faces = smpl.faces
    verts = smpl().vertices[0].detach().numpy()

    m = trimesh.Trimesh(verts, faces, process=False)

    if mode == 'smplpix':
        # mano_segm_labels = m.triangles_center
        face_labels = m.triangles_center
        face_colors = (face_labels - face_labels.min()) / np.ptp(face_labels)
        texture = np.zeros((1, faces.shape[0], 1, 1, 1, 3), dtype=np.float32)
        texture[0, :, 0, 0, 0, :] = face_colors[:, :3]
        texture = torch.from_numpy(texture).float()

        vert_colors = face_to_vertex_color(m, face_colors).astype(float) / 255.0

    elif mode == 'decomr':
        texture = np.load('data/body_models/smpl_uv_20200910/data/vertex_texture.npy')
        texture = torch.from_numpy(texture).float()
        face_colors = texture[0, :, 0, 0, 0, :]
        vert_colors = face_to_vertex_color(m, face_colors).astype(float) / 255.0

    elif mode == 'colorwheel':
        face_colors = np.load('data/body_models/smpl/color_wheel_face_colors.npy')
        texture = np.zeros((1, faces.shape[0], 1, 1, 1, 3), dtype=np.float32)
        texture[0, :, 0, 0, 0, :] = face_colors[:, :3]
        texture = torch.from_numpy(texture).float()
        face_colors[:, :3] *= 255
        vert_colors = face_to_vertex_color(m, face_colors).astype(float) / 255.0
    else:
        raise ValueError(f'{mode} is not defined!')

    print(vert_colors.shape, vert_colors.max())
    np.save(f'data/body_models/smpl/{mode}_vertex_colors.npy', vert_colors)
    return texture