# Standard modules.
import gc
import sys
import time
from functools import partial
from pathlib import Path

# CV/ML.
import cv2

import numpy as np

# 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

print(f"Python:{sys.version}")
print(f"OpenCV:{cv2.__version__}")
print(f"Numpy:{np.__version__}")

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

trackfile = Path("./finger_far0_non_static.npy")
outvideo = Path("./finger_far0_track_orig.mp4")

drawing(videopath_orig, trackfile, outvideo)

def get_perf_str(val):
    token_si = ["", "m", "µ", "n", "p"]
    exp_si = [1, 1e3, 1e6, 1e9, 1e12]
    perf_str = f"{val:3g}s"
    si = ""
    sval = val
    for token, exp in zip(token_si, exp_si):
        if val * exp > 1.0:
            si = token
            sval = val * exp
            break
    perf_str = f"{sval:3g}{si}s"
    return perf_str

def print_perf_time(intervals, top_k=None):
    if top_k is not None:
        intervals = np.sort(intervals)[:top_k]
    min = intervals.min()
    max = intervals.max()
    mean = intervals.mean()
    std = intervals.std()

    smin = get_perf_str(min)
    smax = get_perf_str(max)
    mean = get_perf_str(mean)
    std = get_perf_str(std)
    if top_k:
        print(f"Top {top_k} summary: Max {smax}, Min {smin}, Mean +/- Std {mean} +/- {std}")
    else:
        print(f"Overall summary: Max {smax}, Min {smin}, Mean +/- Std {mean} +/- {std}")

class PerfMeasure():
    def __init__(self,
                 trials=100,
                 top_k=10):
        self.trials = trials
        self.top_k = top_k

    def __call__(self, func):
        gc.collect()
        gc.disable()
        intervals = []
        for _ in range(self.trials):
            start = time.perf_counter()
            func()
            end = time.perf_counter()
            intervals.append(end - start)
        intervals = np.array(intervals)
        print_perf_time(intervals)
        if self.top_k:
            print_perf_time(intervals, self.top_k)
        gc.enable()
        gc.collect()

TRIALS = 100
TOPK = 10
pmeasure = PerfMeasure(TRIALS, TOPK)

def simple_interp(trackdata):
    tlength, num_joints, _ = trackdata.shape
    newtrack = np.zeros_like(trackdata)
    for i in range(num_joints):
        temp = trackdata[:, i, :]
        mask = temp[:, -1] != 0
        valid = mask.sum()
        if valid == tlength:
            newtrack[:, i] = temp
            continue
        xs = np.where(mask != 0)[0]
        ys = temp[xs, :]
        newys = np.zeros_like(temp)
        for j in range(temp.shape[-1]):
            newy = np.interp(np.arange(tlength), xs, ys[:, j])
            newys[:, j] = newy
        newtrack[:, i] = newys
    return newtrack

# Load track data.
trackdata = np.load(trackfile)
print(trackdata.shape)

# Remove person axis.
trackdata = trackdata.squeeze()
print(trackdata.shape)

# Print joints including missing frames.
tlength, num_joints, num_featurs = trackdata.shape
for i in range(num_joints):
    temp = trackdata[:, i, :]
    mask = temp[:, -1] != 0
    valid = mask.sum()
    if valid != tlength:
        print(i, valid)

# Apply interpolation.
newtrack = simple_interp(trackdata)

# Add person axis.
newtrack = np.expand_dims(newtrack, 0)

outpath_track = trackfile.stem + "_interp.npy"
np.save(outpath_track, newtrack)

# Draw result.
outvideo = Path("./finger_far0_track_interp.mp4")

drawing(videopath_orig, outpath_track, outvideo)

show_video(outvideo.stem + ".webm")

target_fn = partial(simple_interp, trackdata=trackdata)
pmeasure(target_fn)