/**
 * engines/c/main.c
 * -----------------
 * C 语言动力学模拟引擎。
 * 与 Python 版 (compute.py) 算法保持一致。
 *
 * 输入:  param.json      数值参数（Python 从 YAML 转换得来）
 *       <input_dir>/coord.txt
 *       <input_dir>/connection.txt
 *       <input_dir>/bond.txt
 * 输出: <output_dir>/trajectory.txt  (JSON 格式，与 Python 版兼容)
 *
 * 编译: cmake --build build --target dynamics_c
 * 用法: ./build/dynamics_c <input_dir> <output_dir> <param_json_path>
 */

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

/* ========================================================================
 * 配置参数（从 param.json 读取）
 * ======================================================================== */
typedef struct {
    double  box_a;           /* 盒子半边长 */
    int     NT;              /* 总步数 */
    double  DT;              /* 时间步长 */
    int     NSTEP;           /* 抽帧间隔 */
    int     warmup_steps;    /* 预热步数 */
    char    method[32];      /* 算法名称 */
    double  G[3];            /* 重力分量 */
    double  B[3];            /* 阻尼分量 */
    int     gravity_field;       /* 均匀重力场开关 */
    int     gravity_interaction; /* 原子间万有引力开关 */
    int     elastic_force;       /* 弹簧键力开关 */
    int     damping_force;       /* 阻尼开关 */
    double  gravity_strength;    /* 万有引力强度 */
    int     driving_force;       /* 驱动力开关 */
    int     save_trajectory;     /* 是否保存完整轨迹文件 */
    double  alpha[6];            /* 盒子透明度 */
    double  ball_radius;
    double  ball_color[3];
    double  box_color[3];
    int     use_marker;
    double  camera_distance, camera_elevation, camera_azimuth;
} SimParams;

/* ========================================================================
 * 原子数据
 * ======================================================================== */
typedef struct {
    int    n_atoms;
    int    *atom_ids;
    double *masses;
    double *radii;
    double *pos_0;           /* 初始位置 (n_atoms*3) */
    double *vel_0;           /* 初始速度 (n_atoms*3) */
    int    *fixed;           /* 固定约束 (n_atoms*3) */
} AtomData;

/* ========================================================================
 * 成键数据
 * ======================================================================== */
typedef struct {
    int    n_bonds;
    int    *pairs;           /* (n_bonds*2) */
    double *stiffness;
    double *rest_lengths;
} BondData;

/* 前向声明 */
static void *xmalloc(size_t sz);

/* ========================================================================
 * 驱动力数据
 * ======================================================================== */
typedef struct {
    int    n_drivers;
    int    *atom_idx;
    double *amp_x, *amp_y, *amp_z;
    double *freq_x, *freq_y, *freq_z;
    double *phi_x, *phi_y, *phi_z;   /* radians */
    int    *has_period;    /* 0=all, 1=limited cycles */
    double *period_cycles; /* number of cycles */
    double *eq_x, *eq_y, *eq_z;       /* 平衡位置（初始坐标） */
    double *freeze_x, *freeze_y, *freeze_z;
} DriverData;

/* 读取 driver.txt */
static DriverData read_driver(const char *input_dir, const AtomData *atoms) {
    DriverData d;
    memset(&d, 0, sizeof(d));

    char path[512];
    snprintf(path, sizeof(path), "%s/driver.txt", input_dir);
    FILE *f = fopen(path, "r");
    if (!f) return d;

    char line[1024];
    if (!fgets(line, sizeof(line), f)) { fclose(f); return d; }

    /* 第一遍：统计行数 */
    int n_lines = 0;
    while (fgets(line, sizeof(line), f)) {
        char trimmed[1024];
        int j = 0;
        for (int i = 0; line[i]; i++) {
            if (line[i] != ' ' && line[i] != '\t' && line[i] != '\n' && line[i] != '\r')
                trimmed[j++] = line[i];
        }
        trimmed[j] = '\0';
        if (strlen(trimmed) > 0 && trimmed[0] != '#') n_lines++;
    }

    if (n_lines == 0) { fclose(f); return d; }

    /* 分配内存 */
    d.n_drivers = n_lines;
    d.atom_idx     = (int*)xmalloc(n_lines * sizeof(int));
    d.amp_x = (double*)xmalloc(n_lines * sizeof(double));
    d.amp_y = (double*)xmalloc(n_lines * sizeof(double));
    d.amp_z = (double*)xmalloc(n_lines * sizeof(double));
    d.freq_x = (double*)xmalloc(n_lines * sizeof(double));
    d.freq_y = (double*)xmalloc(n_lines * sizeof(double));
    d.freq_z = (double*)xmalloc(n_lines * sizeof(double));
    d.phi_x = (double*)xmalloc(n_lines * sizeof(double));
    d.phi_y = (double*)xmalloc(n_lines * sizeof(double));
    d.phi_z = (double*)xmalloc(n_lines * sizeof(double));
    d.has_period = (int*)xmalloc(n_lines * sizeof(int));
    d.period_cycles = (double*)xmalloc(n_lines * sizeof(double));
    d.eq_x = (double*)xmalloc(n_lines * sizeof(double));
    d.eq_y = (double*)xmalloc(n_lines * sizeof(double));
    d.eq_z = (double*)xmalloc(n_lines * sizeof(double));
    d.freeze_x = (double*)xmalloc(n_lines * sizeof(double));
    d.freeze_y = (double*)xmalloc(n_lines * sizeof(double));
    d.freeze_z = (double*)xmalloc(n_lines * sizeof(double));

    /* 初始化 freeze/eq 数组 */
    for (int i = 0; i < n_lines; i++) {
        d.eq_x[i] = d.eq_y[i] = d.eq_z[i] = 0.0;
        d.freeze_x[i] = d.freeze_y[i] = d.freeze_z[i] = 0.0;
    }

    /* 第二遍：解析 */
    rewind(f);
    fgets(line, sizeof(line), f); /* 跳过表头 */

    int idx = 0;
    while (idx < n_lines && fgets(line, sizeof(line), f)) {
        char trimmed[1024];
        int j = 0;
        for (int i = 0; line[i]; i++) {
            if (line[i] != ' ' && line[i] != '\t' && line[i] != '\n' && line[i] != '\r')
                trimmed[j++] = line[i];
        }
        trimmed[j] = '\0';
        if (strlen(trimmed) == 0 || trimmed[0] == '#') continue;

        int atom_id;
        double amp_x, amp_y, amp_z;
        double freq_x, freq_y, freq_z;
        double phi_x, phi_y, phi_z;
        char period_str[256] = {0};

        int n_parsed = sscanf(line,
            "%d %lf %lf %lf %lf %lf %lf %lf %lf %lf %255s",
            &atom_id,
            &amp_x, &amp_y, &amp_z,
            &freq_x, &freq_y, &freq_z,
            &phi_x, &phi_y, &phi_z,
            period_str);

        if (n_parsed < 11) continue;

        /* 通过原子 ID 匹配内部索引（线性搜索）*/
        int ii = -1;
        for (int k = 0; k < atoms->n_atoms; k++) {
            if (atoms->atom_ids[k] == atom_id) { ii = k; break; }
        }
        if (ii < 0) continue;

        d.atom_idx[idx] = ii;
        d.eq_x[idx] = atoms->pos_0[ii*3+0];
        d.eq_y[idx] = atoms->pos_0[ii*3+1];
        d.eq_z[idx] = atoms->pos_0[ii*3+2];
        d.amp_x[idx] = amp_x;
        d.amp_y[idx] = amp_y;
        d.amp_z[idx] = amp_z;
        d.freq_x[idx] = freq_x;
        d.freq_y[idx] = freq_y;
        d.freq_z[idx] = freq_z;
        /* 角度 → 弧度 */
        d.phi_x[idx] = phi_x * M_PI / 180.0;
        d.phi_y[idx] = phi_y * M_PI / 180.0;
        d.phi_z[idx] = phi_z * M_PI / 180.0;

        if (strcmp(period_str, "all") == 0 || strcmp(period_str, "-1") == 0) {
            d.has_period[idx] = 0;
            d.period_cycles[idx] = -1.0;
        } else {
            d.has_period[idx] = 1;
            d.period_cycles[idx] = strtod(period_str, NULL);
        }
        idx++;
    }
    d.n_drivers = idx;

    fclose(f);
    return d;
}

/* ========================================================================
 * 轨迹缓冲区
 * ======================================================================== */
typedef struct {
    int    n_steps;
    int    n_atoms;
    double *x, *y, *z;
    double *vx, *vy, *vz;
} Trajectory;

/* ========================================================================
 * 辅助函数
 * ======================================================================== */

static void die(const char *msg) {
    fprintf(stderr, "[C-engine] 错误: %s\n", msg);
    exit(1);
}

static void *xmalloc(size_t sz) {
    void *p = malloc(sz);
    if (!p) die("内存分配失败");
    return p;
}

/* 从 JSON 中读取 double 值 */
static double json_read_double(const char *json, const char *key) {
    char search[256];
    snprintf(search, sizeof(search), "\"%s\"", key);
    const char *p = strstr(json, search);
    if (!p) return 0.0;
    p = strchr(p, ':');
    if (!p) return 0.0;
    p++;
    while (*p == ' ' || *p == '\t' || *p == '\n') p++;
    return strtod(p, NULL);
}

static int json_read_int(const char *json, const char *key) {
    return (int)json_read_double(json, key);
}

/* 从 JSON 中读取字符串值（写入 dst，最多 dst_sz 字节） */
static void json_read_string(const char *json, const char *key, char *dst, int dst_sz) {
    char search[256];
    snprintf(search, sizeof(search), "\"%s\"", key);
    const char *p = strstr(json, search);
    if (!p) { dst[0] = '\0'; return; }
    p = strchr(p, ':');
    if (!p) { dst[0] = '\0'; return; }
    p++;
    while (*p == ' ' || *p == '\t' || *p == '\n') p++;
    if (*p != '"') { dst[0] = '\0'; return; }
    p++;
    int i = 0;
    while (*p && *p != '"' && i < dst_sz - 1) { dst[i++] = *p++; }
    dst[i] = '\0';
}

/* 读取 JSON 数组 (如 "G": [0, 0, -9.8]) 到 double[3] */
static void json_read_double3(const char *json, const char *key, double out[3]) {
    char search[256];
    snprintf(search, sizeof(search), "\"%s\"", key);
    const char *p = strstr(json, search);
    if (!p) { out[0]=out[1]=out[2]=0; return; }
    p = strchr(p, '[');
    if (!p) { out[0]=out[1]=out[2]=0; return; }
    p++;
    for (int i = 0; i < 3; i++) {
        while (*p == ' ' || *p == '\t' || *p == '\n' || *p == ',') p++;
        out[i] = strtod(p, (char**)&p);
    }
}

static void json_read_double6(const char *json, const char *key, double out[6]) {
    char search[256];
    snprintf(search, sizeof(search), "\"%s\"", key);
    const char *p = strstr(json, search);
    if (!p) { for (int i=0;i<6;i++) out[i]=0; return; }
    p = strchr(p, '[');
    if (!p) { for (int i=0;i<6;i++) out[i]=0; return; }
    p++;
    for (int i = 0; i < 6; i++) {
        while (*p == ' ' || *p == '\t' || *p == '\n' || *p == ',' || *p == ']') p++;
        out[i] = strtod(p, (char**)&p);
    }
}

/* 读取 param.json */
static int    g_gravity_field = 1;
static int    g_gravity_interaction = 0;
static int    g_elastic_force = 1;
static int    g_damping_force = 0;
static double g_gravity_strength = 1.0;

static SimParams read_params(const char *path) {
    FILE *f = fopen(path, "rb");
    if (!f) die("无法打开 param.json");
    fseek(f, 0, SEEK_END);
    long sz = ftell(f);
    fseek(f, 0, SEEK_SET);
    char *buf = (char*)xmalloc(sz + 1);
    fread(buf, 1, sz, f);
    buf[sz] = '\0';
    fclose(f);

    SimParams p;
    p.box_a       = json_read_double(buf, "box_a");
    p.NT          = json_read_int(buf, "NT");
    p.DT          = json_read_double(buf, "DT");
    p.NSTEP       = json_read_int(buf, "NSTEP");
    p.warmup_steps = json_read_int(buf, "warmup_steps");
    strcpy(p.method, "leapfrog"); /* 默认 */
    json_read_string(buf, "method", p.method, sizeof(p.method));
    json_read_double3(buf, "G", p.G);
    json_read_double3(buf, "B", p.B);
    p.gravity_field       = json_read_int(buf, "gravity_field");
    p.gravity_interaction = json_read_int(buf, "gravity_interaction");
    p.elastic_force       = json_read_int(buf, "elastic_force");
    p.damping_force       = json_read_int(buf, "damping_force");
    p.gravity_strength    = json_read_double(buf, "gravity_strength");
    p.driving_force       = json_read_int(buf, "driving_force");
    p.save_trajectory     = json_read_int(buf, "save_trajectory");
    /* 渲染参数 */
    json_read_double6(buf, "alpha", p.alpha);
    p.ball_radius         = json_read_double(buf, "ball_radius");
    json_read_double3(buf, "ball_color", p.ball_color);
    json_read_double3(buf, "box_color", p.box_color);
    p.use_marker          = json_read_int(buf, "use_marker");
    p.camera_distance     = json_read_double(buf, "camera_distance");
    p.camera_elevation    = json_read_double(buf, "camera_elevation");
    p.camera_azimuth      = json_read_double(buf, "camera_azimuth");
    g_gravity_field       = p.gravity_field;
    g_gravity_interaction = p.gravity_interaction;
    g_elastic_force       = p.elastic_force;
    g_damping_force       = p.damping_force;
    g_gravity_strength    = p.gravity_strength;

    free(buf);
    return p;
}

/* 读取 coord.txt */
static AtomData read_coord(const char *input_dir) {
    char path[512];
    snprintf(path, sizeof(path), "%s/coord.txt", input_dir);
    FILE *f = fopen(path, "r");
    if (!f) die("无法打开 coord.txt");

    /* 跳过第一行表头 */
    char line[1024];
    if (!fgets(line, sizeof(line), f)) die("coord.txt 为空");

    int capacity = 16;
    AtomData a;
    a.n_atoms = 0;
    a.atom_ids  = (int*)xmalloc(capacity * sizeof(int));
    a.masses    = (double*)xmalloc(capacity * sizeof(double));
    a.radii     = (double*)xmalloc(capacity * sizeof(double));
    a.pos_0     = (double*)xmalloc(capacity * 3 * sizeof(double));
    a.vel_0     = (double*)xmalloc(capacity * 3 * sizeof(double));
    a.fixed     = (int*)xmalloc(capacity * 3 * sizeof(int));

    while (fgets(line, sizeof(line), f)) {
        if (a.n_atoms >= capacity) {
            capacity *= 2;
            a.atom_ids = realloc(a.atom_ids, capacity * sizeof(int));
            a.masses   = realloc(a.masses,   capacity * sizeof(double));
            a.radii    = realloc(a.radii,    capacity * sizeof(double));
            a.pos_0    = realloc(a.pos_0,    capacity * 3 * sizeof(double));
            a.vel_0    = realloc(a.vel_0,    capacity * 3 * sizeof(double));
            a.fixed    = realloc(a.fixed,    capacity * 3 * sizeof(int));
        }
        int id, fx, fy, fz;
        double mass, rad, px, py, pz, vx, vy, vz;
        int n_parsed = sscanf(line, "%d %lf %lf %lf %lf %lf %lf %lf %lf %d %d %d",
                   &id, &mass, &rad, &px, &py, &pz, &vx, &vy, &vz, &fx, &fy, &fz);
        if (n_parsed == 9) {
            fx = fy = fz = 0;
        } else if (n_parsed != 12) {
            continue;
        }
        int i = a.n_atoms;
        a.atom_ids[i] = id;
        a.masses[i]   = mass;
        a.radii[i]    = rad;
        a.pos_0[i*3+0] = px; a.pos_0[i*3+1] = py; a.pos_0[i*3+2] = pz;
        a.vel_0[i*3+0] = vx; a.vel_0[i*3+1] = vy; a.vel_0[i*3+2] = vz;
        a.fixed[i*3+0] = fx; a.fixed[i*3+1] = fy; a.fixed[i*3+2] = fz;
        a.n_atoms++;
    }
    fclose(f);

    if (a.n_atoms <= 0) die("coord.txt 原子数无效");
    return a;
}

/* 读取 connection.txt */
static BondData read_bonds(const char *input_dir, const AtomData *atoms) {
    char path[512];
    BondData b;
    b.n_bonds = 0;
    b.pairs = NULL;
    b.stiffness = NULL;
    b.rest_lengths = NULL;

    snprintf(path, sizeof(path), "%s/connection.txt", input_dir);
    FILE *f = fopen(path, "r");
    if (!f) return b;

    char line[256];
    if (!fgets(line, sizeof(line), f)) { fclose(f); return b; }

    int n_lines = 0, tmp_a, tmp_b;
    char bond_name[256];
    while (fscanf(f, "%d %d %s", &tmp_a, &tmp_b, bond_name) == 3) n_lines++;
    rewind(f);
    fgets(line, sizeof(line), f);  // 再次跳过表头

    if (n_lines == 0) { fclose(f); return b; }

    b.n_bonds = n_lines;
    b.pairs = (int*)xmalloc(n_lines * 2 * sizeof(int));
    b.stiffness = (double*)xmalloc(n_lines * sizeof(double));
    b.rest_lengths = (double*)xmalloc(n_lines * sizeof(double));

    char bond_path[512];
    snprintf(bond_path, sizeof(bond_path), "%s/bond.txt", input_dir);
    FILE *fb = fopen(bond_path, "r");

    for (int i = 0; i < n_lines; i++) {
        fscanf(f, "%d %d %s", &tmp_a, &tmp_b, bond_name);
        b.pairs[i*2+0] = tmp_a - 1;
        b.pairs[i*2+1] = tmp_b - 1;
        b.stiffness[i] = 1.0;
        b.rest_lengths[i] = 2.0;
        if (fb) {
            char name[256], header[256];
            double k, r0;
            rewind(fb);
            fgets(header, sizeof(header), fb);  // 跳过表头行
            while (fscanf(fb, "%s %lf %lf", name, &k, &r0) == 3) {
                if (strcmp(name, bond_name) == 0) {
                    b.stiffness[i] = k;
                    b.rest_lengths[i] = r0;
                    break;
                }
            }
        }
    }
    fclose(f);
    if (fb) fclose(fb);
    return b;
}

/* ========================================================================
 * 物理核心（与 Python compute.py 对应）
 * ======================================================================== */

/* 加速度计算（各力独立开关控制） */
static void compute_acceleration(
    int n, const double *x, const double *y, const double *z,
    const double *vx, const double *vy, const double *vz,
    const double *m, const double G[3], const double B[3],
    const BondData *bonds,
    double *ax, double *ay, double *az)
{
    /* 先清零 */
    for (int i = 0; i < n; i++) {
        ax[i] = 0.0; ay[i] = 0.0; az[i] = 0.0;
    }

    /* 均匀重力场 */
    if (g_gravity_field) {
        for (int i = 0; i < n; i++) {
            ax[i] += G[0];
            ay[i] += G[1];
            az[i] += G[2];
        }
    }

    /* 阻尼 */
    if (g_damping_force) {
        for (int i = 0; i < n; i++) {
            ax[i] -= B[0] * vx[i] / m[i];
            ay[i] -= B[1] * vy[i] / m[i];
            az[i] -= B[2] * vz[i] / m[i];
        }
    }

    /* 弹簧键力 */
    if (g_elastic_force) {
        for (int b = 0; b < bonds->n_bonds; b++) {
            int i = bonds->pairs[b*2+0];
            int j = bonds->pairs[b*2+1];
            double dx = x[j] - x[i];
            double dy = y[j] - y[i];
            double dz = z[j] - z[i];
            double dist = sqrt(dx*dx + dy*dy + dz*dz);
            if (dist < 1e-12) continue;
            double stretch = dist - bonds->rest_lengths[b];
            double fmag = bonds->stiffness[b] * stretch;
            double ux = dx / dist, uy = dy / dist, uz = dz / dist;
            double fx = fmag * ux, fy = fmag * uy, fz = fmag * uz;
            ax[i] += fx / m[i];  ay[i] += fy / m[i];  az[i] += fz / m[i];
            ax[j] -= fx / m[j];  ay[j] -= fy / m[j];  az[j] -= fz / m[j];
        }
    }

    /* 万有引力（所有原子对之间） */
    if (g_gravity_interaction) {
        for (int i = 0; i < n; i++) {
            for (int j = i + 1; j < n; j++) {
                double dx = x[j] - x[i];
                double dy = y[j] - y[i];
                double dz = z[j] - z[i];
                double r2 = dx*dx + dy*dy + dz*dz;
                if (r2 <= 1e-12) continue;
                double r = sqrt(r2);
                double f_mag = g_gravity_strength * m[i] * m[j] / r2;
                double fx_g = f_mag * dx / r;
                double fy_g = f_mag * dy / r;
                double fz_g = f_mag * dz / r;
                ax[i] += fx_g / m[i]; ay[i] += fy_g / m[i]; az[i] += fz_g / m[i];
                ax[j] -= fx_g / m[j]; ay[j] -= fy_g / m[j]; az[j] -= fz_g / m[j];
            }
        }
    }
}

/* 保守力加速度（不含阻尼），供真蛙跳法专用。
   通过传入零速度调用 compute_acceleration，阻尼项 -B*v/m 自动为零。 */
static void compute_accel_conservative(
    int n, const double *x, const double *y, const double *z,
    const double *m, const double G[3],
    const BondData *bonds,
    double *ax, double *ay, double *az)
{
    double *v0 = (double*)alloca(n * sizeof(double));
    for (int i = 0; i < n; i++) v0[i] = 0.0;
    double Bzero[3] = {0.0, 0.0, 0.0};
    compute_acceleration(n, x, y, z, v0, v0, v0, m, G, Bzero, bonds, ax, ay, az);
}

/* 边界条件：clamp 位置 + 速度反转 ——与 Python Limit_in_box 一致 */
static void limit_in_box(double *pos, double *vel, double lo, double hi) {
    if (*pos > hi)  { *pos = hi;  *vel = -*vel; }
    if (*pos < lo)  { *pos = lo;  *vel = -*vel; }
}

/* 周期边界回绕：超出边界的位置绕到对侧（与 Python wrap_position 一致）*/
static void wrap_position(double *pos, double lo, double hi) {
    if (*pos > hi)  *pos = lo;
    if (*pos < lo)  *pos = hi;
}

/* ── 显式欧拉法 ──────────── */
static void explicit_euler_step(
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    const double *m, const double G[3], const double B[3],
    const BondData *bonds, const int *fixed, double dt)
{
    double *ax = (double*)alloca(n * sizeof(double));
    double *ay = (double*)alloca(n * sizeof(double));
    double *az = (double*)alloca(n * sizeof(double));
    compute_acceleration(n, x, y, z, vx, vy, vz, m, G, B, bonds, ax, ay, az);

    for (int i = 0; i < n; i++) {
        if (fixed[i*3+0] && 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;
    }
}

/* ── 隐式欧拉法 ──────────── */
static void implicit_euler_step(
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    const double *m, const double G[3], const double B[3],
    const BondData *bonds, const int *fixed, double dt)
{
    double *vxn = (double*)alloca(n * sizeof(double));
    double *vyn = (double*)alloca(n * sizeof(double));
    double *vzn = (double*)alloca(n * sizeof(double));

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

    double *ax = (double*)alloca(n * sizeof(double));
    double *ay = (double*)alloca(n * sizeof(double));
    double *az = (double*)alloca(n * sizeof(double));
    compute_acceleration(n, x, y, z, vxn, vyn, vzn, m, G, B, bonds, ax, ay, az);

    for (int i = 0; i < n; i++) {
        if (fixed[i*3+0] && 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;
    }
}

/* ── 中点法 ──────────── */
static void midpoint_step(
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    const double *m, const double G[3], const double B[3],
    const BondData *bonds, const int *fixed, double dt)
{
    double *ax = (double*)alloca(n * sizeof(double));
    double *ay = (double*)alloca(n * sizeof(double));
    double *az = (double*)alloca(n * sizeof(double));
    compute_acceleration(n, x, y, z, vx, vy, vz, m, G, B, bonds, ax, ay, az);

    double *xm = (double*)alloca(n * sizeof(double));
    double *ym = (double*)alloca(n * sizeof(double));
    double *zm = (double*)alloca(n * sizeof(double));
    double *vxm = (double*)alloca(n * sizeof(double));
    double *vym = (double*)alloca(n * sizeof(double));
    double *vzm = (double*)alloca(n * sizeof(double));

    for (int i = 0; i < n; i++) {
        if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) {
            xm[i]=ym[i]=zm[i]=vxm[i]=vym[i]=vzm[i]=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;
    }

    compute_acceleration(n, xm, ym, zm, vxm, vym, vzm, m, G, B, bonds, ax, ay, az);

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

/* ── 蛙跳法（Velocity-Verlet）── */
/* 真蛙跳一步：x(t), v(t-dt/2) → x(t+dt), v(t+dt/2)
 *
 * 无阻尼：纯保守蛙跳，每步 1 次力计算，辛积分器。
 *   v(t+dt/2) = v(t-dt/2) + a_c(t)·dt
 *
 * 有阻尼：半隐式处理，仍 1 次力计算，对任意阻尼无条件稳定。
 *   利用 v(t) ≈ [v(t-dt/2) + v(t+dt/2)] / 2 解析求解：
 *   v(t+dt/2) = [v(t-dt/2)·(1-α) + a_c(t)·dt] / (1+α)，α = 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 double G[3], const double B[3],
    const BondData *bonds, const int *fixed, double dt)
{
    double *ax = (double*)alloca(n * sizeof(double));
    double *ay = (double*)alloca(n * sizeof(double));
    double *az = (double*)alloca(n * sizeof(double));

    /* 1 次保守力计算（不含阻尼） */
    compute_accel_conservative(n, x, y, z, m, G, bonds, ax, ay, az);

    int has_damping = g_damping_force && (B[0] != 0.0 || B[1] != 0.0 || B[2] != 0.0);

    for (int i = 0; i < n; i++) {
        if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue;
        if (has_damping) {
            double alphax = B[0] * dt / (2.0 * m[i]);
            double alphay = B[1] * dt / (2.0 * m[i]);
            double alphaz = B[2] * dt / (2.0 * m[i]);
            vx[i] = (vx[i] * (1.0 - alphax) + ax[i] * dt) / (1.0 + alphax);
            vy[i] = (vy[i] * (1.0 - alphay) + ay[i] * dt) / (1.0 + alphay);
            vz[i] = (vz[i] * (1.0 - alphaz) + az[i] * dt) / (1.0 + alphaz);
        } 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;
    }
}

/* ── 驱动力（与 Python apply_driving_force 一致）──────────────── */
static void apply_driving_force(
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    double t, int step, double dt,
    const DriverData *drivers)
{
    if (!drivers || drivers->n_drivers == 0) return;
    for (int d = 0; d < drivers->n_drivers; d++) {
        int idx = drivers->atom_idx[d];
        /* 检查周期限制 */
        if (drivers->has_period[d]) {
            double max_freq = fmax(fabs(drivers->freq_x[d]),
                                   fmax(fabs(drivers->freq_y[d]), fabs(drivers->freq_z[d])));
            int period_steps = 0;
            if (max_freq > 1e-12) {
                period_steps = (int)(drivers->period_cycles[d] / max_freq / dt);
            }
            if (step > period_steps) {
                /* 冻结 */
                if (drivers->freeze_x) {
                    x[idx] = drivers->freeze_x[d];
                    y[idx] = drivers->freeze_y[d];
                    z[idx] = drivers->freeze_z[d];
                }
                vx[idx] = vy[idx] = vz[idx] = 0.0;
                continue;
            }
        }

        double px = drivers->eq_x[d] + drivers->amp_x[d] * cos(2*M_PI*drivers->freq_x[d]*t + drivers->phi_x[d]);
        double py = drivers->eq_y[d] + drivers->amp_y[d] * cos(2*M_PI*drivers->freq_y[d]*t + drivers->phi_y[d]);
        double pz = drivers->eq_z[d] + drivers->amp_z[d] * cos(2*M_PI*drivers->freq_z[d]*t + drivers->phi_z[d]);
        double vpx = -drivers->amp_x[d]*2*M_PI*drivers->freq_x[d]*sin(2*M_PI*drivers->freq_x[d]*t + drivers->phi_x[d]);
        double vpy = -drivers->amp_y[d]*2*M_PI*drivers->freq_y[d]*sin(2*M_PI*drivers->freq_y[d]*t + drivers->phi_y[d]);
        double vpz = -drivers->amp_z[d]*2*M_PI*drivers->freq_z[d]*sin(2*M_PI*drivers->freq_z[d]*t + drivers->phi_z[d]);

        x[idx] = px; y[idx] = py; z[idx] = pz;
        vx[idx] = vpx; vy[idx] = vpy; vz[idx] = vpz;

        /* 记录冻结位置（周期结束时） */
        if (drivers->has_period[d]) {
            double max_freq = fmax(fabs(drivers->freq_x[d]),
                                   fmax(fabs(drivers->freq_y[d]), fabs(drivers->freq_z[d])));
            int period_steps = 0;
            if (max_freq > 1e-12) {
                period_steps = (int)(drivers->period_cycles[d] / max_freq / dt);
            }
            if (step == period_steps) {
                drivers->freeze_x[d] = px;
                drivers->freeze_y[d] = py;
                drivers->freeze_z[d] = pz;
            }
        }
    }
}

/* ── 分发器：调用对应积分方法 + 边界条件 + 自由度约束（与 Python 一致）── */
static void apply_step(
    const char *method,
    int n, double *x, double *y, double *z,
    double *vx, double *vy, double *vz,
    const double *m, const double G[3], const double B[3],
    const BondData *bonds, const int *fixed,
    const double *pos_0,
    double box_a, double dt)
{
    if (strcmp(method, "explicit_euler") == 0) {
        explicit_euler_step(n, x, y, z, vx, vy, vz, m, G, B, bonds, fixed, dt);
    } else if (strcmp(method, "implicit_euler") == 0) {
        implicit_euler_step(n, x, y, z, vx, vy, vz, m, G, B, bonds, fixed, dt);
    } else if (strcmp(method, "midpoint") == 0) {
        midpoint_step(n, x, y, z, vx, vy, vz, m, G, B, bonds, fixed, dt);
    } else if (strcmp(method, "leapfrog") == 0) {
        leapfrog_step(n, x, y, z, vx, vy, vz, m, G, B, bonds, fixed, dt);
    } else {
        fprintf(stderr, "[C-engine] 未知算法: %s\n", method);
        exit(1);
    }

    /* 边界条件（与 Python Limit_in_box 一致） */
    for (int i = 0; i < n; i++) {
        if (fixed[i*3+0] && fixed[i*3+1] && fixed[i*3+2]) continue;
        limit_in_box(&x[i],  &vx[i], -box_a, box_a);
        limit_in_box(&y[i],  &vy[i], -box_a, box_a);
        limit_in_box(&z[i],  &vz[i], -box_a, box_a);
    }

    /* 周期边界回绕（与 Python wrap_position 一致）*/
    for (int i = 0; i < n; i++) {
        wrap_position(&x[i], -box_a, box_a);
        wrap_position(&y[i], -box_a, box_a);
        wrap_position(&z[i], -box_a, box_a);
    }

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

// ========================================================================
// JSON 输出
// ========================================================================

static void write_trajectory_json(const char *path, const Trajectory *traj,
                                  const SimParams *params, const AtomData *atoms,
                                  const BondData *bonds)
{
    FILE *f = fopen(path, "w");
    if (!f) die("无法写入 trajectory.txt");

    fprintf(f, "{\n");

    const char *names[] = {"traj_x","traj_y","traj_z","traj_vx","traj_vy","traj_vz"};
    double *arrs[] = {traj->x, traj->y, traj->z, traj->vx, traj->vy, traj->vz};

    printf("[C-engine] 正在写入轨迹数据…\n");
    fflush(stdout);
    for (int a = 0; a < 6; a++) {
        fprintf(f, "  \"%s\": [\n", names[a]);
        for (int t = 0; t < traj->n_steps; t++) {
            fprintf(f, "    [");
            for (int i = 0; i < traj->n_atoms; i++) {
                fprintf(f, "%.8g", arrs[a][t * traj->n_atoms + i]);
                if (i < traj->n_atoms - 1) fputc(',', f);
            }
            fprintf(f, "]");
            if (t < traj->n_steps - 1) fputc(',', f);
            fputc('\n', f);
        }
        fprintf(f, "  ]");
        fputc(',', f);
        fputc('\n', f);
    }

    /* 标量参数 */
    fprintf(f, "  \"NT\": %d,\n", params->NT);
    fprintf(f, "  \"DT\": %.8g,\n", params->DT);
    fprintf(f, "  \"NSTEP\": %d,\n", params->NSTEP);
    fprintf(f, "  \"method\": \"%s\",\n", params->method);
    fprintf(f, "  \"warmup_steps\": %d,\n", params->warmup_steps);
    fprintf(f, "  \"G\": [%.8g, %.8g, %.8g],\n", params->G[0], params->G[1], params->G[2]);
    fprintf(f, "  \"B\": [%.8g, %.8g, %.8g],\n", params->B[0], params->B[1], params->B[2]);

    fprintf(f, "  \"atom_ids\": [");
    for (int i = 0; i < atoms->n_atoms; i++) {
        if (i > 0) fputc(',', f);
        fprintf(f, "%d", atoms->atom_ids[i]);
    }
    fprintf(f, "],\n");

    fprintf(f, "  \"atom_masses\": [");
    for (int i = 0; i < atoms->n_atoms; i++) {
        if (i > 0) fputc(',', f);
        fprintf(f, "%.8g", atoms->masses[i]);
    }
    fprintf(f, "],\n");

    fprintf(f, "  \"bond_pairs\": [");
    for (int b = 0; b < bonds->n_bonds; b++) {
        if (b > 0) fputc(',', f);
        fprintf(f, "[%d, %d]", bonds->pairs[b*2], bonds->pairs[b*2+1]);
    }
    fprintf(f, "],\n");

    fprintf(f, "  \"bond_stiffness\": [");
    for (int b = 0; b < bonds->n_bonds; b++) {
        if (b > 0) fputc(',', f);
        fprintf(f, "%.8g", bonds->stiffness[b]);
    }
    fprintf(f, "],\n");

    fprintf(f, "  \"bond_rest_lengths\": [");
    for (int b = 0; b < bonds->n_bonds; b++) {
        if (b > 0) fputc(',', f);
        fprintf(f, "%.8g", bonds->rest_lengths[b]);
    }
    fprintf(f, "],\n");
    fprintf(f, "  \"driving_force\": %d\n", params->driving_force);

    fprintf(f, "}\n");
    fclose(f);
}

static void write_display_txt(const char *path, const Trajectory *traj,
                               const SimParams *params, const AtomData *atoms)
{
    FILE *f = fopen(path, "w");
    if (!f) die("无法写入 display.txt");

    int n_frames    = traj->n_steps;   /* 实际采样帧数，用于下面的帧循环 */
    int n_particles = traj->n_atoms;
    int dynamic_steps = params->NT - params->warmup_steps;
    double T_total = dynamic_steps * params->DT;

    /* number of frames 写总积分步数（与 draw.py NT 对应），不是采样帧数 */
    fprintf(f, "number of frames: %d\n", dynamic_steps);
    fprintf(f, "number of particles: %d\n", n_particles);
    fprintf(f, "DT: %.16g\n", params->DT);
    fprintf(f, "NSTEP: %d\n", params->NSTEP);
    fprintf(f, "method: %s\n", params->method);
    fprintf(f, "warmup_steps: %d\n", params->warmup_steps);
    fprintf(f, "dynamic_steps: %d\n", dynamic_steps);
    fprintf(f, "T_total: %.16g\n", T_total);
    fprintf(f, "box_a: %.16g\n", params->box_a);
    fprintf(f, "alpha: %.16g,%.16g,%.16g,%.16g,%.16g,%.16g\n",
            params->alpha[0], params->alpha[1], params->alpha[2],
            params->alpha[3], params->alpha[4], params->alpha[5]);
    fprintf(f, "ball_radius: %.16g\n", params->ball_radius);
    fprintf(f, "ball_color_r: %.16g\n", params->ball_color[0]);
    fprintf(f, "ball_color_g: %.16g\n", params->ball_color[1]);
    fprintf(f, "ball_color_b: %.16g\n", params->ball_color[2]);
    fprintf(f, "box_color_r: %.16g\n", params->box_color[0]);
    fprintf(f, "box_color_g: %.16g\n", params->box_color[1]);
    fprintf(f, "box_color_b: %.16g\n", params->box_color[2]);
    fprintf(f, "use_marker: %d\n", params->use_marker);
    fprintf(f, "camera_distance: %.16g\n", params->camera_distance);
    fprintf(f, "camera_elevation: %.16g\n", params->camera_elevation);
    fprintf(f, "camera_azimuth: %.16g\n", params->camera_azimuth);
    fprintf(f, "\n");

    if (params->driving_force) {
        fprintf(f, "driving_force: 1\n");
    } else {
        fprintf(f, "driving_force: 0\n");
    }
    fprintf(f, "\n");

    for (int t = 0; t < n_frames; t++) {
        fprintf(f, "frame: %3d\n", t + 1);
        fprintf(f, "n   x            y            z           vx         vy         vz\n");
        for (int i = 0; i < n_particles; i++) {
            int idx = t * n_particles + i;
            fprintf(f, "%4d %12.6f %12.6f %12.6f %10.6f %10.6f %10.6f\n",
                    atoms->atom_ids[i],
                    traj->x[idx], traj->y[idx], traj->z[idx],
                    traj->vx[idx], traj->vy[idx], traj->vz[idx]);
        }
    }

    fclose(f);
}

// ========================================================================
// 主函数
// ========================================================================

int main(int argc, char **argv) {
    if (argc < 4) {
        fprintf(stderr, "用法: %s <input_dir> <output_dir> <param_json>\n", argv[0]);
        return 1;
    }
    const char *input_dir  = argv[1];
    const char *output_dir = argv[2];
    const char *param_path = argv[3];

    clock_t t0 = clock();

    SimParams params = read_params(param_path);
    AtomData  atoms  = read_coord(input_dir);
    BondData  bonds  = read_bonds(input_dir, &atoms);

    DriverData drivers;
    drivers.n_drivers = 0;
    if (params.driving_force) {
        drivers = read_driver(input_dir, &atoms);
    }

    printf("[C-engine] 原子数=%d, 键数=%d, 驱动=%d, NT=%d, DT=%.6g, method=%s\n",
           atoms.n_atoms, bonds.n_bonds, drivers.n_drivers, params.NT, params.DT, params.method);

    int n = atoms.n_atoms;
    double *x  = (double*)xmalloc(n * sizeof(double));
    double *y  = (double*)xmalloc(n * sizeof(double));
    double *z  = (double*)xmalloc(n * sizeof(double));
    double *vx = (double*)xmalloc(n * sizeof(double));
    double *vy = (double*)xmalloc(n * sizeof(double));
    double *vz = (double*)xmalloc(n * sizeof(double));
    for (int i = 0; i < n; i++) {
        x[i]  = atoms.pos_0[i*3+0];
        y[i]  = atoms.pos_0[i*3+1];
        z[i]  = atoms.pos_0[i*3+2];
        vx[i] = atoms.vel_0[i*3+0];
        vy[i] = atoms.vel_0[i*3+1];
        vz[i] = atoms.vel_0[i*3+2];
    }

    /* 分配轨迹缓冲区 */
    int record_steps = params.NT - params.warmup_steps;
    Trajectory traj;
    traj.n_atoms = n;
    if (params.save_trajectory) {
        traj.n_steps = record_steps;
        traj.x  = (double*)xmalloc(record_steps * n * sizeof(double) * 6);
        traj.y  = traj.x  + record_steps * n;
        traj.z  = traj.y  + record_steps * n;
        traj.vx = traj.z  + record_steps * n;
        traj.vy = traj.vx + record_steps * n;
        traj.vz = traj.vy + record_steps * n;
    } else {
        int sampled_steps = (record_steps + params.NSTEP - 1) / params.NSTEP;
        if (sampled_steps < 1) sampled_steps = 1;
        traj.n_steps = sampled_steps;
        traj.x  = (double*)xmalloc(sampled_steps * n * sizeof(double) * 6);
        traj.y  = traj.x  + sampled_steps * n;
        traj.z  = traj.y  + sampled_steps * n;
        traj.vx = traj.z  + sampled_steps * n;
        traj.vy = traj.vx + sampled_steps * n;
        traj.vz = traj.vy + sampled_steps * n;
    }

    /* 真蛙跳初始化：v(0) 反推 v(-dt/2) = v(0) - 0.5*a_c(0)*dt */
    if (strcmp(params.method, "leapfrog") == 0) {
        double *ax0 = (double*)alloca(n * sizeof(double));
        double *ay0 = (double*)alloca(n * sizeof(double));
        double *az0 = (double*)alloca(n * sizeof(double));
        compute_accel_conservative(n, x, y, z, atoms.masses, params.G, &bonds, ax0, ay0, az0);
        for (int i = 0; i < n; i++) {
            if (atoms.fixed[i*3+0] && atoms.fixed[i*3+1] && atoms.fixed[i*3+2]) continue;
            vx[i] -= 0.5 * ax0[i] * params.DT;
            vy[i] -= 0.5 * ay0[i] * params.DT;
            vz[i] -= 0.5 * az0[i] * params.DT;
        }
    }

    /* 预热 */
    /* 初始时刻 t=0 驱动力（与 Python run_simulation 一致）*/
    if (params.driving_force) apply_driving_force(n, x, y, z, vx, vy, vz, 0.0, 0, params.DT, &drivers);

    for (int s = 0; s < params.warmup_steps; s++) {
        double tw = (s + 1) * params.DT;
        if (params.driving_force) apply_driving_force(n, x, y, z, vx, vy, vz, tw, s, params.DT, &drivers);
        apply_step(params.method, n, x, y, z, vx, vy, vz,
                   atoms.masses, params.G, params.B, &bonds, atoms.fixed,
                   atoms.pos_0,
                   params.box_a, params.DT);
    }

    /* 记录 */
    int _prog_interval = record_steps / 100;
    if (_prog_interval < 1) _prog_interval = 1;
    int sample_idx = 0;
    for (int s = 0; s < record_steps; s++) {
        if (s % _prog_interval == 0 && s > 0) {
            printf("[C-engine] progress: %d/%d\n", s, record_steps);
            fflush(stdout);
        }
        double t = (s + params.warmup_steps) * params.DT;
        if (params.driving_force) apply_driving_force(n, x, y, z, vx, vy, vz, t, s, params.DT, &drivers);
        int do_record = params.save_trajectory || (s % params.NSTEP == 0);
        if (do_record) {
            int idx = params.save_trajectory ? s : sample_idx;
            for (int i = 0; i < n; i++) {
                traj.x[ idx * n + i]  = x[i];
                traj.y[ idx * n + i]  = y[i];
                traj.z[ idx * n + i]  = z[i];
                traj.vx[idx * n + i]  = vx[i];
                traj.vy[idx * n + i]  = vy[i];
                traj.vz[idx * n + i]  = vz[i];
            }
            if (!params.save_trajectory) sample_idx++;
        }
        apply_step(params.method, n, x, y, z, vx, vy, vz,
                   atoms.masses, params.G, params.B, &bonds, atoms.fixed,
                   atoms.pos_0,
                   params.box_a, params.DT);
    }

    char out_path[512];
    if (params.save_trajectory) {
        snprintf(out_path, sizeof(out_path), "%s/trajectory.txt", output_dir);
        write_trajectory_json(out_path, &traj, &params, &atoms, &bonds);
    } else {
        snprintf(out_path, sizeof(out_path), "%s/display.txt", output_dir);
        write_display_txt(out_path, &traj, &params, &atoms);
    }

    clock_t t1 = clock();
    double elapsed = (double)(t1 - t0) / CLOCKS_PER_SEC;
    printf("[C-engine] 计算完成: %d 步, %.3f s\n", record_steps, elapsed);

    /* 清理 */
    free(traj.x);
    free(x); free(y); free(z);
    free(vx); free(vy); free(vz);
    free(atoms.atom_ids);
    free(atoms.masses); free(atoms.radii);
    free(atoms.pos_0); free(atoms.vel_0);
    free(atoms.fixed);
    if (bonds.pairs) { free(bonds.pairs); free(bonds.stiffness); free(bonds.rest_lengths); }

    return 0;
}