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

# 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 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):
    """Draw landmarks over input video.

    # Args:
      - videofile: The path of input video. This is used to draw background image.
      - trackfile: The path of tracking data.
      - outvideo: The path of output video file.
    """
    trackdata = np.load(trackfile)
    if videofile.is_file():
        capture = cv2.VideoCapture(str(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 = 1200
        height = 1200

    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(str(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, 3], 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()

def show_video(video_path, video_width=500, video_height=500):
    video = io.open(str(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))

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

drawing(Path(""), 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 get_affine_matrix_2d(center,
                         trans,
                         scale,
                         rot,
                         skew,
                         to_radians=True,
                         order=["center", "scale", "rot", "skew", "trans"],
                         dtype=np.float32):
    center = np.array(center)
    trans = np.array(trans)
    scale = np.array(scale)
    center_m = np.array([[1, 0, -center[0]],
                         [0, 1, -center[1]],
                         [0, 0, 1]])
    scale_m = np.array([[scale[0], 0, 0],
                        [0, scale[1], 0],
                        [0, 0, 1]])
    rot = np.radians(rot) if to_radians else rot
    _cos = np.cos(rot)
    _sin = np.sin(rot)
    rot_m = np.array([[_cos, -_sin, 0],
                      [_sin, _cos, 0],
                      [0, 0, 1]])
    _tan = np.tan(np.radians(skew)) if to_radians else np.tan(skew)
    skew_m = np.array([[1, _tan[0], 0],
                       [_tan[1], 1, 0],
                       [0, 0, 1]])
    move = center + trans
    trans_m = np.array([[1, 0, move[0]],
                        [0, 1, move[1]],
                        [0, 0, 1]])
    mat = np.identity(3, dtype=dtype)
    for name in order:
        if name == "center":
            mat = np.matmul(center_m, mat)
        if name == "scale":
            mat = np.matmul(scale_m, mat)
        if name == "rot":
            mat = np.matmul(rot_m, mat)
        if name == "skew":
            mat = np.matmul(skew_m, mat)
        if name == "trans":
            mat = np.matmul(trans_m, mat)
    return mat.astype(dtype)

def apply_affine(inputs, mat):
    xy = inputs[:, :, :2]
    xy = np.concatenate([xy, np.ones([xy.shape[0], xy.shape[1], 1])], axis=-1)
    xy = np.einsum("...j,ij", xy, mat)
    inputs[:, :, :2] = xy[:, :, :-1]
    return inputs

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

# Remove person axis.
trackdata = trackdata[0]

# Get affine matrix.
center = (638.0, 389.0)
trans = (100.0, 0.0)
scale = (2.0, 0.5)
rot = 15.0
skew = (15.0, 15.0)
print("Parameters")
print("Center:", center)
print("Trans:", trans)
print("Scale:", scale)
print("Rot:", rot)
print("Skew:", skew)
mat = get_affine_matrix_2d(center, trans, scale, rot, skew, to_radians=True, dtype=trackdata.dtype)
print("Affine matrix:", mat)

newtrack = apply_affine(trackdata.copy(), mat)

# Add person axis.
newtrack = np.expand_dims(newtrack, axis=0)
print("Shape of newtrack:", newtrack.shape)

# Save.
np.save("newtrack.npy", newtrack)

trackfile = Path("./newtrack.npy")
outvideo = Path("./newtrack.mp4")

drawing(Path(""), trackfile, outvideo)

def perf_wrap_func(trackdata, center, trans, scale, rot, skew, to_radians, dtype):
    mat = get_affine_matrix_2d(center, trans, scale, rot, skew, to_radians=to_radians, dtype=dtype)
    newtrack = apply_affine(trackdata, mat)

testtrack = trackdata.copy()
target_fn = partial(perf_wrap_func,
                    trackdata=testtrack,
                    center=center, trans=trans, scale=scale, rot=rot, skew=skew,
                    to_radians=True, dtype=testtrack.dtype)
pmeasure(target_fn)

class RandomAffineTransform2D():
    def __init__(self,
                 center_joints,
                 apply_ratio,
                 trans_range,
                 scale_range,
                 rot_range,
                 skew_range,
                 random_seed=None,
                 order=["center", "scale", "rot", "skew", "trans"],
                 dtype=np.float32):
        self.center_joints = center_joints
        self.apply_ratio = apply_ratio
        self.trans_range = trans_range
        self.scale_range = scale_range
        self.rot_range = np.radians(rot_range)
        self.skew_range = np.radians(skew_range)
        self.order = order
        self.dtype = dtype
        if random_seed is not None:
            self.rng = np.random.default_rng(random_seed)
        else:
            self.rng = np.random.default_rng()

    def __call__(self, inputs):
        if self.rng.uniform() >= self.apply_ratio:
            return inputs

        # Calculate center position.
        temp = inputs[:, self.center_joints, :]
        temp = temp.reshape([inputs.shape[0], -1, inputs.shape[-1]])
        mask = np.sum(temp, axis=(1, 2)) != 0
        # Use x and y only.
        center = temp[mask].mean(axis=0).mean(axis=0)[:2]

        trans = self.rng.uniform(self.trans_range[0], self.trans_range[1], 2)
        scale = self.rng.uniform(self.scale_range[0], self.scale_range[1], 2)
        rot = self.rng.uniform(self.rot_range[0], self.rot_range[1])
        skew = self.rng.uniform(self.skew_range[0], self.skew_range[1], 2)

        # Calculate matrix.
        mat = get_affine_matrix_2d(center, trans, scale, rot, skew,
            to_radians=False, order=self.order, dtype=self.dtype)

        # Apply transform.
        inputs = apply_affine(inputs, mat)
        return inputs

aug_fn = RandomAffineTransform2D(
    center_joints=[11, 12],
    apply_ratio=1.0,
    trans_range=[-100.0, 100.0],
    scale_range=[0.5, 2.0],
    rot_range=[-30.0, 30.0],
    skew_range=[-30.0, 30.0])

augtracks = []
for _ in range(3):
    augtracks.append(aug_fn(trackdata.copy()))

# Save trakcs and videos.
for i in range(3):
  track = augtracks[i]
  track = np.expand_dims(track, axis=0)
  outtrack = Path(f"./newtrack_{i}.npy")
  outvideo = Path(f"./newtrack_{i}.mp4")
  np.save(outtrack, track)
  drawing(Path(""), outtrack, outvideo)

testtrack = trackdata.copy()
target_fn = partial(aug_fn, inputs=testtrack)
pmeasure(target_fn)