! engines/fortran/main.f90
! ------------------------
! Fortran 动力学模拟引擎。
! 与 Python 版 (compute.py) 算法保持一致。
!
! 输入: param.json, <input_dir>/coord.txt, connection.txt, bond.txt
! 输出: <output_dir>/trajectory.txt (JSON)
!
! 编译: cmake --build build --target dynamics_f90
! 用法: ./build/dynamics_f90 <input_dir> <output_dir> <param_json>

program dynamics_f90
    implicit none

    character(len=256) :: input_dir, output_dir, param_path
    integer :: narg, i, n, s
    integer :: NT, NSTEP, warmup_steps, n_atoms, n_bonds
    double precision :: DT, box_a
    double precision :: G(3), B(3)
    integer :: gravity_field, gravity_interaction, elastic_force, damping_force
    integer :: driving_force
    double precision :: gravity_strength
    character(len=32) :: method
    double precision :: t0, t1, elapsed, tw

    ! 原子数据
    integer, allocatable :: atom_ids(:)
    double precision, allocatable :: masses(:), radii(:)
    double precision, allocatable :: pos_0(:, :), vel_0(:, :)
    integer, allocatable :: fixed(:, :)

    ! 成键数据
    integer, allocatable :: bond_pairs(:, :)
    double precision, allocatable :: bond_stiffness(:), bond_rest_lengths(:)

    ! 驱动力数据
    integer :: n_drivers, prog_step
    integer, allocatable :: drv_atom_idx(:)
    double precision, allocatable :: drv_amp_x(:), drv_amp_y(:), drv_amp_z(:)
    double precision, allocatable :: drv_freq_x(:), drv_freq_y(:), drv_freq_z(:)
    double precision, allocatable :: drv_phi_x(:), drv_phi_y(:), drv_phi_z(:)
    integer, allocatable :: drv_has_period(:)
    double precision, allocatable :: drv_period_cycles(:)
    double precision, allocatable :: drv_freeze_x(:), drv_freeze_y(:), drv_freeze_z(:)

    ! 运行时位置/速度
    double precision, allocatable :: x(:), y(:), z(:)
    double precision, allocatable :: vx(:), vy(:), vz(:)

    ! 轨迹缓冲区
    integer :: record_steps
    double precision, allocatable :: traj_x(:, :), traj_y(:, :), traj_z(:, :)
    double precision, allocatable :: traj_vx(:, :), traj_vy(:, :), traj_vz(:, :)

    ! ========================================================================
    ! 主流程
    ! ========================================================================
    narg = command_argument_count()
    if (narg < 3) then
        write(*,*) "用法: dynamics_f90 <input_dir> <output_dir> <param_json>"
        stop
    end if

    call get_command_argument(1, input_dir)
    call get_command_argument(2, output_dir)
    call get_command_argument(3, param_path)

    call cpu_time(t0)

    ! 读取 param.json
    call read_params(param_path, box_a, NT, DT, NSTEP, warmup_steps, method, G, B, &
                     gravity_field, gravity_interaction, &
                     elastic_force, damping_force, gravity_strength, driving_force)

    ! 读取 coord.txt
    call read_coord(input_dir, n_atoms, atom_ids, masses, radii, pos_0, vel_0, fixed)

    ! 读取成键信息
    call read_bonds(input_dir, n_atoms, atom_ids, pos_0, n_bonds, &
                    bond_pairs, bond_stiffness, bond_rest_lengths)

    ! 读取驱动力
    n_drivers = 0
    if (driving_force /= 0) then
        call read_driver(input_dir, n_atoms, atom_ids, n_drivers, &
                         drv_atom_idx, drv_amp_x, drv_amp_y, drv_amp_z, &
                         drv_freq_x, drv_freq_y, drv_freq_z, &
                         drv_phi_x, drv_phi_y, drv_phi_z, &
                         drv_has_period, drv_period_cycles, &
                         drv_freeze_x, drv_freeze_y, drv_freeze_z)
    end if

    write(*, '("[Fortran-engine] 原子数=", i0, ", 键数=", i0, &
           &", NT=", i0, ", DT=", f0.6, ", method=", a)') &
        n_atoms, n_bonds, NT, DT, trim(method)

    ! 初始化位置/速度
    n = n_atoms
    allocate(x(n), y(n), z(n), vx(n), vy(n), vz(n))
    do i = 1, n
        x(i)  = pos_0(i, 1); y(i)  = pos_0(i, 2); z(i)  = pos_0(i, 3)
        vx(i) = vel_0(i, 1); vy(i) = vel_0(i, 2); vz(i) = vel_0(i, 3)
    end do

    ! 分配轨迹缓冲区
    record_steps = NT - warmup_steps
    allocate(traj_x(record_steps, n), traj_y(record_steps, n), traj_z(record_steps, n))
    allocate(traj_vx(record_steps, n), traj_vy(record_steps, n), traj_vz(record_steps, n))

    ! 预热
    ! 初始时刻 t=0 驱动力（与 Python run_simulation 一致）
    if (driving_force /= 0 .and. n_drivers > 0) then
        call apply_driving(n, x, y, z, vx, vy, vz, 0.0d0, 0, DT, &
                           n_drivers, drv_atom_idx, &
                           drv_amp_x, drv_amp_y, drv_amp_z, &
                           drv_freq_x, drv_freq_y, drv_freq_z, &
                           drv_phi_x, drv_phi_y, drv_phi_z, &
                           drv_has_period, drv_period_cycles, &
                           drv_freeze_x, drv_freeze_y, drv_freeze_z)
    end if

    do s = 1, warmup_steps
        if (driving_force /= 0 .and. n_drivers > 0) then
            tw = (s * 1.0d0) * DT
            call apply_driving(n, x, y, z, vx, vy, vz, tw, s-1, DT, &
                               n_drivers, drv_atom_idx, &
                               drv_amp_x, drv_amp_y, drv_amp_z, &
                               drv_freq_x, drv_freq_y, drv_freq_z, &
                               drv_phi_x, drv_phi_y, drv_phi_z, &
                               drv_has_period, drv_period_cycles, &
                               drv_freeze_x, drv_freeze_y, drv_freeze_z)
        end if
        call apply_step(method, n, x, y, z, vx, vy, vz, masses, G, B, &
                        n_bonds, bond_pairs, bond_stiffness, bond_rest_lengths, &
                        fixed, box_a, DT, &
                        gravity_field, gravity_interaction, &
                        elastic_force, damping_force, gravity_strength, &
                        pos_0)
    end do

    ! 记录
    prog_step = record_steps / 100
    if (prog_step < 1) prog_step = 1
    do s = 1, record_steps
        if (mod(s, prog_step) == 0 .and. s > 0) then
            write(*, '("[Fortran-engine] progress: ", i0, "/", i0)') s, record_steps
        end if
        if (driving_force /= 0 .and. n_drivers > 0) then
            tw = ((s-1 + warmup_steps) * 1.0d0) * DT
            call apply_driving(n, x, y, z, vx, vy, vz, tw, s-1, DT, &
                               n_drivers, drv_atom_idx, &
                               drv_amp_x, drv_amp_y, drv_amp_z, &
                               drv_freq_x, drv_freq_y, drv_freq_z, &
                               drv_phi_x, drv_phi_y, drv_phi_z, &
                               drv_has_period, drv_period_cycles, &
                               drv_freeze_x, drv_freeze_y, drv_freeze_z)
        end if
        traj_x(s, :)  = x;  traj_y(s, :)  = y;  traj_z(s, :)  = z
        traj_vx(s, :) = vx; traj_vy(s, :) = vy; traj_vz(s, :) = vz
        call apply_step(method, n, x, y, z, vx, vy, vz, masses, G, B, &
                        n_bonds, bond_pairs, bond_stiffness, bond_rest_lengths, &
                        fixed, box_a, DT, &
                        gravity_field, gravity_interaction, &
                        elastic_force, damping_force, gravity_strength, &
                        pos_0)
    end do

    ! 输出轨迹
    write(*, '("[Fortran-engine] 正在写入轨迹数据…")')
    call write_json(output_dir, traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz, &
                    record_steps, n_atoms, atom_ids, masses, &
                    NT, DT, NSTEP, warmup_steps, method, G, B, &
                    n_bonds, bond_pairs, bond_stiffness, bond_rest_lengths, &
                    driving_force)

    call cpu_time(t1)
    elapsed = t1 - t0
    write(*, '("[Fortran-engine] 计算完成: ", i0, " 步, ", f0.6, " s")') record_steps, elapsed

    deallocate(x, y, z, vx, vy, vz)
    deallocate(traj_x, traj_y, traj_z, traj_vx, traj_vy, traj_vz)
    deallocate(atom_ids, masses, radii, pos_0, vel_0, fixed)
    if (allocated(bond_pairs)) deallocate(bond_pairs, bond_stiffness, bond_rest_lengths)
    if (allocated(drv_atom_idx)) then
        deallocate(drv_atom_idx)
        deallocate(drv_amp_x, drv_amp_y, drv_amp_z)
        deallocate(drv_freq_x, drv_freq_y, drv_freq_z)
        deallocate(drv_phi_x, drv_phi_y, drv_phi_z)
        deallocate(drv_has_period, drv_period_cycles)
        deallocate(drv_freeze_x, drv_freeze_y, drv_freeze_z)
    end if

contains

! ========================================================================
! 读取 param.json
! ========================================================================
subroutine read_params(path, box_a, NT, DT, NSTEP, warmup_steps, method, G, B, &
                       gravity_field, gravity_interaction, &
                       elastic_force, damping_force, gravity_strength, driving_force)
    character(len=*), intent(in) :: path
    double precision, intent(out) :: box_a, DT, G(3), B(3), gravity_strength
    integer, intent(out) :: NT, NSTEP, warmup_steps
    integer, intent(out) :: gravity_field, gravity_interaction, elastic_force, damping_force, driving_force
    character(len=*), intent(out) :: method
    character(len=8096) :: buf
    character(len=256) :: line
    integer :: u, ios

    box_a = 10.0d0; NT = 10000; DT = 0.001d0; NSTEP = 100; warmup_steps = 0
    method = 'leapfrog'
    G = (/ 0.0d0, 0.0d0, -9.8d0 /); B = 0.0d0
    gravity_field = 1; gravity_interaction = 0
    elastic_force = 1; damping_force = 0; gravity_strength = 1.0d0
    driving_force = 0

    open(newunit=u, file=trim(path), status='old', action='read', iostat=ios)
    if (ios /= 0) then
        write(*, '("[Fortran-engine] 警告: 无法打开 ", a, ", 使用默认值")') trim(path)
        return
    end if
    buf = ''
    do
        read(u, '(a)', iostat=ios) line
        if (ios /= 0) exit
        buf = trim(buf) // trim(line) // char(10)
    end do
    close(u)
    if (ios > 0) return

    box_a = json_get_double(buf, 'box_a', box_a)
    NT = json_get_int(buf, 'NT', NT)
    DT = json_get_double(buf, 'DT', DT)
    NSTEP = json_get_int(buf, 'NSTEP', NSTEP)
    warmup_steps = json_get_int(buf, 'warmup_steps', warmup_steps)
    call json_get_string(buf, 'method', method)
    if (len_trim(method) == 0) method = 'leapfrog'
    call json_get_double3(buf, 'G', G)
    call json_get_double3(buf, 'B', B)
    gravity_field       = json_get_int(buf, 'gravity_field', 1)
    gravity_interaction = json_get_int(buf, 'gravity_interaction', 0)
    elastic_force       = json_get_int(buf, 'elastic_force', 1)
    damping_force       = json_get_int(buf, 'damping_force', 0)
    gravity_strength    = json_get_double(buf, 'gravity_strength', 1.0d0)
    driving_force       = json_get_int(buf, 'driving_force', 0)
end subroutine read_params

! ========================================================================
! JSON 辅助解析
! ========================================================================
function json_get_double(buf, key, default) result(val)
    character(len=*), intent(in) :: buf, key
    double precision, intent(in) :: default
    double precision :: val
    integer :: p, ios
    val = default
    p = index(buf, '"' // key // '"')
    if (p == 0) return
    p = index(buf(p:), ':') + p
    do while (p <= len(buf) .and. (buf(p:p) == ':' .or. buf(p:p) == ' ' &
           .or. buf(p:p) == char(9) .or. buf(p:p) == char(10)))
        p = p + 1
    end do
    read(buf(p:), *, iostat=ios) val
end function json_get_double

function json_get_int(buf, key, default) result(val)
    character(len=*), intent(in) :: buf, key
    integer, intent(in) :: default
    integer :: val
    val = nint(json_get_double(buf, key, dble(default)))
end function json_get_int

subroutine json_get_double3(buf, key, vals)
    character(len=*), intent(in) :: buf, key
    double precision, intent(out) :: vals(3)
    integer :: p, i, ios
    vals = 0.0d0
    p = index(buf, '"' // key // '"')
    if (p == 0) return
    p = index(buf(p:), '[') + p
    do i = 1, 3
        do while (buf(p:p) == ' ' .or. buf(p:p) == ',' .or. &
                  buf(p:p) == char(9) .or. buf(p:p) == char(10))
            p = p + 1
        end do
        read(buf(p:), *, iostat=ios) vals(i)
        if (ios /= 0) exit
        do while (p <= len(buf) .and. buf(p:p) /= ',' .and. buf(p:p) /= ']')
            p = p + 1
        end do
    end do
end subroutine json_get_double3

! 从 JSON 中读取字符串值
subroutine json_get_string(buf, key, val)
    character(len=*), intent(in) :: buf, key
    character(len=*), intent(out) :: val
    integer :: p, i, j
    val = ''
    p = index(buf, '"' // key // '"')
    if (p == 0) return
    p = index(buf(p:), ':') + p
    do while (p <= len(buf) .and. (buf(p:p) == ':' .or. buf(p:p) == ' ' &
           .or. buf(p:p) == char(9) .or. buf(p:p) == char(10)))
        p = p + 1
    end do
    if (buf(p:p) /= '"') return
    p = p + 1
    i = 1
    do while (p <= len(buf) .and. buf(p:p) /= '"' .and. i <= len(val))
        val(i:i) = buf(p:p)
        i = i + 1; p = p + 1
    end do
end subroutine json_get_string

! ========================================================================
! 读取 coord.txt
! ========================================================================
subroutine read_coord(input_dir, n_atoms, atom_ids, masses, radii, pos_0, vel_0, fixed)
    character(len=*), intent(in) :: input_dir
    integer, intent(out) :: n_atoms
    integer, allocatable, intent(out) :: atom_ids(:)
    double precision, allocatable, intent(out) :: masses(:), radii(:)
    double precision, allocatable, intent(out) :: pos_0(:, :), vel_0(:, :)
    integer, allocatable, intent(out) :: fixed(:, :)

    character(len=512) :: path, line
    integer :: u, ios, ncols, n_cap, id, fx, fy, fz, i
    double precision :: mass, rad, px, py, pz, vx, vy, vz
    integer, parameter :: MX = 4096
    integer :: ids_tmp(MX), fix_tmp(MX, 3)
    double precision :: m_tmp(MX), r_tmp(MX), p_tmp(MX, 3), v_tmp(MX, 3)

    n_atoms = 0
    path = trim(input_dir) // '/coord.txt'
    open(newunit=u, file=trim(path), status='old', action='read', iostat=ios)
    if (ios /= 0) then; write(*, '("[Fortran-engine] 错误: 无法打开 ", a)') trim(path); stop; end if

    read(u, '(a)', iostat=ios) line
    ncols = 0
    do i = 1, len_trim(line)
        if (line(i:i) /= ' ') then
            if (i == 1) then
                ncols = ncols + 1
            else if (line(i-1:i-1) == ' ') then
                ncols = ncols + 1
            end if
        end if
    end do

    do
        read(u, '(a)', iostat=ios) line
        if (ios /= 0) exit
        if (len_trim(line) == 0 .or. line(1:1) == '#') cycle
        n_atoms = n_atoms + 1
        if (n_atoms > MX) exit
        if (ncols >= 12) then
            read(line, *) id, mass, rad, px, py, pz, vx, vy, vz, fx, fy, fz
        else
            read(line, *) id, mass, rad, px, py, pz, vx, vy, vz
            fx = 0; fy = 0; fz = 0
        end if
        ids_tmp(n_atoms) = id
        m_tmp(n_atoms) = mass; r_tmp(n_atoms) = rad
        p_tmp(n_atoms, :) = (/ px, py, pz /)
        v_tmp(n_atoms, :) = (/ vx, vy, vz /)
        fix_tmp(n_atoms, :) = (/ fx, fy, fz /)
    end do
    close(u)
    if (n_atoms <= 0) then; write(*, '("[Fortran-engine] 错误: coord.txt 为空")') ; stop; end if

    allocate(atom_ids(n_atoms), masses(n_atoms), radii(n_atoms))
    allocate(pos_0(n_atoms, 3), vel_0(n_atoms, 3), fixed(n_atoms, 3))
    atom_ids = ids_tmp(1:n_atoms)
    masses = m_tmp(1:n_atoms); radii = r_tmp(1:n_atoms)
    pos_0 = p_tmp(1:n_atoms, :); vel_0 = v_tmp(1:n_atoms, :)
    fixed = fix_tmp(1:n_atoms, :)
end subroutine read_coord

! ========================================================================
! 读取成键信息 (connection.txt + bond.txt)
! ========================================================================
subroutine read_bonds(input_dir, n_atoms, atom_ids, pos_0, n_bonds, &
                      bond_pairs, bond_stiffness, bond_rest_lengths)
    character(len=*), intent(in) :: input_dir
    integer, intent(in) :: n_atoms, atom_ids(n_atoms)
    double precision, intent(in) :: pos_0(n_atoms, 3)
    integer, intent(out) :: n_bonds
    integer, allocatable, intent(out) :: bond_pairs(:, :)
    double precision, allocatable, intent(out) :: bond_stiffness(:), bond_rest_lengths(:)

    character(len=512) :: conn_path
    integer :: u, ios, a1, a2, idx1, idx2, i
    integer, parameter :: MX = 4096
    integer :: ptmp(MX, 2), idx_map(9999)
    double precision :: ktmp(MX), r0tmp(MX), dx, dy, dz
    character(len=64) :: name

    n_bonds = 0; idx_map = -1
    do i = 1, n_atoms
        if (atom_ids(i) > 0 .and. atom_ids(i) <= 9998) idx_map(atom_ids(i)) = i
    end do

    conn_path = trim(input_dir) // '/connection.txt'
    open(newunit=u, file=trim(conn_path), status='old', action='read', iostat=ios)
    if (ios /= 0) return

    read(u, *, iostat=ios)
    do
        read(u, *, iostat=ios) a1, a2, name
        if (ios /= 0) exit
        call find_bond(trim(input_dir), trim(name), ktmp(n_bonds+1), r0tmp(n_bonds+1))
        idx1 = idx_map(a1); idx2 = idx_map(a2)
        if (idx1 < 1 .or. idx2 < 1) cycle
        if (r0tmp(n_bonds+1) <= 0.0d0) then
            dx = pos_0(idx2, 1) - pos_0(idx1, 1)
            dy = pos_0(idx2, 2) - pos_0(idx1, 2)
            dz = pos_0(idx2, 3) - pos_0(idx1, 3)
            r0tmp(n_bonds+1) = sqrt(dx*dx + dy*dy + dz*dz)
        end if
        n_bonds = n_bonds + 1
        ptmp(n_bonds, :) = (/ idx1 - 1, idx2 - 1 /)
        if (n_bonds >= MX) exit
    end do
    close(u)

    if (n_bonds <= 0) return
    allocate(bond_pairs(n_bonds, 2), bond_stiffness(n_bonds), bond_rest_lengths(n_bonds))
    bond_pairs(:, :) = ptmp(1:n_bonds, :)
    bond_stiffness = ktmp(1:n_bonds)
    bond_rest_lengths = r0tmp(1:n_bonds)
end subroutine read_bonds

subroutine find_bond(bond_dir, bond_name, k, r0)
    character(len=*), intent(in) :: bond_dir, bond_name
    double precision, intent(out) :: k, r0
    character(len=512) :: path, bn
    integer :: u, ios
    double precision :: bk, br
    k = 1.0d0; r0 = -1.0d0
    path = trim(bond_dir) // '/bond.txt'
    open(newunit=u, file=trim(path), status='old', action='read', iostat=ios)
    if (ios /= 0) return
    read(u, *, iostat=ios)
    do
        read(u, *, iostat=ios) bn, bk, br
        if (ios /= 0) exit
        if (trim(bn) == trim(bond_name)) then
            k = bk
            if (br > 0.0d0) r0 = br
            exit
        end if
    end do
    close(u)
end subroutine find_bond

! ========================================================================
! 物理核心
! ========================================================================

! 加速度计算（各力独立开关控制）
pure subroutine accel(n, x, y, z, vx, vy, vz, m, G, B, &
                      nb, bp, bk, br, &
                      gravity_field, gravity_interaction, &
                      elastic_force, damping_force, gravity_strength, &
                      ax, ay, az)
    integer, intent(in) :: n, nb, bp(nb, 2)
    integer, intent(in) :: gravity_field, gravity_interaction, elastic_force, damping_force
    double precision, intent(in) :: x(n), y(n), z(n), vx(n), vy(n), vz(n)
    double precision, intent(in) :: m(n), G(3), B(3), bk(nb), br(nb), gravity_strength
    double precision, intent(out) :: ax(n), ay(n), az(n)
    integer :: i, j, ib
    double precision :: dx, dy, dz, dist, s, fm, ux, uy, uz, r2, fg

    ! 清零
    ax = 0.0d0; ay = 0.0d0; az = 0.0d0

    ! 均匀重力场
    if (gravity_field /= 0) then
        do i = 1, n
            ax(i) = G(1); ay(i) = G(2); az(i) = G(3)
        end do
    end if

    ! 阻尼
    if (damping_force /= 0) then
        do i = 1, n
            ax(i) = ax(i) - B(1) * vx(i) / m(i)
            ay(i) = ay(i) - B(2) * vy(i) / m(i)
            az(i) = az(i) - B(3) * vz(i) / m(i)
        end do
    end if

    ! 弹簧键力
    if (elastic_force /= 0) then
        do ib = 1, nb
            i = bp(ib, 1) + 1; j = bp(ib, 2) + 1
            dx = x(j) - x(i); dy = y(j) - y(i); dz = z(j) - z(i)
            dist = sqrt(dx*dx + dy*dy + dz*dz)
            if (dist < 1.0d-12) cycle
            s = (dist - br(ib)) / dist
            fm = bk(ib) * s
            ux = fm * dx; uy = fm * dy; uz = fm * dz
            ax(i) = ax(i) + ux / m(i); ay(i) = ay(i) + uy / m(i); az(i) = az(i) + uz / m(i)
            ax(j) = ax(j) - ux / m(j); ay(j) = ay(j) - uy / m(j); az(j) = az(j) - uz / m(j)
        end do
    end if

    ! 万有引力（所有原子对之间）
    if (gravity_interaction /= 0) then
        do i = 1, n - 1
            do j = i + 1, n
                dx = x(j) - x(i); dy = y(j) - y(i); dz = z(j) - z(i)
                r2 = dx*dx + dy*dy + dz*dz
                if (r2 <= 1.0d-12) cycle
                dist = sqrt(r2)
                fg = gravity_strength * m(i) * m(j) / r2
                ux = fg * dx / dist; uy = fg * dy / dist; uz = fg * dz / dist
                ax(i) = ax(i) + ux / m(i); ay(i) = ay(i) + uy / m(i); az(i) = az(i) + uz / m(i)
                ax(j) = ax(j) - ux / m(j); ay(j) = ay(j) - uy / m(j); az(j) = az(j) - uz / m(j)
            end do
        end do
    end if
end subroutine accel

! 边界条件：clamp 位置 + 速度反转
subroutine limit_in_box(pos, vel, lo, hi)
    double precision, intent(inout) :: pos, vel
    double precision, intent(in) :: lo, hi
    if (pos > hi) then
        pos = hi; vel = -vel
    else if (pos < lo) then
        pos = lo; vel = -vel
    end if
end subroutine limit_in_box

! 周期边界回绕（与 Python wrap_position 一致）
subroutine wrap_position(pos, lo, hi)
    double precision, intent(inout) :: pos
    double precision, intent(in) :: lo, hi
    if (pos > hi) pos = lo
    if (pos < lo) pos = hi
end subroutine wrap_position

! ── 显式欧拉法 ────────────
subroutine explicit_euler_step(n, x, y, z, vx, vy, vz, m, G, B, &
                               nb, bp, bk, br, fixed, dt, &
                               gravity_field, gravity_interaction, &
                               elastic_force, damping_force, gravity_strength)
    integer, intent(in) :: n, nb, bp(nb, 2), fixed(n, 3)
    integer, intent(in) :: gravity_field, gravity_interaction, elastic_force, damping_force
    double precision, intent(inout) :: x(n), y(n), z(n), vx(n), vy(n), vz(n)
    double precision, intent(in) :: m(n), G(3), B(3), bk(nb), br(nb), dt, gravity_strength
    double precision :: ax(n), ay(n), az(n)
    integer :: i

    call accel(n, x, y, z, vx, vy, vz, m, G, B, nb, bp, bk, br, &
               gravity_field, gravity_interaction, &
               elastic_force, damping_force, gravity_strength, ax, ay, az)
    do i = 1, n
        if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle
        x(i)  = x(i) + vx(i) * dt
        y(i)  = y(i) + vy(i) * dt
        z(i)  = z(i) + vz(i) * dt
        vx(i) = vx(i) + ax(i) * dt
        vy(i) = vy(i) + ay(i) * dt
        vz(i) = vz(i) + az(i) * dt
    end do
end subroutine explicit_euler_step

! ── 隐式欧拉法 ────────────
subroutine implicit_euler_step(n, x, y, z, vx, vy, vz, m, G, B, &
                               nb, bp, bk, br, fixed, dt, &
                               gravity_field, gravity_interaction, &
                               elastic_force, damping_force, gravity_strength)
    integer, intent(in) :: n, nb, bp(nb, 2), fixed(n, 3)
    integer, intent(in) :: gravity_field, gravity_interaction, elastic_force, damping_force
    double precision, intent(inout) :: x(n), y(n), z(n), vx(n), vy(n), vz(n)
    double precision, intent(in) :: m(n), G(3), B(3), bk(nb), br(nb), dt, gravity_strength
    double precision :: ax(n), ay(n), az(n)
    double precision :: vxn(n), vyn(n), vzn(n)
    double precision :: gx, gy, gz
    integer :: i

    ! 隐式速度（重力 + 阻尼）
    do i = 1, n
        if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) then
            vxn(i) = 0; vyn(i) = 0; vzn(i) = 0; cycle
        end if
        gx = B(1) / m(i); gy = B(2) / m(i); gz = B(3) / m(i)
        vxn(i) = (vx(i) + G(1) * dt) / (1.0d0 + gx * dt)
        vyn(i) = (vy(i) + G(2) * dt) / (1.0d0 + gy * dt)
        vzn(i) = (vz(i) + G(3) * dt) / (1.0d0 + gz * dt)
    end do

    ! 用当前位 + 隐式速度算完整加速度
    call accel(n, x, y, z, vxn, vyn, vzn, m, G, B, nb, bp, bk, br, &
               gravity_field, gravity_interaction, &
               elastic_force, damping_force, gravity_strength, ax, ay, az)

    do i = 1, n
        if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle
        vx(i) = vx(i) + ax(i) * dt
        vy(i) = vy(i) + ay(i) * dt
        vz(i) = vz(i) + az(i) * dt
        x(i)  = x(i) + vx(i) * dt
        y(i)  = y(i) + vy(i) * dt
        z(i)  = z(i) + vz(i) * dt
    end do
end subroutine implicit_euler_step

! ── 中点法 ────────────
subroutine midpoint_step(n, x, y, z, vx, vy, vz, m, G, B, &
                         nb, bp, bk, br, fixed, dt, &
                         gravity_field, gravity_interaction, &
                         elastic_force, damping_force, gravity_strength)
    integer, intent(in) :: n, nb, bp(nb, 2), fixed(n, 3)
    integer, intent(in) :: gravity_field, gravity_interaction, elastic_force, damping_force
    double precision, intent(inout) :: x(n), y(n), z(n), vx(n), vy(n), vz(n)
    double precision, intent(in) :: m(n), G(3), B(3), bk(nb), br(nb), dt, gravity_strength
    double precision :: ax(n), ay(n), az(n)
    double precision :: xm(n), ym(n), zm(n), vxm(n), vym(n), vzm(n)
    integer :: i

    call accel(n, x, y, z, vx, vy, vz, m, G, B, nb, bp, bk, br, &
               gravity_field, gravity_interaction, &
               elastic_force, damping_force, gravity_strength, ax, ay, az)
    do i = 1, n
        if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) then
            xm(i)=0; ym(i)=0; zm(i)=0; vxm(i)=0; vym(i)=0; vzm(i)=0; cycle
        end if
        xm(i)  = x(i) + 0.5d0 * vx(i) * dt
        ym(i)  = y(i) + 0.5d0 * vy(i) * dt
        zm(i)  = z(i) + 0.5d0 * vz(i) * dt
        vxm(i) = vx(i) + 0.5d0 * ax(i) * dt
        vym(i) = vy(i) + 0.5d0 * ay(i) * dt
        vzm(i) = vz(i) + 0.5d0 * az(i) * dt
        x(i)   = x(i) + vxm(i) * dt
        y(i)   = y(i) + vym(i) * dt
        z(i)   = z(i) + vzm(i) * dt
    end do

    call accel(n, xm, ym, zm, vxm, vym, vzm, m, G, B, nb, bp, bk, br, &
               gravity_field, gravity_interaction, &
               elastic_force, damping_force, gravity_strength, ax, ay, az)
    do i = 1, n
        if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle
        vx(i) = vx(i) + ax(i) * dt
        vy(i) = vy(i) + ay(i) * dt
        vz(i) = vz(i) + az(i) * dt
    end do
end subroutine midpoint_step

! ── 蛙跳法（Velocity-Verlet）──
subroutine leapfrog_full(n, x, y, z, vx, vy, vz, m, G, B, &
                         nb, bp, bk, br, fixed, dt, &
                         gravity_field, gravity_interaction, &
                         elastic_force, damping_force, gravity_strength)
    integer, intent(in) :: n, nb, bp(nb, 2), fixed(n, 3)
    integer, intent(in) :: gravity_field, gravity_interaction, elastic_force, damping_force
    double precision, intent(inout) :: x(n), y(n), z(n), vx(n), vy(n), vz(n)
    double precision, intent(in) :: m(n), G(3), B(3), bk(nb), br(nb), dt, gravity_strength
    double precision :: ax(n), ay(n), az(n)
    double precision :: dmp_vx(n), dmp_vy(n), dmp_vz(n)
    double precision :: gx, gy, gz
    integer :: i

    call accel(n, x, y, z, vx, vy, vz, m, G, B, nb, bp, bk, br, &
               gravity_field, gravity_interaction, &
               elastic_force, damping_force, gravity_strength, ax, ay, az)

    ! 速度半步推
    do i = 1, n
        if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle
        vx(i) = vx(i) + ax(i) * dt * 0.5d0
        vy(i) = vy(i) + ay(i) * dt * 0.5d0
        vz(i) = vz(i) + az(i) * dt * 0.5d0
    end do

    ! 全推位置（不含边界）
    do i = 1, n
        if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle
        x(i) = x(i) + vx(i) * dt
        y(i) = y(i) + vy(i) * dt
        z(i) = z(i) + vz(i) * dt
    end do

    ! 隐式阻尼处理（用临时数组 dmp_v，不覆盖 vx/vy/vz）
    do i = 1, n
        if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) then
            dmp_vx(i) = 0; dmp_vy(i) = 0; dmp_vz(i) = 0; cycle
        end if
        gx = B(1) / m(i); gy = B(2) / m(i); gz = B(3) / m(i)
        dmp_vx(i) = (vx(i) + 0.5d0 * G(1) * dt) / (1.0d0 + 0.5d0 * gx * dt)
        dmp_vy(i) = (vy(i) + 0.5d0 * G(2) * dt) / (1.0d0 + 0.5d0 * gy * dt)
        dmp_vz(i) = (vz(i) + 0.5d0 * G(3) * dt) / (1.0d0 + 0.5d0 * gz * dt)
    end do

    ! 用新位置 + 阻尼速度重算加速度
    call accel(n, x, y, z, dmp_vx, dmp_vy, dmp_vz, m, G, B, nb, bp, bk, br, &
               gravity_field, gravity_interaction, &
               elastic_force, damping_force, gravity_strength, ax, ay, az)

    ! 速度后半步：v = v_half + 0.5*a_next*dt（vx 仍为 v_half）
    do i = 1, n
        if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle
        vx(i) = vx(i) + ax(i) * dt * 0.5d0
        vy(i) = vy(i) + ay(i) * dt * 0.5d0
        vz(i) = vz(i) + az(i) * dt * 0.5d0
    end do
end subroutine leapfrog_full

! ── 分发器 + 边界条件 + 自由度约束 ──
subroutine apply_step(method, n, x, y, z, vx, vy, vz, m, G, B, &
                      nb, bp, bk, br, fixed, box_a, dt, &
                      gravity_field, gravity_interaction, &
                      elastic_force, damping_force, gravity_strength, &
                      pos_0)
    character(len=*), intent(in) :: method
    integer, intent(in) :: n, nb, bp(nb, 2), fixed(n, 3)
    integer, intent(in) :: gravity_field, gravity_interaction, elastic_force, damping_force
    double precision, intent(inout) :: x(n), y(n), z(n), vx(n), vy(n), vz(n)
    double precision, intent(in) :: m(n), G(3), B(3), bk(nb), br(nb), box_a, dt, gravity_strength
    double precision, intent(in) :: pos_0(n, 3)
    double precision :: mm(n)
    integer :: i

    mm = m

    if (trim(method) == 'explicit_euler') then
        call explicit_euler_step(n, x, y, z, vx, vy, vz, mm, G, B, &
                                 nb, bp, bk, br, fixed, dt, &
                                 gravity_field, gravity_interaction, &
                                 elastic_force, damping_force, gravity_strength)
    else if (trim(method) == 'implicit_euler') then
        call implicit_euler_step(n, x, y, z, vx, vy, vz, mm, G, B, &
                                 nb, bp, bk, br, fixed, dt, &
                                 gravity_field, gravity_interaction, &
                                 elastic_force, damping_force, gravity_strength)
    else if (trim(method) == 'midpoint') then
        call midpoint_step(n, x, y, z, vx, vy, vz, mm, G, B, &
                           nb, bp, bk, br, fixed, dt, &
                           gravity_field, gravity_interaction, &
                           elastic_force, damping_force, gravity_strength)
    else if (trim(method) == 'leapfrog') then
        call leapfrog_full(n, x, y, z, vx, vy, vz, mm, G, B, &
                           nb, bp, bk, br, fixed, dt, &
                           gravity_field, gravity_interaction, &
                           elastic_force, damping_force, gravity_strength)
    else
        write(*, '("[Fortran-engine] 未知算法: ", a)') trim(method)
        stop
    end if

    ! 边界条件
    do i = 1, n
        if (fixed(i,1) /= 0 .and. fixed(i,2) /= 0 .and. fixed(i,3) /= 0) cycle
        call limit_in_box(x(i),  vx(i), -box_a, box_a)
        call limit_in_box(y(i),  vy(i), -box_a, box_a)
        call limit_in_box(z(i),  vz(i), -box_a, box_a)
    end do

    ! 周期边界回绕（与 Python wrap_position 一致）
    do i = 1, n
        call wrap_position(x(i), -box_a, box_a)
        call wrap_position(y(i), -box_a, box_a)
        call wrap_position(z(i), -box_a, box_a)
    end do

    ! 逐自由度固定约束（与 Python apply_fixed_constraints 一致）
    do i = 1, n
        if (fixed(i,1) /= 0) then; x(i) = pos_0(i, 1); vx(i) = 0.0d0; end if
        if (fixed(i,2) /= 0) then; y(i) = pos_0(i, 2); vy(i) = 0.0d0; end if
        if (fixed(i,3) /= 0) then; z(i) = pos_0(i, 3); vz(i) = 0.0d0; end if
    end do
end subroutine apply_step

! ========================================================================
! 读取驱动力 driver.txt
! ========================================================================
subroutine read_driver(input_dir, n_atoms, atom_ids, n_drivers, &
                       drv_atom_idx, drv_amp_x, drv_amp_y, drv_amp_z, &
                       drv_freq_x, drv_freq_y, drv_freq_z, &
                       drv_phi_x, drv_phi_y, drv_phi_z, &
                       drv_has_period, drv_period_cycles, &
                       drv_freeze_x, drv_freeze_y, drv_freeze_z)
    character(len=*), intent(in) :: input_dir
    integer, intent(in) :: n_atoms, atom_ids(n_atoms)
    integer, intent(out) :: n_drivers
    integer, allocatable, intent(out) :: drv_atom_idx(:)
    double precision, allocatable, intent(out) :: drv_amp_x(:), drv_amp_y(:), drv_amp_z(:)
    double precision, allocatable, intent(out) :: drv_freq_x(:), drv_freq_y(:), drv_freq_z(:)
    double precision, allocatable, intent(out) :: drv_phi_x(:), drv_phi_y(:), drv_phi_z(:)
    integer, allocatable, intent(out) :: drv_has_period(:)
    double precision, allocatable, intent(out) :: drv_period_cycles(:)
    double precision, allocatable, intent(out) :: drv_freeze_x(:), drv_freeze_y(:), drv_freeze_z(:)

    character(len=512) :: path, line, period_str
    integer :: u, ios, i, n, idx, n_cap
    double precision :: ax, ay, az, fx, fy, fz, px, py, pz
    integer, parameter :: MX = 256
    integer :: idx_tmp(MX), per_tmp(MX)
    double precision :: ax_tmp(MX), ay_tmp(MX), az_tmp(MX)
    double precision :: fxx_tmp(MX), fyy_tmp(MX), fzz_tmp(MX)
    double precision :: pxx_tmp(MX), pyy_tmp(MX), pzz_tmp(MX)
    double precision :: pc_tmp(MX)
    double precision :: fzx_tmp(MX), fzy_tmp(MX), fzz_tmp2(MX)

    n_drivers = 0
    path = trim(input_dir) // '/driver.txt'
    open(newunit=u, file=trim(path), status='old', action='read', iostat=ios)
    if (ios /= 0) then
        write(*, '("[Fortran-engine] 警告: 无法打开 ", a)') trim(path)
        return
    end if

    read(u, '(a)', iostat=ios)  ! skip header

    do
        read(u, *, iostat=ios) n, ax, ay, az, fx, fy, fz, px, py, pz, period_str
        if (ios /= 0) exit
        n_drivers = n_drivers + 1
        if (n_drivers > MX) exit

        ! Find atom index by id
        idx = -1
        do i = 1, n_atoms
            if (atom_ids(i) == n) then
                idx = i
                exit
            end if
        end do
        if (idx < 0) cycle

        idx_tmp(n_drivers) = idx - 1  ! 0-based index
        ax_tmp(n_drivers) = ax; ay_tmp(n_drivers) = ay; az_tmp(n_drivers) = az
        fxx_tmp(n_drivers) = fx; fyy_tmp(n_drivers) = fy; fzz_tmp(n_drivers) = fz
        ! degrees to radians
        pxx_tmp(n_drivers) = px * 3.141592653589793d0 / 180.0d0
        pyy_tmp(n_drivers) = py * 3.141592653589793d0 / 180.0d0
        pzz_tmp(n_drivers) = pz * 3.141592653589793d0 / 180.0d0

        if (trim(period_str) == 'all') then
            per_tmp(n_drivers) = 0
            pc_tmp(n_drivers) = -1.0d0
        else
            per_tmp(n_drivers) = 1
            read(period_str, *) pc_tmp(n_drivers)
        end if
        fzx_tmp(n_drivers) = 0.0d0
        fzy_tmp(n_drivers) = 0.0d0
        fzz_tmp2(n_drivers) = 0.0d0
    end do
    close(u)

    if (n_drivers <= 0) return

    allocate(drv_atom_idx(n_drivers))
    allocate(drv_amp_x(n_drivers), drv_amp_y(n_drivers), drv_amp_z(n_drivers))
    allocate(drv_freq_x(n_drivers), drv_freq_y(n_drivers), drv_freq_z(n_drivers))
    allocate(drv_phi_x(n_drivers), drv_phi_y(n_drivers), drv_phi_z(n_drivers))
    allocate(drv_has_period(n_drivers), drv_period_cycles(n_drivers))
    allocate(drv_freeze_x(n_drivers), drv_freeze_y(n_drivers), drv_freeze_z(n_drivers))

    drv_atom_idx = idx_tmp(1:n_drivers)
    drv_amp_x = ax_tmp(1:n_drivers); drv_amp_y = ay_tmp(1:n_drivers); drv_amp_z = az_tmp(1:n_drivers)
    drv_freq_x = fxx_tmp(1:n_drivers); drv_freq_y = fyy_tmp(1:n_drivers); drv_freq_z = fzz_tmp(1:n_drivers)
    drv_phi_x = pxx_tmp(1:n_drivers); drv_phi_y = pyy_tmp(1:n_drivers); drv_phi_z = pzz_tmp(1:n_drivers)
    drv_has_period = per_tmp(1:n_drivers)
    drv_period_cycles = pc_tmp(1:n_drivers)
    drv_freeze_x = fzx_tmp(1:n_drivers); drv_freeze_y = fzy_tmp(1:n_drivers); drv_freeze_z = fzz_tmp2(1:n_drivers)

    write(*, '("[Fortran-engine] 已加载驱动力: ", i0, " 条定义")') n_drivers
end subroutine read_driver

! 施加驱动力
subroutine apply_driving(n, x, y, z, vx, vy, vz, t, step, dt, &
                         n_drivers, drv_atom_idx, &
                         drv_amp_x, drv_amp_y, drv_amp_z, &
                         drv_freq_x, drv_freq_y, drv_freq_z, &
                         drv_phi_x, drv_phi_y, drv_phi_z, &
                         drv_has_period, drv_period_cycles, &
                         drv_freeze_x, drv_freeze_y, drv_freeze_z)
    integer, intent(in) :: n, step, n_drivers
    integer, intent(in) :: drv_atom_idx(n_drivers), drv_has_period(n_drivers)
    double precision, intent(inout) :: x(n), y(n), z(n), vx(n), vy(n), vz(n)
    double precision, intent(in) :: t, dt
    double precision, intent(in) :: drv_amp_x(n_drivers), drv_amp_y(n_drivers), drv_amp_z(n_drivers)
    double precision, intent(in) :: drv_freq_x(n_drivers), drv_freq_y(n_drivers), drv_freq_z(n_drivers)
    double precision, intent(in) :: drv_phi_x(n_drivers), drv_phi_y(n_drivers), drv_phi_z(n_drivers)
    double precision, intent(in) :: drv_period_cycles(n_drivers)
    double precision, intent(inout) :: drv_freeze_x(n_drivers), drv_freeze_y(n_drivers), drv_freeze_z(n_drivers)

    integer :: d, idx, period_steps
    double precision :: max_freq, px, py, pz, vpx, vpy, vpz, TWO_PI

    TWO_PI = 2.0d0 * 3.141592653589793d0

    do d = 1, n_drivers
        idx = drv_atom_idx(d) + 1  ! Fortran 1-based indexing

        ! Check period
        if (drv_has_period(d) /= 0) then
            max_freq = max(abs(drv_freq_x(d)), abs(drv_freq_y(d)), abs(drv_freq_z(d)))
            period_steps = 0
            if (max_freq > 1.0d-12) then
                period_steps = int(drv_period_cycles(d) / max_freq / dt)
            end if
            if (step > period_steps) then
                x(idx) = drv_freeze_x(d)
                y(idx) = drv_freeze_y(d)
                z(idx) = drv_freeze_z(d)
                vx(idx) = 0.0d0; vy(idx) = 0.0d0; vz(idx) = 0.0d0
                cycle
            end if
        end if

        px = drv_amp_x(d) * cos(TWO_PI * drv_freq_x(d) * t + drv_phi_x(d))
        py = drv_amp_y(d) * cos(TWO_PI * drv_freq_y(d) * t + drv_phi_y(d))
        pz = drv_amp_z(d) * cos(TWO_PI * drv_freq_z(d) * t + drv_phi_z(d))
        vpx = -drv_amp_x(d) * TWO_PI * drv_freq_x(d) * sin(TWO_PI * drv_freq_x(d) * t + drv_phi_x(d))
        vpy = -drv_amp_y(d) * TWO_PI * drv_freq_y(d) * sin(TWO_PI * drv_freq_y(d) * t + drv_phi_y(d))
        vpz = -drv_amp_z(d) * TWO_PI * drv_freq_z(d) * sin(TWO_PI * drv_freq_z(d) * t + drv_phi_z(d))

        x(idx) = px; y(idx) = py; z(idx) = pz
        vx(idx) = vpx; vy(idx) = vpy; vz(idx) = vpz

        ! Record freeze position
        if (drv_has_period(d) /= 0) then
            max_freq = max(abs(drv_freq_x(d)), abs(drv_freq_y(d)), abs(drv_freq_z(d)))
            period_steps = 0
            if (max_freq > 1.0d-12) then
                period_steps = int(drv_period_cycles(d) / max_freq / dt)
            end if
            if (step == period_steps) then
                drv_freeze_x(d) = px
                drv_freeze_y(d) = py
                drv_freeze_z(d) = pz
            end if
        end if
    end do
end subroutine apply_driving

! ========================================================================
! JSON 输出
! ========================================================================

subroutine write_json(outdir, tx, ty, tz, tvx, tvy, tvz, &
                      nsteps, nat, aid, amass, &
                      NT, DT, NSTEP, warmup, method, G, B, &
                      nb, bp, bk, br, driving_force)
    character(len=*), intent(in) :: outdir, method
    integer, intent(in) :: nsteps, nat, NT, NSTEP, warmup, nb, bp(nb, 2), aid(nat), driving_force
    double precision, intent(in) :: tx(nsteps, nat), ty(nsteps, nat), tz(nsteps, nat)
    double precision, intent(in) :: tvx(nsteps, nat), tvy(nsteps, nat), tvz(nsteps, nat)
    double precision, intent(in) :: DT, G(3), B(3), bk(nb), br(nb), amass(nat)

    character(len=512) :: path, buf
    integer :: u, s, i, ib, ios

    path = trim(outdir) // '/trajectory.txt'
    open(newunit=u, file=trim(path), status='replace', action='write', iostat=ios)
    if (ios /= 0) then
        write(*, '("[Fortran-engine] 错误: 无法写入 ", a)') trim(path)
        stop
    end if

    write(u, '(a)') '{'

    ! traj_x
    write(u, '(a)') '  "traj_x": ['
    do s = 1, nsteps
        call json_arr(u, tx(s, :), nat, s < nsteps, '    ')
    end do
    write(u, '(a)') '  ],'

    ! traj_y
    write(u, '(a)') '  "traj_y": ['
    do s = 1, nsteps
        call json_arr(u, ty(s, :), nat, s < nsteps, '    ')
    end do
    write(u, '(a)') '  ],'

    ! traj_z
    write(u, '(a)') '  "traj_z": ['
    do s = 1, nsteps
        call json_arr(u, tz(s, :), nat, s < nsteps, '    ')
    end do
    write(u, '(a)') '  ],'

    ! traj_vx
    write(u, '(a)') '  "traj_vx": ['
    do s = 1, nsteps
        call json_arr(u, tvx(s, :), nat, s < nsteps, '    ')
    end do
    write(u, '(a)') '  ],'

    ! traj_vy
    write(u, '(a)') '  "traj_vy": ['
    do s = 1, nsteps
        call json_arr(u, tvy(s, :), nat, s < nsteps, '    ')
    end do
    write(u, '(a)') '  ],'

    ! traj_vz
    write(u, '(a)') '  "traj_vz": ['
    do s = 1, nsteps
        call json_arr(u, tvz(s, :), nat, s < nsteps, '    ')
    end do
    write(u, '(a)') '  ],'

    ! 标量参数
    write(buf, '(a, i0, a)') '  "NT": ', NT, ','
    write(u, '(a)') trim(buf)
    write(buf, '(a, g0, a)') '  "DT": ', DT, ','
    write(u, '(a)') trim(buf)
    write(buf, '(a, i0, a)') '  "NSTEP": ', NSTEP, ','
    write(u, '(a)') trim(buf)
    write(buf, '(a, a, a)') '  "method": "', trim(method), '",'
    write(u, '(a)') trim(buf)
    write(buf, '(a, i0, a)') '  "warmup_steps": ', warmup, ','
    write(u, '(a)') trim(buf)

    write(buf, '(a, g0, a, g0, a, g0, a)') &
        '  "G": [', G(1), ', ', G(2), ', ', G(3), '],'
    write(u, '(a)') trim(buf)
    write(buf, '(a, g0, a, g0, a, g0, a)') &
        '  "B": [', B(1), ', ', B(2), ', ', B(3), '],'
    write(u, '(a)') trim(buf)

    ! 原子信息
    write(u, '(a)', advance='no') '  "atom_ids": ['
    do i = 1, nat
        if (i > 1) write(u, '(a)', advance='no') ','
        write(u, '(i0)', advance='no') aid(i)
    end do
    write(u, '(a)') '],'

    write(u, '(a)', advance='no') '  "atom_masses": ['
    do i = 1, nat
        if (i > 1) write(u, '(a)', advance='no') ','
        write(u, '(g0)', advance='no') amass(i)
    end do
    write(u, '(a)') '],'

    ! 成键
    if (nb > 0) then
        call write_int2_arr(u, 'bond_pairs', bp, nb, .true.)
        call write_dbl_arr(u, 'bond_stiffness', bk, nb, .true.)
        call write_dbl_arr(u, 'bond_rest_lengths', br, nb, .true.)
    else
        write(u, '(a)') '  "bond_pairs": [],'
        write(u, '(a)') '  "bond_stiffness": [],'
        write(u, '(a)') '  "bond_rest_lengths": [],'
    end if

    write(buf, '(a, i0)') '  "driving_force": ', driving_force
    write(u, '(a)') trim(buf)

    write(u, '(a)') '}'
    close(u)
end subroutine write_json

! 写出单行 JSON 数组 [v1, v2, ...]
subroutine json_arr(u, vals, n, has_next, indent)
    integer, intent(in) :: u, n
    double precision, intent(in) :: vals(n)
    logical, intent(in) :: has_next
    character(len=*), intent(in) :: indent
    integer :: i
    write(u, '(a)', advance='no') indent // '['
    do i = 1, n
        if (i > 1) write(u, '(a)', advance='no') ','
        write(u, '(g0.8)', advance='no') vals(i)
    end do
    if (has_next) then
        write(u, '(a)') '],'
    else
        write(u, '(a)') ']'
    end if
end subroutine json_arr

subroutine write_int2_arr(u, name, arr, n, has_next)
    integer, intent(in) :: u, n, arr(n, 2)
    character(len=*), intent(in) :: name
    logical, intent(in) :: has_next
    character(len=65536) :: buf
    integer :: i, pos
    write(u, '(a)', advance='no') '  "' // trim(name) // '": ['
    do i = 1, n
        if (i > 1) write(u, '(a)', advance='no') ','
        write(buf, '(a, i0, a, i0, a)') '[', arr(i, 1), ',', arr(i, 2), ']'
        write(u, '(a)', advance='no') trim(buf)
    end do
    if (has_next) then
        write(u, '(a)') '],'
    else
        write(u, '(a)') ']'
    end if
end subroutine write_int2_arr

subroutine write_dbl_arr(u, name, arr, n, has_next)
    integer, intent(in) :: u, n
    double precision, intent(in) :: arr(n)
    character(len=*), intent(in) :: name
    logical, intent(in) :: has_next
    integer :: i
    write(u, '(a)', advance='no') '  "' // trim(name) // '": ['
    do i = 1, n
        if (i > 1) write(u, '(a)', advance='no') ','
        write(u, '(g0.8)', advance='no') arr(i)
    end do
    if (has_next) then
        write(u, '(a)') '],'
    else
        write(u, '(a)') ']'
    end if
end subroutine write_dbl_arr

end program dynamics_f90