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

# CV/ML.
import numpy as np

import tensorflow as tf

print(f"Python:{sys.version}")
print(f"Numpy:{np.__version__}")
print(f"Tensorflow:{tf.__version__}")

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 matrix_interp_tf_eager(track):
    orig_shape = tf.shape(track)
    tlength = orig_shape[0]
    mask = track[:, 0, -1] != 0
    valid = tf.reduce_sum(tf.cast(mask, dtype=tf.int32))
    if valid == tlength:
        y = track
    else:
        xs = tf.where(mask)
        xs = tf.reshape(xs, [valid])
        # determine the output data type
        ys = tf.reshape(track, [tlength, -1])
        ys = tf.gather(ys, xs, axis=0)
        x = tf.range(tlength)
        dtype_ys = ys.dtype

        # normalize data types
        xs = tf.cast(xs, dtype_ys)
        x = tf.cast(x, dtype_ys)

        # pad control points for extrapolation
        xs = tf.concat([[xs.dtype.min], xs, [xs.dtype.max]], axis=0)
        ys = tf.concat([ys[:1], ys, ys[-1:]], axis=0)

        # compute slopes, pad at the edges to flatten
        sloops = (ys[1:] - ys[:-1]) / tf.expand_dims((xs[1:] - xs[:-1]), axis=-1)
        sloops = tf.pad(sloops[:-1], [(1, 1), (0, 0)])

        # solve for intercepts
        intercepts = ys - sloops * tf.expand_dims(xs, axis=-1)

        # search for the line parameters at each input data point
        # create a grid of the inputs and piece breakpoints for thresholding
        # rely on argmax stopping on the first true when there are duplicates,
        # which gives us an index into the parameter vectors
        idx = tf.math.argmax(tf.expand_dims(xs, axis=-2) > tf.expand_dims(x, axis=-1), axis=-1)
        sloop = tf.gather(sloops, idx, axis=0)
        intercept = tf.gather(intercepts, idx, axis=0)

        # apply the linear mapping at each input data point
        y = sloop * tf.expand_dims(x, axis=-1) + intercept
        y = tf.cast(y, dtype_ys)
        y = tf.reshape(y, orig_shape)
    return y

def partsbased_interp_tf_eager(trackdata):
    num_joints = trackdata.shape[1]
    trackdata = tf.convert_to_tensor(trackdata)
    pose = tf.gather(trackdata, tf.range(0, 33), axis=1)
    lhand = tf.gather(trackdata, tf.range(33, 33+21), axis=1)
    rhand = tf.gather(trackdata, tf.range(33+21, 33+21+21), axis=1)
    face = tf.gather(trackdata, tf.range(33+21+21, num_joints), axis=1)

    pose = matrix_interp_tf_eager(pose)
    lhand = matrix_interp_tf_eager(lhand)
    rhand = matrix_interp_tf_eager(rhand)
    face = matrix_interp_tf_eager(face)
    return tf.concat([pose, lhand, rhand, face], axis=1)

trackdata = np.load("finger_far0_non_static.npy")
reftrack = np.load("finger_far0_non_static_interp.npy")
# Remove person axis.
trackdata = trackdata[0]
reftrack = reftrack[0]

# Tensorflow.
# The 1st call may be slow because of the computation graph construction.
print(f"Time of first call.")
start = time.perf_counter()
newtrack = partsbased_interp_tf_eager(trackdata)
interval = time.perf_counter() - start
print_perf_time(np.array([interval]))

diff = (reftrack - newtrack.numpy()).sum()
print(f"Sum of error:{diff}")

print("Time after second call.")
target_fn = partial(partsbased_interp_tf_eager, trackdata=trackdata)
pmeasure(target_fn)

# Tensorflow.
# The 1st call may be slow because of the computation graph construction.
print(f"Time of first call.")
start = time.perf_counter()
newtrack = partsbased_interp_tf_eager(trackdata[:-1])
interval = time.perf_counter() - start
print_perf_time(np.array([interval]))

diff = (reftrack[:-1] - newtrack.numpy()).sum()
print(f"Sum of error:{diff}")

print("Time after second call.")
target_fn = partial(partsbased_interp_tf_eager, trackdata=trackdata[:-1])
pmeasure(target_fn)

# If input_signature is omitted, the re-tracing is performed when a tensor's shape is changed.
@tf.function
def matrix_interp_tf(track):
    ...

# Tensorflow.
# The 1st call may be slow because of the computation graph construction.
print(f"Time of first call.")
start = time.perf_counter()
newtrack = partsbased_interp_tf(trackdata)
interval = time.perf_counter() - start
print_perf_time(np.array([interval]))

diff = (reftrack - newtrack.numpy()).sum()
print(f"Sum of error:{diff}")

print("Time after second call.")
target_fn = partial(partsbased_interp_tf, trackdata=trackdata)
pmeasure(target_fn)

# Tensorflow.
# The 1st call may be slow because of the computation graph construction.
print(f"Time of first call.")
start = time.perf_counter()
newtrack = partsbased_interp_tf(trackdata[:-1])
interval = time.perf_counter() - start
print_perf_time(np.array([interval]))

diff = (reftrack[:-1] - newtrack.numpy()).sum()
print(f"Sum of error:{diff}")

print("Time after second call.")
target_fn = partial(partsbased_interp_tf, trackdata=trackdata[:-1])
pmeasure(target_fn)

# If input_signature is omitted, the re-tracing is performed when a tensor's shape is changed.
@tf.function(input_signature=(tf.TensorSpec(shape=[None, None, 4], dtype=tf.float64),))
def matrix_interp_tf(track):
    ...

# Tensorflow.
# The 1st call may be slow because of the computation graph construction.
print(f"Time of first call.")
start = time.perf_counter()
newtrack = partsbased_interp_tf(trackdata)
interval = time.perf_counter() - start
print_perf_time(np.array([interval]))

diff = (reftrack - newtrack.numpy()).sum()
print(f"Sum of error:{diff}")

print("Time after second call.")
target_fn = partial(partsbased_interp_tf, trackdata=trackdata)
pmeasure(target_fn)

# Tensorflow.
# The 1st call may be slow because of the computation graph construction.
print(f"Time of first call.")
start = time.perf_counter()
newtrack = partsbased_interp_tf(trackdata[:-1])
interval = time.perf_counter() - start
print_perf_time(np.array([interval]))

diff = (reftrack[:-1] - newtrack.numpy()).sum()
print(f"Sum of error:{diff}")

print("Time after second call.")
target_fn = partial(partsbased_interp_tf, trackdata=trackdata[:-1])
pmeasure(target_fn)