import numpy as np
import matplotlib.pyplot as plt
from scipy.special import erfc
from scipy.optimize import fsolve

# --- Physical Parameters, arbitrarily chosen for illustration by H. K. D. H. Bhadeshia ---
C_inf = 0.5
m = -100.0
T_melt = 1000.0
T_start = 998.0
cooling_rate = 15.0  # Increased for visibility (K/s)
k = 0.3
gamma = 1e-7
D_ref = 5e-19
sigma_star = 0.025

# --- Simulation Setup ---
dt = 0.05
steps = 150
T = T_start
tip_z_act = 0.0
V_act = 0.0 

history_T = []
history_V_iso = []
history_V_act = []

def solve_kinetics(temp):
    C_L = (temp - T_melt) / m
    C_S = k * C_L
    Omega = np.clip((C_inf - C_L) / (C_S - C_L), 0.01, 0.95)
    D = D_ref * np.exp(-2000/temp) # Stronger Arrhenius effect
    res = lambda p: Omega - np.sqrt(np.pi * p) * np.exp(p) * erfc(np.sqrt(p))
    Pe = fsolve(res, 0.1)[0]
    R = np.sqrt(gamma * D / (sigma_star * Pe * abs(m) * (C_L - C_S)))
    V = (2 * D * Pe) / R
    return V, R, D, Pe

plt.ion()
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))

for i in range(steps):
    T -= cooling_rate * dt
    V_iso, R_iso, D, Pe = solve_kinetics(T)
    
    # RELAXATION LOGIC:
    # Define a larger tau to force a visible lag.
    # This represents the time for the diffusion field to reorganize.
    tau = 1.5  # Fixed relaxation time (seconds) to create clear gap
    
    if i == 0:
        V_act = V_iso
    else:
        # The growth rate 'chases' the ideal but lags behind
        V_act += (V_iso - V_act) * (dt / (tau + dt))

    tip_z_act += V_act * dt
    history_T.append(T)
    history_V_iso.append(V_iso * 1e6)
    history_V_act.append(V_act * 1e6)
    
    if i % 5 == 0: # Update plot every 5 steps for speed
        ax1.clear()
        zoom = 12 * R_iso
        x_grid = np.linspace(-zoom, zoom, 100)
        z_grid = np.linspace(tip_z_act - zoom*1.5, tip_z_act + zoom*0.5, 100)
        X, Z = np.meshgrid(x_grid, z_grid)
        dz = tip_z_act - Z
        xi = np.sqrt(X**2 + dz**2) - dz
        U = erfc(np.sqrt(xi / (2*R_iso/Pe))) / erfc(np.sqrt(Pe))
        U[(X**2 < 2*R_iso*dz) & (Z < tip_z_act)] = 1.2
        ax1.contourf(X*1e6, Z*1e6, U, levels=20, cmap='viridis')
        ax1.set_title(f"Morphology (T={T:.1f}K)")

        ax2.clear()
        ax2.plot(history_T, history_V_iso, 'r--', linewidth=2, label='Isothermal Steady-State')
        ax2.plot(history_T, history_V_act, 'b-', linewidth=2, label='Actual Cooling (Lag)')
        ax2.set_xlabel("Temperature (K)")
        ax2.set_ylabel("Velocity (µm/s)")
        ax2.invert_xaxis()
        ax2.legend()
        ax2.grid(True, alpha=0.3)
        plt.pause(0.01)

plt.ioff()
plt.show()