import argparse
import cv2
import keyboard
import numpy as np
import open3d as o3d
import os
import pygame
from transforms3d.axangles import axangle2mat

import config

from hand_mesh import HandMesh
from kinematics import mpii_to_mano
from utils import OneEuroFilter, imresize
from wrappers import ModelPipeline
from utils import *

def video_to_images(vid_file, img_folder=None, return_info=False):
    if img_folder is None:
        img_folder = osp.join('/tmp', osp.basename(vid_file).replace('.', '_'))

    os.makedirs(img_folder, exist_ok=True)
    command = ['ffmpeg',
               '-i', vid_file,
               '-f', 'image2',
               '-v', 'error',
               f'{img_folder}/%06d.png']
    print(f'Running \"{" ".join(command)}\"')
    subprocess.call(command)

    print(f'Images saved to \"{img_folder}\"')

    img_shape = cv2.imread(osp.join(img_folder, '000001.png')).shape

    if return_info:
        return img_folder, len(os.listdir(img_folder)), img_shape
    else:
        return img_folder


def run(args):
    ############ output visualization ############
    # view_mat = axangle2mat([1, 0, 0], np.pi) # align different coordinate systems
    # window_size = 1080

    # hand_mesh = HandMesh(config.HAND_MESH_MODEL_PATH)
    # mesh = o3d.geometry.TriangleMesh()
    # mesh.triangles = o3d.utility.Vector3iVector(hand_mesh.faces)
    # mesh.vertices = \
    # o3d.utility.Vector3dVector(np.matmul(view_mat, hand_mesh.verts.T).T * 1000)
    # mesh.compute_vertex_normals()

    # viewer = o3d.visualization.Visualizer()
    # viewer.create_window(
    # width=window_size + 1, height=window_size + 1,
    # window_name='Minimal Hand - output'
    # )
    # viewer.add_geometry(mesh)

    # view_control = viewer.get_view_control()
    # cam_params = view_control.convert_to_pinhole_camera_parameters()
    # extrinsic = cam_params.extrinsic.copy()
    # extrinsic[0:3, 3] = 0
    # cam_params.extrinsic = extrinsic
    # cam_params.intrinsic.set_intrinsics(
    # window_size + 1, window_size + 1, config.CAM_FX, config.CAM_FY,
    # window_size // 2, window_size // 2
    # )
    # view_control.convert_from_pinhole_camera_parameters(cam_params)
    # view_control.set_constant_z_far(1000)

    # render_option = viewer.get_render_option()
    # render_option.load_from_json('./render_option.json')
    # viewer.update_renderer()

    # ############ input visualization ############
    # pygame.init()
    # display = pygame.display.set_mode((window_size, window_size))
    # pygame.display.set_caption('Minimal Hand - input')

    # ############ misc ############
    # mesh_smoother = OneEuroFilter(4.0, 0.0)
    # clock = pygame.time.Clock()

    ############ Move all of above code to local to render ###########

    video_file = args.vid_file
    if not os.path.isfile(video_file):
        exit(f'Input video \"{video_file}\" does not exist!')

    output_path = os.path.join(args.output_folder, os.path.basename(video_file).replace('.mp4', ''))
    os.makedirs(output_path, exist_ok=True)

    image_folder, num_frames, img_shape = video_to_images(video_file, return_info=True)

    print(f'Input video number of frames {num_frames}')
    orig_height, orig_width = img_shape[:2]

    # total_time = time.time()
    import pdb; pdb.set_trace()

    image_file_names = [
            osp.join(image_folder, x)
            for x in os.listdir(image_folder)
            if x.endswith('.png') or x.endswith('.jpg')
        ]

    model = ModelPipeline()

    for i in image_file_names:
        # What do all these conditions check for?
        frame_large = x
        if frame_large is None:
            continue
        if frame_large.shape[0] > frame_large.shape[1]:
            margin = int((frame_large.shape[0] - frame_large.shape[1]) / 2)
            frame_large = frame_large[margin:-margin]
        else:
            margin = int((frame_large.shape[1] - frame_large.shape[0]) / 2)
            frame_large = frame_large[:, margin:-margin]

        frame_large = np.flip(frame_large, axis=1).copy() # why? Camera flip?
        frame = imresize(frame_large, (128, 128))  # needed

        ######## Golden lines, run this here #########
        _, theta_mpii = model.process(frame) 
        theta_mano = mpii_to_mano(theta_mpii)

        ######## Save theta_mano and pass as input to local ######## 
        v = hand_mesh.set_abs_quat(theta_mano)
        v *= 2 # for better visualization
        v = v * 1000 + np.array([0, 0, 400])
        v = mesh_smoother.process(v)

        mesh.triangles = o3d.utility.Vector3iVector(hand_mesh.faces)
        mesh.vertices = o3d.utility.Vector3dVector(np.matmul(view_mat, v.T).T)
        mesh.paint_uniform_color(config.HAND_COLOR)
        mesh.compute_triangle_normals()
        mesh.compute_vertex_normals()
        # for some version of open3d you may need `viewer.update_geometry(mesh)`
        viewer.update_geometry()

        viewer.poll_events()

        display.blit(
            pygame.surfarray.make_surface(
            np.transpose(
                imresize(frame_large, (window_size, window_size)
            ), (1, 0, 2))
            ),
            (0, 0)
        )
        pygame.display.update()

        if keyboard.is_pressed("esc"):
            break

        clock.tick(30) # What's this do? If it adds delay remove it


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--vid_file', type=str,
                    help='input video path or youtube link')
                    
    args = parser.parse_args()
    run(args)

[((22, 4, 22, 42), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((92, 4, 92, 43), 'os.makedirs', 'os.makedirs', (), '', False, 'import os\n'), ((100, 16, 100, 31), 'pdb.set_trace', 'pdb.set_trace', ({}, {}), '()', False, 'import pdb\n'), ((108, 12, 108, 27), 'wrappers.ModelPipeline', 'ModelPipeline', ({}, {}), '()', False, 'from wrappers import ModelPipeline\n'), ((162, 13, 162, 38), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((88, 11, 88, 37), 'os.path.isfile', 'os.path.isfile', ({(88, 26, 88, 36): 'video_file'}, {}), '(video_file)', False, 'import os\n'), ((123, 16, 123, 49), 'utils.imresize', 'imresize', ({(123, 25, 123, 36): 'frame_large', (123, 38, 123, 48): '(128, 128)'}, {}), '(frame_large, (128, 128))', False, 'from utils import OneEuroFilter, imresize\n'), ((127, 21, 127, 45), 'kinematics.mpii_to_mano', 'mpii_to_mano', ({(127, 34, 127, 44): 'theta_mpii'}, {}), '(theta_mpii)', False, 'from kinematics import mpii_to_mano\n'), ((135, 25, 135, 68), 'open3d.utility.Vector3iVector', 'o3d.utility.Vector3iVector', ({(135, 52, 135, 67): 'hand_mesh.faces'}, {}), '(hand_mesh.faces)', True, 'import open3d as o3d\n'), ((153, 8, 153, 31), 'pygame.display.update', 'pygame.display.update', ({}, {}), '()', False, 'import pygame\n'), ((155, 11, 155, 37), 'keyboard.is_pressed', 'keyboard.is_pressed', ({(155, 31, 155, 36): '"""esc"""'}, {}), "('esc')", False, 'import keyboard\n'), ((104, 21, 104, 45), 'os.listdir', 'os.listdir', ({(104, 32, 104, 44): 'image_folder'}, {}), '(image_folder)', False, 'import os\n'), ((132, 23, 132, 44), 'numpy.array', 'np.array', ({(132, 32, 132, 43): '[0, 0, 400]'}, {}), '([0, 0, 400])', True, 'import numpy as np\n'), ((36, 31, 36, 53), 'os.listdir', 'os.listdir', ({(36, 42, 36, 52): 'img_folder'}, {}), '(img_folder)', False, 'import os\n'), ((91, 51, 91, 79), 'os.path.basename', 'os.path.basename', ({(91, 68, 91, 78): 'video_file'}, {}), '(video_file)', False, 'import os\n'), ((122, 22, 122, 50), 'numpy.flip', 'np.flip', (), '', True, 'import numpy as np\n'), ((136, 51, 136, 75), 'numpy.matmul', 'np.matmul', ({(136, 61, 136, 69): 'view_mat', (136, 71, 136, 74): 'v.T'}, {}), '(view_mat, v.T)', True, 'import numpy as np\n'), ((148, 16, 149, 13), 'utils.imresize', 'imresize', ({(148, 25, 148, 36): 'frame_large', (148, 38, 148, 64): '(window_size, window_size)'}, {}), '(frame_large, (window_size, window_size))', False, 'from utils import OneEuroFilter, imresize\n')]