#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
fig_gen.py
===============
Figure-generation script for the manuscript:

  "Perturbative modelling of Arrhenius reactive transport in non-isothermal
   Gordon-Schowalter Couette-Poiseuille flow"

Generates:
  - Figure 1: physical configuration.
  - Figure 2: thermal shape functions and centreline amplification.
  - Figure 3: eigenvalue-shift collapse and route decomposition.
  - figdata.json: numerical data behind the plotted curves and Table 2.

The script is self-contained and includes an independent second-order
finite-difference check of the shooting Sturm-Liouville solver.

Authors: Leonardo D. Soria R. and Anthony A. Harrup G.
Target journal: Journal of Engineering Mathematics
Version: v2-clean, June 2026
"""
import numpy as np, json
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from scipy.optimize import brentq
from scipy.linalg import eigh

plt.rcParams.update({"font.family":"serif","font.size":10,"axes.linewidth":0.8,
                     "figure.dpi":200})
RNG_TOL = dict(method="RK45", rtol=1e-11, atol=1e-13)

# ---------- SL machinery (same as validate.py) ----------
def shoot_residual(lam, Pe, Da_fun, phi_fun, symmetric=True):
    def rhs(eta,y): return [y[1],(Da_fun(eta)-lam**2*Pe*phi_fun(eta))*y[0]]
    if symmetric:
        sol=solve_ivp(rhs,[0.5,1.0],[1.0,0.0],**RNG_TOL)
    else:
        sol=solve_ivp(rhs,[0.0,1.0],[0.0,1.0],**RNG_TOL)
    return sol.y[0,-1]

def lam1_of(Pe,Da_fun,phi_fun,lam_max=8.0,dlam=0.05):
    grid=np.arange(dlam,lam_max,dlam); rp=shoot_residual(grid[0],Pe,Da_fun,phi_fun)
    for lam in grid[1:]:
        r=shoot_residual(lam,Pe,Da_fun,phi_fun)
        if rp*r<0:
            return brentq(shoot_residual,lam-dlam,lam,args=(Pe,Da_fun,phi_fun,True),
                          xtol=1e-12,rtol=1e-13)
        rp=r
    raise RuntimeError("no eigenvalue")

def eigfun(lam,Pe,Da_fun,phi_fun,N=2001):
    eta=np.linspace(0,1,N)
    def rhs(t,y): return [y[1],(Da_fun(t)-lam**2*Pe*phi_fun(t))*y[0]]
    half=eta[eta>=0.5]
    sol=solve_ivp(rhs,[0.5,1.0],[1.0,0.0],t_eval=half,**RNG_TOL)
    Hr=sol.y[0]; H=np.concatenate([Hr[:0:-1],Hr])
    return eta,H/np.max(np.abs(H))

# ---------- shape functions, P*=3 reference ----------
P, Pe, Da0, Br, gA = 3.0, 5.0, 1.0, 1.0, 5.0
phi0=lambda e:1.0+(P/2.0)*e*(1-e)
phi1=lambda e:(P**3/4.0)*(1/16.0-(e-0.5)**4)
thB =lambda e:(P**2/12.0)*(1/16.0-(e-0.5)**4)
thBW=lambda e:(P**4/30.0)*(1/64.0-(e-0.5)**6)

# ---------- FD independent cross-check (self-contained validation) ----------
def lam1_FD(Pe,Da_fun,phi_fun,N=4001):
    """Generalized FD eigenproblem: -H'' + Da H = lam^2 Pe phi H, H(0)=H(1)=0."""
    eta=np.linspace(0,1,N); h=eta[1]-eta[0]; ei=eta[1:-1]
    main=2/h**2+Da_fun(ei); off=-np.ones(len(ei)-1)/h**2
    A=np.diag(main)+np.diag(off,1)+np.diag(off,-1)
    B=np.diag(Pe*phi_fun(ei))
    w=eigh(A,B,eigvals_only=True,subset_by_index=[0,0])[0]
    return np.sqrt(w)

print("FD cross-check of the shooting/Brent solver (base problem, Pe=5, Da=1, P*=3):")
fd_errs=[]
for (pe,da) in [(5.0,1.0),(5.0,5.0),(1.0,0.0)]:
    ls=lam1_of(pe,lambda e:da,phi0); lf=lam1_FD(pe,lambda e:np.full_like(e,da),
        lambda e:1.0+(P/2.0)*e*(1-e))
    err=abs(ls-lf)/ls; fd_errs.append(err)
    print(f"  Pe={pe} Da={da}: shooting={ls:.6f}  FD(N=4001)={lf:.6f}  rel dev={err:.2e}")
fd_worst=max(fd_errs)

# ---------- base mode, integrals, route coefficients ----------
lam10=lam1_of(Pe,lambda e:Da0,phi0)
eta,H1=eigfun(lam10,Pe,lambda e:Da0,phi0)
def integ(f): return np.trapezoid(f,eta)
N1=integ(phi0(eta)*H1**2)
A1=integ(phi1(eta)*H1**2); TBW=integ(thBW(eta)*H1**2); TB=integ(thB(eta)*H1**2)
c_T = Br*gA*Da0*TBW/(Pe*N1)        # thermal-reactive per unit eps
c_A = -lam10**2*A1/N1              # advective per unit eps
dl_B = Br*gA*Da0*TB/(Pe*N1)        # Newtonian thermal shift
K = c_T/(-c_A)
print(f"\nlam1={lam10:.5f} lam1^2={lam10**2:.5f}  dl_B={dl_B:.4e}")
print(f"c_thermal={c_T:.4e}  c_advect={c_A:.4e}  K={K:.4f}")

# ---------- exact perturbed eigenvalues over (a,Ws) grid ----------
grid=[]
for a in [0.0,0.3,0.5,0.7,0.9]:
    for Ws in [0.05,0.10,0.15,0.20,0.25]:
        eps=(1-a**2)*Ws**2
        if eps==0: continue
        phi_p=lambda e,ep=eps: phi0(e)+ep*phi1(e)
        Da_p=lambda e,ep=eps: Da0*(1.0+gA*Br*(thB(e)+ep*thBW(e)))
        Da_b=lambda e: Da0*(1.0+gA*Br*thB(e))
        lam_b=lam1_of(Pe,Da_b,phi0)           # thermally corrected Newtonian base
        lam_e=lam1_of(Pe,Da_p,phi_p)
        dl_exact=lam_e**2-lam_b**2
        dl_first=eps*(c_T+c_A)
        grid.append(dict(a=a,Ws=Ws,eps=eps,dl_exact=dl_exact,dl_first=dl_first,
                         err_pc=100*abs(dl_exact-dl_first)/abs(dl_exact)))
worst=max(g["err_pc"] for g in grid)
print(f"first-order vs exact: worst error {worst:.2f}% over grid (n={len(grid)})")
# UCM check
phi_u=lambda e: phi0(e); Da_u=lambda e: Da0*(1.0+gA*Br*thB(e))
lam_u=lam1_of(Pe,Da_u,phi_u); ucm_dev=abs(lam_u**2-lam1_of(Pe,Da_u,phi0)**2)
print(f"UCM/LCM invariance (eps=0 path identity): {ucm_dev:.1e}")

# ================= FIGURE 1: schematic =================
fig,ax=plt.subplots(figsize=(6.0,2.9))
ax.set_xlim(0,10); ax.set_ylim(-0.6,1.7); ax.axis("off")
ax.add_patch(plt.Rectangle((0.8,0),8.4,1.0,facecolor="#eef3fa",edgecolor="none"))
for yw in [0,1.0]:
    ax.plot([0.8,9.2],[yw,yw],"k-",lw=2.2)
    ax.annotate("",xy=(9.0,yw+(0.22 if yw==0 else -0.22)+ (0 if yw==0 else 0)),
                xytext=(8.2,yw+(0.22 if yw==0 else -0.22)),
                arrowprops=dict(arrowstyle="->",lw=1.4))
ax.text(9.25,1.05,r"$v_x=V,\ T=T_0,\ \Omega=0$",fontsize=9,va="bottom",ha="right")
ax.text(9.25,-0.15,r"$v_x=V,\ T=T_0,\ \Omega=0$",fontsize=9,va="top",ha="right")
# velocity profile phi0 + eps phi1
yy=np.linspace(0,1,100)
prof=1.0+(P/2.0)*yy*(1-yy); prof_v=prof+0.35*phi1(yy)   # exaggerated eps for visibility
x0=2.6
ax.plot(x0+prof*0.9,yy,"b-",lw=1.5,label="Newtonian")
ax.plot(x0+prof_v*0.9,yy,"r--",lw=1.5,label=r"GS, $\varepsilon>0$")
for yq in np.linspace(0.08,0.92,8):
    ax.annotate("",xy=(x0+(1.0+(P/2)*yq*(1-yq))*0.9,yq),xytext=(x0,yq),
                arrowprops=dict(arrowstyle="->",lw=0.6,color="0.45"))
ax.plot([x0,x0],[0,1],"k-",lw=0.8)
ax.text(1.0,0.5,r"$\Omega(0,\eta)=1$",rotation=90,va="center",fontsize=9)
ax.annotate("",xy=(1.9,1.32),xytext=(0.9,1.32),arrowprops=dict(arrowstyle="->",lw=1.2))
ax.text(1.4,1.4,r"$-\partial p/\partial x>0$",ha="center",fontsize=9)
ax.text(6.9,0.52,"Gordon–Schowalter fluid\n"+r"$r_A=k_{v0}e^{\gamma_A\Theta}\,C_A$"+"\nviscous heating "+r"$\Phi=\tau\dot\gamma$",
        ha="center",va="center",fontsize=9)
ax.text(0.55,0.5,r"$H_0$",rotation=90,va="center",fontsize=10)
ax.annotate("",xy=(0.72,1.0),xytext=(0.72,0.0),arrowprops=dict(arrowstyle="<->",lw=0.9))
ax.legend(loc="upper right",fontsize=8,frameon=False,bbox_to_anchor=(1.0,1.32))
plt.tight_layout(); plt.savefig("fig1_schematic.png",bbox_inches="tight"); plt.close()

# ================= FIGURE 2: shape functions & Theta_max(a) =================
fig,ax=plt.subplots(1,2,figsize=(7.6,3.0))
ee=np.linspace(0,1,400)
ax[0].plot(ee,thB(ee)/thB(0.5),"b-",lw=1.6,label=r"$\theta_B/\theta_B(\frac{1}{2})$")
ax[0].plot(ee,thBW(ee)/thBW(0.5),"r--",lw=1.6,label=r"$\theta_{BW}/\theta_{BW}(\frac{1}{2})$")
ax[0].plot(ee,phi1(ee)/phi1(0.5),"g-.",lw=1.6,label=r"$\varphi_1/\varphi_1(\frac{1}{2})$")
ax[0].set_xlabel(r"$\eta$"); ax[0].set_ylabel("normalised shape")
ax[0].set_title("(a) Shape functions (even in $s$)",fontsize=10)
ax[0].legend(fontsize=8,frameon=False); ax[0].set_xlim(0,1); ax[0].set_ylim(0,1.05)
aa=np.linspace(-1,1,401)
for Ws,st in [(0.1,"b-"),(0.2,"g--"),(0.3,"r-.")]:
    amp=1.0+0.1*(1-aa**2)*Ws**2*P**2
    ax[1].plot(aa,amp,st,lw=1.6,label=fr"$Ws={Ws}$")
ax[1].axhline(1.0,color="0.6",lw=0.8)
ax[1].plot([-1,1],[1,1],"ko",ms=5)
ax[1].set_xlabel(r"$a$"); ax[1].set_ylabel(r"$\Theta_{\max}/\Theta_{\max}^{\rm Newt}$")
ax[1].set_title(r"(b) Centreline amplification, $P^*=3$ (Result 1)",fontsize=10)
ax[1].legend(fontsize=8,frameon=False)
plt.tight_layout(); plt.savefig("fig2_thermal.png",bbox_inches="tight"); plt.close()

# ================= FIGURE 3: collapse + routes =================
fig,ax=plt.subplots(1,2,figsize=(7.6,3.0))
eps_line=np.linspace(0,0.0660,50)
mk={0.0:"o",0.3:"s",0.5:"^",0.7:"D",0.9:"v"}
for a in [0.0,0.3,0.5,0.7,0.9]:
    pts=[g for g in grid if g["a"]==a]
    ax[0].plot([g["eps"] for g in pts],[g["dl_exact"] for g in pts],mk[a],ms=5,
               mfc="none",label=fr"$a={a}$ (exact)")
ax[0].plot(eps_line,(c_T+c_A)*eps_line,"k-",lw=1.4,label="first order, eq. (15)")
ax[0].set_xlabel(r"$\varepsilon=(1-a^2)\,Ws^2$"); ax[0].set_ylabel(r"$\delta\lambda_1^2$")
ax[0].set_title("(a) Collapse onto the similarity variable",fontsize=10)
ax[0].legend(fontsize=7,frameon=False)
ax[1].plot(eps_line,c_T*eps_line,"r-",lw=1.6,label=r"thermal–reactive $\delta\lambda^2_{TR}$")
ax[1].plot(eps_line,c_A*eps_line,"b--",lw=1.6,label=r"advective $\delta\lambda^2_{adv}$")
ax[1].plot(eps_line,(c_T+c_A)*eps_line,"k-",lw=1.4,label="net")
ax[1].axhline(0,color="0.6",lw=0.8)
ax[1].set_xlabel(r"$\varepsilon$"); ax[1].set_ylabel(r"$\delta\lambda_1^2$ by route")
ax[1].set_title(f"(b) Route decomposition ($K={K:.3f}$)",fontsize=10)
ax[1].legend(fontsize=8,frameon=False)
plt.tight_layout(); plt.savefig("fig3_collapse.png",bbox_inches="tight"); plt.close()

# ================= FIGURE 4: K design map =================
Das=np.logspace(np.log10(0.2),np.log10(10),22)
F=[]
for Da in Das:
    lb=lam1_of(Pe,lambda e:Da,phi0); _,Hb=eigfun(lb,Pe,lambda e:Da,phi0)
    Nb=integ(phi0(eta)*Hb**2); Ab=integ(phi1(eta)*Hb**2); Tb=integ(thBW(eta)*Hb**2)
    F.append(Da*Tb/(lb**2*Pe*Ab))
F=np.array(F)
BG=np.logspace(np.log10(0.5),np.log10(200),200)
DD,GG=np.meshgrid(Das,BG)
KK=GG*np.interp(DD[0],Das,F)[None,:]
fig,ax=plt.subplots(figsize=(4.6,3.4))
cs=ax.contourf(DD,GG,np.log10(KK),levels=np.linspace(-2.5,1.0,15),cmap="RdBu_r")
c1=ax.contour(DD,GG,KK,levels=[1.0],colors="k",linewidths=2.0)
ax.clabel(c1,fmt={1.0:"K = 1"},fontsize=9)
c2=ax.contour(DD,GG,KK,levels=[0.1],colors="k",linewidths=0.9,linestyles="--")
ax.clabel(c2,fmt={0.1:"K = 0.1"},fontsize=8)
# Table-3 verification points
for (da,bg,sgn) in [(1,2,"-"),(1,60,"-"),(5,30,"-"),(5,60,"+")]:
    ax.plot(da,bg,"k"+("s" if sgn=="+" else "o"),ms=7,mfc=("k" if sgn=="+" else "none"))
ax.set_xscale("log"); ax.set_yscale("log")
ax.set_xlabel(r"$Da_{II}$"); ax.set_ylabel(r"$Br\,\gamma_A$")
ax.set_title(r"$K$ map ($Pe_M=5$, $P^*=3$); $\circ$/$\blacksquare$: verified sign",fontsize=9)
cb=plt.colorbar(cs,ax=ax); cb.set_label(r"$\log_{10}K$",fontsize=9)
ax.text(0.30,1.0,"advection-dominated\n(K<1): longer reactor",fontsize=7.5,va="bottom")
ax.text(2.6,120,"kinetics-dominated\n(K>1): shorter reactor",fontsize=7.5,ha="center")
plt.tight_layout(); plt.savefig("fig4_Kmap.png",bbox_inches="tight"); plt.close()

json.dump(dict(fd_worst=fd_worst,lam10=lam10,c_T=c_T,c_A=c_A,K=K,dl_B=dl_B,
               worst_err_pc=worst,ucm_dev=ucm_dev,grid=grid,
               F_of_Da={f"{d:.3f}":f for d,f in zip(Das,F)}),
          open("figdata.json","w"),indent=1)
print("\nFiguras 1-4 generadas. F(Da=1)=%.4f, F(Da=5)=%.4f"%(np.interp(1,Das,F),np.interp(5,Das,F)))