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

try:
    from numba import njit, prange, set_num_threads, get_num_threads
    NUMBA_AVAILABLE = True
except Exception:
    print("\nError: Failed to import numba")
    print("Install: pip install numba if needed\n")
    try:
        input("Run without numba")
    except EOFError:
        pass
    NUMBA_AVAILABLE = False

"""
Planet simulator with Numba-parallel acceleration (in-place accel):
- Integrators: Euler / Verlet / velocity-Verlet ("velet" accepted)
- O(N^2) force & energy via @njit(parallel=...)
- Generator yields every `yield_every` steps (default 10)
- Ctrl-C behavior selectable (stop/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         # 1: stop on Ctrl-C, 0: ignore
num_threads = 0          # 0 = numba default
parallel = True          # decorator-time flag (set before definitions)
#parallel = False          # decorator-time flag (set before definitions)

# ===================
# I/O 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)
    for d in rows:
        for k in list(d.keys()):
            if k != 'Name':
                d[k] = float(d[k])
    return rows

# ===================
# physics (pure NumPy fallbacks; will be used if numba is unavailable)
# ===================
def normalize_momentum_numpy(m, v):
    P = (m[:, None] * v).sum(axis=0)
    v -= P[None, :] / m.sum()

def pairwise_accel_numpy(m, r):
    D = r[None, :, :] - r[:, None, :]
    r2 = np.einsum('ijk,ijk->ij', D, D)
    np.fill_diagonal(r2, np.inf)
    inv_r3 = 1.0 / (r2 * np.sqrt(r2))
    W = (m[None, :] * inv_r3)[:, :, None]
    return G1 * np.sum(W * D, axis=1)

def potential_energy_numpy(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 kinetic_energy_numpy(m, v):
    return 0.5 * (m * (v*v).sum(axis=1)).sum()

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

# ===================
# numba-accelerated kernels (in-place force write)
# ===================
if NUMBA_AVAILABLE:
    @njit(parallel=parallel, fastmath=True, cache=True)
    def pairwise_accel_inplace_jit(m, r, a_out):
        """
        a_out に加速度を書き込む in-place 版。
        r2==0 の際は寄与をスキップしてゼロ割を回避（ウォームアップ等の安全策）。
        """
        N = r.shape[0]
        for i in prange(N):
            xi = r[i,0]; yi = r[i,1]; zi = r[i,2]
            ax = 0.0; ay = 0.0; az = 0.0
            for j in range(N):
                if i == j:
                    continue
                dx = r[j,0]-xi; dy = r[j,1]-yi; dz = r[j,2]-zi
                r2 = dx*dx + dy*dy + dz*dz
                if r2 == 0.0:
                    continue
                invr  = 1.0 / np.sqrt(r2)
                invr3 = invr * invr * invr
                s = G1 * m[j] * invr3
                ax += s*dx; ay += s*dy; az += s*dz
            a_out[i,0] = ax; a_out[i,1] = ay; a_out[i,2] = az

    @njit(fastmath=True, cache=True)
    def potential_energy_jit(m, r):
        N = r.shape[0]
        U = 0.0
        for i in range(N):
            xi = r[i,0]; yi = r[i,1]; zi = r[i,2]
            mi = m[i]
            for j in range(i+1, N):
                dx = r[j,0]-xi; dy = r[j,1]-yi; dz = r[j,2]-zi
                dist = np.sqrt(dx*dx + dy*dy + dz*dz)
                U += G1 * mi * m[j] / dist
        return U

    @njit(cache=True)
    def kinetic_energy_jit(m, v):
        K = 0.0
        for i in range(m.shape[0]):
            vi2 = v[i,0]*v[i,0] + v[i,1]*v[i,1] + v[i,2]*v[i,2]
            K += 0.5 * m[i] * vi2
        return K

    @njit(cache=True)
    def momentum_jit(m, v):
        Px=0.0; Py=0.0; Pz=0.0
        for i in range(m.shape[0]):
            mi = m[i]
            Px += mi*v[i,0]; Py += mi*v[i,1]; Pz += mi*v[i,2]
        P2 = Px*Px + Py*Py + Pz*Pz
        Pmsm = np.sqrt(P2/(3.0*max(1, m.shape[0])))
        return Px, Py, Pz, Pmsm

    @njit(cache=True)
    def normalize_momentum_jit(m, v):
        Px=0.0; Py=0.0; Pz=0.0
        Mtot = 0.0
        for i in range(m.shape[0]):
            mi = m[i]; Mtot += mi
            Px += mi*v[i,0]; Py += mi*v[i,1]; Pz += mi*v[i,2]
        for i in range(m.shape[0]):
            v[i,0] -= Px/Mtot
            v[i,1] -= Py/Mtot
            v[i,2] -= Pz/Mtot

    # --- JIT steppers: advance K steps in-place ---
    @njit(cache=True)
    def step_euler_k_jit(m, r, v, a, dt, k_steps):
        for _ in range(k_steps):
            v[:,0] += dt*a[:,0]; v[:,1] += dt*a[:,1]; v[:,2] += dt*a[:,2]
            r[:,0] += dt*v[:,0]; r[:,1] += dt*v[:,1]; r[:,2] += dt*v[:,2]
            pairwise_accel_inplace_jit(m, r, a)
        return a

    @njit(cache=True)
    def step_verlet_k_jit(m, r_prev, r, v, a, dt, k_steps):
        # r_prev, r, v, a updated in-place; v is estimated from positions
        for _ in range(k_steps):
            r_new_x = 2.0*r[:,0] - r_prev[:,0] + (dt*dt)*a[:,0]
            r_new_y = 2.0*r[:,1] - r_prev[:,1] + (dt*dt)*a[:,1]
            r_new_z = 2.0*r[:,2] - r_prev[:,2] + (dt*dt)*a[:,2]
            v[:,0] = (r_new_x - r_prev[:,0])/(2.0*dt)
            v[:,1] = (r_new_y - r_prev[:,1])/(2.0*dt)
            v[:,2] = (r_new_z - r_prev[:,2])/(2.0*dt)
            r_prev[:,0], r[:,0] = r[:,0], r_new_x
            r_prev[:,1], r[:,1] = r[:,1], r_new_y
            r_prev[:,2], r[:,2] = r[:,2], r_new_z
            pairwise_accel_inplace_jit(m, r, a)
        return a

    @njit(cache=True)
    def step_vverlet_k_jit(m, r, v, a, dt, k_steps):
        for _ in range(k_steps):
            v[:,0] += 0.5*dt*a[:,0]; v[:,1] += 0.5*dt*a[:,1]; v[:,2] += 0.5*dt*a[:,2]
            r[:,0] += dt*v[:,0];     r[:,1] += dt*v[:,1];     r[:,2] += dt*v[:,2]
            pairwise_accel_inplace_jit(m, r, a)
            v[:,0] += 0.5*dt*a[:,0]; v[:,1] += 0.5*dt*a[:,1]; v[:,2] += 0.5*dt*a[:,2]
        return a

# ===================
# initialization
# ===================
def initialize(planets):
    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))
    # initial states
    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)
    # momentum normalization & initial accel
    if NUMBA_AVAILABLE:
        normalize_momentum_jit(m, v)
        pairwise_accel_inplace_jit(m, r, a)
    else:
        normalize_momentum_numpy(m, v)
        a[:] = pairwise_accel_numpy(m, r)
    return names, m, r, v, a

# ===================
# generator (calls kernels K steps at a time)
# ===================
def md_generator(planets, solver, dt, nt, yield_every=10):
    s = (solver or '').lower()
    if s == 'velet':
        s = 'vverlet'

    names, m, r, v, a = initialize(planets)

    # t=0 yield
    if NUMBA_AVAILABLE:
        U = potential_energy_jit(m, r)
        K = kinetic_energy_jit(m, v)
        Px, Py, Pz, Pmsm = momentum_jit(m, v)
    else:
        U = potential_energy_numpy(m, r)
        K = kinetic_energy_numpy(m, v)
        Px, Py, Pz, Pmsm = momentum_numpy(m, v)
    yield (0.0, r, v, a, (U, K, U+K), (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 NUMBA_AVAILABLE:
            if s == 'euler':
                a = step_euler_k_jit(m, r, v, a, dt, k)
            elif s == 'verlet':
                a = step_verlet_k_jit(m, r_prev, r, v, a, dt, k)
            else:
                a = step_vverlet_k_jit(m, r, v, a, dt, k)
            U = potential_energy_jit(m, r)
            K = kinetic_energy_jit(m, v)
            Px, Py, Pz, Pmsm = momentum_jit(m, v)
        else:
            # NumPy fallback
            for _ in range(k):
                if s == 'euler':
                    a[:] = pairwise_accel_numpy(m, r)
                    v += dt * a
                    r += dt * v
                elif s == 'verlet':
                    a[:] = pairwise_accel_numpy(m, r)
                    r_new = 2.0*r - r_prev + (dt*dt) * a
                    v[:] = (r_new - r_prev) / (2.0*dt)
                    r_prev, r = r, r_new
                else:
                    a[:] = pairwise_accel_numpy(m, r)
                    v += 0.5*dt*a
                    r += dt*v
                    a[:] = pairwise_accel_numpy(m, r)
                    v += 0.5*dt*a
            U = potential_energy_numpy(m, r)
            K = kinetic_energy_numpy(m, v)
            Px, Py, Pz, Pmsm = momentum_numpy(m, v)

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

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

    # 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, 0: ignore
    if len(argv) >= 8: num_threads = int(argv[7])       # 0: default

    # Numba threads
    if NUMBA_AVAILABLE and num_threads > 0:
        try:
            set_num_threads(num_threads)
        except Exception:
            pass

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

    print("Planet simulator ({}), dt={}, nt={}, yield_every={}, Ctrl-C={}, Numba={}".format(
        solver, dt, nt, yield_every, "stop" if enable_ctrlc else "ignore",
        "ON" if NUMBA_AVAILABLE else "OFF"))
    if NUMBA_AVAILABLE:
        try:
            print(f"Numba threads: {get_num_threads()}")
        except Exception:
            pass
    print("G1 = {}".format(G1))
    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)
    ftraj = open(out_traj, 'w', newline='', encoding='utf-8')
    fcon  = open(out_cons, 'w', newline='', encoding='utf-8')
    wtraj = csv.writer(ftraj, lineterminator='\n')
    wcon  = csv.writer(fcon,  lineterminator='\n')

    # headers: t, x(name1), y(name1), x(name2), y(name2), ...
    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)
    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:
            xy = r[:, :2]
            # trajectory CSV (t, x1,y1, x2,y2, ...)
            row = [t]
            for i in range(xy.shape[0]):
                row.append(float(xy[i,0]))
                row.append(float(xy[i,1]))
            wtraj.writerow(row)
            # conservation CSV
            wcon.writerow([t, float(U), float(K), float(E), float(Px), float(Py), float(Pz), float(Pmsm)])

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

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

    finally:
        ftraj.close()
        fcon.close()

# ===================
# entry
# ===================
if __name__ == '__main__':
    # --- Warm up numba kernels (avoid 0-division by separating particles) ---
    start_jit = time.perf_counter()
    if NUMBA_AVAILABLE:
        _m = np.array([1.0, 1.0], dtype=np.float64)
        _r = np.zeros((2,3), dtype=np.float64)
        _r[1,0] = 1.0   # 1 AU (normed) separation along x to avoid r2==0
        _v = np.zeros((2,3), dtype=np.float64)
        _a = np.zeros((2,3), dtype=np.float64)
        pairwise_accel_inplace_jit(_m, _r, _a)      # force kernels compile (in-place)
        _ = potential_energy_jit(_m, _r)
        _ = kinetic_energy_jit(_m, _v)

    # --- Timed run ---
    start = time.perf_counter()
    main()
    end = time.perf_counter()
    print(f"\nElapsed time:")
    print(f"   JIT first run   : {start - start_jit:.3f} seconds")
    print(f"   Calculation+draw: {end - start:.3f} seconds")

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