import ring
import mediapy as media
import jax
from diodem import load_data

sys_str = """
<x_xy>
<worldbody>
<geom dim="0.1" type="xyz"></geom>
<body joint="free" name="seg" pos="0 0 .5">
<geom color="dustin_exp_blue" dim="0.15 0.075 0.05" mass="0.2" pos="0.03 0 0" type="box"></geom>
<body joint="frozen" name="imu" pos="0.0 0.0 0.03">
<geom color="dustin_exp_orange" dim="0.05 0.03 0.02" mass="0.1" type="box"></geom>
</body>
</body>
</worldbody>
</x_xy>
"""

sys = ring.System.create(sys_str)

(X, y), (key, q, x, _) = ring.RCMG(sys, keep_output_extras=True).to_list()[0]

media.show_video(sys.render(x, width=640, height=480, camera="target", render_every_nth=4), fps=25)

exp_data = load_data(exp_id=1, motion_start="fast")

exp_data.keys()

exp_data["seg1"].keys()

segment = "seg2"
omc_data_sys = {
    "seg": {
        "pos": exp_data[segment]["marker1"],
        "quat": exp_data[segment]["quat"],
    },
    "imu": {
        "quat": exp_data[segment]["quat"],
    }
}

omc_data_sys

x = ring.sim2real.xs_from_raw(sys, omc_data_sys)

# vectorize this function over time
@jax.vmap
def update_position_vector_of_imu(x):
    state = ring.State.create(sys, x=x)
    # populate minimal coordinates `state.q` from maximal coordinates `state.x`
    state = ring.algorithms.inverse_kinematics(sys, state)
    # re-calculate maximal coordiantes `state.x` from minimal coordinates `state.q`
    # this uses the position vector specified in the system (and so the xml file)
    # to produce an offset between IMu and segment geom box
    _, state = ring.algorithms.forward_kinematics(sys, state)
    return state.x

x = update_position_vector_of_imu(x)

media.show_video(sys.render(x, width=640, height=480, camera="target", render_every_nth=4), fps=25)