import sys
import csv
import numpy as np
import matplotlib.pyplot as plt
import signal
import time
import os
import ctypes as ct

"""
Planet simulator (NumPy + C DLL): Euler / Verlet / velocity-Verlet with generator
- yield_every=10 (default)
- Ctrl-C behavior selectable: enable_ctrlc=1 (default) to stop, 0 to ignore
- Trails are lines only (no markers)
"""

# ===================
# constants
# ===================
G = 6.67259e-11                       # Nm2/kg2
DayToSecond      = 60 * 60 * 24       # s
AstronomicalUnit = 1.49597870e11      # m
AU               = AstronomicalUnit
G1 = G * DayToSecond * DayToSecond / AU / AU / AU

# ===================
# defaults / params
# ===================
solver = 'velet'      # 'Euler' | 'Verlet' | 'velet' (= velocity-Verlet)
fplot = 1
dbfile   = 'planet_db.csv'
dt = 0.1
nt = 20000
iprint_interval    = 100
nprint_planets     = 4
xgrange = (-5.0, 5.0)
ygrange = (-5.0, 5.0)
yield_every = 10         # ★ default 10
enable_ctrlc = 1         # ★ Ctrl-C で停止(1) / 無視(0)

# ===================
# C DLL loader via ctypes
# ===================
def _default_lib_path():
    here = os.path.dirname(os.path.abspath(__file__))
    if os.name == 'nt':
        cand = os.path.join(here, 'nbody.dll')
    elif sys.platform == 'darwin':
        cand = os.path.join(here, 'libnbody.dylib')
    else:
        cand = os.path.join(here, 'libnbody.so')
    return cand

class NBodyCLib:
    def __init__(self, path=None):
        self.path = path or _default_lib_path()
        self.lib = ct.CDLL(self.path)

        # 型定義
        self.lib.accel_inplace.argtypes = [
            ct.POINTER(ct.c_double),  # m[N]
            ct.POINTER(ct.c_double),  # r[N*3]
            ct.POINTER(ct.c_double),  # a[N*3]
            ct.c_int,                 # N
            ct.c_double               # G1
        ]
        self.lib.accel_inplace.restype = None

        self.lib.step_euler_k.argtypes = [
            ct.POINTER(ct.c_double),  # m
            ct.POINTER(ct.c_double),  # r
            ct.POINTER(ct.c_double),  # v
            ct.POINTER(ct.c_double),  # a
            ct.c_double, ct.c_int, ct.c_int, ct.c_double
        ]
        self.lib.step_euler_k.restype = None

        self.lib.step_verlet_k.argtypes = [
            ct.POINTER(ct.c_double),  # m
            ct.POINTER(ct.c_double),  # r_prev
            ct.POINTER(ct.c_double),  # r
            ct.POINTER(ct.c_double),  # v
            ct.POINTER(ct.c_double),  # a
            ct.c_double, ct.c_int, ct.c_int, ct.c_double
        ]
        self.lib.step_verlet_k.restype = None

        self.lib.step_vverlet_k.argtypes = [
            ct.POINTER(ct.c_double),  # m
            ct.POINTER(ct.c_double),  # r
            ct.POINTER(ct.c_double),  # v
            ct.POINTER(ct.c_double),  # a
            ct.c_double, ct.c_int, ct.c_int, ct.c_double
        ]
        self.lib.step_vverlet_k.restype = None

        self.lib.kinetic_energy.argtypes = [
            ct.POINTER(ct.c_double),  # m
            ct.POINTER(ct.c_double),  # v
            ct.c_int
        ]
        self.lib.kinetic_energy.restype = ct.c_double

        self.lib.potential_energy.argtypes = [
            ct.POINTER(ct.c_double),  # m
            ct.POINTER(ct.c_double),  # r
            ct.c_int, ct.c_double
        ]
        self.lib.potential_energy.restype = ct.c_double

        self.lib.momentum.argtypes = [
            ct.POINTER(ct.c_double),  # m
            ct.POINTER(ct.c_double),  # v
            ct.c_int,
            ct.POINTER(ct.c_double), ct.POINTER(ct.c_double),
            ct.POINTER(ct.c_double), ct.POINTER(ct.c_double)
        ]
        self.lib.momentum.restype = None

    @staticmethod
    def _ptr(a: np.ndarray):
        return a.ctypes.data_as(ct.POINTER(ct.c_double))

    def accel_inplace(self, m, r, a):
        N = r.shape[0]
        self.lib.accel_inplace(self._ptr(m), self._ptr(r), self._ptr(a), ct.c_int(N), ct.c_double(G1))

    def step_euler_k(self, m, r, v, a, dt, k):
        N = r.shape[0]
        self.lib.step_euler_k(self._ptr(m), self._ptr(r), self._ptr(v), self._ptr(a),
                              ct.c_double(dt), ct.c_int(k), ct.c_int(N), ct.c_double(G1))

    def step_verlet_k(self, m, r_prev, r, v, a, dt, k):
        N = r.shape[0]
        self.lib.step_verlet_k(self._ptr(m), self._ptr(r_prev), self._ptr(r), self._ptr(v), self._ptr(a),
                               ct.c_double(dt), ct.c_int(k), ct.c_int(N), ct.c_double(G1))

    def step_vverlet_k(self, m, r, v, a, dt, k):
        N = r.shape[0]
        self.lib.step_vverlet_k(self._ptr(m), self._ptr(r), self._ptr(v), self._ptr(a),
                                ct.c_double(dt), ct.c_int(k), ct.c_int(N), ct.c_double(G1))

    def kinetic_energy(self, m, v):
        N = v.shape[0]
        return float(self.lib.kinetic_energy(self._ptr(m), self._ptr(v), ct.c_int(N)))

    def potential_energy(self, m, r):
        N = r.shape[0]
        return float(self.lib.potential_energy(self._ptr(m), self._ptr(r), ct.c_int(N), ct.c_double(G1)))

    def momentum(self, m, v):
        N = v.shape[0]
        Px = ct.c_double(); Py = ct.c_double(); Pz = ct.c_double(); Pmsm = ct.c_double()
        self.lib.momentum(self._ptr(m), self._ptr(v), ct.c_int(N),
                          ct.byref(Px), ct.byref(Py), ct.byref(Pz), ct.byref(Pmsm))
        return Px.value, Py.value, Pz.value, Pmsm.value

# ===================
# helpers
# ===================
def readdb(dbfile):
    rows = []
    with open(dbfile, "r", newline='', encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            rows.append(row)
    # cast
    for d in rows:
        for k in list(d.keys()):
            if k != 'Name':
                d[k] = float(d[k])
    return rows

def normalize_momentum(m, v):
    P = (m[:, None] * v).sum(axis=0)
    v -= P[None, :] / m.sum()

def kinetic_energy_py(m, v):  # 参照: C と同じ
    return 0.5 * (m * (v*v).sum(axis=1)).sum()

def potential_energy_py(m, r):
    D = r[None, :, :] - r[:, None, :]
    r2 = np.einsum('ijk,ijk->ij', D, D)
    N = r.shape[0]
    iu = np.triu_indices(N, k=1)
    dist = np.sqrt(r2[iu])
    mi_mj = (m[iu[0]] * m[iu[1]])
    return (G1 * mi_mj / dist).sum()

def total_momentum_py(m, v):
    P = (m[:, None] * v).sum(axis=0)
    Pmsm = np.sqrt((P*P).sum() / (3.0 * max(1, len(m))))
    return P[0], P[1], P[2], Pmsm

def initialize(planets, clib: NBodyCLib):
    names = np.array([p['Name'] for p in planets], dtype=object)
    m  = np.ascontiguousarray(np.array([p['Mass'] for p in planets], dtype=np.float64))
    r  = np.ascontiguousarray(np.zeros((len(planets), 3), dtype=np.float64))
    v  = np.ascontiguousarray(np.zeros((len(planets), 3), dtype=np.float64))
    a  = np.ascontiguousarray(np.zeros((len(planets), 3), dtype=np.float64))

    r[:, 0] = np.array([p['Revolution Radius'] / AU for p in planets], dtype=np.float64)
    v[:, 1] = np.array([p['Revolution Velocity'] * DayToSecond / AU for p in planets], dtype=np.float64)

    normalize_momentum(m, v)
    clib.accel_inplace(m, r, a)  # 初期加速度
    return names, m, r, v, a

def md_generator(planets, solver, dt, nt, yield_every=1, libpath=None):
    """
    C DLL を呼ぶジェネレータ。
    Yield: (t, r(N,3), v(N,3), a(N,3), (U,K,E), (Px,Py,Pz,Pmsm))
    """
    s = (solver or '').lower()
    if s == 'velet':
        s = 'vverlet'

    clib = NBodyCLib(libpath)
    names, m, r, v, a = initialize(planets, clib)

    # first yield at t=0
    U = clib.potential_energy(m, r)
    K = clib.kinetic_energy(m, v)
    E = U + K
    Px, Py, Pz, Pmsm = clib.momentum(m, v)
    yield (0.0, r, v, a, (U, K, E), (Px, Py, Pz, Pmsm))

    # for position-Verlet
    if s == 'verlet':
        r_prev = np.ascontiguousarray(r - dt * v)

    steps_done = 0
    while steps_done < nt:
        k = min(yield_every, nt - steps_done)
        if s == 'euler':
            clib.step_euler_k(m, r, v, a, dt, k)
        elif s == 'verlet':
            clib.step_verlet_k(m, r_prev, r, v, a, dt, k)
        else:
            clib.step_vverlet_k(m, r, v, a, dt, k)

        U = clib.potential_energy(m, r)
        K = clib.kinetic_energy(m, v)
        E = U + K
        Px, Py, Pz, Pmsm = clib.momentum(m, v)

        steps_done += k
        t = steps_done * dt
        yield (t, r, v, a, (U, K, E), (Px, Py, Pz, Pmsm))

# ===================
# main
# ===================
def main():
    global solver, dt, nt, fplot, yield_every, enable_ctrlc

    # CLI
    argv = sys.argv
    if len(argv) >= 2: solver = argv[1]
    if len(argv) >= 3: dt = float(argv[2])
    if len(argv) >= 4: nt = int(argv[3])
    if len(argv) >= 5: fplot = int(argv[4])
    if len(argv) >= 6: yield_every = int(argv[5])        # ★
    if len(argv) >= 7: enable_ctrlc = int(argv[6])       # ★ 1: stop on Ctrl-C, 0: ignore

    # Ctrl-C 挙動
    if enable_ctrlc:
        signal.signal(signal.SIGINT, signal.default_int_handler)
    else:
        signal.signal(signal.SIGINT, signal.SIG_IGN)

    print("Planet simulator C-DLL ({})".format(solver))
    print("G1 = {}".format(G1))
    print("dt = {}, nt = {}, yield_every = {}, Ctrl-C={}".format(dt, nt, yield_every, "stop" if enable_ctrlc else "ignore"))
    print("")

    planets = readdb(dbfile)
    names = [p['Name'] for p in planets]
    print("Planets:", ", ".join(names))

    out_traj  = "diffeq2nd_Planet_{}.csv".format(solver)
    out_cons  = "diffeq2nd_Planet_{}_conservation.csv".format(solver)
    wtraj = csv.writer(open(out_traj, 'w', newline='', encoding='utf-8'), lineterminator='\n')
    wcon  = csv.writer(open(out_cons, 'w', newline='', encoding='utf-8'),  lineterminator='\n')

    # headers
    lab = ['t'] + sum(([f"x({nm})", f"y({nm})"] for nm in names), [])
    wtraj.writerow(lab)
    wcon.writerow(['t','U','K','E','Px','Py','Pz','Pmsm'])

    # plot setup (lines only, no markers)
    if fplot:
        fig, ax = plt.subplots(1,1)
        ax.set_aspect('equal', adjustable='box')
        ax.set_xlim(xgrange); ax.set_ylim(ygrange)
        ax.set_xlabel('x [AU]'); ax.set_ylabel('y [AU]')
        nshow = min(7, len(names))
        trails_x = [[] for _ in range(nshow)]
        trails_y = [[] for _ in range(nshow)]
        lines = [ax.plot([], [], linewidth=0.6)[0] for _ in range(nshow)]

    gen = md_generator(planets, solver, dt, nt, yield_every=yield_every)

    print("{:^7}".format('t'), end='')
    for i in range(0, min(nprint_planets, len(names))):
        print("  {:^12}  {:^12}".format(f"x({names[i]})", f"y({names[i]})"), end='')
    print("")

    try:
        for (t, r, v, a, (U,K,E), (Px,Py,Pz,Pmsm)) in gen:
            # CSV trajectory
            row = [t]
            xy = r[:, :2]
            row = [t]
            for i in range(xy.shape[0]):
                row.append(float(xy[i,0]))
                row.append(float(xy[i,1]))
            wtraj.writerow(row)

            # CSV conservation
            wcon.writerow([t, U, K, E, Px, Py, Pz, Pmsm])

            # print
            istep = int(round(t/dt))
            if (istep % iprint_interval) == 0:
                print(f"{t:^7.2f}", end='')
                for i in range(0, min(nprint_planets, xy.shape[0])):
                    print("  {:>12.4g}  {:>12.4g}".format(xy[i,0], xy[i,1]), end='')
                print("")

            # plot
            if fplot:
                nshow = min(len(lines), xy.shape[0])
                for i in range(nshow):
                    trails_x[i].append(xy[i,0])
                    trails_y[i].append(xy[i,1])
                    if len(trails_x[i]) > 2000:
                        trails_x[i] = trails_x[i][-1000:]
                        trails_y[i] = trails_y[i][-1000:]
                    lines[i].set_data(trails_x[i], trails_y[i])
                ax.set_xlim(xgrange); ax.set_ylim(ygrange)
                plt.pause(1e-10)

    except KeyboardInterrupt:
        if enable_ctrlc:
            print("\nInterrupted by user (Ctrl-C).")

if __name__ == '__main__':
    start = time.perf_counter()
    main()
    end = time.perf_counter()
    print(f"\nElapsed time: {end - start:.3f} seconds")

    print("Press ENTER to exit>>", end='')
    try:
        input()
    except EOFError:
        pass
