mujoco

// Copyright 2022 DeepMind Technologies Limited
//
// 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 engine/engine_derivative.c.

#include "src/engine/engine_derivative.h"

#include <cmath>
#include <cstddef>
#include <random>
#include <string>
#include <vector>

#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <mujoco/mjmodel.h>
#include <mujoco/mujoco.h>
#include "src/engine/engine_core_smooth.h"
#include "src/engine/engine_core_util.h"
#include "src/engine/engine_derivative_fd.h"
#include "src/engine/engine_forward.h"
#include "src/engine/engine_io.h"
#include "src/engine/engine_util_blas.h"
#include "src/engine/engine_util_sparse.h"
#include "test/fixture.h"

namespace mujoco {
namespace {

using ::std::vector;
using ::testing::DoubleNear;
using ::testing::Each;
using ::testing::Eq;
using ::testing::HasSubstr;
using ::testing::NotNull;
using ::testing::Pointwise;
using DerivativeTest = MujocoTest;

// errors smaller than this are ignored
static const mjtNum absolute_tolerance = MjTol(1e-9, 1e-3);

// corrected relative error
static mjtNum RelativeError(mjtNum a, mjtNum b) {
  mjtNum nominator = mjMAX(0, mju_abs(a - b) - absolute_tolerance);
  mjtNum denominator = (mju_abs(a) + mju_abs(b) + absolute_tolerance);
  return nominator / denominator;
}

// expect two 2D arrays to have elementwise relative error smaller than eps
// return maximum absolute error
static mjtNum CompareMatrices(mjtNum* Actual, mjtNum* Expected, int nrow,
                              int ncol, mjtNum eps) {
  mjtNum max_error = 0;
  for (int i = 0; i < nrow; i++) {
    for (int j = 0; j < ncol; j++) {
      mjtNum actual = Actual[i * ncol + j];
      mjtNum expected = Expected[i * ncol + j];
      EXPECT_LT(RelativeError(actual, expected), eps)
          << "error at position (" << i << ", " << j << ")"
          << "\nexpected = " << expected << "\nactual   = " << actual
          << "\ndiff     = " << expected - actual;
      max_error = mjMAX(mju_abs(actual - expected), max_error);
    }
  }
  return max_error;
}

static const char* const kEnergyConservingPendulumPath =
    "engine/testdata/derivative/energy_conserving_pendulum.xml";
static const char* const kTumblingThinObjectPath =
    "engine/testdata/derivative/tumbling_thin_object.xml";
static const char* const kTumblingThinObjectEllipsoidPath =
    "engine/testdata/derivative/tumbling_thin_object_ellipsoid.xml";
static const char* const kDampedActuatorsPath =
    "engine/testdata/derivative/damped_actuators.xml";
static const char* const kDamperActuatorsPath =
    "engine/testdata/actuation/damper.xml";
static const char* const kDampedPendulumPath =
    "engine/testdata/derivative/damped_pendulum.xml";
static const char* const kLinearPath = "engine/testdata/derivative/linear.xml";
static const char* const kDCMotorPath =
    "engine/testdata/derivative/dcmotor.xml";
static const char* const kModelPath = "testdata/model.xml";
static const char* const kSleepEqualityPath =
    "engine/testdata/sleep/equality.xml";

// compare analytic and finite-difference d_smooth/d_qvel
TEST_F(DerivativeTest, SmoothDvel) {
  // run test on all models
  for (const char* local_path :
       {kEnergyConservingPendulumPath, kTumblingThinObjectPath,
        kDampedActuatorsPath, kDamperActuatorsPath, kDCMotorPath}) {
    const std::string xml_path = GetTestDataFilePath(local_path);
    char error[1024] = "";
    mjModel* model =
        mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
    ASSERT_THAT(model, testing::NotNull()) << "Failed to load model: " << error;
    int nD = model->nD;
    mjData* data = mj_makeData(model);

    for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) {
      // set sparsity
      model->opt.jacobian = sparsity;

      // take 100 steps so we have some velocities, then call forward
      mj_resetData(model, data);
      if (model->nu) {
        data->ctrl[0] = 0.1;
      }
      for (int i = 0; i < 100; i++) {
        mj_step(model, data);
      }
      mj_forward(model, data);

      // construct sparse structure in d->D_xxx, compute analytical qDeriv
      mju_zero(data->qDeriv, nD);
      mjd_smooth_vel(model, data, /*flg_bias=*/true);

      // expect derivatives to be non-zero, make copy of qDeriv as a vector
      EXPECT_GT(mju_norm(data->qDeriv, nD), 0);
      vector<mjtNum> qDerivAnalytic = AsVector(data->qDeriv, nD);

      // compute finite-difference derivatives
      mjtNum eps = MjEps(1e-7, 1e-3);
      mju_zero(data->qDeriv, nD);
      mjd_smooth_velFD(model, data, eps);

      // expect FD and analytic derivatives to be numerically different
      EXPECT_NE(mju_norm(data->qDeriv, nD),
                mju_norm(qDerivAnalytic.data(), nD));

      // expect FD and analytic derivatives to be similar to eps precision
      EXPECT_THAT(AsVector(data->qDeriv, nD),
                  Pointwise(MjNear(1e-7, 3e-3), qDerivAnalytic));
    }
    mj_deleteData(data);
    mj_deleteModel(model);
  }
}

// mjd_freeBias_vel: 6x6 bias-derivative block for a standalone free body
//   validated against mjd_rne_vel and against finite-differenced mj_rne
TEST_F(DerivativeTest, FreeBiasVel) {
  // free body with offset CoM, rotated inertia, non-identity orientation
  static constexpr char xml[] = R"(
  <mujoco>
    <worldbody>
      <body pos="0.1 -0.2 0.3" euler="20 -30 40">
        <freejoint/>
        <geom type="box" size=".1 .2 .3" mass="2" pos=".04 -.02 .03" euler="10 20 30"/>
      </body>
    </worldbody>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  MjDataPtr data = MakeData(model);
  mjModel* m = model.get();
  mjData* d = data.get();

  // set fast, fully populated velocity
  mjtNum qvel[6] = {0.4, -0.3, 0.2, 5, -3, 2};
  mju_copy(d->qvel, qvel, 6);
  mj_forward(m, d);

  // analytic block
  mjtNum B[36];
  mjd_freeBias_vel(m, d, /*jnt=*/0, B);

  // linear columns are zero by construction
  for (int r = 0; r < 6; r++) {
    for (int c = 0; c < 3; c++) {
      EXPECT_EQ(B[6 * r + c], 0);
    }
  }

  // compare with mjd_rne_vel: B == -(qDeriv(flg_bias=1) - qDeriv(flg_bias=0))
  mju_zero(d->qDeriv, m->nD);
  mjd_smooth_vel(m, d, /*flg_bias=*/1);
  vector<mjtNum> qDeriv_bias = AsVector(d->qDeriv, m->nD);
  mju_zero(d->qDeriv, m->nD);
  mjd_smooth_vel(m, d, /*flg_bias=*/0);
  for (int r = 0; r < 6; r++) {
    int rowadr = m->D_rowadr[r];
    ASSERT_EQ(m->D_rownnz[r], 6);
    for (int k = 0; k < 6; k++) {
      int c = m->D_colind[rowadr + k];
      mjtNum rne_val = -(qDeriv_bias[rowadr + k] - d->qDeriv[rowadr + k]);
      EXPECT_NEAR(B[6 * r + c], rne_val, MjTol(1e-14, 1e-6))
          << "mismatch at (" << r << ", " << c << ")";
    }
  }

  // compare with central finite differences of mj_rne
  mjtNum eps = MjEps(1e-6, 1e-3);
  for (int c = 0; c < 6; c++) {
    mjtNum bias_plus[6], bias_minus[6];

    d->qvel[c] = qvel[c] + eps;
    mj_comVel(m, d);
    mj_rne(m, d, /*flg_acc=*/0, bias_plus);

    d->qvel[c] = qvel[c] - eps;
    mj_comVel(m, d);
    mj_rne(m, d, /*flg_acc=*/0, bias_minus);

    d->qvel[c] = qvel[c];

    for (int r = 0; r < 6; r++) {
      mjtNum fd = (bias_plus[r] - bias_minus[r]) / (2 * eps);
      EXPECT_NEAR(B[6 * r + c], fd, MjTol(1e-7, 1e-2))
          << "FD mismatch at (" << r << ", " << c << ")";
    }
  }
}

// fixed descendants contribute inertia about the free joint, not their own body
// frames
TEST_F(DerivativeTest, FreeBiasVelFixedDescendants) {
  static constexpr char xml[] = R"(
  <mujoco>
    <default><geom contype="0" conaffinity="0"/></default>
    <worldbody>
      <body pos=".1 -.2 .3" euler="20 -30 40">
        <freejoint/>
        <geom type="box" size=".1 .2 .3" mass="2" pos=".04 -.02 .03" euler="10 20 30"/>
        <body pos=".2 -.1 .3" euler="30 20 -10">
          <geom type="box" size=".2 .1 .1" mass="1" pos=".03 .01 -.02"/>
          <body pos="-.1 .2 .1" euler="10 -20 30">
            <geom type="box" size=".1 .1 .2" mass=".5"/>
          </body>
        </body>
      </body>
    </worldbody>
  </mujoco>
  )";
  char error[1024];
  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  ASSERT_EQ(model->nbody, 4);
  MjDataPtr data = MakeData(model);
  mjModel* m = model.get();
  mjData* d = data.get();
  mjtNum qvel[6] = {0.4, -0.3, 0.2, 5, -3, 2};
  mju_copy(d->qvel, qvel, 6);
  mj_forward(m, d);

  mjtNum B[36];
  mjd_freeBias_vel(m, d, /*jnt=*/0, B);

  // compare the aggregate derivative with the full recursive Newton-Euler
  // derivative
  mju_zero(d->qDeriv, m->nD);
  mjd_smooth_vel(m, d, /*flg_bias=*/1);
  for (int r = 0; r < 6; r++) {
    ASSERT_EQ(m->D_rownnz[r], 6);
    for (int k = 0; k < 6; k++) {
      int adr = m->D_rowadr[r] + k;
      EXPECT_NEAR(B[6 * r + m->D_colind[adr]], -d->qDeriv[adr],
                  MjTol(1e-14, 1e-5));
    }
  }

  // independent central finite differences, including the zero linear-velocity
  // columns
  mjtNum eps = MjEps(1e-6, 1e-3);
  for (int c = 0; c < 6; c++) {
    mjtNum plus[6], minus[6];
    d->qvel[c] = qvel[c] + eps;
    mj_comVel(m, d);
    mj_rne(m, d, /*flg_acc=*/0, plus);
    d->qvel[c] = qvel[c] - eps;
    mj_comVel(m, d);
    mj_rne(m, d, /*flg_acc=*/0, minus);
    d->qvel[c] = qvel[c];
    for (int r = 0; r < 6; r++) {
      EXPECT_NEAR(B[6 * r + c], (plus[r] - minus[r]) / (2 * eps),
                  MjTol(1e-7, 1e-2));
    }
  }
}

// fixed-body fusion preserves the bias derivative of a massless free root
TEST_F(DerivativeTest, FreeBiasVelFusedInertia) {
  static constexpr char xml[] = R"(
  <mujoco>
    <compiler fusestatic="false" alignfree="false"/>
    <default><geom contype="0" conaffinity="0"/></default>
    <worldbody>
      <body pos=".1 -.2 .3" euler="20 -30 40">
        <freejoint/>
        <body pos=".2 -.1 .3" euler="30 20 -10">
          <geom type="box" size=".2 .1 .1" mass="1" pos=".03 .01 -.02"/>
          <body pos="-.1 .2 .1" euler="10 -20 30">
            <geom type="box" size=".1 .1 .2" mass=".5"/>
          </body>
        </body>
        <body pos="-.1 .3 -.2" euler="-20 10 40">
          <geom type="box" size=".1 .2 .3" mass="2" pos=".04 -.02 .03"/>
        </body>
      </body>
    </worldbody>
  </mujoco>
  )";
  std::string fused_xml = xml;
  fused_xml.replace(fused_xml.find("false"), 5, "true");
  char error[1024];
  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  MjModelPtr fused =
      LoadModelFromString(fused_xml.c_str(), error, sizeof(error));
  ASSERT_THAT(fused.get(), NotNull()) << error;
  ASSERT_EQ(model->nbody, 5);
  ASSERT_EQ(fused->nbody, 2);
  ASSERT_EQ(model->body_mass[1], 0);
  MjDataPtr data = MakeData(model);
  MjDataPtr fused_data = MakeData(fused);

  mjtNum qvel[6] = {0.4, -0.3, 0.2, 5, -3, 2};
  mju_copy(data->qvel, qvel, 6);
  mju_copy(fused_data->qvel, qvel, 6);
  mj_forward(model.get(), data.get());
  mj_forward(fused.get(), fused_data.get());

  mjtNum B[36], B_fused[36];
  mjd_freeBias_vel(model.get(), data.get(), /*jnt=*/0, B);
  mjd_freeBias_vel(fused.get(), fused_data.get(), /*jnt=*/0, B_fused);
  EXPECT_GT(mju_norm(B, 36), 1);

  // fusing diagonalizes the inertia at compiler precision
  EXPECT_THAT(AsVector(B, 36),
              Pointwise(MjNear(1e-5, 1e-5), AsVector(B_fused, 36)));
}

// a jointed descendant must still exclude the free root from the local six-DOF
// solve
TEST_F(DerivativeTest, FreeMhatRejectsArticulatedSubtree) {
  static constexpr char xml[] = R"(
  <mujoco>
    <worldbody>
      <body>
        <freejoint/>
        <geom size=".1"/>
        <body pos="0 0 1">
          <joint/>
          <geom size=".1"/>
        </body>
      </body>
    </worldbody>
  </mujoco>
  )";
  char error[1024];
  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  MjDataPtr data = MakeData(model);
  mj_forward(model.get(), data.get());
  mjtNum A[36];
  for (int discrete : {0, 1}) {
    EXPECT_EQ(mjd_freeMhat(model.get(), data.get(), /*jnt=*/0,
                           model->opt.timestep, A, discrete),
              0);
  }
}

// disabled actuators do not contribute to d_qfrc_actuator/d_qvel
TEST_F(DerivativeTest, DisabledActuators) {
  // model with only a position actuator
  static constexpr char xml1[] = R"(
  <mujoco>
    <option integrator="implicitfast"/>

    <worldbody>
      <body>
        <joint name="joint" type="slide"/>
        <geom size=".1"/>
      </body>
    </worldbody>

    <actuator>
      <position joint="joint" group="1" kp="2000" kv="200"/>
    </actuator>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr m1 = LoadModelFromString(xml1, error, sizeof(error));
  ASSERT_THAT(m1.get(), NotNull()) << error;
  MjDataPtr d1 = MakeData(m1);

  d1->ctrl[0] = 6;
  while (d1->time < 1) mj_step(m1.get(), d1.get());

  // model with a position actuator and an intvelocity actuator
  static constexpr char xml2[] = R"(
  <mujoco>
    <option integrator="implicitfast" actuatorgroupdisable="2"/>

    <worldbody>
      <body>
        <joint name="joint" type="slide"/>
        <geom size=".1"/>
      </body>
    </worldbody>

    <actuator>
      <position joint="joint" group="1" kp="2000" kv="200"/>
      <intvelocity joint="joint" group="2" kp="2000" kv="200" actrange="-6 6"/>
    </actuator>
  </mujoco>
  )";

  MjModelPtr m2 = LoadModelFromString(xml2);
  MjDataPtr d2 = MakeData(m2);

  d2->ctrl[0] = 6;
  d2->ctrl[1] = 6;

  while (d2->time < 1) mj_step(m2.get(), d2.get());

  // expect same qvel in both models
  EXPECT_EQ(d1->qvel[0], d2->qvel[0]);
}

// actuator order has no effect
TEST_F(DerivativeTest, ActuatorOrder) {
  // model with stateful actuator first
  static constexpr char xml1[] = R"(
  <mujoco>
    <option integrator="implicitfast"/>

    <worldbody>
      <body>
        <joint name="0" type="slide" range="-1 1"/>
        <geom size=".1"/>
      </body>
      <body pos="1 0 0">
        <joint name="1" type="slide" range="-1 1"/>
        <geom size=".1"/>
      </body>
    </worldbody>

    <actuator>
      <muscle joint="0" ctrlrange="0 6"/>
      <damper joint="1" kv="200" ctrlrange="0 6"/>
    </actuator>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr m1 = LoadModelFromString(xml1, error, sizeof(error));
  ASSERT_THAT(m1.get(), NotNull()) << "Failed to load model: " << error;
  MjDataPtr d1 = MakeData(m1);

  d1->ctrl[0] = 6;
  d1->ctrl[1] = 6;

  while (d1->time < 1) mj_step(m1.get(), d1.get());

  // model with stateful actuator second
  static constexpr char xml2[] = R"(
  <mujoco>
    <option integrator="implicitfast"/>

    <worldbody>
      <body>
        <joint name="0" type="slide" range="-1 1"/>
        <geom size=".1"/>
      </body>
      <body pos="1 0 0">
        <joint name="1" type="slide" range="-1 1"/>
        <geom size=".1"/>
      </body>
    </worldbody>

    <actuator>
      <damper joint="1" kv="200" ctrlrange="0 6"/>
      <muscle joint="0" ctrlrange="0 6"/>
    </actuator>
  </mujoco>
  )";

  MjModelPtr m2 = LoadModelFromString(xml2, error, sizeof(error));
  ASSERT_THAT(m2.get(), NotNull()) << "Failed to load model: " << error;
  MjDataPtr d2 = MakeData(m2);

  d2->ctrl[0] = 6;
  d2->ctrl[1] = 6;

  while (d2->time < 1) mj_step(m2.get(), d2.get());

  // expect same qvel in both models
  EXPECT_EQ(d1->qvel[0], d2->qvel[0]);
  EXPECT_EQ(d1->qvel[1], d2->qvel[1]);
}

// compare analytic and fin-diff d_qfrc_passive/d_qvel
TEST_F(DerivativeTest, PassiveDvel) {
  for (const char* local_path :
       {kTumblingThinObjectPath, kTumblingThinObjectEllipsoidPath}) {
    // load model
    const std::string xml_path = GetTestDataFilePath(local_path);
    char error[1024];
    mjModel* model =
        mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
    ASSERT_THAT(model, NotNull()) << error;
    int nD = model->nD;
    mjData* data = mj_makeData(model);
    // allocate Jacobians
    mjtNum* qDerivAnalytic = (mjtNum*)mju_malloc(sizeof(mjtNum) * nD);
    mjtNum* qDerivFD = (mjtNum*)mju_malloc(sizeof(mjtNum) * nD);

    for (mjtJacobian sparsity : {mjJAC_DENSE, mjJAC_SPARSE}) {
      // set sparsity
      model->opt.jacobian = sparsity;

      // take 100 steps so we have some velocities, then call forward
      mj_resetData(model, data);
      for (int i = 0; i < 100; i++) {
        mj_step(model, data);
      }
      mj_forward(model, data);

      // get analytic derivatives
      mju_zero(data->qDeriv, model->nD);
      mjd_passive_vel(model, data);
      mju_copy(qDerivAnalytic, data->qDeriv, nD);

      // clear qDeriv, get finite-difference derivatives
      mju_zero(data->qDeriv, nD);
      mju_zero(qDerivFD, nD);
      mjtNum eps = MjEps(1e-6, 1e-4);
      mjd_passive_velFD(model, data, eps);

      // expect FD and analytic derivatives to be similar to tol precision
      EXPECT_THAT(AsVector(data->qDeriv, nD),
                  Pointwise(MjNear(1e-6, 1e-4), AsVector(qDerivAnalytic, nD)));
    }

    mju_free(qDerivFD);
    mju_free(qDerivAnalytic);
    mj_deleteData(data);
    mj_deleteModel(model);
  }
}

// ----------------------- derivatives of mj_step() ----------------------------

// mj_stepSkip computes the same next state as mj_step
TEST_F(DerivativeTest, StepSkip) {
  const std::string xml_path = GetTestDataFilePath(kDampedPendulumPath);
  char error[1024];
  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
  ASSERT_THAT(model, NotNull()) << error;
  mjData* data = mj_makeData(model);
  int nq = model->nq;
  int nv = model->nv;

  // disable warm-starts so we don't need to save qacc_warmstart
  model->opt.disableflags |= mjDSBL_WARMSTART;

  for (const mjtIntegrator integrator :
       {mjINT_EULER, mjINT_IMPLICIT, mjINT_IMPLICITFAST, mjINT_DISCRETE}) {
    model->opt.integrator = integrator;

    // reset, take 20 steps
    mj_resetData(model, data);
    for (int i = 0; i < 20; i++) {
      mj_step(model, data);
    }

    // denormalize the quat, just to see that it doesn't make a difference
    for (int j = 0; j < model->njnt; j++) {
      if (model->jnt_type[j] == mjJNT_BALL) {
        int adr = model->jnt_qposadr[j];
        for (int k = 0; k < 4; k++) {
          data->qpos[adr + k] *= 8;
        }
      }
    }

    // save state
    vector<mjtNum> qpos = AsVector(data->qpos, nq);
    vector<mjtNum> qvel = AsVector(data->qvel, nv);

    // take one more step, save next state
    mj_step(model, data);
    vector<mjtNum> qpos_next = AsVector(data->qpos, nq);
    vector<mjtNum> qvel_next = AsVector(data->qvel, nv);

    // reset state, take step again, compare (assert mj_step is deterministic)
    mju_copy(data->qpos, qpos.data(), nq);
    mju_copy(data->qvel, qvel.data(), nv);
    mj_step(model, data);
    EXPECT_THAT(AsVector(data->qpos, nq), Pointwise(Eq(), qpos_next));
    EXPECT_THAT(AsVector(data->qvel, nv), Pointwise(Eq(), qvel_next));

    // reset state, change ctrl, call mj_stepSkip, save next state
    mju_copy(data->qpos, qpos.data(), nq);
    mju_copy(data->qvel, qvel.data(), nv);
    data->ctrl[0] = 1;
    mj_stepSkip(model, data, mjSTAGE_VEL, 0);  // skipping both POS and VEL
    vector<mjtNum> qpos_next_dctrl = AsVector(data->qpos, nq);
    vector<mjtNum> qvel_next_dctrl = AsVector(data->qvel, nv);

    // reset state (ctrl remains unchanged), call full mj_step, compare
    mju_copy(data->qpos, qpos.data(), nq);
    mju_copy(data->qvel, qvel.data(), nv);
    mj_step(model, data);
    EXPECT_THAT(AsVector(data->qpos, nq), Pointwise(Eq(), qpos_next_dctrl));
    EXPECT_THAT(AsVector(data->qvel, nv), Pointwise(Eq(), qvel_next_dctrl));

    // reset state, change velocity, call mj_stepSkip, save next state
    mju_copy(data->qpos, qpos.data(), nq);
    mju_copy(data->qvel, qvel.data(), nv);
    data->qvel[0] += 1;
    mj_stepSkip(model, data, mjSTAGE_POS, 0);  // skipping POS
    vector<mjtNum> qpos_next_dvel = AsVector(data->qpos, nq);
    vector<mjtNum> qvel_next_dvel = AsVector(data->qvel, nv);

    // reset state, change velocity, call full mj_step, compare
    mju_copy(data->qpos, qpos.data(), nq);
    mju_copy(data->qvel, qvel.data(), nv);
    data->qvel[0] += 1;
    mj_step(model, data);
    EXPECT_THAT(AsVector(data->qpos, nq), Pointwise(Eq(), qpos_next_dvel));
    EXPECT_THAT(AsVector(data->qvel, nv), Pointwise(Eq(), qvel_next_dvel));
  }

  mj_deleteData(data);
  mj_deleteModel(model);
}

// Analytic transition matrices for linear dynamical system xn = A*x + B*u
//   given modified mass matrix H (`data->qH`) and
//   Ac = H^-1 [diag(-stiffness) diag(-damping)]
//   we have
//   A  = eye(2*nv) + dt [dt*Ac + [zeros(3) eye(3)]; Ac]
//   given the moment arm matrix K (`data->actuator_moment`) and Bc = H^-1 K
//   B  = dt*[Bc*dt; Bc]
static void LinearSystem(const mjModel* m, mjData* d, mjtNum* A, mjtNum* B) {
  int nv = m->nv, nu = m->nu;
  mjtNum dt = m->opt.timestep;
  mj_markStack(d);

  // === state-transition matrix A
  if (A) {
    mjtNum* Ac = mj_stackAllocNum(d, 2 * nv * nv);
    // Ac = H^-1 [diag(-stiffness) diag(-damping)]
    mju_zero(Ac, 2 * nv * nv);
    for (int i = 0; i < nv; i++) {
      Ac[i * nv + i] = -m->jnt_stiffness[i];
      Ac[nv * nv + i * nv + i] = -m->dof_damping[i];
    }
    mj_solveLD(Ac, d->qH, d->qHDiagInv, nv, 2 * nv, m->M_rownnz, m->M_rowadr,
               m->M_colind, nullptr);

    // A = [dt*Ac; Ac]
    mju_transpose(A, Ac, 2 * nv, nv);
    mju_scl(A, A, dt, nv * 2 * nv);
    mju_transpose(A + 2 * nv * nv, Ac, 2 * nv, nv);

    // Add eye(nv) to top right quadrant of A
    for (int i = 0; i < nv; i++) {
      A[i * 2 * nv + nv + i] += 1;
    }

    // A *= dt
    mju_scl(A, A, dt, 2 * nv * 2 * nv);

    // A += eye(2*nv)
    for (int i = 0; i < 2 * nv; i++) {
      A[i * 2 * nv + i] += 1;
    }
  }

  // === control-transition matrix B
  if (B) {
    mjtNum* Bc = mj_stackAllocNum(d, nu * nv);
    mjtNum* BcT = mj_stackAllocNum(d, nv * nu);
    mju_sparse2dense(Bc, d->actuator_moment, nu, nv, d->moment_rownnz,
                     d->moment_rowadr, d->moment_colind);
    mj_solveLD(Bc, d->qH, d->qHDiagInv, nv, nu, m->M_rownnz, m->M_rowadr,
               m->M_colind, nullptr);
    mju_transpose(BcT, Bc, nu, nv);
    mju_scl(B, BcT, dt * dt, nu * nv);
    mju_scl(B + nu * nv, BcT, dt, nu * nv);
  }

  mj_freeStack(d);
}

// compare FD derivatives to analytic derivatives of linear dynamical system
TEST_F(DerivativeTest, LinearSystem) {
  const std::string xml_path = GetTestDataFilePath(kLinearPath);
  char error[1024];
  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
  ASSERT_THAT(model, NotNull()) << error;
  mjData* data = mj_makeData(model);
  int nv = model->nv, nu = model->nu;

  // set ctrl, integrate for 20 steps
  data->ctrl[0] = .1;
  data->ctrl[1] = -.1;
  for (int i = 0; i < 20; i++) {
    mj_step(model, data);
  }

  // analytic A and B
  mjtNum* A = (mjtNum*)mju_malloc(sizeof(mjtNum) * 2 * nv * 2 * nv);
  mjtNum* B = (mjtNum*)mju_malloc(sizeof(mjtNum) * 2 * nv * nu);

  LinearSystem(model, data, A, B);

  // uncomment for debugging:
  // PrintMatrix(A, 2*nv, 2*nv);
  // PrintMatrix(B, 2*nv, nu);

  // forward differenced A and B
  mjtNum eps = MjEps(1e-6, 1e-3);
  mjtNum* AFD = (mjtNum*)mju_malloc(sizeof(mjtNum) * 2 * nv * 2 * nv);
  mjtNum* BFD = (mjtNum*)mju_malloc(sizeof(mjtNum) * 2 * nv * nu);

  mjd_transitionFD(model, data, eps, /*centered=*/0, AFD, BFD, nullptr,
                   nullptr);

  // uncomment for debugging:
  // PrintMatrix(AFD, 2*nv, 2*nv);
  // PrintMatrix(BFD, 2*nv, nu);

  // expect FD and analytic derivatives to be similar to eps precision
  CompareMatrices(A, AFD, 2 * nv, 2 * nv, eps);
  CompareMatrices(B, BFD, 2 * nv, nu, eps);

  // central differenced A and B
  mjtNum* AFDc = (mjtNum*)mju_malloc(sizeof(mjtNum) * 2 * nv * 2 * nv);
  mjtNum* BFDc = (mjtNum*)mju_malloc(sizeof(mjtNum) * 2 * nv * nu);
  mjd_transitionFD(model, data, eps, /*centered=*/1, AFDc, BFDc, nullptr,
                   nullptr);

  // expect central derivatives to be equal to forward differences
  CompareMatrices(AFD, AFDc, 2 * nv, 2 * nv, eps);
  CompareMatrices(BFD, BFDc, 2 * nv, nu, eps);

  mju_free(BFDc);
  mju_free(AFDc);
  mju_free(BFD);
  mju_free(AFD);
  mju_free(B);
  mju_free(A);
  mj_deleteData(data);
  mj_deleteModel(model);
}

// check ctrl derivatives at the range limit
TEST_F(DerivativeTest, ClampedCtrlDerivatives) {
  const std::string xml_path = GetTestDataFilePath(kLinearPath);
  char error[1024];
  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
  ASSERT_THAT(model, NotNull()) << error;
  mjData* data = mj_makeData(model);
  int nv = model->nv, nu = model->nu;

  // set ctrl, integrate for 20 steps
  data->ctrl[0] = .1;
  data->ctrl[1] = -.1;
  for (int i = 0; i < 20; i++) {
    mj_step(model, data);
  }

  // analytic B
  mjtNum* B = (mjtNum*)mju_malloc(sizeof(mjtNum) * 2 * nv * nu);

  LinearSystem(model, data, nullptr, B);

  // forward differenced A and B
  mjtNum eps = MjEps(1e-6, 1e-3);
  mjtNum* BFD = (mjtNum*)mju_malloc(sizeof(mjtNum) * 2 * nv * nu);

  // set ctrl to the limits, request forward differences
  data->ctrl[0] = 1;
  data->ctrl[1] = -1;
  mjd_transitionFD(model, data, eps, /*centered=*/0, nullptr, BFD, nullptr,
                   nullptr);
  // expect FD and analytic derivatives to be similar to eps precision
  CompareMatrices(B, BFD, 2 * nv, nu, eps);

  // ctrl remains at limits, request central differences
  mjd_transitionFD(model, data, eps, /*centered=*/1, nullptr, BFD, nullptr,
                   nullptr);
  // expect FD and analytic derivatives to be similar to eps precision
  CompareMatrices(B, BFD, 2 * nv, nu, eps);

  // set ctrl beyond limits, request forward differences
  data->ctrl[0] = 2;
  data->ctrl[1] = -2;
  mjd_transitionFD(model, data, eps, /*centered=*/0, nullptr, BFD, nullptr,
                   nullptr);
  // expect derivatives to be 0
  EXPECT_THAT(AsVector(BFD, 2 * nv * nu), Each(Eq(0.0)));

  // expect ctrl to remain unchanged (despite internal clamping)
  EXPECT_EQ(data->ctrl[0], 2.0);
  EXPECT_EQ(data->ctrl[1], -2.0);

  // ctrl remains beyond limits, request centered differences
  mjd_transitionFD(model, data, eps, /*centered=*/1, nullptr, BFD, nullptr,
                   nullptr);
  // expect derivatives to be 0
  EXPECT_THAT(AsVector(BFD, 2 * nv * nu), Each(Eq(0.0)));

  mju_free(BFD);
  mju_free(B);
  mj_deleteData(data);
  mj_deleteModel(model);
}

// compare FD sensor derivatives to analytic derivatives
TEST_F(DerivativeTest, SensorDerivatives) {
  static constexpr char xml[] = R"(
  <mujoco>
    <worldbody>
      <body>
        <joint name="joint" type="slide"/>
        <geom size=".1"/>
      </body>
    </worldbody>

    <actuator>
      <general name="actuator" joint="joint" gainprm="3"/>
    </actuator>

    <sensor>
      <jointpos joint="joint"/>
      <jointvel joint="joint"/>
      <actuatorfrc actuator="actuator"/>
    </sensor>
  </mujoco>
  )";

  MjModelPtr model = LoadModelFromString(xml);
  int nv = model->nv, nu = model->nu, ns = model->nsensordata;
  MjDataPtr data = MakeData(model);

  // finite differenced C and D
  mjtNum eps = 1e-6;
  mjtNum* CFD = (mjtNum*)mju_malloc(sizeof(mjtNum) * ns * 2 * nv);
  mjtNum* DFD = (mjtNum*)mju_malloc(sizeof(mjtNum) * ns * nu);
  mjd_transitionFD(model.get(), data.get(), eps, /*centered=*/0, nullptr,
                   nullptr, CFD, DFD);

  // expected analytic C and D
  mjtNum C[6] = {1, 0, 0, 1, 0, 0};

  mjtNum D[3] = {
      0,
      0,
      3,
  };

  // compare expected and actual values
  CompareMatrices(CFD, C, ns, 2 * nv, eps);
  CompareMatrices(DFD, D, ns, nu, eps);

  mju_free(DFD);
  mju_free(CFD);
}

// if sensor derivatives aren't requested, don't compute sensors
TEST_F(DerivativeTest, SensorSkip) {
  static constexpr char xml[] = R"(
  <mujoco>
    <worldbody>
      <body>
        <joint name="joint" type="slide"/>
        <geom size=".1"/>
      </body>
    </worldbody>

    <actuator>
      <general name="actuator" joint="joint" gainprm="3"/>
    </actuator>

    <sensor>
      <jointpos joint="joint"/>
    </sensor>
  </mujoco>
  )";

  MjModelPtr model = LoadModelFromString(xml);
  int nv = model->nv, nu = model->nu;
  MjDataPtr data = MakeData(model);

  // set a sentinel value in the sensor
  data->sensordata[0] = 1337;

  // finite differenced B
  mjtNum eps = 1e-6;
  mjtNum* BFD = (mjtNum*)mju_malloc(sizeof(mjtNum) * 2 * nv * nu);
  mjd_transitionFD(model.get(), data.get(), eps, /*centered=*/0, nullptr, BFD,
                   nullptr, nullptr);

  EXPECT_EQ(data->sensordata[0], 1337) << "sensors should not be recomputed";

  mju_free(BFD);
}

// derivatives don't mutate the state
TEST_F(DerivativeTest, NoStateMutation) {
  const std::string xml_path = GetTestDataFilePath(kModelPath);
  char error[1024];
  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
  ASSERT_THAT(model, NotNull()) << error;
  mjData* data0 = mj_makeData(model);
  mjData* data = mj_makeData(model);
  int nv = model->nv, nu = model->nu, na = model->na, ns = model->nsensordata;

  // set time
  data->time = data0->time = 0.5;

  for (int i = 0; i < nv; i++) {
    data->qpos[i] = data0->qpos[i] = (mjtNum)i + 1;
    data->qvel[i] = data0->qvel[i] = (mjtNum)i + 2;
  }

  // set ctrl
  for (int i = 0; i < nu; i++) {
    data->ctrl[i] = data0->ctrl[i] = (mjtNum)i + 1;
  }

  // set act
  for (int i = 0; i < na; i++) {
    data->act[i] = data0->act[i] = (mjtNum)i + 1;
  }

  // allocate Jacobians, call derivatives
  int ndx = nv + nv + na;
  mjtNum* A = (mjtNum*)mju_malloc(sizeof(mjtNum) * ndx * ndx);
  mjtNum* B = (mjtNum*)mju_malloc(sizeof(mjtNum) * ndx * nu);
  mjtNum* C = (mjtNum*)mju_malloc(sizeof(mjtNum) * ns * ndx);
  mjtNum* D = (mjtNum*)mju_malloc(sizeof(mjtNum) * ns * nu);
  mjtNum eps = 1e-6;
  mjd_transitionFD(model, data, eps, /*centered=*/0, A, B, C, D);

  // compare states in data and data0
  EXPECT_EQ(data->time, data0->time);
  EXPECT_EQ(AsVector(data->qpos, model->nq), AsVector(data0->qpos, model->nq));
  EXPECT_EQ(AsVector(data->qvel, nv), AsVector(data0->qvel, nv));
  EXPECT_EQ(AsVector(data->act, na), AsVector(data0->act, na));
  EXPECT_EQ(AsVector(data->ctrl, nu), AsVector(data0->ctrl, nu));

  mju_free(D);
  mju_free(C);
  mju_free(B);
  mju_free(A);
  mj_deleteData(data0);
  mj_deleteData(data);
  mj_deleteModel(model);
}

// finite differencing is not supported with sleeping enabled
TEST_F(DerivativeTest, FiniteDifferenceRejectsSleep) {
  const std::string xml_path = GetTestDataFilePath(kSleepEqualityPath);
  char error[1024];
  MjModelPtr model(mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error)));
  ASSERT_THAT(model.get(), NotNull()) << error;
  MjDataPtr data = MakeData(model);
  int nv = model->nv, ndx = 2 * nv + model->na;
  vector<mjtNum> A(ndx * ndx);
  vector<mjtNum> DfDq(nv * nv);
  mjtNum eps = 1e-6;

  // all trees are awake, but the evaluations would change the sleep state
  auto transition_error = MjuErrorMessageFrom(mjd_transitionFD);
  EXPECT_THAT(transition_error(model.get(), data.get(), eps, /*centered=*/1,
                               A.data(), nullptr, nullptr, nullptr),
              HasSubstr("sleeping is not supported"));
  auto inverse_error = MjuErrorMessageFrom(mjd_inverseFD);
  EXPECT_THAT(inverse_error(model.get(), data.get(), eps, /*flg_actuation=*/0,
                            DfDq.data(), nullptr, nullptr, nullptr, nullptr,
                            nullptr, nullptr),
              HasSubstr("sleeping is not supported"));

  // disabling sleep allows finite differencing
  model->opt.enableflags &= ~mjENBL_SLEEP;
  EXPECT_EQ(transition_error(model.get(), data.get(), eps, /*centered=*/1,
                             A.data(), nullptr, nullptr, nullptr),
            "");
  EXPECT_EQ(inverse_error(model.get(), data.get(), eps, /*flg_actuation=*/0,
                          DfDq.data(), nullptr, nullptr, nullptr, nullptr,
                          nullptr, nullptr),
            "");
}

// compare dense and sparse derivatives of qfrc_bias (RNE)
TEST_F(DerivativeTest, DenseSparseRneEquivalent) {
  // run test on all models
  for (const char* local_path :
       {kEnergyConservingPendulumPath, kTumblingThinObjectPath,
        kDampedActuatorsPath, kDamperActuatorsPath, kDCMotorPath}) {
    const std::string xml_path = GetTestDataFilePath(local_path);
    char error[1024] = "";
    mjModel* model =
        mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
    ASSERT_THAT(model, testing::NotNull()) << "Failed to load model: " << error;
    int nD = model->nD;
    mjtNum* qDeriv = (mjtNum*)mju_malloc(sizeof(mjtNum) * nD);
    mjData* data = mj_makeData(model);

    // take 100 steps so we have some velocities, then call forward
    mj_resetData(model, data);
    if (model->nu) {
      data->ctrl[0] = 0.1;
    }
    for (int i = 0; i < 100; i++) {
      mj_step(model, data);
    }
    mj_forward(model, data);

    // compute qDeriv with sparse function, make local copy
    mjd_smooth_vel(model, data, /*flg_bias=*/1);
    mju_copy(qDeriv, data->qDeriv, nD);

    // re-compute with dense function
    mju_zero(data->qDeriv, model->nD);
    mjd_actuator_vel(model, data);
    mjd_passive_vel(model, data);
    mjd_rne_vel_dense(model, data);

    // expect dense and sparse derivatives to be similar to precision
    EXPECT_THAT(AsVector(data->qDeriv, nD),
                Pointwise(MjNear(1e-12, 5e-5), AsVector(qDeriv, nD)));

    mju_free(qDeriv);
    mj_deleteData(data);
    mj_deleteModel(model);
  }
}

// compare FD inverse derivatives to analytic derivatives of linear system
TEST_F(DerivativeTest, LinearSystemInverse) {
  const std::string xml_path = GetTestDataFilePath(kLinearPath);
  char error[1024];
  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
  ASSERT_THAT(model, NotNull()) << error;
  mjData* data = mj_makeData(model);

  int nv = model->nv;
  int ns = model->nsensordata;
  int nC = model->nC;

  vector<mjtNum> DfDq(nv * nv);
  vector<mjtNum> DfDv(nv * nv);
  vector<mjtNum> DfDa(nv * nv);
  vector<mjtNum> DsDq(nv * ns);
  vector<mjtNum> DsDv(nv * ns);
  vector<mjtNum> DsDa(nv * ns);
  vector<mjtNum> DmDq(nv * nC);

  // call mj_forward to get accelerations at initial state
  mj_forward(model, data);

  // get derivatives
  mjtNum eps = 1e-6;
  mjtByte flg_actuation = 0;
  mjd_inverseFD(model, data, eps, flg_actuation, DfDq.data(), DfDv.data(),
                DfDa.data(), DsDq.data(), DsDv.data(), DsDa.data(),
                DmDq.data());

  // expect that position derivatives are the stiffnesses
  vector<mjtNum> DfDq_expect = {model->jnt_stiffness[0], 0, 0, 0,
                                model->jnt_stiffness[1], 0, 0, 0,
                                model->jnt_stiffness[2]};
  EXPECT_THAT(DfDq, Pointwise(DoubleNear(eps), DfDq_expect));

  // expect that velocity derivatives are the dampings
  vector<mjtNum> DfDv_expect = {model->dof_damping[0], 0, 0, 0,
                                model->dof_damping[1], 0, 0, 0,
                                model->dof_damping[2]};
  EXPECT_THAT(DfDv, Pointwise(DoubleNear(eps), DfDv_expect));

  // expect that acceleration derivatives are the mass matrix
  vector<mjtNum> DfDa_expect(nv * nv, 0);
  mj_fullM(model, data, DfDa_expect.data());
  EXPECT_THAT(DfDa, Pointwise(DoubleNear(eps), DfDa_expect));

  // expect that sensor derivatives w.r.t position only see sensor 1 at dof 0
  vector<mjtNum> DsDq_expect(nv * ns, 0);
  int dof_index = 0;
  int sensordata_index = model->sensor_adr[1];
  DsDq_expect[dof_index * ns + sensordata_index] = 1;
  EXPECT_THAT(DsDq, Pointwise(DoubleNear(eps), DsDq_expect));

  // expect that sensor derivatives w.r.t velocity only see sensor 0 at dof 1
  vector<mjtNum> DsDv_expect(nv * ns, 0);
  dof_index = 1;
  sensordata_index = model->sensor_adr[0];
  DsDv_expect[dof_index * ns + sensordata_index] = 1;
  EXPECT_THAT(DsDv, Pointwise(DoubleNear(eps), DsDv_expect));

  // expect that sensor derivatives w.r.t acceleration see the accelerometer
  // in the y-axis, affected by both dof 0 and dof 1
  vector<mjtNum> DsDa_expect(nv * ns, 0);
  dof_index = 0;
  sensordata_index = model->sensor_adr[2] + 1;
  DsDa_expect[dof_index * ns + sensordata_index] = 1;
  dof_index = 1;
  DsDa_expect[dof_index * ns + sensordata_index] = 1;
  EXPECT_THAT(DsDa, Pointwise(DoubleNear(eps), DsDa_expect));

  // expect that mass matrix derivatives are zero
  vector<mjtNum> DmDq_expect(nv * nC, 0);
  EXPECT_THAT(DmDq, Pointwise(DoubleNear(eps), DmDq_expect));

  mj_deleteData(data);
  mj_deleteModel(model);
}

// utility: generate two random quaternions with a given angle difference
void randomQuatPair(mjtNum qa[4], mjtNum qb[4], mjtNum angle, int seed) {
  // make distribution using seed
  std::mt19937_64 rng;
  rng.seed(seed);
  std::normal_distribution<double> dist(0, 1);

  // sample qa = qb
  for (int i = 0; i < 4; i++) {
    qa[i] = qb[i] = dist(rng);
  }
  mju_normalize4(qa);
  mju_normalize4(qb);

  // integrate qb in random direction by angle
  mjtNum dir[3];
  for (int i = 0; i < 3; i++) {
    dir[i] = dist(rng);
  }
  mju_normalize3(dir);
  mju_quatIntegrate(qb, dir, angle);
}

// utility: finite-difference Jacobians of mju_subQuat
static void subQuatFD(mjtNum Da[9], mjtNum Db[9], const mjtNum qa[4],
                      const mjtNum qb[4], mjtNum eps) {
  // subQuat
  mjtNum y[3];
  mju_subQuat(y, qa, qb);

  mjtNum dq[3];   // nudge input direction
  mjtNum dqa[4];  // nudged qa input
  mjtNum dqb[4];  // nudged qb input
  mjtNum dy[3];   // nudged output
  mjtNum DaT[9];  // Da transposed
  mjtNum DbT[9];  // Db transposed

  for (int i = 0; i < 3; i++) {
    // perturbation
    mju_zero3(dq);
    dq[i] = 1.0;

    // Jacobian: d_y / d_qa
    mju_copy4(dqa, qa);
    mju_quatIntegrate(dqa, dq, eps);
    mju_subQuat(dy, dqa, qb);

    mju_sub3(DaT + i * 3, dy, y);
    mju_scl3(DaT + i * 3, DaT + i * 3, 1.0 / eps);

    // Jacobian: d_y / d_qb
    mju_copy4(dqb, qb);
    mju_quatIntegrate(dqb, dq, eps);
    mju_subQuat(dy, qa, dqb);

    mju_sub3(DbT + i * 3, dy, y);
    mju_scl3(DbT + i * 3, DbT + i * 3, 1.0 / eps);
  }

  // transpose result
  mju_transpose(Da, DaT, 3, 3);
  mju_transpose(Db, DbT, 3, 3);
}

TEST_F(DerivativeTest, SubQuat) {
  const int nrepeats = 10;  // number of repeats
  const mjtNum eps =
      MjTol(1e-7, 1e-3);  // epsilon for finite-differencing and comparison

  int seed = 1;
  for (int i = 0; i < nrepeats; i++) {
    for (mjtNum angle : {0.0, 1e-9, 1e-5, 1e-2, 1.0, 4.0}) {
      // random quaternions
      mjtNum qa[4];
      mjtNum qb[4];

      // make random quaternion pair with given relative angle
      randomQuatPair(qa, qb, angle, seed++);

      // analytic Jacobians
      mjtNum Da[9];  // d_subQuat(qa, qb) / d_qa
      mjtNum Db[9];  // d_subQuat(qa, qb) / d_qb
      mjd_subQuat(qa, qb, Da, Db);

      // finite-differenced Jacobians
      mjtNum DaFD[9];
      mjtNum DbFD[9];
      subQuatFD(DaFD, DbFD, qa, qb, eps);

      // expect numerical equality
      EXPECT_THAT(AsVector(DaFD, 9),
                  Pointwise(MjNear(1e-7, 1e-3), AsVector(Da, 9)));
      EXPECT_THAT(AsVector(DbFD, 9),
                  Pointwise(MjNear(1e-7, 1e-3), AsVector(Db, 9)));
    }
  }
}

// utility: random quaternion, 3D velocity
static void randomQuatVel(mjtNum quat[4], mjtNum vel[3], int seed) {
  // make distribution using seed
  std::mt19937_64 rng;
  rng.seed(seed);
  std::normal_distribution<double> dist(0, 1);

  // sample quat
  for (int i = 0; i < 4; i++) {
    quat[i] = dist(rng);
  }
  mju_normalize4(quat);

  // sample vel
  for (int i = 0; i < 3; i++) {
    vel[i] = dist(rng);
  }
}

// utility: finite-difference Jacobians of mju_quatIntegrate
void mjd_quatIntegrateFD(mjtNum Dquat[9], mjtNum Ds[9], mjtNum Dvel[9],
                         mjtNum Dh[3], const mjtNum quat[4],
                         const mjtNum vel[3], mjtNum h, mjtNum eps) {
  // compute y, output of mju_quatIntegrate(quat, vel, h)
  mjtNum y[4] = {quat[0], quat[1], quat[2], quat[3]};
  mju_quatIntegrate(y, vel, h);

  mjtNum dx[3];      // nudged tangent-space input
  mjtNum dq[4];      // quat output
  mjtNum dy[3];      // nudged tangent-space output
  mjtNum DquatT[9];  // Dquat transposed
  mjtNum DsT[9];     // Ds transposed
  mjtNum DvelT[9];   // Dvel transposed

  for (int i = 0; i < 3; i++) {
    // perturbation
    mju_zero3(dx);
    dx[i] = 1.0;

    // d_y / d_quat
    mju_copy4(dq, quat);
    mju_quatIntegrate(dq, dx, eps);  // nudge dq
    mju_quatIntegrate(dq, vel, h);   // compute nudged
    mju_subQuat(dy, dq, y);          // subtract
    mju_scl3(DquatT + i * 3, dy, 1.0 / eps);

    // d_y / d_sv (scaled velocity)
    mju_copy4(dq, quat);
    mjtNum dsv[3] = {vel[0] * h, vel[1] * h, vel[2] * h};
    mju_addToScl3(dsv, dx, eps);      // nudge dsv
    mju_quatIntegrate(dq, dsv, 1.0);  // compute nudged
    mju_subQuat(dy, dq, y);           // subtract
    mju_scl3(DsT + i * 3, dy, 1.0 / eps);

    // d_y / d_v (unscaled velocity)
    mju_copy4(dq, quat);
    mjtNum dv[3] = {vel[0], vel[1], vel[2]};
    mju_addToScl3(dv, dx, eps);    // nudge dv
    mju_quatIntegrate(dq, dv, h);  // compute nudged
    mju_subQuat(dy, dq, y);        // subtract
    mju_scl3(DvelT + i * 3, dy, 1.0 / eps);
  }

  // d_y / d_h (unscaled velocity)
  mju_copy4(dq, quat);
  mju_quatIntegrate(dq, vel, h + eps);  // compute nudged
  mju_subQuat(dy, dq, y);               // subtract
  mju_scl3(Dh, dy, 1.0 / eps);

  // transpose
  mju_transpose(Dquat, DquatT, 3, 3);
  mju_transpose(Ds, DsT, 3, 3);
  mju_transpose(Dvel, DsT, 3, 3);
}

TEST_F(DerivativeTest, quatIntegrate) {
  const int nrepeats = 10;  // number of repeats
  const mjtNum eps =
      MjTol(1e-7, 1e-3);  // epsilon for finite-differencing and comparison

  int seed = 1;
  for (int i = 0; i < nrepeats; i++) {
    for (mjtNum h : {0.0, 1e-9, 1e-5, 1e-2, 1.0, 4.0}) {
      // make random quaternion and velocity
      mjtNum quat[4];
      mjtNum vel[3];
      randomQuatVel(quat, vel, seed++);

      // analytic Jacobians
      mjtNum Dquat[9];  // d_quatIntegrate(quat, vel, h) / d_quat
      mjtNum Dvel[9];   // d_quatIntegrate(quat, vel, h) / d_vel
      mjtNum Dh[3];     // d_quatIntegrate(quat, vel, h) / d_h
      mjd_quatIntegrate(vel, h, Dquat, Dvel, Dh);

      // finite-differenced Jacobians
      mjtNum DquatFD[9];
      mjtNum DsFD[9];
      mjtNum DvelFD[9];
      mjtNum DhFD[3];
      mjd_quatIntegrateFD(DquatFD, DsFD, DvelFD, DhFD, quat, vel, h, eps);

      // expect numerical equality of un/scaled velocity derivatives
      EXPECT_THAT(AsVector(DvelFD, 9), Pointwise(MjNear(1e-7, 1e-2), DsFD));

      // expect numerical equality of analytic and FD derivatives
      EXPECT_THAT(AsVector(DquatFD, 9), Pointwise(MjNear(1e-7, 1e-2), Dquat));
      EXPECT_THAT(AsVector(DvelFD, 9), Pointwise(MjNear(1e-7, 1e-2), Dvel));
      EXPECT_THAT(AsVector(DhFD, 3), Pointwise(MjNear(1e-7, 1e-2), Dh));
    }
  }
}

// implicit integration is better than Euler with active forcerange clamping
TEST_F(DerivativeTest, ForcerangeClampedDerivative) {
  static constexpr char xml[] = R"(
  <mujoco>
    <option timestep="0.01" integrator="implicitfast"/>

    <worldbody>
      <geom name="plane" type="plane" size="2 2 0.1"/>
      <light pos="0 0 3"/>
      <body name="1" pos="0 0 1">
        <joint name="1" type="slide" axis="1 0 0"/>
        <geom type="sphere" size="0.1" mass="1"/>
      </body>
    </worldbody>

    <actuator>
      <position joint="1" kp="10000" kv="1000" forcerange="-10 10"/>
    </actuator>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr m = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(m.get(), NotNull()) << error;

  mjtNum dt_small = 1e-4;
  mjtNum dt_large = 1e-2;
  mjtNum duration = 1.0;

  MjDataPtr d_gt = MakeData(m);
  MjDataPtr d_implicit = MakeData(m);
  MjDataPtr d_euler = MakeData(m);
  MjDataPtr d_discrete = MakeData(m);

  mj_resetData(m.get(), d_gt.get());
  mj_resetData(m.get(), d_implicit.get());
  mj_resetData(m.get(), d_euler.get());
  mj_resetData(m.get(), d_discrete.get());

  d_gt.get()->ctrl[0] = 0.5;
  d_implicit->ctrl[0] = 0.5;
  d_euler->ctrl[0] = 0.5;
  d_discrete->ctrl[0] = 0.5;

  mjtNum error_implicit = 0;
  mjtNum error_discrete = 0;
  mjtNum error_euler = 0;
  int nsteps_large = static_cast<int>(duration / dt_large);
  int substeps = static_cast<int>(dt_large / dt_small);

  m->opt.timestep = dt_large;

  m->opt.integrator = mjINT_IMPLICITFAST;
  mj_resetData(m.get(), d_gt.get());
  d_gt.get()->ctrl[0] = 0.5;
  m->opt.timestep = dt_small;
  m->opt.integrator = mjINT_EULER;

  for (int i = 0; i < nsteps_large; i++) {
    // ground truth: small steps with Euler
    m->opt.integrator = mjINT_EULER;
    m->opt.timestep = dt_small;
    for (int j = 0; j < substeps; j++) {
      mj_step(m.get(), d_gt.get());
    }

    // euler at large timestep
    m->opt.timestep = dt_large;
    mj_step(m.get(), d_euler.get());

    // implicitfast at large timestep
    m->opt.integrator = mjINT_IMPLICITFAST;
    mj_step(m.get(), d_implicit.get());

    // discrete at large timestep
    m->opt.integrator = mjINT_DISCRETE;
    mj_step(m.get(), d_discrete.get());

    // accumulate errors
    mjtNum diff_implicit = d_gt.get()->qpos[0] - d_implicit->qpos[0];
    mjtNum diff_euler = d_gt.get()->qpos[0] - d_euler->qpos[0];
    mjtNum diff_discrete = d_gt.get()->qpos[0] - d_discrete->qpos[0];
    error_implicit += diff_implicit * diff_implicit;
    error_euler += diff_euler * diff_euler;
    error_discrete += diff_discrete * diff_discrete;
  }

  // expect implicitfast to be more accurate than Euler
  EXPECT_LT(error_implicit, error_euler)
      << "implicitfast should be more accurate than Euler at large timestep "
      << "when forcerange derivatives are correctly handled";

  // the discrete arm: while the force is clamped, actuatorDerivSkip keeps the
  // gains out of the metric and discrete coincides with Euler; unclamped
  // stretches trade Euler's explicit overshoot for the metric's implicit
  // damping, so the errors are near-equal with a platform-dependent sign.
  // Assert comparable accuracy; the stability advantage at large h*omega is
  // pinned by DiscreteStiffLegPress
  EXPECT_LT(error_discrete, 1.5 * error_euler)
      << "discrete should be comparable to Euler under a clamped servo";
}

// under discrete, qacc is the step map: finite differences of mj_step through
// mjd_transitionFD (which reuses stages via mj_stepSkip) must match naive
// finite differences with a full reset and a fresh step per perturbation
TEST_F(DerivativeTest, DiscreteStepMapDerivatives) {
  static constexpr char xml[] = R"(
  <mujoco>
    <option timestep="0.01" integrator="discrete"/>
    <worldbody>
      <geom name="floor" type="plane" size="1 1 .1"/>
      <body pos="0 0 0.079">
        <joint name="slide" type="slide" axis="0 0 1" stiffness="2000" damping="10" springref="-0.02"/>
        <joint name="hinge" type="hinge" axis="0 1 0" stiffness="500" damping="2"/>
        <geom type="capsule" size="0.04" fromto="0 0 0 0.2 0 -0.04"/>
      </body>
    </worldbody>
    <actuator>
      <motor joint="hinge"/>
    </actuator>
  </mujoco>
  )";
  char error[1024];
  MjModelPtr m = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(m.get(), NotNull()) << error;
  MjDataPtr d = MakeData(m);
  int nv = m->nv, nu = m->nu, ns = 2 * nv;

  // state with an active contact and nonzero velocity
  const mjtNum qpos0[2] = {-0.01, 0.05};
  const mjtNum qvel0[2] = {-0.2, 0.3};
  const mjtNum ctrl0[1] = {0.5};
  mju_copy(d->qpos, qpos0, nv);
  mju_copy(d->qvel, qvel0, nv);
  mju_copy(d->ctrl, ctrl0, nu);
  mj_forward(m.get(), d.get());
  ASSERT_GT(d->ncon, 0);

  // transition derivatives via stage-reusing finite differences
  // (eps is a perturbation size, not a tolerance: independent of MJTOL_SCALE)
  mjtNum eps = sizeof(mjtNum) == 8 ? 1e-6 : 1e-3;
  std::vector<mjtNum> A(ns * ns), B(ns * nu);
  mjd_transitionFD(m.get(), d.get(), eps, /*flg_centered=*/1, A.data(),
                   B.data(), nullptr, nullptr);

  // naive centered differences: full reset and one fresh mj_step per
  // perturbation
  auto naive = [&](int what, int idx, mjtNum* dx) {
    mjtNum x[2][4];
    for (int sgn = 0; sgn < 2; sgn++) {
      mj_resetData(m.get(), d.get());
      mju_copy(d->qpos, qpos0, nv);
      mju_copy(d->qvel, qvel0, nv);
      mju_copy(d->ctrl, ctrl0, nu);
      mjtNum e = sgn ? eps : -eps;
      if (what == 0) d->qpos[idx] += e;
      if (what == 1) d->qvel[idx] += e;
      if (what == 2) d->ctrl[idx] += e;
      mj_step(m.get(), d.get());
      mju_copy(x[sgn], d->qpos, nv);
      mju_copy(x[sgn] + nv, d->qvel, nv);
    }
    for (int i = 0; i < ns; i++) {
      dx[i] = (x[1][i] - x[0][i]) / (2 * eps);
    }
  };

  // guard against a vacuous pass: A contains the identity block and h-scale
  // dynamics
  EXPECT_GT(mju_norm(A.data(), ns * ns), 1);

  mjtNum maxdiff = 0;
  mjtNum dx[4];
  for (int j = 0; j < ns; j++) {
    naive(j < nv ? 0 : 1, j < nv ? j : j - nv, dx);
    for (int i = 0; i < ns; i++) {
      maxdiff = mju_max(maxdiff, mju_abs(A[i * ns + j] - dx[i]));
    }
  }
  for (int j = 0; j < nu; j++) {
    naive(2, j, dx);
    for (int i = 0; i < ns; i++) {
      maxdiff = mju_max(maxdiff, mju_abs(B[i * nu + j] - dx[i]));
    }
  }
  EXPECT_LT(maxdiff, MjTol(1e-8, 2e-5));
}

// forcelimited actuator following a multi-output SO3 actuator: the derivative
// skip for saturated actuators must index forcerange per actuator, not per
// output.
TEST_F(DerivativeTest, ForcerangeClampedAfterSO3) {
  static constexpr char xml[] = R"(
  <mujoco>
    <option>
      <flag contact="disable" gravity="disable"/>
    </option>
    <worldbody>
      <body>
        <joint name="ball" type="ball"/>
        <geom type="box" size=".05 .07 .03"/>
      </body>
      <body pos="0 0 .3">
        <joint name="hinge"/>
        <geom size=".05"/>
      </body>
    </worldbody>
    <actuator>
      <orientation joint="ball" kp="1" kv="1"/>
      <velocity joint="hinge" kv="10" forcerange="-1 1"/>
    </actuator>
  </mujoco>
  )";
  char error[1024];
  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  MjDataPtr data = MakeData(model);
  mjModel* m = model.get();
  mjData* d = data.get();

  // spin the hinge so the velocity actuator saturates: force -50, clamped -1
  mjtNum qvel[4] = {0.1, 0.2, 0.3, 5};
  mju_copy(d->qvel, qvel, 4);
  mj_forward(m, d);
  ASSERT_EQ(d->actuator_force[3], -1);

  // analytic qDeriv
  mju_zero(d->qDeriv, m->nD);
  mjd_smooth_vel(m, d, /*flg_bias=*/1);
  vector<mjtNum> qDerivAnalytic = AsVector(d->qDeriv, m->nD);
  EXPECT_GT(mju_norm(qDerivAnalytic.data(), m->nD), 0);

  // expect match with finite differences: the saturated actuator contributes
  // nothing, the SO3 actuator's damping is unaffected by its neighbor
  mjtNum eps = MjEps(1e-7, 1e-3);
  mju_zero(d->qDeriv, m->nD);
  mjd_smooth_velFD(m, d, eps);
  EXPECT_THAT(AsVector(d->qDeriv, m->nD),
              Pointwise(MjNear(1e-7, 3e-3), qDerivAnalytic));
}

TEST_F(DerivativeTest, NonlinearDampingDerivative) {
  static constexpr char xml[] = R"(
  <mujoco>
    <worldbody>
      <body>
        <joint type="slide" damping="2 3 4"/>
        <geom size="1" mass="1"/>
      </body>
    </worldbody>

    <keyframe>
      <key qvel="3"/>
    </keyframe>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr m = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(m.get(), NotNull()) << error;

  mjtNum dt_small = 1e-4;
  mjtNum dt_large = 1e-2;
  mjtNum duration = 1.0;

  MjDataPtr d_gt = MakeData(m);
  MjDataPtr d_enabled = MakeData(m);
  MjDataPtr d_disabled = MakeData(m);

  mj_resetDataKeyframe(m.get(), d_gt.get(), 0);
  mj_resetDataKeyframe(m.get(), d_enabled.get(), 0);
  mj_resetDataKeyframe(m.get(), d_disabled.get(), 0);

  m->opt.integrator = mjINT_EULER;
  mjtNum error_enabled = 0;
  mjtNum error_disabled = 0;
  int nsteps_large = static_cast<int>(duration / dt_large);
  int substeps = static_cast<int>(dt_large / dt_small);

  for (int i = 0; i < nsteps_large; i++) {
    m->opt.timestep = dt_small;
    m->opt.disableflags |= mjDSBL_EULERDAMP;  // disable implicit damping
    for (int j = 0; j < substeps; j++) {
      mj_step(m.get(), d_gt.get());
    }

    m->opt.timestep = dt_large;
    mj_step(m.get(), d_disabled.get());

    m->opt.disableflags &= ~mjDSBL_EULERDAMP;  // enable implicit damping
    mj_step(m.get(), d_enabled.get());

    mjtNum diff_enabled = d_gt.get()->qvel[0] - d_enabled.get()->qvel[0];
    mjtNum diff_disabled = d_gt.get()->qvel[0] - d_disabled.get()->qvel[0];
    error_enabled += diff_enabled * diff_enabled;
    error_disabled += diff_disabled * diff_disabled;
  }

  EXPECT_LT(error_enabled, error_disabled)
      << "Euler with implicit damping should be more accurate than without "
      << "when nonlinear damping derivatives are correctly handled";
}

// implicit derivatives should use next activation when actearly is set
TEST_F(DerivativeTest, ActearlyDerivative) {
  static constexpr char xml[] = R"(
  <mujoco>
    <option timestep="1" integrator="implicitfast"/>

    <worldbody>
      <body>
        <joint name="early" type="slide"/>
        <geom type="sphere" size="0.1" mass="1"/>
      </body>
      <body pos="1 0 0">
        <joint name="late" type="slide"/>
        <geom type="sphere" size="0.1" mass="1"/>
      </body>
    </worldbody>

    <actuator>
      <general joint="early" dyntype="integrator" gaintype="affine"
               gainprm="1 0 1" actearly="true"/>
      <general joint="late" dyntype="integrator" gaintype="affine"
               gainprm="1 0 1" actearly="false"/>
    </actuator>
  </mujoco>
  )";

  // the affine velocity gain of 1 makes M - h*qDeriv exactly singular at h=1:
  // expect a clamped-pivot warning from the implicitfast qH factorization
  mock_warning_handler.ExpectWarnings(
      "Inertia matrix is too close to singular");

  char error[1024];
  MjModelPtr m = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(m.get(), NotNull()) << error;
  MjDataPtr d = MakeData(m);

  // set identical ctrl with zero initial activation
  d->ctrl[0] = 1.0;
  d->ctrl[1] = 1.0;
  d->act[0] = 0.0;
  d->act[1] = 0.0;

  // step computes derivatives during implicit integration
  mj_step(m.get(), d.get());

  // both should have same act_dot
  EXPECT_EQ(d->act_dot[0], d->act_dot[1]);

  // with actearly=true and nonzero act_dot, derivative should differ
  // because actearly uses next activation: act + act_dot*dt
  // for our model: next_act = 0 + 1*1 = 1, current_act = 0
  // derivative adds gain_vel * act to qDeriv diagonal
  // for independent bodies, D is diagonal, so diag[i] is at D_rowadr[i]
  int diag0 = m->D_rowadr[0];  // first joint's diagonal
  int diag1 = m->D_rowadr[1];  // second joint's diagonal
  EXPECT_NE(d->qDeriv[diag0], d->qDeriv[diag1])
      << "actearly=true should use next activation in derivative";

  // verify specific values: gain_vel=1, next_act=1, current_act=0
  EXPECT_NEAR(d->qDeriv[diag0], 1.0, 1e-10)
      << "actearly=true should use next_act=1";
  EXPECT_NEAR(d->qDeriv[diag1], 0.0, 1e-10)
      << "actearly=false should use current_act=0";
}

// verify stateful DC motor derivative matches analytical formula
TEST_F(DerivativeTest, DCMotorStatefulDerivative) {
  static constexpr char xml[] = R"(
  <mujoco>
    <option timestep="0.002"/>
    <worldbody>
      <body>
        <joint name="j" type="slide"/>
        <geom type="sphere" size="0.1" mass="1"/>
      </body>
    </worldbody>
    <actuator>
      <dcmotor name="dc" joint="j" motorconst="2.0" resistance="0.5"
               inductance="0 0.001" input="pos vel" controller="10 0 5"/>
    </actuator>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr m = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(m.get(), NotNull()) << error;
  MjDataPtr d = MakeData(m);

  // set nonzero velocity and ctrl
  d->qvel[0] = 1.0;
  d->ctrl[0] = 0.5;

  // forward to compute act_dot, etc.
  mj_forward(m.get(), d.get());

  // compute analytical derivatives
  mjd_smooth_vel(m.get(), d.get(), /* flg_bias = */ 1);

  // extract diagonal of qDeriv
  mjtNum qDeriv_diag = d->qDeriv[m->D_rowadr[0] + m->D_rownnz[0] - 1];

  // expected: K*(dVdw - K)*(1 - exp(-h/te))/R with the torque-space map
  // dVdw = -kd*R/K + K, so the expression reduces to -kd*(1 - exp(-h/te))
  mjtNum te = 0.001, h = 0.002, kd = 5.0;
  mjtNum expected = -kd * (1 - mju_exp(-h / te));
  EXPECT_NEAR(qDeriv_diag, expected, 1e-10)
      << "stateful DC motor derivative should match analytical formula";
}

// verify that stateful DC motor derivative converges to stateless as te -> 0
TEST_F(DerivativeTest, DCMotorStatefulConvergesToStateless) {
  // stateless DC motor with position controller
  static constexpr char xml_stateless[] = R"(
  <mujoco>
    <option timestep="0.002"/>
    <worldbody>
      <body>
        <joint name="j" type="slide"/>
        <geom type="sphere" size="0.1" mass="1"/>
      </body>
    </worldbody>
    <actuator>
      <dcmotor name="dc" joint="j" motorconst="1.0" resistance="1.0"
               input="pos vel" controller="10 0 5"/>
    </actuator>
  </mujoco>
  )";

  // stateful DC motor with very small te
  static constexpr char xml_stateful[] = R"(
  <mujoco>
    <option timestep="0.002"/>
    <worldbody>
      <body>
        <joint name="j" type="slide"/>
        <geom type="sphere" size="0.1" mass="1"/>
      </body>
    </worldbody>
    <actuator>
      <dcmotor name="dc" joint="j" motorconst="1.0" resistance="1.0"
               inductance="0 1e-8" input="pos vel" controller="10 0 5"/>
    </actuator>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr m_sl = LoadModelFromString(xml_stateless, error, sizeof(error));
  ASSERT_THAT(m_sl.get(), NotNull()) << error;
  MjDataPtr d_sl = MakeData(m_sl);

  MjModelPtr m_sf = LoadModelFromString(xml_stateful, error, sizeof(error));
  ASSERT_THAT(m_sf.get(), NotNull()) << error;
  MjDataPtr d_sf = MakeData(m_sf);

  // set identical state
  d_sl.get()->qvel[0] = d_sf.get()->qvel[0] = 1.0;
  d_sl.get()->ctrl[0] = d_sf.get()->ctrl[0] = 0.5;

  // forward and compute derivatives
  mj_forward(m_sl.get(), d_sl.get());
  mj_forward(m_sf.get(), d_sf.get());
  mjd_smooth_vel(m_sl.get(), d_sl.get(), 1);
  mjd_smooth_vel(m_sf.get(), d_sf.get(), 1);

  // extract diagonals
  mjtNum diag_sl =
      d_sl.get()->qDeriv[m_sl->D_rowadr[0] + m_sl->D_rownnz[0] - 1];
  mjtNum diag_sf =
      d_sf.get()->qDeriv[m_sf->D_rowadr[0] + m_sf->D_rownnz[0] - 1];

  EXPECT_NEAR(diag_sf, diag_sl, 1e-6)
      << "stateful derivative should converge to stateless as te -> 0";
}

// verify hot winding resistance is used for stateless and stateful derivatives
TEST_F(DerivativeTest, DCMotorThermalDerivative) {
  static constexpr char xml[] = R"(
  <mujoco>
    <option timestep="0.002"/>
    <worldbody>
      <body>
        <joint name="stateless" type="slide"/>
        <geom type="sphere" size="0.1" mass="1"/>
      </body>
      <body>
        <joint name="stateful" type="slide"/>
        <geom type="sphere" size="0.1" mass="1"/>
      </body>
    </worldbody>
    <actuator>
      <dcmotor joint="stateless" motorconst="1" resistance="1"
               thermal="1 1 0 0.004 25 25"/>
      <dcmotor joint="stateful" motorconst="1" resistance="1"
               inductance="0 0.01" thermal="1 1 0 0.004 25 25"/>
    </actuator>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr m = LoadModelFromString(xml, error, sizeof(error));
  ASSERT_THAT(m.get(), NotNull()) << error;
  MjDataPtr d = MakeData(m);

  // A 250-degree rise doubles both winding resistances.
  d->act[m->actuator_actadr[0]] = 250;
  d->act[m->actuator_actadr[1]] = 250;
  d->qvel[0] = 0.5;
  d->qvel[1] = 0.5;
  mj_forward(m.get(), d.get());

  mju_zero(d->qDeriv, m->nD);
  mjd_smooth_vel(m.get(), d.get(), /*flg_bias=*/1);
  vector<mjtNum> analytic = AsVector(d->qDeriv, m->nD);

  mju_zero(d->qDeriv, m->nD);
  mjd_smooth_velFD(m.get(), d.get(), MjTol(1e-7, 1e-3));
  EXPECT_THAT(AsVector(d->qDeriv, m->nD),
              Pointwise(MjNear(1e-7, 3e-3), analytic));
}

// Utility: Rotate flex grid
void RotateFlexGrid(mjModel* model, mjData* data, const char* flex_name,
                    double angle) {
  int flex_id = mj_name2id(model, mjOBJ_FLEX, flex_name);
  ASSERT_NE(flex_id, -1);
  int node_adr = model->flex_nodeadr[flex_id];
  int* node_bodies = model->flex_nodebodyid + node_adr;
  int nodenum = model->flex_nodenum[flex_id];

  // Make deterministic quaternion for rotation inside helper
  mjtNum quat[4] = {1, 0, 0, 0};
  if (angle != 0) {
    mjtNum vel[3] = {1, 1, 1};
    mju_normalize3(vel);
    mju_quatIntegrate(quat, vel, angle);
  }

  // reset first to get initial positions
  mj_resetData(model, data);
  mj_forward(model, data);  // Compute initial xpos

  // Update qpos
  for (int i = 0; i < nodenum; i++) {
    int bodyid = node_bodies[i];

    // Only process nodes with valid bodies (FlexInterpDamping assumes this)
    if (bodyid >= 0) {
      mjtNum xpos0[3];
      mju_copy3(xpos0, data->xpos + 3 * bodyid);  // Initial absolute position

      mjtNum xpos_new[3];
      mju_rotVecQuat(xpos_new, xpos0, quat);  // Rotate absolute position

      mjtNum delta[3];
      mju_sub3(delta, xpos_new, xpos0);

      // Find the qpos address for this node/body
      int jnt = model->body_jntadr[bodyid];
      if (jnt >= 0) {
        int qadr = model->jnt_qposadr[jnt];
        mju_addTo3(data->qpos + qadr, delta);
      }
    }
  }
}

// Helper: assemble flex stiffness into dense matrix via matrix-vector products.
// Builds K column-by-column using mjd_flexInterp_mul.
// Result is -(h^2 + h*damping) * J'KJ (negative sign matches the old addH
// convention where stiffness is subtracted from the system matrix).
static void mulKD_dense(mjModel* m, mjData* d, mjtNum* H_dense, int nv,
                        mjtNum h) {
  std::vector<mjtNum> e_i(nv, 0);
  std::vector<mjtNum> col(nv, 0);
  for (int i = 0; i < nv; i++) {
    mju_zero(e_i.data(), nv);
    mju_zero(col.data(), nv);
    e_i[i] = 1.0;
    mjd_flexInterp_mul(m, d, col.data(), e_i.data(), h * h, h, NULL);
    // col = +(h^2 + h*damp)*K*e_i, negate to match addH convention (H -= K)
    for (int j = 0; j < nv; j++) {
      H_dense[j * nv + i] = -col[j];
    }
  }
}

// compare analytic and fin-diff d_qfrc_passive/d_qvel for flex interp
// Combined test for verify mjd_flexInterp_mulK (stiffness) and damping
TEST_F(DerivativeTest, FlexInterpDerivatives) {
  static const char* const kXml = R"(
  <mujoco>
    <option integrator="discrete"/>
    <worldbody>
      <flexcomp name="flex" type="grid" count="3 3 3" spacing="0.1 0.2 0.3"
                radius=".01" dim="3" mass="1" dof="trilinear">
        <contact selfcollide="none"/>
        <elasticity young="1e4" poisson="0.3" damping="50"/>
      </flexcomp>
    </worldbody>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr model = LoadModelFromString(kXml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  int nD = model->nD;
  int nv = model->nv;
  ASSERT_EQ(model->nq, 24);  // 8 corners * 3 dofs

  MjDataPtr data = MakeData(model);

  // iterate over rotations
  for (mjtNum angle : {0.0, 0.5, 1.0, mjPI / 2, mjPI, 2.0 * mjPI}) {
    RotateFlexGrid(model.get(), data.get(), "flex", angle);
    mj_forward(model.get(), data.get());

    // part 1: stiffness verification
    {
      std::vector<mjtNum> vec(nv);
      std::vector<mjtNum> res(nv);
      mju_zero(vec.data(), nv);
      // use deterministic random perturbation to verify full stiffness matrix
      // behavior
      for (int i = 0; i < nv; i++) {
        vec[i] = mju_Halton(i, 2) - 0.5;
      }

      // use mulKD to compute K * vec
      // mulKD adds (h^2*K + h*D)*vec to res
      // if we set h=1, damping=0, we get K*vec
      mjtNum save_damping = model->flex_damping[0];
      model->flex_damping[0] = 0;

      std::vector<mjtNum> H(nv * nv, 0);

      // assemble K into H column-by-column
      mulKD_dense(model.get(), data.get(), H.data(), nv, 1.0);

      // restore damping
      model->flex_damping[0] = save_damping;

      // compute res = K * vec
      mju_mulMatVec(res.data(), H.data(), vec.data(), nv, nv);

      // finite difference of mj_passive for stiffness
      mjtNum eps = MjEps(1e-6, 1e-3);
      mjData* data_perturbed = mj_copyData(NULL, model.get(), data.get());

      // apply perturbation
      mju_addToScl(data_perturbed->qpos, vec.data(), eps, nv);

      // recompute geometry/passive
      mj_forward(model.get(), data_perturbed);

      // compute FD estimate of K * vec
      // qfrc_passive = -dV/dq => d(qfrc)/dq = -K
      // (qfrc_new - qfrc)/eps ~= -K * vec
      std::vector<mjtNum> fd_res(nv);
      for (int i = 0; i < nv; ++i) {
        fd_res[i] =
            -(data_perturbed->qfrc_passive[i] - data->qfrc_passive[i]) / eps;
      }

      // compare analytical result (H*vec) with FD result
      for (int i = 0; i < nv; ++i) {
        EXPECT_THAT(res[i], MjNear(fd_res[i], 5e-3, 5.0))
            << "Stiffness Mismatch at DOF " << i;
      }

      mj_deleteData(data_perturbed);

      // check symmetry: K[i,j] == K[j,i]
      std::vector<mjtNum>& K_full = H;
      mjtNum max_asymmetry = 0;
      for (int i = 0; i < nv; i++) {
        for (int j = 0; j < i; j++) {
          mjtNum diff = mju_abs(K_full[i * nv + j] - K_full[j * nv + i]);
          max_asymmetry = mju_max(max_asymmetry, diff);
        }
      }
      EXPECT_THAT(max_asymmetry, MjNear(0, 1e-10, 5e-4))
          << "K matrix is not symmetric at angle " << angle;

      // check positive semi-definiteness: v^T * K * v >= 0
      for (int trial = 0; trial < 5; trial++) {
        std::vector<mjtNum> v(nv);
        for (int i = 0; i < nv; i++) {
          v[i] = mju_Halton(i + trial * nv, 3) - 0.5;
        }
        mjtNum vKv = 0;
        for (int i = 0; i < nv; i++) {
          for (int j = 0; j < nv; j++) {
            vKv += v[i] * K_full[i * nv + j] * v[j];
          }
        }
        EXPECT_GE(vKv, MjTol(-1e-8, -1e-5))
            << "K matrix is not PSD at angle " << angle;
      }
    }

    // part 2: damping verification
    {
      // set velocity non-zero to test damping
      data->qvel[0] = 1.0;

      mj_forward(model.get(), data.get());

      // get analytic derivatives (without Flex Damping currently)
      std::vector<mjtNum> qDerivAnalytic(nD);
      mju_zero(data->qDeriv, nD);
      mjd_passive_vel(model.get(), data.get());
      mju_copy(qDerivAnalytic.data(), data->qDeriv, nD);

      // finite-difference derivatives
      std::vector<mjtNum> qDerivFD(nD);
      mju_zero(data->qDeriv, nD);
      mjtNum eps = MjEps(1e-6, 1e-3);

      mjd_passive_velFD(model.get(), data.get(), eps);
      mju_copy(qDerivFD.data(), data->qDeriv, nD);

      // check that we have non-zero damping (FD should find it)
      EXPECT_GT(mju_norm(qDerivFD.data(), nD), 1e-3);

      // compute expected flex damping using mulKD_dense
      // D = 4*H(0.5) - H(1)
      vector<mjtNum> H1(nv * nv, 0);
      mulKD_dense(model.get(), data.get(), H1.data(), nv, 1.0);

      vector<mjtNum> H2(nv * nv, 0);
      mulKD_dense(model.get(), data.get(), H2.data(), nv, 0.5);

      vector<mjtNum> D(nv * nv);
      for (int i = 0; i < nv * nv; i++) {
        D[i] = 4.0 * H2[i] - H1[i];
      }

      // subtract D from qDerivAnalytic using sparse indexing
      // d(force)/d(vel) = -D
      for (int i = 0; i < nv; i++) {
        int rownnz = model->D_rownnz[i];
        int rowadr = model->D_rowadr[i];
        for (int k = 0; k < rownnz; k++) {
          int index = rowadr + k;
          int j = model->D_colind[index];
          qDerivAnalytic[index] -= D[i * nv + j];
        }
      }

      // expect FD and corrected analytic derivatives to match
      EXPECT_THAT(qDerivAnalytic, Pointwise(MjNear(1e-4, 1e4), qDerivFD))
          << "Damping Mismatch at angle: " << angle;
    }
  }
}

// Test Jacobian under deformation to highlight approximation error
TEST_F(DerivativeTest, FlexInterpDerivativesDeformed) {
  static const char* const kXml = R"(
  <mujoco>
    <option integrator="discrete"/>
    <worldbody>
      <flexcomp name="flex" type="grid" count="3 3 3" spacing="0.1 0.2 0.3"
                radius=".01" dim="3" mass="1" dof="trilinear">
        <contact selfcollide="none"/>
        <elasticity young="1e4" poisson="0.3" damping="0"/>
      </flexcomp>
    </worldbody>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr model = LoadModelFromString(kXml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  int nv = model->nv;

  MjDataPtr data = MakeData(model);

  // Apply rotation
  RotateFlexGrid(model.get(), data.get(), "flex", 1.0);  // 1 radian rotation

  // Apply deformation (stretch along X)
  // qpos is initialized by RotateFlexGrid.
  // Add a random perturbation to qpos that represents deformation.
  // We use a deterministic sequence to ensure reproducibility.
  std::vector<mjtNum> deformation(nv);
  for (int i = 0; i < nv; i++) {
    // Large deformation to make sure terms are significant
    deformation[i] = (mju_Halton(i, 3) - 0.5) * 0.2;
  }
  mju_addTo(data->qpos, deformation.data(), nv);

  mj_forward(model.get(), data.get());

  // 1. Compute Analytic Jacobian (Approximate)
  // We use mulKD_dense to get K_approx
  std::vector<mjtNum> H_approx(nv * nv, 0);

  // h=1, damping=0 => gives K
  mulKD_dense(model.get(), data.get(), H_approx.data(), nv, 1.0);

  // 2. Compute Finite Difference Jacobian (Ground Truth)
  // qfrc_passive = -dV/dq
  // d(qfrc)/dq = -K_true
  std::vector<mjtNum> K_true(nv * nv, 0);
  mjtNum eps = 1e-6;

  for (int i = 0; i < nv; i++) {
    mjData* data_p = mj_copyData(NULL, model.get(), data.get());
    data_p->qpos[i] += eps;
    mj_forward(model.get(), data_p);

    for (int j = 0; j < nv; j++) {
      // d(force_j)/d(q_i)
      mjtNum df = data_p->qfrc_passive[j] - data->qfrc_passive[j];
      // K_true[j, i] = -df/eps
      K_true[j * nv + i] = -df / eps;
    }
    mj_deleteData(data_p);
  }

  // 3. Compare and check for significant mismatch
  mjtNum max_error = 0;
  for (int i = 0; i < nv * nv; i++) {
    max_error = mju_max(max_error, mju_abs(H_approx[i] - K_true[i]));
  }

  // We expect significant error because of deformation + rotation.
  // The missing term (geometric stiffness) is proportional to stress.
  // We assert that the error is relatively large to confirm the approximation
  // exists.
  EXPECT_GT(max_error, 1e-3)
      << "Jacobian approximation should differ from FD when deformed";
}

// Helper: assemble the standard-flex stretch stiffness into a dense matrix,
// column-by-column using mjd_flexStretch_mul with scale (s1 + s2*damping).
static void stretchK_dense(mjModel* m, mjData* d, mjtNum* K, int nv, mjtNum s1,
                           mjtNum s2) {
  std::vector<mjtNum> e_i(nv, 0);
  std::vector<mjtNum> col(nv, 0);
  for (int i = 0; i < nv; i++) {
    mju_zero(e_i.data(), nv);
    mju_zero(col.data(), nv);
    e_i[i] = 1.0;
    mjd_flexStretch_mul(m, d, col.data(), e_i.data(), s1, s2);
    for (int j = 0; j < nv; j++) {
      K[j * nv + i] = col[j];
    }
  }
}

// The same frame mismatch seen through the derivative: mjd_flexBend_mul is the
// Jacobian of the bending force only if operator and force agree on the frame.
// An unrotated flex cannot catch it, because the stencil is one scalar per
// vertex pair applied coordinate-wise and so commutes with a rotation shared by
// every vertex. The pin also covers the zero-dof vertex path, where body_dofadr
// is negative and an unguarded stencil indexes out of bounds.
TEST_F(DerivativeTest, FlexBendDerivativesRotated) {
  static const char* const kXml = R"(
  <mujoco>
    <option integrator="discrete" solver="CG"/>
    <worldbody>
      <body name="turned" euler="90 35 20">
        <flexcomp name="rot" type="grid" count="3 3 1" spacing="0.1 0.1 0.1"
                  radius=".01" dim="2" mass="1">
          <pin id="0"/>
          <contact selfcollide="none" contype="0" conaffinity="0"/>
          <elasticity young="1e4" poisson="0.3" thickness="0.01"
                      elastic2d="bend" damping="0"/>
        </flexcomp>
      </body>
      <flexcomp name="flat" type="grid" count="3 3 1" spacing="0.1 0.1 0.1"
                radius=".01" dim="2" mass="1" pos="1 0 0">
        <pin id="0"/>
          <contact selfcollide="none" contype="0" conaffinity="0"/>
        <elasticity young="1e4" poisson="0.3" thickness="0.01"
                    elastic2d="bend" damping="0"/>
      </flexcomp>
    </worldbody>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr model = LoadModelFromString(kXml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  int nv = model->nv;
  ASSERT_EQ(model->nq, nv);  // all slide dofs
  ASSERT_EQ(model->nflex, 2);
  MjDataPtr data = MakeData(model);

  // deform both flexes out of plane so the bending stencil carries real force
  for (int i = 0; i < nv; i++) {
    data->qpos[i] += 2e-3 * (mju_Halton(i, 2) - 0.5);
  }
  mj_forward(model.get(), data.get());

  std::vector<mjtNum> vec(nv), res(nv, 0);
  for (int i = 0; i < nv; i++) {
    vec[i] = mju_Halton(i, 3) - 0.5;
  }
  mjd_flexBend_mul(model.get(), data.get(), res.data(), vec.data(), 1, 0);

  mjtNum eps = MjEps(1e-7, 1e-4);
  mjData* perturbed = mj_copyData(NULL, model.get(), data.get());
  mju_addToScl(perturbed->qpos, vec.data(), eps, nv);
  mj_forward(model.get(), perturbed);

  // check each flex on its own: the unrotated one is the control that isolates
  // the rotation
  for (int f = 0; f < model->nflex; f++) {
    SCOPED_TRACE(model->names + model->name_flexadr[f]);
    mjtNum max_err = 0, scale = 0;
    for (int k = 0; k < model->flex_vertnum[f]; k++) {
      int body = model->flex_vertbodyid[model->flex_vertadr[f] + k];
      int adr = model->body_dofadr[body];
      for (int x = 0; x < model->body_dofnum[body]; x++) {
        mjtNum fd =
            -(perturbed->qfrc_passive[adr + x] - data->qfrc_passive[adr + x]) /
            eps;
        max_err = mju_max(max_err, mju_abs(res[adr + x] - fd));
        scale = mju_max(scale, mju_abs(fd));
      }
    }
    // bending is linear in position, so the difference is exact up to roundoff
    EXPECT_GT(scale, 1e-3)
        << "test should exercise nontrivial bending stiffness";
    EXPECT_LT(max_err, MjTol(1e-4, 1e-2) * scale)
        << "mjd_flexBend_mul is not the Jacobian of the flex bending force";
  }
  mj_deleteData(perturbed);
}

struct FlexFrameCase {
  const char* name;
  bool welded;
  bool rotated_siblings;
  int nflex;
};

class FlexFrameTest : public MujocoTest,
                      public testing::WithParamInterface<FlexFrameCase> {
 protected:
  void SetUp() override { LoadPatch(); }

  // Two equilateral triangles; only the body frames and attachment topology
  // vary.
  void LoadPatch(const char* elastic2d = "both") {
    std::string xml = R"(
      <mujoco>
        <option timestep="0.002" integrator="discrete" solver="CG"/>
        <worldbody><body name="parent" euler="90 35 20">
    )";
    const char* positions[] = {"0 0 0", "0.1 0 0", "0.05 0.0866025404 0",
                               "0.15 0.0866025404 0"};
    const char* rotations[] = {"0 0 0", "30 -15 20", "-20 40 10", "45 10 -35"};
    for (int v = 0; v < 4; v++) {
      const char* rotation =
          GetParam().rotated_siblings ? rotations[v] : "0 0 0";
      if (GetParam().welded && v == 3) {
        rotation = "25 -15 40";
      }
      xml += "<body name='v" + std::to_string(v) + "' pos='" + positions[v] +
             "' euler='" + rotation + R"('>
          <inertial pos="0 0 0" mass="0.1" diaginertia="1e-4 1e-4 1e-4"/>
          <joint type="slide" axis="1 0 0" armature="0.01"/>
          <joint type="slide" axis="0 1 0" armature="0.02"/>
          <joint type="slide" axis="0 0 1" armature="0.03"/>
      )";
      if (GetParam().welded && v == 3) {
        xml += R"(<body name="welded" euler="20 10 30">
          <inertial pos="0 0 0" mass="0.1" diaginertia="1e-4 1e-4 1e-4"/>
        </body>)";
      }
      xml += "</body>";
    }
    xml += "</body></worldbody><deformable>";
    for (int f = 0; f < GetParam().nflex; f++) {
      xml += std::string("<flex dim='2' body='v0 v1 v2 ") +
             (GetParam().welded ? "welded" : "v3") + R"(' element="0 1 2 1 3 2">
          <contact selfcollide="none" contype="0" conaffinity="0"/>
          <elasticity young="1e4" poisson="0.3" thickness="0.01" damping="0.02"
      )";
      xml += std::string("elastic2d='") + elastic2d + "'/></flex>";
    }
    xml += "</deformable></mujoco>";
    char error[1024];
    data.reset();
    model = LoadModelFromString(xml, error, sizeof(error));
    ASSERT_THAT(model.get(), NotNull()) << error;
    ASSERT_EQ(model->nv, 12);
    if (!GetParam().welded) {
      for (int v = 0; v < 4; v++) {
        EXPECT_EQ(model->body_simple[model->flex_vertbodyid[v]], 2);
      }
    }
    data = MakeData(model);
    mj_forward(model.get(), data.get());
  }

  std::vector<mjtNum> Direction() const {
    std::vector<mjtNum> vec(model->nv);
    for (int i = 0; i < model->nv; i++) {
      vec[i] = mju_Halton(i, 3) - 0.5;
    }
    return vec;
  }

  // Dilate in world coordinates, then add a small out-of-plane perturbation.
  // All edges stay in tension, where the stretch Hessian is exact.
  void Deform() {
    for (int v = 0; v < 4; v++) {
      int body = model->body_weldid[model->flex_vertbodyid[v]];
      mjtNum delta[3];
      mju_scl3(delta, data->flexvert_xpos + 3 * v, 0.05);
      delta[2] += 1e-4 * (mju_Halton(v, 2) - 0.5);
      mju_mulMatTVec3(data->qpos + model->body_dofadr[body],
                      data->xmat + 9 * body, delta);
    }
    mj_forward(model.get(), data.get());
  }

  std::vector<mjtNum> Apply(const std::vector<mjtNum>& vec, mjtNum s1,
                            mjtNum s2) {
    std::vector<mjtNum> result(model->nv, 0);
    mjd_flexBend_mul(model.get(), data.get(), result.data(), vec.data(), s1,
                     s2);
    mjd_flexStretch_mul(model.get(), data.get(), result.data(), vec.data(), s1,
                        s2);
    return result;
  }

  void ExpectDerivative(bool velocity) {
    // The stretch operator includes geometric stiffness; its velocity part
    // equals the force Jacobian at the unstretched reference, where that extra
    // term vanishes.
    if (!velocity) {
      Deform();
    }
    auto vec = Direction();
    auto result = Apply(vec, !velocity, velocity);
    MjDataPtr perturbed(mj_copyData(nullptr, model.get(), data.get()));
    mjtNum eps = MjEps(1e-7, 1e-4);
    mju_addToScl(velocity ? perturbed->qvel : perturbed->qpos, vec.data(), eps,
                 model->nv);
    mj_forward(model.get(), perturbed.get());
    const mjtNum* force = velocity ? data->qfrc_damper : data->qfrc_spring;
    const mjtNum* pert_force =
        velocity ? perturbed->qfrc_damper : perturbed->qfrc_spring;
    for (int i = 0; i < model->nv; i++) {
      mjtNum fd = -(pert_force[i] - force[i]) / eps;
      EXPECT_NEAR(result[i], fd, MjTol(1e-3, 5e-2) * mju_max(1, mju_abs(fd)))
          << i;
    }
  }

  void ExpectAssembly(mjtNum s1, mjtNum s2) {
    Deform();
    auto vec = Direction();
    auto result = Apply(vec, s1, s2);
    std::vector<int> rownnz(model->nv), rowadr(model->nv);
    int nnz = mjd_flexStiff_assemble(model.get(), data.get(), rownnz.data(),
                                     rowadr.data(), nullptr, nullptr, s1, s2, 1,
                                     1, nullptr);
    std::vector<int> colind(nnz);
    std::vector<mjtNum> val(nnz);
    mjd_flexStiff_assemble(model.get(), data.get(), rownnz.data(),
                           rowadr.data(), colind.data(), val.data(), s1, s2, 1,
                           1, nullptr);
    for (int i = 0; i < model->nv; i++) {
      mjtNum assembled = mju_dotSparse(val.data() + rowadr[i], vec.data(),
                                       rownnz[i], colind.data() + rowadr[i]);
      EXPECT_THAT(assembled, MjNear(result[i], 1e-10, 1e-5)) << i;
    }
  }

  void ExpectBendingFactorInverse() {
    // Bending-only models use the constant factor, without per-step block
    // preconditioning.
    ASSERT_EQ(data->nefmdof, 0);
    ASSERT_EQ(model->nefm0dof, model->nv);
    int nv = model->nv;
    std::vector<mjtNum> rhs(nv), solution(nv), product(nv);
    mjtNum h = model->opt.timestep;
    for (int col = 0; col < nv; col++) {
      mju_zero(rhs.data(), nv);
      rhs[col] = 1;
      mjd_effPrec(model.get(), data.get(), solution.data(), rhs.data());
      mj_mulM(model.get(), data.get(), product.data(), solution.data());
      mjd_flexBend_mul(model.get(), data.get(), product.data(), solution.data(),
                       h * h, h);
      for (int row = 0; row < nv; row++) {
        EXPECT_THAT(product[row], MjNear(rhs[row], 1e-12, 2e-6))
            << row << ", " << col;
      }
    }
  }

  MjModelPtr model;
  MjDataPtr data;
};

TEST_P(FlexFrameTest, RigidTranslationHasNoDamping) {
  const mjtNum velocity[3] = {0.7, -0.4, 1.2};
  for (int v = 0; v < 4; v++) {
    int body = model->body_weldid[model->flex_vertbodyid[v]];
    mju_mulMatTVec3(data->qvel + model->body_dofadr[body],
                    data->xmat + 9 * body, velocity);
  }
  mj_forward(model.get(), data.get());
  EXPECT_THAT(AsVector(data->qfrc_damper, model->nv),
              Each(MjNear(0, 1e-10, 1e-5)));
}

TEST_P(FlexFrameTest, PositionDerivative) { ExpectDerivative(false); }

TEST_P(FlexFrameTest, VelocityDerivativeAtReference) { ExpectDerivative(true); }

TEST_P(FlexFrameTest, PositionAssembly) { ExpectAssembly(1, 0); }

TEST_P(FlexFrameTest, VelocityAssembly) { ExpectAssembly(0, 1); }

TEST_P(FlexFrameTest, ConstantBendingFactor) {
  ASSERT_NO_FATAL_FAILURE(LoadPatch("bend"));
  ExpectBendingFactorInverse();
}

TEST_P(FlexFrameTest, ConstantBendingFactorAfterFrameChange) {
  ASSERT_NO_FATAL_FAILURE(LoadPatch("bend"));
  int body = mj_name2id(model.get(), mjOBJ_BODY, "v1");
  const mjtNum rotation[4] = {0.5, 0.5, 0.5, 0.5};
  mju_copy4(model->body_quat + 4 * body, rotation);
  mj_setConst(model.get(), data.get());
  mj_forward(model.get(), data.get());
  ExpectBendingFactorInverse();
}

INSTANTIATE_TEST_SUITE_P(
    Frames, FlexFrameTest,
    testing::Values(FlexFrameCase{"Aligned", false, false, 1},
                    FlexFrameCase{"Welded", true, false, 1},
                    FlexFrameCase{"RotatedSiblings", false, true, 1},
                    FlexFrameCase{"SharedFlexBodies", true, true, 2}),
    [](const testing::TestParamInfo<FlexFrameCase>& info) {
      return info.param.name;
    });

// The exact SNH Hessian remains finite through degeneracy and inversion,
// including negative curvature. Assembly, matrix-free products, damping, and
// force derivatives agree.
TEST_F(DerivativeTest, SNHStiffnessThroughInversion) {
  static constexpr char xml[] = R"(
  <mujoco>
    <option gravity="0 0 0" integrator="discrete"/>
    <worldbody>
      <flexcomp name="tet" type="direct" dim="3" mass="1"
                point="0 0 0  1 0 0  .2 .9 0  -.1 .3 1.1" element="0 1 2 3">
        <contact contype="0" conaffinity="0" selfcollide="none"/>
        <elasticity young="1000" poisson=".3" damping=".02"/>
      </flexcomp>
    </worldbody>
  </mujoco>
  )";
  mjSpec* spec = mj_parseXMLString(xml, nullptr, nullptr, 0);
  ASSERT_THAT(spec, NotNull());
  mjs_asFlex(mjs_findElement(spec, mjOBJ_FLEX, "tet"))->elastic3d = 1;
  MjModelPtr m(mj_compile(spec, nullptr));
  mj_deleteSpec(spec);
  ASSERT_THAT(m.get(), NotNull());
  MjDataPtr d = MakeData(m);
  mj_forward(m.get(), d.get());
  std::vector<mjtNum> rest(d->flexvert_xpos, d->flexvert_xpos + 12);
  int nv = m->nv;
  ASSERT_EQ(nv, 12);
  bool negative_curvature = false;
  for (int rank = 0; rank <= 3; rank++) {
    for (mjtNum scale :
         {mjtNum(-1), mjtNum(-.01), mjtNum(0), mjtNum(.01), mjtNum(1.05)}) {
      for (int v = 0; v < 4; v++) {
        int adr = m->body_dofadr[m->flex_vertbodyid[v]];
        mjtNum x = rank > 0 ? 1.05 * rest[3 * v] : 0;
        mjtNum y = rank > 1 ? 1.05 * rest[3 * v + 1] : 0;
        mjtNum z = rank > 2 ? scale * rest[3 * v + 2] : 0;
        if (rank == 3 && scale > 1) {
          // Distinct singular values and nontrivial left/right singular
          // vectors.
          x = 1.06 * (.8 * rest[3 * v] - .6 * rest[3 * v + 1]);
          y = 1.09 * (.6 * rest[3 * v] + .8 * rest[3 * v + 1]);
          z = 1.12 * rest[3 * v + 2];
        }
        d->qpos[adr] = .8 * x - .6 * y - rest[3 * v];
        d->qpos[adr + 1] = .6 * x + .8 * y - rest[3 * v + 1];
        d->qpos[adr + 2] = z - rest[3 * v + 2];
      }
      for (int i = 0; i < nv; i++) d->qvel[i] = mju_Halton(i, 2) - .5;
      mj_forward(m.get(), d.get());
      std::vector<mjtNum> K(nv * nv, 0);
      stretchK_dense(m.get(), d.get(), K.data(), nv, 1, 0);
      std::vector<int> rownnz(nv), rowadr(nv);
      int nnz =
          mjd_flexStiff_assemble(m.get(), d.get(), rownnz.data(), rowadr.data(),
                                 nullptr, nullptr, 1, 0, 0, 1, nullptr);
      std::vector<int> colind(nnz);
      std::vector<mjtNum> values(nnz);
      mjd_flexStiff_assemble(m.get(), d.get(), rownnz.data(), rowadr.data(),
                             colind.data(), values.data(), 1, 0, 0, 1, nullptr);
      for (int i = 0; i < nv; i++) {
        for (int j = 0; j < nv; j++) {
          EXPECT_TRUE(std::isfinite(K[i * nv + j]));
          EXPECT_NEAR(K[i * nv + j], K[j * nv + i], MjTol(1e-10, 1e-3));
        }
        for (int j = rowadr[i]; j < rowadr[i] + rownnz[i]; j++) {
          EXPECT_NEAR(values[j], K[i * nv + colind[j]], MjTol(1e-10, 1e-3));
        }
      }
      std::vector<mjtNum> L = K;
      for (int i = 0; i < nv; i++) L[i * nv + i] += MjTol(1e-8, .01);
      negative_curvature |= mju_cholFactor(L.data(), nv, 0) < nv;
      std::vector<mjtNum> damp(nv, 0);
      mjd_flexStretch_mul(m.get(), d.get(), damp.data(), d->qvel, 0, 1);
      for (int i = 0; i < nv; i++) {
        EXPECT_NEAR(damp[i], -d->qfrc_damper[i], MjTol(1e-10, 1e-3));
      }
      // Central differences verify the full tangent even at rank zero and under
      // reflection.
      m->flex_damping[0] = 0;
      mjtNum eps = MjEps(1e-6, 1e-3);
      for (int j = 0; j < nv; j++) {
        mjtNum saved = d->qpos[j];
        d->qpos[j] = saved + eps;
        mj_forward(m.get(), d.get());
        std::vector<mjtNum> plus(d->qfrc_spring, d->qfrc_spring + nv);
        d->qpos[j] = saved - eps;
        mj_forward(m.get(), d.get());
        for (int i = 0; i < nv; i++) {
          mjtNum fd = -(plus[i] - d->qfrc_spring[i]) / (2 * eps);
          EXPECT_NEAR(K[i * nv + j], fd, MjTol(1e-6, .2));
        }
        d->qpos[j] = saved;
      }
      m->flex_damping[0] = .02;
    }
  }
  EXPECT_TRUE(negative_curvature)
      << "the exact Hessian must not silently clamp negative modes";
}

// K_stretch must be the full Hessian of the stretch force, not just its
// Gauss-Newton part: the geometric (stress-proportional) term is what makes it
// the Jacobian at finite strain. Uniformly dilating the mesh puts every edge in
// tension, so the tensile clamp is inactive and the operator is exact -- with
// only the Gauss-Newton term the finite-difference error is a large fraction of
// the force. FlexStretchDerivatives covers the near-rest limit, where the two
// agree anyway.
TEST_F(DerivativeTest, FlexStretchDerivativesTensile) {
  static const char* const kXml = R"(
  <mujoco>
    <option integrator="discrete"/>
    <worldbody>
      <flexcomp name="cloth" type="grid" count="4 4 1" spacing="0.1 0.1 0.1"
                radius=".01" dim="2" mass="1" pos="0 0 1">
        <contact selfcollide="none" contype="0" conaffinity="0"/>
        <elasticity young="1e4" poisson="0.3" thickness="0.01"
                    elastic2d="stretch" damping="0"/>
      </flexcomp>
    </worldbody>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr model = LoadModelFromString(kXml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  int nv = model->nv;
  ASSERT_EQ(model->nq, nv);  // all slide dofs
  MjDataPtr data = MakeData(model);
  mj_forward(model.get(), data.get());

  // dilate about the flex centroid: every edge stretches by 5%, so every Me is
  // strictly positive
  mjtNum centroid[3] = {0, 0, 0};
  int nvert = model->flex_vertnum[0];
  for (int v = 0; v < nvert; v++) {
    mju_addTo3(centroid, data->flexvert_xpos + 3 * v);
  }
  mju_scl3(centroid, centroid, 1.0 / nvert);
  for (int v = 0; v < nvert; v++) {
    int body = model->flex_vertbodyid[model->flex_vertadr[0] + v];
    int adr = model->body_dofadr[body];
    for (int x = 0; x < 3; x++) {
      data->qpos[adr + x] +=
          0.05 * (data->flexvert_xpos[3 * v + x] - centroid[x]);
    }
  }
  mj_forward(model.get(), data.get());

  std::vector<mjtNum> vec(nv), res(nv, 0);
  for (int i = 0; i < nv; i++) {
    vec[i] = mju_Halton(i, 2) - 0.5;
  }
  mjd_flexStretch_mul(model.get(), data.get(), res.data(), vec.data(), 1, 0);

  mjtNum eps = MjEps(1e-7, 1e-4);
  mjData* perturbed = mj_copyData(NULL, model.get(), data.get());
  mju_addToScl(perturbed->qpos, vec.data(), eps, nv);
  mj_forward(model.get(), perturbed);

  mjtNum max_err = 0, scale = 0;
  for (int i = 0; i < nv; ++i) {
    mjtNum fd = -(perturbed->qfrc_passive[i] - data->qfrc_passive[i]) / eps;
    max_err = mju_max(max_err, mju_abs(res[i] - fd));
    scale = mju_max(scale, mju_abs(fd));
  }
  mj_deleteData(perturbed);

  EXPECT_GT(scale, 1e-3) << "test should exercise nontrivial stretch stiffness";
  EXPECT_LT(max_err, MjTol(1e-4, 1e-3) * scale)
      << "mjd_flexStretch_mul is not the Jacobian of the flex stretch force";
}

// verify mjd_flexStretch_mul (stiffness of the standard-flex stretch force)
// against finite differences of qfrc_passive, plus symmetry, positive
// semi-definiteness and (s1, s2) scale linearity. The model covers both
// element edge tables (dim=2 triangles and dim=3 tets) and a pinned vertex
// (zero-dof body guard).
TEST_F(DerivativeTest, FlexStretchDerivatives) {
  static const char* const kXml = R"(
  <mujoco>
    <option integrator="discrete"/>
    <worldbody>
      <flexcomp name="cloth" type="grid" count="4 4 1" spacing="0.1 0.1 0.1"
                radius=".01" dim="2" mass="1" pos="0 0 1">
        <contact selfcollide="none" contype="0" conaffinity="0"/>
        <elasticity young="1e4" poisson="0.3" thickness="0.01"
                    elastic2d="stretch" damping="50"/>
        <pin id="0"/>
      </flexcomp>
      <flexcomp name="solid" type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
                radius=".01" dim="3" mass="1" pos="1 0 1">
        <contact selfcollide="none" contype="0" conaffinity="0"/>
        <elasticity young="1e4" poisson="0.3" damping="10"/>
      </flexcomp>
    </worldbody>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr model = LoadModelFromString(kXml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  int nv = model->nv;
  ASSERT_EQ(model->nq, nv);  // all slide dofs

  MjDataPtr data = MakeData(model);

  // deform both flexes deterministically, at small strain.
  // FlexStretchDerivativesTensile covers finite strain, where the geometric
  // term of K_stretch is what carries the accuracy; here some edges are
  // compressed and the operator drops their geometric term to stay PSD, so the
  // tolerance stays loose.
  for (int i = 0; i < nv; i++) {
    data->qpos[i] += 5e-4 * (mju_Halton(i, 2) - 0.5);
  }
  mj_forward(model.get(), data.get());

  // part 1: FD verification of K*vec against qfrc_passive (qvel = 0, so the
  // kD elongation term vanishes and qfrc_passive is the pure stretch spring)
  {
    std::vector<mjtNum> vec(nv), res(nv, 0);
    for (int i = 0; i < nv; i++) {
      vec[i] = mju_Halton(i, 2) - 0.5;
    }
    mjd_flexStretch_mul(model.get(), data.get(), res.data(), vec.data(), 1, 0);

    mjtNum eps = MjEps(1e-7, 1e-4);
    mjData* data_perturbed = mj_copyData(NULL, model.get(), data.get());
    mju_addToScl(data_perturbed->qpos, vec.data(), eps, nv);
    mj_forward(model.get(), data_perturbed);

    // qfrc_passive = -dV/dq  =>  -(qfrc_new - qfrc)/eps ~= K * vec.
    // Compare max error against the force scale rather than entrywise:
    // individual near-zero entries are not meaningful, and the dropped
    // compressive geometric term leaves a residual in both flexes.
    mjtNum max_err = 0, scale = 0;
    for (int i = 0; i < nv; ++i) {
      mjtNum fd =
          -(data_perturbed->qfrc_passive[i] - data->qfrc_passive[i]) / eps;
      max_err = mju_max(max_err, mju_abs(res[i] - fd));
      scale = mju_max(scale, mju_abs(fd));
    }
    EXPECT_GT(scale, 1.0) << "test should exercise nontrivial stiffness";
    EXPECT_LT(max_err, MjTol(5e-3, 5e-2) * scale)
        << "stretch stiffness mismatch: max_err " << max_err
        << " at force scale " << scale;
    mj_deleteData(data_perturbed);
  }

  // part 2: symmetry and positive semi-definiteness of the assembled K
  {
    std::vector<mjtNum> K(nv * nv, 0);
    stretchK_dense(model.get(), data.get(), K.data(), nv, 1, 0);

    mjtNum max_asymmetry = 0;
    for (int i = 0; i < nv; i++) {
      for (int j = 0; j < i; j++) {
        max_asymmetry =
            mju_max(max_asymmetry, mju_abs(K[i * nv + j] - K[j * nv + i]));
      }
    }
    EXPECT_THAT(max_asymmetry, MjNear(0, 1e-10, 5e-4))
        << "K_stretch is not symmetric";

    for (int trial = 0; trial < 5; trial++) {
      std::vector<mjtNum> v(nv);
      for (int i = 0; i < nv; i++) {
        v[i] = mju_Halton(i + trial * nv, 3) - 0.5;
      }
      mjtNum vKv = 0;
      for (int i = 0; i < nv; i++) {
        for (int j = 0; j < nv; j++) {
          vKv += v[i] * K[i * nv + j] * v[j];
        }
      }
      EXPECT_GE(vKv, MjTol(-1e-8, -1e-5)) << "K_stretch is not PSD";
    }
  }

  // part 3: (s1, s2) scale linearity across flexes with different damping:
  // mul(s1, s2) == s1*mul(1, 0) + s2*mul(0, 1)
  {
    std::vector<mjtNum> vec(nv), a(nv, 0), b(nv, 0), c(nv, 0);
    for (int i = 0; i < nv; i++) {
      vec[i] = mju_Halton(i, 5) - 0.5;
    }
    mjtNum h = 1e-3;
    mjd_flexStretch_mul(model.get(), data.get(), a.data(), vec.data(), h * h,
                        h);
    mjd_flexStretch_mul(model.get(), data.get(), b.data(), vec.data(), 1, 0);
    mjd_flexStretch_mul(model.get(), data.get(), c.data(), vec.data(), 0, 1);
    for (int i = 0; i < nv; i++) {
      EXPECT_THAT(a[i], MjNear(h * h * b[i] + h * c[i], 1e-12, 1e-5))
          << "scale linearity mismatch at DOF " << i;
    }
  }
}

// verify mjd_flexStiff_assemble against the matrix-free operators: the
// assembled CSR applied to test vectors must reproduce mjd_flexBend_mul +
// mjd_flexStretch_mul at the same state
TEST_F(DerivativeTest, FlexStiffAssemble) {
  static const char* const kXml = R"(
  <mujoco>
    <option integrator="discrete"/>
    <worldbody>
      <flexcomp name="cloth" type="grid" count="4 4 1" spacing="0.1 0.1 0.1"
                radius=".01" dim="2" mass="1" pos="0 0 1">
        <contact selfcollide="none" contype="0" conaffinity="0"/>
        <elasticity young="1e4" poisson="0.3" thickness="0.01"
                    elastic2d="both" damping="7"/>
      </flexcomp>
      <flexcomp name="solid" type="grid" count="3 3 3" spacing="0.1 0.1 0.1"
                radius=".01" dim="3" mass="1" pos="1 0 1">
        <contact selfcollide="none" contype="0" conaffinity="0"/>
        <elasticity young="1e4" poisson="0.3" damping="10"/>
        <pin id="0"/>
      </flexcomp>
    </worldbody>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr model = LoadModelFromString(kXml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  int nv = model->nv;
  MjDataPtr data = MakeData(model);

  // deform deterministically
  for (int i = 0; i < nv; i++) {
    data->qpos[i] += 2e-3 * (mju_Halton(i, 2) - 0.5);
  }
  mj_forward(model.get(), data.get());

  // assemble both terms with a mixed (s1, s2) scale
  mjtNum s1 = 4e-6, s2 = 2e-3;
  std::vector<int> rownnz(nv), rowadr(nv);
  int nnz = mjd_flexStiff_assemble(model.get(), data.get(), rownnz.data(),
                                   rowadr.data(), NULL, NULL, s1, s2,
                                   /*flg_bend=*/1, /*flg_stretch=*/1, NULL);
  ASSERT_GT(nnz, 0);
  std::vector<int> colind(nnz);
  std::vector<mjtNum> val(nnz);
  mjd_flexStiff_assemble(model.get(), data.get(), rownnz.data(), rowadr.data(),
                         colind.data(), val.data(), s1, s2, /*flg_bend=*/1,
                         /*flg_stretch=*/1, NULL);

  // compare CSR apply vs operators on test vectors
  for (int trial = 0; trial < 3; trial++) {
    std::vector<mjtNum> vec(nv), res_op(nv, 0), res_csr(nv, 0);
    for (int i = 0; i < nv; i++) {
      vec[i] = mju_Halton(i + trial * nv, 3) - 0.5;
    }
    mjd_flexBend_mul(model.get(), data.get(), res_op.data(), vec.data(), s1,
                     s2);
    mjd_flexStretch_mul(model.get(), data.get(), res_op.data(), vec.data(), s1,
                        s2);
    for (int i = 0; i < nv; i++) {
      mjtNum sum = 0;
      for (int k = 0; k < rownnz[i]; k++) {
        sum += val[rowadr[i] + k] * vec[colind[rowadr[i] + k]];
      }
      res_csr[i] = sum;
    }
    for (int i = 0; i < nv; i++) {
      EXPECT_THAT(res_csr[i], MjNear(res_op[i], 1e-12, 2e-5))
          << "assembly/operator mismatch at DOF " << i << " trial " << trial;
    }
  }
}

// verify the interp assembly mode: with the K_rot cache supplied, the assembled
// CSR applied to test vectors must reproduce mjd_flexInterp_mul, whose sign
// convention is negated
TEST_F(DerivativeTest, FlexStiffAssembleInterp) {
  static const char* const kXml = R"(
  <mujoco>
    <option integrator="discrete"/>
    <worldbody>
      <flexcomp name="soft" type="grid" count="4 4 4" spacing="0.1 0.1 0.1"
                radius=".01" dim="3" mass="1" pos="0 0 1" dof="trilinear">
        <contact selfcollide="none" contype="0" conaffinity="0"/>
        <elasticity young="1e4" poisson="0.3" damping="2"/>
      </flexcomp>
    </worldbody>
  </mujoco>
  )";

  char error[1024];
  MjModelPtr model = LoadModelFromString(kXml, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  int nv = model->nv;
  MjDataPtr data = MakeData(model);
  ASSERT_EQ(mjd_flexInterpAssemblable(model.get()), 1);

  // deform deterministically, refresh kinematics, cache the corotated stiffness
  for (int i = 0; i < nv; i++) {
    data->qpos[i] += 2e-3 * (mju_Halton(i, 2) - 0.5);
  }
  mj_forward(model.get(), data.get());
  std::vector<mjtNum> krot(model->nflexstiffness, 0);
  mjd_flexInterp_cacheKrot(model.get(), data.get(), krot.data());

  // assemble interp only
  mjtNum s1 = 4e-6, s2 = 2e-3;
  std::vector<int> rownnz(nv), rowadr(nv);
  int nnz = mjd_flexStiff_assemble(
      model.get(), data.get(), rownnz.data(), rowadr.data(), NULL, NULL, s1, s2,
      /*flg_bend=*/0, /*flg_stretch=*/0, krot.data());
  ASSERT_GT(nnz, 0);
  std::vector<int> colind(nnz);
  std::vector<mjtNum> val(nnz);
  mjd_flexStiff_assemble(model.get(), data.get(), rownnz.data(), rowadr.data(),
                         colind.data(), val.data(), s1, s2, /*flg_bend=*/0,
                         /*flg_stretch=*/0, krot.data());

  // compare CSR apply vs the operator called with negated scales (its
  // convention)
  for (int trial = 0; trial < 3; trial++) {
    std::vector<mjtNum> vec(nv), res_op(nv, 0), res_csr(nv, 0);
    for (int i = 0; i < nv; i++) {
      vec[i] = mju_Halton(i + trial * nv, 3) - 0.5;
    }
    mjd_flexInterp_mul(model.get(), data.get(), res_op.data(), vec.data(), -s1,
                       -s2, krot.data());
    for (int i = 0; i < nv; i++) {
      mjtNum sum = 0;
      for (int k = 0; k < rownnz[i]; k++) {
        sum += val[rowadr[i] + k] * vec[colind[rowadr[i] + k]];
      }
      res_csr[i] = sum;
    }
    for (int i = 0; i < nv; i++) {
      EXPECT_THAT(res_csr[i], MjNear(res_op[i], 1e-12, 2e-5))
          << "interp assembly/operator mismatch at DOF " << i << " trial "
          << trial;
    }
  }
}

struct FlexAttachmentCase {
  const char* name;
  const char* joints;
  bool sliding_child = false;
  bool shared_body = false;
};

class FlexAttachmentTest
    : public MujocoTest,
      public testing::WithParamInterface<FlexAttachmentCase> {
 protected:
  void SetUp() override { Load(); }

  std::string Xml(bool independent, const char* elasticity,
                  bool mixed = false) {
    const char* inertia =
        R"(<inertial pos="0 0 0" mass="0.1" diaginertia=".01 .02 .03"/>)";
    const char* slides = R"(<joint type="slide" axis="1 0 0"/>
      <joint type="slide" axis="0 1 0"/><joint type="slide" axis="0 0 1"/>)";
    const char* positions[] = {"0 0 0", ".1 0 0", "0 .1 0", ".1 .1 0"};
    std::string xml = R"(<mujoco><option integrator="discrete" solver="CG"
      gravity="0 0 0" iterations="500"/><worldbody>)";
    int count = independent ? 4 : (GetParam().shared_body ? 2 : 3);
    for (int v = 0; v < count; v++) {
      xml += "<body name='v" + std::to_string(v) + "' pos='" + positions[v] +
             "'>" + inertia + slides + "</body>";
    }
    if (!independent) {
      xml += std::string("<body name='carrier' pos='.08 .08 -.02'>") + inertia +
             GetParam().joints;
      if (GetParam().shared_body) xml += "<body name='v2' pos='-.08 .02 .02'/>";
      xml += "<body name='v3' pos='.02 .02 .02'>";
      if (GetParam().sliding_child) xml += std::string(inertia) + slides;
      xml += "</body></body>";
    }
    if (mixed) {
      xml += R"(<flexcomp name="ordinary" dim="2" type="grid" count="3 3 1"
        spacing=".1 .1 .1" pos="2 0 0" mass="1">
        <contact contype="0" conaffinity="0" selfcollide="none"/>
        <elasticity young="100" thickness=".02" damping=".02" elastic2d="stretch"/>
      </flexcomp>)";
    }
    xml +=
        R"(</worldbody><deformable><flex name="patch" dim="2" body="v0 v1 v2 v3"
      element="0 1 2 1 3 2"><contact contype="0" conaffinity="0" selfcollide="none"/>
      <elasticity young="100" poisson=".3" thickness=".02" damping=".02" elastic2d=")";
    return xml + elasticity + R"("/></flex></deformable></mujoco>)";
  }

  void Load(const char* elasticity = "both", bool mixed = false) {
    char error[1024];
    data.reset();
    model = LoadModelFromString(Xml(false, elasticity, mixed), error,
                                sizeof(error));
    ASSERT_THAT(model.get(), NotNull()) << error;
    data = MakeData(model);
    mj_forward(model.get(), data.get());
    flex = mj_name2id(model.get(), mjOBJ_FLEX, "patch");
    ASSERT_FALSE(mj_flexSimple(model.get(), flex));
  }

  std::vector<mjtNum> Direction(int base = 2) {
    std::vector<mjtNum> result(model->nv);
    for (int i = 0; i < model->nv; i++) result[i] = mju_Halton(i, base) - .5;
    return result;
  }

  void Deform() {
    auto direction = Direction();
    mj_integratePos(model.get(), data->qpos, direction.data(), .05);
    mju_copy(data->qvel, direction.data(), model->nv);
    mj_forward(model.get(), data.get());
  }

  // Independent oracle: public point Jacobians and force application, not the
  // flex helpers.
  std::vector<mjtNum> Gather(const std::vector<mjtNum>& vec) {
    std::vector<mjtNum> result(12), jac(3 * model->nv);
    for (int v = 0; v < 4; v++) {
      int gv = model->flex_vertadr[flex] + v;
      mj_jac(model.get(), data.get(), jac.data(), nullptr,
             data->flexvert_xpos + 3 * gv, model->flex_vertbodyid[gv]);
      mju_mulMatVec(result.data() + 3 * v, jac.data(), vec.data(), 3,
                    model->nv);
    }
    return result;
  }

  std::vector<mjtNum> Scatter(const mjtNum* force) {
    std::vector<mjtNum> result(model->nv, 0);
    for (int v = 0; v < 4; v++) {
      int gv = model->flex_vertadr[flex] + v;
      mj_applyFT(model.get(), data.get(), force + 3 * v, nullptr,
                 data->flexvert_xpos + 3 * gv, model->flex_vertbodyid[gv],
                 result.data());
    }
    return result;
  }

  void Reference(const char* elasticity) {
    char error[1024];
    reference_data.reset();
    reference =
        LoadModelFromString(Xml(true, elasticity), error, sizeof(error));
    ASSERT_THAT(reference.get(), NotNull()) << error;
    reference_data = MakeData(reference);
    auto velocity = Gather(AsVector(data->qvel, model->nv));
    mju_copy(reference_data->qvel, velocity.data(), 12);
    for (int v = 0; v < 4; v++) {
      int body = reference->flex_vertbodyid[v];
      mju_sub3(reference_data->qpos + 3 * v,
               data->flexvert_xpos + 3 * (model->flex_vertadr[flex] + v),
               reference->body_pos + 3 * body);
    }
    mj_forward(reference.get(), reference_data.get());
  }

  void ExpectForce(const char* elasticity, bool damping) {
    ASSERT_NO_FATAL_FAILURE(Load(elasticity));
    Deform();
    ASSERT_NO_FATAL_FAILURE(Reference(elasticity));
    auto expected = Scatter(damping ? reference_data->qfrc_damper
                                    : reference_data->qfrc_spring);
    const mjtNum* actual = damping ? data->qfrc_damper : data->qfrc_spring;
    EXPECT_GT(mju_norm(expected.data(), model->nv), MjTol(1e-7, 1e-4));
    EXPECT_THAT(AsVector(actual, model->nv),
                Pointwise(MjNear(1e-10, 1e-5), expected));
  }

  void ExpectOperator(bool bend) {
    Deform();
    ASSERT_NO_FATAL_FAILURE(Reference("both"));
    auto vec = Direction(3), world = Gather(vec);
    auto apply = bend ? mjd_flexBend_mul : mjd_flexStretch_mul;
    std::vector<mjtNum> world_result(12, 0), result(model->nv + 3, 0);
    apply(reference.get(), reference_data.get(), world_result.data(),
          world.data(), 1, 0);
    result[model->nv] = result[model->nv + 1] = result[model->nv + 2] = 123;
    apply(model.get(), data.get(), result.data(), vec.data(), 1, 0);
    EXPECT_THAT(AsVector(result.data(), model->nv),
                Pointwise(MjNear(1e-10, 2e-5), Scatter(world_result.data())));
    EXPECT_THAT(AsVector(result.data() + model->nv, 3), Each(123));
  }

  void ExpectMetric(bool mixed) {
    ASSERT_NO_FATAL_FAILURE(Load("both", mixed));
    Deform();
    EXPECT_EQ(data->nefmK > 0, mixed);
    auto vec = Direction(3);
    std::vector<mjtNum> expected(model->nv, 0), actual(model->nv, 0);
    mjtNum h = model->opt.timestep;
    mjd_flexBend_mul(model.get(), data.get(), expected.data(), vec.data(),
                     h * h, h);
    mjd_flexStretch_mul(model.get(), data.get(), expected.data(), vec.data(),
                        h * h, h);
    mjd_effMulAdd(model.get(), data.get(), actual.data(), vec.data(), 1);
    EXPECT_THAT(actual, Pointwise(MjNear(1e-12, 1e-7), expected));
  }

  int flex;
  MjModelPtr model, reference;
  MjDataPtr data, reference_data;
};

TEST_P(FlexAttachmentTest, GatherMatchesPointJacobian) {
  Deform();
  auto vec = Direction();
  std::vector<mjtNum> actual(12);
  mj_flexGather(model.get(), data.get(), flex, actual.data(), vec.data());
  EXPECT_THAT(actual, Pointwise(MjNear(1e-12, 1e-6), Gather(vec)));
}

TEST_P(FlexAttachmentTest, ScatterPreservesVirtualWork) {
  Deform();
  auto vec = Direction(), world = Gather(vec);
  std::vector<mjtNum> force(12), actual(model->nv, 0);
  for (int i = 0; i < 12; i++) force[i] = mju_Halton(i, 5) - .5;
  mj_flexScatter(model.get(), data.get(), flex, actual.data(), force.data(), 1);
  EXPECT_THAT(actual, Pointwise(MjNear(1e-12, 1e-6), Scatter(force.data())));
  EXPECT_THAT(mju_dot(vec.data(), actual.data(), model->nv),
              MjNear(mju_dot(world.data(), force.data(), 12), 1e-12, 1e-6));
}

TEST_P(FlexAttachmentTest, BendingSpringIncludesAttachmentReaction) {
  ExpectForce("bend", false);
}
TEST_P(FlexAttachmentTest, BendingDampingIncludesAttachmentReaction) {
  ExpectForce("bend", true);
}
TEST_P(FlexAttachmentTest, StretchSpringIncludesAttachmentReaction) {
  ExpectForce("stretch", false);
}
TEST_P(FlexAttachmentTest, StretchDampingIncludesAttachmentReaction) {
  ExpectForce("stretch", true);
}
TEST_P(FlexAttachmentTest, BendingOperatorMatchesWorldPatch) {
  ExpectOperator(true);
}
TEST_P(FlexAttachmentTest, StretchOperatorMatchesWorldPatch) {
  ExpectOperator(false);
}
TEST_P(FlexAttachmentTest, MatrixFreeMetric) { ExpectMetric(false); }
TEST_P(FlexAttachmentTest, MixedAssembledAndMatrixFreeMetric) {
  ExpectMetric(true);
}

TEST_P(FlexAttachmentTest, BendingDampingDerivative) {
  ASSERT_NO_FATAL_FAILURE(Load("bend"));
  Deform();
  auto vec = Direction(3);
  std::vector<mjtNum> expected(model->nv, 0);
  mjd_flexBend_mul(model.get(), data.get(), expected.data(), vec.data(), 0, 1);
  auto initial = AsVector(data->qfrc_damper, model->nv);
  mjtNum eps = MjEps(1e-6, 1e-3);
  mju_addToScl(data->qvel, vec.data(), eps, model->nv);
  mj_forward(model.get(), data.get());
  for (int i = 0; i < model->nv; i++) {
    EXPECT_THAT(-(data->qfrc_damper[i] - initial[i]) / eps,
                MjNear(expected[i], 1e-9, 1e-5));
  }
}

TEST_P(FlexAttachmentTest, EffectiveSolveResidual) {
  ASSERT_NO_FATAL_FAILURE(Load("both", true));
  Deform();
  auto rhs = Direction(3);
  std::vector<mjtNum> solution(model->nv), product(model->nv);
  mjd_effSolve(model.get(), data.get(), solution.data(), rhs.data());
  mj_mulM(model.get(), data.get(), product.data(), solution.data());
  mjd_effMulAdd(model.get(), data.get(), product.data(), solution.data(), 1);
  EXPECT_THAT(product, Pointwise(MjNear(2e-8, 2e-5), rhs));
}

TEST_P(FlexAttachmentTest, MetricIsSymmetricAndPositive) {
  Deform();
  auto u = Direction(3), v = Direction(5);
  std::vector<mjtNum> Ku(model->nv, 0), Kv(model->nv, 0);
  mjd_effMulAdd(model.get(), data.get(), Ku.data(), u.data(), 1);
  mjd_effMulAdd(model.get(), data.get(), Kv.data(), v.data(), 1);
  EXPECT_THAT(mju_dot(u.data(), Kv.data(), model->nv),
              MjNear(mju_dot(v.data(), Ku.data(), model->nv), 1e-12, 1e-7));
  EXPECT_GE(mju_dot(u.data(), Ku.data(), model->nv), -MjTol(1e-12, 1e-7));
}

TEST_P(FlexAttachmentTest, AttachmentCannotUseFreeGyroShortcut) {
  int body = mj_name2id(model.get(), mjOBJ_BODY, "carrier");
  EXPECT_FALSE(
      mjd_freeGyroPossible(model.get(), data.get(), model->body_jntadr[body]));
}

TEST_P(FlexAttachmentTest, RequiresCGForDiscreteIntegrator) {
  std::string xml = Xml(false, "both");
  xml.replace(xml.find("solver=\"CG\""), 11, "solver=\"Newton\"");
  char error[1024];
  auto rejected = LoadModelFromString(xml, error, sizeof(error));
  EXPECT_EQ(rejected.get(), nullptr);
  EXPECT_THAT(error, HasSubstr("general attachments requires solver='CG'"));
}

INSTANTIATE_TEST_SUITE_P(
    Attachments, FlexAttachmentTest,
    testing::Values(
        FlexAttachmentCase{"Hinge", "<joint axis='0 1 0'/>"},
        FlexAttachmentCase{"Ball", "<joint type='ball'/>"},
        FlexAttachmentCase{"Free", "<freejoint/>"},
        FlexAttachmentCase{"ReorderedSlides", R"(
                      <joint type="slide" axis="0 1 0"/>
                      <joint type="slide" axis="1 0 0"/>
                      <joint type="slide" axis="0 0 -1"/>)"},
        FlexAttachmentCase{"MovingAncestor", "<joint type='ball'/>", true},
        FlexAttachmentCase{"SharedBody", "<joint type='ball'/>", false, true}),
    [](const testing::TestParamInfo<FlexAttachmentCase>& info) {
      return info.param.name;
    });

TEST_F(DerivativeTest, EffSolve) {
  // relative residual of (M+K)x - b after mjd_effSolve
  auto solve_residual = [](const mjModel* m, mjData* d) {
    int nv = m->nv;
    std::vector<mjtNum> b(nv), x(nv), r(nv);
    for (int i = 0; i < nv; i++) {
      b[i] = mju_Halton(i, 3) - 0.5;
    }
    mjd_effSolve(m, d, x.data(), b.data());
    mju_mulSymVecSparse(r.data(), d->M, x.data(), nv, m->M_rownnz, m->M_rowadr,
                        m->M_colind);
    mjd_effMulAdd(m, d, r.data(), x.data(), /*flg_contact=*/1);
    mju_subFrom(r.data(), b.data(), nv);
    return mju_norm(r.data(), nv) / mju_norm(b.data(), nv);
  };

  // stretch + bending cloth on world: per-step factor, exact
  static const char* const kXmlBoth = R"(
  <mujoco>
    <option solver="CG" integrator="discrete"/>
    <worldbody>
      <flexcomp name="cloth" type="grid" count="6 6 1" spacing="0.05 0.05 0.05"
                radius=".005" dim="2" mass="0.5" pos="0 0 1" dof="full">
        <contact selfcollide="none" contype="0" conaffinity="0"/>
        <elasticity young="1e3" poisson="0.2" damping="0.1" elastic2d="both" thickness="0.01"/>
      </flexcomp>
    </worldbody>
  </mujoco>
  )";
  char error[1024];
  MjModelPtr model = LoadModelFromString(kXmlBoth, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  MjDataPtr data = MakeData(model);
  mj_forward(model.get(), data.get());
  ASSERT_GE(data->efm_active, 1);
  EXPECT_GT(data->nefmK, 0);
  EXPECT_GT(data->nefmdof, 0);
  EXPECT_LT(solve_residual(model.get(), data.get()), MjTol(1e-8, 1e-4));

  // bending-only cloth: no CSR or per-step factor, constant factor covers,
  // exact
  static const char* const kXmlBend = R"(
  <mujoco>
    <option solver="CG" integrator="discrete"/>
    <worldbody>
      <flexcomp name="cloth" type="grid" count="6 6 1" spacing="0.05 0.05 0.05"
                radius=".005" dim="2" mass="0.5" pos="0 0 1" dof="full">
        <contact selfcollide="none" contype="0" conaffinity="0"/>
        <elasticity young="1e3" poisson="0.2" damping="0.1" elastic2d="bend" thickness="0.01"/>
      </flexcomp>
    </worldbody>
  </mujoco>
  )";
  model = LoadModelFromString(kXmlBend, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  data = MakeData(model);
  mj_forward(model.get(), data.get());
  ASSERT_GE(data->efm_active, 1);
  EXPECT_EQ(data->nefmK, 0);
  EXPECT_EQ(data->nefmdof, 0);
  EXPECT_GT(model->nefm0dof, 0);
  EXPECT_LT(solve_residual(model.get(), data.get()), MjTol(1e-8, 1e-4));

  // A cloth under a jointed parent uses the matrix-free attachment mapping;
  // the iterative solve must still meet its tolerance.
  static const char* const kXmlMoving = R"(
  <mujoco>
    <option solver="CG" integrator="discrete"/>
    <worldbody>
      <body name="base" pos="0 0 1">
        <joint type="slide" axis="0 0 1"/>
        <geom type="sphere" size=".01" mass="1" contype="0" conaffinity="0"/>
        <flexcomp name="cloth" type="grid" count="6 6 1" spacing="0.05 0.05 0.05"
                  radius=".005" dim="2" mass="0.5" pos="0 0 0" dof="full">
          <contact selfcollide="none" contype="0" conaffinity="0"/>
          <elasticity young="1e3" poisson="0.2" damping="0.1" elastic2d="stretch" thickness="0.01"/>
        </flexcomp>
      </body>
    </worldbody>
  </mujoco>
  )";
  model = LoadModelFromString(kXmlMoving, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  data = MakeData(model);
  mj_forward(model.get(), data.get());
  ASSERT_GE(data->efm_active, 1);
  EXPECT_LT(solve_residual(model.get(), data.get()), 1e-4);
}

// A cloth with per-step stretch stiffness, used by the two tests below.
// Independent flex slides with an articulated descendant: the mass matrix
// couples the covered vertex triple to an uncovered hinge in the same kinematic
// tree.
static const char* const kStretchCloth = R"(
<mujoco>
  <option solver="CG" integrator="discrete"/>
  <default>
    <geom type="sphere" size=".01" mass=".1" contype="0" conaffinity="0"/>
  </default>
  <worldbody>
    <body name="v0">
      <joint type="slide" axis="1 0 0"/>
      <joint type="slide" axis="0 1 0"/>
      <joint type="slide" axis="0 0 1"/>
      <geom/>
      <body pos=".02 0 0">
        <joint axis="0 1 0"/>
        <geom pos=".1 0 0"/>
      </body>
    </body>
    <body name="v1" pos=".1 0 0">
      <joint type="slide" axis="1 0 0"/>
      <joint type="slide" axis="0 1 0"/>
      <joint type="slide" axis="0 0 1"/>
      <geom/>
    </body>
    <body name="v2" pos="0 .1 0">
      <joint type="slide" axis="1 0 0"/>
      <joint type="slide" axis="0 1 0"/>
      <joint type="slide" axis="0 0 1"/>
      <geom/>
    </body>
    <body name="v3" pos=".1 .1 0">
      <joint type="slide" axis="1 0 0"/>
      <joint type="slide" axis="0 1 0"/>
      <joint type="slide" axis="0 0 1"/>
      <geom/>
    </body>
  </worldbody>
  <deformable>
    <flex dim="2" body="v0 v1 v2 v3" element="0 1 2 1 3 2">
      <contact selfcollide="none" contype="0" conaffinity="0"/>
      <elasticity young="1e3" poisson=".2" damping=".1" elastic2d="stretch" thickness=".01"/>
    </flex>
  </deformable>
</mujoco>
)";

// A metric too ill-conditioned for the 3x3 blocks exhausts mjd_effSolve's
// iteration budget; it must report that rather than return an under-converged
// qacc_smooth silently. Forced by conditioning rather than by an unreachable
// opt.tolerance, which cannot be expressed in single precision: there the
// residual reaches exactly zero and CG breaks down on a converged solve.
TEST_F(DerivativeTest, EffSolveCapWarns) {
  static const char* const kStiffCloth = R"(
  <mujoco>
    <option solver="CG" integrator="discrete"/>
    <worldbody>
      <body name="base" pos="0 0 1">
        <joint type="slide" axis="0 0 1"/>
        <geom type="sphere" size=".01" mass="1" contype="0" conaffinity="0"/>
        <flexcomp name="cloth" type="grid" count="24 24 1" dim="2"
                  spacing="0.02 0.02 0.02" radius=".002" mass="0.02" dof="full">
          <contact selfcollide="none" contype="0" conaffinity="0"/>
          <elasticity young="1e9" poisson="0.45" damping="0" elastic2d="stretch"
                      thickness="0.02"/>
        </flexcomp>
      </body>
    </worldbody>
  </mujoco>
  )";
  char error[1024];
  MjModelPtr model = LoadModelFromString(kStiffCloth, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  MjDataPtr data = MakeData(model);
  // Excite in-plane stretch modes; gravity alone translates this flat cloth.
  for (int i = 0; i < model->nv; i++) {
    data->qfrc_applied[i] = mju_Halton(i, 2) - 0.5;
  }

  MockWarningHandler warning_handler;
  // both: the specific cause, then the mjWARN_INERTIA it is reported through
  warning_handler.ExpectWarnings("Flex stiffness is too ill-conditioned");
  warning_handler.ExpectWarnings("Inertia matrix is too close to singular");
  mj_forward(model.get(), data.get());
  testing::Mock::VerifyAndClearExpectations(&warning_handler);
}

// PCG requires a symmetric preconditioner. mjd_effPrec must satisfy
// u.P(v) == v.P(u); it did not when the covered and uncovered dofs shared a
// kinematic tree, which is what the articulated vertex body here exercises.
TEST_F(DerivativeTest, EffPrecIsSymmetric) {
  char error[1024];
  MjModelPtr model = LoadModelFromString(kStretchCloth, error, sizeof(error));
  ASSERT_THAT(model.get(), NotNull()) << error;
  MjDataPtr data = MakeData(model);
  mj_forward(model.get(), data.get());
  ASSERT_GE(data->efm_active, 1);
  ASSERT_EQ(data->nefmdof, 4);
  ASSERT_EQ(model->nv, 13);

  int nv = model->nv;
  std::vector<mjtNum> u(nv), v(nv), Pu(nv), Pv(nv);
  for (int trial = 0; trial < 5; trial++) {
    for (int i = 0; i < nv; i++) {
      u[i] = mju_Halton(i + trial * nv, 2) - 0.5;
      v[i] = mju_Halton(i + trial * nv, 5) - 0.5;
    }
    mjd_effPrec(model.get(), data.get(), Pu.data(), u.data());
    mjd_effPrec(model.get(), data.get(), Pv.data(), v.data());
    mjtNum a = mju_dot(v.data(), Pu.data(), nv);
    mjtNum b = mju_dot(u.data(), Pv.data(), nv);
    EXPECT_THAT(a, MjNear(b, 1e-10, 1e-4))
        << "preconditioner is not symmetric, trial " << trial;
  }
}

}  // namespace
}  // namespace mujoco