viewer/builder.py

"""3D skeleton (glTF) visualization for GaitDynamics results.

Converts the final OpenSim results (model + kinematics .mot + GRF .mot) into a single
self-contained .gltf showing three overlaid, semi-transparent skeletons (original + the
two generated trials), each in its own color, with color-matched GRF force arrows.

Uses the vendored opensim-org/opensim-viewer-backend converters in ./osimConverters.
opensim and osimConverters are imported lazily so this module can be imported without
the OpenSim API present (it is installed at startup by app.py).
"""
import os
import tempfile

import numpy as np
import pandas as pd

from utils import (first_gait_cycle_indices, heel_strike_threshold, left_vy_column,
                   read_mot, WCB_GC_FRAMES, WCB_GC_DURATION, WCB_COLOR_LO, WCB_COLOR_HI)

# Overlay colors (RGB 0-1). lo/hi come from the shared palette (utils.WCB_COLOR_*) so the 3D
# overlay, the 2D plots, and the HTML legend can't drift; 'original' is intentionally a
# medium-specific gray (the 2D line and legend use lighter grays for their own roles).
WCB_GLB_COLORS = {
    'original':   [0.30, 0.30, 0.30],   # gray
    'cadence_lo': WCB_COLOR_LO[1],      # steelblue
    'cadence_hi': WCB_COLOR_HI[1],      # tomato
}
WCB_GLB_ALPHA = 0.30   # 70% transparent
WCB_GRF_VIS_SCALE = 4.5   # multiply GRF arrow length (force magnitude) in the viewer only

# Gait-cycle normalization: every skeleton is resampled to one LEFT heel-strike-to-heel-strike
# cycle of WCB_GC_FRAMES samples mapped onto a shared [0, WCB_GC_DURATION]s timeline, so all
# three animate in lockstep (frame k == same gait % in every skeleton). The heel-strike
# threshold / min-duration guard and the left-foot column pick live in utils.py, shared with
# app.first_left_gc so the 2D plots and this 3D viewer key off the same left-foot gait cycle.


def _write_mot_meta_and_df(path, meta, df):
    """Write a .mot file, preserving header units/flags verbatim and fixing nRows/nColumns."""
    new_meta = []
    for l in meta:
        s = l.strip().lower()
        if s.startswith('ncolumns='):
            new_meta.append(f'nColumns={df.shape[1]}\n')
        elif s.startswith('nrows='):
            new_meta.append(f'nRows={df.shape[0]}\n')
        else:
            new_meta.append(l)
    with open(path, 'w', newline='') as f:
        f.writelines(new_meta)
        df.to_csv(f, sep='\t', index=False)


def _slice_resample_by_time(df, t0, t1, n, duration):
    """Slice a .mot DataFrame to the time span [t0, t1], linearly resample every column to
    `n` frames, and replace the time column with a fresh [0, duration] ramp."""
    t = df['time'].values.astype(float)
    mask = (t >= t0) & (t <= t1)
    sub = df.loc[mask].reset_index(drop=True)
    if len(sub) < 2:
        sub = df.reset_index(drop=True)
    src = np.linspace(0.0, 1.0, len(sub))
    dst = np.linspace(0.0, 1.0, n)
    out = pd.DataFrame({col: np.interp(dst, src, sub[col].values.astype(float))
                        for col in sub.columns})
    out['time'] = np.linspace(0.0, duration, n)
    return out


def _resample_to_gait_cycle(kin_in, grf_in, kin_out, grf_out):
    """Normalize one trial's kinematics + GRF to a single gait cycle on a shared timeline.

    Detects one LEFT heel-strike-to-heel-strike cycle from the left-foot vertical GRF column
    (`*_l_vy` / `ground_force_2_vy`), matching app.first_left_gc, then slices BOTH files to that
    time span (by time, not row index, since the original trial's kin/grf may differ in sampling
    rate) and resamples each to WCB_GC_FRAMES frames over [0, WCB_GC_DURATION]s. Falls back to the
    full window if no left-foot column or fewer than two heel strikes are found."""
    kmeta, kdf = read_mot(kin_in)
    gmeta, gdf = read_mot(grf_in)

    gtime = gdf['time'].values.astype(float)
    vy_cols = [c for c in gdf.columns if c.endswith('_vy')]
    vcol = left_vy_column(vy_cols)
    if vcol is None and vy_cols:                 # unexpected naming → fall back to peak foot
        vcol = max(vy_cols, key=lambda c: np.abs(gdf[c].values.astype(float)).max())
    if vcol is not None:
        v = gdf[vcol].values.astype(float)
        pair = first_gait_cycle_indices(v, time=gtime)
        t0, t1 = (gtime[pair[0]], gtime[pair[1]]) if pair else (gtime[0], gtime[-1])
    else:
        t0, t1 = gtime[0], gtime[-1]

    kout = _slice_resample_by_time(kdf, t0, t1, WCB_GC_FRAMES, WCB_GC_DURATION)
    gout = _slice_resample_by_time(gdf, t0, t1, WCB_GC_FRAMES, WCB_GC_DURATION)
    _write_mot_meta_and_df(kin_out, kmeta, kout)
    _write_mot_meta_and_df(grf_out, gmeta, gout)


def _coloc_kin_and_grf(kin_in, grf_in, kin_out, grf_out):
    """Treadmill co-location: pin pelvis_tx/tz to their frame-0 value so all skeletons
    overlay at a shared origin, and shift the GRF centre-of-pressure (_px/_pz) by the same
    per-frame horizontal offset so the force arrows stay attached to the feet."""
    kmeta, kdf = read_mot(kin_in)
    ktime = kdf['time'].values.astype(float)
    tx0 = float(kdf['pelvis_tx'].iloc[0])
    tz0 = float(kdf['pelvis_tz'].iloc[0])
    dx = kdf['pelvis_tx'].values.astype(float) - tx0
    dz = kdf['pelvis_tz'].values.astype(float) - tz0
    kdf['pelvis_tx'] = tx0
    kdf['pelvis_tz'] = tz0
    _write_mot_meta_and_df(kin_out, kmeta, kdf)

    gmeta, gdf = read_mot(grf_in)
    gtime = gdf['time'].values.astype(float)
    gdx = np.interp(gtime, ktime, dx)
    gdz = np.interp(gtime, ktime, dz)

    # Swing-phase suppression: per foot, mark frames whose (raw) vertical GRF is below the
    # heel-strike threshold (max 20 N, 5% of peak vGRF — same cutoff as gait detection). Those
    # are off-the-ground frames; zeroing the foot's force makes the converter draw no arrow
    # (its mag<1e-10 branch) instead of a tiny distracting swing-phase stub. Compute masks from
    # raw values first, before the visibility scaling below.
    swing_masks = {}   # foot prefix (e.g. 'force_l') -> boolean frame mask
    for col in gdf.columns:
        if col.endswith('_vy'):
            vy_raw = np.abs(gdf[col].values.astype(float))
            swing_masks[col[:-3]] = vy_raw < heel_strike_threshold(vy_raw)

    for col in gdf.columns:
        if col.endswith('_px'):
            gdf[col] = gdf[col].values.astype(float) - gdx
        elif col.endswith('_pz'):
            gdf[col] = gdf[col].values.astype(float) - gdz
        elif col.endswith(('_vx', '_vy', '_vz')):
            # Scale force magnitude (arrow length) for visibility; viewer copy only.
            vals = gdf[col].values.astype(float) * WCB_GRF_VIS_SCALE
            mask = swing_masks.get(col[:-3])
            if mask is not None:
                vals[mask] = 0.0   # hide this foot's arrow during its swing phase
            gdf[col] = vals
    _write_mot_meta_and_df(grf_out, gmeta, gdf)


def _append_force_arrows(gltf, grf_mot, color):
    """Append GRF force-arrow nodes + animation for one trial into an existing skeleton glTF,
    reusing the vendored force converter helpers and merging into the skeleton's animation[0]."""
    import opensim as osim
    from . import osimConverters as osimC
    from pygltflib import Node
    table = osim.TimeSeriesTable(grf_mot)
    table_vec3 = table.packVec3()
    labels = table_vec3.getColumnLabels()
    fdict = {}
    osimC.createForceDictionary(labels, fdict)
    if len(fdict) == 0 or table_vec3.getNumRows() == 0:
        return
    first_frame = table_vec3.getRowAtIndex(0)
    top_node = Node()
    top_node.name = 'ForceData'
    if top_node.children is None:
        top_node.children = []
    gltf.nodes.append(top_node)
    gltf.scenes[0].nodes.append(len(gltf.nodes) - 1)
    first_node_index = osimC.createForceNodes(
        'arrow', fdict, 1.0, False, osimC.getForceMeshScale(), first_frame, gltf, top_node)
    osimC.convertForcesTableToGltfAnimation(gltf, table_vec3, 1.0, fdict, first_node_index)


def _recolor_gltf(gltf, color, alpha=WCB_GLB_ALPHA):
    """Set every material of one skeleton instance to its overlay color at the given opacity."""
    from pygltflib import PbrMetallicRoughness
    r, g, b = color
    for mat in gltf.materials:
        if mat.pbrMetallicRoughness is None:
            mat.pbrMetallicRoughness = PbrMetallicRoughness()
        mat.pbrMetallicRoughness.baseColorFactor = [r, g, b, alpha]
        mat.alphaMode = 'BLEND'
        mat.alphaCutoff = None
        mat.doubleSided = True


_GLTF_ATTRS = ['POSITION', 'NORMAL', 'TANGENT', 'TEXCOORD_0', 'TEXCOORD_1',
               'COLOR_0', 'JOINTS_0', 'WEIGHTS_0']


def _merge_gltf(base, add):
    """Merge glTF `add` into `base`, offsetting every cross-index reference. Animation channels
    are folded into base.animations[0] so all skeletons + arrows play on one timeline."""
    from pygltflib import Animation
    node_off = len(base.nodes)
    mesh_off = len(base.meshes)
    acc_off = len(base.accessors)
    bv_off = len(base.bufferViews)
    buf_off = len(base.buffers)
    mat_off = len(base.materials)
    cam_off = len(base.cameras)

    base.buffers.extend(add.buffers)
    for bv in add.bufferViews:
        if bv.buffer is not None:
            bv.buffer += buf_off
        base.bufferViews.append(bv)
    for acc in add.accessors:
        if acc.bufferView is not None:
            acc.bufferView += bv_off
        base.accessors.append(acc)
    base.materials.extend(add.materials)
    base.cameras.extend(add.cameras)
    for mesh in add.meshes:
        for prim in mesh.primitives:
            attr = prim.attributes
            for fld in _GLTF_ATTRS:
                v = getattr(attr, fld, None)
                if v is not None:
                    setattr(attr, fld, v + acc_off)
            if prim.indices is not None:
                prim.indices += acc_off
            if prim.material is not None:
                prim.material += mat_off
        base.meshes.append(mesh)
    for node in add.nodes:
        if node.mesh is not None:
            node.mesh += mesh_off
        if node.camera is not None:
            node.camera += cam_off
        if node.children:
            node.children = [c + node_off for c in node.children]
        base.nodes.append(node)
    for ni in add.scenes[0].nodes:
        base.scenes[0].nodes.append(ni + node_off)
    if add.animations:
        if not base.animations:
            base.animations.append(Animation())
        banim = base.animations[0]
        for aanim in add.animations:
            samp_off = len(banim.samplers)
            for s in aanim.samplers:
                if s.input is not None:
                    s.input += acc_off
                if s.output is not None:
                    s.output += acc_off
                banim.samplers.append(s)
            for ch in aanim.channels:
                if ch.target is not None and ch.target.node is not None:
                    ch.target.node += node_off
                if ch.sampler is not None:
                    ch.sampler += samp_off
                banim.channels.append(ch)


def build_results_glb(out_path, instances):
    """Build one self-contained .gltf overlaying the given skeleton instances.
    `instances`: list of dicts with keys name, osim, kin_mot, grf_mot, color."""
    import opensim  # noqa: F401  (ensures the OpenSim API is importable before converting)
    from . import osimConverters as osimC
    from .osim_viewer_options import osimViewerOptions

    geom_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'osimConverters', 'Geometry')
    options = osimViewerOptions()
    options.setShowMuscles(False)
    tmpdir = tempfile.mkdtemp(prefix='wcb_glb_')

    merged = None
    for inst in instances:
        gc_kin = os.path.join(tmpdir, inst['name'] + '_gc_kin.mot')
        gc_grf = os.path.join(tmpdir, inst['name'] + '_gc_grf.mot')
        kin_out = os.path.join(tmpdir, inst['name'] + '_kin.mot')
        grf_out = os.path.join(tmpdir, inst['name'] + '_grf.mot')
        _resample_to_gait_cycle(inst['kin_mot'], inst['grf_mot'], gc_kin, gc_grf)
        _coloc_kin_and_grf(gc_kin, gc_grf, kin_out, grf_out)
        g = osimC.convertOsim2Gltf(inst['osim'], geom_dir, [kin_out], options)
        _append_force_arrows(g, grf_out, inst['color'])
        _recolor_gltf(g, inst['color'], inst.get('alpha', WCB_GLB_ALPHA))
        if merged is None:
            merged = g
        else:
            _merge_gltf(merged, g)

    merged.save(out_path)
    return out_path