import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

class BlackHoleMeshSim:
    def __init__(self, grid_size=12):
        self.grid_size = grid_size
        
        # 1. Initial Grid
        x = np.linspace(-6, 6, grid_size)
        y = np.linspace(-6, 6, grid_size)
        X, Y = np.meshgrid(x, y)
        self.pos = np.stack([X.flatten(), Y.flatten()], axis=-1).astype(float)
        self.orig_pos = self.pos.copy()
        self.velocity = np.zeros_like(self.pos)
        
        # The Central Truth Singularity
        self.center = np.array([0.0, 0.0])
        
        # The "Matter" (New Information) Particles
        self.matter_particles = [] # List of [x, y, mass]
        self.matter_timer = 0

    def spawn_matter(self):
        # Randomly spawn a piece of "information matter" at the edge
        angle = np.random.uniform(0, 2*np.pi)
        dist = 6.0
        pos = np.array([np.cos(angle)*dist, np.sin(angle)*dist])
        mass = np.random.uniform(0.2, 0.8)
        self.matter_particles.append({'pos': pos, 'mass': mass, 'vel': -0.05 * pos})

    def compute_forces(self):
        curr_pos = self.pos
        total_force = np.zeros_like(curr_pos)
        
        # A. Central Singularity Pull
        diff_c = self.center - curr_pos
        dist_c = np.linalg.norm(diff_c, axis=1, keepdims=True)
        force_c = (0.6 / (dist_c**2 + 0.5)) * (diff_c / (dist_c + 1e-5))
        total_force += force_c
        
        # B. Matter Particle Pull (The Perturbations)
        for m in self.matter_particles:
            diff_m = m['pos'] - curr_pos
            dist_m = np.linalg.norm(diff_m, axis=1, keepdims=True)
            force_m = (m['mass'] / (dist_m**2 + 0.5)) * (diff_m / (dist_m + 1e-5))
            total_force += force_m
            
        # C. Stability (Rubber Sheet Tension)
        force_s = 0.04 * (self.orig_pos - curr_pos)
        total_force += force_s
        
        return total_force

    def update(self, dt=0.1):
        # Update Matter Particles (they fall toward the center)
        for m in self.matter_particles:
            dir_to_center = self.center - m['pos']
            m['pos'] += dir_to_center * 0.01 # Slow drift inward
            
        # Update Mesh
        forces = self.compute_forces()
        self.velocity = (self.velocity + forces * dt) * 0.9
        self.pos += self.velocity * dt

# --- Setup ---
sim = BlackHoleMeshSim()
fig, ax = plt.subplots(figsize=(8, 8))
ax.set_xlim(-7, 7); ax.set_ylim(-7, 7)
ax.set_aspect('equal'); ax.set_facecolor('#0a0a0a')
fig.patch.set_facecolor('#0a0a0a')
ax.axis('off')

# Elements
center_dot, = ax.plot(0, 0, 'yo', markersize=8, label="Truth Singularity", zorder=5)
points, = ax.plot([], [], 'wo', markersize=2, alpha=0.5, zorder=4)
matter_dots, = ax.plot([], [], 'ro', markersize=5, label="Information Matter", zorder=6)
lines = []

edges = []
for i in range(sim.grid_size):
    for j in range(sim.grid_size):
        idx = i * sim.grid_size + j
        if j < sim.grid_size - 1: edges.append((idx, idx + 1))
        if i < sim.grid_size - 1: edges.append((idx, idx + sim.grid_size))

for _ in edges:
    l, = ax.plot([], [], 'w-', alpha=0.15, linewidth=0.5, zorder=3)
    lines.append(l)

def animate(frame):
    if frame % 50 == 0: sim.spawn_matter()
    sim.update()
    
    points.set_data(sim.pos[:, 0], sim.pos[:, 1])
    
    m_pos = np.array([m['pos'] for m in sim.matter_particles])
    if m_pos.size > 0:
        matter_dots.set_data(m_pos[:, 0], m_pos[:, 1])
    
    for i, (p1, p2) in enumerate(edges):
        lines[i].set_data([sim.pos[p1, 0], sim.pos[p2, 0]], [sim.pos[p1, 1], sim.pos[p2, 1]])
    
    return [points, matter_dots] + lines

ani = FuncAnimation(fig, animate, frames=500, interval=30, blit=True)
plt.legend(loc='upper right')
plt.show()