import pickle
import os
import h5py
import sys
import numpy as np
import open3d as o3d
from snapshot_smpl.smpl import Smpl
import cv2
import tqdm


def read_pickle(pkl_path):
    with open(pkl_path, 'rb') as f:
        u = pickle._Unpickler(f)
        u.encoding = 'latin1'
        return u.load()


def get_KRTD(camera):
    K = np.zeros([3, 3])
    K[0, 0] = camera['camera_f'][0]
    K[1, 1] = camera['camera_f'][1]
    K[:2, 2] = camera['camera_c']
    K[2, 2] = 1
    R = np.eye(3)
    T = np.zeros([3])
    D = camera['camera_k']
    return K, R, T, D


def get_o3d_mesh(vertices, faces):
    mesh = o3d.geometry.TriangleMesh()
    mesh.vertices = o3d.utility.Vector3dVector(vertices)
    mesh.triangles = o3d.utility.Vector3iVector(faces)
    mesh.compute_vertex_normals()
    return mesh


def get_smpl(base_smpl, betas, poses, trans):
    base_smpl.betas = betas
    base_smpl.pose = poses
    base_smpl.trans = trans
    vertices = np.array(base_smpl)

    faces = base_smpl.f
    mesh = get_o3d_mesh(vertices, faces)

    return vertices, mesh


def render_smpl(mesh, img, K, R, T):
    vertices = np.array(mesh.vertices)
    rendered_img = renderer.render_multiview(vertices, K[None], R[None],
                                             T[None, None], [img])[0]
    return rendered_img


def extract_image(data_path):
    data_root = os.path.dirname(data_path)
    img_dir = os.path.join(data_root, 'image')
    os.system('mkdir -p {}'.format(img_dir))

    if len(os.listdir(img_dir)) >= 200:
        return

    cap = cv2.VideoCapture(data_path)

    ret, frame = cap.read()
    i = 0

    while ret:
        cv2.imwrite(os.path.join(img_dir, '{}.jpg'.format(i)), frame)
        ret, frame = cap.read()
        i = i + 1

    cap.release()


def extract_mask(masks, mask_dir):
    if len(os.listdir(mask_dir)) >= len(masks):
        return

    for i in tqdm.tqdm(range(len(masks))):
        mask = masks[i].astype(np.uint8)

        # erode the mask
        border = 4
        kernel = np.ones((border, border), np.uint8)
        mask = cv2.erode(mask.copy(), kernel)

        cv2.imwrite(os.path.join(mask_dir, '{}.png'.format(i)), mask)


data_root = 'data/people_snapshot'
videos = ['female-3-casual']

model_paths = [
    'basicModel_f_lbs_10_207_0_v1.0.0.pkl',
    'basicmodel_m_lbs_10_207_0_v1.0.0.pkl'
]

for video in videos:
    camera_path = os.path.join(data_root, video, 'camera.pkl')
    camera = read_pickle(camera_path)
    K, R, T, D = get_KRTD(camera)

    # process video
    video_path = os.path.join(data_root, video, video + '.mp4')
    extract_image(video_path)

    # process mask
    mask_path = os.path.join(data_root, video, 'masks.hdf5')
    masks = h5py.File(mask_path)['masks']
    mask_dir = os.path.join(data_root, video, 'mask')
    os.system('mkdir -p {}'.format(mask_dir))
    extract_mask(masks, mask_dir)

    smpl_path = os.path.join(data_root, video, 'reconstructed_poses.hdf5')
    smpl = h5py.File(smpl_path)
    betas = smpl['betas']
    pose = smpl['pose']
    trans = smpl['trans']

    pose = pose[len(pose) - len(masks):]
    trans = trans[len(trans) - len(masks):]

    # process smpl parameters
    params = {'beta': np.array(betas), 'pose': pose, 'trans': trans}
    params_path = os.path.join(data_root, video, 'params.npy')
    np.save(params_path, params)

    if 'female' in video:
        model_path = model_paths[0]
    else:
        model_path = model_paths[1]
    model_data = read_pickle(model_path)

    img_dir = os.path.join(data_root, video, 'image')
    vertices_dir = os.path.join(data_root, video, 'vertices')
    os.system('mkdir -p {}'.format(vertices_dir))

    num_img = len(os.listdir(img_dir))
    for i in tqdm.tqdm(range(num_img)):
        base_smpl = Smpl(model_data)
        vertices, mesh = get_smpl(base_smpl, betas, pose[i], trans[i])
        np.save(os.path.join(vertices_dir, '{}.npy'.format(i)), vertices)