# Copyright 2025 The Newton Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Tests for solver functions."""

import mujoco
import numpy as np
import warp as wp
from absl.testing import absltest
from absl.testing import parameterized

import mujoco_warp as mjw
from mujoco_warp import ConeType
from mujoco_warp import SolverType
from mujoco_warp import test_data
from mujoco_warp._src import solver

# tolerance for difference between MuJoCo and MJWarp solver calculations - mostly
# due to float precision
_TOLERANCE = 5e-3


def _assert_eq(a, b, name):
  tol = _TOLERANCE * 20  # avoid test noise
  err_msg = f"mismatch: {name}"
  np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)


class SolverTest(parameterized.TestCase):
  @parameterized.product(cone=tuple(ConeType), solver_=tuple(SolverType))
  def test_constraint_update(self, cone, solver_):
    """Tests _update_constraint function is correct."""
    for keyframe in range(3):
      mjm, mjd, m, d = test_data.fixture(
        "constraints.xml",
        keyframe=keyframe,
        overrides={"opt.solver": solver_, "opt.cone": cone, "opt.iterations": 0},
      )

      def cost(qacc):
        jaref = np.zeros(mjd.nefc, dtype=float)
        cost = np.zeros(1)
        mujoco.mj_mulJacVec(mjm, mjd, jaref, qacc)
        mujoco.mj_constraintUpdate(mjm, mjd, jaref - mjd.efc_aref, cost, 0)
        return cost

      mjd_cost = cost(mjd.qacc)

      # solve with 0 iterations just initializes constraints and costs and then exits
      d.efc.force.zero_()
      d.qfrc_constraint.zero_()
      ctx = solver.create_solver_context(m, d)
      solver._solve(m, d, ctx)

      # Get the ordering indices based on efc_force, efc_state for MJWarp
      nefc = d.nefc.numpy()[0]
      efc_force = d.efc.force.numpy()[0, :nefc]
      efc_state = d.efc.state.numpy()[0, :nefc]
      # Get the ordering indices based on efc_force, efc_state for MuJoCo
      mjd_efc_force = mjd.efc_force[:nefc]
      mjd_efc_state = mjd.efc_state[:nefc]

      # Create sorting keys using lexsort (more efficient for multiple keys)
      d_sort_indices = np.lexsort((efc_force, efc_state))
      mjd_sort_indices = np.lexsort((mjd_efc_force, mjd_efc_state))

      solver.init_context(m, d, ctx, grad=False)
      ctx_cost = ctx.cost.numpy()[0] - ctx.gauss.numpy()[0]
      qfrc_constraint = d.qfrc_constraint.numpy()[0]

      efc_sorted_force = efc_force[d_sort_indices]
      efc_sorted_state = efc_state[d_sort_indices]
      mjd_sorted_force = mjd_efc_force[mjd_sort_indices]
      mjd_sorted_state = mjd_efc_state[mjd_sort_indices]

      _assert_eq(efc_sorted_state, mjd_sorted_state, "efc_state")
      _assert_eq(efc_sorted_force, mjd_sorted_force, "efc_force")
      _assert_eq(ctx_cost, mjd_cost, "cost")
      _assert_eq(qfrc_constraint, mjd.qfrc_constraint, "qfrc_constraint")

  @parameterized.product(ls_parallel=(True, False), cone=(ConeType.PYRAMIDAL, ConeType.ELLIPTIC))
  def test_init_linesearch(self, ls_parallel, cone):
    """Test linesearch initialization.

    Parallel linesearch has separate prep kernels that write quad, quad_gauss, jv.
    Iterative linesearch fuses these in-kernel: quad_gauss is internal, quad is
    only written for elliptic cones.
    """
    for keyframe in range(3):
      mjm, mjd, m, d = test_data.fixture(
        "constraints.xml",
        keyframe=keyframe,
        overrides={
          "opt.iterations": 0,
          "opt.ls_iterations": 1,
          "opt.ls_parallel": ls_parallel,
          "opt.cone": cone,
        },
      )

      # One step to obtain more non-zeros results
      mjw.step(m, d)

      # Create a SolverContext to access internal solver arrays
      ctx = solver.create_solver_context(m, d)
      solver._solve(m, d, ctx)

      # Calculate target values
      nefc = d.nefc.numpy()[0]
      ctx_search_np = ctx.search.numpy()[0]
      efc_J_np = d.efc.J.numpy()[0][:nefc, : m.nv]
      ctx_gauss_np = ctx.gauss.numpy()[0]
      efc_Ma_np = d.efc.Ma.numpy()[0]
      ctx_Jaref_np = ctx.Jaref.numpy()[0][:nefc]
      efc_D_np = d.efc.D.numpy()[0][:nefc]
      qfrc_smooth_np = d.qfrc_smooth.numpy()[0]

      target_mv = np.zeros(mjm.nv)
      mujoco.mj_mulM(mjm, mjd, target_mv, ctx_search_np)
      target_jv = efc_J_np @ ctx_search_np
      target_quad_gauss = np.array(
        [
          ctx_gauss_np,
          np.dot(ctx_search_np, efc_Ma_np - qfrc_smooth_np),
          0.5 * np.dot(ctx_search_np, target_mv),
        ]
      )
      target_quad = np.transpose(
        np.vstack(
          [
            0.5 * ctx_Jaref_np * ctx_Jaref_np * efc_D_np,
            target_jv * ctx_Jaref_np * efc_D_np,
            0.5 * target_jv * target_jv * efc_D_np,
          ]
        )
      )

      # Reset and launch linesearch
      ctx.jv.zero_()
      ctx.quad.zero_()
      ctx.quad_gauss.zero_()
      step_size_cost = wp.empty((d.nworld, m.opt.ls_iterations), dtype=float)
      solver._linesearch(m, d, ctx, step_size_cost)

      # mv and jv are always written
      ctx_mv = ctx.mv.numpy()[0]
      ctx_jv = ctx.jv.numpy()[0]
      _assert_eq(ctx_mv, target_mv, "mv")
      _assert_eq(ctx_jv[:nefc], target_jv[:nefc], "jv")

      if ls_parallel and cone == ConeType.PYRAMIDAL:
        # Parallel pyramidal has separate prep kernels that write quad_gauss and quad
        # (Elliptic quad uses special quad1/quad2 format that target_quad doesn't compute)
        ctx_quad_gauss = ctx.quad_gauss.numpy()[0]
        ctx_quad = ctx.quad.numpy()[0]
        _assert_eq(ctx_quad_gauss, target_quad_gauss, "quad_gauss")
        _assert_eq(ctx_quad[:nefc], target_quad[:nefc], "quad")
      elif ls_parallel and cone == ConeType.ELLIPTIC:
        # Parallel elliptic: only check quad_gauss (quad uses special format)
        ctx_quad_gauss = ctx.quad_gauss.numpy()[0]
        _assert_eq(ctx_quad_gauss, target_quad_gauss, "quad_gauss")

  @parameterized.product(
    cone=(ConeType.PYRAMIDAL, ConeType.ELLIPTIC), jacobian=(mujoco.mjtJacobian.mjJAC_SPARSE, mujoco.mjtJacobian.mjJAC_DENSE)
  )
  def test_update_gradient_CG(self, cone, jacobian):
    """Test _update_gradient function is correct for the CG solver."""
    mjm, mjd, m, d = test_data.fixture(
      "humanoid/humanoid.xml",
      keyframe=0,
      overrides={"opt.cone": cone, "opt.solver": SolverType.CG, "opt.jacobian": jacobian, "opt.iterations": 0},
    )

    # Create SolverContext and initialize
    ctx = solver.create_solver_context(m, d)
    solver.init_context(m, d, ctx, grad=True)

    # Calculate Mgrad with Mujoco C
    mj_Mgrad = np.zeros(shape=(1, mjm.nv), dtype=float)
    mj_grad = np.tile(ctx.grad.numpy()[:, : mjm.nv], (1, 1))
    mujoco.mj_solveM(mjm, mjd, mj_Mgrad, mj_grad)

    ctx_Mgrad = ctx.Mgrad.numpy()[0, : mjm.nv]
    _assert_eq(ctx_Mgrad, mj_Mgrad[0], name="Mgrad")

  @parameterized.parameters(ConeType.PYRAMIDAL, ConeType.ELLIPTIC)
  def test_parallel_linesearch(self, cone):
    """Test that iterative and parallel linesearch leads to equivalent results."""
    _, _, m, d = test_data.fixture(
      "humanoid/humanoid.xml",
      qpos_noise=0.01,
      overrides={"opt.cone": cone, "opt.iterations": 50, "opt.ls_iterations": 50},
    )

    # One step to obtain more non-zeros results
    mjw.step(m, d)

    # Preparing for linesearch
    m.opt.iterations = 0
    mjw.fwd_velocity(m, d)
    mjw.fwd_acceleration(m, d, factorize=True)
    ctx = solver.create_solver_context(m, d)
    solver._solve(m, d, ctx)

    # Storing some initial values
    d_efc_Ma = d.efc.Ma.numpy().copy()
    ctx_Jaref = ctx.Jaref.numpy().copy()
    d_qacc = d.qacc.numpy().copy()

    # Launching iterative linesearch
    m.opt.ls_parallel = False
    step_size_cost = wp.empty((d.nworld, 0), dtype=float)
    solver._linesearch(m, d, ctx, step_size_cost)
    # Iterative computes alpha internally and directly updates outputs
    qacc_iterative = d.qacc.numpy().copy()
    Ma_iterative = d.efc.Ma.numpy().copy()
    Jaref_iterative = ctx.Jaref.numpy().copy()

    # Launching parallel linesearch with 50 testing points
    m.opt.ls_parallel = True
    m.opt.ls_iterations = 50
    d.efc.Ma = wp.array2d(d_efc_Ma)
    ctx.Jaref = wp.array2d(ctx_Jaref)
    d.qacc = wp.array2d(d_qacc)
    step_size_cost = wp.empty((d.nworld, m.opt.ls_iterations), dtype=float)
    solver._linesearch(m, d, ctx, step_size_cost)
    qacc_parallel = d.qacc.numpy().copy()
    Ma_parallel = d.efc.Ma.numpy().copy()
    Jaref_parallel = ctx.Jaref.numpy().copy()

    # Check that iterative and parallel linesearch produce equivalent outputs
    _assert_eq(qacc_iterative, qacc_parallel, name="qacc")
    _assert_eq(Ma_iterative, Ma_parallel, name="Ma")
    _assert_eq(Jaref_iterative, Jaref_parallel, name="Jaref")

  @parameterized.parameters(
    (ConeType.PYRAMIDAL, SolverType.CG, 10, 5, mujoco.mjtJacobian.mjJAC_DENSE, False),
    (ConeType.ELLIPTIC, SolverType.CG, 10, 5, mujoco.mjtJacobian.mjJAC_DENSE, False),
    (ConeType.PYRAMIDAL, SolverType.NEWTON, 5, 10, mujoco.mjtJacobian.mjJAC_DENSE, False),
    (ConeType.ELLIPTIC, SolverType.NEWTON, 5, 10, mujoco.mjtJacobian.mjJAC_DENSE, False),
    (ConeType.PYRAMIDAL, SolverType.NEWTON, 5, 64, mujoco.mjtJacobian.mjJAC_SPARSE, True),
    (ConeType.ELLIPTIC, SolverType.NEWTON, 5, 64, mujoco.mjtJacobian.mjJAC_SPARSE, True),
  )
  def test_solve(self, cone, solver_, iterations, ls_iterations, jacobian, ls_parallel):
    """Tests solve."""
    for keyframe in range(3):
      mjm, mjd, m, d = test_data.fixture(
        "constraints.xml",
        keyframe=keyframe,
        overrides={
          "opt.jacobian": jacobian,
          "opt.cone": cone,
          "opt.solver": solver_,
          "opt.iterations": iterations,
          "opt.ls_iterations": ls_iterations,
          "opt.ls_parallel": ls_parallel,
        },
      )

      mujoco.mj_forward(mjm, mjd)

      d.qacc.zero_()
      d.qfrc_constraint.zero_()
      d.efc.force.zero_()

      if solver_ == mujoco.mjtSolver.mjSOL_CG:
        mjw.factor_m(m, d)
      mjw.solve(m, d)

      def cost(qacc):
        jaref = np.zeros(mjd.nefc, dtype=float)
        cost = np.zeros(1)
        mujoco.mj_mulJacVec(mjm, mjd, jaref, qacc)
        mujoco.mj_constraintUpdate(mjm, mjd, jaref - mjd.efc_aref, cost, 0)
        return cost

      mj_cost = cost(mjd.qacc)
      mjwarp_cost = cost(d.qacc.numpy()[0])
      self.assertLessEqual(mjwarp_cost, mj_cost * 1.025)

      if m.opt.solver == mujoco.mjtSolver.mjSOL_NEWTON:
        _assert_eq(d.qacc.numpy()[0], mjd.qacc, "qacc")
        _assert_eq(d.qfrc_constraint.numpy()[0], mjd.qfrc_constraint, "qfrc_constraint")
        _assert_eq(d.efc.force.numpy()[0, : mjd.nefc], mjd.efc_force, "efc_force")

  @parameterized.parameters(
    (ConeType.PYRAMIDAL, SolverType.CG, 25, 5),
    (ConeType.PYRAMIDAL, SolverType.NEWTON, 2, 4),
  )
  def test_solve_batch(self, cone, solver_, iterations, ls_iterations):
    """Tests solve (batch)."""
    mjm0, mjd0, _, _ = test_data.fixture(
      "humanoid/humanoid.xml",
      keyframe=0,
      overrides={"opt.cone": cone, "opt.solver": solver_, "opt.iterations": iterations, "opt.ls_iterations": ls_iterations},
    )
    qacc_warmstart0 = mjd0.qacc_warmstart.copy()
    mujoco.mj_forward(mjm0, mjd0)
    mjd0.qacc_warmstart = qacc_warmstart0

    mjm1, mjd1, _, _ = test_data.fixture(
      "humanoid/humanoid.xml",
      keyframe=2,
      overrides={"opt.cone": cone, "opt.solver": solver_, "opt.iterations": iterations, "opt.ls_iterations": ls_iterations},
    )
    qacc_warmstart1 = mjd1.qacc_warmstart.copy()
    mujoco.mj_forward(mjm1, mjd1)
    mjd1.qacc_warmstart = qacc_warmstart1

    mjm2, mjd2, _, _ = test_data.fixture(
      "humanoid/humanoid.xml",
      keyframe=1,
      overrides={"opt.cone": cone, "opt.solver": solver_, "opt.iterations": iterations, "opt.ls_iterations": ls_iterations},
    )
    qacc_warmstart2 = mjd2.qacc_warmstart.copy()
    mujoco.mj_forward(mjm2, mjd2)
    mjd2.qacc_warmstart = qacc_warmstart2

    nefc_active = mjd0.nefc + mjd1.nefc + mjd2.nefc
    ne_active = mjd0.ne + mjd1.ne + mjd2.ne

    mjm, mjd, m, _ = test_data.fixture(
      "humanoid/humanoid.xml",
      overrides={"opt.cone": cone, "opt.solver": solver_, "opt.iterations": iterations, "opt.ls_iterations": ls_iterations},
    )
    d = mjw.put_data(mjm, mjd, nworld=3, njmax=2 * nefc_active)

    d.nefc = wp.array([nefc_active, nefc_active, nefc_active], dtype=wp.int32, ndim=1)
    d.ne = wp.array([ne_active, ne_active, ne_active], dtype=wp.int32, ndim=1)

    qacc_warmstart = np.vstack(
      [
        np.expand_dims(qacc_warmstart0, axis=0),
        np.expand_dims(qacc_warmstart1, axis=0),
        np.expand_dims(qacc_warmstart2, axis=0),
      ]
    )

    qM0 = np.zeros((mjm0.nv, mjm0.nv))
    mujoco.mj_fullM(mjm0, qM0, mjd0.qM)
    qM1 = np.zeros((mjm1.nv, mjm1.nv))
    mujoco.mj_fullM(mjm1, qM1, mjd1.qM)
    qM2 = np.zeros((mjm2.nv, mjm2.nv))
    mujoco.mj_fullM(mjm2, qM2, mjd2.qM)

    qM = np.vstack(
      [
        np.expand_dims(qM0, axis=0),
        np.expand_dims(qM1, axis=0),
        np.expand_dims(qM2, axis=0),
      ]
    )
    qacc_smooth = np.vstack(
      [
        np.expand_dims(mjd0.qacc_smooth, axis=0),
        np.expand_dims(mjd1.qacc_smooth, axis=0),
        np.expand_dims(mjd2.qacc_smooth, axis=0),
      ]
    )
    qfrc_smooth = np.vstack(
      [
        np.expand_dims(mjd0.qfrc_smooth, axis=0),
        np.expand_dims(mjd1.qfrc_smooth, axis=0),
        np.expand_dims(mjd2.qfrc_smooth, axis=0),
      ]
    )

    # Reshape the Jacobians
    efc_J0 = mjd0.efc_J.reshape((mjd0.nefc, mjm0.nv))
    efc_J1 = mjd1.efc_J.reshape((mjd1.nefc, mjm1.nv))
    efc_J2 = mjd2.efc_J.reshape((mjd2.nefc, mjm2.nv))

    efc_J_fill = np.zeros((3, d.njmax, m.nv))
    efc_J_fill[0, : mjd0.nefc, :] = efc_J0
    efc_J_fill[1, : mjd1.nefc, :] = efc_J1
    efc_J_fill[2, : mjd2.nefc, :] = efc_J2

    # Similarly for D and aref values
    efc_D0 = mjd0.efc_D[: mjd0.nefc]
    efc_D1 = mjd1.efc_D[: mjd1.nefc]
    efc_D2 = mjd2.efc_D[: mjd2.nefc]

    efc_D_fill = np.zeros((3, d.njmax))
    efc_D_fill[0, : mjd0.nefc] = efc_D0
    efc_D_fill[1, : mjd1.nefc] = efc_D1
    efc_D_fill[2, : mjd2.nefc] = efc_D2

    efc_aref0 = mjd0.efc_aref[: mjd0.nefc]
    efc_aref1 = mjd1.efc_aref[: mjd1.nefc]
    efc_aref2 = mjd2.efc_aref[: mjd2.nefc]

    efc_aref_fill = np.zeros((3, d.njmax))
    efc_aref_fill[0, : mjd0.nefc] = efc_aref0
    efc_aref_fill[1, : mjd1.nefc] = efc_aref1
    efc_aref_fill[2, : mjd2.nefc] = efc_aref2

    d.qacc_warmstart = wp.from_numpy(qacc_warmstart, dtype=wp.float32)
    d.qM = wp.from_numpy(qM, dtype=wp.float32)
    d.qacc_smooth = wp.from_numpy(qacc_smooth, dtype=wp.float32)
    d.qfrc_smooth = wp.from_numpy(qfrc_smooth, dtype=wp.float32)
    d.efc.J = wp.from_numpy(efc_J_fill, dtype=wp.float32)
    d.efc.D = wp.from_numpy(efc_D_fill, dtype=wp.float32)
    d.efc.aref = wp.from_numpy(efc_aref_fill, dtype=wp.float32)

    if solver_ == SolverType.CG:
      m0 = mjw.put_model(mjm0)
      d0 = mjw.put_data(mjm0, mjd0)
      mjw.factor_m(m0, d0)
      qLD0 = d0.qLD.numpy()

      m1 = mjw.put_model(mjm1)
      d1 = mjw.put_data(mjm1, mjd1)
      mjw.factor_m(m1, d1)
      qLD1 = d1.qLD.numpy()

      m2 = mjw.put_model(mjm2)
      d2 = mjw.put_data(mjm2, mjd2)
      mjw.factor_m(m2, d2)
      qLD2 = d2.qLD.numpy()

      qLD = np.vstack([qLD0, qLD1, qLD2])
      d.qLD = wp.from_numpy(qLD, dtype=wp.float32)

    d.qacc.zero_()
    d.qfrc_constraint.zero_()
    d.efc.force.zero_()
    solver.solve(m, d)

    def cost(m, d, qacc):
      jaref = np.zeros(d.nefc, dtype=float)
      cost = np.zeros(1)
      mujoco.mj_mulJacVec(m, d, jaref, qacc)
      mujoco.mj_constraintUpdate(m, d, jaref - d.efc_aref, cost, 0)
      return cost

    mj_cost0 = cost(mjm0, mjd0, mjd0.qacc)
    mjwarp_cost0 = cost(mjm0, mjd0, d.qacc.numpy()[0])
    self.assertLessEqual(mjwarp_cost0, mj_cost0 * 1.025)

    mj_cost1 = cost(mjm1, mjd1, mjd1.qacc)
    mjwarp_cost1 = cost(mjm1, mjd1, d.qacc.numpy()[1])
    self.assertLessEqual(mjwarp_cost1, mj_cost1 * 1.025)

    mj_cost2 = cost(mjm2, mjd2, mjd2.qacc)
    mjwarp_cost2 = cost(mjm2, mjd2, d.qacc.numpy()[2])
    self.assertLessEqual(mjwarp_cost2, mj_cost2 * 1.025)

    if m.opt.solver == SolverType.NEWTON:
      _assert_eq(d.qacc.numpy()[0], mjd0.qacc, "qacc0")
      _assert_eq(d.qacc.numpy()[1], mjd1.qacc, "qacc1")
      _assert_eq(d.qacc.numpy()[2], mjd2.qacc, "qacc2")

      _assert_eq(d.qfrc_constraint.numpy()[0], mjd0.qfrc_constraint, "qfrc_constraint0")
      _assert_eq(d.qfrc_constraint.numpy()[1], mjd1.qfrc_constraint, "qfrc_constraint1")
      _assert_eq(d.qfrc_constraint.numpy()[2], mjd2.qfrc_constraint, "qfrc_constraint2")

      # Get world 0 forces - equality constraints at start, inequality constraints later
      nieq0 = mjd0.nefc - mjd0.ne
      nieq1 = mjd1.nefc - mjd1.ne
      nieq2 = mjd2.nefc - mjd2.ne
      world0_eq_forces = d.efc.force.numpy()[0, : mjd0.ne]
      world0_ineq_forces = d.efc.force.numpy()[0, ne_active : ne_active + nieq0]
      world0_forces = np.concatenate([world0_eq_forces, world0_ineq_forces])
      _assert_eq(world0_forces, mjd0.efc_force, "efc_force0")

      # Get world 1 forces
      world1_eq_forces = d.efc.force.numpy()[1, : mjd1.ne]
      world1_ineq_forces = d.efc.force.numpy()[1, ne_active : ne_active + nieq1]
      world1_forces = np.concatenate([world1_eq_forces, world1_ineq_forces])
      _assert_eq(world1_forces, mjd1.efc_force, "efc_force1")

      # Get world 2 forces
      world2_eq_forces = d.efc.force.numpy()[2, : mjd2.ne]
      world2_ineq_forces = d.efc.force.numpy()[2, ne_active : ne_active + nieq2]
      world2_forces = np.concatenate([world2_eq_forces, world2_ineq_forces])
      _assert_eq(world2_forces, mjd2.efc_force, "efc_force2")

  def test_frictionloss(self):
    """Tests solver with frictionloss."""
    for keyframe in range(3):
      _, mjd, m, d = test_data.fixture("constraints.xml", keyframe=keyframe)
      mjw.solve(m, d)

      _assert_eq(d.nf.numpy()[0], mjd.nf, "nf")
      _assert_eq(d.qacc.numpy()[0], mjd.qacc, "qacc")
      _assert_eq(d.qfrc_constraint.numpy()[0], mjd.qfrc_constraint, "qfrc_constraint")
      _assert_eq(d.efc.force.numpy()[0, : mjd.nefc], mjd.efc_force, "efc_force")

  def test_parallel_linesearch_threads_per_efc_gt_1(self):
    """Test parallel linesearch with threads_per_efc > 1."""
    xml = """
    <mujoco>
      <worldbody>
        <body>
          <freejoint/>
          <geom size="0.1"/>
        </body>
        <body>
          <freejoint/>
          <geom size="0.1"/>
        </body>
        <body>
          <freejoint/>
          <geom size="0.1"/>
        </body>
        <body>
          <freejoint/>
          <geom size="0.1"/>
        </body>
        <body>
          <freejoint/>
          <geom size="0.1"/>
        </body>
        <body>
          <freejoint/>
          <geom size="0.1"/>
        </body>
        <body>
          <freejoint/>
          <geom size="0.1"/>
        </body>
        <body>
          <freejoint/>
          <geom size="0.1"/>
        </body>
        <body>
          <freejoint/>
          <geom size="0.1"/>
        </body>
      </worldbody>
    </mujoco>
    """
    mjm, mjd, m, d = test_data.fixture(xml=xml)
    self.assertEqual(mjm.nv, 54)  # 9 freejoints * 6 dofs each

    # parallel linesearch path with nv > 50 -> threads_per_efc > 1
    m.opt.ls_parallel = True
    m.opt.iterations = 1
    m.opt.ls_iterations = 1
    mjw.step(m, d)


if __name__ == "__main__":
  wp.init()
  absltest.main()
