import os
import sys
import math
import argparse
from pathlib import Path
from multiprocessing.pool import ThreadPool
import subprocess

# Parse arguments before importing libraries
def parse_args(args=None):
    parser = argparse.ArgumentParser(argument_default=None)
    parser.add_argument('--data', help='CSV animation data files', nargs='*',
                        metavar='[NAME:]PATH', type=named_path)
    parser.add_argument('--models', help='model files (default mode: CUBIC)', nargs='*',
                        metavar='[NAME:]PATH[?MODE]', type=named_path_mode('CUBIC'))
    parser.add_argument('--out-dir', help='output results directory', type=Path)
    parser.add_argument('--per-data', help='write results for each data file', type=str2bool, default='true')
    parser.add_argument('--anim-dir', help='output animation directory', type=Path)
    parser.add_argument('--anim-layout', help='animation layout', metavar='ROWS,COLS', type=layout_shape)
    parser.add_argument('--anim-order', help='direction in which animations are laid out (default: %(default)s)',
                        choices=['HORIZONTAL', 'VERTICAL'], default='HORIZONTAL', type=lambda s: s.strip().upper())
    parser.add_argument('--anim-dpi', help='animation DPI (default: %(default)d)', metavar='DPI', type=int, default=100)
    parser.add_argument('--anim-error', help='whether to display error per bone or per frame (default: %(default)s)',
                        choices=['BONE', 'FRAME'], default='BONE', type=lambda s: s.strip().upper())
    parsed_args = parser.parse_args(args=args)
    data = parsed_args.data or []
    if len(set(name for name, _ in data)) != len(data):
        raise ValueError('data file names must be unique.')
    models = parsed_args.models or []
    if len(set(name for name, _, _ in models)) != len(models):
        raise ValueError('model names must be unique.')
    out_dir = parsed_args.out_dir
    per_data = parsed_args.per_data
    anim_dir = parsed_args.anim_dir
    anim_layout = parsed_args.anim_layout
    if anim_dir and anim_layout and anim_layout[0] * anim_layout[1] < 1 + len(models):
        raise ValueError('animation layout is too small.')
    anim_layout_vert = (parsed_args.anim_order == 'VERTICAL')
    anim_dpi = parsed_args.anim_dpi
    error_bone = (parsed_args.anim_error == 'BONE')
    return data, models, out_dir, per_data, anim_dir, anim_layout, anim_layout_vert, anim_dpi, error_bone

def named_path(v):
    try:
        name = ''
        path = v
        if ':' in v: name, path = v.split(':', 1)
        if not path: raise ValueError
        path = Path(path)
        if not name: name = path.stem
        if not name: raise ValueError
        return name, Path(path)
    except BaseException as e:
        raise argparse.ArgumentTypeError('value must have the format [NAME:]PATH.')

def named_path_mode(default_mode):
    def f(v):
        try:
            mode = default_mode
            name_path = v
            if '?' in v: name_path, mode = v.split('?', 1)
            if not name_path: raise ValueError
            name, path = named_path(name_path)
            return name, path, mode
        except BaseException as e:
            raise argparse.ArgumentTypeError('value must have the format [NAME:]PATH[?MODE].')
    return f

def layout_shape(v):
    try:
        rows, cols = map(int, v.split(','))
        if rows <= 0 or cols <= 0: raise ValueError
        return rows, cols
    except BaseException as e:
        raise argparse.ArgumentTypeError('value must have the format ROWS,COLS.')

def str2bool(v):
    v = str(v).strip().lower()
    if v in {'yes', 'true', 't', 'y', '1'}:
        return True
    elif v in {'no', 'false', 'f', 'n', '0'}:
        return False
    else:
        raise argparse.ArgumentTypeError('boolean value expected.')

if __name__ == '__main__':
    ARGS = parse_args()

import numpy as np
import pandas as pd
import quaternion
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf

import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as mplanim
import seaborn as sns

from seg2seg import SegmentsClosestPoints


ROOT = 'root'
HEAD = 'head'
NPC_JOINTS = [
    'COG',
    'spine_00',
    'spine_01',
    'spine_02',
    'spine_03',
    'sternum',
    'neck',
    'head',
    'l_shoulder',
    'l_arm',
    'l_elbow',
    'l_wrist',
    'r_shoulder',
    'r_arm',
    'r_elbow',
    'r_wrist',
    'l_leg',
    'l_knee',
    'l_ankle',
    'l_ball',
    'r_leg',
    'r_knee',
    'r_ankle',
    'r_ball',
    'r_weapon_01',
]
NPC_JOINT_PARENTS = {
    'root': None,
    'COG': 'root',
    'spine_00': 'COG',
    'spine_01': 'spine_00',
    'spine_02': 'spine_01',
    'spine_03': 'spine_02',
    'sternum': 'spine_03',
    'neck': 'sternum',
    'head': 'neck',
    'l_shoulder': 'sternum',
    'l_arm': 'l_shoulder',
    'l_elbow': 'l_arm',
    'l_wrist': 'l_elbow',
    'l_weapon_01': 'l_wrist',
    'r_shoulder': 'sternum',
    'r_arm': 'r_shoulder',
    'r_elbow': 'r_arm',
    'r_wrist': 'r_elbow',
    'r_weapon_01': 'r_wrist',
    'l_leg': 'COG',
    'l_knee': 'l_leg',
    'l_ankle': 'l_knee',
    'l_ball': 'l_ankle',
    'r_leg': 'COG',
    'r_knee': 'r_leg',
    'r_ankle': 'r_knee',
    'r_ball': 'r_ankle',
}
NPC_SWORD_TIP = 'r_weapon_01'
NPC_SWORD_VECTOR = [90., 0., 0.]
NPC_SWORD_OFFSET = [-20., 0., 0.]
PLAYER_SWORD_TIP = 'l_weapon_01'
PLAYER_SWORD_VECTOR = [-90., 0., 0.]
PLAYER_SWORD_OFFSET = [0., 0., 0.]
COG = 'COG'
STYLES = [
    'stand',
    'block',
]
TRAJECTORY_WINDOW = 3
TRAJECTORY_STEP_SIZE = 3
COORD_DEPTH_RANGE = 40, 110
COORD_WIDTH = 120
COORD_HEIGHT = 200
JOINT_SMOOTH = 0.5
ERROR_BONE_SATURATION = 1_000
ERROR_FRAME_SATURATION = 10_000

TIME_COLUMN = ':time'
STANDING_COLUMN = ':notify:PfnnStanding'
STYLE_COLUMN = ':notify:PfnnStyle:Name'
SWORD_BLOCK = ':notify:SwordBlock'


# Transforms are multiplied FROM THE LEFT: Transform * Object
# Rightmost transform is applied first: Transform3 * Transform2 * Transform1
def make_transform_relative(tx, root):
    tx_rel = np.linalg.inv(root) @ tx
    return normalize_transform(tx_rel)

def normalize_transform(tx):
    tx = np.asarray(tx)
    return tx / tx[..., 3:, 3:]

def rotate_vector(v, tx):
    tx = normalize_transform(tx)
    v = np.asarray(v)
    v_rot = tx[..., :3, :3] @ v[..., np.newaxis]
    return v_rot[..., 0]

def transform_vector(v, tx):
    v = np.asarray(v)
    v = np.concatenate([v, np.ones_like(v[..., :1])], axis=-1)
    v_rot = tx @ v[..., np.newaxis]
    return v_rot[..., :3, 0] / v_rot[..., 3:, 0]

def triangles_area(triangles):
    triangles = np.asarray(triangles)
    ab = triangles[..., 1, :] - triangles[..., 0, :]
    ac = triangles[..., 2, :] - triangles[..., 0, :]
    return np.linalg.norm(np.cross(ab, ac), axis=-1) / 2

def segments_error(segments1, segments2):
    segments1 = np.asarray(segments1)
    segments2 = np.asarray(segments2)
    triangles = np.empty(segments1.shape[:-2] + (4, 3, 3), dtype=segments1.dtype)
    mid_points = np.mean(np.concatenate([segments1, segments2], axis=-2), axis=-2)
    triangles[..., 0, 0, :] = segments1[..., 0, :]
    triangles[..., 0, 1, :] = segments1[..., 1, :]
    triangles[..., 0, 2, :] = mid_points
    triangles[..., 1, 0, :] = segments1[..., 0, :]
    triangles[..., 1, 1, :] = segments2[..., 0, :]
    triangles[..., 1, 2, :] = mid_points
    triangles[..., 2, 0, :] = segments2[..., 0, :]
    triangles[..., 2, 1, :] = segments2[..., 1, :]
    triangles[..., 2, 2, :] = mid_points
    triangles[..., 3, 0, :] = segments1[..., 1, :]
    triangles[..., 3, 1, :] = segments2[..., 1, :]
    triangles[..., 3, 2, :] = mid_points
    areas = triangles_area(triangles)
    return areas.sum(axis=-1)

def make_joint_transform(pos, x, y):
    pos = np.asarray(pos)
    s = np.shape(x)
    d = s[-1]
    xs = np.reshape(x, (-1, d))
    ys = np.reshape(y, (-1, d))
    xs /= np.maximum(np.linalg.norm(xs, axis=-1, keepdims=True), 1e-8)
    ys /= np.maximum(np.linalg.norm(ys, axis=-1, keepdims=True), 1e-8)
    zs = np.cross(xs, ys, axis=-1)
    zs /= np.maximum(np.linalg.norm(zs, axis=-1, keepdims=True), 1e-8)
    ys = np.cross(zs, xs, axis=-1)
    x = np.reshape(xs, s)
    y = np.reshape(ys, s)
    z = np.reshape(zs, s)
    tx_h = np.stack([x, y, z, pos], axis=-1)
    h = np.tile([[0, 0, 0, 1]], s[:-1] + (1, 1))
    tx = np.concatenate([tx_h, h], axis=-2)
    #assert np.allclose(np.linalg.det(tx[..., :3, :3]), 1), f'Invalid rotation matrix\n{tx[..., :3, :3]}'
    return tx

def get_transforms(df, joints, dtype=None):
    dtype = dtype or np.float32
    if not joints:
        joints = [col.split(':')[0] for col in df if col.endswith(':cs:position:x')]
    unpack = False
    if np.isscalar(joints):
        joints = [joints]
        unpack = True
    transforms = []
    for joint in joints:
        pos = df[[f'{joint}:cs:position:x',
                  f'{joint}:cs:position:y',
                  f'{joint}:cs:position:z']].values
        quat = df[[f'{joint}:cs:quaternion:w',
                   f'{joint}:cs:quaternion:x',
                   f'{joint}:cs:quaternion:y',
                   f'{joint}:cs:quaternion:z']].values
        quat = quaternion.as_quat_array(quat)
        rot = quaternion.as_rotation_matrix(quat)
        tx = np.concatenate([rot, pos[..., np.newaxis]], axis=2)
        tx = np.concatenate([tx, np.tile([0, 0, 0, 1], (len(tx), 1, 1))], axis=1)
        transforms.append(tx)
    return transforms[0] if unpack else transforms

def fix_maya_rotation(tx):
    tx = normalize_transform(tx)
    # Axes given in columns
    axes_ue = np.array([[ 1,  0,  0],
                        [ 0,  1,  0],
                        [ 0,  0,  1]], dtype=tx.dtype)
    axes_maya = np.array([[ 0,  0,  1],
                          [-1,  0,  0],
                          [ 0, -1,  0]], dtype=tx.dtype)
    tx[..., :3, :3] = tx[..., :3, :3] @ axes_ue @ np.linalg.inv(axes_maya)
    return tx

def get_root_transforms(df):
    return fix_maya_rotation(get_transforms(df, ROOT))

def get_npc_joint_transforms(df):
    npc_tx = get_transforms(df, NPC_JOINTS)
    if PLAYER_SWORD_TIP in NPC_JOINTS:
        idx = NPC_JOINTS.index(PLAYER_SWORD_TIP)
        npc_tx[idx][..., :3, 3] += rotate_vector(PLAYER_SWORD_OFFSET, npc_tx[idx])
    if NPC_SWORD_TIP in NPC_JOINTS:
        idx = NPC_JOINTS.index(NPC_SWORD_TIP)
        npc_tx[idx][..., :3, 3] += rotate_vector(NPC_SWORD_OFFSET, npc_tx[idx])
    return np.stack(npc_tx, axis=1)

def get_sword_transforms(df):
    player_tx = get_transforms(df, PLAYER_SWORD_TIP)
    npc_tx = get_transforms(df, NPC_SWORD_TIP)
    player_tx[..., :3, 3] += rotate_vector(PLAYER_SWORD_OFFSET, player_tx)
    npc_tx[..., :3, 3] += rotate_vector(NPC_SWORD_OFFSET, npc_tx)
    return player_tx, npc_tx

def get_styles(df, dtype=None):
    dtype = dtype or np.float32
    style_data = np.zeros((len(df), len(STYLES)), dtype)
    if 'stand' in STYLES and STANDING_COLUMN in df:
        style_data[:, STYLES.index('stand')] = df[STANDING_COLUMN].values
    if STYLE_COLUMN not in df:
        return style_data
    style_values = df[STYLE_COLUMN].fillna('').values
    for i, style in enumerate(style_values):
        if not style: continue
        style_tags = style.split(':')
        style_data[i, [STYLES.index(tag) for tag in style_tags]] = 1
    # Overlapping areas
    overlap = np.sum(style_data, axis=1)
    if np.all(overlap <= 1):
        # No overlapping
        return style_data
    if np.any(overlap > 2):
        raise ValueError('no more than two styles can overlap.')
    # Find overlap start and end points
    overlap, = np.nonzero(overlap > 1)
    overlap_border = np.diff(overlap) > 1
    overlap_begin = overlap[np.insert(overlap_border, 0, True)]
    overlap_end = overlap[np.append(overlap_border, True)] + 1
    assert len(overlap_begin) == len(overlap_end)
    # Linear blend in overlapping areas
    for begin, end in zip(overlap_begin, overlap_end):
        alpha = np.linspace(0, 1, end - begin)[:, np.newaxis]
        style_data[begin:end] = style_data[max(begin - 1, 0)] * (1 - alpha) + style_data[min(end, len(df) - 1)] * alpha
    return style_data

def joint_parent_index():
    joint_parents = []
    for joint in NPC_JOINTS:
        parent_name = NPC_JOINT_PARENTS.get(joint, None)
        if parent_name and parent_name in NPC_JOINTS:
            joint_parents.append(NPC_JOINTS.index(parent_name))
        else:
            joint_parents.append(-1)
    return joint_parents

def simulate(player_sword_tx, styles, npc_pose, root_tx, model_file, mode):
    with tf.Graph().as_default(), tf.Session() as sess:
        gd = tf.GraphDef()
        with Path(model_file).open('rb') as f:
            gd.MergeFromString(f.read())
        model_in, model_coords, model_out = tf.import_graph_def(
            gd, name='', return_elements=['Input:0', 'Coords:0', f'Output{mode.title()}:0'])
        del gd
        num_frames = len(player_sword_tx)
        coords_size = model_coords.shape[-1].value if model_coords.shape.ndims > 0 else None
        trajectory_size = TRAJECTORY_STEP_SIZE * (TRAJECTORY_WINDOW - 1) + 1
        player_sword_trajectory = np.tile(player_sword_tx[0], (trajectory_size, 1, 1))
        npc_sword_trajectory = np.tile(npc_pose[NPC_JOINTS.index(NPC_SWORD_TIP)], (trajectory_size, 1, 1))
        npc_simulation = np.tile(npc_pose, (num_frames, 1, 1, 1))
        root_simulation = np.tile(root_tx, (num_frames, 1, 1))
        coords_simulation = np.zeros((num_frames, coords_size) if coords_size is not None else num_frames,
                                     dtype=player_sword_tx.dtype)
        prev_root_tx = root_tx
        npc_joint_vel = np.zeros_like(npc_pose[:, :3, 3])
        for i, (current_player_sword_tx, current_style) in enumerate(zip(player_sword_tx[:-1], styles[:-1])):
            # Update trajectories
            player_sword_trajectory[:-1] = player_sword_trajectory[1:]
            player_sword_trajectory[-1] = current_player_sword_tx
            npc_sword_trajectory[:-1] = npc_sword_trajectory[1:]
            npc_sword_trajectory[-1] = npc_pose[NPC_JOINTS.index(NPC_SWORD_TIP)]
            # Make input
            in_vector, in_coords = make_input_data(player_sword_trajectory, npc_sword_trajectory,
                                                   current_style, npc_pose, npc_joint_vel, root_tx,
                                                   coords_size)
            coords_simulation[i] = in_coords
            # Run model
            out_vector = sess.run(model_out, feed_dict={model_in: in_vector, model_coords: in_coords})
            # TODO New root tx
            root_vel = out_vector[0:2]
            root_angvel = out_vector[2]
            root_tx, prev_root_tx =  root_tx, root_tx
            root_simulation[i + 1] = root_tx
            # New pose
            npc_pos_out = out_vector[3:3 + 3 * len(NPC_JOINTS)].reshape((-1, 3))
            npc_vel_out = out_vector[3 + 3 * len(NPC_JOINTS):3 + 6 * len(NPC_JOINTS)].reshape((-1, 3))
            npc_prev_pos_rel = make_transform_relative(npc_pose, prev_root_tx)[:, :3, 3]
            npc_pos_out = (1 - JOINT_SMOOTH) * (npc_prev_pos_rel + npc_vel_out) + JOINT_SMOOTH * npc_pos_out
            npc_xy_out = out_vector[3 + 6 * len(NPC_JOINTS):3 + 12 * len(NPC_JOINTS)].reshape((-1, 2, 3))
            npc_x_out = npc_xy_out[:, 0]
            npc_y_out = npc_xy_out[:, 1]
            npc_tx_out = make_joint_transform(npc_pos_out, npc_x_out, npc_y_out)
            npc_pose = root_tx @ npc_tx_out
            npc_joint_vel = rotate_vector(npc_vel_out, root_tx)
            npc_simulation[i + 1] = npc_pose
    if len(coords_simulation) > 1:
        coords_simulation[-1] = coords_simulation[-2]
    return npc_simulation, root_simulation, coords_simulation

def make_input_data(player_sword_trajectory, npc_sword_trajectory, style, npc_pose, npc_joint_vel, root_tx, coords_size):
    input_size = (4 * 3 * TRAJECTORY_WINDOW + 7 + 4 + len(STYLES) + 2 * 3 * len(NPC_JOINTS))
    # Sword relative to COG
    cog_tx = root_tx
    if COG:
        cog_tx = cog_tx.copy()
        cog_tx[2, 3] = npc_pose[NPC_JOINTS.index(COG)][2, 3]
    # assert np.allclose(np.linalg.det(cog_tx[:3, :3]), 1)
    player_sword_rel_tx = make_transform_relative(player_sword_trajectory[::TRAJECTORY_STEP_SIZE], cog_tx)
    npc_sword_rel_tx = make_transform_relative(npc_sword_trajectory[::TRAJECTORY_STEP_SIZE], cog_tx)
    # Sword position
    player_sword_pos = player_sword_rel_tx[:, :3, 3]
    player_hilt_pos = transform_vector(PLAYER_SWORD_VECTOR, player_sword_rel_tx)
    npc_sword_pos = npc_sword_rel_tx[:, :3, 3]
    npc_hilt_pos = transform_vector(NPC_SWORD_VECTOR, npc_sword_rel_tx)
    # Sword direction
    player_sword_dir = player_hilt_pos - player_sword_pos
    player_sword_dir /= np.linalg.norm(player_sword_dir, axis=-1, keepdims=True)
    npc_sword_dir = npc_hilt_pos - npc_sword_pos
    npc_sword_dir /= np.linalg.norm(npc_sword_dir, axis=-1, keepdims=True)
    # Sword close points
    player_sword_close, npc_sword_close = SegmentsClosestPoints(player_sword_pos[-1],
                                                                player_hilt_pos[-1],
                                                                npc_sword_pos[-1],
                                                                npc_hilt_pos[-1])
    sword_close_dist = np.linalg.norm(player_sword_close - npc_sword_close, axis=-1)
    # Sword angles
    player_sword_elevation = np.arctan2(player_sword_pos[..., -1, 2], player_sword_pos[..., -1, 0])
    player_sword_azimuth = np.arctan2(player_sword_pos[..., -1, 1], player_sword_pos[..., -1, 0])
    player_hilt_elevation = np.arctan2(player_hilt_pos[..., -1, 2], player_hilt_pos[..., -1, 0])
    player_hilt_azimuth = np.arctan2(player_hilt_pos[..., -1, 1], player_hilt_pos[..., -1, 0])
    # Joint position
    joint_pos = make_transform_relative(npc_pose, root_tx)[:, :3, 3]
    # Joint velocity
    joint_vel = rotate_vector(npc_joint_vel, np.linalg.inv(root_tx))
    # Make vector
    vector =  np.r_[
        npc_sword_pos.ravel(),
        npc_sword_dir.ravel(),
        player_sword_pos.ravel(),
        player_sword_dir.ravel(),
        npc_sword_close,
        player_sword_close,
        sword_close_dist,
        player_sword_elevation,
        player_sword_azimuth,
        player_hilt_elevation,
        player_hilt_azimuth,
        style,
        joint_pos.ravel(),
        joint_vel.ravel(),
    ]
    assert len(vector) == input_size, f'Expected len(vector) = {input_size}, was {len(vector)}.'
    # Grid coordinates
    if coords_size is None:
        coords = (np.arctan2(player_sword_pos[..., -1, 1], player_sword_pos[..., -1, 2]) + 2 * np.pi) % (2 * np.pi)
    if coords_size == 1:
        coords = np.r_[
            (np.arctan2(player_sword_pos[..., -1, 1], player_sword_pos[..., -1, 2]) + 2 * np.pi) % (2 * np.pi)
        ]
    elif coords_size == 2:
        coords = np.r_[
            player_sword_pos[..., -1, 1] / COORD_WIDTH + 0.5,
            player_sword_pos[..., -1, 2] / COORD_HEIGHT + 0.5,
        ]
    elif coords_size == 3:
        coords = np.r_[
            (player_sword_pos[..., -1, 0] - COORD_DEPTH_RANGE[0]) / (COORD_DEPTH_RANGE[1] - COORD_DEPTH_RANGE[0]),
            player_sword_pos[..., -1, 1] / COORD_WIDTH + 0.5,
            player_sword_pos[..., -1, 2] / COORD_HEIGHT + 0.5,
        ]
    else:
        raise ValueError('invalid coordinate size {}.'.format(coords_size))
    return vector, coords

def simulation_error(npc_joint_tx, root_tx, npc_simulation_tx, root_simulation_tx):
    npc_joint_rel = npc_joint_tx @ np.linalg.inv(root_tx[:, np.newaxis])
    npc_simulation_rel = npc_simulation_tx @ np.linalg.inv(root_simulation_tx[:, np.newaxis])
    joint_parents = joint_parent_index()
    segment_idx = [(i_joint, i_parent if i_parent >= 0 else i_joint) for i_joint, i_parent in enumerate(joint_parents)]
    segment_idx = np.asarray(segment_idx)
    segments1 = np.stack([npc_joint_rel[:, segment_idx[:, 0]][..., :3, 3],
                          npc_joint_rel[:, segment_idx[:, 1]][..., :3, 3]], axis=2)
    segments2 = np.stack([npc_simulation_rel[:, segment_idx[:, 0]][..., :3, 3],
                          npc_simulation_rel[:, segment_idx[:, 1]][..., :3, 3]], axis=2)
    return segments_error(segments1, segments2)

def error_hist(error, ax=None):
    error = pd.Series(error / 10_000, name='Pose difference (m^2)')
    return sns.distplot(error, ax=ax)

def make_anim_function(ax, player_sword_tx, npc_joint_tx, root_tx, temperature=None):
    num_frames = min(len(player_sword_tx), len(npc_joint_tx), len(root_tx))
    temperature = temperature if temperature is not None else np.zeros((num_frames, len(NPC_JOINTS)))
    color_good = mpl.colors.to_rgb('#42BF42')
    color_bad = mpl.colors.to_rgb('#BC4242')
    color_sword = mpl.colors.to_rgb('#4343BA')
    joint_plt = [ax.plot3D([], [], [], c=color_sword, lw=1)[0] for _ in NPC_JOINTS]
    player_sword_plt = ax.plot3D([], [], [], c=color_sword, lw=1)[0]
    player_sword_length = np.linalg.norm(PLAYER_SWORD_VECTOR)
    player_sword_rel_tx = make_transform_relative(player_sword_tx[:num_frames], root_tx[:num_frames])
    player_sword_tip = player_sword_rel_tx[:, :3, 3]
    player_sword_hilt = transform_vector(PLAYER_SWORD_VECTOR, player_sword_rel_tx)
    np_joint_rel_tx = make_transform_relative(npc_joint_tx[:num_frames], root_tx[:num_frames, np.newaxis])
    joint_pos = np_joint_rel_tx[:, :, :3, 3]
    joint_parents = joint_parent_index()
    temperature = temperature[:, :, np.newaxis]
    color = np.array(color_good) * (1 - temperature) + np.array(color_bad) * temperature

    def anim_frame(frame):
        nonlocal player_sword_rel_tx, np_joint_rel_tx, color
        nonlocal joint_plt, player_sword_plt
        nonlocal joint_parents
        # Joints
        for i_joint, (i_plt, i_parent) in enumerate(zip(joint_plt, joint_parents)):
            if i_parent < 0: continue
            i_plt.set_data(joint_pos[frame, [i_parent, i_joint], 0], joint_pos[frame, [i_parent, i_joint], 1])
            i_plt.set_3d_properties(joint_pos[frame, [i_parent, i_joint], 2])
            i_plt.set_color(color[frame, i_joint])
        # Sword
        player_sword_plt.set_data([player_sword_tip[frame, 0], player_sword_hilt[frame, 0]],
                                  [player_sword_tip[frame, 1], player_sword_hilt[frame, 1]])
        player_sword_plt.set_3d_properties([player_sword_tip[frame, 2], player_sword_hilt[frame, 2]])
        # Updated plots
        return joint_plt + [player_sword_plt]

    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    ax.set_xlim(-100, 100)
    ax.set_ylim(-100, 100)
    ax.set_zlim(0, 200)
    ax.invert_yaxis()
    ax.set_aspect('equal')
    ax.view_init(elev=25, azim=-40)

    return anim_frame

def merge_anim_functions(anim_funcs):
    return lambda frame: [art for anim_func in anim_funcs for art in anim_func(frame)]

def make_empty_dataframe(columns, dtypes, index=None):
    return pd.concat([pd.Series(name=col, dtype=dt, index=index) for col, dt in zip(columns, dtypes)], axis=1)

# Model analysis

def num_nn_params(input_size, output_size, hidden_layers, use_bias=True):
    nn_params = 0
    last = input_size
    for layer in [*hidden_layers, output_size]:
        nn_params += last * layer
        if use_bias:
            nn_params += layer
        last = layer
    return nn_params

def num_gfnn_params(spline_points, input_size, output_size, hidden_layers, use_bias=True):
    nn_params = num_nn_params(input_size, output_size, hidden_layers, use_bias)
    spline_size = 1
    for p in spline_points:
        spline_size *= p
    return nn_params * spline_size

def flops_spline(spline_points, nn_params):
    flops = 0
    # Bigger dimensions are reduced to 4, dimensions are sorted by size
    spline_points = sorted((min(p, 4) for p in spline_points), reverse=True)
    # Weight computation cost
    for p in spline_points:
        if p > 1:
            # Weight computation fixed cost
            flops += 17
            # Dimensions with size 2 or 3 aggregate weights
            if p < 4:
                flops += 4 * p + 3 * p
    # Interpolation cost
    flops_per_param = 0
    for i, p in enumerate(spline_points):
        if p > 1:
            # Weight multiplication and reduction
            a = np.prod(spline_points[i:])
            flops_per_param += a + (a // p) * (p - 1)
    flops += flops_per_param * nn_params
    return flops

def flops_linear(spline_points, nn_params):
    n_dims = sum(1 for d in spline_points if d > 1)
    return 3 * n_dims + (1 << n_dims) * nn_params

def flops_constant(spline_points):
    n_dims = sum(1 for d in spline_points if d > 1)
    return 1 * n_dims

def flops_nn(input_size, output_size, hidden_layers, use_bias=True):
    flops = 0
    last = input_size
    for layer in [*hidden_layers, output_size]:
        flops += 2 * (last * layer)
        if use_bias:
            flops += layer
        last = layer
    return flops

def test_run_time(model_file, input_size, spline_dim, mode):
    test_exe = Path(Path(__file__).parent, 'test_model', 'test_model.exe').resolve()
    # test_exe = Path('workspace', 'test_model', 'test_model.exe').resolve()
    model_path = Path(model_file).resolve()
    args = [str(test_exe), str(model_path), str(input_size), str(spline_dim), str(mode)]
    p = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
    if p.returncode != 0:
        return math.nan
    try:
        return float(p.stdout)
    except ValueError:
        return math.nan

def get_model_metadata(model_files):
    model_names = [model_name for model_name, _, _ in model_files]
    df = make_empty_dataframe(*zip(('input_size', int),
                                   ('output_size', int),
                                   ('hidden_layers', object),
                                   ('use_bias', bool),
                                   ('parameters_nn', int),
                                   ('spline_points', object),
                                   ('parameters_spline', int),
                                   ('mode', object),
                                   ('grid_evaluations', object),
                                   ('parameters_grid', object),
                                   ('flops_interp', int),
                                   ('flops_nn', int),
                                   ('flops_total', int),
                                   ('run_time', float)),
                              index=model_names)
    df.index.name = 'model'
    for i, (model_name, model_file, model_mode) in enumerate(model_files):
        model_mode = model_mode.strip().upper()
        with tf.Graph().as_default(), tf.Session() as sess:
            # Load graph
            gd = tf.GraphDef()
            with Path(model_file).open('rb') as f:
                gd.MergeFromString(f.read())
            inp, coords, out, *metadata = tf.import_graph_def(gd, name='', return_elements=[
                'Input:0', 'Coords:0', f'Output{model_mode.title()}:0',
                'Meta/NumParameters:0', 'Meta/SplinePoints:0', f'Meta/GridEvaluations{model_mode.title()}:0',
                'Meta/HiddenLayers:0', 'Meta/UseBias:0', 'Meta/InputSize:0', 'Meta/OutputSize:0'
            ])
            # Get metadata
            params_spline, spline_points, grid_evals, hidden_layers, use_bias, input_size, output_size = sess.run(metadata)
            params_nn = params_spline // np.prod(spline_points)
            params_grid = params_nn * np.prod(grid_evals)
            # Flops estimation
            if model_mode == 'CUBIC':
                flops_interp = flops_spline(spline_points, params_nn)
            elif model_mode == 'LINEAR':
                flops_interp = flops_linear(spline_points, params_nn)
            elif model_mode == 'CONSTANT':
                flops_interp = flops_constant(spline_points)
            else:
                raise ValueError(f'unknown model mode "{model_mode}"')
            flops_net = flops_nn(input_size, output_size, hidden_layers, use_bias)
            flops_total = flops_interp + flops_net
            run_time = test_run_time(model_file, input_size, len(spline_points), model_mode.title())
            df.loc[model_name, :] = [input_size, output_size, ' '.join(map(str, hidden_layers)), use_bias, params_nn,
                                     ' '.join(map(str, spline_points)), params_spline,
                                     model_mode, ' '.join(map(str, grid_evals)), params_grid,
                                     flops_interp, flops_net, flops_total, run_time]
    return df

def get_error_statistics(errors):
    stats_columns = ['mean', 'std', 'min', 'max', 'p25', 'p50', 'p75']
    #errors = np.asarray(errors)[1:]  # Skip first common frame
    return pd.DataFrame([[np.mean(errors, axis=0), np.std(errors, axis=0),
                          np.min(errors, axis=0), np.max(errors, axis=0),
                          np.percentile(errors, 25, axis=0),
                          np.percentile(errors, 50, axis=0),
                          np.percentile(errors, 75, axis=0)]],
                          columns=stats_columns)

def num_rows_cols(n):
    for a in range(int(math.sqrt(n)), 0, -1):
        if n % a == 0:
            break
    b = n // a
    return min(a, b), max(a, b)

def main(data_files, model_files, out_dir, per_data, anim_dir, anim_layout, anim_layout_vert, anim_dpi, error_bone):
    data_names = [data_name for data_name, _ in data_files]
    model_names = [model_name for model_name, _, _ in model_files]
    data_df = make_empty_dataframe(['num_frames'], [int], index=data_names)
    data_df.index.name = 'data'
    model_df = get_model_metadata(model_files)
    data_fps = []
    data_frame_errors = []
    data_anim = []
    sim_errors = [[] for _ in model_files]
    for data_name, data_file in data_files:
        print(f'Reading data {data_name}...')
        df = pd.read_csv(data_file, index_col=0)
        num_frames = len(df)
        data_df.loc[data_name, :] = num_frames
        npc_joint_tx = get_npc_joint_transforms(df)
        root_tx = get_root_transforms(df)
        player_sword_tx, npc_sword_tx = get_sword_transforms(df)
        styles = get_styles(df)
        anim_fps = 1. / df[TIME_COLUMN].diff().mean()
        data_fps.append(anim_fps)
        error_df = make_empty_dataframe(model_names, [float] * len(model_files), index=range(num_frames))
        npc_simulation_tx = []
        root_simulation_tx = []
        for i, (model_name, model_file, model_mode) in enumerate(model_files):
            print(f'Simulating {model_name} on {data_name}...')
            npc_model_tx, root_model_tx, coords_model = simulate(
                player_sword_tx, styles, npc_joint_tx[0], root_tx[0], model_file, model_mode)
            model_error = simulation_error(npc_joint_tx, root_tx, npc_model_tx, root_model_tx)
            model_frame_error = model_error.sum(1)
            error_df[model_name] = model_frame_error
            npc_simulation_tx.append(npc_model_tx)
            root_simulation_tx.append(root_model_tx)
            sim_errors[i].append(model_error)
        data_frame_errors.append(error_df.melt(var_name='model', value_name='error'))
        data_frame_errors[-1].insert(0, 'data', data_name)
        data_frame_errors[-1].insert(1, 'frame', range(len(error_df)))
        # Animation
        if anim_dir:
            animation_funcs = []
            if anim_layout:
                rows, cols = anim_layout
            else:
                rows, cols = num_rows_cols(len(model_files) + 1)
            fig = plt.figure(figsize=(5 * cols, 5 * rows), dpi=anim_dpi)
            ax = fig.add_subplot(rows, cols, 1, projection='3d')
            ax.set_title(data_name)
            animation_funcs.append(make_anim_function(ax, player_sword_tx, npc_joint_tx, root_tx))
            for i, (model_name, model_file, _) in enumerate(model_files):
                plot_idx = i + 2
                if anim_layout_vert:
                    plot_idx = (i % rows) * cols + (i // rows) + 1
                simulation_ax = fig.add_subplot(rows, cols, plot_idx, projection='3d')
                simulation_ax.set_title(model_name)
                if error_bone:
                    anim_simulation_temperature = np.clip(sim_errors[i][-1] / ERROR_BONE_SATURATION, 0, 1)
                else:
                    anim_simulation_temperature = sim_errors[i][-1].copy()
                    anim_simulation_temperature[:] = anim_simulation_temperature.sum(1, keepdims=True)
                    anim_simulation_temperature = np.clip(anim_simulation_temperature / ERROR_FRAME_SATURATION, 0, 1)
                animation_funcs.append(make_anim_function(
                    simulation_ax, player_sword_tx, npc_simulation_tx[i], root_simulation_tx[i], anim_simulation_temperature))
            anim_interval = 1000 / anim_fps
            anim_all_func = merge_anim_functions(animation_funcs)
            fig.tight_layout(pad=1.5)
            data_anim.append(mplanim.FuncAnimation(fig, anim_all_func, range(num_frames), interval=anim_interval, repeat=False, blit=True))
    # Save results
    if out_dir:
        print('Saving results...')
        out_dir = Path(out_dir)
        if not out_dir.exists():
            out_dir.mkdir(parents=True)
        with pd.ExcelWriter(str(Path(out_dir, 'results.xlsx'))) as writer:
            if model_files:
                model_df.to_excel(writer, sheet_name='Models')
            if data_files:
                data_df.to_excel(writer, sheet_name='Data')
            if data_files and model_files:
                error_df = pd.concat(data_frame_errors, ignore_index=True)
                model_group = error_df.groupby(['model'])['error']
                overall_df = (model_group
                              .apply(get_error_statistics)
                              .reset_index(level=-1, drop=True)
                              .loc[model_names, :])
                overall_df.to_excel(writer, sheet_name='Overall')
                if per_data:
                    data_group = error_df.groupby(['data', 'model'])['error']
                    data_stats = (data_group
                                  .apply(get_error_statistics)
                                  .reset_index(level=-1, drop=True))
                    for data_name, _ in data_files:
                        data_df = data_stats.loc[data_name].loc[model_names]
                        data_df.to_excel(writer, sheet_name=data_name)
                        for model_name, _, _ in model_files:
                            plt.subplot(111)
                            error_hist(data_group.get_group((data_name, model_name)))
                            plt.savefig(str(Path(out_dir, f'error_{model_name}_{data_name}.png')))
                            plt.close()
                error_df.to_excel(writer, index=False, sheet_name='Error')
                for model_name, _, _ in model_files:
                    plt.subplot(111)
                    error_hist(model_group.get_group(model_name))
                    plt.savefig(str(Path(out_dir, f'error_{model_name}.png')))
                    plt.close()
    # Write animations
    if anim_dir:
        print('Animating...')
        anim_dir = Path(anim_dir)
        if not anim_dir.exists():
            anim_dir.mkdir(parents=True)
        with ThreadPool() as pool:
            pool.starmap(lambda data, fps, anim: anim.save(str(Path(anim_dir, f'{data[0]}.mp4')),
                                                           writer=mplanim.writers['ffmpeg'](fps=fps),
                                                           dpi=anim_dpi),
                         zip(data_files, data_fps, data_anim))
    print('Done.')


if __name__ == '__main__':
    main(*ARGS)
