"""
engines/python/main.py
-----------------------
独立 Python 计算引擎。

与 main.c / main.cpp / main.f90 结构一致：
  输入: <input_dir>/coord.txt, connection.txt, bond.txt, [driver.txt]
        <param_json> (同 engines/c/param.json 格式)
  输出: <output_dir>/display.txt  (+ display.npz)
        <output_dir>/trajectory.txt (若 save_trajectory=1)

用法:
  python main.py <input_dir> <output_dir> <param_json>

内部调用 dynamics_lib.run_dynamics()，算法与 compute.py 完全一致。
"""

import json
import os
import sys
import time
import numpy as np

# 将父目录（engines/python 的上级 engines）加入 sys.path，
# 以便在独立运行时也能找到 dynamics_lib
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)

from dynamics_lib import run_dynamics

# 为读取 coord/bond/display，复用 compute.py 中的 I/O 函数
_COMPUTE = os.path.join(_HERE, "..", "..")
sys.path.insert(0, _COMPUTE)
import compute as _c

_METHOD_ID = {
    "explicit_euler": 0,
    "euler":          0,
    "implicit_euler": 1,
    "midpoint":       2,
    "leapfrog":       3,
}


def _load_params(param_path):
    """读取 param.json（与 C 引擎格式相同）."""
    with open(param_path, "r", encoding="utf-8") as f:
        p = json.load(f)
    return p


def main():
    if len(sys.argv) < 4:
        print("用法: python main.py <input_dir> <output_dir> <param_json>")
        sys.exit(1)

    input_dir  = sys.argv[1]
    output_dir = sys.argv[2]
    param_path = sys.argv[3]

    os.makedirs(output_dir, exist_ok=True)

    # ── 读取参数 ─────────────────────────────────────────────
    p = _load_params(param_path)
    box_a         = float(p.get("box_a", 10.0))
    NT            = int(p.get("NT", 10000))
    dt            = float(p.get("DT", 0.001))
    NSTEP         = int(p.get("NSTEP", 100))
    warmup_steps  = int(p.get("warmup_steps", 0))
    method_str    = str(p.get("method", "leapfrog")).lower().replace(" ", "_")
    method_id     = _METHOD_ID.get(method_str, 3)
    G             = p.get("G", [0.0, 0.0, -9.8])
    B             = p.get("B", [0.0, 0.0,  0.0])
    gravity_field    = int(p.get("gravity_field", 1))
    elastic_force    = int(p.get("elastic_force", 1))
    damping_force    = int(p.get("damping_force", 0))
    gravity_strength = float(p.get("gravity_strength", 1.0))
    driving_force    = int(p.get("driving_force", 0))
    save_traj        = int(p.get("save_trajectory", 0))

    # ── 读取原子数据 ──────────────────────────────────────────
    coord_path = os.path.join(input_dir, "coord.txt")
    atom_ids, masses, radii, positions, velocities, fixed = _c.load_coord_file(coord_path)

    # ── 读取键数据 ────────────────────────────────────────────
    conn_path = os.path.join(input_dir, "connection.txt")
    bond_path = os.path.join(input_dir, "bond.txt")
    bond_map  = _c.load_bond_parameters(bond_path)
    bond_pairs, bond_names, bond_stiffness, bond_rest_lengths = \
        _c.load_bond_connections(conn_path, atom_ids, positions, bond_map)
    n_bonds = len(bond_pairs)

    # ── 读取驱动力 ────────────────────────────────────────────
    drv_list = []
    if driving_force:
        driver_path = os.path.join(input_dir, "driver.txt")
        raw_drivers = _c.load_driver_file(driver_path, atom_ids)
        if raw_drivers:
            atom_id_map = {int(aid): i for i, aid in enumerate(atom_ids)}
            for d in raw_drivers:
                aid = int(d["atom_id"])
                if aid not in atom_id_map:
                    continue
                lidx = atom_id_map[aid]
                eq   = positions[lidx].tolist()
                d["eq_pos"] = np.array(eq)
                pc   = d.get("period_cycles")
                nc   = float(pc) if pc is not None else 0.0
                hp   = 1 if nc > 0 else 0
                drv_list.append({
                    "local_idx": lidx,
                    "amp":  d["amp"].tolist(),
                    "freq": d["freq"].tolist(),
                    "phi":  d["phi"].tolist(),   # radians
                    "eq":   eq,
                    "nc":   nc,
                    "hp":   hp,
                })

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

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

    # ── 进度回调 ──────────────────────────────────────────────
    def _progress(step, total):
        pct = step * 100 // total
        print(f"[python-engine] progress: {step}/{total} ({pct}%)", flush=True)

    # ── 运行计算 ──────────────────────────────────────────────
    t0 = time.time()
    print(f"[python-engine] NT={NT}  NSTEP={NSTEP}  method={method_str}  "
          f"n_atoms={len(atom_ids)}  n_bonds={n_bonds}")

    out_x, out_y, out_z, out_vx, out_vy, out_vz = run_dynamics(
        n_atoms=len(atom_ids),
        pos_init=positions,
        vel_init=velocities,
        masses=masses,
        fixed=fixed,
        n_bonds=n_bonds,
        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_steps,
        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,
    )
    elapsed = time.time() - t0
    print(f"[python-engine] 完成: {n_frames} 帧  {elapsed:.3f} s")

    # ── 构建 display header ───────────────────────────────────
    ball_radius    = float(p.get("ball_radius",  0.5))
    ball_color     = p.get("ball_color",   [0.9, 0.2, 0.2])
    box_color      = p.get("box_color",    [0.8, 0.8, 0.85])
    use_marker     = int(p.get("use_marker",     0))
    alpha_val      = p.get("alpha",        0.2)
    cam_dist       = float(p.get("camera_distance",   40.0))
    cam_elev       = float(p.get("camera_elevation",   0.0))
    cam_azim       = float(p.get("camera_azimuth",     0.0))
    cam_cx         = float(p.get("camera_center_x",    0.0))
    cam_cy         = float(p.get("camera_center_y",    0.0))
    cam_cz         = float(p.get("camera_center_z",    0.0))

    header = {
        "DT":                  str(dt),
        "NSTEP":               str(NSTEP),
        "method":              method_str,
        "NT":                  str(NT),
        "warmup_steps":        str(warmup_steps),
        "dynamic_steps":       str(record_steps),
        "T_total":             str(NT * dt),
        "box_a":               str(box_a),
        "gravity_field":       str(gravity_field),
        "elastic_force":       str(elastic_force),
        "damping_force":       str(damping_force),
        "driving_force":       str(driving_force),
        "gravity_strength":    str(gravity_strength),
        "G":                   json.dumps([float(v) for v in G]),
        "B":                   json.dumps([float(v) for v in B]),
        "number_of_frames":    str(n_frames),
        "number_of_particles": str(len(atom_ids)),
        "use_marker":          str(use_marker),
        "display_color":       json.dumps(p.get("display_color",
            {"x":[0,[1.0,0.0,0.0]],"y":[0,[0.0,1.0,0.0]],"z":[0,[0.0,0.0,1.0]],
             "xy":[0,[1.0,1.0,0.0]],"yz":[0,[0.0,1.0,1.0]],"zx":[0,[1.0,0.0,1.0]],
             "xyz":[1,[1.0,1.0,1.0]]})),
        "ball_radius":         str(ball_radius),
        "ball_color_r":        str(ball_color[0]),
        "ball_color_g":        str(ball_color[1]),
        "ball_color_b":        str(ball_color[2]),
        "box_color_r":         str(box_color[0]),
        "box_color_g":         str(box_color[1]),
        "box_color_b":         str(box_color[2]),
        "alpha":               str(alpha_val) if not isinstance(alpha_val, list)
                               else ",".join(str(a) for a in alpha_val),
        "atom_radii":          ",".join(str(r) for r in radii),
        "atom_masses":         json.dumps([float(m) for m in masses]),
        "atom_positions":      json.dumps(positions.tolist()),
        "bond_pairs":          json.dumps(bond_pairs.tolist() if n_bonds else []),
        "bond_stiffness":      json.dumps(bond_stiffness.tolist() if n_bonds else []),
        "bond_rest_lengths":   json.dumps(bond_rest_lengths.tolist() if n_bonds else []),
        "X_MIN": str(-box_a), "X_MAX": str(box_a),
        "Y_MIN": str(-box_a), "Y_MAX": str(box_a),
        "Z_MIN": str(-box_a), "Z_MAX": str(box_a),
        "camera_distance":  str(cam_dist),
        "camera_elevation": str(cam_elev),
        "camera_azimuth":   str(cam_azim),
        "camera_center_x":  str(cam_cx),
        "camera_center_y":  str(cam_cy),
        "camera_center_z":  str(cam_cz),
        "camera_keyframes": "",
    }

    # ── 保存 display.txt + display.npz ───────────────────────
    disp_txt = os.path.join(output_dir, "display.txt")
    _c.save_display_txt(
        disp_txt,
        out_x, out_y, out_z, out_vx, out_vy, out_vz,
        atom_ids, record_steps, len(atom_ids),
        header_fields=header,
    )
    print(f"[python-engine] display.txt 已保存: {disp_txt}")

    disp_npz = os.path.join(output_dir, "display.npz")
    _c.save_display_npz(
        disp_npz,
        out_x, out_y, out_z, out_vx, out_vy, out_vz,
        atom_ids, header_fields=header,
    )
    print(f"[python-engine] display.npz 已保存: {disp_npz}")

    # ── 可选：保存 trajectory.txt ─────────────────────────────
    if save_traj:
        traj_payload = {
            "traj_x":  out_x, "traj_y":  out_y, "traj_z":  out_z,
            "traj_vx": out_vx, "traj_vy": out_vy, "traj_vz": out_vz,
            "NT": record_steps, "DT": dt, "NSTEP": NSTEP,
            "method": method_str,
            "atom_ids": atom_ids,
            "atom_masses": masses,
            "atom_radii": radii,
            "atom_positions": positions,
            "bond_pairs": bond_pairs,
            "bond_stiffness": bond_stiffness,
            "bond_rest_lengths": bond_rest_lengths,
            "G": [float(v) for v in G],
            "B": [float(v) for v in B],
        }
        traj_path = os.path.join(output_dir, "trajectory.txt")
        _c.save_text_data(traj_path, traj_payload)
        print(f"[python-engine] trajectory.txt 已保存: {traj_path}")


if __name__ == "__main__":
    main()
