#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
validate.py
=====================
Numerical validation suite for the manuscript:

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

Purpose
-------
This script reproduces the main validation results reported in the manuscript:

S0  Sign adjudication of the viscoelastic thermal correction using the exact
    rational Johnson-Segalman/Gordon-Schowalter shear law at prescribed stress.
S1  Cross-check of the Sturm-Liouville shooting solver against an independent
    finite-difference/generalised-eigenvalue implementation.
S2  Closed-form temperature fields theta_B and theta_BW against direct
    double-quadrature integration of the local energy balance.
S3  Exact perturbed fundamental eigenvalue versus the first-order formula,
    including route decomposition, additivity, UCM/LCM invariance and the
    crossover criterion K.

Conventions
-----------
  phi0(eta)  = 1 + (P/2) eta (1-eta)
  tau(eta)   = -P s,  s = eta - 1/2
  JS law     : tau = gdot / (1 + eps_w * gdot^2),
               eps_w = (1-a^2) Ws^2
  Theta      : d2Theta/deta2 = -Br * tau * gdot
  Reactive   : phi dOmega/dxi = (1/Pe) d2Omega/deta2
               - (Da(eta)/Pe) Omega,
               Da(eta) = Da0 * exp(gammaA * Theta(eta))
  Modal      : H'' - (Da(eta) - lam^2 Pe phi(eta)) H = 0,
               H(0)=H(1)=0

Authors
-------
Leonardo D. Soria R. and Anthony A. Harrup G.
Target journal: Journal of Engineering Mathematics
Version: v2-clean, June 2026
"""

import json
import numpy as np
from scipy.integrate import solve_ivp, simpson
from scipy.optimize import brentq
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

RNG_TOL = dict(rtol=1e-11, atol=1e-13)
REPORT = []


def log(*args):
    line = " ".join(str(a) for a in args)
    print(line)
    REPORT.append(line)


# ----------------------------------------------------------------------
# S0. SIGN ADJUDICATION OF THE VISCOELASTIC THERMAL CORRECTION
# ----------------------------------------------------------------------
# Exact inversion of the rational JS law at prescribed stress:
#   tau = g/(1+e g^2)  =>  e*tau*g^2 - g + tau = 0
#   g = (1 - sqrt(1 - 4 e tau^2)) / (2 e tau)   (branch -> Newtonian as e->0)

def gdot_exact_JS(tau, eps_w):
    if eps_w == 0.0:
        return tau
    tau = np.asarray(tau, dtype=float)
    out = np.empty_like(tau)
    small = np.abs(tau) < 1e-12
    disc = 1.0 - 4.0 * eps_w * tau**2
    if np.any(disc < 0):
        raise ValueError("Beyond JS flow-curve maximum: 4*eps*tau^2 > 1.")
    out[~small] = (1.0 - np.sqrt(disc[~small])) / (2.0 * eps_w * tau[~small])
    out[small] = tau[small]
    return out


def theta_exact_JS(P, Br, eps_w, N=4001):
    """Integrate d2T/deta2 = -Br tau gdot, T(0)=T(1)=0, by double quadrature."""
    eta = np.linspace(0.0, 1.0, N)
    s = eta - 0.5
    tau = -P * s
    Phi = Br * tau * gdot_exact_JS(tau, eps_w)
    # T(eta) = -int_0^eta int_0^y Phi + C*eta ; enforce T(1)=0
    inner = np.concatenate([[0.0], np.cumsum((Phi[1:] + Phi[:-1]) / 2) * np.diff(eta)])
    dbl = np.concatenate([[0.0], np.cumsum((inner[1:] + inner[:-1]) / 2) * np.diff(eta)])
    C = dbl[-1]
    return eta, -dbl + C * eta


def theta_pert(eta, P, Br, eps_w, sign=+1):
    """Closed-form perturbative Theta with cross-term sign = +1 (ours) or -1 (Art.1)."""
    s = eta - 0.5
    tB = (P**2 / 12.0) * (1.0 / 16.0 - s**4)
    tBW = (P**4 / 30.0) * (1.0 / 64.0 - s**6)
    return Br * tB + sign * Br * eps_w * tBW


def section0():
    log("=" * 78)
    log("S0  SIGN ADJUDICATION: exact rational JS vs perturbative candidates")
    log("    Config: stress-controlled Couette-Poiseuille, P=2, Br=1, a=0 (Jaumann)")
    log("-" * 78)
    P, Br = 2.0, 1.0
    log(f"{'Ws':>6} {'eps_w':>8} {'Tmax_exactJS':>14} {'Tmax_Newt':>12} "
        f"{'pert(+)':>12} {'pert(-)':>12} {'err(+)%':>9} {'err(-)%':>9}")
    rows = []
    for Ws in [0.1, 0.2, 0.3, 0.4]:
        eps_w = Ws**2  # a = 0
        eta, T = theta_exact_JS(P, Br, eps_w)
        Tmax = T[len(T) // 2]
        Tnewt = Br * P**2 / 192.0
        Tp = theta_pert(np.array([0.5]), P, Br, eps_w, +1)[0]
        Tm = theta_pert(np.array([0.5]), P, Br, eps_w, -1)[0]
        ep = 100 * abs(Tp - Tmax) / Tmax
        em = 100 * abs(Tm - Tmax) / Tmax
        log(f"{Ws:6.2f} {eps_w:8.4f} {Tmax:14.8f} {Tnewt:12.8f} "
            f"{Tp:12.8f} {Tm:12.8f} {ep:9.4f} {em:9.4f}")
        rows.append(dict(Ws=Ws, Tmax_exact=Tmax, Tnewt=Tnewt, pert_plus=Tp,
                         pert_minus=Tm, err_plus_pc=ep, err_minus_pc=em))
    log("-" * 78)
    inc = all(r["Tmax_exact"] > r["Tnewt"] for r in rows)
    log(f"VERDICT: exact JS Theta_max {'EXCEEDS' if inc else 'is below'} the Newtonian "
        f"value at fixed pressure gradient for |a|<1.")
    log("         => correct perturbative sign is '+': Theta = Pi f0 + (1-a^2) Psi f1.")
    log("")
    return rows


# ----------------------------------------------------------------------
# Sturm-Liouville machinery (shooting + Brent)
# ----------------------------------------------------------------------

def shoot_residual(lam, Pe, Da_fun, phi_fun, symmetric):
    """Integrate transverse ODE; return H at far wall."""
    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, dense_output=False)
    else:
        sol = solve_ivp(rhs, [0.0, 1.0], [0.0, 1.0], **RNG_TOL, dense_output=False)
    return sol.y[0, -1]


def first_eigenvalues(Pe, Da_fun, phi_fun, n=1, lam_max=40.0, dlam=0.02,
                      symmetric=True):
    """Bracket sign changes of the shooting residual, refine with Brent."""
    lams, found = [], []
    grid = np.arange(dlam, lam_max, dlam)
    r_prev = shoot_residual(grid[0], Pe, Da_fun, phi_fun, symmetric)
    for lam in grid[1:]:
        r = shoot_residual(lam, Pe, Da_fun, phi_fun, symmetric)
        if r_prev * r < 0:
            root = brentq(shoot_residual, lam - dlam, lam,
                          args=(Pe, Da_fun, phi_fun, symmetric),
                          xtol=1e-12, rtol=1e-13)
            found.append(root)
            if len(found) >= n:
                return found
        r_prev = r
    return found


def eigenfunction(lam, Pe, Da_fun, phi_fun, N=2001, symmetric=True):
    """Return (eta, H) on a uniform grid, H normalized to max=1."""
    eta = np.linspace(0.0, 1.0, N)

    def rhs(t, y):
        return [y[1], (Da_fun(t) - lam**2 * Pe * phi_fun(t)) * y[0]]
    if symmetric:
        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])  # mirror
        if len(H) != N:  # N odd guarantees exact mirroring
            raise RuntimeError("Use odd N for symmetric reconstruction.")
    else:
        sol = solve_ivp(rhs, [0.0, 1.0], [0.0, 1.0], t_eval=eta, **RNG_TOL)
        H = sol.y[0]
    return eta, H / np.max(np.abs(H))


# ----------------------------------------------------------------------
# S1. Cross-check against an independent finite-difference solver
# ----------------------------------------------------------------------

def section1():
    log("=" * 78)
    log("S1  SOLVER CROSS-CHECK vs independent finite-difference eigenproblem")
    log("-" * 78)
    checks = []
    # Table 1, linear case phi=eta (non-symmetric modes): (Pe, Da) -> lam1..3
    t1 = {(1, 0): [4.354, 9.049, 13.756], (5, 1): [2.040, 4.098, 6.188],
          (10, 0): [1.377, 2.862, 4.350], (50, 1): [0.645, 1.296, 1.957]}
    for (Pe, Da), ref in t1.items():
        lams = first_eigenvalues(Pe, lambda e: float(Da), lambda e: e, n=3,
                                 lam_max=ref[-1] + 2, dlam=0.05, symmetric=False)
        err = max(abs(a - b) / b for a, b in zip(lams, ref))
        log(f"  linear   Pe={Pe:>3} Da={Da}      computed={['%.3f' % x for x in lams]}"
            f"  ref={ref}  max rel err={err:.2e}")
        checks.append(err)
    # Table 2, parabolic, Art-2 convention phi = 1 + (P*/2)(eta^2-eta), Pe=1, Da=0
    t2 = {-2: 2.847, 0: 3.142, +2: 3.551}
    for Pstar2, ref in t2.items():
        phi = lambda e, p=Pstar2: 1.0 + (p / 2.0) * (e**2 - e)
        lam1 = first_eigenvalues(1, lambda e: 0.0, phi, n=1, lam_max=6,
                                 dlam=0.05, symmetric=True)[0]
        err = abs(lam1 - ref) / ref
        log(f"  parabolic Pe=1 Da=0 P*_A2={Pstar2:+d}  computed={lam1:.3f}  "
            f"ref={ref}  rel err={err:.2e}")
        checks.append(err)
    ok = max(checks) < 5e-3
    log(f"  -> solver {'REPRODUCES' if ok else 'FAILS'} finite-difference benchmark "
        f"(worst {max(checks):.2e}).")
    log("")
    return max(checks)


# ----------------------------------------------------------------------
# S2. Closed-form Theta check
# ----------------------------------------------------------------------

def section2():
    log("=" * 78)
    log("S2  CLOSED-FORM THERMAL FIELD vs direct integration (consistent order)")
    log("-" * 78)
    P, Br = 3.0, 1.0
    eta = np.linspace(0, 1, 4001)
    s = eta - 0.5
    worst = 0.0
    for (a, Ws) in [(0.0, 0.2), (0.5, 0.3), (0.9, 0.25)]:
        eps_w = (1 - a**2) * Ws**2
        # consistent-order dissipation: Phi = tau^2 + eps tau^4 with tau=-Ps
        tau = -P * s
        Phi = Br * (tau**2 + eps_w * tau**4)
        inner = np.concatenate([[0], np.cumsum((Phi[1:] + Phi[:-1]) / 2) * np.diff(eta)])
        dbl = np.concatenate([[0], np.cumsum((inner[1:] + inner[:-1]) / 2) * np.diff(eta)])
        Tnum = -dbl + dbl[-1] * eta
        Tcf = theta_pert(eta, P, Br, eps_w, +1)
        err = np.max(np.abs(Tnum - Tcf)) / np.max(Tcf)
        worst = max(worst, err)
        log(f"  a={a:4.2f} Ws={Ws:4.2f}  max|Theta_num - Theta_cf|/Theta_max = {err:.2e}")
    log(f"  -> closed forms theta_B=(P^2/12)(1/16-s^4), theta_BW=(P^4/30)(1/64-s^6) "
        f"verified ({worst:.1e}).")
    log("")
    return worst


# ----------------------------------------------------------------------
# S3. Eigenvalue perturbation: exact vs first-order, routes, collapse, K
# ----------------------------------------------------------------------

def section3():
    log("=" * 78)
    log("S3  CROSS-TERM VALIDATION: lambda_1^2(eps) exact vs perturbation")
    log("-" * 78)
    Pe, Da0, P, Br, gA = 5.0, 1.0, 3.0, 1.0, 5.0
    log(f"    Pe={Pe} Da0={Da0} P={P} Br={Br} gammaA={gA}")
    phi0 = lambda e: 1.0 + (P / 2.0) * e * (1 - e)
    phi1f = lambda e: (P**3 / 4.0) * (1.0 / 16.0 - (e - 0.5)**4)
    thBf = lambda e: (P**2 / 12.0) * (1.0 / 16.0 - (e - 0.5)**4)
    thBWf = lambda e: (P**4 / 30.0) * (1.0 / 64.0 - (e - 0.5)**6)

    # Baseline mode and integrals
    lam10 = first_eigenvalues(Pe, lambda e: Da0, phi0, n=1, lam_max=8, dlam=0.05)[0]
    eta, H1 = eigenfunction(lam10, Pe, lambda e: Da0, phi0)
    H2 = H1**2
    N1 = simpson(phi0(eta) * H2, x=eta)
    A = simpson(phi1f(eta) * H2, x=eta)
    TB = simpson(thBf(eta) * H2, x=eta)
    TBW = simpson(thBWf(eta) * H2, x=eta)
    log(f"    lambda_1(base) = {lam10:.8f}   lambda_1^2 = {lam10**2:.8f}")
    log(f"    <phi1> = {A:.6e}  <thB> = {TB:.6e}  <thBW> = {TBW:.6e}  N1 = {N1:.6e}")

    dl_B = Br * gA * Da0 * TB / (Pe * N1)            # O(Br), a-independent
    c_thermal = Br * gA * Da0 * TBW / (Pe * N1)       # cross, per unit eps
    c_advect = -lam10**2 * A / N1                      # advective, per unit eps
    K = c_thermal / (-c_advect)
    log(f"    Newtonian-thermal shift  dl_B            = {dl_B:+.6e}")
    log(f"    cross thermal-reactive   slope (per eps)  = {c_thermal:+.6e}")
    log(f"    advective                slope (per eps)  = {c_advect:+.6e}")
    log(f"    predicted net slope d(lam^2)/d(eps)       = {c_thermal + c_advect:+.6e}")
    log(f"    crossover number K = {K:.4f}  "
        f"({'thermal-reactive dominates' if K > 1 else 'advective dominates'})")
    log("")

    # --- exact eigenvalues over (a, Ws) grid, three model variants -----------
    a_grid = [0.0, 0.3, 0.5, 0.7, 0.9, 1.0]
    Ws_grid = [0.05, 0.10, 0.15, 0.20, 0.25]
    results = []
    log(f"{'a':>5} {'Ws':>5} {'eps':>8} {'dl2_full':>12} {'dl2_pert':>12} "
        f"{'err%':>7} {'dl2_adv':>12} {'dl2_thm':>12} {'add_err%':>9}")
    for a in a_grid:
        for Ws in Ws_grid:
            eps = (1 - a**2) * Ws**2
            phi_e = lambda e: phi0(e) + eps * phi1f(e)
            Th = lambda e: Br * thBf(e) + Br * eps * thBWf(e)
            Da_full = lambda e: Da0 * (1 + gA * Th(e))
            Da_thm = lambda e: Da0 * (1 + gA * Br * (thBf(e) + eps * thBWf(e)))
            Da_B = lambda e: Da0 * (1 + gA * Br * thBf(e))
            lam_full = first_eigenvalues(Pe, Da_full, phi_e, n=1, lam_max=8,
                                         dlam=0.05)[0]
            lam_adv = first_eigenvalues(Pe, Da_B, phi_e, n=1, lam_max=8,
                                        dlam=0.05)[0]
            lam_thm = first_eigenvalues(Pe, Da_thm, phi0, n=1, lam_max=8,
                                        dlam=0.05)[0]
            lam_base = first_eigenvalues(Pe, Da_B, phi0, n=1, lam_max=8,
                                         dlam=0.05)[0]
            dl_full = lam_full**2 - lam_base**2
            dl_adv = lam_adv**2 - lam_base**2
            dl_thm = lam_thm**2 - lam_base**2
            dl_pred = (c_thermal + c_advect) * eps
            err = (100 * abs(dl_full - dl_pred) / abs(dl_pred)
                   if eps > 0 else 0.0)
            add = (100 * abs(dl_full - dl_adv - dl_thm) / max(abs(dl_full), 1e-30)
                   if eps > 0 else 0.0)
            log(f"{a:5.2f} {Ws:5.2f} {eps:8.5f} {dl_full:+12.3e} {dl_pred:+12.3e} "
                f"{err:7.3f} {dl_adv:+12.3e} {dl_thm:+12.3e} {add:9.4f}")
            results.append(dict(a=a, Ws=Ws, eps=eps, dl_full=dl_full,
                                dl_pred=dl_pred, dl_adv=dl_adv, dl_thm=dl_thm,
                                lam_base2=lam_base**2))
    # UCM/LCM invariance check
    ucm = [r for r in results if r["a"] == 1.0]
    inv = max(abs(r["dl_full"]) for r in ucm)
    log("-" * 78)
    log(f"  UCM (a=1) viscoelastic shift max |dl2| = {inv:.2e}  "
        f"(must vanish: parity of GS correction).")

    # collapse figure -------------------------------------------------------
    fig, ax = plt.subplots(1, 2, figsize=(11, 4.2))
    mk = dict(zip(a_grid, ["o", "s", "^", "v", "D", "x"]))
    for a in a_grid:
        pts = [r for r in results if r["a"] == a and r["eps"] > 0]
        if pts:
            ax[0].plot([r["eps"] for r in pts], [r["dl_full"] for r in pts],
                       mk[a], ms=5, label=f"a={a}")
    e = np.linspace(0, max(r["eps"] for r in results), 50)
    ax[0].plot(e, (c_thermal + c_advect) * e, "k-", lw=1,
               label="first-order prediction")
    ax[0].set_xlabel(r"$\varepsilon=(1-a^2)\,\mathrm{Ws}^2$")
    ax[0].set_ylabel(r"$\delta\lambda_1^2$ (full)")
    ax[0].set_title("Collapse of the viscoelastic eigenvalue shift")
    ax[0].legend(fontsize=8)
    pts = sorted([r for r in results if r["eps"] > 0], key=lambda r: r["eps"])
    ax[1].plot([r["eps"] for r in pts], [r["dl_adv"] for r in pts], "b.",
               label="advective route")
    ax[1].plot([r["eps"] for r in pts], [r["dl_thm"] for r in pts], "r.",
               label="thermal-reactive route")
    ax[1].plot(e, c_advect * e, "b-", lw=0.8)
    ax[1].plot(e, c_thermal * e, "r-", lw=0.8)
    ax[1].axhline(0, color="k", lw=0.5)
    ax[1].set_xlabel(r"$\varepsilon$")
    ax[1].set_ylabel(r"$\delta\lambda_1^2$ by route")
    ax[1].set_title(f"Route decomposition (K = {K:.3f})")
    ax[1].legend(fontsize=8)
    fig.tight_layout()
    fig.savefig("fig_collapse.png", dpi=160)
    log("  Figure saved: fig_collapse.png")
    log("")
    return dict(lam10=lam10, K=K, c_thermal=c_thermal, c_advect=c_advect,
                dl_B=dl_B, results=results, ucm_invariance=inv)


# ----------------------------------------------------------------------
# S3b. Crossover demonstration: sweep K through 1
# ----------------------------------------------------------------------

def section3b():
    log("=" * 78)
    log("S3b CROSSOVER NUMBER K: sign reversal of the net viscoelastic effect")
    log("-" * 78)
    Pe, P, Ws, a = 5.0, 3.0, 0.25, 0.0
    eps = (1 - a**2) * Ws**2
    phi0 = lambda e: 1.0 + (P / 2.0) * e * (1 - e)
    phi1f = lambda e: (P**3 / 4.0) * (1.0 / 16.0 - (e - 0.5)**4)
    thBf = lambda e: (P**2 / 12.0) * (1.0 / 16.0 - (e - 0.5)**4)
    thBWf = lambda e: (P**4 / 30.0) * (1.0 / 64.0 - (e - 0.5)**6)
    log(f"    fixed: Pe={Pe} P={P} a={a} Ws={Ws} (eps={eps:.4f}); sweep Br*gammaA, Da0")
    log(f"{'Da0':>5} {'Br*gA':>7} {'K':>8} {'dl2_net_exact':>15} {'sign':>6}")
    sweep = []
    for Da0 in [1.0, 5.0]:
        for BrgA in [2.0, 10.0, 30.0, 60.0]:
            Br, gA = 1.0, BrgA
            lam_b = first_eigenvalues(Pe, lambda e: Da0 * (1 + gA * Br * thBf(e)),
                                      phi0, n=1, lam_max=10, dlam=0.05)[0]
            eta, H1 = eigenfunction(lam_b, Pe,
                                    lambda e: Da0 * (1 + gA * Br * thBf(e)), phi0)
            H2 = H1**2
            N1 = simpson(phi0(eta) * H2, x=eta)
            A = simpson(phi1f(eta) * H2, x=eta)
            TBW = simpson(thBWf(eta) * H2, x=eta)
            K = Br * gA * Da0 * TBW / (lam_b**2 * Pe * A)
            Da_full = lambda e: Da0 * (1 + gA * Br * (thBf(e) + eps * thBWf(e)))
            phi_e = lambda e: phi0(e) + eps * phi1f(e)
            lam_f = first_eigenvalues(Pe, Da_full, phi_e, n=1, lam_max=10,
                                      dlam=0.05)[0]
            dl = lam_f**2 - lam_b**2
            sgn = "+" if dl > 0 else "-"
            log(f"{Da0:5.1f} {BrgA:7.1f} {K:8.4f} {dl:+15.6e} {sgn:>6}")
            sweep.append(dict(Da0=Da0, BrgA=BrgA, K=K, dl_net=dl))
    flips = [(r["K"] > 1) == (r["dl_net"] > 0) for r in sweep]
    log(f"  -> K>1 <=> net shift positive: {'CONFIRMED' if all(flips) else 'FAILED'} "
        f"on all {len(sweep)} cases.")
    log("")
    return sweep


if __name__ == "__main__":
    s0 = section0()
    s1 = section1()
    s2 = section2()
    s3 = section3()
    s3b = section3b()
    log("=" * 78)
    log("SUMMARY")
    log(f"  S0 sign verdict: '+' branch correct; opposite sign error grows "
        f"with Ws (see table).")
    log(f"  S1 shooting vs finite-difference benchmark: worst rel err {s1:.2e}")
    log(f"  S2 closed-form Theta: worst rel err {s2:.2e}")
    log(f"  S3 cross-term: K = {s3['K']:.4f}; UCM invariance {s3['ucm_invariance']:.1e}")
    with open("results.json", "w") as f:
        json.dump(dict(S0=s0, S3=dict(K=s3["K"], c_thermal=s3["c_thermal"],
                                      c_advect=s3["c_advect"], dl_B=s3["dl_B"],
                                      grid=s3["results"]), S3b=s3b),
                  f, indent=1)
    with open("report.txt", "w") as f:
        f.write("\n".join(REPORT))
