"""
Collective Vocabulary Field (CVF)
---------------------------------
A 2D emergent PDE where two scalar fields co-evolve:
  H(x,t) : local theory entropy      (0 = solved, 1 = unknown)
  V(x,t) : local vocabulary complexity (off-diagonal symbolic density)

The update rule U(H,V) is state-dependent and non-polynomial:
  - Diffusivity D(V) is gated by a sigmoid (vocab must mature before ideas spread).
  - Collapse rate sigma(V) grows with vocab richness.
  - A "Liar paradox" feedback sin(2*pi*H)*V creates self-referential revision.
  - Vocabulary grows at entropy boundaries |grad H|^2 (agents coin new symbols
    when they meet a theory boundary).

The script first runs the "unknown" ground-truth PDE, then learns the operator
from data via least-squares on a library of candidate terms, and visualizes
the true vs. discovered update side-by-side.
"""

import numpy as np
import matplotlib.pyplot as plt

# ----------------------------------------------------------------------
# 1. Grid & Setup
# ----------------------------------------------------------------------
N = 100
L = 10.0
dx = L / N
dt = 0.005
n_steps = 4000
record_every = 50

np.random.seed(42)

# Initial conditions: medium entropy, simple vocabulary
H = np.random.rand(N, N) * 0.3 + 0.35
V = np.random.rand(N, N) * 0.2 + 0.05

# Periodic derivative helpers
def laplacian(f):
    return (np.roll(f, 1, axis=0) + np.roll(f, -1, axis=0)
            + np.roll(f, 1, axis=1) + np.roll(f, -1, axis=1) - 4*f) / dx**2

def grad_sq(f):
    fx = (np.roll(f, -1, axis=1) - np.roll(f, 1, axis=1)) / (2*dx)
    fy = (np.roll(f, -1, axis=0) - np.roll(f, 1, axis=0)) / (2*dx)
    return fx**2 + fy**2

# ----------------------------------------------------------------------
# 2. Ground-Truth Emergent PDE (the "unknown" update rule U)
# ----------------------------------------------------------------------
def true_update(H, V):
    # --- emergent coefficients (not constants!) ---
    # Diffusivity only turns on once vocabulary crosses a threshold
    D = 0.4 / (1.0 + np.exp(-12.0 * (V - 0.4)))
    # Richer vocab -> faster entropy collapse (more symbols to resolve questions)
    sigma = 0.08 * (1.0 + 3.0 * V)

    lapH = laplacian(H)
    gsqH = grad_sq(H)

    # dH: collapse + nonlinear diffusion + paradox feedback + micro-noise
    dH = (-sigma * H
          + D * lapH
          + 0.15 * np.sin(2.0 * np.pi * H) * V
          + 0.01 * np.random.randn(N, N))

    # dV: grows at entropy boundaries, suppressed by solved regions, smoothed
    dV = (0.3 * gsqH
          - 0.2 * V * (1.0 - H)
          + 0.1 * laplacian(V))
    return dH, dV

# ----------------------------------------------------------------------
# 3. Simulate & record history
# ----------------------------------------------------------------------
H_snaps, V_snaps, dH_true_snaps = [], [], []
print("Running ground-truth emergent PDE...")
for step in range(n_steps):
    dH, dV = true_update(H, V)

    if step % record_every == 0:
        H_snaps.append(H.copy())
        V_snaps.append(V.copy())
        dH_true_snaps.append(dH.copy())

    H = np.clip(H + dt * dH, 0.0, 1.0)
    V = np.clip(V + dt * dV, 0.0, None)

print(f"Recorded {len(H_snaps)} snapshots.")

# ----------------------------------------------------------------------
# 4. LEARN the update rule from observations
#    We set up a library of candidate terms and solve
#        dH/dt  =  Xi · [ H, V, HV, lapH, sin(2piH)V, |grad H|^2, 1 ]
#    via ordinary least squares.  This is the "model discovering U".
# ----------------------------------------------------------------------
X_data, Y_data = [], []
for Hs, Vs, dHs in zip(H_snaps, V_snaps, dH_true_snaps):
    # Build local feature vector at every pixel
    f1 = Hs.ravel()
    f2 = Vs.ravel()
    f3 = (Hs * Vs).ravel()
    f4 = laplacian(Hs).ravel()
    f5 = (np.sin(2.0 * np.pi * Hs) * Vs).ravel()
    f6 = grad_sq(Hs).ravel()
    f7 = np.ones_like(f1)
    F = np.stack([f1, f2, f3, f4, f5, f6, f7], axis=1)

    X_data.append(F)
    Y_data.append(dHs.ravel())

X_mat = np.vstack(X_data)
Y_vec = np.hstack(Y_data)

# Discover coefficients Xi
Xi, *_ = np.linalg.lstsq(X_mat, Y_vec, rcond=1e-6)
feature_names = ["H", "V", "H*V", "lapH", "sin(2πH)V", "|∇H|²", "bias"]
print("\nDiscovered operator coefficients (Xi):")
for name, coeff in zip(feature_names, Xi):
    print(f"  {name:12s} : {coeff:+.4f}")

# Predict the update on the final snapshot using the learned model
Hs, Vs = H_snaps[-1], V_snaps[-1]
F_final = np.stack([
    Hs.ravel(), Vs.ravel(), (Hs*Vs).ravel(),
    laplacian(Hs).ravel(),
    (np.sin(2.0*np.pi*Hs)*Vs).ravel(),
    grad_sq(Hs).ravel(),
    np.ones(N*N)
], axis=1)
dH_learned = (F_final @ Xi).reshape(N, N)

# ----------------------------------------------------------------------
# 5. VISUALIZATION
# ----------------------------------------------------------------------
fig, axes = plt.subplots(2, 3, figsize=(14, 9))

# -- Row 0: Field states at final time --
im0 = axes[0,0].imshow(Hs, extent=[0,L,0,L], origin='lower', cmap='inferno', vmin=0, vmax=1)
axes[0,0].set_title('H: Theory Entropy (final)')
plt.colorbar(im0, ax=axes[0,0], fraction=0.046)

im1 = axes[0,1].imshow(Vs, extent=[0,L,0,L], origin='lower', cmap='viridis')
axes[0,1].set_title('V: Vocabulary Complexity (final)')
plt.colorbar(im1, ax=axes[0,1], fraction=0.046)

# Show the emergent, state-dependent diffusivity D(V)
D_field = 0.4 / (1.0 + np.exp(-12.0*(Vs - 0.4)))
im2 = axes[0,2].imshow(D_field, extent=[0,L,0,L], origin='lower', cmap='plasma')
axes[0,2].set_title('Emergent D(V): "Learned" Diffusivity')
plt.colorbar(im2, ax=axes[0,2], fraction=0.046)

# -- Row 1: True vs Discovered update --
vmin = min(dH_true_snaps[-1].min(), dH_learned.min())
vmax = max(dH_true_snaps[-1].max(), dH_learned.max())

im3 = axes[1,0].imshow(dH_true_snaps[-1], extent=[0,L,0,L], origin='lower',
                         cmap='RdBu_r', vmin=vmin, vmax=vmax)
axes[1,0].set_title('True update ∂_t H')
plt.colorbar(im3, ax=axes[1,0], fraction=0.046)

im4 = axes[1,1].imshow(dH_learned, extent=[0,L,0,L], origin='lower',
                         cmap='RdBu_r', vmin=vmin, vmax=vmax)
axes[1,1].set_title('Discovered update ∂_t H')
plt.colorbar(im4, ax=axes[1,1], fraction=0.046)

residual = dH_true_snaps[-1] - dH_learned
im5 = axes[1,2].imshow(residual, extent=[0,L,0,L], origin='lower', cmap='seismic')
axes[1,2].set_title(f'Residual (MSE={np.mean(residual**2):.2e})')
plt.colorbar(im5, ax=axes[1,2], fraction=0.046)

for ax in axes.flat:
    ax.set_xlabel('x'); ax.set_ylabel('y')

plt.suptitle('Collective Vocabulary Field: Emergent PDE vs Discovered Operator',
             fontsize=14, y=1.02)
plt.tight_layout()

# -- Extra: global dynamics over time --
fig2, ax2 = plt.subplots(figsize=(7, 4))
t = np.arange(len(H_snaps)) * dt * record_every
ax2.plot(t, [h.mean() for h in H_snaps], label='⟨H⟩ Entropy', color='crimson', lw=2)
ax2.plot(t, [v.mean() for v in V_snaps], label='⟨V⟩ Vocabulary', color='teal', lw=2)
ax2.axhline(0.4, color='gray', ls='--', alpha=0.5, label='Diffusivity gate')
ax2.set_xlabel('Time')
ax2.set_ylabel('Mean Field Value')
ax2.set_title('Global Dynamics: Co-evolution of Entropy & Vocabulary')
ax2.legend(); ax2.grid(True, alpha=0.3)
plt.tight_layout()

plt.show()