!pip install numba -q

import numba as nb
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML
import time

# ============================================================
# ПАРАМЕТРЫ (БОЛЬШЕ ТОЧЕК!)
# ============================================================
N = 10000
G = 1.0
dt = 0.002
steps = 400
eps = 0.05
C = 0.3  # скорость гравитации

# ============================================================
# JIT-ФУНКЦИИ
# ============================================================
@nb.njit(parallel=True, fastmath=True)
def update_nonlocal(pos, vel, G, dt, eps):
    N = pos.shape[0]
    acc = np.zeros_like(pos)
    eps2 = eps**2
    
    for i in nb.prange(N):
        ax = 0.0; ay = 0.0; az = 0.0
        xi = pos[i, 0]; yi = pos[i, 1]; zi = pos[i, 2]
        
        for j in range(N):
            dx = pos[j, 0] - xi
            dy = pos[j, 1] - yi
            dz = pos[j, 2] - zi
            r2 = dx*dx + dy*dy + dz*dz + eps2
            inv_r3 = 1.0 / (r2 * np.sqrt(r2))
            
            ax += dx * inv_r3
            ay += dy * inv_r3
            az += dz * inv_r3
            
        acc[i, 0] = G * ax
        acc[i, 1] = G * ay
        acc[i, 2] = G * az
        
    pos += vel * dt + 0.5 * acc * dt**2
    vel += acc * dt
    vel[:, 2] *= 0.995
    return pos, vel

@nb.njit(parallel=True, fastmath=True)
def update_local(pos, vel, history, step, G, dt, eps, C):
    N = pos.shape[0]
    acc = np.zeros_like(pos)
    eps2 = eps**2
    
    for i in nb.prange(N):
        ax = 0.0; ay = 0.0; az = 0.0
        xi = pos[i, 0]; yi = pos[i, 1]; zi = pos[i, 2]
        
        for j in range(N):
            dx = pos[j, 0] - xi
            dy = pos[j, 1] - yi
            dz = pos[j, 2] - zi
            r = np.sqrt(dx*dx + dy*dy + dz*dz + eps2)
            
            delay = int(r / C / dt)
            delay = max(1, min(delay, step))
            idx = max(0, step - delay)
            
            pj = history[idx][j]
            dx = pj[0] - xi
            dy = pj[1] - yi
            dz = pj[2] - zi
            r2 = dx*dx + dy*dy + dz*dz + eps2
            inv_r3 = 1.0 / (r2 * np.sqrt(r2))
            
            ax += dx * inv_r3
            ay += dy * inv_r3
            az += dz * inv_r3
            
        acc[i, 0] = G * ax
        acc[i, 1] = G * ay
        acc[i, 2] = G * az
        
    pos += vel * dt + 0.5 * acc * dt**2
    vel += acc * dt
    vel[:, 2] *= 0.995
    return pos, vel

# ============================================================
# НАЧАЛЬНЫЕ УСЛОВИЯ
# ============================================================
np.random.seed(42)

r = np.random.exponential(3, N)
theta = np.random.uniform(0, 2*np.pi, N)
z = np.random.normal(0, 0.2, N)

pos = np.zeros((N, 3))
pos[:, 0] = r * np.cos(theta)
pos[:, 1] = r * np.sin(theta)
pos[:, 2] = z

vel = np.zeros((N, 3))
v_circ = np.sqrt(G * N / r) * 0.8
vel[:, 0] = -pos[:, 1] / r * v_circ
vel[:, 1] = pos[:, 0] / r * v_circ

# ============================================================
# ЗАПУСК НЕЛОКАЛЬНОЙ
# ============================================================
print("Запуск НЕЛОКАЛЬНОЙ (Ньютон, 10k частиц)...")
pos_nl = pos.copy()
vel_nl = vel.copy()
history_nl = [pos_nl.copy()]
start_time = time.time()

for step in range(steps):
    pos_nl, vel_nl = update_nonlocal(pos_nl, vel_nl, G, dt, eps)
    if step % 5 == 0:
        history_nl.append(pos_nl.copy())

print(f"Готово за {time.time() - start_time:.2f} сек")

# ============================================================
# ЗАПУСК ЛОКАЛЬНОЙ
# ============================================================
print(f"Запуск ЛОКАЛЬНОЙ (C={C}, 10k частиц)...")
pos_l = pos.copy()
vel_l = vel.copy()
history_l = [pos_l.copy()]
history_all = [pos_l.copy()]
start_time = time.time()

for step in range(steps):
    pos_l, vel_l = update_local(pos_l, vel_l, history_all, step, G, dt, eps, C)
    history_all.append(pos_l.copy())
    if step % 5 == 0:
        history_l.append(pos_l.copy())

print(f"Готово за {time.time() - start_time:.2f} сек")

# ============================================================
# СОЗДАНИЕ ВИДЕО (ДВЕ ПАНЕЛИ)
# ============================================================
print("\nСоздание видео...")

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 7))

ax1.set_xlim(-12, 12); ax1.set_ylim(-12, 12)
ax1.set_aspect('equal')
ax1.set_facecolor('#0a0a1a')
ax1.set_title('НЕЛОКАЛЬНАЯ (Ньютон)', color='white', fontsize=14)
ax1.grid(True, alpha=0.1)

ax2.set_xlim(-12, 12); ax2.set_ylim(-12, 12)
ax2.set_aspect('equal')
ax2.set_facecolor('#0a0a1a')
ax2.set_title(f'ЛОКАЛЬНАЯ (C={C})', color='white', fontsize=14)
ax2.grid(True, alpha=0.1)

scat1 = ax1.scatter([], [], s=0.3, alpha=0.6, cmap='viridis', vmin=0, vmax=12)
scat2 = ax2.scatter([], [], s=0.3, alpha=0.6, cmap='viridis', vmin=0, vmax=12)

time_text = fig.text(0.02, 0.98, '', color='white', fontsize=12, verticalalignment='top')

def update(frame):
    pos_nl_frame = history_nl[frame]
    pos_l_frame = history_l[frame]
    r_nl = np.linalg.norm(pos_nl_frame, axis=1)
    r_l = np.linalg.norm(pos_l_frame, axis=1)
    
    scat1.set_offsets(pos_nl_frame[:, :2])
    scat1.set_array(r_nl)
    scat2.set_offsets(pos_l_frame[:, :2])
    scat2.set_array(r_l)
    
    time_text.set_text(f'Шаг: {frame * 5}')
    return scat1, scat2, time_text

ani = FuncAnimation(fig, update, frames=len(history_nl), interval=50, blit=True)
plt.close()
display(HTML(ani.to_html5_video()))

# ============================================================
# ФИНАЛЬНЫЕ КАДРЫ
# ============================================================
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(14, 14))

ax1.set_xlim(-12, 12); ax1.set_ylim(-12, 12)
ax1.set_aspect('equal')
ax1.set_title('НЕЛОКАЛЬНАЯ (финал)', fontsize=14)
ax1.scatter(history_nl[-1][:, 0], history_nl[-1][:, 1], s=0.3, alpha=0.5, color='blue')

ax2.set_xlim(-12, 12); ax2.set_ylim(-12, 12)
ax2.set_aspect('equal')
ax2.set_title(f'ЛОКАЛЬНАЯ (финал, C={C})', fontsize=14)
ax2.scatter(history_l[-1][:, 0], history_l[-1][:, 1], s=0.3, alpha=0.5, color='red')

ax3.set_xlim(-12, 12); ax3.set_ylim(-12, 12)
ax3.set_aspect('equal')
ax3.set_title('РАЗНИЦА (синий - красный)', fontsize=14)
diff = np.linalg.norm(history_nl[-1] - history_l[-1], axis=1)
ax3.scatter(history_nl[-1][:, 0], history_nl[-1][:, 1], c=diff, cmap='hot', s=0.3, alpha=0.6)

# Кривые вращения
r_nl = np.linalg.norm(history_nl[-1][:, :2], axis=1)
r_l = np.linalg.norm(history_l[-1][:, :2], axis=1)
v_nl = np.linalg.norm(vel_nl[:, :2], axis=1)
v_l = np.linalg.norm(vel_l[:, :2], axis=1)

bins = np.linspace(0, 10, 30)
r_bins_nl, v_bins_nl = [], []
r_bins_l, v_bins_l = [], []

for i in range(len(bins)-1):
    mask_nl = (r_nl > bins[i]) & (r_nl <= bins[i+1])
    mask_l = (r_l > bins[i]) & (r_l <= bins[i+1])
    if np.sum(mask_nl) > 20:
        r_bins_nl.append((bins[i] + bins[i+1]) / 2)
        v_bins_nl.append(np.mean(v_nl[mask_nl]))
    if np.sum(mask_l) > 20:
        r_bins_l.append((bins[i] + bins[i+1]) / 2)
        v_bins_l.append(np.mean(v_l[mask_l]))

ax4.plot(r_bins_nl, v_bins_nl, 'o-', color='blue', label='Нелокальная')
ax4.plot(r_bins_l, v_bins_l, 'o-', color='red', label='Локальная')
ax4.set_xlabel('Радиус')
ax4.set_ylabel('Скорость')
ax4.set_title('КРИВЫЕ ВРАЩЕНИЯ')
ax4.legend()
ax4.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

print("\n✅ Сравнение завершено!")