app.py

# Define functions and run GaitDynamics.

# NOTE: the OpenSim Python API (used for the 3D skeleton glTF export) is installed at
# BUILD time from a vendored, metadata-relaxed opensim 4.6 wheel (see requirements.txt).
# It is therefore a normal site-packages package — importable in the main process and
# inherited by the ZeroGPU @spaces.GPU worker — so no runtime pip install is needed.
import numpy as np
import pandas as pd
from utils import first_gait_cycle_indices, read_mot, WCB_COLOR_LO, WCB_COLOR_HI
# `spaces` must be imported before torch so its monkey-patch is active: on ZeroGPU it
# makes torch.cuda.is_available() return True and defers .to("cuda") until a
# @spaces.GPU-decorated function attaches the GPU. Pre-installed in the ZeroGPU image.
import spaces
from torch.utils.data import Dataset
import torch.nn.functional as F
import torch
from torch import nn
from torch import Tensor
import argparse
from typing import Any
import nimblephysics as nimble
import scipy.interpolate as interpo
from scipy.interpolate import interp1d
import random
from scipy.signal import filtfilt, butter
from accelerate import Accelerator, DistributedDataParallelKwargs
import os
from inspect import isfunction
from math import pi
from einops import rearrange, repeat
import math
import copy
import matplotlib.pyplot as plt
from tqdm import tqdm
from functools import partial
from accelerate.state import AcceleratorState


device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.manual_seed(0)
random.seed(0)
np.random.seed(0)

""" ============================ Start scaler.py ============================ """


def _handle_zeros_in_scale(scale, copy=True, constant_mask=None):
    # if we are fitting on 1D arrays, scale might be a scalar
    if constant_mask is None:
        # Detect near constant values to avoid dividing by a very small
        # value that could lead to surprising results and numerical
        # stability issues.
        constant_mask = scale < 10 * torch.finfo(scale.dtype).eps

    if copy:
        # New array to avoid side-effects
        scale = scale.clone()
    scale[constant_mask] = 1.0
    return scale


class MinMaxScaler:
    _parameter_constraints: dict = {
        "feature_range": [tuple],
        "copy": ["boolean"],
        "clip": ["boolean"],
    }

    def __init__(self, feature_range=(-1, 1), *, copy=True, clip=True):
        self.feature_range = feature_range
        self.copy = copy
        self.clip = clip

    def _reset(self):
        if hasattr(self, "scale_"):
            del self.scale_
            del self.min_
            del self.n_samples_seen_
            del self.data_min_
            del self.data_max_
            del self.data_range_

    def fit(self, X):
        # Reset internal state before fitting
        self._reset()
        return self.partial_fit(X)

    def partial_fit(self, X):
        feature_range = self.feature_range
        if feature_range[0] >= feature_range[1]:
            raise ValueError(
                "Minimum of desired feature range must be smaller than maximum. Got %s."
                % str(feature_range)
            )

        data_min = torch.min(X, axis=0)[0]
        data_max = torch.max(X, axis=0)[0]

        self.n_samples_seen_ = X.shape[0]

        data_range = data_max - data_min
        self.scale_ = (feature_range[1] - feature_range[0]) / _handle_zeros_in_scale(
            data_range, copy=True
        )
        self.min_ = feature_range[0] - data_min * self.scale_
        self.data_min_ = data_min
        self.data_max_ = data_max
        self.data_range_ = data_range
        return self

    def transform(self, X):
        X *= self.scale_.to(X.device)
        X += self.min_.to(X.device)
        return X

    def inverse_transform(self, X):
        X -= self.min_.to(X.device)
        X /= self.scale_.to(X.device)
        return X


class Normalizer:
    def __init__(self, scaler, cols_to_normalize):
        self.scaler = scaler
        self.cols_to_normalize = cols_to_normalize

    def normalize(self, x):
        x = x.clone()
        x[:, self.cols_to_normalize] = self.scaler.transform(x[:, self.cols_to_normalize])
        return x

    def unnormalize(self, x):
        batch, seq, ch = x.shape
        x = x.clone()
        x = x.reshape(-1, ch)
        x[:, self.cols_to_normalize] = self.scaler.inverse_transform(x[:, self.cols_to_normalize])
        return x.reshape((batch, seq, ch))


""" ============================ End scaler.py ============================ """


""" ============================ Start consts.py ============================ """

JOINTS_3D_ALL = {
    'pelvis': ['pelvis_tilt', 'pelvis_list', 'pelvis_rotation'],
    'hip_r': ['hip_flexion_r', 'hip_adduction_r', 'hip_rotation_r'],
    'hip_l': ['hip_flexion_l', 'hip_adduction_l', 'hip_rotation_l'],
    'lumbar': ['lumbar_extension', 'lumbar_bending', 'lumbar_rotation'],
    # 'arm_r': ['arm_flex_r', 'arm_add_r', 'arm_rot_r'],
    # 'arm_l': ['arm_flex_l', 'arm_add_l', 'arm_rot_l']
    }

OSIM_DOF_ALL = [
    'pelvis_tilt', 'pelvis_list', 'pelvis_rotation', 'pelvis_tx', 'pelvis_ty', 'pelvis_tz', 'hip_flexion_r',
    'hip_adduction_r', 'hip_rotation_r', 'knee_angle_r', 'ankle_angle_r', 'subtalar_angle_r', 'mtp_angle_r',
    'hip_flexion_l', 'hip_adduction_l', 'hip_rotation_l', 'knee_angle_l', 'ankle_angle_l', 'subtalar_angle_l',
    'mtp_angle_l', 'lumbar_extension', 'lumbar_bending', 'lumbar_rotation', 'arm_flex_r', 'arm_add_r', 'arm_rot_r',
    'elbow_flex_r', 'pro_sup_r', 'wrist_flex_r', 'wrist_dev_r', 'arm_flex_l', 'arm_add_l', 'arm_rot_l',
    'elbow_flex_l', 'pro_sup_l', 'wrist_flex_l', 'wrist_dev_l']


KINETICS_ALL = [body + modality for body in ['calcn_r', 'calcn_l'] for modality in
                ['_force_vx', '_force_vy', '_force_vz', '_force_normed_cop_x', '_force_normed_cop_y', '_force_normed_cop_z']]

_GRF_COL_MAP = {
    'calcn_r_force_vx':           'ground_force_1_vx',
    'calcn_r_force_vy':           'ground_force_1_vy',
    'calcn_r_force_vz':           'ground_force_1_vz',
    'calcn_r_force_normed_cop_x': 'ground_force_1_px',
    'calcn_r_force_normed_cop_y': 'ground_force_1_py',
    'calcn_r_force_normed_cop_z': 'ground_force_1_pz',
    'calcn_l_force_vx':           'ground_force_2_vx',
    'calcn_l_force_vy':           'ground_force_2_vy',
    'calcn_l_force_vz':           'ground_force_2_vz',
    'calcn_l_force_normed_cop_x': 'ground_force_2_px',
    'calcn_l_force_normed_cop_y': 'ground_force_2_py',
    'calcn_l_force_normed_cop_z': 'ground_force_2_pz',
}

MODEL_STATES_COLUMN_NAMES_WITH_ARM = [
                                         'pelvis_tx', 'pelvis_ty', 'pelvis_tz', 'knee_angle_r', 'ankle_angle_r', 'subtalar_angle_r',
                                         'knee_angle_l', 'ankle_angle_l', 'subtalar_angle_l', 'elbow_flex_r', 'pro_sup_r', 'elbow_flex_l', 'pro_sup_l'
                                     ] + KINETICS_ALL + [
                                         'pelvis_0', 'pelvis_1', 'pelvis_2', 'pelvis_3', 'pelvis_4', 'pelvis_5',
                                         'hip_r_0', 'hip_r_1', 'hip_r_2', 'hip_r_3', 'hip_r_4', 'hip_r_5',
                                         'hip_l_0', 'hip_l_1', 'hip_l_2', 'hip_l_3', 'hip_l_4', 'hip_l_5',
                                         'lumbar_0', 'lumbar_1', 'lumbar_2', 'lumbar_3', 'lumbar_4', 'lumbar_5',
                                         'arm_r_0', 'arm_r_1', 'arm_r_2', 'arm_r_3', 'arm_r_4', 'arm_r_5',        # only for with arm
                                         'arm_l_0', 'arm_l_1', 'arm_l_2', 'arm_l_3', 'arm_l_4', 'arm_l_5'        # only for with arm
                                     ]

MODEL_STATES_COLUMN_NAMES_NO_ARM = copy.deepcopy(MODEL_STATES_COLUMN_NAMES_WITH_ARM)
for name_ in ['elbow_flex_r', 'pro_sup_r', 'elbow_flex_l', 'pro_sup_l', 'arm_r_0', 'arm_r_1', 'arm_r_2', 'arm_r_3',
              'arm_r_4', 'arm_r_5', 'arm_l_0', 'arm_l_1', 'arm_l_2', 'arm_l_3', 'arm_l_4', 'arm_l_5']:
    MODEL_STATES_COLUMN_NAMES_NO_ARM.remove(name_)

FROZEN_DOFS = ['mtp_angle_r', 'mtp_angle_l',
               'wrist_flex_r', 'wrist_dev_r', 'wrist_flex_l', 'wrist_dev_l']

FULL_OSIM_DOF = ['pelvis_tilt', 'pelvis_list', 'pelvis_rotation', 'pelvis_tx', 'pelvis_ty', 'pelvis_tz',
                 'hip_flexion_r', 'hip_adduction_r', 'hip_rotation_r', 'knee_angle_r', 'ankle_angle_r',
                 'subtalar_angle_r', 'mtp_angle_r', 'hip_flexion_l', 'hip_adduction_l', 'hip_rotation_l',
                 'knee_angle_l', 'ankle_angle_l', 'subtalar_angle_l', 'mtp_angle_l', 'lumbar_extension',
                 'lumbar_bending', 'lumbar_rotation']

JOINTS_1D_ALL = ['pelvis_tx', 'pelvis_ty', 'pelvis_tz',
                 'knee_angle_r', 'ankle_angle_r', 'subtalar_angle_r', 'mtp_angle_r',
                 'knee_angle_l', 'ankle_angle_l', 'subtalar_angle_l', 'mtp_angle_l']

""" ============================ End consts.py ============================ """


""" ============================ Start model.py ============================ """


class GaussianDiffusion(nn.Module):
    def __init__(
            self,
            model,
            horizon,
            repr_dim,
            opt,
            n_timestep=1000,
            schedule="linear",
            loss_type="l1",
            clip_denoised=False,
            predict_epsilon=True,
            guidance_weight=1,
            use_p2=False,
            cond_drop_prob=0.,
    ):
        super().__init__()
        self.horizon = horizon
        self.transition_dim = repr_dim
        self.model = model
        self.ema = EMA(0.99)
        self.master_model = copy.deepcopy(self.model)
        self.opt = opt

        self.cond_drop_prob = cond_drop_prob

        # make a SMPL instance for FK module

        betas = torch.Tensor(
            make_beta_schedule(schedule=schedule, n_timestep=n_timestep)
        )
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, axis=0)
        alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]])

        self.n_timestep = int(n_timestep)
        self.clip_denoised = clip_denoised
        self.predict_epsilon = predict_epsilon

        self.register_buffer("betas", betas)
        self.register_buffer("alphas_cumprod", alphas_cumprod)
        self.register_buffer("alphas_cumprod_prev", alphas_cumprod_prev)

        self.guidance_weight = guidance_weight

        # calculations for diffusion q(x_t | x_{t-1}) and others
        self.register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod))
        self.register_buffer(
            "sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod)
        )
        self.register_buffer(
            "log_one_minus_alphas_cumprod", torch.log(1.0 - alphas_cumprod)
        )
        self.register_buffer(
            "sqrt_recip_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod)
        )
        self.register_buffer(
            "sqrt_recipm1_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod - 1)
        )

        # calculations for posterior q(x_{t-1} | x_t, x_0)
        posterior_variance = (
                betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
        )
        self.register_buffer("posterior_variance", posterior_variance)

        ## log calculation clipped because the posterior variance
        ## is 0 at the beginning of the diffusion chain
        self.register_buffer(
            "posterior_log_variance_clipped",
            torch.log(torch.clamp(posterior_variance, min=1e-20)),
        )
        self.register_buffer(
            "posterior_mean_coef1",
            betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod),
            )
        self.register_buffer(
            "posterior_mean_coef2",
            (1.0 - alphas_cumprod_prev) * np.sqrt(alphas) / (1.0 - alphas_cumprod),
            )

        # p2 weighting
        self.p2_loss_weight_k = 1
        self.p2_loss_weight_gamma = 0.5 if use_p2 else 0
        self.register_buffer(
            "p2_loss_weight",
            (self.p2_loss_weight_k + alphas_cumprod / (1 - alphas_cumprod))
            ** -self.p2_loss_weight_gamma,
            )

        ## get loss coefficients and initialize objective
        self.loss_fn = F.mse_loss if loss_type == "l2" else F.l1_loss

    def set_normalizer(self, normalizer):
        self.normalizer = normalizer

    # ------------------------------------------ sampling ------------------------------------------#

    def predict_start_from_noise(self, x_t, t, noise):
        """
            if self.predict_epsilon, model output is (scaled) noise;
            otherwise, model predicts x0 directly
        """
        if self.predict_epsilon:
            return (
                    extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
                    - extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
            )
        else:
            return noise

    def predict_noise_from_start(self, x_t, t, x0):
        return (
                (extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0) / \
                extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
        )

    def model_predictions(self, x, cond, time_cond, weight=None, clip_x_start=False):
        weight = weight if weight is not None else self.guidance_weight
        model_output = self.model.guided_forward(x, cond, time_cond, weight)
        maybe_clip = partial(torch.clamp, min=-1., max=1.) if clip_x_start else identity

        x_start = model_output
        x_start = maybe_clip(x_start)
        pred_noise = self.predict_noise_from_start(x, time_cond, x_start)

        return pred_noise, x_start

    def q_posterior(self, x_start, x_t, t):
        posterior_mean = (
                extract(self.posterior_mean_coef1, t, x_t.shape) * x_start
                + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
        )
        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
        posterior_log_variance_clipped = extract(
            self.posterior_log_variance_clipped, t, x_t.shape
        )
        return posterior_mean, posterior_variance, posterior_log_variance_clipped

    def p_mean_variance(self, x, cond, t):
        # guidance clipping
        if t[0] > 1.0 * self.n_timestep:
            weight = min(self.guidance_weight, 0)
        elif t[0] < 0.1 * self.n_timestep:
            weight = min(self.guidance_weight, 1)
        else:
            weight = self.guidance_weight

        x_recon = self.predict_start_from_noise(
            x, t=t, noise=self.model.guided_forward(x, cond, t, weight)
        )

        if self.clip_denoised:
            x_recon.clamp_(-1.0, 1.0)
        else:
            assert RuntimeError()

        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
            x_start=x_recon, x_t=x, t=t
        )
        return model_mean, posterior_variance, posterior_log_variance, x_recon

    @torch.no_grad()
    def p_sample(self, x, cond, t):
        b, *_, device = *x.shape, x.device
        model_mean, _, model_log_variance, x_start = self.p_mean_variance(
            x=x, cond=cond, t=t
        )
        noise = torch.randn_like(model_mean)
        # no noise when t == 0
        nonzero_mask = (1 - (t == 0).float()).reshape(
            b, *((1,) * (len(noise.shape) - 1))
        )
        x_out = model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
        return x_out, x_start

    @torch.no_grad()
    def p_sample_loop(          # Only used during inference
            self,
            shape,
            cond,
            noise=None,
            constraint=None,
            return_diffusion=False,
            start_point=None,
    ):
        device = self.betas.device

        # default to diffusion over whole timescale
        start_point = self.n_timestep if start_point is None else start_point
        batch_size = shape[0]
        x = torch.randn(shape, device=device) if noise is None else noise.to(device)
        cond = cond.to(device)

        if return_diffusion:
            diffusion = [x]

        for i in tqdm(reversed(range(0, start_point))):
            # fill with i
            timesteps = torch.full((batch_size,), i, device=device, dtype=torch.long)
            x, _ = self.p_sample(x, cond, timesteps)

            if return_diffusion:
                diffusion.append(x)

        if return_diffusion:
            return x, diffusion
        else:
            return x

    @torch.no_grad()
    def inpaint_ddim_guided(self, shape, noise=None, constraint=None, return_diffusion=False, start_point=None):
        batch, device, total_timesteps, sampling_timesteps, eta = shape[0], self.betas.device, self.n_timestep, 50, 1

        times = torch.linspace(-1, total_timesteps - 1, steps=sampling_timesteps + 1)   # [-1, 0, 1, 2, ..., T-1] when sampling_timesteps == total_timesteps
        times = list(reversed(times.int().tolist()))
        time_pairs = list(zip(times[:-1], times[1:])) # [(T-1, T-2), (T-2, T-3), ..., (1, 0), (0, -1)]

        x = torch.randn(shape, device=device)
        cond = constraint["cond"].to(device)

        mask = constraint["mask"].to(device)  # batch x horizon x channels
        value = constraint["value"].to(device)  # batch x horizon x channels
        value_diff_thd = constraint["value_diff_thd"].to(device)  # channels
        value_diff_weight = constraint["value_diff_weight"].to(device)  # channels

        for time, time_next in time_pairs:
            time_cond = torch.full((batch,), time, device=device, dtype=torch.long)

            if self.opt.guide_x_start_the_end_step <= time <= self.opt.guide_x_start_the_beginning_step:
                x.requires_grad_()
                with torch.enable_grad():
                    for step_ in range(self.opt.n_guided_steps):
                        pred_noise, x_start, *_ = self.model_predictions(x, cond, time_cond, clip_x_start=self.clip_denoised)
                        value_diff = torch.subtract(x_start, value)
                        loss = torch.relu(value_diff.abs() - value_diff_thd) * value_diff_weight
                        grad = torch.autograd.grad([loss.sum()], [x])[0]
                        x = x - self.opt.guidance_lr * grad

            pred_noise, x_start, *_ = self.model_predictions(x, cond, time_cond, clip_x_start=self.clip_denoised)
            if time_next < 0:
                x = x_start
                x = value * mask + (1.0 - mask) * x
                return x

            alpha = self.alphas_cumprod[time]
            alpha_next = self.alphas_cumprod[time_next]

            sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
            c = (1 - alpha_next - sigma ** 2).sqrt()

            noise = torch.randn_like(x)

            x = x_start * alpha_next.sqrt() + \
                c * pred_noise + \
                sigma * noise

            timesteps = torch.full((batch,), time_next, device=device, dtype=torch.long)
            value_ = self.q_sample(value, timesteps) if (time > 0) else x
            x = value_ * mask + (1.0 - mask) * x
        return x

    @torch.no_grad()
    def inpaint_ddim_loop(self, shape, noise=None, constraint=None, return_diffusion=False, start_point=None):
        batch, device, total_timesteps, sampling_timesteps, eta = shape[0], self.betas.device, self.n_timestep, 50, 0

        times = torch.linspace(-1, total_timesteps - 1, steps=sampling_timesteps + 1)
        times = list(reversed(times.int().tolist()))
        time_pairs = list(zip(times[:-1], times[1:]))

        x = torch.randn(shape, device=device)
        cond = constraint["cond"].to(device)

        mask = constraint["mask"].to(device)  # batch x horizon x channels
        value = constraint["value"].to(device)  # batch x horizon x channels

        for time, time_next in time_pairs:
            time_cond = torch.full((batch,), time, device=device, dtype=torch.long)

            pred_noise, x_start, *_ = self.model_predictions(x, cond, time_cond, clip_x_start=self.clip_denoised)
            if time_next < 0:
                x = x_start
                return x

            alpha = self.alphas_cumprod[time]
            alpha_next = self.alphas_cumprod[time_next]

            sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
            c = (1 - alpha_next - sigma ** 2).sqrt()

            noise = torch.randn_like(x)

            x = x_start * alpha_next.sqrt() + \
                c * pred_noise + \
                sigma * noise

            timesteps = torch.full((batch,), time_next, device=device, dtype=torch.long)
            value_ = self.q_sample(value, timesteps) if (time > 0) else x
            x = value_ * mask + (1.0 - mask) * x
        return x

    def q_sample(self, x_start, t, noise=None):     # blend noise into state variables
        if noise is None:
            noise = torch.randn_like(x_start)

        sample = (
                extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
                + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
        )
        return sample

    def forward(self, x, cond, t_override=None):
        return self.loss(x, cond, t_override)

    def generate_samples(
            self,
            shape,
            normalizer,
            opt,
            mode,
            noise=None,
            constraint=None,
            start_point=None,
    ):
        torch.manual_seed(torch.randint(0, 2 ** 32, (1,)).item())
        if isinstance(shape, tuple):
            if mode == "inpaint":
                func_class = self.inpaint_ddim_loop
            elif mode == "inpaint_ddim_guided":
                func_class = self.inpaint_ddim_guided
            else:
                assert False, "Unrecognized inference mode"
            samples = (
                func_class(
                    shape,
                    noise=noise,
                    constraint=constraint,
                    start_point=start_point,
                )
            )
        else:
            samples = shape

        samples = normalizer.unnormalize(samples.detach().cpu())
        samples = samples.detach().cpu()
        return samples


class FillingBase:
    def fill_param(self, windows, diffusion_model_for_filling, progress_callback=None):
        return self.filling(windows, diffusion_model_for_filling, self.update_kinematics_and_masks_for_masking_column, progress_callback=progress_callback)

    @staticmethod
    def update_kinematics_and_masks_for_masking_column(windows, samples, i_win, masks):
        unmasked_samples_in_temporal_dim = (masks.sum(axis=2)).bool()
        for j_win in range(len(samples)):
            windows[j_win+i_win].pose = samples[j_win]
            updated_mask = windows[j_win+i_win].mask
            updated_mask[unmasked_samples_in_temporal_dim[j_win], :] = 1
            updated_mask[:, opt.kinetic_diffusion_col_loc] = 0
            windows[j_win+i_win].mask = updated_mask
        return windows

    def update_kinematics_and_masks_for_masking_temporal(self, windows, samples, i_win, masks):
        for j_win in range(len(samples)):
            windows[j_win+i_win].pose = samples[j_win]
            windows[j_win+i_win].mask = self.mask_original[j_win+i_win]
        return windows

    @staticmethod
    def filling(windows, diffusion_model_for_filling, windows_update_func):
        raise NotImplementedError


class DiffusionFilling(FillingBase):
    @staticmethod
    def filling(windows, diffusion_model_for_filling, windows_update_func, progress_callback=None):
        windows = copy.deepcopy(windows)
        total_iterations = len(range(0, len(windows), opt.batch_size_inference))
        for iteration_idx, i_win in enumerate(range(0, len(windows), opt.batch_size_inference)):
            state_true = torch.stack([win.pose for win in windows[i_win:i_win+opt.batch_size_inference]])
            masks = torch.stack([win.mask for win in windows[i_win:i_win+opt.batch_size_inference]])
            cond = torch.ones([6])

            constraint = {'mask': masks, 'value': state_true.clone(), 'cond': cond}
            shape = (state_true.shape[0], state_true.shape[1], state_true.shape[2])
            samples = (diffusion_model_for_filling.diffusion.inpaint_ddim_loop(shape, constraint=constraint))
            samples = state_true * masks + (1.0 - masks) * samples.to(state_true.device)
            # samples[:, :, opt.kinetic_diffusion_col_loc] = state_true[:, :, opt.kinetic_diffusion_col_loc]

            windows = windows_update_func(windows, samples, i_win, masks)
            
            # Report progress: 10% for first iteration, 80% for last iteration
            if progress_callback:
                progress = 0.1 + 0.7 * (iteration_idx + 1) / total_iterations
                progress_callback(progress)
        return windows

    def __str__(self):
        return 'diffusion_filling'


def wrap(x):
    return {f"module.{key}": value for key, value in x.items()}


def maybe_wrap(x, num):
    return x if num == 1 else wrap(x)


def exists(val):
    return val is not None


def broadcat(tensors, dim=-1):
    num_tensors = len(tensors)
    shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
    assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
    shape_len = list(shape_lens)[0]

    dim = (dim + shape_len) if dim < 0 else dim
    dims = list(zip(*map(lambda t: list(t.shape), tensors)))

    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
    assert all(
        [*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]
    ), "invalid dimensions for broadcastable concatentation"
    max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
    expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
    expanded_dims.insert(dim, (dim, dims[dim]))
    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
    tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
    return torch.cat(tensors, dim=dim)


def rotate_half(x):
    x = rearrange(x, "... (d r) -> ... d r", r=2)
    x1, x2 = x.unbind(dim=-1)
    x = torch.stack((-x2, x1), dim=-1)
    return rearrange(x, "... d r -> ... (d r)")


def apply_rotary_emb(freqs, t, start_index=0):
    freqs = freqs.to(t)
    rot_dim = freqs.shape[-1]
    end_index = start_index + rot_dim
    assert (
            rot_dim <= t.shape[-1]
    ), f"feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}"
    t_left, t, t_right = (
        t[..., :start_index],
        t[..., start_index:end_index],
        t[..., end_index:],
    )
    t = (t * freqs.cos()) + (rotate_half(t) * freqs.sin())
    return torch.cat((t_left, t, t_right), dim=-1)


class RotaryEmbedding(nn.Module):
    def __init__(
            self,
            dim,
            custom_freqs=None,
            freqs_for="lang",
            theta=10000,
            max_freq=10,
            num_freqs=1,
            learned_freq=False,
    ):
        super().__init__()
        if exists(custom_freqs):
            freqs = custom_freqs
        elif freqs_for == "lang":
            freqs = 1.0 / (
                    theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)
            )
        elif freqs_for == "pixel":
            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
        elif freqs_for == "constant":
            freqs = torch.ones(num_freqs).float()
        else:
            raise ValueError(f"unknown modality {freqs_for}")

        self.cache = dict()

        if learned_freq:
            self.freqs = nn.Parameter(freqs)
        else:
            self.register_buffer("freqs", freqs)

    def rotate_queries_or_keys(self, t, seq_dim=-2):
        device = t.device
        seq_len = t.shape[seq_dim]
        freqs = self.forward(
            lambda: torch.arange(seq_len, device=device), cache_key=seq_len
        )
        return apply_rotary_emb(freqs, t)

    def forward(self, t, cache_key=None):
        if exists(cache_key) and cache_key in self.cache:
            return self.cache[cache_key]

        if isfunction(t):
            t = t()

        freqs = self.freqs

        freqs = torch.einsum("..., f -> ... f", t.type(freqs.dtype), freqs)
        freqs = repeat(freqs, "... n -> ... (n r)", r=2)

        if exists(cache_key):
            self.cache[cache_key] = freqs

        return freqs

class EMA:
    def __init__(self, beta):
        super().__init__()
        self.beta = beta

    def update_model_average(self, ma_model, current_model):
        for current_params, ma_params in zip(
                current_model.parameters(), ma_model.parameters()
        ):
            old_weight, up_weight = ma_params.data, current_params.data
            ma_params.data = self.update_average(old_weight, up_weight)

    def update_average(self, old, new):
        if old is None:
            return new
        return old * self.beta + (1 - self.beta) * new


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len, dropout: float = 0.):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        position = torch.arange(max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(max_len * 2) / d_model))
        pe = torch.zeros(max_len, 1, d_model)
        pe[:, 0, 0::2] = torch.sin(position * div_term)
        try:
            pe[:, 0, 1::2] = torch.cos(position * div_term)
        except RuntimeError:
            pe[:, 0, 1::2] = torch.cos(position * div_term)[:, :-1]

        pe = pe.transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x: Tensor) -> Tensor:
        """
        Args:
            x: Tensor, shape [batch_size, seq_len, embedding_dim]
        """
        x = x + self.pe[:, :x.shape[1], :]
        return self.dropout(x)


class PositionWiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super(PositionWiseFeedForward, self).__init__()
        self.fc1 = nn.Linear(d_model, d_ff)
        self.fc2 = nn.Linear(d_ff, d_model)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))


class EncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads=8, d_ff=512, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.rotary = RotaryEmbedding(dim=d_model)
        self.use_rotary = self.rotary is not None

        self.self_attn = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
        self.feed_forward = PositionWiseFeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        qk = self.rotary.rotate_queries_or_keys(x) if self.use_rotary else x
        attn_output, _ = self.self_attn(qk, qk, x, need_weights=False)
        x = self.norm1(x + self.dropout(attn_output))
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))
        return x


class TransformerEncoderArchitecture(nn.Module):
    def __init__(self, repr_dim, opt, nlayers=6):
        super(TransformerEncoderArchitecture, self).__init__()
        self.input_dim = len(opt.kinematic_diffusion_col_loc)
        self.output_dim = repr_dim - self.input_dim
        embedding_dim = 192
        self.input_to_embedding = nn.Linear(self.input_dim, embedding_dim)
        self.encoder_layers = nn.Sequential(*[EncoderLayer(embedding_dim) for _ in range(nlayers)])
        self.embedding_to_output = nn.Linear(embedding_dim, self.output_dim)
        self.opt = opt
        self.input_col_loc = opt.kinematic_diffusion_col_loc
        self.output_col_loc = [i for i in range(repr_dim) if i not in self.input_col_loc]

    def loss_fun(self, output_pred, output_true):
        return F.mse_loss(output_pred, output_true, reduction='none')

    def end_to_end_prediction(self, x):
        input = x[0][:, :, self.input_col_loc]
        sequence = self.input_to_embedding(input)
        sequence = self.encoder_layers(sequence)
        output_pred = self.embedding_to_output(sequence)
        return output_pred

    def __str__(self):
        return 'tf'


class DiffusionShellForAdaptingTheOriginalFramework(nn.Module):
    def __init__(self, model):
        super(DiffusionShellForAdaptingTheOriginalFramework, self).__init__()
        self.device = device
        self.model = model
        self.ema = EMA(0.99)
        self.master_model = copy.deepcopy(self.model)

    def set_normalizer(self, normalizer):
        self.normalizer = normalizer

    def predict_samples(self, x, constraint):
        x[0] = x[0] * constraint['mask']
        output_pred = self.model.end_to_end_prediction(x)
        x[0][:, :, self.model.output_col_loc] = output_pred
        return x[0]

    def forward(self, x, cond, t_override):
        output_true = x[0][:, :, self.model.output_col_loc]
        output_pred = self.model.end_to_end_prediction(x)
        loss_simple = torch.zeros(x[0].shape).to(x[0].device)
        loss_simple[:, :, self.model.output_col_loc] = self.model.loss_fun(output_pred, output_true)

        losses = [
            1. * loss_simple.mean(),
            torch.tensor(0.).to(loss_simple.device),
            torch.tensor(0.).to(loss_simple.device),
            torch.tensor(0.).to(loss_simple.device),
            torch.tensor(0).to(loss_simple.device)]
        return sum(losses), losses + [loss_simple]


def featurewise_affine(x, scale_shift):
    scale, shift = scale_shift
    return (scale + 1) * x + shift


class DenseFiLM(nn.Module):
    """Feature-wise linear modulation (FiLM) generator."""

    def __init__(self, embed_channels):
        super().__init__()
        self.embed_channels = embed_channels
        self.block = nn.Sequential(
            nn.Mish(), nn.Linear(embed_channels, embed_channels * 2)
        )

    def forward(self, position):
        pos_encoding = self.block(position)
        pos_encoding = rearrange(pos_encoding, "b c -> b 1 c")
        scale_shift = pos_encoding.chunk(2, dim=-1)
        return scale_shift


class FiLMTransformerDecoderLayer(nn.Module):
    def __init__(
            self,
            d_model: int,
            nhead: int,
            dim_feedforward=2048,
            dropout=0.1,
            activation=F.relu,
            layer_norm_eps=1e-5,
            batch_first=False,
            norm_first=True,
            device=None,
            dtype=None,
            rotary=None,
    ):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(
            d_model, nhead, dropout=dropout, batch_first=batch_first
        )
        self.multihead_attn = nn.MultiheadAttention(
            d_model, nhead, dropout=dropout, batch_first=batch_first
        )
        # Feedforward
        self.linear1 = nn.Linear(d_model, dim_feedforward)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(dim_feedforward, d_model)

        self.norm_first = norm_first
        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps)
        self.norm3 = nn.LayerNorm(d_model, eps=layer_norm_eps)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.activation = activation

        self.film1 = DenseFiLM(d_model)
        self.film3 = DenseFiLM(d_model)

        self.rotary = rotary
        self.use_rotary = rotary is not None

    # x, cond, t
    def forward(
            self,
            tgt,
            memory,
            t,
            tgt_mask=None,
            memory_mask=None,
            tgt_key_padding_mask=None,
            memory_key_padding_mask=None,
    ):
        x = tgt
        if self.norm_first:
            x_1 = self._sa_block(self.norm1(x), tgt_mask, tgt_key_padding_mask)
            x = x + featurewise_affine(x_1, self.film1(t))
            x_3 = self._ff_block(self.norm3(x))
            x = x + featurewise_affine(x_3, self.film3(t))
        else:
            x = self.norm1(
                x
                + featurewise_affine(
                    self._sa_block(x, tgt_mask, tgt_key_padding_mask), self.film1(t)
                )
            )
            x = self.norm3(x + featurewise_affine(self._ff_block(x), self.film3(t)))
        return x

    # self-attention block
    # qkv
    def _sa_block(self, x, attn_mask, key_padding_mask):
        qk = self.rotary.rotate_queries_or_keys(x) if self.use_rotary else x
        x = self.self_attn(
            qk,
            qk,
            x,
            attn_mask=attn_mask,
            key_padding_mask=key_padding_mask,
            need_weights=False,
        )[0]
        return self.dropout1(x)

    # multihead attention block
    # qkv
    def _mha_block(self, x, mem, attn_mask, key_padding_mask):
        q = self.rotary.rotate_queries_or_keys(x) if self.use_rotary else x
        k = self.rotary.rotate_queries_or_keys(mem) if self.use_rotary else mem
        x = self.multihead_attn(
            q,
            k,
            mem,
            attn_mask=attn_mask,
            key_padding_mask=key_padding_mask,
            need_weights=False,
        )[0]
        return self.dropout2(x)

    # feed forward block
    def _ff_block(self, x):
        x = self.linear2(self.dropout(self.activation(self.linear1(x))))
        return self.dropout3(x)


class DecoderLayerStack(nn.Module):
    def __init__(self, stack):
        super().__init__()
        self.stack = stack

    def forward(self, x, cond, t):
        for layer in self.stack:
            x = layer(x, cond, t)
        return x


class DanceDecoder(nn.Module):
    def __init__(
            self,
            nfeats: int,
            seq_len: int = 150,  # 5 seconds, 30 fps
            latent_dim: int = 512,
            ff_size: int = 1024,
            num_layers: int = 8,
            num_heads: int = 8,
            dropout: float = 0.1,
            # cond_feature_dim: int = 6,
            activation=F.gelu,
            use_rotary=True,
            **kwargs
    ) -> None:

        super().__init__()

        output_feats = nfeats

        # positional embeddings
        self.rotary = None
        self.abs_pos_encoding = nn.Identity()
        # if rotary, replace absolute embedding with a rotary embedding instance (absolute becomes an identity)
        if use_rotary:
            self.rotary = RotaryEmbedding(dim=latent_dim)
        else:
            self.abs_pos_encoding = PositionalEncoding(
                latent_dim, dropout, batch_first=True
            )

        # time embedding processing
        self.time_mlp = nn.Sequential(
            SinusoidalPosEmb(latent_dim),
            nn.Linear(latent_dim, latent_dim * 4),
            nn.Mish(),
        )
        # input projection
        self.input_projection = nn.Linear(nfeats, latent_dim)

        self.to_time_cond = nn.Sequential(nn.Linear(latent_dim * 4, latent_dim),)

        decoderstack = nn.ModuleList([])
        for _ in range(num_layers):
            decoderstack.append(
                FiLMTransformerDecoderLayer(        # decoder layers
                    latent_dim,
                    num_heads,
                    dim_feedforward=ff_size,
                    dropout=dropout,
                    activation=activation,
                    batch_first=True,
                    rotary=self.rotary,
                )
            )

        self.seqTransDecoder = DecoderLayerStack(decoderstack)

        self.final_layer = nn.Linear(latent_dim, output_feats)

    def guided_forward(self, x, cond_embed, time_cond, guidance_weight):
        return self.forward(x, cond_embed, time_cond)

    def __str__(self):
        return 'diffusion'

    # No conditioning version
    def forward(self, x: Tensor, cond_embed: Tensor, time_cond: Tensor, cond_drop_prob: float = 0.0):
        x = self.input_projection(x)
        x = self.abs_pos_encoding(x)
        t_hidden = self.time_mlp(time_cond)
        t = self.to_time_cond(t_hidden)
        output = self.seqTransDecoder(x, None, t)
        output = self.final_layer(output)
        return output


class MotionModel:
    def __init__(
            self,
            opt,
            normalizer=None,
            EMA=True,
            learning_rate=4e-4,
            weight_decay=0.02,
    ):
        self.opt = opt
        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
        self.accelerator = Accelerator(cpu=(device != 'cuda'), kwargs_handlers=[ddp_kwargs])
        print(f"Diffusion model running on {self.accelerator.device}")
        state = AcceleratorState()
        num_processes = state.num_processes

        self.repr_dim = len(opt.model_states_column_names)

        self.horizon = horizon = opt.window_len

        self.accelerator.wait_for_everyone()

        checkpoint = None
        if opt.checkpoint != "":
            checkpoint = torch.load(
                opt.checkpoint, map_location=self.accelerator.device,
                weights_only=False
            )
            self.normalizer = checkpoint["normalizer"]

        model = DanceDecoder(
            nfeats=self.repr_dim,
            seq_len=horizon,
            latent_dim=512,
            ff_size=1024,
            num_layers=8,
            num_heads=8,
            dropout=0.1,
            activation=F.gelu,
        )

        diffusion = GaussianDiffusion(
            model,
            horizon,
            self.repr_dim,
            opt,
            # schedule="cosine",
            n_timestep=1000,
            predict_epsilon=False,
            loss_type="l2",
            use_p2=False,
            cond_drop_prob=0.,
            guidance_weight=2,
        )

        self.model = self.accelerator.prepare(model)
        self.diffusion = diffusion.to(self.accelerator.device)

        if opt.checkpoint != "":
            self.model.load_state_dict(
                maybe_wrap(
                    checkpoint["ema_state_dict" if EMA else "model_state_dict"],
                    num_processes,
                )
            )

    def eval(self):
        self.diffusion.eval()

    def prepare(self, objects):
        return self.accelerator.prepare(*objects)

    def eval_loop(self, opt, state_true, masks, value_diff_thd=None, value_diff_weight=None, cond=None,
                  num_of_generation_per_window=1, mode="inpaint"):
        self.eval()
        if value_diff_thd is None:
            value_diff_thd = torch.zeros([state_true.shape[2]])
        if value_diff_weight is None:
            value_diff_weight = torch.ones([state_true.shape[2]])
        if cond is None:
            cond = torch.ones([6])

        constraint = {'mask': masks, 'value': state_true.clone(), 'value_diff_thd': value_diff_thd,
                      'value_diff_weight': value_diff_weight, 'cond': cond}

        shape = (state_true.shape[0], self.horizon, self.repr_dim)
        state_pred_list = [self.diffusion.generate_samples(
            shape,
            self.normalizer,
            opt,
            mode=mode,
            constraint=constraint)
            for _ in range(num_of_generation_per_window)]
        return torch.stack(state_pred_list)


class BaselineModel:
    def __init__(
            self,
            opt,
            model_architecture_class,
            EMA=True,
    ):
        self.opt = opt
        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
        self.accelerator = Accelerator(cpu=(device != 'cuda'), kwargs_handlers=[ddp_kwargs])
        print(f"Baseline model running on {self.accelerator.device}")
        self.repr_dim = len(opt.model_states_column_names)
        self.horizon = horizon = opt.window_len
        self.accelerator.wait_for_everyone()

        self.model = model_architecture_class(self.repr_dim, opt)
        self.diffusion = DiffusionShellForAdaptingTheOriginalFramework(self.model)
        self.diffusion = self.accelerator.prepare(self.diffusion)

        checkpoint = None
        if opt.checkpoint_bl != "":
            checkpoint = torch.load(
                opt.checkpoint_bl, map_location=self.accelerator.device,
                weights_only=False
            )
            self.normalizer = checkpoint["normalizer"]

        if opt.checkpoint_bl != "":
            self.model.load_state_dict(
                maybe_wrap(
                    checkpoint["ema_state_dict" if EMA else "model_state_dict"],
                    1,
                )
            )

    def eval_loop(self, opt, state_true, masks, value_diff_thd=None, value_diff_weight=None, cond=None,
                  num_of_generation_per_window=1, mode="inpaint"):
        self.eval()
        constraint = {'mask': masks.to(self.accelerator.device), 'value': state_true, 'cond': cond}
        state_true = state_true.to(self.accelerator.device)
        state_pred_list = [self.diffusion.predict_samples([state_true], constraint)
                           for _ in range(num_of_generation_per_window)]
        state_pred_list = [self.normalizer.unnormalize(state_pred.detach().cpu()) for state_pred in state_pred_list]
        return torch.stack(state_pred_list)

    def eval(self):
        self.diffusion.eval()

    def train(self):
        self.diffusion.train()

    def prepare(self, objects):
        return self.accelerator.prepare(*objects)


""" ============================ End model.py ============================ """


""" ============================ Start util.py ============================ """


def load_diffusion_model(opt):
    opt.checkpoint = DIFFUSION_CKPT
    model = MotionModel(opt)
    model_key = 'diffusion'
    return model, model_key


class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        device = x.device
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = x[:, None] * emb[None, :]
        emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
        return emb


def convertDfToGRFMot(df, out_path, dt, time_column):
    numFrames = df.shape[0]
    for key in df.keys():
        if key == 'TimeStamp':
            continue

    out_file = open(out_path, 'w')
    out_file.write('nColumns=9\n')
    out_file.write('nRows='+str(numFrames)+'\n')
    out_file.write('DataType=double\n')
    out_file.write('version=3\n')
    out_file.write('OpenSimVersion=4.1\n')
    out_file.write('endheader\n')
    out_file.write('time')

    plate_num = 2

    for i, side in zip(range(1, 1 + plate_num), ['r', 'l']):
        out_file.write('\t' + f'force_{side}_vx')
        out_file.write('\t' + f'force_{side}_vy')
        out_file.write('\t' + f'force_{side}_vz')
        out_file.write('\t' + f'force_{side}_px')
        out_file.write('\t' + f'force_{side}_py')
        out_file.write('\t' + f'force_{side}_pz')
        out_file.write('\t' + f'torque_{side}_x')
        out_file.write('\t' + f'torque_{side}_y')
        out_file.write('\t' + f'torque_{side}_z')

    out_file.write('\n')
    for i in range(numFrames):
        out_file.write(str(round(dt * i + time_column[0], 5)))
        for side in ['r', 'l']:
            out_file.write('\t' + str(df[f'calcn_{side}_force_vx'][i]))
            out_file.write('\t' + str(df[f'calcn_{side}_force_vy'][i]))
            out_file.write('\t' + str(df[f'calcn_{side}_force_vz'][i]))
            out_file.write('\t' + str(df[f'calcn_{side}_force_normed_cop_x'][i]))
            out_file.write('\t' + str(df[f'calcn_{side}_force_normed_cop_y'][i]))
            out_file.write('\t' + str(df[f'calcn_{side}_force_normed_cop_z'][i]))
            out_file.write('\t' + str(0))
            out_file.write('\t' + str(0))
            out_file.write('\t' + str(0))
        out_file.write('\n')
    out_file.close()
    print('Ground reaction forces exported to ' + out_path)


def convertDataframeToMot(df, out_path, dt, time_column):
    numFrames = df.shape[0]
    out_file = open(out_path, 'w')
    out_file.write('Coordinates\n')
    out_file.write('version=1\n')
    out_file.write(f'nRows={numFrames}\n')
    out_file.write(f'nColumns={len(df.columns)+1}\n')
    out_file.write('inDegrees=no\n\n')
    out_file.write('If the header above contains a line with \'inDegrees\', this indicates whether rotational values are in degrees (yes) or radians (no).\n\n')
    out_file.write('endheader\n')

    out_file.write('time')
    for i in range(len(df.columns)):
        out_file.write('\t' + df.columns[i])
    out_file.write('\n')

    for i in range(numFrames):
        out_file.write(str(round(dt * i + time_column[0], 5)))
        for j in range(len(df.columns)):
            out_file.write('\t' + str(df.iloc[i, j]))
        out_file.write('\n')
    out_file.close()
    print('Missing kinematics exported to ' + out_path)


def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
    a1, a2 = d6[..., :3], d6[..., 3:]
    b1 = F.normalize(a1, dim=-1)
    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
    b2 = F.normalize(b2, dim=-1)
    b3 = torch.cross(b1, b2, dim=-1)
    return torch.stack((b1, b2, b3), dim=-2)


def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
    batch_dim = matrix.size()[:-2]
    return matrix[..., :2, :].clone().reshape(batch_dim + (6,))


def euler_from_6v(q, convention="XYZ"):
    assert q.shape[-1] == 6
    mat = rotation_6d_to_matrix(q)
    eul = matrix_to_euler_angles(mat, convention)
    return eul


def euler_to_6v(q, convention="XYZ"):
    assert q.shape[-1] == 3
    mat = euler_angles_to_matrix(q, convention)
    mat = matrix_to_rotation_6d(mat)
    return mat


def second_order_poly(coeff, x):
    y = coeff[...,0] * x**2 + coeff[...,1] * x + coeff[...,2]
    return y


def batch_identity(batch_shape, size):
    batch_identity = torch.eye(size)
    output_shape = batch_shape.copy()
    output_shape.append(size)
    output_shape.append(size)
    batch_identity_out = batch_identity.view(*(1,) * (len(output_shape) - batch_identity.ndim),*batch_identity.shape).expand(output_shape)
    return batch_identity_out.clone()


def get_knee_rotation_coefficients():

    knee_Z_rotation_function = np.array([[0, 0.174533, 0.349066, 0.523599, 0.698132, 0.872665, 1.0472, 1.22173,
                                          1.39626, 1.5708, 1.74533, 1.91986, 2.0944],
                                         [0, 0.0126809, 0.0226969, 0.0296054, 0.0332049, 0.0335354, 0.0308779,
                                          0.0257548, 0.0189295, 0.011407, 0.00443314, -0.00050475, -0.0016782]]).T
    polyfit_knee_Z_rotation = np.polyfit(knee_Z_rotation_function[:,0], knee_Z_rotation_function[:,1], deg=2, full = True)
    coefficients_knee_Z_rotation = polyfit_knee_Z_rotation[0]

    knee_Y_rotation_function = np.array([[0, 0.174533, 0.349066, 0.523599, 0.698132, 0.872665, 1.0472, 1.22173,
                                          1.39626, 1.5708, 1.74533, 1.91986, 2.0944],
                                         [0, 0.059461, 0.109399, 0.150618, 0.18392, 0.210107, 0.229983, 0.24435, 0.254012, 0.25977, 0.262428, 0.262788, 0.261654]]).T
    polyfit_knee_Y_rotation = np.polyfit(knee_Y_rotation_function[:, 0], knee_Y_rotation_function[:, 1], deg=2, full=True)
    coefficients_knee_Y_rotation = polyfit_knee_Y_rotation[0]

    knee_X_translation_function = np.array([[0, 0.174533, 0.349066, 0.523599, 0.698132, 0.872665, 1.0472, 1.22173,
                                             1.39626, 1.5708, 1.74533, 1.91986, 2.0944],
                                            [0, 5.3e-05, 0.000188, 0.000378, 0.000597, 0.000825, 0.001045, 0.001247, 0.00142, 0.001558, 0.001661, 0.001728, 0.00176]]).T
    polyfit_knee_X_translation = np.polyfit(knee_X_translation_function[:, 0], knee_X_translation_function[:, 1], deg=2,
                                            full=True)
    coefficients_knee_X_translation = polyfit_knee_X_translation[0]


    knee_Y_translation_function = np.array([[0, 0.174533, 0.349066, 0.523599, 0.698132, 0.872665, 1.0472, 1.22173,
                                             1.39626, 1.5708, 1.74533, 1.91986, 2.0944],
                                            [0, 0.000301, 0.000143, -0.000401, -0.001233, -0.002243, -0.003316, -0.004346, -0.005239, -0.005924, -0.006361, -0.006539, -0.00648]]).T
    polyfit_knee_Y_translation = np.polyfit(knee_Y_translation_function[:, 0], knee_Y_translation_function[:, 1], deg=2,
                                            full=True)
    coefficients_knee_Y_translation = polyfit_knee_Y_translation[0]

    knee_Z_translation_function = np.array([[0, 0.174533, 0.349066, 0.523599, 0.698132, 0.872665, 1.0472, 1.22173,
                                             1.39626, 1.5708, 1.74533, 1.91986, 2.0944],
                                            [0, 0.001055, 0.002061, 0.00289, 0.003447, 0.003676, 0.003559, 0.00311, 0.002373, 0.001418, 0.000329, -0.000805, -0.001898]]).T
    polyfit_knee_Z_translation = np.polyfit(knee_Z_translation_function[:, 0], knee_Z_translation_function[:, 1], deg=2,
                                            full=True)
    coefficients_knee_Z_translation = polyfit_knee_Z_translation[0]

    walker_knee_coefficients = np.stack((coefficients_knee_Y_rotation, coefficients_knee_Z_rotation, coefficients_knee_X_translation, coefficients_knee_Y_translation, coefficients_knee_Z_translation), axis=1)

    return walker_knee_coefficients


walker_knee_coefficients = get_knee_rotation_coefficients()
# Kept on CPU: its only reader is forward_kinematics (currently unused). Avoids the sole
# import-time .to("cuda"), which would otherwise rely on the spaces monkey-patch.
walker_knee_coefficients = torch.tensor(walker_knee_coefficients)


def forward_kinematics(pose, offsets, with_arm=False):
    """
    Pose indices
    0-5: pelvis orientation + translation
    6-8: hip_r
    9: knee_r
    10: ankle_r
    11: subtalar_r
    12: mtp_r
    13-15: hip_l
    16: knee_l
    17: ankle_l
    18: subtalar_l
    19: mtp_l
    20-22: lumbar
    23-25: shoulder_r
    26: elbow_r
    27: radioulnar
    28-30: shoulder_l
    31: elbow_l
    32: radioulnar_l
    """
    if isinstance(pose, np.ndarray):
        pose = torch.from_numpy(pose)
    pose = pose.to(torch.device(device))
    if len(pose.shape) == 2:
        pose = pose[None, ...]
    if len(offsets.shape) == 3:
        offsets = offsets[None, ...]

    offsets = offsets.to(pose.device)
    offsets = offsets[:, None, ...]
    batch_shape = pose.shape[:-1]
    batch_shape_list = []
    for i in range(pose.dim()-1):
        batch_shape_list.append(int(batch_shape[i]))
    batch_shape = batch_shape_list
    if batch_shape == ():
        batch_shape = (1,)

    coefficients_knee_Y_rotation = walker_knee_coefficients[..., 0]
    coefficients_knee_Z_rotation = walker_knee_coefficients[..., 1]
    coefficients_knee_X_translation = walker_knee_coefficients[..., 2]
    coefficients_knee_Y_translation = walker_knee_coefficients[..., 3]
    coefficients_knee_Z_translation = walker_knee_coefficients[..., 4]

    knee_r_Y_rot = second_order_poly(coefficients_knee_Y_rotation, pose[..., 9])
    knee_r_Z_rot = second_order_poly(coefficients_knee_Z_rotation, pose[..., 9])
    knee_r_X_trans = second_order_poly(coefficients_knee_X_translation, pose[..., 9])
    knee_r_Y_trans = second_order_poly(coefficients_knee_Y_translation, pose[..., 9])
    knee_r_Z_trans = second_order_poly(coefficients_knee_Z_translation, pose[..., 9])

    knee_l_Y_rot = second_order_poly(coefficients_knee_Y_rotation, pose[..., 16])
    knee_l_Z_rot = second_order_poly(coefficients_knee_Z_rotation, pose[..., 16])
    knee_l_X_trans = second_order_poly(coefficients_knee_X_translation, pose[..., 16])
    knee_l_Y_trans = second_order_poly(coefficients_knee_Y_translation, pose[..., 16])
    knee_l_Z_trans = second_order_poly(coefficients_knee_Z_translation, pose[..., 16])

    # Pelvis
    pelvis_transform = batch_identity(batch_shape, 4).to(pose.device)
    pelvis_transform[..., :3, :3] = euler_angles_to_matrix(pose[..., 0:3], 'ZXY')
    pelvis_transform[..., :3, 3] = pose[..., 3:6].clone().detach()

    # Get offsets (model and model scaling dependent)
    offset_hip_pelvis_r = offsets[..., 0]
    femur_offset_in_hip_r = offsets[..., 1]
    knee_offset_in_femur_r = offsets[..., 2]
    tibia_offset_in_knee_r = offsets[..., 3]
    ankle_offset_in_tibia_r = offsets[..., 4]
    talus_offset_in_ankle_r = offsets[..., 5]
    subtalar_offset_in_talus_r = offsets[..., 6]
    calcaneus_offset_in_subtalar_r = offsets[..., 7]
    mtp_offset_in_calcaneus_r = offsets[..., 8]
    offset_hip_pelvis_l = offsets[..., 9]
    femur_offset_in_hip_l = offsets[..., 10]
    knee_offset_in_femur_l = offsets[..., 11]
    tibia_offset_in_knee_l = offsets[..., 12]
    ankle_offset_in_tibia_l = offsets[..., 13]
    talus_offset_in_ankle_l = offsets[..., 14]
    subtalar_offset_in_talus_l = offsets[..., 15]
    calcaneus_offset_in_subtalar_l = offsets[..., 16]
    mtp_offset_in_calcaneus_l = offsets[..., 17]
    lumbar_offset_in_pelvis = offsets[..., 18]
    torso_offset_in_lumbar = offsets[..., 19]
    if with_arm:
        shoulder_offset_in_torso_r = offsets[..., 20]
        humerus_offset_in_shoulder_r = offsets[..., 21]
        elbow_offset_in_humerus_r = offsets[..., 22]
        ulna_offset_in_elbow_r = offsets[..., 23]
        radioulnar_offset_in_radius_r = offsets[..., 24]
        radius_offset_in_radioulnar_r = offsets[..., 25]
        wrist_offset_in_radius_r = offsets[..., 26]
        hand_offset_in_wrist_r = offsets[..., 27]
        shoulder_offset_in_torso_l = offsets[..., 28]
        humerus_offset_in_shoulder_l = offsets[..., 29]
        elbow_offset_in_humerus_l = offsets[..., 30]
        ulna_offset_in_elbow_l = offsets[..., 31]
        radioulnar_offset_in_radius_l = offsets[..., 32]
        radius_offset_in_radioulnar_l = offsets[..., 33]
        wrist_offset_in_radius_l = offsets[..., 34]
        hand_offset_in_wrist_l = offsets[..., 35]

    # Coordinates to transformation matrix
    hip_coordinates_transform_r = batch_identity(batch_shape, 4).to(pose.device)
    knee_coordinates_transform_r = batch_identity(batch_shape, 4).to(pose.device)
    ankle_coordinates_transform_r = batch_identity(batch_shape, 4).to(pose.device)
    subtalar_coordinates_transform_r = batch_identity(batch_shape, 4).to(pose.device)
    mtp_coordinates_transform_r = batch_identity(batch_shape, 4).to(pose.device)
    hip_coordinates_transform_l = batch_identity(batch_shape, 4).to(pose.device)
    knee_coordinates_transform_l = batch_identity(batch_shape, 4).to(pose.device)
    ankle_coordinates_transform_l = batch_identity(batch_shape, 4).to(pose.device)
    subtalar_coordinates_transform_l = batch_identity(batch_shape, 4).to(pose.device)
    mtp_coordinates_transform_l = batch_identity(batch_shape, 4).to(pose.device)
    lumbar_coordinates_transform = batch_identity(batch_shape, 4).to(pose.device)
    if with_arm:
        shoulder_coordinates_transform_r = batch_identity(batch_shape, 4).to(pose.device)
        elbow_coordinates_transform_r = batch_identity(batch_shape, 4).to(pose.device)
        radioulnar_coordinates_transform_r = batch_identity(batch_shape, 4).to(pose.device)
        wrist_coordinates_transform_r = batch_identity(batch_shape, 4).to(pose.device)
        shoulder_coordinates_transform_l = batch_identity(batch_shape, 4).to(pose.device)
        elbow_coordinates_transform_l = batch_identity(batch_shape, 4).to(pose.device)
        radioulnar_coordinates_transform_l = batch_identity(batch_shape, 4).to(pose.device)
        wrist_coordinates_transform_l = batch_identity(batch_shape, 4).to(pose.device)

    # Knee axis translation
    knee_coordinates_transform_r[..., :3, -1] = torch.stack((knee_r_X_trans, knee_r_Y_trans, knee_r_Z_trans), dim=-1)
    knee_coordinates_transform_l[..., :3, -1] = torch.stack((knee_l_X_trans, knee_l_Y_trans, -knee_l_Z_trans), dim=-1)

    # Joint rotations
    zero_2_shape = batch_shape.copy()
    zero_2_shape.append(2)
    zero_2 = torch.zeros(tuple(zero_2_shape), device=pose.device)
    zero_3_shape = batch_shape.copy()
    zero_3_shape.append(3)
    zero_3 = torch.zeros(tuple(zero_3_shape), device=pose.device)
    hip_coordinates_transform_r[..., :3, :3] = euler_angles_to_matrix((pose[..., 6:9]), 'ZXY')
    knee_coordinates_transform_r[..., :3, :3] = euler_angles_to_matrix(
        torch.stack((pose[..., 9], knee_r_Y_rot, knee_r_Z_rot), dim=-1), 'XYZ')
    ankle_coordinates_transform_r[..., :3, :3] = euler_angles_to_matrix(
        torch.cat((pose[..., 10:11], zero_2), dim=-1), 'ZXY')
    subtalar_coordinates_transform_r[..., :3, :3] = euler_angles_to_matrix(
        torch.cat((pose[..., 11:12], zero_2), dim=-1), 'ZXY')
    mtp_coordinates_transform_r[..., :3, :3] = euler_angles_to_matrix(
        torch.cat((pose[..., 12:13], zero_2), dim=-1), 'ZXY')
    hip_coordinates_transform_l[..., :3, :3] = euler_angles_to_matrix(
        torch.cat(([pose[..., 13:14], -pose[..., 14:15], -pose[..., 15:16]]), dim=-1), 'ZXY')
    knee_coordinates_transform_l[..., :3, :3] = euler_angles_to_matrix(
        torch.stack((-pose[..., 16], -knee_l_Y_rot, knee_l_Z_rot), dim=-1), 'XYZ')
    ankle_coordinates_transform_l[..., :3, :3] = euler_angles_to_matrix(
        torch.cat((pose[..., 17:18], zero_2), dim=-1), 'ZXY')
    subtalar_coordinates_transform_l[..., :3, :3] = euler_angles_to_matrix(
        torch.cat((pose[..., 18:19], zero_2), dim=-1), 'ZXY')
    mtp_coordinates_transform_l[..., :3, :3] = euler_angles_to_matrix(
        torch.cat((pose[..., 19:20], zero_2), dim=-1), 'ZXY')
    lumbar_coordinates_transform[..., :3, :3] = euler_angles_to_matrix(
        torch.cat((pose[..., 20:21], pose[..., 21:22], pose[..., 22:23]), dim=-1), 'ZXY')
    if with_arm:
        shoulder_coordinates_transform_r[..., :3, :3] = euler_angles_to_matrix(
            torch.cat((pose[..., 23:24], pose[..., 24:25], pose[..., 25:26]), dim=-1), 'ZXY')
        elbow_coordinates_transform_r[..., :3, :3] = euler_angles_to_matrix(
            torch.cat((pose[..., 26:27], zero_2), dim=-1), 'ZXY')
        radioulnar_coordinates_transform_r[..., :3, :3] = euler_angles_to_matrix(
            torch.cat((pose[..., 27:28], zero_2), dim=-1), 'ZXY')
        wrist_coordinates_transform_r[..., :3, :3] = euler_angles_to_matrix(
            zero_3, 'ZXY')
        shoulder_coordinates_transform_l[..., :3, :3] = euler_angles_to_matrix(
            torch.cat((pose[..., 28:29], -pose[..., 29:30], -pose[..., 30:31]), dim=-1), 'ZXY')
        elbow_coordinates_transform_l[..., :3, :3] = euler_angles_to_matrix(
            torch.cat((pose[..., 31:32], zero_2), dim=-1), 'ZXY')
        radioulnar_coordinates_transform_l[..., :3, :3] = euler_angles_to_matrix(
            torch.cat((pose[..., 32:33], zero_2), dim=-1), 'ZXY')
        wrist_coordinates_transform_l[..., :3, :3] = euler_angles_to_matrix(
            zero_3, 'ZXY')

    # Forward kinematics for the lower body
    hip_transform_r = torch.matmul(torch.matmul(pelvis_transform, offset_hip_pelvis_r), hip_coordinates_transform_r)
    femur_transform_r = torch.matmul(hip_transform_r, femur_offset_in_hip_r)
    knee_transform_r = torch.matmul(torch.matmul(femur_transform_r, knee_offset_in_femur_r),
                                    knee_coordinates_transform_r)
    tibia_transform_r = torch.matmul(knee_transform_r, tibia_offset_in_knee_r)
    ankle_transform_r = torch.matmul(torch.matmul(tibia_transform_r, ankle_offset_in_tibia_r),
                                     ankle_coordinates_transform_r)
    talus_transform_r = torch.matmul(ankle_transform_r, talus_offset_in_ankle_r)
    subtalar_transform_r = torch.matmul(torch.matmul(talus_transform_r, subtalar_offset_in_talus_r),
                                        subtalar_coordinates_transform_r)
    calcaneus_transform_r = torch.matmul(subtalar_transform_r, calcaneus_offset_in_subtalar_r)
    mtp_offset_transform_r = torch.matmul(torch.matmul(calcaneus_transform_r, mtp_offset_in_calcaneus_r),
                                          mtp_coordinates_transform_r)

    hip_transform_l = torch.matmul(torch.matmul(pelvis_transform, offset_hip_pelvis_l), hip_coordinates_transform_l)
    femur_transform_l = torch.matmul(hip_transform_l, femur_offset_in_hip_l)

    knee_transform_l = torch.matmul(torch.matmul(femur_transform_l, knee_offset_in_femur_l),
                                    knee_coordinates_transform_l)
    tibia_transform_l = torch.matmul(knee_transform_l, tibia_offset_in_knee_l)

    ankle_transform_l = torch.matmul(torch.matmul(tibia_transform_l, ankle_offset_in_tibia_l),
                                     ankle_coordinates_transform_l)
    talus_transform_l = torch.matmul(ankle_transform_l, talus_offset_in_ankle_l)

    subtalar_transform_l = torch.matmul(torch.matmul(talus_transform_l, subtalar_offset_in_talus_l),
                                        subtalar_coordinates_transform_l)
    calcaneus_transform_l = torch.matmul(subtalar_transform_l, calcaneus_offset_in_subtalar_l)

    mtp_offset_transform_l = torch.matmul(torch.matmul(calcaneus_transform_l, mtp_offset_in_calcaneus_l),
                                          mtp_coordinates_transform_l)

    # Forward kinematics for the upper body
    lumbar_transform = torch.matmul(torch.matmul(pelvis_transform, lumbar_offset_in_pelvis), lumbar_coordinates_transform)
    torso_transform = torch.matmul(lumbar_transform, torso_offset_in_lumbar)
    if with_arm:
        shoulder_transform_r = torch.matmul(torch.matmul(torso_transform, shoulder_offset_in_torso_r), shoulder_coordinates_transform_r)
        humerus_transform_r = torch.matmul(shoulder_transform_r, humerus_offset_in_shoulder_r)
        elbow_transform_r = torch.matmul(torch.matmul(humerus_transform_r, elbow_offset_in_humerus_r), elbow_coordinates_transform_r)
        ulna_transform_r = torch.matmul(elbow_transform_r, ulna_offset_in_elbow_r)
        radioulnar_transform_r = torch.matmul(torch.matmul(ulna_transform_r, radioulnar_offset_in_radius_r), radioulnar_coordinates_transform_r)
        radius_transform_r = torch.matmul(radioulnar_transform_r, radius_offset_in_radioulnar_r)
        wrist_transform_r = torch.matmul(torch.matmul(ulna_transform_r, wrist_offset_in_radius_r), wrist_coordinates_transform_r)
        hand_transform_r = torch.matmul(wrist_transform_r, hand_offset_in_wrist_r)
        shoulder_transform_l = torch.matmul(torch.matmul(torso_transform, shoulder_offset_in_torso_l), shoulder_coordinates_transform_l)
        humerus_transform_l = torch.matmul(shoulder_transform_l, humerus_offset_in_shoulder_l)
        elbow_transform_l = torch.matmul(torch.matmul(humerus_transform_l, elbow_offset_in_humerus_l), elbow_coordinates_transform_l)
        ulna_transform_l = torch.matmul(elbow_transform_l, ulna_offset_in_elbow_l)
        radioulnar_transform_l = torch.matmul(torch.matmul(ulna_transform_l, radioulnar_offset_in_radius_l), radioulnar_coordinates_transform_l)
        radius_transform_l = torch.matmul(radioulnar_transform_l, radius_offset_in_radioulnar_l)
        wrist_transform_l = torch.matmul(torch.matmul(ulna_transform_l, wrist_offset_in_radius_l), wrist_coordinates_transform_l)
        hand_transform_l = torch.matmul(wrist_transform_l, hand_offset_in_wrist_l)

    joint_locations = torch.stack((pelvis_transform[..., :3, 3],
                                   hip_transform_r[..., :3, 3], knee_transform_r[..., :3, 3],
                                   ankle_transform_r[..., :3, 3], calcaneus_transform_r[..., :3, 3],
                                   mtp_offset_transform_r[..., :3, 3],
                                   hip_transform_l[..., :3, 3], knee_transform_l[..., :3, 3],
                                   ankle_transform_l[..., :3, 3], calcaneus_transform_l[..., :3, 3],
                                   mtp_offset_transform_l[..., :3, 3]))
    if with_arm:
        joint_locations = torch.stack((*[joint_locations[i] for i in range(joint_locations.shape[0])],
                                       lumbar_transform[..., :3, 3],
                                       shoulder_transform_r[..., :3, 3],
                                       elbow_transform_r[..., :3, 3], wrist_transform_r[..., :3, 3],
                                       shoulder_transform_l[..., :3, 3],
                                       elbow_transform_l[..., :3, 3], wrist_transform_l[..., :3, 3]))
    if torch.isnan(joint_locations).any():
        print('NAN in joint locations')

    foot_locations = torch.stack((calcaneus_transform_r[..., :3, 3], mtp_offset_transform_r[..., :3, 3],
                                  calcaneus_transform_l[..., :3, 3], mtp_offset_transform_l[..., :3, 3]))

    segment_orientations = torch.stack((pelvis_transform[..., :3, :3],
                                        femur_transform_r[..., :3, :3], tibia_transform_r[..., :3, :3],
                                        talus_transform_r[..., :3, :3], calcaneus_transform_r[..., :3, :3],
                                        femur_transform_l[..., :3, :3], tibia_transform_l[..., :3, :3],
                                        talus_transform_l[..., :3, :3], calcaneus_transform_l[..., :3, :3],
                                        torso_transform[..., :3, :3]))
    if with_arm:
        segment_orientations = torch.stack((*[segment_orientations[i] for i in range(segment_orientations.shape[0])],
                                            humerus_transform_r[..., :3, :3], ulna_transform_r[..., :3, :3],
                                            radius_transform_r[..., :3, :3],
                                            humerus_transform_l[..., :3, :3], ulna_transform_l[..., :3, :3],
                                            radius_transform_l[..., :3, :3]))

    joint_names = ['pelvis', 'hip_r', 'knee_r', 'ankle_r', 'calcn_r', 'mtp_r',
                   'hip_l', 'knee_l', 'ankle_l', 'calcn_l', 'mtp_l']
    return foot_locations, joint_locations, joint_names, segment_orientations


def get_model_offsets(skeleton, with_arm=False):
    pelvis = skeleton.getBodyNode(0)
    hip_r_joint = pelvis.getChildJoint(0)
    femur_r = hip_r_joint.getChildBodyNode()
    knee_r_joint = femur_r.getChildJoint(0)
    tibia_r = knee_r_joint.getChildBodyNode()
    ankle_r_joint = tibia_r.getChildJoint(0)
    talus_r = ankle_r_joint.getChildBodyNode()
    subtalar_r_joint = talus_r.getChildJoint(0)
    calcn_r = subtalar_r_joint.getChildBodyNode()
    mtp_r_joint = calcn_r.getChildJoint(0)

    hip_l_joint = pelvis.getChildJoint(1)
    femur_l = hip_l_joint.getChildBodyNode()
    knee_l_joint = femur_l.getChildJoint(0)
    tibia_l = knee_l_joint.getChildBodyNode()
    ankle_l_joint = tibia_l.getChildJoint(0)
    talus_l = ankle_l_joint.getChildBodyNode()
    subtalar_l_joint = talus_l.getChildJoint(0)
    calcn_l = subtalar_l_joint.getChildBodyNode()
    mtp_l_joint = calcn_l.getChildJoint(0)

    lumbar_joint = pelvis.getChildJoint(2)
    torso = lumbar_joint.getChildBodyNode()

    if with_arm:
        shoulder_r_joint = torso.getChildJoint(0)
        humerus_r = shoulder_r_joint.getChildBodyNode()
        elbow_r_joint = humerus_r.getChildJoint(0)
        ulna_r = elbow_r_joint.getChildBodyNode()
        radioulnar_r_joint = ulna_r.getChildJoint(0)
        radius_r = radioulnar_r_joint.getChildBodyNode()
        wrist_r_joint = radius_r.getChildJoint(0)
        hand_r = wrist_r_joint.getChildBodyNode()

        shoulder_l_joint = torso.getChildJoint(1)
        humerus_l = shoulder_l_joint.getChildBodyNode()
        elbow_l_joint = humerus_l.getChildJoint(0)
        ulna_l = elbow_l_joint.getChildBodyNode()
        radioulnar_l_joint = ulna_l.getChildJoint(0)
        radius_l = radioulnar_l_joint.getChildBodyNode()
        wrist_l_joint = radius_l.getChildJoint(0)
        hand_l = wrist_l_joint.getChildBodyNode()

    # hip offset
    hip_offset_r = torch.eye(4)
    hip_offset_r[:3, :3] = torch.tensor(hip_r_joint.getTransformFromParentBodyNode().rotation())
    hip_offset_r[:3, 3] = torch.tensor(hip_r_joint.getTransformFromParentBodyNode().translation())
    hip_offset_l = torch.eye(4)
    hip_offset_l[:3, :3] = torch.tensor(hip_l_joint.getTransformFromParentBodyNode().rotation())
    hip_offset_l[:3, 3] = torch.tensor(hip_l_joint.getTransformFromParentBodyNode().translation())

    # femur offset
    femur_offset_to_knee_in_femur_r = -torch.tensor(hip_r_joint.getTransformFromChildBodyNode().translation())
    femur_rotation_to_knee_in_femur_r = torch.inverse(torch.tensor(hip_r_joint.getTransformFromChildBodyNode().rotation()))
    femur_offset_rotation_r = torch.eye(4)
    femur_offset_rotation_r[:3, :3] = femur_rotation_to_knee_in_femur_r
    femur_offset_translation_r = torch.eye(4)
    femur_offset_translation_r[:3, 3] = femur_offset_to_knee_in_femur_r
    femur_offset_r = torch.matmul(femur_offset_rotation_r, femur_offset_translation_r)

    femur_offset_to_knee_in_femur_l = -torch.tensor(hip_l_joint.getTransformFromChildBodyNode().translation())
    femur_rotation_to_knee_in_femur_l = torch.inverse(torch.tensor(hip_l_joint.getTransformFromChildBodyNode().rotation()))
    femur_offset_rotation_l = torch.eye(4)
    femur_offset_rotation_l[:3, :3] = femur_rotation_to_knee_in_femur_l
    femur_offset_translation_l = torch.eye(4)
    femur_offset_translation_l[:3, 3] = femur_offset_to_knee_in_femur_l
    femur_offset_l = torch.matmul(femur_offset_rotation_l, femur_offset_translation_l)

    # knee offset
    knee_offset_r = torch.eye(4)
    knee_offset_r[:3, :3] = torch.tensor(knee_r_joint.getTransformFromParentBodyNode().rotation())
    knee_offset_r[:3, 3] = torch.tensor(knee_r_joint.getTransformFromParentBodyNode().translation())

    knee_offset_l = torch.eye(4)
    knee_offset_l[:3, :3] = torch.tensor(knee_l_joint.getTransformFromParentBodyNode().rotation())
    knee_offset_l[:3, 3] = torch.tensor(knee_l_joint.getTransformFromParentBodyNode().translation())

    # tibia offset
    tibia_offset_to_knee_in_tibia_r = -torch.tensor(knee_r_joint.getTransformFromChildBodyNode().translation())
    tibia_rotation_to_knee_in_tibia_r = torch.inverse(torch.tensor(knee_r_joint.getTransformFromChildBodyNode().rotation()))
    tibia_offset_rotation_r = torch.eye(4)
    tibia_offset_rotation_r[:3, :3] = tibia_rotation_to_knee_in_tibia_r
    tibia_offset_translation_r = torch.eye(4)
    tibia_offset_translation_r[:3, 3] = tibia_offset_to_knee_in_tibia_r
    tibia_offset_r = torch.matmul(tibia_offset_rotation_r, tibia_offset_translation_r)

    tibia_offset_to_knee_in_tibia_l = -torch.tensor(knee_l_joint.getTransformFromChildBodyNode().translation())
    tibia_rotation_to_knee_in_tibia_l = torch.inverse(torch.tensor(knee_l_joint.getTransformFromChildBodyNode().rotation()))
    tibia_offset_rotation_l = torch.eye(4)
    tibia_offset_rotation_l[:3, :3] = tibia_rotation_to_knee_in_tibia_l
    tibia_offset_translation_l = torch.eye(4)
    tibia_offset_translation_l[:3, 3] = tibia_offset_to_knee_in_tibia_l
    tibia_offset_l = torch.matmul(tibia_offset_rotation_l, tibia_offset_translation_l)

    # ankle offset
    ankle_offset_r = torch.eye(4)
    ankle_offset_r[:3,:3] = torch.tensor(ankle_r_joint.getTransformFromParentBodyNode().rotation())
    ankle_offset_r[:3, 3] = torch.tensor(ankle_r_joint.getTransformFromParentBodyNode().translation())

    ankle_offset_l = torch.eye(4)
    ankle_offset_l[:3, :3] = torch.tensor(ankle_l_joint.getTransformFromParentBodyNode().rotation())
    ankle_offset_l[:3, 3] = torch.tensor(ankle_l_joint.getTransformFromParentBodyNode().translation())

    # talus offset
    talus_offset_to_ankle_in_talus_r = -torch.tensor(ankle_r_joint.getTransformFromChildBodyNode().translation())
    talus_rotation_to_ankle_in_talus_r = torch.inverse(torch.tensor(ankle_r_joint.getTransformFromChildBodyNode().rotation()))
    talus_offset_rotation_r = torch.eye(4)
    talus_offset_rotation_r[:3, :3] = talus_rotation_to_ankle_in_talus_r
    talus_offset_translation_r = torch.eye(4)
    talus_offset_translation_r[:3, 3] = talus_offset_to_ankle_in_talus_r
    talus_offset_r = torch.matmul(talus_offset_rotation_r, talus_offset_translation_r)

    talus_offset_to_ankle_in_talus_l = -torch.tensor(ankle_l_joint.getTransformFromChildBodyNode().translation())
    talus_rotation_to_ankle_in_talus_l = torch.inverse(torch.tensor(ankle_l_joint.getTransformFromChildBodyNode().rotation()))
    talus_offset_rotation_l = torch.eye(4)
    talus_offset_rotation_l[:3, :3] = talus_rotation_to_ankle_in_talus_l
    talus_offset_translation_l = torch.eye(4)
    talus_offset_translation_l[:3, 3] = talus_offset_to_ankle_in_talus_l
    talus_offset_l = torch.matmul(talus_offset_rotation_l, talus_offset_translation_l)

    # subtalar offset
    subtalar_offset_r = torch.eye(4)
    subtalar_offset_r[:3,:3] = torch.tensor(subtalar_r_joint.getTransformFromParentBodyNode().rotation())
    subtalar_offset_r[:3, 3] = torch.tensor(subtalar_r_joint.getTransformFromParentBodyNode().translation())

    subtalar_offset_l = torch.eye(4)
    subtalar_offset_l[:3, :3] = torch.tensor(subtalar_l_joint.getTransformFromParentBodyNode().rotation())
    subtalar_offset_l[:3, 3] = torch.tensor(subtalar_l_joint.getTransformFromParentBodyNode().translation())

    # calcaneus offset
    calcaneus_offset_to_subtalar_in_calcaneus_r = -torch.tensor(subtalar_r_joint.getTransformFromChildBodyNode().translation())
    calcaneus_rotation_to_subtalar_in_calcaneus_r = torch.inverse(torch.tensor(subtalar_r_joint.getTransformFromChildBodyNode().rotation()))
    calcaneus_offset_rotation_r = torch.eye(4)
    calcaneus_offset_rotation_r[:3, :3] = calcaneus_rotation_to_subtalar_in_calcaneus_r
    calcaneus_offset_translation_r = torch.eye(4)
    calcaneus_offset_translation_r[:3, 3] = calcaneus_offset_to_subtalar_in_calcaneus_r
    calcaneus_offset_r = torch.matmul(calcaneus_offset_rotation_r, calcaneus_offset_translation_r)

    calcaneus_offset_to_subtalar_in_calcaneus_l = -torch.tensor(subtalar_l_joint.getTransformFromChildBodyNode().translation())
    calcaneus_rotation_to_subtalar_in_calcaneus_l = torch.inverse(torch.tensor(subtalar_l_joint.getTransformFromChildBodyNode().rotation()))
    calcaneus_offset_rotation_l = torch.eye(4)
    calcaneus_offset_rotation_l[:3, :3] = calcaneus_rotation_to_subtalar_in_calcaneus_l
    calcaneus_offset_translation_l = torch.eye(4)
    calcaneus_offset_translation_l[:3, 3] = calcaneus_offset_to_subtalar_in_calcaneus_l
    calcaneus_offset_l = torch.matmul(calcaneus_offset_rotation_l, calcaneus_offset_translation_l)

    # mtp offset
    mtp_offset_r = torch.eye(4)
    mtp_offset_r[:3, :3] = torch.tensor(mtp_r_joint.getTransformFromParentBodyNode().rotation())
    mtp_offset_r[:3, 3] = torch.tensor(mtp_r_joint.getTransformFromParentBodyNode().translation())

    mtp_offset_l = torch.eye(4)
    mtp_offset_l[:3, :3] = torch.tensor(mtp_l_joint.getTransformFromParentBodyNode().rotation())
    mtp_offset_l[:3, 3] = torch.tensor(mtp_l_joint.getTransformFromParentBodyNode().translation())

    # toes offset
    toes_offset_to_mtp_in_toes_r = -torch.tensor(mtp_r_joint.getTransformFromChildBodyNode().translation())
    toes_rotation_to_mtp_in_toes_r = torch.inverse(torch.tensor(mtp_r_joint.getTransformFromChildBodyNode().rotation()))
    toes_offset_rotation_r = torch.eye(4)
    toes_offset_rotation_r[:3, :3] = toes_rotation_to_mtp_in_toes_r
    toes_offset_translation_r = torch.eye(4)
    toes_offset_translation_r[:3, 3] = toes_offset_to_mtp_in_toes_r
    toes_offset_r = torch.matmul(toes_offset_rotation_r, toes_offset_translation_r)

    toes_offset_to_mtp_in_toes_l = -torch.tensor(mtp_l_joint.getTransformFromChildBodyNode().translation())
    toes_rotation_to_mtp_in_toes_l = torch.inverse(torch.tensor(mtp_l_joint.getTransformFromChildBodyNode().rotation()))
    toes_offset_rotation_l = torch.eye(4)
    toes_offset_rotation_l[:3, :3] = toes_rotation_to_mtp_in_toes_l
    toes_offset_translation_l = torch.eye(4)
    toes_offset_translation_l[:3, 3] = toes_offset_to_mtp_in_toes_l
    toes_offset_l = torch.matmul(toes_offset_rotation_l, toes_offset_translation_l)


    # lumbar offset
    lumbar_offset = torch.eye(4)
    lumbar_offset[:3, :3] = torch.tensor(lumbar_joint.getTransformFromParentBodyNode().rotation())
    lumbar_offset[:3, 3] = torch.tensor(lumbar_joint.getTransformFromParentBodyNode().translation())

    # torso offset
    torso_offset_to_lumbar_in_torso = -torch.tensor(lumbar_joint.getTransformFromChildBodyNode().translation())
    torso_offset_rotation_to_lumbar_in_torso = torch.inverse(torch.tensor(lumbar_joint.getTransformFromChildBodyNode().rotation()))
    torso_offset_rotation = torch.eye(4)
    torso_offset_rotation[:3, :3] = torso_offset_rotation_to_lumbar_in_torso
    torso_offset_translation = torch.eye(4)
    torso_offset_translation[:3, 3] = torso_offset_to_lumbar_in_torso
    torso_offset = torch.matmul(torso_offset_rotation, torso_offset_translation)

    if with_arm:
        # shoulder offset
        shoulder_offset_r = torch.eye(4)
        shoulder_offset_r[:3, :3] = torch.tensor(shoulder_r_joint.getTransformFromParentBodyNode().rotation())
        shoulder_offset_r[:3, 3] = torch.tensor(shoulder_r_joint.getTransformFromParentBodyNode().translation())

        shoulder_offset_l = torch.eye(4)
        shoulder_offset_l[:3, :3] = torch.tensor(shoulder_l_joint.getTransformFromParentBodyNode().rotation())
        shoulder_offset_l[:3, 3] = torch.tensor(shoulder_l_joint.getTransformFromParentBodyNode().translation())

        # humerus offset
        humerus_offset_to_shoulder_in_humerus_r = -torch.tensor(shoulder_r_joint.getTransformFromChildBodyNode().translation())
        humerus_offset_rotation_to_shoulder_in_humerus_r = torch.inverse(torch.tensor(shoulder_r_joint.getTransformFromChildBodyNode().rotation()))
        humerus_offset_rotation_r = torch.eye(4)
        humerus_offset_rotation_r[:3, :3] = humerus_offset_rotation_to_shoulder_in_humerus_r
        humerus_offset_translation_r = torch.eye(4)
        humerus_offset_translation_r[:3, 3] = humerus_offset_to_shoulder_in_humerus_r
        humerus_offset_r = torch.matmul(humerus_offset_rotation_r, humerus_offset_translation_r)

        humerus_offset_to_shoulder_in_humerus_l = -torch.tensor(shoulder_l_joint.getTransformFromChildBodyNode().translation())
        humerus_offset_rotation_to_shoulder_in_humerus_l = torch.inverse(torch.tensor(shoulder_l_joint.getTransformFromChildBodyNode().rotation()))
        humerus_offset_rotation_l = torch.eye(4)
        humerus_offset_rotation_l[:3, :3] = humerus_offset_rotation_to_shoulder_in_humerus_l
        humerus_offset_translation_l = torch.eye(4)
        humerus_offset_translation_l[:3, 3] = humerus_offset_to_shoulder_in_humerus_l
        humerus_offset_l = torch.matmul(humerus_offset_rotation_l, humerus_offset_translation_l)

        # elbow offset
        elbow_offset_r = torch.eye(4)
        elbow_offset_r[:3, :3] = torch.tensor(elbow_r_joint.getTransformFromParentBodyNode().rotation())
        elbow_offset_r[:3, 3] = torch.tensor(elbow_r_joint.getTransformFromParentBodyNode().translation())

        elbow_offset_l = torch.eye(4)
        elbow_offset_l[:3, :3] = torch.tensor(elbow_l_joint.getTransformFromParentBodyNode().rotation())
        elbow_offset_l[:3, 3] = torch.tensor(elbow_l_joint.getTransformFromParentBodyNode().translation())

        # ulna offset
        ulna_offset_to_elbow_in_ulna_r = -torch.tensor(elbow_r_joint.getTransformFromChildBodyNode().translation())
        ulna_offset_rotation_to_elbow_in_ulna_r = torch.inverse(torch.tensor(elbow_r_joint.getTransformFromChildBodyNode().rotation()))
        ulna_offset_rotation_r = torch.eye(4)
        ulna_offset_rotation_r[:3, :3] = ulna_offset_rotation_to_elbow_in_ulna_r
        ulna_offset_translation_r = torch.eye(4)
        ulna_offset_translation_r[:3, 3] = ulna_offset_to_elbow_in_ulna_r
        ulna_offset_r = torch.matmul(ulna_offset_rotation_r, ulna_offset_translation_r)

        ulna_offset_to_elbow_in_ulna_l = -torch.tensor(elbow_l_joint.getTransformFromChildBodyNode().translation())
        ulna_offset_rotation_to_elbow_in_ulna_l = torch.inverse(torch.tensor(elbow_l_joint.getTransformFromChildBodyNode().rotation()))
        ulna_offset_rotation_l = torch.eye(4)
        ulna_offset_rotation_l[:3, :3] = ulna_offset_rotation_to_elbow_in_ulna_l
        ulna_offset_translation_l = torch.eye(4)
        ulna_offset_translation_l[:3, 3] = ulna_offset_to_elbow_in_ulna_l
        ulna_offset_l = torch.matmul(ulna_offset_rotation_l, ulna_offset_translation_l)

        # radioulnar offset
        radioulnar_offset_r = torch.eye(4)
        radioulnar_offset_r[:3, :3] = torch.tensor(radioulnar_r_joint.getTransformFromParentBodyNode().rotation())
        radioulnar_offset_r[:3, 3] = torch.tensor(radioulnar_r_joint.getTransformFromParentBodyNode().translation())

        radioulnar_offset_l = torch.eye(4)
        radioulnar_offset_l[:3, :3] = torch.tensor(radioulnar_l_joint.getTransformFromParentBodyNode().rotation())
        radioulnar_offset_l[:3, 3] = torch.tensor(radioulnar_l_joint.getTransformFromParentBodyNode().translation())

        # radius offset
        radius_offset_to_radioulnar_in_radius_r = -torch.tensor(radioulnar_r_joint.getTransformFromChildBodyNode().translation())
        radius_offset_rotation_to_radioulnar_in_radius_r = torch.inverse(torch.tensor(radioulnar_r_joint.getTransformFromChildBodyNode().rotation()))
        radius_offset_rotation_r = torch.eye(4)
        radius_offset_rotation_r[:3, :3] = radius_offset_rotation_to_radioulnar_in_radius_r
        radius_offset_translation_r = torch.eye(4)
        radius_offset_translation_r[:3, 3] = radius_offset_to_radioulnar_in_radius_r
        radius_offset_r = torch.matmul(radius_offset_rotation_r, radius_offset_translation_r)

        radius_offset_to_radioulnar_in_radius_l = -torch.tensor(radioulnar_l_joint.getTransformFromChildBodyNode().translation())
        radius_offset_rotation_to_radioulnar_in_radius_l = torch.inverse(torch.tensor(radioulnar_l_joint.getTransformFromChildBodyNode().rotation()))
        radius_offset_rotation_l = torch.eye(4)
        radius_offset_rotation_l[:3, :3] = radius_offset_rotation_to_radioulnar_in_radius_l
        radius_offset_translation_l = torch.eye(4)
        radius_offset_translation_l[:3, 3] = radius_offset_to_radioulnar_in_radius_l
        radius_offset_l = torch.matmul(radius_offset_rotation_l, radius_offset_translation_l)

        # wrist offset
        wrist_offset_r = torch.eye(4)
        wrist_offset_r[:3, :3] = torch.tensor(wrist_r_joint.getTransformFromParentBodyNode().rotation())
        wrist_offset_r[:3, 3] = torch.tensor(wrist_r_joint.getTransformFromParentBodyNode().translation())

        wrist_offset_l = torch.eye(4)
        wrist_offset_l[:3, :3] = torch.tensor(wrist_l_joint.getTransformFromParentBodyNode().rotation())
        wrist_offset_l[:3, 3] = torch.tensor(wrist_l_joint.getTransformFromParentBodyNode().translation())

        # hand offset
        hand_offset_to_wrist_in_hand_r = -torch.tensor(wrist_r_joint.getTransformFromChildBodyNode().translation())
        hand_offset_rotation_to_wrist_in_hand_r = torch.inverse(torch.tensor(wrist_r_joint.getTransformFromChildBodyNode().rotation()))
        hand_offset_rotation_r = torch.eye(4)
        hand_offset_rotation_r[:3, :3] = hand_offset_rotation_to_wrist_in_hand_r
        hand_offset_translation_r = torch.eye(4)
        hand_offset_translation_r[:3, 3] = hand_offset_to_wrist_in_hand_r
        hand_offset_r = torch.matmul(hand_offset_rotation_r, hand_offset_translation_r)

        hand_offset_to_wrist_in_hand_l = -torch.tensor(wrist_l_joint.getTransformFromChildBodyNode().translation())
        hand_offset_rotation_to_wrist_in_hand_l = torch.inverse(torch.tensor(wrist_l_joint.getTransformFromChildBodyNode().rotation()))
        hand_offset_rotation_l = torch.eye(4)
        hand_offset_rotation_l[:3, :3] = hand_offset_rotation_to_wrist_in_hand_l
        hand_offset_translation_l = torch.eye(4)
        hand_offset_translation_l[:3, 3] = hand_offset_to_wrist_in_hand_l
        hand_offset_l = torch.matmul(hand_offset_rotation_l, hand_offset_translation_l)

    offsets = torch.stack((hip_offset_r, femur_offset_r, knee_offset_r, tibia_offset_r,
                           ankle_offset_r, talus_offset_r, subtalar_offset_r,
                           calcaneus_offset_r, mtp_offset_r,
                           hip_offset_l, femur_offset_l, knee_offset_l, tibia_offset_l,
                           ankle_offset_l, talus_offset_l, subtalar_offset_l,
                           calcaneus_offset_l, mtp_offset_l,
                           lumbar_offset, torso_offset), dim=2)
    if with_arm:
        offsets = torch.stack((*[offsets[i] for i in range(offsets.shape[0])],
                               shoulder_offset_r, humerus_offset_r, elbow_offset_r, ulna_offset_r,
                               radioulnar_offset_r, radius_offset_r, wrist_offset_r, hand_offset_r,
                               shoulder_offset_l, humerus_offset_l, elbow_offset_l, ulna_offset_l,
                               radioulnar_offset_l, radius_offset_l, wrist_offset_l, hand_offset_l,toes_offset_r, toes_offset_l), dim=2)
    return offsets


def euler_to_angular_velocity(q, sampling_fre, convention="XYZ"):
    assert q.shape[-1] == 3
    mat = euler_angles_to_matrix(q, convention)
    mat_diff = mat.clone()
    mat_diff[:-1, :, :] = mat_diff[:-1, :, :] - mat_diff[1:, :, :]
    mat_diff[-1, :, :] = mat_diff[-2, :, :]

    angular_velocity = torch.stack([mat_diff[:, 2, 1] - mat_diff[:, 1, 2],
                                    mat_diff[:, 0, 2] - mat_diff[:, 2, 0],
                                    mat_diff[:, 1, 0] - mat_diff[:, 0, 1]], dim=1) * sampling_fre * 0.5
    return angular_velocity


def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor:
    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
        raise ValueError("Invalid input euler angles.")
    if len(convention) != 3:
        raise ValueError("Convention must have 3 letters.")
    if convention[1] in (convention[0], convention[2]):
        raise ValueError(f"Invalid convention {convention}.")
    for letter in convention:
        if letter not in ("X", "Y", "Z"):
            raise ValueError(f"Invalid letter {letter} in convention string.")
    matrices = [
        _axis_angle_rotation(c, e)
        for c, e in zip(convention, torch.unbind(euler_angles, -1))
    ]
    return torch.matmul(torch.matmul(matrices[0], matrices[1]), matrices[2])


def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
    cos = torch.cos(angle)
    sin = torch.sin(angle)
    one = torch.ones_like(angle)
    zero = torch.zeros_like(angle)

    if axis == "X":
        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
    elif axis == "Y":
        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
    elif axis == "Z":
        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
    else:
        raise ValueError("letter must be either X, Y or Z.")

    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))


def matrix_to_euler_angles(matrix: torch.Tensor, convention: str) -> torch.Tensor:
    if len(convention) != 3:
        raise ValueError("Convention must have 3 letters.")
    if convention[1] in (convention[0], convention[2]):
        raise ValueError(f"Invalid convention {convention}.")
    for letter in convention:
        if letter not in ("X", "Y", "Z"):
            raise ValueError(f"Invalid letter {letter} in convention string.")
    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
    i0 = _index_from_letter(convention[0])
    i2 = _index_from_letter(convention[2])
    tait_bryan = i0 != i2
    if tait_bryan:
        central_angle = torch.asin(
            matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)
        )
    else:
        central_angle = torch.acos(matrix[..., i0, i0])

    o = (
        _angle_from_tan(
            convention[0], convention[1], matrix[..., i2], False, tait_bryan
        ),
        central_angle,
        _angle_from_tan(
            convention[2], convention[1], matrix[..., i0, :], True, tait_bryan
        ),
    )
    return torch.stack(o, -1)


def _angle_from_tan(
        axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool
) -> torch.Tensor:

    i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
    if horizontal:
        i2, i1 = i1, i2
    even = (axis + other_axis) in ["XY", "YZ", "ZX"]
    if horizontal == even:
        return torch.atan2(data[..., i1], data[..., i2])
    if tait_bryan:
        return torch.atan2(-data[..., i2], data[..., i1])
    return torch.atan2(data[..., i2], -data[..., i1])


def _index_from_letter(letter: str) -> int:
    if letter == "X":
        return 0
    if letter == "Y":
        return 1
    if letter == "Z":
        return 2
    raise ValueError("letter must be either X, Y or Z.")


def get_multi_body_loc_using_nimble_by_body_names(body_names, skel, poses):
    body_ids = [skel.getBodyNode(name) for name in body_names]
    return get_multi_body_loc_using_nimble_by_body_nodes(body_ids, skel, poses)


def get_multi_body_loc_using_nimble_by_body_nodes(body_nodes, skel, poses):
    body_loc = []
    for i_frame in range(len(poses)):
        skel.setPositions(poses[i_frame])
        body_loc.append(np.concatenate([body_node.getWorldTransform().translation() for body_node in body_nodes]))
    body_loc = np.array(body_loc)
    return body_loc


def norm_cops(skel, states, opt, sub_mass, height_m, check_cop_to_calcn_distance, wrong_cop_ratio=0.002):
    grf_thd_to_zero_vector = 20
    states = torch.from_numpy(states)
    poses  = states[:, opt.kinematic_osim_col_loc]
    forces = states[:, opt.grf_osim_col_loc]
    cops   = states[:, opt.cop_osim_col_loc]
    heel_centers = torch.from_numpy(get_multi_body_loc_using_nimble_by_body_names(('calcn_r', 'calcn_l'), skel, poses)).to(states.dtype)
    for i_plate in range(2):
        force_vector = forces[:, 3*i_plate:3*(i_plate+1)]
        vector = (cops - heel_centers)[:, 3*i_plate:3*(i_plate+1)]
        stance_phase = force_vector[:, 1] > (50 / sub_mass)
        if check_cop_to_calcn_distance:
            if ((vector[stance_phase, :].abs() > 0.3).sum() / 3) > wrong_cop_ratio * len(stance_phase):
                return False
        normed_vector = vector * force_vector[:, 1:2] / height_m
        if grf_thd_to_zero_vector:
            normed_vector[force_vector[:, 1] * sub_mass < grf_thd_to_zero_vector] = 0
        states[:, opt.cop_osim_col_loc[3*i_plate:3*(i_plate+1)]] = normed_vector.to(states.dtype)
    return states


def inverse_norm_cops(skel, states, opt, sub_mass, height_m, grf_thd_to_zero_cop=20):
    poses = states[:, opt.kinematic_osim_col_loc]
    forces = states[:, opt.grf_osim_col_loc]
    normed_cops = states[:, opt.cop_osim_col_loc]
    if len(skel.getDofs()) != poses.shape[1]:        # With Arm model
        print('With Arm model is used. Adding 6 zeros to the end of the poses.')
        poses = np.concatenate([poses, np.zeros((poses.shape[0], 10))], axis=-1)
    foot_loc = get_multi_body_loc_using_nimble_by_body_names(('calcn_r', 'calcn_l'), skel, poses)

    for i_plate in range(2):
        force_v = forces[:, 3*i_plate:3*(i_plate+1)]
        force_v[force_v == 0] = 1e-6
        vector = normed_cops[:, 3 * i_plate:3 * (i_plate + 1)] / force_v[:, 1:2] * height_m
        vector = np.nan_to_num(vector, posinf=0, neginf=0)
        # vector.clip(min=-0.4, max=0.4, out=vector)      # CoP should be within 0.4 m from the foot
        cops = vector + foot_loc[:, 3*i_plate:3*(i_plate+1)]

        if grf_thd_to_zero_cop:
            cops[force_v[:, 1] * sub_mass < grf_thd_to_zero_cop] = 0

        if isinstance(states, torch.Tensor):
            cops = torch.from_numpy(cops).to(states.dtype)
        else:
            cops = cops.astype(states.dtype)
        states[:, opt.cop_osim_col_loc[3*i_plate:3*(i_plate+1)]] = cops
    return states


def inverse_align_moving_direction(results_pred, column_names, rot_mat):
    if isinstance(results_pred, np.ndarray):
        poses = torch.from_numpy(results_pred)
    results_pred_clone = poses.clone().float()
    pelvis_orientation_col_loc = [column_names.index(col) for col in JOINTS_3D_ALL['pelvis']]
    p_pos_col_loc = [column_names.index(col) for col in [f'pelvis_t{x}' for x in ['x', 'y', 'z']]]
    r_grf_col_loc = [column_names.index(col) for col in ['calcn_r_force_vx', 'calcn_r_force_vy', 'calcn_r_force_vz']]
    l_grf_col_loc = [column_names.index(col) for col in ['calcn_l_force_vx', 'calcn_l_force_vy', 'calcn_l_force_vz']]
    r_cop_col_loc = [column_names.index(col) for col in [f'calcn_r_force_normed_cop_{x}' for x in ['x', 'y', 'z']]]
    l_cop_col_loc = [column_names.index(col) for col in [f'calcn_l_force_normed_cop_{x}' for x in ['x', 'y', 'z']]]

    if len(pelvis_orientation_col_loc) != 3 or len(p_pos_col_loc) != 3 or len(r_grf_col_loc) != 3 or len(
            l_grf_col_loc) != 3:
        raise ValueError('check column names')

    pelvis_orientation = results_pred_clone[:, pelvis_orientation_col_loc]
    pelvis_orientation = euler_angles_to_matrix(pelvis_orientation, "ZXY")
    p_pos = results_pred_clone[:, p_pos_col_loc]
    r_grf = results_pred_clone[:, r_grf_col_loc]
    l_grf = results_pred_clone[:, l_grf_col_loc]
    r_cop = results_pred_clone[:, r_cop_col_loc]
    l_cop = results_pred_clone[:, l_cop_col_loc]

    rot_mat = rot_mat.T

    pelvis_orientation_rotated = torch.matmul(rot_mat, pelvis_orientation)
    p_pos_rotated = torch.matmul(rot_mat, p_pos.unsqueeze(2)).squeeze(2)
    r_grf_rotated = torch.matmul(rot_mat, r_grf.unsqueeze(2)).squeeze(2)
    l_grf_rotated = torch.matmul(rot_mat, l_grf.unsqueeze(2)).squeeze(2)
    r_cop_rotated = torch.matmul(rot_mat, r_cop.unsqueeze(2)).squeeze(2)
    l_cop_rotated = torch.matmul(rot_mat, l_cop.unsqueeze(2)).squeeze(2)

    results_pred_clone[:, pelvis_orientation_col_loc] = matrix_to_euler_angles(pelvis_orientation_rotated.float(), "ZXY")
    results_pred_clone[:, p_pos_col_loc] = p_pos_rotated.float()
    results_pred_clone[:, r_grf_col_loc] = r_grf_rotated.float()
    results_pred_clone[:, l_grf_col_loc] = l_grf_rotated.float()
    results_pred_clone[:, r_cop_col_loc] = r_cop_rotated.float()
    results_pred_clone[:, l_cop_col_loc] = l_cop_rotated.float()
    return results_pred_clone


def convert_addb_state_to_model_input(pose_df, joints_3d, sampling_fre):
    # shift root position to start in (x,y) = (0,0)
    pos_vec = [pose_df['pelvis_tx'][0], pose_df['pelvis_ty'][0], pose_df['pelvis_tz'][0]]
    pose_df['pelvis_tx'] = pose_df['pelvis_tx'] - pose_df['pelvis_tx'][0]
    pose_df['pelvis_ty'] = pose_df['pelvis_ty'] - pose_df['pelvis_ty'][0]
    pose_df['pelvis_tz'] = pose_df['pelvis_tz'] - pose_df['pelvis_tz'][0]

    # remove frozen dof
    for frozen_col in FROZEN_DOFS:
        if frozen_col in pose_df.columns:
            pose_df = pose_df.drop(frozen_col, axis=1)

    # convert euler to 6v
    for joint_name, joints_with_3_dof in joints_3d.items():
        joint_6v = euler_to_6v(torch.tensor(pose_df[joints_with_3_dof].values), "ZXY").numpy()
        for i in range(6):
            pose_df[joint_name + '_' + str(i)] = joint_6v[:, i]

    for joint_name, joints_with_3_dof in joints_3d.items():
        joint_angular_v = euler_to_angular_velocity(torch.tensor(pose_df[joints_with_3_dof].values), sampling_fre, "ZXY").numpy()
        joint_angular_v = data_filter(joint_angular_v, 15, sampling_fre, 4)
        for joints_euler_name in joints_with_3_dof:
            pose_df = pose_df.drop(joints_euler_name, axis=1)
        for i, axis in enumerate(['x', 'y', 'z']):
            pose_df[joint_name + '_' + axis + '_angular' + '_vel'] = joint_angular_v[:, i]

    vel_col_loc = [i for i, col in enumerate(pose_df.columns) if not np.sum([term in col for term in ['force', 'pelvis_', '_vel', '_0', '_1', '_2', '_3', '_4', '_5']])]
    vel_col_names = [f'{col}_vel' for i, col in enumerate(pose_df.columns) if not np.sum([term in col for term in ['force', 'pelvis_', '_vel', '_0', '_1', '_2', '_3', '_4', '_5']])]
    kinematics_np = pose_df.iloc[:, vel_col_loc].to_numpy().copy()
    kinematics_np_filtered = data_filter(kinematics_np, 15, sampling_fre, 4)
    kinematics_vel = np.stack([spline_fitting_1d(kinematics_np_filtered[:, i_col], range(kinematics_np_filtered.shape[0]), 1).ravel()
                               for i_col in range(kinematics_np_filtered.shape[1])]).T

    pose_vel_df = pd.DataFrame(np.concatenate([pose_df.values, kinematics_vel], axis=1), columns=list(pose_df.columns)+vel_col_names)
    return pose_vel_df, pos_vec


def inverse_convert_addb_state_to_model_input(
        model_states, model_states_column_names, treadmill_speed, joints_3d, osim_dof_columns, pos_vec, height_m, sampling_fre=100):
    model_states_dict = {col: model_states[..., i] for i, col in enumerate(model_states_column_names) if
                         col in osim_dof_columns}

    for i_col, col in enumerate(['pelvis_tx', 'pelvis_ty', 'pelvis_tz']):
        model_states_dict[col] = model_states_dict[col] * height_m.unsqueeze(-1).expand(model_states_dict[col].shape)
        if col == 'pelvis_tx':
            model_states_dict[col] -= treadmill_speed
        model_states_dict[col] = torch.cumsum(model_states_dict[col], dim=-1) / sampling_fre

    # convert 6v to euler
    for joint_name, joints_with_3_dof in joints_3d.items():
        joint_name_6v = [joint_name + '_' + str(i) for i in range(6)]
        index_ = [model_states_column_names.index(joint_name_6v[i]) for i in range(6)]
        joint_euler = euler_from_6v(model_states[..., index_], "ZXY")

        for i, joints_euler_name in enumerate(joints_with_3_dof):
            model_states_dict[joints_euler_name] = joint_euler[..., i]

    # add frozen dof back
    for frozen_col in FROZEN_DOFS:
        model_states_dict[frozen_col] = torch.zeros(model_states.shape[:len(model_states.shape)-1]).to(model_states.device)

    pos_vec = torch.tensor(pos_vec)
    pos_vec_torch = pos_vec.unsqueeze(-1).repeat(*[1 for _ in range(len(pos_vec.shape))], model_states.shape[-2]).to(model_states.device)
    for i_col, col in enumerate(['pelvis_tx', 'pelvis_ty', 'pelvis_tz']):
        model_states_dict[col] += pos_vec_torch[..., i_col, :]

    osim_states = torch.stack([model_states_dict[col] for col in osim_dof_columns], dim=len(model_states.shape)-1).float()
    return osim_states


def align_moving_direction(poses, column_names):
    if isinstance(poses, np.ndarray):
        poses = torch.from_numpy(poses)
    pose_clone = poses.clone().float()
    pelvis_orientation_col_loc = [column_names.index(col) for col in JOINTS_3D_ALL['pelvis']]
    p_pos_col_loc = [column_names.index(col) for col in [f'pelvis_t{x}' for x in ['x', 'y', 'z']]]
    r_grf_col_loc = [column_names.index(col) for col in ['calcn_r_force_vx', 'calcn_r_force_vy', 'calcn_r_force_vz']]
    l_grf_col_loc = [column_names.index(col) for col in ['calcn_l_force_vx', 'calcn_l_force_vy', 'calcn_l_force_vz']]
    r_cop_col_loc = [column_names.index(col) for col in [f'calcn_r_force_normed_cop_{x}' for x in ['x', 'y', 'z']]]
    l_cop_col_loc = [column_names.index(col) for col in [f'calcn_l_force_normed_cop_{x}' for x in ['x', 'y', 'z']]]

    if len(pelvis_orientation_col_loc) != 3 or len(p_pos_col_loc) != 3 or len(r_grf_col_loc) != 3 or len(
            l_grf_col_loc) != 3:
        raise ValueError('check column names')

    pelvis_orientation = pose_clone[:, pelvis_orientation_col_loc]
    pelvis_orientation = euler_angles_to_matrix(pelvis_orientation, "ZXY")
    p_pos = pose_clone[:, p_pos_col_loc]
    r_grf = pose_clone[:, r_grf_col_loc]
    l_grf = pose_clone[:, l_grf_col_loc]
    r_cop = pose_clone[:, r_cop_col_loc]
    l_cop = pose_clone[:, l_cop_col_loc]

    angles = np.arctan2(- pelvis_orientation[:, 0, 2], pelvis_orientation[:, 2, 2])
    if np.rad2deg(angles.max() - angles.min()) > 45:
        return False, None
    angle = angles.median()
    rot_mat = torch.tensor([[np.cos(angle), 0, np.sin(angle)], [0, 1, 0], [-np.sin(angle), 0, np.cos(angle)]]).float()

    pelvis_orientation_rotated = torch.matmul(rot_mat, pelvis_orientation)
    p_pos_rotated = torch.matmul(rot_mat, p_pos.unsqueeze(2)).squeeze(2)
    r_grf_rotated = torch.matmul(rot_mat, r_grf.unsqueeze(2)).squeeze(2)
    l_grf_rotated = torch.matmul(rot_mat, l_grf.unsqueeze(2)).squeeze(2)
    r_cop_rotated = torch.matmul(rot_mat, r_cop.unsqueeze(2)).squeeze(2)
    l_cop_rotated = torch.matmul(rot_mat, l_cop.unsqueeze(2)).squeeze(2)

    pose_clone[:, pelvis_orientation_col_loc] = matrix_to_euler_angles(pelvis_orientation_rotated.float(), "ZXY")
    pose_clone[:, p_pos_col_loc] = p_pos_rotated.float()
    pose_clone[:, r_grf_col_loc] = r_grf_rotated.float()
    pose_clone[:, l_grf_col_loc] = l_grf_rotated.float()
    pose_clone[:, r_cop_col_loc] = r_cop_rotated.float()
    pose_clone[:, l_cop_col_loc] = l_cop_rotated.float()

    return pose_clone, rot_mat


def data_filter(data, cut_off_fre, sampling_fre, filter_order=4):
    fre = cut_off_fre / (sampling_fre / 2)
    b, a = butter(filter_order, fre, 'lowpass')
    if len(data.shape) == 1:
        data_filtered = filtfilt(b, a, data)
    else:
        data_filtered = filtfilt(b, a, data, axis=0)
    return data_filtered


def fix_seed(seed):
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)
    np.random.seed(0)


def linear_resample_data(trial_data, original_fre, target_fre):
    x, step = np.linspace(0., 1., trial_data.shape[0], retstep=True)
    new_x = np.arange(0., 1., step * original_fre / target_fre)
    f = interp1d(x, trial_data, axis=0)
    trial_data_resampled = f(new_x)
    return trial_data_resampled


def spline_fitting_1d(data_, step_to_resample, der=0):
    assert len(data_.shape) == 1
    data_ = data_.reshape(1, -1)
    tck, step = interpo.splprep(data_, u=range(data_.shape[1]), s=0)
    data_resampled = interpo.splev(step_to_resample, tck, der=der)
    data_resampled = np.column_stack(data_resampled)
    return data_resampled


def convert_overlapped_list_to_array(trial_len, win_list, s_, e_, fun=np.nanmedian, max_size=150):
    array_val_expand = np.full((max_size, trial_len, win_list[0].shape[1]), np.nan)
    for i_win, (win, s, e) in enumerate(zip(win_list, s_, e_)):
        array_val_expand[i_win%max_size, s:e] = win[:e-s]
    array_val = fun(array_val_expand, axis=0)
    std_val = np.nanstd(array_val_expand, axis=0)
    return array_val, std_val


def identity(t, *args, **kwargs):
    return t


def extract(a, t, x_shape):
    b, *_ = t.shape
    out = a.gather(-1, t)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))


def make_beta_schedule(
        schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3
):
    if schedule == "linear":
        betas = (
                torch.linspace(
                    linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64
                )
                ** 2
        )

    elif schedule == "cosine":
        timesteps = (
                torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s
        )
        alphas = timesteps / (1 + cosine_s) * np.pi / 2
        alphas = torch.cos(alphas).pow(2)
        alphas = alphas / alphas[0]
        betas = 1 - alphas[1:] / alphas[:-1]
        betas = np.clip(betas, a_min=0, a_max=0.999)

    elif schedule == "sqrt_linear":
        betas = torch.linspace(
            linear_start, linear_end, n_timestep, dtype=torch.float64
        )
    elif schedule == "sqrt":
        betas = (
                torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
                ** 0.5
        )
    else:
        raise ValueError(f"schedule '{schedule}' unknown.")
    return betas.numpy()


""" ============================ End util.py ============================ """


""" ============================ Start args.py ============================ """


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument("--exp_name", default="new_orientation_alignment_diffusion", help="save to project/name")
    parser.add_argument("--with_arm", type=bool, default=False, help="whether osim model has arm DoFs")
    parser.add_argument("--with_kinematics_vel", type=bool, default=True, help="whether to include 1st derivative of kinematics")
    parser.add_argument("--epochs", type=int, default=7680)
    parser.add_argument("--target_sampling_rate", type=int, default=100)
    parser.add_argument("--window_len", type=int, default=150)
    parser.add_argument("--guide_x_start_the_beginning_step", type=int, default=-10)      # negative value means no guidance

    parser.add_argument("--project", default="runs/train", help="project/name")
    parser.add_argument(
        "--processed_data_dir",
        type=str,
        default="dataset_backups/",
        help="Dataset backup path",
    )

    parser.add_argument("--feature_type", type=str, default="jukebox")
    parser.add_argument("--batch_size", type=int, default=256, help="batch size")
    parser.add_argument("--batch_size_inference", type=int, default=5, help="batch size during inference")
    parser.add_argument(
        "--checkpoint", type=str, default="", help="trained checkpoint path (optional)"
    )
    # parser.add_argument(
    #     "--checkpoint_bl", type=str, default="", help="trained checkpoint path (optional)"
    # )
    opt = parser.parse_args(args=[])
    set_no_arm_opt(opt)
    current_folder = os.getcwd()
    opt.subject_data_path = current_folder
    opt.geometry_folder = current_folder + '/Geometry/'
    opt.checkpoint_bl = current_folder + '/GaitDynamicsRefinement.pt'
    return opt


def set_no_arm_opt(opt):
    opt.with_arm = False
    opt.osim_dof_columns = copy.deepcopy(OSIM_DOF_ALL[:23] + KINETICS_ALL)
    opt.joints_3d = {key_: value_ for key_, value_ in JOINTS_3D_ALL.items() if key_ in ['pelvis', 'hip_r', 'hip_l', 'lumbar']}
    opt.model_states_column_names = copy.deepcopy(MODEL_STATES_COLUMN_NAMES_NO_ARM)
    for joint_name, joints_with_3_dof in opt.joints_3d.items():
        opt.model_states_column_names = opt.model_states_column_names + [
            joint_name + '_' + axis + '_angular_vel' for axis in ['x', 'y', 'z']]

    if opt.with_kinematics_vel:
        opt.model_states_column_names = opt.model_states_column_names + [
            f'{col}_vel' for i_col, col in enumerate(opt.model_states_column_names)
            if not sum([term in col for term in ['force', 'pelvis_', '_vel', '_0', '_1', '_2', '_3', '_4', '_5']])]


    opt.knee_diffusion_col_loc = [i_col for i_col, col in enumerate(opt.model_states_column_names) if 'knee' in col]
    opt.ankle_diffusion_col_loc = [i_col for i_col, col in enumerate(opt.model_states_column_names) if 'ankle' in col]
    opt.hip_diffusion_col_loc = [i_col for i_col, col in enumerate(opt.model_states_column_names) if 'hip' in col]
    opt.kinematic_diffusion_col_loc = [i_col for i_col, col in enumerate(opt.model_states_column_names) if 'force' not in col]
    opt.kinetic_diffusion_col_loc = [i_col for i_col, col in enumerate(opt.model_states_column_names) if i_col not in opt.kinematic_diffusion_col_loc]
    opt.grf_osim_col_loc = [i_col for i_col, col in enumerate(opt.osim_dof_columns) if 'force' in col and '_cop_' not in col]
    opt.cop_osim_col_loc = [i_col for i_col, col in enumerate(opt.osim_dof_columns) if '_cop_' in col]
    opt.kinematic_osim_col_loc = [i_col for i_col, col in enumerate(opt.osim_dof_columns) if 'force' not in col]


""" ============================ Start dataset.py ============================ """


def linear_resample_data_as_num_of_dp(data_raw, target_dp):
    x = np.linspace(0., 1., data_raw.shape[0])
    new_x = np.linspace(0., 1., target_dp)
    return interp1d(x, data_raw, axis=0)(new_x)


class MotionDataset(Dataset):
    def __init__(self, opt, normalizer, cadence_scale):
        self.data_path = opt.subject_data_path
        self.subject_osim_model = opt.subject_osim_model
        if opt.target_sampling_rate != 100:
            raise ValueError('100 Hz sampling rate is not confirmed. Confirm by setting opt.target_sampling_rate = 100')
        self.target_sampling_rate = opt.target_sampling_rate
        self.window_len = opt.window_len
        self.opt = opt
        self.skel = None
        self.dset_set = set()
        self.cadence_scale = cadence_scale

        print("Loading dataset...")
        self.load_addb(opt)

        self.guess_vel_and_replace_txtytz()

        # Cadence manipulation in the diffusion-state (velocity) domain: re-time each trial to
        # n_target = round(len / cadence_scale) frames. Because pelvis translation is now a
        # velocity, re-timing changes the stride frequency (cadence) while preserving the
        # speed magnitude (the constant treadmill mean is untouched) — no scaling correction
        # needed. Only pelvis_tx conditions the model; every other (masked) column is
        # regenerated, so linear interpolation of the model rep is harmless here.
        if self.cadence_scale != 1.0:
            for trial in self.trials:
                n_target = int(round(trial.converted_pose.shape[0] / self.cadence_scale))
                resampled = linear_resample_data_as_num_of_dp(trial.converted_pose.numpy(), n_target)
                trial.converted_pose = torch.from_numpy(resampled).float()
                trial.length = n_target

        if not len(self.trials):
            print("No trials loaded")
            return
        self.normalizer = normalizer
        for i_trial in range(len(self.trials)):
            self.trials[i_trial].converted_pose = self.normalizer.normalize(self.trials[i_trial].converted_pose).clone().detach().float()

    def guess_vel_and_replace_txtytz(self):
        pelvis_pos_loc = [self.opt.model_states_column_names.index(col) for col in [f'pelvis_t{x}' for x in ['x', 'y', 'z']]]
        for i_trial, trial in enumerate(self.trials):
            body_center = trial.converted_pose[:, 0:3]
            body_center = data_filter(body_center, 10, self.target_sampling_rate).astype(np.float32)
            vel_from_t = np.diff(body_center, axis=0) * self.target_sampling_rate
            vel_from_t = np.concatenate([vel_from_t, vel_from_t[-1][None, :]], axis=0)
            if self.opt.treadmill_speed is None:
                raise ValueError('Treadmill speed is not set. For overground walking, set opt.treadmill_speed = 0')
            vel_from_t[:, 0] = vel_from_t[:, 0] + self.opt.treadmill_speed
            walking_vel = vel_from_t

            self.trials[i_trial].converted_pose[:, pelvis_pos_loc] = torch.from_numpy(walking_vel) / self.trials[i_trial].height_m

    def __len__(self):
        return self.opt.pseudo_dataset_len

    def get_one_win_from_the_beginning_of_each_trial(self, col_loc_to_unmask):
        windows = []
        cond = torch.ones([6])
        for i_trial, trial_ in enumerate(self.trials):
            len_ = min(self.opt.window_len, trial_.length)
            mask = torch.zeros([self.opt.window_len, len(self.opt.model_states_column_names)])
            mask[:len_, col_loc_to_unmask] = 1
            pose = torch.zeros([self.opt.window_len, len(self.opt.model_states_column_names)])
            pose[:len_] = trial_.converted_pose[:len_]
            windows.append(WindowData(pose, trial_.model_offsets, i_trial, None, mask, trial_.height_m,
                                      trial_.weight_kg, cond))
        return windows

    def get_overlapping_wins(self, col_loc_to_unmask, step_len, start_trial=0, end_trial=None, including_shorter_than_window_len=False):
        if end_trial is None:
            end_trial = len(self.trials)
        windows, s_list, e_list = [], [], []
        for i_trial in range(start_trial, end_trial):
            trial_ = self.trials[i_trial]
            col_loc_to_unmask_trial = copy.deepcopy(col_loc_to_unmask)
            for col in trial_.missing_col:
                col_index = self.opt.model_states_column_names.index(col)
                col_loc_to_unmask_trial.remove(col_index)
            trial_len = trial_.converted_pose.shape[0]
            if including_shorter_than_window_len:
                e_of_trial = trial_len
            else:
                e_of_trial = trial_len - self.opt.window_len + step_len
            for i in range(0, e_of_trial, step_len):
                s = max(0, i)
                e = min(trial_len, i + self.opt.window_len)
                s_list.append(s)
                e_list.append(e)
                mask = torch.zeros([self.opt.window_len, len(self.opt.model_states_column_names)])
                mask[:, col_loc_to_unmask_trial] = 1
                mask[e-s:, :] = 0
                data_ = torch.zeros([self.opt.window_len, len(self.opt.model_states_column_names)])
                data_[:e-s] = trial_.converted_pose[s:e, ...]
                windows.append(WindowData(data_, trial_.model_offsets, i_trial, None, mask,
                                          trial_.height_m, trial_.weight_kg, trial_.missing_col))
        return windows, s_list, e_list

    def load_addb(self, opt):
        file_paths = opt.file_paths
        customOsim: nimble.biomechanics.OpenSimFile = nimble.biomechanics.OpenSimParser.parseOsim(self.subject_osim_model, self.opt.geometry_folder)
        skel = customOsim.skeleton
        self.skel = skel

        self.trials, self.file_names, self.rot_mat_trials, self.time_column = [], [], [], []
        for i_file, file_path in enumerate(file_paths):
            model_offsets = get_model_offsets(skel).float()
            with open(file_path) as f:
                line_num = 0
                endheader_line, angle_scale = None, None
                while True:
                    header = f.readline()
                    line_num += 1
                    if 'endheader' in header:
                        endheader_line = line_num
                        break
                    if 'inDegrees' in header and angle_scale is None:
                        if 'yes' in header and 'no' not in header:
                            angle_scale = np.pi / 180
                        elif 'no' in header and 'yes' not in header:
                            angle_scale = 1
                        else:
                            raise ValueError('No inDegrees keyword in the header, cannot determine the unit of angles. '
                                             'Here is an example header: \nCoordinates\nversion=1\nnRows=1380'
                                             '\nnColumns=26\ninDegrees=yes\n')
                if endheader_line is None:
                    raise ValueError(f'No \'endheader\' line found in the header of {file_path}.')
            poses_df = pd.read_csv(file_path, sep='\t', skiprows=endheader_line)

            if 'time' not in poses_df.columns:
                raise ValueError(f'{file_path} does not have time column. Necessary for compuing sampling rate')
            sampling_rate = round((poses_df.shape[0] - 1) / (poses_df['time'].iloc[-1] - poses_df['time'].iloc[0]))
            self.time_column.append(poses_df['time'])

            missing_col = []
            for col in JOINTS_1D_ALL:
                if col not in poses_df.columns:
                    poses_df[col] = 0.
                    if col not in FROZEN_DOFS:
                        missing_col.append(col)
                        missing_col.append(col + '_vel')

            for key_, value_ in JOINTS_3D_ALL.items():
                if not all([col in poses_df.columns for col in value_]):
                    columns_not_in_poses = [col for col in value_ if col not in poses_df.columns]
                    if len(columns_not_in_poses) < 3:
                        print(f'Warning {key_} is a 3-D joint. Since {columns_not_in_poses.__str__()[1:-1]} is missing, '
                              f'the other DoFs ({[dof for dof in value_ if dof not in columns_not_in_poses].__str__()[1:-1]}) are filled with zeros too.')
                    missing_col.extend([f'{key_}_{n}' for n in range(6)])
                    missing_col.extend([f'{key_}_{axis}_angular_vel' for axis in ['x', 'y', 'z']])
                    poses_df[value_] = 0.

            poses_df = poses_df.astype(float)
            col_list = list(poses_df.columns)
            col_loc = [col_list.index(col) for col in OSIM_DOF_ALL[:23]]
            angle_col_loc = [col_list.index(col) for col in OSIM_DOF_ALL[:23] if col not in ['pelvis_tx', 'pelvis_ty', 'pelvis_tz']]
            poses_df.iloc[:, angle_col_loc] = poses_df.iloc[:, angle_col_loc] * angle_scale
            poses = poses_df.values[:, col_loc]

            all_zeros = np.zeros([poses_df.shape[0], 12])
            states = np.concatenate([np.array(poses), all_zeros], axis=1)
            # Inject experimental GRF force vectors before direction alignment so
            # align_moving_direction can rotate them into the canonical frame.
            _grf_df = read_mot(GRF_MOT)[1]
            _t_src = np.linspace(0, 1, len(_grf_df))
            _t_dst = np.linspace(0, 1, poses_df.shape[0])
            for _osim_col, _grf_col in _GRF_COL_MAP.items():
                _src_col = _grf_col if _grf_col in _grf_df.columns else _osim_col if _osim_col in _grf_df.columns else None
                if _src_col is None:
                    continue
                _idx = opt.osim_dof_columns.index(_osim_col)
                _vals = interp1d(_t_src, _grf_df[_src_col].values, kind='linear')(_t_dst)
                states[:, _idx] = _vals / opt.weight_kg if 'force_v' in _osim_col else _vals
            states = norm_cops(skel, states, opt, opt.weight_kg, opt.height_m, check_cop_to_calcn_distance=False).numpy()

            states_aligned, rot_mat = align_moving_direction(states, opt.osim_dof_columns)
            if states_aligned is False:
                print(f'Warning: {file_path} Pelvis orientation changed by more than 45 deg, don\'t trust this trial')
                rot_mat = torch.eye(3).float()
            else:
                states = states_aligned
            self.rot_mat_trials.append(rot_mat)

            file_name = file_path.split('/')[-1]
            states_df = pd.DataFrame(states, columns=opt.osim_dof_columns)

            states_df, pos_vec = convert_addb_state_to_model_input(states_df, opt.joints_3d, sampling_rate)

            if sampling_rate != self.target_sampling_rate:
                print('Warning: Sampling rate is not 100 Hz, using linear interpolation to resample data to 100 Hz')
                states_resampled = linear_resample_data(states_df.values, sampling_rate, self.target_sampling_rate)
                states_df = pd.DataFrame(states_resampled, columns=states_df.columns)

            assert self.opt.model_states_column_names == list(states_df.columns)
            converted_states = torch.tensor(states_df.values).float()

            trial_data = TrialData(converted_states, model_offsets, opt.height_m, opt.weight_kg, rot_mat, pos_vec,
                                   self.window_len, missing_col)
            self.trials.append(trial_data)
            self.file_names.append(file_name)


class WindowData:
    def __init__(self, pose, model_offsets, trial_id, gait_phase_label, mask, height_m, weight_kg, missing_col):
        self.pose = pose
        self.model_offsets = model_offsets
        self.trial_id = trial_id
        self.gait_phase_label = gait_phase_label
        self.mask = mask
        self.height_m = height_m
        self.weight_kg = weight_kg
        self.missing_col = missing_col


class TrialData:
    def __init__(self, converted_states, model_offsets, height_m, weight_kg, rot_mat_for_moving_direction_alignment,
                 pos_vec_for_pos_alignment, window_len, missing_col):
        self.converted_pose = converted_states
        self.model_offsets = model_offsets
        self.height_m = height_m
        self.weight_kg = weight_kg
        self.length = converted_states.shape[0]
        self.rot_mat_for_moving_direction_alignment = rot_mat_for_moving_direction_alignment
        self.pos_vec_for_pos_alignment = pos_vec_for_pos_alignment
        self.window_len = window_len
        self.missing_col = missing_col


""" ============================ End dataset.py ============================ """


def _resample_to_n(arr, n):
    t_src = np.linspace(0, 1, len(arr))
    t_dst = np.linspace(0, 1, n)
    return np.stack([interp1d(t_src, arr[:, i], kind='linear')(t_dst)
                     for i in range(arr.shape[1])], axis=1)


def _extract_gait_cycles(windows, v_grf_col, n=101, threshold=1.0):
    cycles = []
    for w in windows:
        grf = w[:, v_grf_col]
        hs = np.where((grf[:-1] < threshold) & (grf[1:] >= threshold))[0]
        for k in range(len(hs) - 1):
            seg = w[hs[k]:hs[k + 1]]
            if len(seg) < 3:
                continue
            cycles.append(_resample_to_n(seg, n))
    return cycles


def load_exp_window(kin_mot_path, grf_mot_path, osim_dof_columns):
    kin_df = read_mot(kin_mot_path)[1]
    grf_df = read_mot(grf_mot_path)[1]
    n = len(kin_df)
    result = np.zeros((n, len(osim_dof_columns)))
    for i, col in enumerate(osim_dof_columns):
        if col in kin_df.columns:
            result[:, i] = kin_df[col].values
        elif col in grf_df.columns:
            result[:, i] = grf_df[col].values
        elif col in _GRF_COL_MAP and _GRF_COL_MAP[col] in grf_df.columns:
            result[:, i] = grf_df[_GRF_COL_MAP[col]].values
    return result


def plot_wcb_results(exp_win, gen_cadence_lo_wins, gen_cadence_hi_wins, osim_dof_columns,
                     out_path_gaitcycle, dt=0.01, treadmill_speed=0.0):
    """exp_win: (n_frames, n_cols); gen_cadence_lo_wins / gen_cadence_hi_wins: list of (n_frames, n_cols) arrays."""
    pelvis_tx_idx  = osim_dof_columns.index('pelvis_tx')
    calcn_l_vy_idx = osim_dof_columns.index('calcn_l_force_vy')
    knee_l_idx     = osim_dof_columns.index('knee_angle_l')
    ankle_l_idx    = osim_dof_columns.index('ankle_angle_l')

    plot_cols = [
        (lambda arr: arr[:, calcn_l_vy_idx],                                   'Left vGRF (N)',        None),
        (lambda arr: np.gradient(arr[:, pelvis_tx_idx], dt) + treadmill_speed, 'Running speed (m/s)', 0),
        (lambda arr: np.rad2deg(arr[:, knee_l_idx]),                           'Left knee flexion (deg)', None),
        (lambda arr: np.rad2deg(arr[:, ankle_l_idx]),                          'Left ankle angle (deg)',  None),
    ]
    n_gc = 101
    groups = [
        (gen_cadence_lo_wins, WCB_COLOR_LO[0], 'Low cadence'),
        (gen_cadence_hi_wins, WCB_COLOR_HI[0], 'High cadence'),
    ]

    pct_axis  = np.linspace(0, 100, n_gc)

    def first_left_gc(arr):
        # Shared left heel-strike detection (utils.first_gait_cycle_indices), so the 2D plots
        # and the 3D viewer key off the same left-foot gait cycle.
        pair = first_gait_cycle_indices(arr[:, calcn_l_vy_idx], dt=dt)
        return _resample_to_n(arr[pair[0]:pair[1]], n_gc) if pair else None

    exp_gc = first_left_gc(exp_win)
    if exp_gc is None:
        exp_gc = _resample_to_n(exp_win, n_gc)

    fig2, axes2 = plt.subplots(1, len(plot_cols), figsize=(16, 4))
    for i_ax, (ax, (extractor, ylabel, ymin)) in enumerate(zip(axes2, plot_cols)):
        for wins, color, label in groups:
            gcs = [g for g in (first_left_gc(w) for w in wins) if g is not None]
            if not gcs:
                continue
            curves = np.stack([extractor(g) for g in gcs])
            mean_c = np.mean(curves, axis=0)
            sd_c   = np.std(curves, axis=0)
            ax.plot(pct_axis, mean_c, color=color, linewidth=1.5, label=label)
            ax.fill_between(pct_axis, mean_c - sd_c, mean_c + sd_c,
                            color=color, alpha=0.25)
        ax.set_xlabel('Left gait cycle (%)')
        ax.set_ylabel(ylabel)
        if ymin is not None:
            ax.set_ylim(bottom=ymin)
        ax.grid(True, alpha=0.3)
    # Single legend centered above the row, so it never overlaps the left-most vGRF subplot.
    handles, labels = axes2[0].get_legend_handles_labels()
    if handles:
        fig2.legend(handles, labels, loc='upper center', ncol=len(labels),
                    bbox_to_anchor=(0.5, 1.03), frameon=False)
    fig2.tight_layout(rect=(0, 0, 1, 0.93))
    fig2.savefig(out_path_gaitcycle, dpi=150, bbox_inches='tight')
    plt.close(fig2)


# 3D skeleton (glTF) visualization helpers live in the viewer/ package (kept separate so the
# heavy/optional OpenSim+glTF path is isolated from the diffusion pipeline).
from viewer import build_results_glb, WCB_GLB_COLORS  # noqa: E402


# @spaces.GPU attaches a ZeroGPU GPU for the duration of this call and releases it on
# return. This is the only path that runs torch model inference; the diffusion loop moves
# its results to CPU (.numpy()) before the CPU-only plot/glTF/zip tail, so the GPU scope is
# clean. duration=120s covers 2 cadence variants x n_generation guided-DDIM runs + model
# load + export; lower it if runs are consistently faster (improves queue ranking).
@spaces.GPU(duration=120)
def predict_with_synthetic_grf(opt, lo_scale=0.9, hi_scale=1.1, progress=None):
    diffusion_model, _ = load_diffusion_model(opt)

    # Forward speed (pelvis_tx) is the ONLY conditioning input; vGRF is now a generated
    # output. The two variants change cadence (stride frequency) at a fixed speed.
    speed_col_loc = [opt.model_states_column_names.index('pelvis_tx')]

    cadence_variants = {'cadence_lo': lo_scale, 'cadence_hi': hi_scale}
    total_seeds  = len(cadence_variants) * opt.n_generation
    seeds_done   = 0

    results_by_variant = {}
    for variant_name, cadence_scale in cadence_variants.items():
        print(f'\n--- Running diffusion with {variant_name} ({opt.n_generation} generations) ---')
        dataset = MotionDataset(opt, diffusion_model.normalizer, cadence_scale)

        windows = dataset.get_one_win_from_the_beginning_of_each_trial(speed_col_loc)
        if not windows:
            continue

        win    = windows[0]
        i_trial = 0
        state_syn       = win.pose.unsqueeze(0)
        masks           = win.mask.unsqueeze(0)
        cond            = torch.ones([6])
        trial_len       = dataset.trials[i_trial].converted_pose.shape[0]
        height_m_tensor = torch.tensor([win.height_m])

        value_diff_thd         = torch.zeros([len(opt.model_states_column_names)])
        do_not_follow_col_loc  = [loc for loc, col in enumerate(opt.model_states_column_names) if '_vel' in col]
        for i_dof in range(len(opt.model_states_column_names)):
            thd = (state_syn[:, :, i_dof].max() - state_syn[:, :, i_dof].min()) * 0.3
            value_diff_thd[i_dof] = thd
        value_diff_thd[speed_col_loc]         = 0
        value_diff_thd[do_not_follow_col_loc] = 999

        value_diff_weight                       = torch.ones([len(opt.model_states_column_names)])
        value_diff_weight[do_not_follow_col_loc] = 0

        preds = []
        for seed in range(opt.n_generation):
            if progress is not None:
                progress(seeds_done / total_seeds,
                         desc=f'{variant_name} seed {seed + 1}/{opt.n_generation}...')
            fix_seed(seed)
            state_pred = diffusion_model.eval_loop(
                opt, state_syn, masks,
                value_diff_thd=value_diff_thd,
                value_diff_weight=value_diff_weight,
                cond=cond,
                num_of_generation_per_window=1,
                mode='inpaint_ddim_guided',
            )[0]

            results_pred = state_pred[0][:trial_len].numpy()
            results_pred = inverse_convert_addb_state_to_model_input(
                torch.from_numpy(results_pred).unsqueeze(0),
                opt.model_states_column_names, opt.treadmill_speed,
                opt.joints_3d, opt.osim_dof_columns,
                dataset.trials[i_trial].pos_vec_for_pos_alignment,
                height_m_tensor)[0].numpy()
            results_pred = inverse_norm_cops(dataset.skel, results_pred, opt,
                                             win.weight_kg, win.height_m)
            results_pred = inverse_align_moving_direction(
                results_pred, opt.osim_dof_columns, dataset.rot_mat_trials[i_trial])
            results_pred[:, -12:-9] *= opt.weight_kg
            results_pred[:, -6:-3]  *= opt.weight_kg
            preds.append(results_pred)
            seeds_done += 1
            print(f'  seed {seed} done')

        results_by_variant[variant_name] = (preds, dataset, round(1 / opt.target_sampling_rate, 3))

    if progress is not None:
        progress(seeds_done / total_seeds, desc='Plotting results...')

    gc_path = os.path.join(DATA_WCB_DIR, 'exports', 'wcb_gaitcycle.png')
    exp_win = gen_cadence_lo_wins = gen_cadence_hi_wins = None

    if 'cadence_lo' in results_by_variant and 'cadence_hi' in results_by_variant:
        exp_win     = load_exp_window(KIN_MOT, GRF_MOT, opt.osim_dof_columns)
        gen_cadence_lo_wins = results_by_variant['cadence_lo'][0]
        gen_cadence_hi_wins = results_by_variant['cadence_hi'][0]
        dt = results_by_variant['cadence_lo'][2]
        plot_wcb_results(
            exp_win,
            gen_cadence_lo_wins,
            gen_cadence_hi_wins,
            opt.osim_dof_columns,
            gc_path,
            dt=dt,
            treadmill_speed=opt.treadmill_speed,
        )

    export_paths = {}
    for variant_name, (preds, dataset, dt) in results_by_variant.items():
        results_pred = preds[0]
        df = pd.DataFrame(results_pred, columns=opt.osim_dof_columns)
        base_name = f'exp_{opt.subject_name}_{opt.subject_name}_run400_segment_0_{variant_name}_pred'
        grf_path = os.path.join(DATA_WCB_DIR, 'exports', f'{base_name}_grf.mot')
        kin_path = os.path.join(DATA_WCB_DIR, 'exports', f'{base_name}_kin.mot')
        convertDfToGRFMot(df, grf_path, dt, dataset.time_column[0])
        convertDataframeToMot(df[OSIM_DOF_ALL[:23]], kin_path, dt, dataset.time_column[0])
        export_paths[variant_name] = {'kin': kin_path, 'grf': grf_path}

    # Build the combined 3D skeleton overlay (.gltf): both generated trials.
    glb_path = None
    if 'cadence_lo' in export_paths and 'cadence_hi' in export_paths:
        if progress is not None:
            progress(seeds_done / total_seeds, desc='Building 3D skeleton...')
        try:
            instances = [
                {'name': 'cadence_lo', 'osim': OSIM_PATH, 'kin_mot': export_paths['cadence_lo']['kin'],
                 'grf_mot': export_paths['cadence_lo']['grf'], 'color': WCB_GLB_COLORS['cadence_lo']},
                {'name': 'cadence_hi', 'osim': OSIM_PATH, 'kin_mot': export_paths['cadence_hi']['kin'],
                 'grf_mot': export_paths['cadence_hi']['grf'], 'color': WCB_GLB_COLORS['cadence_hi']},
            ]
            glb_path = build_results_glb(
                os.path.join(DATA_WCB_DIR, 'exports', 'wcb_results.gltf'), instances)
        except Exception as e:
            import traceback
            print('3D skeleton build failed:', e)
            traceback.print_exc()
            glb_path = None

    # Bundle all generated artifacts (lo/hi kinematics + GRF .mot, gait-cycle PNG, 3D .gltf)
    # into one fixed-name zip so the Part 3 download has a recognizable file name.
    import zipfile
    gen_zip_path = os.path.join(DATA_WCB_DIR, 'exports', 'GaitDynamics_generated_data.zip')
    try:
        with zipfile.ZipFile(gen_zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
            for variant_name, paths in export_paths.items():
                for kind in ('kin', 'grf'):
                    p = paths.get(kind)
                    if p and os.path.exists(p):
                        zf.write(p, os.path.basename(p))
            for extra in (gc_path, glb_path):
                if extra and os.path.exists(extra):
                    zf.write(extra, os.path.basename(extra))
    except Exception as e:
        print('Generated-data zip build failed:', e)
        gen_zip_path = None

    return gc_path, glb_path, gen_zip_path, exp_win, gen_cadence_lo_wins, gen_cadence_hi_wins


# ── WCB Configuration ─────────────────────────────────────────────────────
# Edit these values to change the subject, model paths, and diffusion params.

DATA_WCB_DIR   = os.path.dirname(os.path.abspath(__file__))
OSIM_PATH      = os.path.join(DATA_WCB_DIR, 'opensim_model.osim')
DIFFUSION_CKPT = os.path.join(DATA_WCB_DIR, 'GaitDynamicsDiffusion.pt')
KIN_MOT        = os.path.join(DATA_WCB_DIR, 'run400_segment_0_ik.mot')
GRF_MOT        = os.path.join(DATA_WCB_DIR, 'run400_segment_0_grf.mot')

# Subject anthropometrics and trial metadata
SUBJECT_HEIGHT_M  = 1.74
SUBJECT_WEIGHT_KG = 69.3
TREADMILL_SPEED   = 4.0     # m/s

# Diffusion model parameters
N_GENERATION                 = 3
N_GUIDED_STEPS               = 1
GUIDANCE_LR                  = 0.02
GUIDE_X_START_BEGINNING_STEP = 1000
GUIDE_X_START_END_STEP       = 0

opt = parse_opt()
opt.subject_osim_model               = OSIM_PATH
opt.geometry_folder                  = DATA_WCB_DIR + '/'
opt.height_m                         = SUBJECT_HEIGHT_M
opt.weight_kg                        = SUBJECT_WEIGHT_KG
opt.treadmill_speed                  = TREADMILL_SPEED
opt.n_guided_steps                   = N_GUIDED_STEPS
opt.guidance_lr                      = GUIDANCE_LR
opt.guide_x_start_the_beginning_step = GUIDE_X_START_BEGINNING_STEP
opt.guide_x_start_the_end_step       = GUIDE_X_START_END_STEP
opt.pseudo_dataset_len               = 1000
opt.subject_name                     = 'subject11'
opt.n_generation                     = N_GENERATION
opt.file_paths                       = [KIN_MOT]
# ──────────────────────────────────────────────────────────────────────────

import gradio as gr
import matplotlib.pyplot as plt

os.makedirs(os.path.join(DATA_WCB_DIR, 'exports'), exist_ok=True)

# ── Shared visual language ─────────────────────────────────────────────────
# Data colors used across both the UI accents and the charts so the two agree.
C_LOW  = WCB_COLOR_LO[0]   # low-cadence condition  (shared palette: 'steelblue')
C_HIGH = WCB_COLOR_HI[0]   # high-cadence condition (shared palette: 'tomato')
C_ORIG = 'lightgray'       # original / recorded signal (intentionally a faint, medium-specific gray)

# Global matplotlib theme: every figure (including the composite gait-cycle
# image built inside predict_with_synthetic_grf) inherits this clean, light look.
plt.rcParams.update({
    'figure.dpi':        110,
    'savefig.dpi':       110,
    'font.family':       ['Inter', 'DejaVu Sans', 'sans-serif'],
    'font.size':         10,
    'axes.titlesize':    11,
    'axes.titleweight':  'semibold',
    'axes.labelsize':    10,
    'axes.labelcolor':   '#334155',
    'axes.edgecolor':    '#cbd5e1',
    'axes.linewidth':    0.8,
    'axes.spines.top':   False,
    'axes.spines.right': False,
    'axes.grid':         True,
    'grid.color':        '#cbd5e1',
    'grid.alpha':        0.3,
    'grid.linewidth':    0.7,
    'xtick.color':       '#475569',
    'ytick.color':       '#475569',
    'xtick.labelsize':   9,
    'ytick.labelsize':   9,
    'legend.fontsize':   9,
    'legend.frameon':    False,
    'figure.facecolor':  'white',
    'axes.facecolor':    'white',
})


def _style_ax(ax):
    """Apply the shared chart styling (light grid, despined, muted labels)."""
    ax.grid(True, alpha=0.3)
    ax.set_axisbelow(True)
    for spine in ('top', 'right'):
        ax.spines[spine].set_visible(False)
    return ax


def read_mot_to_dataframe(file_path):
    return read_mot(file_path)[1]


def list_real_mot_columns(file_choice):
    selected = KIN_MOT if file_choice == 'Kinematics' else GRF_MOT
    df = read_mot_to_dataframe(selected)
    if df is None or 'time' not in df.columns:
        return gr.update(choices=[], value=None)
    cols = [c for c in df.columns if c != 'time']
    return gr.update(choices=cols, value=(cols[0] if cols else None))


def plot_real_mot_signal(file_choice, column):
    selected = KIN_MOT if file_choice == 'Kinematics' else GRF_MOT
    df = read_mot_to_dataframe(selected)
    if df is None or 'time' not in df.columns or column is None or column not in df.columns:
        fig, ax = plt.subplots(figsize=(10, 4))
        ax.set_title('No data to display')
        fig.tight_layout()
        return fig
    fig, ax = plt.subplots(figsize=(10, 4))
    ax.plot(df['time'], df[column])
    ax.set_xlabel('time (s)')
    ax.set_ylabel(column)
    ax.set_title(column)
    ax.grid(True)
    fig.tight_layout()
    return fig


def list_kin_columns():
    df = read_mot_to_dataframe(KIN_MOT)
    # Exclude FROZEN_DOFS (e.g. locked MTP angle, constant 0) — nothing to inspect.
    cols = ([c for c in df.columns if c != 'time' and c not in FROZEN_DOFS]
            if df is not None else [])
    return gr.update(choices=[(_human_label(c), c) for c in cols], value=None)


def list_grf_columns():
    df = read_mot_to_dataframe(GRF_MOT)
    if df is None:
        return gr.update(choices=[], value=None)
    skip = ('_mx', '_my', '_mz')
    cols = [c for c in df.columns if c != 'time' and not any(c.endswith(s) for s in skip)]
    return gr.update(choices=[(_human_label(c), c) for c in cols], value=None)


_KIN_BASE_LABELS = {
    'pelvis_tilt': 'Pelvis Tilt', 'pelvis_list': 'Pelvis List', 'pelvis_rotation': 'Pelvis Rotation',
    'pelvis_tx': 'Pelvis Fwd Translation', 'pelvis_ty': 'Pelvis Vertical Translation',
    'pelvis_tz': 'Pelvis Lateral Translation',
    'hip_flexion': 'Hip Flexion', 'hip_adduction': 'Hip Adduction', 'hip_rotation': 'Hip Rotation',
    'knee_angle': 'Knee Flexion', 'ankle_angle': 'Ankle Angle', 'subtalar_angle': 'Subtalar Angle',
    'mtp_angle': 'MTP Angle',
    'lumbar_extension': 'Lumbar Extension', 'lumbar_bending': 'Lumbar Bending',
    'lumbar_rotation': 'Lumbar Rotation',
}
_FORCE_QUANTITY = {'vx': 'Anterior-Posterior GRF', 'vy': 'vGRF', 'vz': 'Medial-Lateral GRF'}
_COP_QUANTITY   = {'x': 'CoP X', 'y': 'CoP Y', 'z': 'CoP Z'}


def _signal_side(col):
    """'Right' / 'Left' / '' for a raw column name (calcn_r, ground_force_1 → Right; 2 → Left)."""
    lc = col.lower()
    if 'calcn_r' in lc or 'ground_force_1' in lc or 'ground_torque_1' in lc or lc.endswith('_r'):
        return 'Right'
    if 'calcn_l' in lc or 'ground_force_2' in lc or 'ground_torque_2' in lc or lc.endswith('_l'):
        return 'Left'
    return ''


def _human_label(col):
    """Human-readable, side-first display name for a raw signal/column name
    (e.g. 'ground_force_1_vy' -> 'Right vGRF'). Falls back to a title-cased
    version of the raw name when unmapped."""
    if col is None:
        return col
    lc   = col.lower()
    side = _signal_side(col)

    # Centre-of-pressure positions (normed CoP, or .mot _px/_py/_pz)
    if 'normed_cop' in lc or lc.endswith(('_px', '_py', '_pz')):
        base = _COP_QUANTITY.get(lc[-1], 'CoP')
        return f'{side} {base}'.strip()
    # GRF force components
    if lc.endswith(('_vx', '_vy', '_vz')):
        base = _FORCE_QUANTITY.get(lc[-2:], 'GRF')
        return f'{side} {base}'.strip()

    # Kinematic degrees of freedom
    base_key = lc[:-2] if lc.endswith(('_r', '_l')) else lc
    if base_key in _KIN_BASE_LABELS:
        name = _KIN_BASE_LABELS[base_key]
        return f'{side} {name}'.strip() if side else name

    return col.replace('_', ' ').title()


def _mot_angles_in_radians(meta):
    """True if a .mot header declares inDegrees=no (angles in radians),
    False if inDegrees=yes, None if the flag is absent."""
    if not meta:
        return None
    for line in meta:
        ls = line.strip().lower().replace(' ', '')
        if ls.startswith('indegrees='):
            return ls.endswith('no')
    return None


def _col_display(col, model_array=False):
    """Return (scale_factor, unit_label) for a column name.

    model_array=True  → angles are in radians (model internal); apply rad→deg.
    model_array=False → reading from a .mot file; angles already in degrees.
    """
    lc = col.lower()
    # Translations and COP positions → metres
    if any(x in lc for x in ['_tx', '_ty', '_tz', '_px', '_py', '_pz']):
        return 1.0, 'm'
    # Force components → Newtons
    if any(x in lc for x in ['_vx', '_vy', '_vz']):
        return 1.0, 'N'
    # Normalised COP (dimensionless combination of force and displacement)
    if 'normed_cop' in lc:
        return 1.0, 'norm.'
    # Angles
    if model_array:
        return 180.0 / np.pi, 'deg'
    return 1.0, 'deg'


def plot_real_signal_auto(column):
    if column is None:
        return None
    for path in [KIN_MOT, GRF_MOT]:
        meta, df = read_mot(path)
        if df is not None and 'time' in df.columns and column in df.columns:
            scale, unit = _col_display(column, model_array=bool(_mot_angles_in_radians(meta)))
            fig, ax = plt.subplots(figsize=(7, 3))
            ax.plot(df['time'], df[column] * scale, color=C_LOW, linewidth=1.6)
            ax.set_xlabel('time (s)')
            ax.set_ylabel(f'{_human_label(column)} ({unit})')
            ax.set_title(_human_label(column))
            _style_ax(ax)
            fig.tight_layout()
            return fig
    return None


_MV_PLACEHOLDER = ("<div style='padding:1em;color:#888'>Run the model to view the "
                   "3D skeleton overlay.</div>")


def wcb_model_viewer_html(gltf_path):
    """Build the <model-viewer> embed + pause/resume/speed controls for a generated .gltf.
    Controls are client-side inline handlers calling the model-viewer JS API (timeScale = speed)."""
    if not gltf_path or not os.path.exists(gltf_path):
        return _MV_PLACEHOLDER
    import time
    abs_path = os.path.abspath(gltf_path).replace('\\', '/')
    ts = int(time.time())   # cache-bust the fixed wcb_results.gltf filename across runs
    return f'''
<div class="wcb-mv-stage">
  <model-viewer id="wcb-mv" src="/gradio_api/file={abs_path}?t={ts}"
      autoplay camera-controls interaction-prompt="none" shadow-intensity="1"
      style="width:100%;height:480px;background-color:#ffffff;"></model-viewer>
  <div class="wcb-mv-controls">
    <button type="button" onclick="document.getElementById('wcb-mv').play()" title="Play" aria-label="Play">
      <svg viewBox="0 0 16 16" fill="currentColor" aria-hidden="true"><path d="M4 2.5v11l9-5.5z"/></svg>
    </button>
    <button type="button" onclick="document.getElementById('wcb-mv').pause()" title="Pause" aria-label="Pause">
      <svg viewBox="0 0 16 16" fill="currentColor" aria-hidden="true"><rect x="3.5" y="2.5" width="3.5" height="11" rx="0.5"/><rect x="9" y="2.5" width="3.5" height="11" rx="0.5"/></svg>
    </button>
    <select class="wcb-mv-speed" title="Play speed" aria-label="Play speed"
        onchange="document.getElementById('wcb-mv').timeScale = parseFloat(this.value)">
      <option value="0.1"> 0.1x </option>
      <option value="0.2"> 0.2x </option>
      <option value="0.5"> 0.5x </option>
      <option value="1" selected> 1.0x </option>
    </select>
  </div>
</div>
<div class="wcb-mv-scrubrow">
  <span class="wcb-mv-scrublabel">Left gait cycle (%)</span>
  <input type="range" class="wcb-mv-scrub" min="0" max="100" step="0.5" value="0"
    oninput="(function(s){{var mv=document.getElementById('wcb-mv');mv.pause();if(mv.duration)mv.currentTime=(parseFloat(s.value)/100)*mv.duration;s.nextElementSibling.textContent=Math.round(parseFloat(s.value))+'%';}})(this)">
  <span class="wcb-mv-pct">0%</span>
</div>
'''


def btn3_generate_and_plot(lo_scale_pct, hi_scale_pct, lo_set, hi_set, progress=gr.Progress()):
    if not lo_set or not hi_set:
        return ("Please set both cadence scale factors using the sliders in Step 2 first.",
                _MV_PLACEHOLDER, None, gr.update(visible=False))
    progress(0.0, desc="Loading diffusion model...")
    try:
        _gc_path, glb_path, gen_zip_path, exp_win, gen_cadence_lo_wins, gen_cadence_hi_wins = predict_with_synthetic_grf(
            opt, lo_scale=lo_scale_pct / 100, hi_scale=hi_scale_pct / 100, progress=progress)
        progress(1.0, desc="Done.")
        results_data = (exp_win, gen_cadence_lo_wins, gen_cadence_hi_wins, opt.osim_dof_columns,
                        round(1 / opt.target_sampling_rate, 3), opt.treadmill_speed)
        dl_update = (gr.update(value=gen_zip_path, visible=True) if gen_zip_path
                     else gr.update(visible=False))
        return ("Generation complete.", wcb_model_viewer_html(glb_path),
                results_data, dl_update)
    except Exception as e:
        import traceback
        return (f"Error: {e}\n\n{traceback.format_exc()}", _MV_PLACEHOLDER, None,
                gr.update(visible=False))


def create_data_zip():
    # Fixed name (under exports/) so the downloaded file is recognizable, not random characters.
    import zipfile
    zip_path = os.path.join(DATA_WCB_DIR, 'exports', 'GaitDynamics_real_data.zip')
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
        zf.write(KIN_MOT,   os.path.basename(KIN_MOT))
        zf.write(GRF_MOT,   os.path.basename(GRF_MOT))
        zf.write(OSIM_PATH, os.path.basename(OSIM_PATH))
    return zip_path


def plot_cadence_scale_preview(scale_pct, color):
    """Preview the cadence re-timing on a recognizable kinematic signal (left knee flexion).
    The trial is temporally resampled to round(len / scale) frames at 100 Hz, so the stride
    frequency changes (faster steps when scale > 1) while the amplitude is left untouched."""
    meta, df = read_mot(KIN_MOT)
    orig = None
    if df is not None and 'knee_angle_l' in df.columns:
        orig = df['knee_angle_l'].values.copy()
        if _mot_angles_in_radians(meta):
            orig = np.rad2deg(orig)
    if orig is None:
        fig, ax = plt.subplots(figsize=(8, 3))
        ax.set_title('Column not found')
        fig.tight_layout()
        return fig
    scale = scale_pct / 100.0
    n_resampled = int(round(len(orig) / scale))
    resampled = linear_resample_data_as_num_of_dp(orig.reshape(-1, 1), n_resampled).ravel()
    # Plot on a seconds axis (100 Hz) so the change in step frequency is visible; amplitude
    # is identical between the two curves — only the timing is compressed/stretched.
    t_orig = np.arange(len(orig)) / 100.0
    t_new  = np.arange(len(resampled)) / 100.0
    fig, ax = plt.subplots(figsize=(8, 3))
    ax.plot(t_orig, orig,      color=C_ORIG, linewidth=1.6, label='Original')
    ax.plot(t_new,  resampled, color=color,  linewidth=1.8, label='Re-timed')
    ax.set_xlabel('time (s)')
    ax.set_ylabel('Left knee flexion (deg)')
    ax.legend()
    _style_ax(ax)
    fig.tight_layout()
    return fig


def plot_param_over_gait_cycle(param, results):
    if results is None or param is None:
        return None
    exp_win, gen_cadence_lo_wins, gen_cadence_hi_wins, osim_dof_columns, dt, treadmill_speed = results
    if param not in osim_dof_columns:
        return None
    col_idx        = osim_dof_columns.index(param)
    calcn_l_vy_idx = osim_dof_columns.index('calcn_l_force_vy')
    n_gc      = 101
    pct_axis  = np.linspace(0, 100, n_gc)

    def first_left_gc(arr):
        # Shared left heel-strike detection (utils.first_gait_cycle_indices), so the 2D plots
        # and the 3D viewer key off the same left-foot gait cycle.
        pair = first_gait_cycle_indices(arr[:, calcn_l_vy_idx], dt=dt)
        return _resample_to_n(arr[pair[0]:pair[1]], n_gc) if pair else None

    scale, unit = _col_display(param, model_array=True)
    fig, ax = plt.subplots(figsize=(7, 3.5))
    for wins, color, label in [(gen_cadence_lo_wins, C_LOW,  'Low cadence'),
                                (gen_cadence_hi_wins, C_HIGH, 'High cadence')]:
        if not wins:
            continue
        gcs = [g for g in (first_left_gc(w) for w in wins) if g is not None]
        if not gcs:
            continue
        curves = np.stack([g[:, col_idx] * scale for g in gcs])
        mean_c = np.mean(curves, axis=0)
        sd_c   = np.std(curves, axis=0)
        ax.plot(pct_axis, mean_c, color=color, linewidth=1.5, label=label)
        ax.fill_between(pct_axis, mean_c - sd_c, mean_c + sd_c, color=color, alpha=0.25)

    ax.set_xlabel('Left gait cycle (%)')
    ax.set_ylabel(f'{_human_label(param)} ({unit})')
    ax.set_title(_human_label(param))
    ax.legend()
    _style_ax(ax)
    fig.tight_layout()
    return fig


# ── Gradio UI ──────────────────────────────────────────────────────────────

_CSS = """
/* ── Layout shell ───────────────────────────────────────────────────────── */
.gradio-container { max-width: 100% !important; margin: 0 auto !important; }

/* ── Hero ───────────────────────────────────────────────────────────────── */
.wcb-hero { padding: 4px 2px 6px; margin-bottom: 4px; }
.wcb-hero h1 { font-size: 1.9rem; font-weight: 700; line-height: 1.15; margin: 0 0 .25rem;
               letter-spacing: -0.01em; color: var(--body-text-color); }
.wcb-hero .wcb-sub { font-size: 1.02rem; color: var(--body-text-color-subdued);
                     margin: 0 0 .9rem; }

/* ── Step cards ─────────────────────────────────────────────────────────── */
.wcb-card { border: 1px solid var(--border-color-primary); border-radius: 14px;
            padding: 18px 20px 20px; margin: 14px 0; background: var(--background-fill-primary);
            box-shadow: 0 1px 2px rgba(15,23,42,0.04), 0 6px 16px rgba(15,23,42,0.05); }
.wcb-step-head { display: flex; align-items: center; gap: .65rem; margin: 0 0 .35rem; }
.wcb-step-head h3, .wcb-step-head h4 { margin: 0; font-size: 1.18rem; font-weight: 650;
                                        color: var(--body-text-color); }
.wcb-step-head .wcb-spacer { flex: 1; }
.wcb-step-badge { display: inline-flex; align-items: center; justify-content: center;
                  width: 1.7rem; height: 1.7rem; flex: 0 0 1.7rem; border-radius: 50%;
                  background: var(--primary-500); color: #fff; font-size: .92rem;
                  font-weight: 700; line-height: 1; }
.wcb-subhead { display: flex; align-items: center; gap: .6rem; margin: 1.4rem 0 .2rem;
               font-size: 1.02rem; font-weight: 600; }

/* ── Plots ──────────────────────────────────────────────────────────────── */
.wcb-plot    { min-height: 320px; }
.wcb-plot > div { min-height: 320px; }
.wcb-plot-sm { min-height: 240px; }
.wcb-plot-sm > div { min-height: 240px; }

/* ── 3D model-viewer controls ───────────────────────────────────────────── */
.wcb-mv-stage { position: relative; border-radius: 12px; overflow: hidden;
                border: 1px solid var(--border-color-primary); background: #ffffff; }
/* The model-viewer sits on a solid white background for skeleton contrast; strip any
   fill/border from the wrapping group so the only box is .wcb-mv-stage. */
.wcb-mv-group { background: #ffffff !important; border: none !important;
                box-shadow: none !important; padding: 0 !important; }
.wcb-mv-controls { position: absolute; top: 10px; left: 10px; z-index: 10;
                   display: flex; flex-direction: column; align-items: center; gap: 0.5em;
                   background: rgba(255,255,255,0.85); backdrop-filter: blur(4px);
                   padding: 5px 8px; border-radius: 8px;
                   box-shadow: 0 1px 4px rgba(15,23,42,0.12); }
.wcb-mv-controls button { display: inline-flex; align-items: center; justify-content: center;
                          width: 2.4em; height: 2.4em; padding: 0; line-height: 1;
                          font-size: 1.15em; cursor: pointer; color: #fff;
                          background: var(--primary-500); border: none; border-radius: 6px; }
.wcb-mv-controls button:hover { background: var(--primary-600); }
.wcb-mv-controls button svg { width: 1.1em; height: 1.1em; fill: #fff; }
.wcb-mv-speed { font-size: 0.95em; padding: 0.3em 0.4em; cursor: pointer; color: #fff;
                background: var(--primary-500); border: none; border-radius: 6px; }
.wcb-mv-speed:hover { background: var(--primary-600); }
.wcb-mv-speed option { color: #fff; background: var(--primary-500); }
.wcb-mv-scrubrow { display: flex; align-items: center; gap: 0.6em; margin-top: 0.6em; }
.wcb-mv-scrublabel { font-size: 0.85em; color: var(--body-text-color-subdued); white-space: nowrap; }
.wcb-mv-scrub { flex: 1; accent-color: var(--primary-500); }
.wcb-mv-pct { font-size: 0.85em; color: var(--body-text-color-subdued);
              min-width: 2.6em; text-align: right; }
"""

# Clean-scientific theme: light neutral surfaces, a single blue accent, Inter type.
_THEME = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="sky",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
).set(
    body_background_fill="#f8fafc",
    block_radius="12px",
    block_border_width="1px",
    button_large_radius="10px",
    button_small_radius="8px",
)

# Loads Google's <model-viewer> web component (used for the 3D skeleton viewer with
# pause/resume/speed controls). Defined once for the whole app.
_HEAD = ('<script type="module" '
         'src="https://unpkg.com/@google/model-viewer@4.0.0/dist/model-viewer.min.js"></script>')

with gr.Blocks(css=_CSS, head=_HEAD, theme=_THEME,
               title="GaitDynamics — Cadence") as demo:
    gr.HTML("""
        <div class="wcb-hero">
            <h1>GaitDynamics Demo</h1>
            <p>
                This demo uses <a href="https://doi.org/10.1038/s41551-025-01565-8" target="_blank">GaitDynamics</a>,
                a diffusion model, to generate running dynamics. The runner's forward speed is prescribed as the input
                ; the full-body kinematics and ground reaction force (GRF) are generated by the model as outputs.
                Here we hold the speed fixed and vary the <b>cadence</b> to see how it leads to different gait parameters.
            </p>
            <p style="margin-bottom:6px;">The demo has three steps:</p>
            <ol style="margin-top:6px; margin-bottom:0; padding-left:1.4em; line-height:1.8;">
                <li><b>Explore real data</b> — Browse the real kinematics and GRF recorded during the experiment.</li>
                <li><b>Prescribe cadence</b> — Set a low and a high cadence at the same fixed speed. The sliders preview how the stride timing is re-timed.</li>
                <li><b>Generate</b> — Run the model to synthesize the motion for both cadences and compare them over the gait cycle.</li>
            </ol>
        </div>
    """)

    # Section 1: Explore real data ────────────────────────────────────────────
    with gr.Column(elem_classes="wcb-card"):
        gr.HTML('<div class="wcb-step-head"><span class="wcb-step-badge">1</span>'
                '<h3>Explore real data</h3></div>')
        gr.DownloadButton("Download real data", value=create_data_zip,
                          variant="primary", size="sm", scale=0)
        gr.Markdown(
            "This is a single real running trial (speed = 4.0 m/s). The subject height is 1.74 m and weight is 69.3 kg."
            "\n\nOn the left, pick any kinematic signal (joint angles, in degrees) to see how it evolves over time; on "
            "the right, pick a GRF or center-of-pressure (CoP) signal. GRF components are shown in "
            "Newtons (N), and CoP positions in meters."
        )
        with gr.Row():
            with gr.Column(scale=1):
                col_select_kin = gr.Dropdown(choices=[], label="Kinematics", interactive=True)
                plot_real_kin  = gr.Plot(show_label=False, elem_classes="wcb-plot-sm")
            with gr.Column(scale=1):
                col_select_grf = gr.Dropdown(choices=[], label="Force / CoP", interactive=True)
                plot_real_grf  = gr.Plot(show_label=False, elem_classes="wcb-plot-sm")

    # Section 2: Prescribe cadence ────────────────────────────────────────────
    with gr.Column(elem_classes="wcb-card"):
        gr.HTML('<div class="wcb-step-head"><span class="wcb-step-badge">2</span>'
                '<h3>Prescribe cadence</h3></div>')
        gr.Markdown(
            "Speed is fixed, identical to the recorded trial. Cadence is what we change. "
            "A cadence scale *s* simply resamples the "
            "stride: the trial is temporally resampled to `T / s` of its original duration, so the runner takes "
            "faster, shorter steps when *s* > 1 and slower, longer steps when *s* < 1. Low = 85–99%, high = 101–115%.\n\n"
            "The previews below show the resampled left-knee flexion: the curve's shape and "
            "amplitude are unchanged, only its timing compresses or stretches with cadence."
        )

        lo_set = gr.State(False)
        hi_set = gr.State(False)

        with gr.Row():
            with gr.Column():
                lo_slider  = gr.Slider(minimum=85,  maximum=99,  step=1, value=90,  label="Low cadence scale (%)")
                lo_preview = gr.Plot(show_label=False, elem_classes="wcb-plot")
            with gr.Column():
                hi_slider  = gr.Slider(minimum=101, maximum=115, step=1, value=110, label="High cadence scale (%)")
                hi_preview = gr.Plot(show_label=False, elem_classes="wcb-plot")

    # Section 3: Generate ─────────────────────────────────────────────────────
    with gr.Column(elem_classes="wcb-card"):
        gr.HTML('<div class="wcb-step-head"><span class="wcb-step-badge">3</span>'
                '<h3>Generate kinematics for each cadence</h3></div>')
        gr.Markdown(
            "Now we run GaitDynamics, the diffusion model, to synthesize the full-body motion "
            "for each of the two cadences. \n\n"
            "Because the process is stochastic, we draw several samples — each one is called a seed. A "
            "seed is just one random starting point, so different seeds give slightly different but equally "
            "plausible motions. The status box reports progress as *seed k / N*; the plots below show the "
            "mean ± standard deviation across seeds, so you can see both the typical motion and how "
            "much it varies. Generation runs the model many times and may take a couple of minutes."
        )
        with gr.Row():
            btn3   = gr.Button("Run GaitDynamics", variant="primary", size="sm", scale=0)
            dl_gen = gr.DownloadButton("Download generated data", variant="primary",
                                       size="sm", scale=0, visible=False)
        gen_status  = gr.Textbox(
            label="Generation status", lines=3,
            info="Shows which seed is currently being generated (e.g. 'cadence_lo seed 2/3').")

        gr.HTML('<div class="wcb-subhead">3D skeleton overlay &mdash; '
                f"<span style='color:{WCB_COLOR_LO[0]}'>low cadence</span> / "
                f"<span style='color:{WCB_COLOR_HI[0]}'>high cadence</span> (kinematics + GRF arrows)</div>")
        gr.Markdown(
            "The two generated cadences are overlaid on one skeleton so you "
            "can compare the strides directly. The arrows represent the generated GRF vectors."
        )
        with gr.Group(elem_classes="wcb-mv-group"):
            model3d_out = gr.HTML(value=_MV_PLACEHOLDER)

        results_state = gr.State(None)
        gr.HTML('<div class="wcb-subhead">Compare cadences over the gait cycle</div>')
        gr.Markdown(
            "After generating, pick any kinematic or force / CoP channel to compare "
            "the low- vs. high-cadence conditions over the gait cycle (mean ± SD across seeds)."
        )
        with gr.Row():
            with gr.Column(scale=1):
                param_sel_kin  = gr.Dropdown(
                    choices=[(_human_label(c), c) for c in OSIM_DOF_ALL[:23]
                             if c not in FROZEN_DOFS], value=None, label="Kinematics")
                param_plot_kin = gr.Plot(show_label=False, elem_classes="wcb-plot-sm")
            with gr.Column(scale=1):
                param_sel_grf  = gr.Dropdown(
                    choices=[(_human_label(c), c) for c in KINETICS_ALL], value=None, label="Force / CoP")
                param_plot_grf = gr.Plot(show_label=False, elem_classes="wcb-plot-sm")

    # Citation ────────────────────────────────────────────────────────────────
    with gr.Column(elem_classes="wcb-card"):
        gr.HTML('<div class="wcb-step-head"><h3>Citation</h3></div>')
        gr.Markdown(
            "Tan, T., Van Wouwe, T., Werling, K. F., Liu, C. K., Delp, S. L., Hicks, J. L., & Chaudhari, A. S. "
            "GaitDynamics: a generative foundation model for analyzing human walking "
            "and running. *Nat. Biomed. Eng* (2026). "
            "[https://doi.org/10.1038/s41551-025-01565-8]"
            "(https://doi.org/10.1038/s41551-025-01565-8)"
        )

    # Wire interactions ───────────────────────────────────────────────────────
    def _load_columns():
        return list_kin_columns(), list_grf_columns()
    demo.load(fn=_load_columns, outputs=[col_select_kin, col_select_grf])

    col_select_kin.change(fn=plot_real_signal_auto, inputs=[col_select_kin], outputs=[plot_real_kin])
    col_select_grf.change(fn=plot_real_signal_auto, inputs=[col_select_grf], outputs=[plot_real_grf])

    def _lo_release(v):
        return plot_cadence_scale_preview(v, C_LOW), True

    def _hi_release(v):
        return plot_cadence_scale_preview(v, C_HIGH), True

    lo_slider.release(fn=_lo_release, inputs=[lo_slider], outputs=[lo_preview, lo_set])
    hi_slider.release(fn=_hi_release, inputs=[hi_slider], outputs=[hi_preview, hi_set])

    btn3.click(
        fn=btn3_generate_and_plot,
        inputs=[lo_slider, hi_slider, lo_set, hi_set],
        outputs=[gen_status, model3d_out, results_state, dl_gen],
        show_progress_on=[gen_status],
    )

    param_sel_kin.change(fn=plot_param_over_gait_cycle,
                         inputs=[param_sel_kin, results_state], outputs=[param_plot_kin])
    param_sel_grf.change(fn=plot_param_over_gait_cycle,
                         inputs=[param_sel_grf, results_state], outputs=[param_plot_grf])


if __name__ == '__main__':
    demo.launch(allowed_paths=[os.path.join(DATA_WCB_DIR, 'exports')])