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

# CV/ML.
import numpy as np

import jax
import jax.numpy as jnp
from jax import jit

# Enable float64.
jax.config.update("jax_enable_x64", True)

print(f"Python:{sys.version}")
print(f"Numpy:{np.__version__}")
print(f"JAX:{jax.__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 simple_interp_jax(trackdata):
    tlength, num_joints, _ = trackdata.shape
    newtrack = jnp.zeros_like(trackdata)
    for i in range(num_joints):
        temp = trackdata[:, i, :]
        mask = temp[:, -1] != 0
        valid = mask.sum()
        if valid == tlength:
            newtrack = newtrack.at[:, i].set(temp)
            continue
        xs = jnp.where(mask != 0, size=valid)[0]
        # ys = temp[xs, :] <- can't be compiled.
        ys = jnp.take(temp, xs, axis=0)
        newys = jnp.zeros_like(temp)
        for j in range(temp.shape[-1]):
            newy = jnp.interp(jnp.arange(tlength), xs, ys[:, j])
            newys = newys.at[:, j].set(newy)
        newtrack = newtrack.at[:, i].set(newys)
    return newtrack

def matrix_interp_jax(track):
    orig_shape = track.shape
    tlength = orig_shape[0]
    mask = track[:, 0, -1] != 0
    valid = mask.sum()
    if valid == tlength:
        return track

    xs = jnp.where(mask != 0, size=valid)[0]
    # ys = track.reshape([tlength, -1])[xs, :] <- can't be compiled
    ys = jnp.take(track.reshape([tlength, -1]), xs, axis=0)
    x = jnp.arange(tlength)

    # ========================================================================
    # Interpolation.
    # ========================================================================
    xs = xs.astype(ys.dtype)
    x = x.astype(ys.dtype)
    # Pad control points for extrapolation.
    xs = jnp.concatenate([jnp.array([jnp.finfo(xs.dtype).min]), xs, jnp.array([jnp.finfo(xs.dtype).max])], axis=0)
    ys = jnp.concatenate([ys[:1], ys, ys[-1:]], axis=0)

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

    # Solve for intercepts.
    intercepts = ys - sloops * jnp.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 = jnp.argmax(jnp.expand_dims(xs, axis=-2) > jnp.expand_dims(x, axis=-1), axis=-1)
    sloop = sloops[idx]
    intercept = intercepts[idx]

    # Apply the linear mapping at each input data point.
    y = sloop * jnp.expand_dims(x, axis=-1) + intercept
    y = y.astype(ys.dtype)
    y = y.reshape(orig_shape)
    return y

def partsbased_interp_jax(trackdata):
    pose = trackdata[:, :33]
    lhand = trackdata[:, 33:33+21]
    rhand = trackdata[:, 33+21:33+21+21]
    face = trackdata[:, 33+21+21:]

    pose = matrix_interp_jax(pose)
    lhand = matrix_interp_jax(lhand)
    rhand = matrix_interp_jax(rhand)
    face = matrix_interp_jax(face)
    return jnp.concatenate([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]

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

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

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

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

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

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

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

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

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

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

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

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

from typing import Generic, TypeVar
from functools import partial

T = TypeVar('T')      # Declare type variable

# Workaround to avoid unhashable error.
# https://github.com/google/jax/issues/4572
class HashableArrayWrapper(Generic[T]):
    def __init__(self, val: T):
        self.val = val

    def __getattribute__(self, prop):
        if prop == 'val' or prop == "__hash__" or prop == "__eq__":
            return super(HashableArrayWrapper, self).__getattribute__(prop)
        return getattr(self.val, prop)

    def __getitem__(self, key):
        return self.val[key]

    def __setitem__(self, key, val):
        self.val[key] = val

    def __hash__(self):
        return hash(self.val.tobytes())

    def __eq__(self, other):
        if isinstance(other, HashableArrayWrapper):
            return self.__hash__() == other.__hash__()

        f = getattr(self.val, "__eq__")
        return f(self, other)

@partial(jit, static_argnums=(0,))
def simple_interp_jax_jit(trackdata):
    ...

def matrix_interp_jax_jit(track):
    ...

@partial(jit, static_argnums=(0,))
def partsbased_interp_jax_jit(trackdata):
    ...

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

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

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

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

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

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

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

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

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

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

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

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