"""
engines/python/dynamics_lib.py
-------------------------------
纯 NumPy 计算引擎：无文件 I/O，所有数据以 NumPy 数组传入，
结果作为 NumPy 数组返回。

接口与 C/C++/Fortran DLL 的 run_dynamics() 完全一致，
算法与 compute.py 的 run_simulation() 保持一致。

用法（由 engine_dll.py 内部调用）：
    from engines.python.dynamics_lib import run_dynamics
    out_x, out_y, out_z, out_vx, out_vy, out_vz = run_dynamics(...)
"""

import numpy as np

TWO_PI = 2.0 * np.pi

# ── method_id 映射 ──────────────────────────────────────────
# 0=euler  1=implicit_euler  2=midpoint  3=leapfrog


# ── 保守加速度（弹簧键 + 均匀重力场，不含阻尼）──────────────
def _accel_conservative(x, y, z, m, Gx, Gy, Gz,
                         gravity_field, elastic_force,
                         bond_pairs, bond_k, bond_r0):
    ax = np.full_like(x, Gx) if gravity_field else np.zeros_like(x)
    ay = np.full_like(y, Gy) if gravity_field else np.zeros_like(y)
    az = np.full_like(z, Gz) if gravity_field else np.zeros_like(z)

    if elastic_force and len(bond_pairs) > 0:
        i1 = bond_pairs[:, 0]
        i2 = bond_pairs[:, 1]
        dx = x[i2] - x[i1]
        dy = y[i2] - y[i1]
        dz = z[i2] - z[i1]
        dist = np.sqrt(dx*dx + dy*dy + dz*dz)
        valid = dist > 1e-12
        fac = np.where(valid, bond_k * (dist - bond_r0) / dist, 0.0)
        fx = fac * dx
        fy = fac * dy
        fz_b = fac * dz
        np.add.at(ax, i1, fx / m[i1]);   np.add.at(ax, i2, -fx / m[i2])
        np.add.at(ay, i1, fy / m[i1]);   np.add.at(ay, i2, -fy / m[i2])
        np.add.at(az, i1, fz_b / m[i1]); np.add.at(az, i2, -fz_b / m[i2])

    return ax, ay, az


# ── 完整加速度（含阻尼）──────────────────────────────────────
def _accel_full(x, y, z, vx, vy, vz, m, Gx, Gy, Gz, Bx, By, Bz,
                gravity_field, elastic_force, damping_force,
                bond_pairs, bond_k, bond_r0):
    ax, ay, az = _accel_conservative(x, y, z, m, Gx, Gy, Gz,
                                      gravity_field, elastic_force,
                                      bond_pairs, bond_k, bond_r0)
    if damping_force:
        ax -= Bx * vx / m
        ay -= By * vy / m
        az -= Bz * vz / m
    return ax, ay, az


# ── 蛙跳法（半隐式阻尼，与 compute.py leapfrog_staggered_step 一致）─
def _leapfrog_step(x, y, z, vx, vy, vz, fixed, m,
                    Gx, Gy, Gz, Bx, By, Bz,
                    gravity_field, elastic_force, damping_force,
                    bond_pairs, bond_k, bond_r0, dt):
    ax, ay, az = _accel_conservative(x, y, z, m, Gx, Gy, Gz,
                                      gravity_field, elastic_force,
                                      bond_pairs, bond_k, bond_r0)
    has_damp = damping_force and (Bx != 0.0 or By != 0.0 or Bz != 0.0)
    if has_damp:
        alpha_x = Bx * dt / (2.0 * m)
        alpha_y = By * dt / (2.0 * m)
        alpha_z = Bz * dt / (2.0 * m)
        vx_new = (vx * (1.0 - alpha_x) + ax * dt) / (1.0 + alpha_x)
        vy_new = (vy * (1.0 - alpha_y) + ay * dt) / (1.0 + alpha_y)
        vz_new = (vz * (1.0 - alpha_z) + az * dt) / (1.0 + alpha_z)
    else:
        vx_new = vx + ax * dt
        vy_new = vy + ay * dt
        vz_new = vz + az * dt
    # 全固定原子保持不变
    all_fixed = np.all(fixed, axis=1)
    vx_new = np.where(all_fixed, vx, vx_new)
    vy_new = np.where(all_fixed, vy, vy_new)
    vz_new = np.where(all_fixed, vz, vz_new)
    x_new = x + vx_new * dt
    y_new = y + vy_new * dt
    z_new = z + vz_new * dt
    return x_new, y_new, z_new, vx_new, vy_new, vz_new


# ── 显式欧拉法 ───────────────────────────────────────────────
def _euler_step(x, y, z, vx, vy, vz, fixed, m,
                 Gx, Gy, Gz, Bx, By, Bz,
                 gravity_field, elastic_force, damping_force,
                 bond_pairs, bond_k, bond_r0, dt):
    ax, ay, az = _accel_full(x, y, z, vx, vy, vz, m, Gx, Gy, Gz, Bx, By, Bz,
                              gravity_field, elastic_force, damping_force,
                              bond_pairs, bond_k, bond_r0)
    all_fixed = np.all(fixed, axis=1)
    mask = ~all_fixed
    x_new  = np.where(mask, x  + vx * dt, x)
    y_new  = np.where(mask, y  + vy * dt, y)
    z_new  = np.where(mask, z  + vz * dt, z)
    vx_new = np.where(mask, vx + ax * dt, vx)
    vy_new = np.where(mask, vy + ay * dt, vy)
    vz_new = np.where(mask, vz + az * dt, vz)
    return x_new, y_new, z_new, vx_new, vy_new, vz_new


# ── 隐式欧拉法（与 compute.py Implicit_Euler_Method 一致）──────
def _implicit_euler_step(x, y, z, vx, vy, vz, fixed, m,
                          Gx, Gy, Gz, Bx, By, Bz,
                          gravity_field, elastic_force, damping_force,
                          bond_pairs, bond_k, bond_r0, dt):
    gamma_x = Bx / m
    gamma_y = By / m
    gamma_z = Bz / m
    vx_next = (vx + Gx * dt) / (1.0 + gamma_x * dt)
    vy_next = (vy + Gy * dt) / (1.0 + gamma_y * dt)
    vz_next = (vz + Gz * dt) / (1.0 + gamma_z * dt)
    ax, ay, az = _accel_full(x, y, z, vx_next, vy_next, vz_next, m,
                              Gx, Gy, Gz, Bx, By, Bz,
                              gravity_field, elastic_force, damping_force,
                              bond_pairs, bond_k, bond_r0)
    all_fixed = np.all(fixed, axis=1)
    mask = ~all_fixed
    vx_new = np.where(mask, vx + ax * dt, vx)
    vy_new = np.where(mask, vy + ay * dt, vy)
    vz_new = np.where(mask, vz + az * dt, vz)
    x_new  = np.where(mask, x  + vx_new * dt, x)
    y_new  = np.where(mask, y  + vy_new * dt, y)
    z_new  = np.where(mask, z  + vz_new * dt, z)
    return x_new, y_new, z_new, vx_new, vy_new, vz_new


# ── 中点法（与 compute.py Midpoint_Method 一致）────────────────
def _midpoint_step(x, y, z, vx, vy, vz, fixed, m,
                    Gx, Gy, Gz, Bx, By, Bz,
                    gravity_field, elastic_force, damping_force,
                    bond_pairs, bond_k, bond_r0, dt):
    ax, ay, az = _accel_full(x, y, z, vx, vy, vz, m, Gx, Gy, Gz, Bx, By, Bz,
                              gravity_field, elastic_force, damping_force,
                              bond_pairs, bond_k, bond_r0)
    all_fixed = np.all(fixed, axis=1)
    mask = ~all_fixed
    xm  = np.where(mask, x  + 0.5*vx*dt, x)
    ym  = np.where(mask, y  + 0.5*vy*dt, y)
    zm  = np.where(mask, z  + 0.5*vz*dt, z)
    vxm = np.where(mask, vx + 0.5*ax*dt, 0.0)
    vym = np.where(mask, vy + 0.5*ay*dt, 0.0)
    vzm = np.where(mask, vz + 0.5*az*dt, 0.0)
    x_new = np.where(mask, x + vxm * dt, x)
    y_new = np.where(mask, y + vym * dt, y)
    z_new = np.where(mask, z + vzm * dt, z)
    axm, aym, azm = _accel_full(xm, ym, zm, vxm, vym, vzm, m, Gx, Gy, Gz, Bx, By, Bz,
                                 gravity_field, elastic_force, damping_force,
                                 bond_pairs, bond_k, bond_r0)
    vx_new = np.where(mask, vx + axm * dt, vx)
    vy_new = np.where(mask, vy + aym * dt, vy)
    vz_new = np.where(mask, vz + azm * dt, vz)
    return x_new, y_new, z_new, vx_new, vy_new, vz_new


# ── 边界：反弹 + 回绕 + 逐自由度固定约束 ───────────────────────
def _apply_bc(x, y, z, vx, vy, vz, fixed, pos_init, box_a):
    lo, hi = -box_a, box_a

    # 反弹（全固定原子跳过）
    all_fixed = np.all(fixed, axis=1)
    do_bc = ~all_fixed
    over_x  = do_bc & (x > hi); under_x = do_bc & (x < lo)
    over_y  = do_bc & (y > hi); under_y = do_bc & (y < lo)
    over_z  = do_bc & (z > hi); under_z = do_bc & (z < lo)
    x  = np.where(over_x, hi, np.where(under_x, lo, x))
    y  = np.where(over_y, hi, np.where(under_y, lo, y))
    z  = np.where(over_z, hi, np.where(under_z, lo, z))
    vx = np.where(over_x | under_x, -np.abs(vx)*np.sign(np.where(over_x, 1, -1)), vx)
    vy = np.where(over_y | under_y, -np.abs(vy)*np.sign(np.where(over_y, 1, -1)), vy)
    vz = np.where(over_z | under_z, -np.abs(vz)*np.sign(np.where(over_z, 1, -1)), vz)
    # 反弹速度简化：越界则取反绝对值（与 C 版 _limit1 一致）
    vx = np.where(over_x, -np.abs(vx), np.where(under_x, np.abs(vx), vx))
    vy = np.where(over_y, -np.abs(vy), np.where(under_y, np.abs(vy), vy))
    vz = np.where(over_z, -np.abs(vz), np.where(under_z, np.abs(vz), vz))

    # 回绕
    x = np.where(x > hi, lo, np.where(x < lo, hi, x))
    y = np.where(y > hi, lo, np.where(y < lo, hi, y))
    z = np.where(z > hi, lo, np.where(z < lo, hi, z))

    # 逐自由度固定约束
    fx = fixed[:, 0].astype(bool)
    fy = fixed[:, 1].astype(bool)
    fz = fixed[:, 2].astype(bool)
    x  = np.where(fx, pos_init[:, 0], x);  vx = np.where(fx, 0.0, vx)
    y  = np.where(fy, pos_init[:, 1], y);  vy = np.where(fy, 0.0, vy)
    z  = np.where(fz, pos_init[:, 2], z);  vz = np.where(fz, 0.0, vz)
    return x, y, z, vx, vy, vz


# ── 驱动力（与 compute.py apply_driving_force 逻辑一致）─────────
def _apply_driving(x, y, z, vx, vy, vz, t, step, dt,
                    drv_idx, drv_amp, drv_freq, drv_phi, drv_eq,
                    drv_ncycles, drv_has_period, freeze):
    """freeze: (n_drivers, 3) mutable array for frozen positions."""
    nd = len(drv_idx)
    for d in range(nd):
        idx = drv_idx[d]
        fx_ = drv_freq[d, 0]; fy_ = drv_freq[d, 1]; fz_ = drv_freq[d, 2]

        if drv_has_period[d]:
            mf = max(abs(fx_), abs(fy_), abs(fz_))
            ps = int(drv_ncycles[d] / mf / dt) if mf > 1e-12 else 0
            if step > ps:
                x[idx] = freeze[d, 0]; y[idx] = freeze[d, 1]; z[idx] = freeze[d, 2]
                vx[idx] = vy[idx] = vz[idx] = 0.0
                continue
            px = drv_eq[d,0] + drv_amp[d,0]*np.cos(TWO_PI*fx_*t + drv_phi[d,0])
            py = drv_eq[d,1] + drv_amp[d,1]*np.cos(TWO_PI*fy_*t + drv_phi[d,1])
            pz = drv_eq[d,2] + drv_amp[d,2]*np.cos(TWO_PI*fz_*t + drv_phi[d,2])
            if step == ps:
                freeze[d, 0] = px; freeze[d, 1] = py; freeze[d, 2] = pz

        x[idx]  = drv_eq[d,0] + drv_amp[d,0]*np.cos(TWO_PI*fx_*t + drv_phi[d,0])
        y[idx]  = drv_eq[d,1] + drv_amp[d,1]*np.cos(TWO_PI*fy_*t + drv_phi[d,1])
        z[idx]  = drv_eq[d,2] + drv_amp[d,2]*np.cos(TWO_PI*fz_*t + drv_phi[d,2])
        vx[idx] = -drv_amp[d,0]*TWO_PI*fx_*np.sin(TWO_PI*fx_*t + drv_phi[d,0])
        vy[idx] = -drv_amp[d,1]*TWO_PI*fy_*np.sin(TWO_PI*fy_*t + drv_phi[d,1])
        vz[idx] = -drv_amp[d,2]*TWO_PI*fz_*np.sin(TWO_PI*fz_*t + drv_phi[d,2])


def _do_step(x, y, z, vx, vy, vz, fixed, masses, method_id,
              Gx, Gy, Gz, Bx, By, Bz,
              gravity_field, elastic_force, damping_force,
              bond_pairs, bond_k, bond_r0, dt, pos_init, box_a):
    if method_id == 0:
        x, y, z, vx, vy, vz = _euler_step(
            x, y, z, vx, vy, vz, fixed, masses,
            Gx, Gy, Gz, Bx, By, Bz,
            gravity_field, elastic_force, damping_force,
            bond_pairs, bond_k, bond_r0, dt)
    elif method_id == 1:
        x, y, z, vx, vy, vz = _implicit_euler_step(
            x, y, z, vx, vy, vz, fixed, masses,
            Gx, Gy, Gz, Bx, By, Bz,
            gravity_field, elastic_force, damping_force,
            bond_pairs, bond_k, bond_r0, dt)
    elif method_id == 2:
        x, y, z, vx, vy, vz = _midpoint_step(
            x, y, z, vx, vy, vz, fixed, masses,
            Gx, Gy, Gz, Bx, By, Bz,
            gravity_field, elastic_force, damping_force,
            bond_pairs, bond_k, bond_r0, dt)
    else:
        x, y, z, vx, vy, vz = _leapfrog_step(
            x, y, z, vx, vy, vz, fixed, masses,
            Gx, Gy, Gz, Bx, By, Bz,
            gravity_field, elastic_force, damping_force,
            bond_pairs, bond_k, bond_r0, dt)
    x, y, z, vx, vy, vz = _apply_bc(x, y, z, vx, vy, vz, fixed, pos_init, box_a)
    return x, y, z, vx, vy, vz


# ══════════════════════════════════════════════════════════════
# 主函数：run_dynamics
# 接口与 C/C++/Fortran DLL 的 run_dynamics() 对应，
# 参数格式：numpy 数组（替代 ctypes 指针）。
#
# method_id: 0=euler 1=implicit_euler 2=midpoint 3=leapfrog
# drv_amp/freq/phi/eq: (n_drivers, 3) float64
# drv_ncycles: (n_drivers,) float64   0=不限
# drv_has_period: (n_drivers,) int
#
# 返回：(out_x, out_y, out_z, out_vx, out_vy, out_vz)
#       各 shape=(n_frames, n_atoms)
# ══════════════════════════════════════════════════════════════
def run_dynamics(
    n_atoms, pos_init, vel_init, masses, fixed,
    n_bonds, bond_pairs, bond_k, bond_r0,
    box_a, dt,
    NT, NSTEP, warmup_steps, method_id,
    Gx, Gy, Gz, Bx, By, Bz,
    gravity_field, elastic_force, damping_force, gravity_strength,
    n_drivers, drv_idx, drv_amp, drv_freq, drv_phi, drv_eq,
    drv_ncycles, drv_has_period,
    n_frames,
    progress_cb=None,
):
    """运行动力学模拟，返回抽帧轨迹数组。

    Args:
        pos_init:   (n_atoms, 3) float64
        vel_init:   (n_atoms, 3) float64
        masses:     (n_atoms,)   float64
        fixed:      (n_atoms, 3) int  — 1=固定
        bond_pairs: (n_bonds, 2) int  — 0-based 局部索引
        bond_k:     (n_bonds,)   float64
        bond_r0:    (n_bonds,)   float64
        drv_idx:    (n_drivers,) int  — 0-based
        drv_amp/freq/phi/eq: (n_drivers, 3) float64
        drv_ncycles:  (n_drivers,) float64
        drv_has_period: (n_drivers,) int
        n_frames:   预分配的输出帧数

    Returns:
        out_x, out_y, out_z, out_vx, out_vy, out_vz — 各 (n_frames, n_atoms)
    """
    pos_init  = np.asarray(pos_init,  dtype=np.float64)
    vel_init  = np.asarray(vel_init,  dtype=np.float64)
    masses    = np.asarray(masses,    dtype=np.float64)
    fixed     = np.asarray(fixed,     dtype=np.int32)
    bond_pairs = np.asarray(bond_pairs, dtype=np.int64).reshape(-1, 2) if n_bonds else np.zeros((0,2), dtype=np.int64)
    bond_k    = np.asarray(bond_k,    dtype=np.float64) if n_bonds else np.zeros(0)
    bond_r0   = np.asarray(bond_r0,   dtype=np.float64) if n_bonds else np.zeros(0)

    n = n_atoms
    x  = pos_init[:, 0].copy()
    y  = pos_init[:, 1].copy()
    z  = pos_init[:, 2].copy()
    vx = vel_init[:, 0].copy()
    vy = vel_init[:, 1].copy()
    vz = vel_init[:, 2].copy()

    # 驱动力数据（保证正确形状）
    nd = n_drivers
    if nd > 0:
        drv_idx   = np.asarray(drv_idx,        dtype=np.int64)
        drv_amp   = np.asarray(drv_amp,         dtype=np.float64).reshape(nd, 3)
        drv_freq  = np.asarray(drv_freq,        dtype=np.float64).reshape(nd, 3)
        drv_phi   = np.asarray(drv_phi,         dtype=np.float64).reshape(nd, 3)
        drv_eq    = np.asarray(drv_eq,          dtype=np.float64).reshape(nd, 3)
        drv_nc    = np.asarray(drv_ncycles,     dtype=np.float64)
        drv_hp    = np.asarray(drv_has_period,  dtype=np.int32)
        freeze    = np.zeros((nd, 3), dtype=np.float64)
    else:
        drv_idx = drv_amp = drv_freq = drv_phi = drv_eq = drv_nc = drv_hp = freeze = None

    def _drive(t_, step_):
        if nd > 0:
            _apply_driving(x, y, z, vx, vy, vz, t_, step_, dt,
                           drv_idx, drv_amp, drv_freq, drv_phi, drv_eq,
                           drv_nc, drv_hp, freeze)

    # ── 蛙跳法：初始化 v(-dt/2) ─────────────────────────────
    if method_id == 3:
        ax0, ay0, az0 = _accel_conservative(x, y, z, masses, Gx, Gy, Gz,
                                             gravity_field, elastic_force,
                                             bond_pairs, bond_k, bond_r0)
        all_fixed = np.all(fixed, axis=1)
        vx = np.where(all_fixed, vx, vx - 0.5 * ax0 * dt)
        vy = np.where(all_fixed, vy, vy - 0.5 * ay0 * dt)
        vz = np.where(all_fixed, vz, vz - 0.5 * az0 * dt)

    # ── 初始驱动 t=0 ─────────────────────────────────────────
    _drive(0.0, 0)

    # ── 预热 ─────────────────────────────────────────────────
    for s in range(warmup_steps):
        tw = (s + 1) * dt
        _drive(tw, s)
        x, y, z, vx, vy, vz = _do_step(
            x, y, z, vx, vy, vz, fixed, masses, method_id,
            Gx, Gy, Gz, Bx, By, Bz,
            gravity_field, elastic_force, damping_force,
            bond_pairs, bond_k, bond_r0, dt, pos_init, box_a)

    # ── 记录循环 ─────────────────────────────────────────────
    record_steps = NT - warmup_steps
    prog_interval = max(1, record_steps // 100)

    out_x  = np.zeros((n_frames, n), dtype=np.float64)
    out_y  = np.zeros((n_frames, n), dtype=np.float64)
    out_z  = np.zeros((n_frames, n), dtype=np.float64)
    out_vx = np.zeros((n_frames, n), dtype=np.float64)
    out_vy = np.zeros((n_frames, n), dtype=np.float64)
    out_vz = np.zeros((n_frames, n), dtype=np.float64)
    frame_idx = 0

    for s in range(record_steps):
        if progress_cb is not None and s % prog_interval == 0 and s > 0:
            progress_cb(s, record_steps)

        t = (s + warmup_steps) * dt
        _drive(t, s)

        if s % NSTEP == 0 and frame_idx < n_frames:
            out_x[frame_idx]  = x
            out_y[frame_idx]  = y
            out_z[frame_idx]  = z
            out_vx[frame_idx] = vx
            out_vy[frame_idx] = vy
            out_vz[frame_idx] = vz
            frame_idx += 1

        x, y, z, vx, vy, vz = _do_step(
            x, y, z, vx, vy, vz, fixed, masses, method_id,
            Gx, Gy, Gz, Bx, By, Bz,
            gravity_field, elastic_force, damping_force,
            bond_pairs, bond_k, bond_r0, dt, pos_init, box_a)

    return out_x, out_y, out_z, out_vx, out_vy, out_vz
