# Standard modules.
import copy
import os

# CV/ML.
import cv2

import numpy as np

from mediapipe.python.solutions import holistic

# For drawing.
from mediapipe.python.solutions.face_mesh_connections import (
    FACEMESH_CONTOURS)
from mediapipe.python.solutions.hands_connections import (
    HAND_CONNECTIONS)
from mediapipe.python.solutions.pose_connections import (
    POSE_CONNECTIONS)

from IPython.display import HTML
from base64 import b64encode
import io

def show_video(video_path, video_width=500, video_height=500):
    video = io.open(video_path, 'r+b').read()
    encoded = b64encode(video)
    decoded = encoded.decode("ascii")
    data = f"<video width={video_width} height={video_height} controls>" \
        + f'<source src="data:video/mp4;base64,{decoded}"' \
        + 'type="video/mp4" /></video>'
    return(HTML(data=data))

def ext_face_tracking_points(result, num_detect_faces):
    # `[P, T, J, C]`
    body_kpts = np.zeros([num_detect_faces, 1, 33, 4])
    lhand_kpts = np.zeros([num_detect_faces, 1, 21, 4])
    rhand_kpts = np.zeros([num_detect_faces, 1, 21, 4])
    face_kpts = np.zeros([num_detect_faces, 1, 478, 4])  # Includes iri's keypoints.

    for i, landmarks in enumerate(result.face_landmarks):
        temp = np.array([
            [lmark.x, lmark.y, lmark.z, _binary_conf(lmark)] for lmark
            in landmarks])
        face_kpts[i, 0, :temp.shape[0]] = temp
    return body_kpts, lhand_kpts, rhand_kpts, face_kpts

def _binary_conf(lmark):
    return float((lmark.x + lmark.y + lmark.z) > 0)

def ext_hand_tracking_points(result, num_detect_hands):
    # `[P, T, J, C]`
    body_kpts = np.zeros([num_detect_hands, 1, 33, 4])
    lhand_kpts = np.zeros([num_detect_hands, 1, 21, 4])
    rhand_kpts = np.zeros([num_detect_hands, 1, 21, 4])
    face_kpts = np.zeros([num_detect_hands, 1, 478, 4])  # Includes iri's keypoints.

    if len(result.handedness) > 0:
        for i, (handedness, landmarks) in enumerate(zip(result.handedness, result.hand_landmarks)):
            temp = np.array([
                    [lmark.x, lmark.y, lmark.z, _binary_conf(lmark)] for lmark
                    in landmarks])
            if handedness[0].display_name == "Left":
                lhand_kpts[i, 0, :temp.shape[0]] = temp
            else:
                rhand_kpts[i, 0, :temp.shape[0]] = temp
    return body_kpts, lhand_kpts, rhand_kpts, face_kpts

def ext_pose_tracking_points(result, num_detect_poses):
    # `[P, T, J, C]`
    body_kpts = np.zeros([num_detect_poses, 1, 33, 4])
    lhand_kpts = np.zeros([num_detect_poses, 1, 21, 4])
    rhand_kpts = np.zeros([num_detect_poses, 1, 21, 4])
    face_kpts = np.zeros([num_detect_poses, 1, 478, 4])  # Includes iri's keypoints.

    for i, landmarks in enumerate(result.pose_landmarks):
        temp = np.array([
            [lmark.x, lmark.y, lmark.z, lmark.visibility] for lmark
            in landmarks])
        body_kpts[i, 0, :temp.shape[0]] = temp
    return body_kpts, lhand_kpts, rhand_kpts, face_kpts

def project(kpts,
            width,
            height,
            inplace=True,
            limit_to_window=False):
    if inplace is True:
        temp = kpts
    else:
        temp = copy.deepcopy(kpts)
    depth = max(width, height)
    xs = temp[:, :, :, 0]
    ys = temp[:, :, :, 1]
    zs = temp[:, :, :, 2]
    xs = np.floor(xs * width)
    ys = np.floor(ys * height)
    zs = np.floor(zs * depth)
    if limit_to_window:
      xs[xs > (width - 1)] = width - 1
      ys[ys > (height - 1)] = height - 1
      zs[zs > (depth - 1)] = depth - 1
      xs[xs < 0] = 0
      ys[ys < 0] = 0
      zx[zx < 0] = 0
    temp[:, :, :, 0] = xs
    temp[:, :, :, 1] = ys
    temp[:, :, :, 2] = zs
    return temp

def tracking(tracker, videofile, static_image_mode, num_detects):
    # Load video and start tracking.
    capture = cv2.VideoCapture(videofile)
    capture.set(cv2.CAP_PROP_POS_FRAMES, 0)
    fps = capture.get(cv2.CAP_PROP_FPS)
    frame_count = 0
    frames = None
    while True:
        status, image = capture.read()
        if status is False:
            break
        try:
            # Track.
            mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
            if static_image_mode:
                result = tracker.detect(mp_image)
            else:
                timestamp_ms = int(frame_count / fps * 1000)
                result = tracker.detect_for_video(mp_image, timestamp_ms)
                frame_count += 1
            if isinstance(tracker, FaceLandmarker):
                body, lhand, rhand, face = ext_face_tracking_points(result, num_detects)
            elif isinstance(tracker, HandLandmarker):
                body, lhand, rhand, face = ext_hand_tracking_points(result, num_detects)
            elif isinstance(tracker, PoseLandmarker):
                body, lhand, rhand, face = ext_pose_tracking_points(result, num_detects)
            else:
                raise ValueError(f"Unknown tracker type: {type(tracker)}.")
            # Project to image plane.
            height, width = image.shape[:2]
            body = project(body, width, height)
            lhand = project(lhand, width, height)
            rhand = project(rhand, width, height)
            face = project(face, width, height)

            keypoints = np.concatenate([body, lhand, rhand, face], axis=2)
            if frames is None:
                frames = keypoints
            else:
                frames = np.concatenate([frames, keypoints], axis=1)
        except Exception as inst:
            raise inst
    return frames

videofile = "./finger_near0.mp4"
output = "out_track/finger_near0_face.npy"
static_image_mode = False
num_detect_faces = 1

options = FaceLandmarkerOptions(
    base_options=BaseOptions(model_asset_path="./face_landmarker.task"),
    num_faces=num_detect_faces,
    running_mode=RunningMode.IMAGE if static_image_mode else RunningMode.VIDEO)

tracker = FaceLandmarker.create_from_options(options)

frames = tracking(tracker, videofile, static_image_mode, num_detect_faces)

print(frames.shape)
parent = os.path.dirname(output)
os.makedirs(parent, exist_ok=True)
np.save(output, frames)

videofile = "./finger_near0.mp4"
output = "out_track/finger_near0_hand.npy"
static_image_mode = False
num_detect_hands = 1

options = HandLandmarkerOptions(
    base_options=BaseOptions(model_asset_path="./hand_landmarker.task"),
    num_hands=num_detect_hands,
    running_mode=RunningMode.IMAGE if static_image_mode else RunningMode.VIDEO)

tracker = HandLandmarker.create_from_options(options)

frames = tracking(tracker, videofile, static_image_mode, num_detect_hands)

print(frames.shape)
parent = os.path.dirname(output)
os.makedirs(parent, exist_ok=True)
np.save(output, frames)

videofile = "./finger_near0.mp4"
output = "out_track/finger_near0_pose.npy"
static_image_mode = False
num_detect_poses = 1

options = PoseLandmarkerOptions(
    base_options=BaseOptions(model_asset_path="./pose_landmarker_full.task"),
    num_poses=num_detect_poses,
    running_mode=RunningMode.IMAGE if static_image_mode else RunningMode.VIDEO)

tracker = PoseLandmarker.create_from_options(options)

frames = tracking(tracker, videofile, static_image_mode, num_detect_poses)

print(frames.shape)
parent = os.path.dirname(output)
os.makedirs(parent, exist_ok=True)
np.save(output, frames)

def draw_landmarks(draw, landmarks, connections,
                   pt_color=(0, 0, 0), line_color=(0, 0, 0),
                   pt_size=3, line_size=3,
                   do_remove_error=True):
    for i, line in enumerate(connections):
        p0 = landmarks[line[0], :2]
        p1 = landmarks[line[1], :2]
        if do_remove_error:
            if np.isnan(p0).any() or np.isnan(p1).any():
                continue
            if (p0==0).any() or (p1==0).any():
                continue
        p0 = (int(p0[0]), int(p0[1]))
        p1 = (int(p1[0]), int(p1[1]))
        cv2.line(draw, p0, p1, line_color, line_size)
        cv2.circle(draw, p0, pt_size, pt_color, -1)
        cv2.circle(draw, p1, pt_size, pt_color, -1)
    return draw

def drawing(videofile, trackfile, outvideo):
    trackdata = np.load(trackfile)
    if os.path.exists(videofile):
        capture = cv2.VideoCapture(videofile)
        capture.set(cv2.CAP_PROP_POS_FRAMES, 0)
        width = int(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
    else:
        capture = None
        width = 500
        height = 500

    xmin_orig = trackdata[:, :, :, 0].min()
    ymin_orig = trackdata[:, :, :, 1].min()
    xmax_orig = trackdata[:, :, :, 0].max()
    ymax_orig = trackdata[:, :, :, 1].max()
    if xmax_orig < 10 and ymax_orig < 10:
        trackdata[:, :, 0] *= width
        trackdata[:, :, 1] *= height
    xmin_proj = int(xmin_orig)
    ymin_proj = int(ymin_orig)
    xmax_proj = int(xmax_orig)
    ymax_proj = int(ymax_orig)

    # Add offset.
    offset_x = 0
    offset_y = 0
    if xmin_proj < 0:
        offset_x = -xmin_proj
        trackdata[:, :, :, 0] += offset_x
    if ymin_proj < 0:
        offset_y = -ymin_proj
        trackdata[:, :, :, 1] += offset_y

    ywin = int(max(ymax_proj + offset_y, height + offset_y))
    xwin = int(max(xmax_proj + offset_x, width + offset_x))

    print("Window size:", xwin, ywin)
    print("Offsets:", offset_x, offset_y)

    writer = cv2.VideoWriter(outvideo, cv2.VideoWriter_fourcc(*"mp4v"),
        30.0, (xwin, ywin))
    if writer.isOpened() is False:
        print("VideoWriter is failed to open.")
        if writer is not None:
            writer.release()
        raise ValueError(f"Can not open {videofile}")

    # `[P, T, J, C] -> [T, P, J, C]`
    trackdata = np.transpose(trackdata, [1, 0, 2, 3])
    for frame in trackdata:
        if capture is not None:
            status, image = capture.read()
            if status is False:
                break
        else:
            image = np.full([height, width], 255, dtype=np.uint8)

        # Draw.
        draw = np.full([int(ywin), int(xwin), 3], 64, dtype=np.uint8)
        draw[offset_y: offset_y+height, offset_x: offset_x+width, :] = image

        for instance in frame:
            body = instance[:33, :]
            lhand = instance[33:33+21, :]
            rhand = instance[33+21:33+21+21, :]
            face = instance[33+21+21:, :]
            draw_landmarks(draw, body, POSE_CONNECTIONS, [0, 255, 0], [0, 255, 0],
                           do_remove_error=True)
            draw_landmarks(draw, lhand, HAND_CONNECTIONS, [255, 0, 0], [255, 0, 0],
                           do_remove_error=True)
            draw_landmarks(draw, rhand, HAND_CONNECTIONS, [0, 0, 255], [0, 0, 255],
                           do_remove_error=True)
            draw_landmarks(draw, face, FACEMESH_CONTOURS, [255, 255, 0], [255, 255, 0],
                           do_remove_error=True)
        writer.write(draw)
    writer.release()
    if capture is not None:
        capture.release()

videofile = "./finger_near0.mp4"
trackfile = "out_track/finger_near0_face.npy"
outvideo = "./finger_near0_face_track.mp4"

drawing(videofile, trackfile, outvideo)

videofile = "./finger_near0.mp4"
trackfile = "out_track/finger_near0_hand.npy"
outvideo = "./finger_near0_hand_track.mp4"

drawing(videofile, trackfile, outvideo)

videofile = "./finger_near0.mp4"
trackfile = "out_track/finger_near0_pose.npy"
outvideo = "./finger_near0_pose_track.mp4"

drawing(videofile, trackfile, outvideo)