/**
 * engines/c/main.c
 * -----------------
 * C 语言动力学模拟引擎。
 * 输入:  param.json      数值参数（Python 从 YAML 转换得来）
 *       <input_dir>/coord.txt
 *       <input_dir>/connection.txt
 *       <input_dir>/bond.txt
 * 输出: <output_dir>/trajectory.txt  (JSON 格式，与 Python 版兼容)
 *
 * 编译: make
 * 用法: ./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;    /* 预热步数 */
    double  G[3];            /* 重力分量 */
    double  B[3];            /* 阻尼分量 */
} 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;

/* ========================================================================
 * 轨迹缓冲区
 * ======================================================================== */
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 值。简单实现：搜索 key 并解析 */
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 数组 (如 "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);
    }
}

/* 读取 param.json */
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");
    json_read_double3(buf, "G", p.G);
    json_read_double3(buf, "B", p.B);

    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);

    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));

    /* 读取 bond.txt 获取劲度系数和平衡长度 */
    char bond_path[512];
    snprintf(bond_path, sizeof(bond_path), "%s/bond.txt", input_dir);
    FILE *fb = fopen(bond_path, "r");
    /* bond.txt 格式: name stiffness rest_length */

    for (int i = 0; i < n_lines; i++) {
        fscanf(f, "%d %d %s", &tmp_a, &tmp_b, bond_name);
        /* 转换原子编号为 0-based 索引 */
        b.pairs[i*2+0] = tmp_a - 1;
        b.pairs[i*2+1] = tmp_b - 1;
        /* 查找 bond.txt 中对应的参数 */
        b.stiffness[i] = 1.0;
        b.rest_lengths[i] = 2.0;  /* 默认值 */
        if (fb) {
            char name[256];
            double k, r0;
            rewind(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;
}

/* ========================================================================
 * 物理核心
 * ======================================================================== */

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++) {
        double inv_m = 1.0 / m[i];
        ax[i] = (m[i] * G[0] - B[0] * vx[i]) * inv_m;
        ay[i] = (m[i] * G[1] - B[1] * vy[i]) * inv_m;
        az[i] = (m[i] * G[2] - B[2] * vz[i]) * inv_m;
    }

    /* 弹簧力 */
    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];
    }
}

/* clamp */
static double clamp(double v, double lo, double hi) {
    return fmin(fmax(v, lo), hi);
}

/* 蛙跳法一步 */
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 box_a, 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;
        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;
        /* wrap */
        x[i] = clamp(x[i], -box_a, box_a);
        y[i] = clamp(y[i], -box_a, box_a);
        z[i] = clamp(z[i], -box_a, box_a);
    }
}

/* ========================================================================
 * 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");

    /* traj_x, traj_y, ... */
    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};

    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, "%.15g", 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\": %.15g,\n", params->DT);
    fprintf(f, "  \"NSTEP\": %d,\n", params->NSTEP);
    fprintf(f, "  \"method\": \"leapfrog\",\n");
    fprintf(f, "  \"warmup_steps\": %d,\n", params->warmup_steps);
    fprintf(f, "  \"G\": [%.15g, %.15g, %.15g],\n", params->G[0], params->G[1], params->G[2]);
    fprintf(f, "  \"B\": [%.15g, %.15g, %.15g],\n", params->B[0], params->B[1], params->B[2]);

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

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

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

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

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

    fprintf(f, "}\n");
    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);

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

    /* 初始化位置/速度 */
    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];
    }

    /* 分配轨迹缓冲区 */
    Trajectory traj;
    traj.n_steps = params.NT;
    traj.n_atoms = n;
    traj.x  = (double*)xmalloc(params.NT * n * sizeof(double) * 6);
    traj.y  = traj.x  + params.NT * n;
    traj.z  = traj.y  + params.NT * n;
    traj.vx = traj.z  + params.NT * n;
    traj.vy = traj.vx + params.NT * n;
    traj.vz = traj.vy + params.NT * n;

    /* 预热 */
    for (int s = 0; s < params.warmup_steps; s++) {
        leapfrog_step(n, x, y, z, vx, vy, vz,
                      atoms.masses, params.G, params.B, &bonds, atoms.fixed,
                      params.box_a, params.DT);
    }

    /* 记录 */
    for (int s = 0; s < params.NT; s++) {
        leapfrog_step(n, x, y, z, vx, vy, vz,
                      atoms.masses, params.G, params.B, &bonds, atoms.fixed,
                      params.box_a, params.DT);
        /* 保存这一帧 */
        for (int i = 0; i < n; i++) {
            traj.x[ s * n + i]  = x[i];
            traj.y[ s * n + i]  = y[i];
            traj.z[ s * n + i]  = z[i];
            traj.vx[s * n + i]  = vx[i];
            traj.vy[s * n + i]  = vy[i];
            traj.vz[s * n + i]  = vz[i];
        }
    }

    /* 输出轨迹 */
    char out_path[512];
    snprintf(out_path, sizeof(out_path), "%s/trajectory.txt", output_dir);
    write_trajectory_json(out_path, &traj, &params, &atoms, &bonds);

    clock_t t1 = clock();
    double elapsed = (double)(t1 - t0) / CLOCKS_PER_SEC;
    printf("[C-engine] 计算完成: %d 步, %.3f s\n", params.NT, 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;
}