import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

# ---- Parameters ----
N1 = 50
om0 = 5
omw = 2
K = 2.5

# Kc (same formula as MATLAB)
print(2 * np.sqrt(2) * omw / np.sqrt(np.pi))

# ---- Random intrinsic frequencies ----
np.random.seed(42)
om = np.random.randn(N1) * omw + om0
om = np.sort(om)

# ---- Kuramoto ODE ----
def Kuramoto_ode(t, th):
    # th is shape (N1,)
    # compute pairwise phase differences th_j - th_i
    diff = th[None, :] - th[:, None]     # shape (N1, N1)
    coupling = (K / N1) * np.sum(np.sin(diff), axis=1)
    return om + coupling

# ---- Solve model (equivalent to ode45) ----
sol = solve_ivp(
    Kuramoto_ode,
    t_span=[0, 100],
    y0=np.zeros(N1),
    method='RK45',
    atol=1e-5,
    rtol=1e-5,
    max_step=0.1
)

t = sol.t
th1 = sol.y.T   # shape (len(t), N1)

# ---- Plot solutions (like MATLAB) ----
plt.figure()
plt.plot(t, np.sin(th1) + np.arange(1, N1+1)*2, linewidth=2)
plt.xlabel("t")
plt.ylabel("sin(th_i) + offset")
plt.title("Kuramoto Oscillators")
plt.show()

# ---- Plot averaged order parameter ----
plt.figure()
plt.plot(t, np.sum(np.sin(th1), axis=1)/N1, linewidth=3)
plt.xlim([5, 20])
plt.xlabel("t")
plt.ylabel("mean(sin(theta))")
plt.title("Mean Field")
plt.show()