"""
PDE solved under the Generalized Derivative Framework
======================================================
df ≡ The Update to f  (not a ratio of limits)

For the 2D Heat Equation:
  Standard:  ∂u/∂t = α∇²u           (rate of change)
  Generalized: du = U(u) = α∇²u · dt  (state determines its own update)

  S_{t+1} = S_t + U(S_t)

The update rule asks: "Given my current temperature distribution,
what should I become?" — and answers via the Laplacian: "Smooth out
sharp gradients. Hot regions cool, cold regions warm."

This is the self-referential loop: the state reads itself to decide
its own next step.
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import matplotlib.colors as mcolors
from scipy.ndimage import laplace
import matplotlib.gridspec as gridspec

# ═══════════════════════════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════════════════════════

Nx, Ny = 128, 128           # grid resolution
alpha = 0.2                 # thermal diffusivity
dt = 2.0e-4                 # time step (stable explicit Euler step)
steps_per_frame = 3         # simulation steps per animation frame
total_frames = 400          # animation frames
export_stride = 4           # save every Nth frame to keep GIF encoding manageable
export_fps = 15             # lower FPS makes GIF encoding faster/smaller
export_dpi = 80             # lower DPI reduces render cost

# ═══════════════════════════════════════════════════════════════════
# INITIAL STATE  S₀
# ═══════════════════════════════════════════════════════════════════

x = np.linspace(-1, 1, Nx)
y = np.linspace(-1, 1, Ny)
X, Y = np.meshgrid(x, y)

# A rich initial temperature distribution — multiple features at
# different spatial scales so we can watch the update rule act
# differently in different regions
u = np.zeros((Ny, Nx))

# Hot spot (top-right)
u += 0.8 * np.exp(-((X - 0.5)**2 + (Y - 0.5)**2) / 0.02)
# Cold spot (bottom-left)  
u -= 0.6 * np.exp(-((X + 0.5)**2 + (Y + 0.5)**2) / 0.03)
# Warm ring
u += 0.4 * np.exp(-((np.sqrt(X**2 + Y**2) - 0.4)**2) / 0.005)
# Cool diagonal band
u -= 0.3 * np.exp(-((X - Y)**2) / 0.08)
# Small hot dots
u += 0.5 * np.exp(-((X + 0.6)**2 + (Y - 0.6)**2) / 0.01)
u += 0.5 * np.exp(-((X - 0.7)**2 + (Y + 0.3)**2) / 0.01)

# Boundary: fixed cold
u[0, :] = u[-1, :] = u[:, 0] = u[:, -1] = 0.0

# Store initial state
u0 = u.copy()

# ═══════════════════════════════════════════════════════════════════
# THE GENERALIZED UPDATE RULE  U(S)
# ═══════════════════════════════════════════════════════════════════

def generalized_update(state, alpha, dt, dx):
    """
    du = U(S)  — the update rule

    The state S asks itself: "Where am I steep?"
    The Laplacian ∇²u answers: "Here are the sharp gradients."
    Then: "Smooth proportional to how steep."

    This is the core of the framework:
      S_{t+1} = S_t + U(S_t)

    Not "rate of change" — just "what update should I apply?"
    """
    # The Laplacian tells the state about its own curvature
    laplacian = laplace(state) / dx**2

    # The update is: move toward smoothness, proportional to
    # how far from smooth you currently are
    du = alpha * laplacian * dt

    # Enforce boundary condition within the update (the state
    # "knows" boundaries are fixed)
    du[0, :] = du[-1, :] = du[:, 0] = du[:, -1] = 0.0

    return du


# ═══════════════════════════════════════════════════════════════════
# TRACK EVOLUTION METRICS (for CCT-style analysis)
# ═══════════════════════════════════════════════════════════════════

# We track the "entropy" proxy: total |∇u| (magnitude of gradients)
def gradient_magnitude(state, dx):
    gy, gx = np.gradient(state, dx)
    return np.sqrt(gx**2 + gy**2)


def compute_entropy(state, dx):
    """H(S) — how much 'structure' remains to be collapsed."""
    return np.sum(gradient_magnitude(state, dx))


# ═══════════════════════════════════════════════════════════════════
# SIMULATE FULL TRAJECTORY
# ═══════════════════════════════════════════════════════════════════

dx = x[1] - x[0]
total_steps = total_frames * steps_per_frame

# Storage for the trajectory — so we can reconstruct S from {df₁, df₂, ...}
snapshots = np.zeros((total_frames, Ny, Nx))
entropy_log = np.zeros(total_frames)
update_magnitudes = np.zeros(total_frames)

state = u0.copy()
snapshots[0] = state
entropy_log[0] = compute_entropy(state, dx)

print("Evolving state via generalized update rule U(S) ...")
print(f"  Grid: {Nx}×{Ny}, Total steps: {total_steps}, Frames: {total_frames}")
print(f"  Initial entropy H(S₀) = {entropy_log[0]:.1f}")

for frame in range(1, total_frames):
    for _ in range(steps_per_frame):
        du = generalized_update(state, alpha, dt, dx)
        state = state + du  # S_{t+1} = S_t + U(S_t)
        # Boundaries
        state[0, :] = state[-1, :] = state[:, 0] = state[:, -1] = 0.0
        if not np.isfinite(state).all():
            print(f"  Warning: state became non-finite at frame {frame}. Stopping early.")
            state = np.nan_to_num(state, nan=0.0, posinf=0.0, neginf=0.0)
            break

    if not np.isfinite(state).all():
        break

    snapshots[frame] = state.copy()
    entropy_log[frame] = compute_entropy(state, dx)
    update_magnitudes[frame] = np.sum(np.abs(
        snapshots[frame] - snapshots[frame - 1]
    ))

finite_entropy = entropy_log[np.isfinite(entropy_log)]
if finite_entropy.size == 0:
    final_entropy = np.nan
    entropy_collapse = np.nan
    compression_ratio = np.nan
else:
    final_entropy = finite_entropy[-1]
    entropy_collapse = entropy_log[0] - final_entropy
    compression_ratio = final_entropy / entropy_log[0]

print(f"  Final entropy   H(S_T) = {final_entropy:.1f}")
print(f"  Entropy collapsed:      {entropy_collapse:.1f}")
print(f"  Compression ratio:      {compression_ratio:.3f}")

# ═══════════════════════════════════════════════════════════════════
# VERIFY: Reconstruct S from df sequence
# ═══════════════════════════════════════════════════════════════════

# The generalized integral: S = S₀ + Σ df_i
# df_i = S_{t+1} - S_t
reconstructed = u0.copy()
for frame in range(1, total_frames):
    df = snapshots[frame] - snapshots[frame - 1]
    reconstructed = reconstructed + df

error = np.max(np.abs(reconstructed - snapshots[-1]))
print(f"  Reconstruction error:   {error:.2e}  (should be ~0)")

# ═══════════════════════════════════════════════════════════════════
# VISUALIZATION
# ═══════════════════════════════════════════════════════════════════

print("\nRendering visualization...")

fig = plt.figure(figsize=(18, 10), facecolor='#0d1117')
gs = gridspec.GridSpec(2, 3, figure=fig, hspace=0.35, wspace=0.35,
                       height_ratios=[1, 0.3])

# --- Main heatmap ---
ax_main = fig.add_subplot(gs[0, :2])
cmap = plt.cm.RdBu_r
norm = mcolors.TwoSlopeNorm(vmin=-0.8, vcenter=0, vmax=0.8)

img = ax_main.imshow(snapshots[0], cmap=cmap, norm=norm,
                     extent=[-1, 1, -1, 1], origin='lower',
                     interpolation='bilinear')
ax_main.set_title('State S(t) — Temperature Distribution',
                  color='white', fontsize=14, fontweight='bold')
ax_main.set_xlabel('x', color='white')
ax_main.set_ylabel('y', color='white')
ax_main.tick_params(colors='white')

# Time annotation
time_text = ax_main.text(0.02, 0.98, '', transform=ax_main.transAxes,
                         color='white', fontsize=12, va='top',
                         fontfamily='monospace',
                         bbox=dict(boxstyle='round', facecolor='#0d1117',
                                   alpha=0.8))

# --- Update field (df at current step) ---
ax_df = fig.add_subplot(gs[0, 2])
df_norm = mcolors.TwoSlopeNorm(vmin=-0.05, vcenter=0, vmax=0.05)
df_img = ax_df.imshow(np.zeros_like(u0), cmap='coolwarm', norm=df_norm,
                      extent=[-1, 1, -1, 1], origin='lower',
                      interpolation='bilinear')
ax_df.set_title('Update df = U(S) — "What changed?"',
                color='white', fontsize=12, fontweight='bold')
ax_df.set_xlabel('x', color='white')
ax_df.set_ylabel('y', color='white')
ax_df.tick_params(colors='white')

# --- Entropy plot ---
ax_ent = fig.add_subplot(gs[1, 0])
ax_ent.set_facecolor('#161b22')
line_ent, = ax_ent.plot([], [], color='#58a6ff', lw=2)
ax_ent.set_title('Entropy H(S) = Σ|∇u|', color='white', fontsize=11)
ax_ent.set_xlabel('Time step', color='white')
ax_ent.set_ylabel('H(S)', color='white')
ax_ent.tick_params(colors='white')
ax_ent.set_xlim(0, total_frames)
ax_ent.set_ylim(0, entropy_log[0] * 1.1)
ax_ent.grid(True, alpha=0.2)
ent_marker, = ax_ent.plot([], [], 'ro', markersize=6)

# --- Update magnitude plot ---
ax_upd = fig.add_subplot(gs[1, 1])
ax_upd.set_facecolor('#161b22')
line_upd, = ax_upd.plot([], [], color='#f0883e', lw=2)
ax_upd.set_title('|df| = Update Magnitude', color='white', fontsize=11)
ax_upd.set_xlabel('Time step', color='white')
ax_upd.set_ylabel('Σ|df|', color='white')
ax_upd.tick_params(colors='white')
ax_upd.set_xlim(0, total_frames)
finite_updates = update_magnitudes[1:][np.isfinite(update_magnitudes[1:])]
update_ylim = finite_updates.max() * 1.1 if finite_updates.size else 1.0
ax_upd.set_ylim(0, update_ylim)
ax_upd.grid(True, alpha=0.2)

# --- Phase portrait: dH vs H ---
ax_phase = fig.add_subplot(gs[1, 2])
ax_phase.set_facecolor('#161b22')
dH = -np.diff(entropy_log)
H_mid = (entropy_log[:-1] + entropy_log[1:]) / 2
ax_phase.scatter(H_mid[::5], dH[::5], c=np.arange(len(H_mid))[::5],
                 cmap='viridis', s=8, alpha=0.8)
ax_phase.set_title('Phase: −dH/dt vs H (update vs state)',
                   color='white', fontsize=11)
ax_phase.set_xlabel('H(S)', color='white')
ax_phase.set_ylabel('−ΔH (collapse rate)', color='white')
ax_phase.tick_params(colors='white')
ax_phase.grid(True, alpha=0.2)
# Arrow showing direction
ax_phase.annotate('time →', xy=(0.05, 0.95),
                  xycoords='axes fraction', color='white', fontsize=9)

# --- Info box ---
info_text = fig.text(0.02, 0.02,
    'Generalized Derivative:  df ≡ The Update to f\n'
    'ODE: dS = U(S)  — state determines its own change\n'
    'CCT: dH = −Δ  — each question collapses entropy',
    color='#8b949e', fontsize=9, fontfamily='monospace',
    transform=fig.transFigure)

plt.tight_layout(rect=[0, 0.07, 1, 1])

# ═══════════════════════════════════════════════════════════════════
# ANIMATION
# ═══════════════════════════════════════════════════════════════════

def animate(frame):
    """Each frame: state evolved via U(S), plus diagnostics."""
    # Main state
    img.set_array(snapshots[frame])

    # Current update df
    if frame > 0:
        df_current = snapshots[frame] - snapshots[frame - 1]
    else:
        df_current = np.zeros_like(u0)
    df_img.set_array(df_current)

    # Time label
    t = frame * steps_per_frame * dt
    time_text.set_text(f't = {t:.2f}  |  frame {frame}/{total_frames}')

    # Entropy plot
    line_ent.set_data(range(frame + 1), entropy_log[:frame + 1])
    ent_marker.set_data([frame], [entropy_log[frame]])

    # Update magnitude
    line_upd.set_data(range(frame + 1), update_magnitudes[:frame + 1])

    return [img, df_img, time_text, line_ent, ent_marker, line_upd]


export_frames = range(0, total_frames, export_stride)

ani = FuncAnimation(fig, animate, frames=export_frames,
                    interval=40, blit=True)

# ═══════════════════════════════════════════════════════════════════
# SAVE & DISPLAY
# ═══════════════════════════════════════════════════════════════════

output_path = 'generalized_pde_heat_equation.gif'
print(f"Saving animation to {output_path} ...")
writer = PillowWriter(fps=export_fps)
ani.save(output_path, writer=writer, dpi=export_dpi)
print(f"Done! Saved {output_path}")
print(f"  File size: {__import__('os').path.getsize(output_path) / 1024:.0f} KB")

plt.close(fig)

# ═══════════════════════════════════════════════════════════════════
# FINAL DIAGNOSTIC: Fixed-point analysis
# ═══════════════════════════════════════════════════════════════════

print("\n" + "=" * 60)
print("DIAGNOSTICS — Generalized Derivative Framework")
print("=" * 60)
print(f"""
  Framework mapping:
    State S      = 2D temperature field ({Nx}×{Ny} grid)
    Update U(S)  = α∇²S · dt  (Laplacian smoothing)
    Evolution    = S_{t+1} = S_t + U(S_t)
    
    Initial entropy H₀  = {entropy_log[0]:.1f}
    Final entropy   H_T  = {final_entropy:.1f}
    Total collapse  ΔH   = {entropy_collapse:.1f}
    
    Fixed point U(S*) = 0  →  ∇²S* = 0  →  S* is harmonic
    (The only state that stops updating is the harmonic state)
    
    Path dependence: Each df depends on current ∇²S,
    which depends on all previous df. The trajectory
    is the accumulating sum of state-aware updates.
""")
