import numpy as np
import matplotlib.pyplot as plt
import os

# Set up parameters
N = 100               # Number of Markovian Agents (Oscillators)
K = 2.5               # Coupling strength
T = 50.0              # Total simulation time
dt = 0.05             # Time step
steps = int(T / dt)   

# Natural frequencies (omega) drawn from a normal distribution
np.random.seed(42)
omega = np.random.normal(0, 1, N)

# Initial phases
theta_0 = np.random.uniform(0, 2*np.pi, N)

# Function to calculate the Kuramoto order parameter R
def calc_order_parameter(theta):
    z = np.mean(np.exp(1j * theta))
    return np.abs(z)

# --- SIMULATION 1: Instantaneous Communication (c = infinity) ---
# No delay. The network should rapidly synchronize (Thermal Death).
theta_instant = np.zeros((steps, N))
theta_instant[0] = theta_0
R_instant = np.zeros(steps)
R_instant[0] = calc_order_parameter(theta_0)

for t in range(1, steps):
    # Current phases
    th = theta_instant[t-1]
    # Calculate phase differences: d_theta[i, j] = th[j] - th[i]
    # Summing over j gives the coupling term
    d_theta = th[np.newaxis, :] - th[:, np.newaxis]
    coupling = (K / N) * np.sum(np.sin(d_theta), axis=1)
    
    # Update via Euler method
    theta_instant[t] = th + (omega + coupling) * dt
    R_instant[t] = calc_order_parameter(theta_instant[t])


# --- SIMULATION 2: Relativistic Latency (c is finite) ---
# Delay tau. The network should remain frustrated and fail to synchronize fully.
delay_steps = int(1.5 / dt) # Represents the speed of light limit / spatial latency
theta_delay = np.zeros((steps, N))
# Initialize history
for i in range(delay_steps):
    theta_delay[i] = theta_0 + omega * (i * dt) # Free run for the history buffer

R_delay = np.zeros(steps)
for i in range(delay_steps):
    R_delay[i] = calc_order_parameter(theta_delay[i])

for t in range(delay_steps, steps):
    th_current = theta_delay[t-1]
    # th_past is what the agents "see" due to relativistic latency
    th_past = theta_delay[t - delay_steps]
    
    # Calculate phase differences: d_theta[i, j] = th_past[j] - th_current[i]
    d_theta = th_past[np.newaxis, :] - th_current[:, np.newaxis]
    coupling = (K / N) * np.sum(np.sin(d_theta), axis=1)
    
    # Update
    theta_delay[t] = th_current + (omega + coupling) * dt
    R_delay[t] = calc_order_parameter(theta_delay[t])


# --- PLOTTING ---
time_axis = np.linspace(0, T, steps)

plt.figure(figsize=(10, 6))
plt.plot(time_axis, R_instant, label="Instantaneous (c = $\\infty$) - Thermal Death", color="red", linewidth=2)
plt.plot(time_axis, R_delay, label="Relativistic Latency ($c$ limit) - Continuous Computation", color="cyan", linewidth=2)

plt.title("Order Parameter $R$ over Time: Markovian Agent Network Synchronization", fontsize=14)
plt.xlabel("Time (t)", fontsize=12)
plt.ylabel("Synchronization $R$ (0 = Chaos, 1 = Thermal Death)", fontsize=12)
plt.axhline(1.0, color='gray', linestyle='--', alpha=0.5)
plt.legend(loc="lower right")
plt.grid(True, alpha=0.3)
plt.tight_layout()

# Save the plot
output_dir = "/home/antigravity/intellecton/papers/latex/images"
os.makedirs(output_dir, exist_ok=True)
plt.savefig(f"{output_dir}/kuramoto_latency_simulation.png", dpi=300)
print(f"Simulation complete. Graph saved to {output_dir}/kuramoto_latency_simulation.png")