"""Reference implementation (numpy) of the deployable walk steering rule used with the bases in this repo.

Per control step (50 Hz), before the policy forward:
    d_plan  = plan pelvis xy(t) - plan pelvis xy(0)           (from the motion plan)
    d_robot = robot pelvis xy(t) - robot pelvis xy(0)         (from odometry; no world alignment needed)
    e       = R(plan heading)^T (d_plan - d_robot)            (forward, lateral) in m; + = robot behind / right of plan
    v       = clip(K * e, -VMAX, VMAX)                        K = 160 /m, VMAX = 15 (cm-equivalent units of A)
    v       = 0 if |e| > 0.5 m (guard) or after the plan's planned lift (walk gate)
    offset  = v @ A                                           A [2, 64]: token change per +1 cm forward / lateral
The offset is added to SONIC's post-FSQ tokens of the g1 encoder (64-D) before they enter the decoder.
"""
import numpy as np

K, VMAX, EMAX = 160.0, 15.0, 0.5


def heading_frame(d, yaw):
    c, s = np.cos(yaw), np.sin(yaw)
    return np.array([c * d[0] + s * d[1], -s * d[0] + c * d[1]])


def token_offset(A, plan_xy_t, plan_xy_0, plan_yaw_t, robot_xy_t, robot_xy_0, walking):
    """A [2, 64]; positions in m (world or odometry frame, each consistent with itself); walking: before the plan's lift."""
    e = heading_frame((plan_xy_t - plan_xy_0) - (robot_xy_t - robot_xy_0), plan_yaw_t)
    if not walking or np.linalg.norm(e) > EMAX:
        return np.zeros(A.shape[1], np.float32)
    v = np.clip(K * e, -VMAX, VMAX)
    return (v @ A).astype(np.float32)


if __name__ == "__main__":
    A = np.load("bases/A_R15_omomo_nopca.npz")["A"]
    print(token_offset(A, np.array([1.0, 0.1]), np.zeros(2), 0.0, np.array([0.95, 0.0]), np.zeros(2), True)[:8])