"""
engines/engine_dll.py
---------------------
Python ctypes 包装器：加载 C/C++/Fortran 动态链接库并调用 run_dynamics()。

用法（由 compute.py 内部调用，不直接运行）：

    from engines.engine_dll import load_dll, run_dynamics_dll

    dll = load_dll("c")          # 自动查找 engines/c/build/dynamics_c.dll/.so/.dylib
    arrays = run_dynamics_dll(dll, config, atom_data, bond_data, driver_data)
    # arrays: dict with keys x, y, z, vx, vy, vz  shape=(n_frames, n_atoms)

DLL 编译（C 版本）：
  Windows:  gcc -O3 -shared -o engines/release/dynamics_c.dll engines/src/c/dynamics_lib.c -lm
  Linux:    gcc -O3 -shared -fPIC -o engines/release/dynamics_c.so  engines/src/c/dynamics_lib.c -lm
  macOS:    gcc -O3 -dynamiclib -o engines/release/dynamics_c.dylib engines/src/c/dynamics_lib.c -lm
或用 make dll 一键编译：
  cd engines/src/c && make dll
"""

import ctypes
import os
import platform
import numpy as np

# ── DLL 文件名后缀 ─────────────────────────────────────────────
_SUFFIX = {
    "windows": ".dll",
    "linux":   ".so",
    "darwin":  ".dylib",
}

# ── method 字符串 → 整数 ID ────────────────────────────────────
_METHOD_ID = {
    "explicit_euler": 0,
    "euler":          0,
    "implicit_euler": 1,
    "midpoint":       2,
    "leapfrog":       3,
}

_HERE = os.path.dirname(os.path.abspath(__file__))


_DLL_NAME = {
    "c":       "dynamics_c",
    "cpp":     "dynamics_cpp",
    "c++":     "dynamics_cpp",
    "fortran": "dynamics_f90",
    "f90":     "dynamics_f90",
    # "python" 引擎通过直接 import 调用，不使用 DLL
}

# 引擎名规范化：将别名统一为目录名
_ENGINE_DIR = {
    "c":       "c",
    "cpp":     "cpp",
    "c++":     "cpp",
    "fortran": "fortran",
    "f90":     "fortran",
    "python":  "python",
}


def _dll_candidates(engine: str) -> list[str]:
    """返回 DLL 候选路径列表（按优先级）。"""
    sys = platform.system().lower()
    ext = _SUFFIX.get(sys, ".so")
    eng_dir = _ENGINE_DIR.get(engine, engine)
    name    = _DLL_NAME.get(engine, f"dynamics_{engine}")
    base = os.path.join(_HERE, "release", name)
    return [
        base + ext,
        base + ".dll",
        base + ".so",
        base + ".dylib",
    ]


def load_dll(engine: str = "c"):
    """加载指定引擎。

    - C/C++/Fortran: 返回 ctypes.CDLL 对象
    - Python: 返回模块对象（直接 import，无需编译）

    Args:
        engine: "c", "cpp", "fortran", 或 "python"
    Raises:
        FileNotFoundError: DLL/模块文件不存在
    """
    if _ENGINE_DIR.get(engine, engine) == "python":
        import importlib.util, sys as _sys
        mod_path = os.path.join(_HERE, "python", "dynamics_lib.py")
        if not os.path.exists(mod_path):
            raise FileNotFoundError(f"Python 引擎未找到: {mod_path}")
        spec = importlib.util.spec_from_file_location(
            "engines.python.dynamics_lib", mod_path)
        mod = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        return mod  # 返回模块，不是 CDLL

    for p in _dll_candidates(engine):
        if os.path.exists(p):
            lib = ctypes.CDLL(p)
            _setup_prototype(lib)
            return lib
    raise FileNotFoundError(
        f"DLL 未找到（引擎 {engine}），候选路径：\n" +
        "\n".join(f"  {p}" for p in _dll_candidates(engine)) +
        f"\n请先编译：cd engines/{engine} && make dll"
    )


def _setup_prototype(lib: ctypes.CDLL) -> None:
    """配置 run_dynamics 的参数类型和返回类型。"""
    c_dbl_p = ctypes.POINTER(ctypes.c_double)
    c_int_p = ctypes.POINTER(ctypes.c_int)
    cb_type  = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_int)

    lib.run_dynamics.restype = ctypes.c_int
    lib.run_dynamics.argtypes = [
        ctypes.c_int,    # n_atoms
        c_dbl_p,         # pos_init    [n_atoms*3]
        c_dbl_p,         # vel_init    [n_atoms*3]
        c_dbl_p,         # masses      [n_atoms]
        c_int_p,         # fixed       [n_atoms*3]
        ctypes.c_int,    # n_bonds
        c_int_p,         # bond_pairs  [n_bonds*2]
        c_dbl_p,         # bond_k      [n_bonds]
        c_dbl_p,         # bond_r0     [n_bonds]
        ctypes.c_double, # box_a
        ctypes.c_double, # dt
        ctypes.c_int,    # NT
        ctypes.c_int,    # NSTEP
        ctypes.c_int,    # warmup_steps
        ctypes.c_int,    # method_id
        ctypes.c_double, # Gx
        ctypes.c_double, # Gy
        ctypes.c_double, # Gz
        ctypes.c_double, # Bx
        ctypes.c_double, # By
        ctypes.c_double, # Bz
        ctypes.c_int,    # gravity_field
        ctypes.c_int,    # elastic_force
        ctypes.c_int,    # damping_force
        ctypes.c_double, # gravity_strength
        ctypes.c_int,    # n_drivers
        c_int_p,         # drv_idx     [n_drivers]
        c_dbl_p,         # drv_amp     [n_drivers*3]
        c_dbl_p,         # drv_freq    [n_drivers*3]
        c_dbl_p,         # drv_phi     [n_drivers*3]
        c_dbl_p,         # drv_eq      [n_drivers*3]
        c_dbl_p,         # drv_ncycles [n_drivers]
        c_int_p,         # drv_has_period [n_drivers]
        ctypes.c_int,    # n_frames
        c_dbl_p,         # out_x
        c_dbl_p,         # out_y
        c_dbl_p,         # out_z
        c_dbl_p,         # out_vx
        c_dbl_p,         # out_vy
        c_dbl_p,         # out_vz
        cb_type,         # progress_cb (可为 NULL)
    ]


def _c_dbl(arr: np.ndarray):
    """返回 float64 C 连续数组的 ctypes 指针。"""
    a = np.ascontiguousarray(arr, dtype=np.float64)
    return a.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), a


def _c_int(arr: np.ndarray):
    """返回 int32 C 连续数组的 ctypes 指针。"""
    a = np.ascontiguousarray(arr, dtype=np.int32)
    return a.ctypes.data_as(ctypes.POINTER(ctypes.c_int)), a


def _is_python_module(lib) -> bool:
    """判断 lib 是否为 Python 引擎模块（而非 ctypes.CDLL）。"""
    return not isinstance(lib, ctypes.CDLL)


def _run_dynamics_python(lib, config, atom_positions, atom_velocities, atom_masses,
                          atom_fixed, bond_pairs, bond_stiffness, bond_rest_lengths,
                          driver_data, atom_ids, progress_cb=None) -> dict:
    """调用 Python 引擎的 run_dynamics()，参数/返回值格式与 ctypes 版相同。"""
    n = len(atom_masses)
    NT          = int(config["NT"])
    NSTEP       = int(config.get("NSTEP", 1))
    warmup      = int(config.get("warmup_steps", 0))
    dt          = float(config["DT"])
    box_a       = float(config.get("box_a", 300.0))
    method_str  = str(config.get("method", "leapfrog")).lower().replace(" ", "_")
    method_id   = _METHOD_ID.get(method_str, 3)

    G = config.get("G", [0.0, 0.0, 0.0])
    B = config.get("B", [0.0, 0.0, 0.0])
    if hasattr(G, "tolist"): G = G.tolist()
    if hasattr(B, "tolist"): B = B.tolist()

    gravity_field    = int(config.get("gravity_field", 0))
    elastic_force    = int(config.get("elastic_force", 1))
    damping_force    = int(config.get("damping_force", 0))
    gravity_strength = float(config.get("gravity_strength", 1.0))

    record_steps = NT - warmup
    n_frames = max(1, record_steps // NSTEP)

    nd = len(driver_data) if driver_data else 0
    if nd > 0:
        drv_idx  = np.array([d["local_idx"] for d in driver_data], dtype=np.int64)
        drv_amp  = np.array([d["amp"]       for d in driver_data], dtype=np.float64)
        drv_freq = np.array([d["freq"]      for d in driver_data], dtype=np.float64)
        drv_phi  = np.array([d["phi"]       for d in driver_data], dtype=np.float64)
        drv_eq   = np.array([d["eq_pos"]    for d in driver_data], dtype=np.float64)
        drv_nc   = np.array([d["n_cycles"]  for d in driver_data], dtype=np.float64)
        drv_hp   = np.array([d["has_period"]for d in driver_data], dtype=np.int32)
    else:
        drv_idx = drv_amp = drv_freq = drv_phi = drv_eq = drv_nc = drv_hp = \
            np.zeros(0, dtype=np.int64)

    out_x, out_y, out_z, out_vx, out_vy, out_vz = lib.run_dynamics(
        n_atoms=n,
        pos_init=atom_positions,
        vel_init=atom_velocities,
        masses=atom_masses,
        fixed=atom_fixed,
        n_bonds=len(bond_pairs),
        bond_pairs=bond_pairs,
        bond_k=bond_stiffness,
        bond_r0=bond_rest_lengths,
        box_a=box_a,
        dt=dt,
        NT=NT,
        NSTEP=NSTEP,
        warmup_steps=warmup,
        method_id=method_id,
        Gx=float(G[0]), Gy=float(G[1]), Gz=float(G[2]),
        Bx=float(B[0]), By=float(B[1]), Bz=float(B[2]),
        gravity_field=gravity_field,
        elastic_force=elastic_force,
        damping_force=damping_force,
        gravity_strength=gravity_strength,
        n_drivers=nd,
        drv_idx=drv_idx,
        drv_amp=drv_amp,
        drv_freq=drv_freq,
        drv_phi=drv_phi,
        drv_eq=drv_eq,
        drv_ncycles=drv_nc,
        drv_has_period=drv_hp,
        n_frames=n_frames,
        progress_cb=progress_cb,
    )

    shape = (n_frames, n)
    t_arr = np.arange(n_frames) * NSTEP * dt + warmup * dt
    return {
        "x":  out_x.reshape(shape), "y":  out_y.reshape(shape),
        "z":  out_z.reshape(shape), "vx": out_vx.reshape(shape),
        "vy": out_vy.reshape(shape), "vz": out_vz.reshape(shape),
        "t":  t_arr,
    }


def run_dynamics_dll(
    lib,
    config: dict,
    atom_positions: np.ndarray,   # (n_atoms, 3)
    atom_velocities: np.ndarray,  # (n_atoms, 3)
    atom_masses: np.ndarray,      # (n_atoms,)
    atom_fixed: np.ndarray,       # (n_atoms, 3) int, 1=固定
    bond_pairs: np.ndarray,       # (n_bonds, 2) int 0-based 局部索引
    bond_stiffness: np.ndarray,   # (n_bonds,)
    bond_rest_lengths: np.ndarray,# (n_bonds,)
    driver_data: list,            # 驱动原子列表（见下文）
    atom_ids: np.ndarray,         # (n_atoms,) 全局 atom id（用于驱动原子查找）
    progress_cb=None,
) -> dict:
    """调用 DLL 的 run_dynamics()，返回抽帧后的轨迹数组。

    driver_data 格式（每个元素对应一个驱动原子）：
        {
          "atom_id":    int,          # 全局 atom id
          "local_idx":  int,          # 在 atom_ids 数组中的位置（0-based）
          "amp":        [ax, ay, az],
          "freq":       [fx, fy, fz],
          "phi":        [px, py, pz],
          "eq_pos":     [ex, ey, ez],
          "n_cycles":   float,        # 0=不限
          "has_period": int,          # 0/1
        }

    返回：
        {
          "x":  np.ndarray (n_frames, n_atoms),
          "y":  ...,
          "z":  ...,
          "vx": ..., "vy": ..., "vz": ...,
          "t":  np.ndarray (n_frames,),   # 时间轴
        }
    """
    # Python 引擎：直接调用模块函数，不走 ctypes
    if _is_python_module(lib):
        return _run_dynamics_python(
            lib, config, atom_positions, atom_velocities, atom_masses,
            atom_fixed, bond_pairs, bond_stiffness, bond_rest_lengths,
            driver_data, atom_ids, progress_cb)

    n = len(atom_masses)
    NT          = int(config["NT"])
    NSTEP       = int(config.get("NSTEP", 1))
    warmup      = int(config.get("warmup_steps", 0))
    dt          = float(config["DT"])
    box_a       = float(config.get("box_a", 300.0))
    method_str  = str(config.get("method", "leapfrog")).lower().replace(" ", "_")
    method_id   = _METHOD_ID.get(method_str, 3)

    G = config.get("G", [0.0, 0.0, 0.0])
    B = config.get("B", [0.0, 0.0, 0.0])
    if hasattr(G, "tolist"): G = G.tolist()
    if hasattr(B, "tolist"): B = B.tolist()

    gravity_field    = int(config.get("gravity_field", 0))
    elastic_force    = int(config.get("elastic_force", 1))
    damping_force    = int(config.get("damping_force", 0))
    gravity_strength = float(config.get("gravity_strength", 1.0))

    # ── 计算帧数 ──────────────────────────────────────────────
    record_steps = NT - warmup
    n_frames = max(1, record_steps // NSTEP)

    # ── 驱动原子数据 ──────────────────────────────────────────
    nd = len(driver_data) if driver_data else 0
    if nd > 0:
        drv_idx_arr  = np.array([d["local_idx"]  for d in driver_data], dtype=np.int32)
        drv_amp_arr  = np.array([d["amp"]         for d in driver_data], dtype=np.float64).ravel()
        drv_freq_arr = np.array([d["freq"]        for d in driver_data], dtype=np.float64).ravel()
        drv_phi_arr  = np.array([d["phi"]         for d in driver_data], dtype=np.float64).ravel()
        drv_eq_arr   = np.array([d["eq_pos"]      for d in driver_data], dtype=np.float64).ravel()
        drv_nc_arr   = np.array([d["n_cycles"]    for d in driver_data], dtype=np.float64)
        drv_hp_arr   = np.array([d["has_period"]  for d in driver_data], dtype=np.int32)
    else:
        drv_idx_arr  = np.zeros(1, dtype=np.int32)
        drv_amp_arr  = np.zeros(3, dtype=np.float64)
        drv_freq_arr = np.zeros(3, dtype=np.float64)
        drv_phi_arr  = np.zeros(3, dtype=np.float64)
        drv_eq_arr   = np.zeros(3, dtype=np.float64)
        drv_nc_arr   = np.zeros(1, dtype=np.float64)
        drv_hp_arr   = np.zeros(1, dtype=np.int32)

    # ── 输出缓冲区 ────────────────────────────────────────────
    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)

    # ── ctypes 指针（保留 arr 引用防止 GC） ──────────────────
    p_pos,   _pos   = _c_dbl(atom_positions.ravel())
    p_vel,   _vel   = _c_dbl(atom_velocities.ravel())
    p_mass,  _mass  = _c_dbl(atom_masses)
    p_fixed, _fixed = _c_int(atom_fixed.ravel())
    p_bp,    _bp    = _c_int(bond_pairs.ravel() if len(bond_pairs) else np.zeros(2, dtype=np.int32))
    p_bk,    _bk    = _c_dbl(bond_stiffness if len(bond_stiffness) else np.zeros(1))
    p_br0,   _br0   = _c_dbl(bond_rest_lengths if len(bond_rest_lengths) else np.zeros(1))
    p_didx,  _didx  = _c_int(drv_idx_arr)
    p_damp,  _damp  = _c_dbl(drv_amp_arr)
    p_dfrq,  _dfrq  = _c_dbl(drv_freq_arr)
    p_dphi,  _dphi  = _c_dbl(drv_phi_arr)
    p_deq,   _deq   = _c_dbl(drv_eq_arr)
    p_dnc,   _dnc   = _c_dbl(drv_nc_arr)
    p_dhp,   _dhp   = _c_int(drv_hp_arr)

    p_ox  = out_x.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
    p_oy  = out_y.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
    p_oz  = out_z.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
    p_ovx = out_vx.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
    p_ovy = out_vy.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
    p_ovz = out_vz.ctypes.data_as(ctypes.POINTER(ctypes.c_double))

    # 进度回调
    cb_type = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_int)
    if progress_cb is not None:
        cb = cb_type(progress_cb)
    else:
        cb = ctypes.cast(None, cb_type)

    ret = lib.run_dynamics(
        n,
        p_pos, p_vel, p_mass, p_fixed,
        len(bond_pairs), p_bp, p_bk, p_br0,
        box_a, dt, NT, NSTEP, warmup, method_id,
        float(G[0]), float(G[1]), float(G[2]),
        float(B[0]), float(B[1]), float(B[2]),
        gravity_field, elastic_force, damping_force, gravity_strength,
        nd, p_didx, p_damp, p_dfrq, p_dphi, p_deq, p_dnc, p_dhp,
        n_frames,
        p_ox, p_oy, p_oz, p_ovx, p_ovy, p_ovz,
        cb,
    )

    if ret != 0:
        raise RuntimeError(f"run_dynamics() returned error code {ret}")

    shape = (n_frames, n)
    t_arr = np.arange(n_frames) * NSTEP * dt + warmup * dt

    return {
        "x":  out_x.reshape(shape),
        "y":  out_y.reshape(shape),
        "z":  out_z.reshape(shape),
        "vx": out_vx.reshape(shape),
        "vy": out_vy.reshape(shape),
        "vz": out_vz.reshape(shape),
        "t":  t_arr,
    }


def is_dll_available(engine: str = "c") -> bool:
    """检查指定引擎是否可用（DLL 已编译或 Python 模块存在）。"""
    if _ENGINE_DIR.get(engine, engine) == "python":
        return os.path.exists(os.path.join(_HERE, "python", "dynamics_lib.py"))
    return any(os.path.exists(p) for p in _dll_candidates(engine))
