import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import mutual_info_score

# Parameters
N = 64               # number of oscillators (8x8 grid)
dt = 0.01
T_steps = 5000
coupling_strength = 1.0
noise_strength = 0.5  # acts like temperature

# Initial phases random
theta = np.random.uniform(0, 2*np.pi, N)
omega = np.random.normal(0, 0.1, N)  # natural frequencies

# Telepathic coupling: each oscillator gets mean-field (global) or local
def global_coupling(theta):
    return coupling_strength * np.mean(np.sin(theta[:, None] - theta), axis=0)

def local_coupling_2d(theta, grid_shape=(8,8)):
    # 4-nearest neighbor coupling on 2D grid
    coupling = np.zeros_like(theta)
    for i in range(N):
        x, y = divmod(i, grid_shape[1])
        neighbors = []
        if x > 0: neighbors.append(i - grid_shape[1])
        if x < grid_shape[0]-1: neighbors.append(i + grid_shape[1])
        if y > 0: neighbors.append(i - 1)
        if y < grid_shape[1]-1: neighbors.append(i + 1)
        for nb in neighbors:
            coupling[i] += np.sin(theta[nb] - theta[i])
    return coupling_strength * coupling

# Collapse mechanism: reset to Xi_phase if divergence too high
Xi_phase = 0.0   # reference phase (anchor)
threshold = 0.5
temp = 0.8       # collapse temperature (higher = more random)

phi_history = []   # store order parameter r
phase_history = [] # store all phases for Φ calculation

for step in range(T_steps):
    # Compute coupling
    # Choose one: global_coupling or local_coupling_2d
    coup = global_coupling(theta)   # mean-field case
    # For local, uncomment next line and comment above
    # coup = local_coupling_2d(theta)
    
    # Kuramoto update
    theta += dt * (omega + coup + noise_strength * np.random.randn(N))
    
    # Collapse step: each oscillator independently may collapse to Xi_phase
    divergence = np.abs(theta - Xi_phase)
    p_collapse = 1 / (1 + np.exp((divergence - threshold) / temp))
    collapse_mask = np.random.rand(N) < p_collapse
    theta[collapse_mask] = Xi_phase + 0.1 * np.random.randn(np.sum(collapse_mask))
    
    # Record order parameter
    r = np.abs(np.mean(np.exp(1j * theta)))
    phi_history.append(r)
    if step % 50 == 0:
        phase_history.append(theta.copy())

# Compute Φ from phase history (split array into two halves)
phase_arr = np.array(phase_history)  # (T, N)
half = N // 2
X = phase_arr[:, :half]
Y = phase_arr[:, half:]

# Discretize phases into 4 bins (0, π/2, π, 3π/2) for mutual info
def discretize_phase(ph):
    return (ph % (2*np.pi) // (np.pi/2)).astype(int)

X_disc = discretize_phase(X)
Y_disc = discretize_phase(Y)
I_whole = mutual_info_score(
    [f"{x}{y}" for x,y in zip(X_disc[:-1].flatten(), Y_disc[:-1].flatten())],
    [f"{x}{y}" for x,y in zip(X_disc[1:].flatten(), Y_disc[1:].flatten())]
)
I_X = mutual_info_score(X_disc[:-1].flatten(), X_disc[1:].flatten())
I_Y = mutual_info_score(Y_disc[:-1].flatten(), Y_disc[1:].flatten())
phi_val = max(0.0, I_whole - (I_X + I_Y))

# Compute average order parameter (steady state)
steady_r = np.mean(phi_history[-1000:])

print(f"Average order parameter r = {steady_r:.3f}")
print(f"Integrated information Φ = {phi_val:.4f}")

# Plot order parameter over time
plt.figure(figsize=(10,4))
plt.plot(phi_history)
plt.xlabel('Time steps')
plt.ylabel('Order parameter r')
plt.title('Kuramoto Telepathic Array: Synchronization and Collapse')
plt.grid(True)
plt.show()