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

class SemanticMeshSim:
    def __init__(self, grid_size=10, singularity_strength=0.5):
        self.grid_size = grid_size
        self.strength = singularity_strength
        
        # 1. Initialize the Mesh (The Stationary Grid)
        # Create a grid of points (x, y)
        x = np.linspace(-5, 5, grid_size)
        y = np.linspace(-5, 5, grid_size)
        X, Y = np.meshgrid(x, y)
        
        self.pos = np.stack([X.flatten(), Y.flatten()], axis=-1)
        self.orig_pos = self.pos.copy()
        self.velocity = np.zeros_like(self.pos)
        
        # The Singularity (The "Truth" object at the center)
        self.singularity = np.array([0.0, 0.0])
        
    def compute_forces(self):
        # Current positions
        curr_pos = self.pos
        
        # A. Singularity Pull (The "Collapse" Force)
        # Force = G * M / r^2
        diff = self.singularity - curr_pos
        dist_sq = np.sum(diff**2, axis=1, keepdims=True)
        dist = np.sqrt(dist_sq)
        
        # Avoid division by zero at singularity
        dist_sq_safe = np.maximum(dist_sq, 0.1)
        
        # The pull force: Strength / r^2 * direction
        force_sing = (self.strength / dist_sq_safe) * (diff / dist)
        
        # B. Mesh Stability Force (The "Rubber Sheet" tension)
        # Points want to stay near their original relative neighbors
        # Simplified as a spring force toward their original positions 
        # but scaled by the global compression to allow the mesh to move
        force_spring = 0.05 * (self.orig_pos - curr_pos)
        
        return force_sing + force_spring

    def update(self, dt=0.1):
        #-- CCT Integration (ODE step) --
        forces = self.compute_forces()
        
        # Update velocity (with some damping to prevent infinite oscillation)
        self.velocity = (self.velocity + forces * dt) * 0.9
        
        # Update position
        self.pos += self.velocity * dt

# --- Visualization Setup ---
grid_res = 12
sim = SemanticMeshSim(grid_size=grid_res)

fig, ax = plt.subplots(figsize=(8, 8))
ax.set_xlim(-6, 6)
ax.set_ylim(-6, 6)
ax.set_aspect('equal')
ax.set_title("ODE-CCT Mesh Deformation: Singularity Collapse")
ax.set_facecolor('#121212')
fig.patch.set_facecolor('#121212')

# Draw the singularity
ax.plot(0, 0, 'yo', markersize=10, label="Singularity (Truth)", markeredgecolor='white')

# Create the plot elements
points, = ax.plot([], [], 'wo', markersize=3, alpha=0.6)
lines = []

# Pre-calculate grid indices for drawing lines (the "mesh" connectivity)
edges = []
for i in range(grid_res):
    for j in range(grid_res):
        idx = i * grid_res + j
        if j < grid_res - 1: edges.append((idx, idx + 1)) # Horizontal
        if i < grid_res - 1: edges.append((idx, idx + grid_res)) # Vertical

# Create line objects for the mesh
for _ in edges:
    l, = ax.plot([], [], 'w-', alpha=0.2, linewidth=0.5)
    lines.append(l)

def init():
    points.set_data([], [])
    for l in lines:
        l.set_data([], [])
    return [points] + lines

def animate(frame):
    sim.update()
    
    # Update points
    points.set_data(sim.pos[:, 0], sim.pos[:, 1])
    
    # Update mesh lines
    for i, (p1_idx, p2_idx) in enumerate(edges):
        p1 = sim.pos[p1_idx]
        p2 = sim.pos[p2_idx]
        lines[i].set_data([p1[0], p2[0]], [p1[1], p2[1]])
    
    return [points] + lines

ani = FuncAnimation(fig, animate, frames=200, init_func=init, interval=30, blit=True)
plt.legend()
plt.show()