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

# ------------------------------
# 2D Kuramoto with collapse
# ------------------------------
class Kuramoto2D:
    def __init__(self, L=32, K=1.0, dt=0.01, T=0.5, 
                 collapse_threshold=0.5, Xi=0.0):
        self.L = L
        self.N = L*L
        self.K = K
        self.dt = dt
        self.T = T
        self.collapse_thr = collapse_threshold
        self.Xi = Xi
        
        # phases
        self.theta = np.random.uniform(0, 2*np.pi, self.N)
        # natural frequencies (set zero for KT transition)
        self.omega = np.zeros(self.N)
        
    def neighbors_2d(self, i):
        """Return indices of nearest neighbors (von Neumann)."""
        x, y = divmod(i, self.L)
        neigh = []
        if x > 0:   neigh.append(i - self.L)
        if x < self.L-1: neigh.append(i + self.L)
        if y > 0:   neigh.append(i - 1)
        if y < self.L-1: neigh.append(i + 1)
        return neigh
    
    def compute_coupling(self):
        """Local coupling term for each oscillator."""
        coupling = np.zeros(self.N)
        for i in range(self.N):
            for j in self.neighbors_2d(i):
                coupling[i] += np.sin(self.theta[j] - self.theta[i])
        return self.K * coupling
    
    def collapse_step(self):
        """Stochastic collapse to Xi based on divergence."""
        divergence = np.abs(self.theta - self.Xi)
        p_collapse = 1.0 / (1.0 + np.exp((divergence - self.collapse_thr) / self.T))
        collapse_mask = np.random.rand(self.N) < p_collapse
        self.theta[collapse_mask] = self.Xi
    
    def step(self):
        """One Euler–Maruyama step."""
        coup = self.compute_coupling()
        noise = np.sqrt(2*self.T/self.dt) * np.random.randn(self.N)  # scaled correctly
        self.theta += self.dt * (self.omega + coup + noise)
        self.collapse_step()
        # keep in [0, 2π)
        self.theta %= (2*np.pi)
    
    def run(self, steps, skip=100, measure=True):
        """Run simulation, record order parameter and correlation."""
        r_history = []
        configs = []  # for correlation at end
        for step in range(steps):
            self.step()
            if measure and step % skip == 0:
                r = np.abs(np.mean(np.exp(1j*self.theta)))
                r_history.append(r)
                configs.append(self.theta.copy())
        return np.array(r_history), configs

# ------------------------------
# Sweep temperature
# ------------------------------
temperatures = np.linspace(0.1, 1.2, 20)
r_means = []
r_stds = []
Tc_estimate = 0.7   # known for 2D XY model ~0.89, but with collapse it shifts

for Ti in temperatures:
    kuramoto = Kuramoto2D(L=32, T=Ti, collapse_threshold=0.5, Xi=0.0)
    r_hist, configs = kuramoto.run(steps=50000, skip=500)
    r_means.append(np.mean(r_hist[-20:]))   # steady state average
    r_stds.append(np.std(r_hist[-20:]))
    
# Plot order parameter vs T
plt.errorbar(temperatures, r_means, yerr=r_stds, fmt='o-')
plt.xlabel('Temperature T')
plt.ylabel('Order parameter r')
plt.title('2D Kuramoto with Collapse: Order parameter')
plt.grid(True)
plt.show()

# Fit correlation length divergence (requires correlation function)
# For a given T near Tc, compute g(r) = <cos(θ_i - θ_j)> vs Manhattan distance
# Then fit to exponential decay to get ξ, then ξ(T) ~ exp(b/√(T-Tc))

def correlation_length(config, L):
    """Estimate ξ from spatial correlation function."""
    theta = config.reshape(L, L)
    dists = []
    corrs = []
    max_dist = L//2
    for dx in range(1, max_dist):
        # average over all pairs with Manhattan distance dx
        c = 0
        count = 0
        for i in range(L):
            for j in range(L):
                if i+dx < L:
                    c += np.cos(theta[i,j] - theta[i+dx,j])
                    count += 1
                if j+dx < L:
                    c += np.cos(theta[i,j] - theta[i,j+dx])
                    count += 1
        corrs.append(c / count)
        dists.append(dx)
    # Fit to exponential A * exp(-dx/ξ)
    from scipy.optimize import curve_fit
    def exp_decay(x, A, xi):
        return A * np.exp(-x / xi)
    popt, _ = curve_fit(exp_decay, dists, corrs, p0=[1.0, 2.0])
    return popt[1]

# Compute ξ for several T above Tc
T_above = temperatures[temperatures > 0.7]  # assume Tc~0.7 from plot
xi_vals = []
for Ti, (Tval) in enumerate(T_above):
    # Run a long simulation at that T and get final config
    kuramoto = Kuramoto2D(L=32, T=Tval, collapse_threshold=0.5)
    _, configs = kuramoto.run(steps=200000, skip=1000)
    final_config = configs[-1]
    xi = correlation_length(final_config, 32)
    xi_vals.append(xi)

# Fit ξ(T) ~ exp(b / √(T - Tc))
def kt_divergence(T, b, Tc):
    return np.exp(b / np.sqrt(np.maximum(T - Tc, 1e-6)))
popt, _ = curve_fit(kt_divergence, T_above, xi_vals, p0=[2.0, 0.7])
b_fit, Tc_fit = popt
print(f"Fitted Tc = {Tc_fit:.3f}, b = {b_fit:.3f}")
# Expected KT: b ~ π / (2 * √2) ≈ 1.11 for pure XY; here with collapse may differ.