# ============================================================
# Planet simulator (NumPy): Euler / Heun / Verlet / velocity-Verlet
# ============================================================

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

matplotlib.use("TkAgg")   # または "QtAgg"


# ===================
# 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 = 'vverlet'      # 'Euler' | 'Heun' | 'Verlet' | 'vverlet'
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
enable_ctrlc = 1

# ===================
# 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

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

def kinetic_energy(m, v):
    return 0.5 * (m * (v*v).sum(axis=1)).sum()

def pairwise_accel(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]
    a = G1 * np.sum(W * D, axis=1)
    return a

def potential_energy(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(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):
    names = np.array([p['Name'] for p in planets], dtype=object)
    m  = np.array([p['Mass'] for p in planets], dtype=float)
    r  = np.zeros((len(planets), 3), dtype=float)
    v  = np.zeros((len(planets), 3), dtype=float)

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

    normalize_momentum(m, v)
    a = pairwise_accel(m, r)
    return names, m, r, v, a

# ===================
# generator
# ===================
def md_generator(planets, solver, dt, nt, yield_every=1):
    s = (solver or '').lower()
    if s == 'velet':
        s = 'vverlet'

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

    U = potential_energy(m, r)
    K = kinetic_energy(m, v)
    E = U + K
    Px, Py, Pz, Pmsm = total_momentum(m, v)
    yield (0.0, r, v, a, (U, K, E), (Px, Py, Pz, Pmsm))

    if s == 'verlet':
        r_prev = r - dt * v

    for it in range(1, nt+1):
        t = it * dt

        if s == 'euler':
            v += dt * a
            r += dt * v
            a = pairwise_accel(m, r)

        elif s == 'heun':
            v_pred = v + dt * a
            r_pred = r + dt * v
            a_pred = pairwise_accel(m, r_pred)
            v += 0.5 * dt * (a + a_pred)
            r += 0.5 * dt * (v + v_pred)
            a = pairwise_accel(m, r)

        elif s == 'verlet':
            r_new = 2.0 * r - r_prev + (dt*dt) * a
            v[:] = (r_new - r_prev) / (2.0 * dt)
            r_prev, r = r, r_new
            a = pairwise_accel(m, r)

        else:
            v += 0.5 * dt * a
            r += dt * v
            a_new = pairwise_accel(m, r)
            v += 0.5 * dt * a_new
            a = a_new

        if (it % yield_every) == 0:
            U = potential_energy(m, r)
            K = kinetic_energy(m, v)
            E = U + K
            Px, Py, Pz, Pmsm = total_momentum(m, v)
            yield (t, r, v, a, (U, K, E), (Px, Py, Pz, Pmsm))

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

    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])

    if enable_ctrlc:
        signal.signal(signal.SIGINT, signal.default_int_handler)
    else:
        signal.signal(signal.SIGINT, signal.SIG_IGN)

    print("Planet simulator NumPy ({})".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')

    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'])

    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:
            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)

            wcon.writerow([t, U, K, E, Px, Py, Pz, 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(0.01)

    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