/**
 * engines/c/dynamics_lib.c
 * -------------------------
 * 纯计算 DLL：无文件 I/O，所有数据由 Python 以 NumPy 数组传入，
 * 结果直接写入 Python 预分配的输出数组。
 * 算法与 main.c 和 compute.py 保持完全一致。
 *
 * 编译（Windows DLL）:
 *   gcc -O3 -march=native -shared -o build/dynamics_c.dll dynamics_lib.c -lm
 * 编译（Linux .so）:
 *   gcc -O3 -march=native -shared -fPIC -o build/dynamics_c.so dynamics_lib.c -lm
 * 编译（macOS .dylib）:
 *   gcc -O3 -march=native -dynamiclib -o build/dynamics_c.dylib dynamics_lib.c -lm
 */

#ifdef _WIN32
#  define EXPORT __declspec(dllexport)
#else
#  define EXPORT __attribute__((visibility("default")))
#endif

#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>

/* ── 驱动力结构体 ─────────────────────────────────────────── */
typedef struct {
    int     n_drivers;
    const int    *idx;         /* [n_drivers] 0-based local atom index */
    const double *amp;         /* [n_drivers*3] (ax,ay,az) interleaved */
    const double *freq;        /* [n_drivers*3] */
    const double *phi;         /* [n_drivers*3] radians */
    const double *eq;          /* [n_drivers*3] equilibrium positions */
    const double *ncycles;     /* [n_drivers] 0=unlimited */
    const int    *has_period;  /* [n_drivers] */
    /* mutable freeze positions (allocated internally) */
    double *freeze;            /* [n_drivers*3] */
} Drivers;

/* ── 加速度：保守力（弹簧键 + 均匀重力场）────────────────── */
static void accel_conservative(
    int n, const double *x, const double *y, const double *z,
    const double *m,
    double Gx, double Gy, double Gz,
    int gravity_field, int elastic_force,
    int n_bonds, const int *bond_pairs,
    const double *bond_k, const double *bond_r0,
    double *ax, double *ay, double *az)
{
    for (int i = 0; i < n; i++) {
        ax[i] = gravity_field ? Gx : 0.0;
        ay[i] = gravity_field ? Gy : 0.0;
        az[i] = gravity_field ? Gz : 0.0;
    }

    if (!elastic_force || n_bonds == 0) return;

    for (int b = 0; b < n_bonds; b++) {
        int  ii = bond_pairs[b*2];
        int  jj = bond_pairs[b*2+1];
        double dx = x[jj] - x[ii];
        double dy = y[jj] - y[ii];
        double dz = z[jj] - z[ii];
        double dist = sqrt(dx*dx + dy*dy + dz*dz);
        if (dist < 1e-12) continue;
        double k  = bond_k[b];
        double r0 = bond_r0[b];
        double fac = k * (dist - r0) / dist;
        double fx = fac * dx, fy = fac * dy, fz_b = fac * dz;
        ax[ii] += fx / m[ii];  ay[ii] += fy / m[ii];  az[ii] += fz_b / m[ii];
        ax[jj] -= fx / m[jj];  ay[jj] -= fy / m[jj];  az[jj] -= fz_b / m[jj];
    }
}

/* ── 完整加速度（含阻尼）────────────────────────────────── */
static void accel_full(
    int n, const double *x, const double *y, const double *z,
    const double *vx, const double *vy, const double *vz,
    const double *m,
    double Gx, double Gy, double Gz,
    double Bx, double By, double Bz,
    int gravity_field, int elastic_force, int damping_force,
    int n_bonds, const int *bond_pairs,
    const double *bond_k, const double *bond_r0,
    double *ax, double *ay, double *az)
{
    accel_conservative(n, x, y, z, m, Gx, Gy, Gz,
                       gravity_field, elastic_force,
                       n_bonds, bond_pairs, bond_k, bond_r0,
                       ax, ay, az);
    if (damping_force) {
        for (int i = 0; i < n; i++) {
            ax[i] -= Bx * vx[i] / m[i];
            ay[i] -= By * vy[i] / m[i];
            az[i] -= Bz * vz[i] / m[i];
        }
    }
}

/* ── 边界：反弹（与 main.c limit_in_box 一致）────────────── */
static inline void _limit1(double *p, double *v, double lo, double hi) {
    if (*p > hi) { *p = hi; *v = -fabs(*v); }
    if (*p < lo) { *p = lo; *v =  fabs(*v); }
}

/* ── 边界：回绕（与 main.c wrap_position 一致）──────────── */
static inline void _wrap1(double *p, double lo, double hi) {
    if (*p > hi) *p = lo;
    if (*p < lo) *p = hi;
}

/* ── 边界 + 固定约束（与 main.c apply_step 末尾一致）──────── */
static void apply_boundary_and_constraints(
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    const int *fixed, const double *pos_init,
    double box_a)
{
    double lo = -box_a, hi = box_a;

    /* 反弹 */
    for (int i = 0; i < n; i++) {
        if (fixed[i*3] && fixed[i*3+1] && fixed[i*3+2]) continue;
        _limit1(&x[i], &vx[i], lo, hi);
        _limit1(&y[i], &vy[i], lo, hi);
        _limit1(&z[i], &vz[i], lo, hi);
    }

    /* 回绕 */
    for (int i = 0; i < n; i++) {
        _wrap1(&x[i], lo, hi);
        _wrap1(&y[i], lo, hi);
        _wrap1(&z[i], lo, hi);
    }

    /* 逐自由度固定约束：与 main.c 和 Python apply_fixed_constraints 一致 */
    for (int i = 0; i < n; i++) {
        if (fixed[i*3+0]) { x[i] = pos_init[i*3+0]; vx[i] = 0.0; }
        if (fixed[i*3+1]) { y[i] = pos_init[i*3+1]; vy[i] = 0.0; }
        if (fixed[i*3+2]) { z[i] = pos_init[i*3+2]; vz[i] = 0.0; }
    }
}

/* ══════════════════════════════════════════════════════════
 * 蛙跳法（与 main.c leapfrog_step 完全一致）
 * x(t), v(t-dt/2) → x(t+dt), v(t+dt/2)
 * 无阻尼：纯辛积分。有阻尼：半隐式处理 α = B·dt/(2m)
 * ══════════════════════════════════════════════════════════ */
static void leapfrog_step(
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    const double *m, const int *fixed,
    double Gx, double Gy, double Gz,
    double Bx, double By, double Bz,
    int gravity_field, int elastic_force, int damping_force,
    int n_bonds, const int *bp, const double *bk, const double *br0,
    double dt)
{
    double *ax = (double*)alloca(n*sizeof(double)*3);
    double *ay = ax+n; double *az = ay+n;

    accel_conservative(n, x, y, z, m, Gx, Gy, Gz,
                       gravity_field, elastic_force,
                       n_bonds, bp, bk, br0, ax, ay, az);

    int has_damp = damping_force && (Bx != 0.0 || By != 0.0 || Bz != 0.0);
    for (int i = 0; i < n; i++) {
        if (fixed[i*3] && fixed[i*3+1] && fixed[i*3+2]) continue;
        if (has_damp) {
            double ax_ = Bx*dt/(2.0*m[i]);
            double ay_ = By*dt/(2.0*m[i]);
            double az_ = Bz*dt/(2.0*m[i]);
            vx[i] = (vx[i]*(1.0-ax_) + ax[i]*dt) / (1.0+ax_);
            vy[i] = (vy[i]*(1.0-ay_) + ay[i]*dt) / (1.0+ay_);
            vz[i] = (vz[i]*(1.0-az_) + az[i]*dt) / (1.0+az_);
        } else {
            vx[i] += ax[i]*dt;
            vy[i] += ay[i]*dt;
            vz[i] += az[i]*dt;
        }
        x[i] += vx[i]*dt;
        y[i] += vy[i]*dt;
        z[i] += vz[i]*dt;
    }
}

/* ══════════════════════════════════════════════════════════
 * 显式欧拉法（与 main.c explicit_euler_step 一致）
 * ══════════════════════════════════════════════════════════ */
static void euler_step(
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    const double *m, const int *fixed,
    double Gx, double Gy, double Gz,
    double Bx, double By, double Bz,
    int gravity_field, int elastic_force, int damping_force,
    int n_bonds, const int *bp, const double *bk, const double *br0,
    double dt)
{
    double *ax = (double*)alloca(n*sizeof(double)*3);
    double *ay = ax+n; double *az = ay+n;
    accel_full(n, x, y, z, vx, vy, vz, m, Gx, Gy, Gz, Bx, By, Bz,
               gravity_field, elastic_force, damping_force,
               n_bonds, bp, bk, br0, ax, ay, az);
    for (int i = 0; i < n; i++) {
        if (fixed[i*3] && fixed[i*3+1] && fixed[i*3+2]) continue;
        x[i] += vx[i]*dt;  y[i] += vy[i]*dt;  z[i] += vz[i]*dt;
        vx[i]+= ax[i]*dt;  vy[i]+= ay[i]*dt;  vz[i]+= az[i]*dt;
    }
}

/* ══════════════════════════════════════════════════════════
 * 隐式欧拉法（与 main.c implicit_euler_step 完全一致）
 *
 * main.c 逻辑：
 *   1. 用 v_next ≈ (v + G·dt)/(1 + γ·dt) 预测（只含重力+阻尼，不含弹簧）
 *   2. 用 (x, v_next) 计算完整加速度 a_next
 *   3. v += a_next·dt;  x += v·dt
 * ══════════════════════════════════════════════════════════ */
static void implicit_euler_step(
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    const double *m, const int *fixed,
    double Gx, double Gy, double Gz,
    double Bx, double By, double Bz,
    int gravity_field, int elastic_force, int damping_force,
    int n_bonds, const int *bp, const double *bk, const double *br0,
    double dt)
{
    double *vxn = (double*)alloca(n*sizeof(double)*3);
    double *vyn = vxn+n; double *vzn = vyn+n;

    for (int i = 0; i < n; i++) {
        if (fixed[i*3] && fixed[i*3+1] && fixed[i*3+2]) {
            vxn[i] = vyn[i] = vzn[i] = 0.0; continue;
        }
        double gx = Bx / m[i], gy = By / m[i], gz = Bz / m[i];
        vxn[i] = (vx[i] + Gx*dt) / (1.0 + gx*dt);
        vyn[i] = (vy[i] + Gy*dt) / (1.0 + gy*dt);
        vzn[i] = (vz[i] + Gz*dt) / (1.0 + gz*dt);
    }

    double *ax = (double*)alloca(n*sizeof(double)*3);
    double *ay = ax+n; double *az = ay+n;
    accel_full(n, x, y, z, vxn, vyn, vzn, m, Gx, Gy, Gz, Bx, By, Bz,
               gravity_field, elastic_force, damping_force,
               n_bonds, bp, bk, br0, ax, ay, az);

    for (int i = 0; i < n; i++) {
        if (fixed[i*3] && fixed[i*3+1] && fixed[i*3+2]) continue;
        vx[i] += ax[i]*dt;
        vy[i] += ay[i]*dt;
        vz[i] += az[i]*dt;
        x[i]  += vx[i]*dt;
        y[i]  += vy[i]*dt;
        z[i]  += vz[i]*dt;
    }
}

/* ══════════════════════════════════════════════════════════
 * 中点法（与 main.c midpoint_step 完全一致）
 *
 * main.c 逻辑：
 *   1. a = accel(x, v)
 *   2. xm = x + 0.5·v·dt;  vm = v + 0.5·a·dt
 *   3. x  = x + vm·dt      (位置更新用 vm，即中点速度)
 *   4. am = accel(xm, vm)
 *   5. v  = v + am·dt
 * ══════════════════════════════════════════════════════════ */
static void midpoint_step(
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    const double *m, const int *fixed,
    double Gx, double Gy, double Gz,
    double Bx, double By, double Bz,
    int gravity_field, int elastic_force, int damping_force,
    int n_bonds, const int *bp, const double *bk, const double *br0,
    double dt)
{
    /* Allocate in one block for cache locality */
    double *buf = (double*)alloca(n*sizeof(double)*9);
    double *ax  = buf;
    double *ay  = ax+n; double *az  = ay+n;
    double *xm  = az+n; double *ym  = xm+n; double *zm  = ym+n;
    double *vxm = zm+n; double *vym = vxm+n; double *vzm = vym+n;

    accel_full(n, x, y, z, vx, vy, vz, m, Gx, Gy, Gz, Bx, By, Bz,
               gravity_field, elastic_force, damping_force,
               n_bonds, bp, bk, br0, ax, ay, az);

    for (int i = 0; i < n; i++) {
        if (fixed[i*3] && fixed[i*3+1] && fixed[i*3+2]) {
            xm[i]=x[i]; ym[i]=y[i]; zm[i]=z[i];
            vxm[i]=vym[i]=vzm[i]=0.0; continue;
        }
        xm[i]  = x[i]  + 0.5*vx[i]*dt;
        ym[i]  = y[i]  + 0.5*vy[i]*dt;
        zm[i]  = z[i]  + 0.5*vz[i]*dt;
        vxm[i] = vx[i] + 0.5*ax[i]*dt;
        vym[i] = vy[i] + 0.5*ay[i]*dt;
        vzm[i] = vz[i] + 0.5*az[i]*dt;
        /* position updated with midpoint velocity (same as main.c) */
        x[i]   = x[i]  + vxm[i]*dt;
        y[i]   = y[i]  + vym[i]*dt;
        z[i]   = z[i]  + vzm[i]*dt;
    }

    double *axm = (double*)alloca(n*sizeof(double)*3);
    double *aym = axm+n; double *azm = aym+n;
    accel_full(n, xm, ym, zm, vxm, vym, vzm, m, Gx, Gy, Gz, Bx, By, Bz,
               gravity_field, elastic_force, damping_force,
               n_bonds, bp, bk, br0, axm, aym, azm);

    for (int i = 0; i < n; i++) {
        if (fixed[i*3] && fixed[i*3+1] && fixed[i*3+2]) continue;
        vx[i] += axm[i]*dt;
        vy[i] += aym[i]*dt;
        vz[i] += azm[i]*dt;
    }
}

/* ── 驱动力（与 main.c apply_driving_force 一致）────────── */
static void apply_driving(
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    double t, int step, double dt, Drivers *drv)
{
    (void)n;
    if (!drv || drv->n_drivers == 0) return;
    const double TWO_PI = 2.0 * 3.14159265358979323846;

    for (int d = 0; d < drv->n_drivers; d++) {
        int idx = drv->idx[d];
        double fx = drv->freq[d*3+0];
        double fy = drv->freq[d*3+1];
        double fz = drv->freq[d*3+2];

        if (drv->has_period[d]) {
            double mf = fabs(fx) > fabs(fy) ? fabs(fx) : fabs(fy);
            if (fabs(fz) > mf) mf = fabs(fz);
            int period_steps = 0;
            if (mf > 1e-12)
                period_steps = (int)(drv->ncycles[d] / mf / dt);
            if (step > period_steps) {
                x[idx] = drv->freeze[d*3+0];
                y[idx] = drv->freeze[d*3+1];
                z[idx] = drv->freeze[d*3+2];
                vx[idx] = vy[idx] = vz[idx] = 0.0;
                continue;
            }

            double px = drv->eq[d*3+0] + drv->amp[d*3+0]*cos(TWO_PI*fx*t + drv->phi[d*3+0]);
            double py = drv->eq[d*3+1] + drv->amp[d*3+1]*cos(TWO_PI*fy*t + drv->phi[d*3+1]);
            double pz = drv->eq[d*3+2] + drv->amp[d*3+2]*cos(TWO_PI*fz*t + drv->phi[d*3+2]);
            if (step == period_steps) {
                drv->freeze[d*3+0] = px;
                drv->freeze[d*3+1] = py;
                drv->freeze[d*3+2] = pz;
            }
        }

        x[idx]  = drv->eq[d*3+0] + drv->amp[d*3+0]*cos(TWO_PI*fx*t + drv->phi[d*3+0]);
        y[idx]  = drv->eq[d*3+1] + drv->amp[d*3+1]*cos(TWO_PI*fy*t + drv->phi[d*3+1]);
        z[idx]  = drv->eq[d*3+2] + drv->amp[d*3+2]*cos(TWO_PI*fz*t + drv->phi[d*3+2]);
        vx[idx] = -drv->amp[d*3+0]*TWO_PI*fx*sin(TWO_PI*fx*t + drv->phi[d*3+0]);
        vy[idx] = -drv->amp[d*3+1]*TWO_PI*fy*sin(TWO_PI*fy*t + drv->phi[d*3+1]);
        vz[idx] = -drv->amp[d*3+2]*TWO_PI*fz*sin(TWO_PI*fz*t + drv->phi[d*3+2]);
    }
}

/* ══════════════════════════════════════════════════════════
 * 导出函数：run_dynamics
 *
 * 与 main.c 的计算顺序完全一致：
 *   1. leapfrog 初始化 v(-dt/2)
 *   2. 初始驱动 t=0
 *   3. 预热循环（不记录）
 *   4. 记录循环：drive → record → step → boundary → constraints
 *
 * 参数说明（所有数组均为 C-contiguous 行优先 float64/int32）：
 *   n_atoms     原子数
 *   pos_init    初始位置 [n_atoms*3]  x0,y0,z0, x1,y1,z1, ...
 *   vel_init    初始速度 [n_atoms*3]
 *   masses      质量     [n_atoms]
 *   fixed       自由度约束 [n_atoms*3]  int32, 1=固定
 *   n_bonds     键数
 *   bond_pairs  键对     [n_bonds*2]  int32, 0-based local index
 *   bond_k      刚度     [n_bonds]
 *   bond_r0     平衡键长 [n_bonds]
 *   box_a       盒子半边长
 *   dt          时间步长
 *   NT          总步数（含预热）
 *   NSTEP       抽帧间隔
 *   warmup_steps 预热步数
 *   method_id   0=euler 1=implicit 2=midpoint 3=leapfrog
 *   Gx/Gy/Gz    均匀重力场加速度分量
 *   Bx/By/Bz    阻尼系数分量
 *   gravity_field / elastic_force / damping_force  力开关
 *   gravity_strength  原子间引力强度（暂未实现，留接口）
 *   n_drivers   驱动原子数
 *   drv_idx     驱动原子局部索引  [n_drivers]  int32
 *   drv_amp     振幅 [n_drivers*3]
 *   drv_freq    频率 [n_drivers*3]
 *   drv_phi     初相（弧度）[n_drivers*3]
 *   drv_eq      平衡位置 [n_drivers*3]
 *   drv_ncycles 周期数  [n_drivers]  0=不限
 *   drv_has_period  [n_drivers]  int32
 *   n_frames    输出帧数（Python 预计算：(NT-warmup)/NSTEP 向上取整）
 *   out_x/y/z/vx/vy/vz  输出数组 [n_frames*n_atoms]  由 Python 预分配
 *   progress_cb 进度回调（可为 NULL）
 *
 * 返回：0=成功，负数=错误
 * ══════════════════════════════════════════════════════════ */
EXPORT int run_dynamics(
    int n_atoms,
    const double *pos_init,
    const double *vel_init,
    const double *masses,
    const int    *fixed,
    int n_bonds,
    const int    *bond_pairs,
    const double *bond_k,
    const double *bond_r0,
    double box_a, double dt,
    int NT, int NSTEP, int warmup_steps, int method_id,
    double Gx, double Gy, double Gz,
    double Bx, double By, double Bz,
    int gravity_field, int elastic_force, int damping_force,
    double gravity_strength,
    int n_drivers,
    const int    *drv_idx,
    const double *drv_amp,
    const double *drv_freq,
    const double *drv_phi,
    const double *drv_eq,
    const double *drv_ncycles,
    const int    *drv_has_period,
    int n_frames,
    double *out_x, double *out_y, double *out_z,
    double *out_vx, double *out_vy, double *out_vz,
    void (*progress_cb)(int step, int total))
{
    (void)gravity_strength;   /* 原子间引力暂未实现 */

    int n = n_atoms;

    /* ── 工作数组 ── */
    double *x  = (double*)malloc(n*sizeof(double));
    double *y  = (double*)malloc(n*sizeof(double));
    double *z  = (double*)malloc(n*sizeof(double));
    double *vx = (double*)malloc(n*sizeof(double));
    double *vy = (double*)malloc(n*sizeof(double));
    double *vz = (double*)malloc(n*sizeof(double));
    if (!x||!y||!z||!vx||!vy||!vz) return -1;

    for (int i = 0; i < n; i++) {
        x[i]=pos_init[i*3+0]; y[i]=pos_init[i*3+1]; z[i]=pos_init[i*3+2];
        vx[i]=vel_init[i*3+0]; vy[i]=vel_init[i*3+1]; vz[i]=vel_init[i*3+2];
    }

    /* ── 驱动结构 ── */
    Drivers drv;
    drv.n_drivers  = n_drivers;
    drv.idx        = drv_idx;
    drv.amp        = drv_amp;
    drv.freq       = drv_freq;
    drv.phi        = drv_phi;
    drv.eq         = drv_eq;
    drv.ncycles    = drv_ncycles;
    drv.has_period = drv_has_period;
    drv.freeze     = NULL;
    if (n_drivers > 0) {
        drv.freeze = (double*)calloc(n_drivers*3, sizeof(double));
        if (!drv.freeze) { free(x);free(y);free(z);free(vx);free(vy);free(vz); return -2; }
    }

    /* ── 内联步进宏 ── */
#define DO_STEP() do { \
    switch (method_id) { \
    case 0: euler_step(n,x,y,z,vx,vy,vz,masses,fixed,Gx,Gy,Gz,Bx,By,Bz, \
                       gravity_field,elastic_force,damping_force, \
                       n_bonds,bond_pairs,bond_k,bond_r0,dt); break; \
    case 1: implicit_euler_step(n,x,y,z,vx,vy,vz,masses,fixed,Gx,Gy,Gz,Bx,By,Bz, \
                       gravity_field,elastic_force,damping_force, \
                       n_bonds,bond_pairs,bond_k,bond_r0,dt); break; \
    case 2: midpoint_step(n,x,y,z,vx,vy,vz,masses,fixed,Gx,Gy,Gz,Bx,By,Bz, \
                       gravity_field,elastic_force,damping_force, \
                       n_bonds,bond_pairs,bond_k,bond_r0,dt); break; \
    default: leapfrog_step(n,x,y,z,vx,vy,vz,masses,fixed,Gx,Gy,Gz,Bx,By,Bz, \
                       gravity_field,elastic_force,damping_force, \
                       n_bonds,bond_pairs,bond_k,bond_r0,dt); break; \
    } \
    apply_boundary_and_constraints(n,x,y,z,vx,vy,vz,fixed,pos_init,box_a); \
} while(0)

    /* ── 蛙跳法：初始化 v(-dt/2) = v(0) - 0.5·a_c(0)·dt ── */
    if (method_id == 3) {
        double *ax0 = (double*)alloca(n*sizeof(double)*3);
        double *ay0 = ax0+n; double *az0 = ay0+n;
        accel_conservative(n, x, y, z, masses, Gx, Gy, Gz,
                           gravity_field, elastic_force,
                           n_bonds, bond_pairs, bond_k, bond_r0,
                           ax0, ay0, az0);
        for (int i = 0; i < n; i++) {
            if (fixed[i*3] && fixed[i*3+1] && fixed[i*3+2]) continue;
            vx[i] -= 0.5*ax0[i]*dt;
            vy[i] -= 0.5*ay0[i]*dt;
            vz[i] -= 0.5*az0[i]*dt;
        }
    }

    /* ── 初始驱动 t=0（与 main.c 一致：leapfrog init 之后施加）── */
    if (n_drivers > 0) apply_driving(n, x, y, z, vx, vy, vz, 0.0, 0, dt, &drv);

    /* ── 预热（不记录）── */
    for (int s = 0; s < warmup_steps; s++) {
        double tw = (s + 1) * dt;
        if (n_drivers > 0) apply_driving(n, x, y, z, vx, vy, vz, tw, s, dt, &drv);
        DO_STEP();
    }

    /* ── 记录循环 ── */
    int record_steps = NT - warmup_steps;
    int prog_interval = record_steps / 100;
    if (prog_interval < 1) prog_interval = 1;
    int frame_idx = 0;

    for (int s = 0; s < record_steps; s++) {
        if (progress_cb && s % prog_interval == 0 && s > 0)
            progress_cb(s, record_steps);

        double t = (s + warmup_steps) * dt;
        if (n_drivers > 0) apply_driving(n, x, y, z, vx, vy, vz, t, s, dt, &drv);

        /* 抽帧记录（drive 之后，step 之前，与 main.c 一致）*/
        if (s % NSTEP == 0 && frame_idx < n_frames) {
            int base = frame_idx * n;
            for (int i = 0; i < n; i++) {
                out_x [base+i] = x[i];   out_y [base+i] = y[i];   out_z [base+i] = z[i];
                out_vx[base+i] = vx[i];  out_vy[base+i] = vy[i];  out_vz[base+i] = vz[i];
            }
            frame_idx++;
        }
        DO_STEP();
    }

#undef DO_STEP

    free(x); free(y); free(z);
    free(vx); free(vy); free(vz);
    if (drv.freeze) free(drv.freeze);
    return 0;
}