/**
 * engines/cpp/dynamics_lib.cpp
 * -----------------------------
 * 纯计算 DLL（C++ 版）：无文件 I/O，所有数据由 Python 以 NumPy 数组传入。
 * 算法与 main.cpp / compute.py 保持完全一致。
 *
 * 编译（Windows）:
 *   g++ -O3 -march=native -std=c++17 -shared -o build/dynamics_cpp.dll dynamics_lib.cpp
 * 编译（Linux）:
 *   g++ -O3 -march=native -std=c++17 -shared -fPIC -o build/dynamics_cpp.so dynamics_lib.cpp
 * 编译（macOS）:
 *   g++ -O3 -march=native -std=c++17 -dynamiclib -o build/dynamics_cpp.dylib dynamics_lib.cpp
 */

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

#include <cmath>
#include <cstring>
#include <cstdlib>
#include <vector>

/* ── 驱动力结构体 ─────────────────────────────────────────── */
struct Drivers {
    int           n_drivers  = 0;
    const int    *idx        = nullptr;
    const double *amp        = nullptr;
    const double *freq       = nullptr;
    const double *phi        = nullptr;
    const double *eq         = nullptr;
    const double *ncycles    = nullptr;
    const int    *has_period = nullptr;
    std::vector<double> freeze;   /* [n_drivers*3] 冻结位置（period 结束时锁定）*/
};

/* ── 加速度：保守力（弹簧键 + 均匀重力场）────────────────── */
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], dy = y[jj]-y[ii], dz = z[jj]-z[ii];
        double dist = std::sqrt(dx*dx + dy*dy + dz*dz);
        if (dist < 1e-12) continue;
        double fac = bond_k[b] * (dist - bond_r0[b]) / 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];
        }
    }
}

/* ── 边界：反弹 ──────────────────────────────────────────── */
static inline void _limit1(double &p, double &v, double lo, double hi) {
    if (p > hi) { p = hi; v = -std::fabs(v); }
    if (p < lo) { p = lo; v =  std::fabs(v); }
}

/* ── 边界：回绕 ──────────────────────────────────────────── */
static inline void _wrap1(double &p, double lo, double hi) {
    if (p > hi) p = lo;
    if (p < lo) p = hi;
}

/* ── 边界 + 固定约束 ────────────────────────────────────── */
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);
    }
    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.cpp leapfrog_step 完全一致）
 * ══════════════════════════════════════════════════════════ */
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)
{
    std::vector<double> buf(n * 3);
    double *ax = buf.data(), *ay = ax+n, *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);
    bool 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;
    }
}

/* ══════════════════════════════════════════════════════════
 * 显式欧拉法
 * ══════════════════════════════════════════════════════════ */
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)
{
    std::vector<double> buf(n * 3);
    double *ax = buf.data(), *ay = ax+n, *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.cpp implicit_euler_step 完全一致）
 * ══════════════════════════════════════════════════════════ */
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)
{
    std::vector<double> vbuf(n * 3), abuf(n * 3);
    double *vxn = vbuf.data(), *vyn = vxn+n, *vzn = vyn+n;
    double *ax  = abuf.data(), *ay  = ax+n,  *az  = ay+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);
    }
    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.cpp midpoint_step 完全一致）
 * ══════════════════════════════════════════════════════════ */
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)
{
    std::vector<double> buf(n * 9);
    double *ax  = buf.data();
    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;
        x[i]   = x[i]  + vxm[i]*dt;
        y[i]   = y[i]  + vym[i]*dt;
        z[i]   = z[i]  + vzm[i]*dt;
    }

    std::vector<double> abuf(n * 3);
    double *axm = abuf.data(), *aym = axm+n, *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;
    }
}

/* ── 驱动力 ─────────────────────────────────────────────── */
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.n_drivers == 0) return;
    constexpr 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 = std::fabs(fx) > std::fabs(fy) ? std::fabs(fx) : std::fabs(fy);
            if (std::fabs(fz) > mf) mf = std::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]*std::cos(TWO_PI*fx*t + drv.phi[d*3+0]);
            double py = drv.eq[d*3+1] + drv.amp[d*3+1]*std::cos(TWO_PI*fy*t + drv.phi[d*3+1]);
            double pz = drv.eq[d*3+2] + drv.amp[d*3+2]*std::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]*std::cos(TWO_PI*fx*t + drv.phi[d*3+0]);
        y[idx]  = drv.eq[d*3+1] + drv.amp[d*3+1]*std::cos(TWO_PI*fy*t + drv.phi[d*3+1]);
        z[idx]  = drv.eq[d*3+2] + drv.amp[d*3+2]*std::cos(TWO_PI*fz*t + drv.phi[d*3+2]);
        vx[idx] = -drv.amp[d*3+0]*TWO_PI*fx*std::sin(TWO_PI*fx*t + drv.phi[d*3+0]);
        vy[idx] = -drv.amp[d*3+1]*TWO_PI*fy*std::sin(TWO_PI*fy*t + drv.phi[d*3+1]);
        vz[idx] = -drv.amp[d*3+2]*TWO_PI*fz*std::sin(TWO_PI*fz*t + drv.phi[d*3+2]);
    }
}

/* ══════════════════════════════════════════════════════════
 * 导出函数：run_dynamics（接口与 C 版完全相同）
 * ══════════════════════════════════════════════════════════ */
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;

    std::vector<double> xv(n), yv(n), zv(n);
    std::vector<double> vxv(n), vyv(n), vzv(n);
    for (int i = 0; i < n; i++) {
        xv[i]=pos_init[i*3+0]; yv[i]=pos_init[i*3+1]; zv[i]=pos_init[i*3+2];
        vxv[i]=vel_init[i*3+0]; vyv[i]=vel_init[i*3+1]; vzv[i]=vel_init[i*3+2];
    }
    double *x=xv.data(), *y=yv.data(), *z=zv.data();
    double *vx=vxv.data(), *vy=vyv.data(), *vz=vzv.data();

    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;
    if (n_drivers > 0)
        drv.freeze.assign(n_drivers * 3, 0.0);

#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) */
    if (method_id == 3) {
        std::vector<double> ibuf(n * 3);
        double *ax0=ibuf.data(), *ay0=ax0+n, *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 */
    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 = std::max(1, record_steps / 100);
    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);

        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
    return 0;
}