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

# ---------- Ising model parameters ----------
J = 1.0       # coupling strength (telepathic gradient)
h_ext = 0.0   # external bias (can be varied)
T = 0.5       # temperature (lower = more ordered)
n_steps = 200
N = 2         # two "spins": state (s0) and memory (s1)

# Initialize spins randomly
spins = np.array([1, -1])  # +1 = A, -1 = B

# Store history for Φ computation
history = []

def local_field(spins, i):
    """Ising local field: J * sum_{j != i} spins[j] + external field.
       For STLPL, this is the TELEPATH gradient."""
    return J * (np.sum(spins) - spins[i]) + h_ext

def metropolis_step(spins, T):
    """Single Metropolis update (like a telepathic collapse decision)."""
    i = np.random.randint(0, N)
    Ei = -spins[i] * local_field(spins, i)
    new_spin = -spins[i]
    Ef = -new_spin * local_field(spins, i)
    delta_E = Ef - Ei
    if delta_E < 0 or np.random.rand() < np.exp(-delta_E / T):
        spins[i] = new_spin
    return spins

def compute_phi_from_spins(spin_history):
    """Estimate Φ as mutual information between partition (spin0, spin1) across time."""
    if len(spin_history) < 2:
        return 0.0
    X = np.array([h[0] for h in spin_history])  # spin0 (state)
    Y = np.array([h[1] for h in spin_history])  # spin1 (memory)
    # Discretize: already ±1
    I_whole = mutual_info_score(
        [f"{x}{y}" for x,y in zip(X[:-1], Y[:-1])],
        [f"{x}{y}" for x,y in zip(X[1:], Y[1:])]
    )
    I_X = mutual_info_score(X[:-1], X[1:])
    I_Y = mutual_info_score(Y[:-1], Y[1:])
    return max(0.0, I_whole - (I_X + I_Y))

# Run simulation at different temperatures
temperatures = np.linspace(0.1, 2.0, 20)
phi_vals = []
magnetization_vals = []

for T in temperatures:
    spins = np.array([1, -1])  # reset
    spin_history = []
    for step in range(n_steps):
        spins = metropolis_step(spins, T)
        spin_history.append(spins.copy())
    # Compute average magnetization (absolute)
    mag = np.mean(np.abs(np.mean(spin_history, axis=1)))
    magnetization_vals.append(mag)
    # Compute Φ
    phi = compute_phi_from_spins(spin_history)
    phi_vals.append(phi)

# Plot Φ vs T and magnetization vs T
plt.figure(figsize=(12,4))
plt.subplot(1,2,1)
plt.plot(temperatures, phi_vals, 'o-', color='red')
plt.xlabel('Temperature T')
plt.ylabel('Integrated Information Φ')
plt.title('STLPL/Ising: Φ vs T')
plt.grid(True)

plt.subplot(1,2,2)
plt.plot(temperatures, magnetization_vals, 's-', color='blue')
plt.xlabel('Temperature T')
plt.ylabel('|Magnetization|')
plt.title('Ising Model: Order Parameter')
plt.grid(True)
plt.tight_layout()
plt.show()

print(f"Φ at low T (0.1): {phi_vals[0]:.4f}")
print(f"Φ at high T (2.0): {phi_vals[-1]:.4f}")
print(f"Critical T (inflection): ~{temperatures[np.argmax(np.gradient(phi_vals))]:.2f}")