# 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_holistic_tracking_points(result):
    # `[P, T, J, C]`
    body_kpts = np.zeros([1, 1, 33, 4])
    lhand_kpts = np.zeros([1, 1, 21, 4])
    rhand_kpts = np.zeros([1, 1, 21, 4])
    face_kpts = np.zeros([1, 1, 478, 4])  # Includes iri's keypoints.
    if result.pose_landmarks is not None:
        temp = np.array([
            [lmark.x, lmark.y, lmark.z, lmark.visibility] for lmark
            in result.pose_landmarks.landmark])
        body_kpts[0, :temp.shape[0]] = temp

    # visibility of left/right hands and face return dummy value (0.0).
    # So, we assign binary visibilities judged by coordinates values.
    def _binary_conf(lmark):
        return float((lmark.x + lmark.y + lmark.z) > 0)

    if result.left_hand_landmarks is not None:
        temp = np.array([
            [lmark.x, lmark.y, lmark.z, _binary_conf(lmark)] for lmark
            in result.left_hand_landmarks.landmark])
        lhand_kpts[0, 0, :temp.shape[0]] = temp
    if result.right_hand_landmarks is not None:
        temp = np.array([
            [lmark.x, lmark.y, lmark.z, _binary_conf(lmark)] for lmark
            in result.right_hand_landmarks.landmark])
        rhand_kpts[0, 0, :temp.shape[0]] = temp
    if result.face_landmarks is not None:
        temp = np.array([
            [lmark.x, lmark.y, lmark.z, _binary_conf(lmark)] for lmark
            in result.face_landmarks.landmark])
        face_kpts[0, 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

# Instantiate Holistic model.
params = {
    "static_image_mode": False,
    "model_complexity": 1,
    "smooth_landmarks": True,
    "enable_segmentation": False,
    "smooth_segmentation": False,
    "refine_face_landmarks": True,
    "min_detection_confidence": 0.5,
    "min_tracking_confidence": 0.5}
tracker = holistic.Holistic(**params)

videofile = "./finger_far0.mp4"
output = "out_track/finger_far0.npy"

# Load video and start tracking.
capture = cv2.VideoCapture(videofile)
capture.set(cv2.CAP_PROP_POS_FRAMES, 0)
frames = None
while True:
    status, image = capture.read()
    if status is False:
        break
    try:
        # Track.
        result = tracker.process(image)
        body, lhand, rhand, face = ext_holistic_tracking_points(result)
        # 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

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

trackfile = output
outvideo = "finger_far0_track.mp4"

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()