! 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
    double precision :: gravity_strength
    character(len=32) :: method
    double precision :: t0, t1, elapsed

    ! 原子数据
    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(:)

    ! 运行时位置/速度
    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)

    ! 读取 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)

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

    ! 预热
    do s = 1, warmup_steps
        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

    ! 记录
    do s = 1, record_steps
        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

    ! 输出轨迹
    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)

    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)

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

    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)
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) /= ' ' .and. (i == 1 .or. line(i-1:i-1) == ' ')) ncols = ncols + 1
    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

! ── 显式欧拉法 ────────────
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 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

! ========================================================================
! 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)
    character(len=*), intent(in) :: outdir, method
    integer, intent(in) :: nsteps, nat, NT, NSTEP, warmup, nb, bp(nb, 2), aid(nat)
    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, .false.)
    else
        write(u, '(a)') '  "bond_pairs": [],'
        write(u, '(a)') '  "bond_stiffness": [],'
        write(u, '(a)') '  "bond_rest_lengths": []'
    end if

    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)', 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)', 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