"""
fea_validation.py
=================

Finite-element validation of the analytical FGM stress model, using the
free CalculiX solver (ccx, v2.21). Reproduces the FEA figures of the paper
and the elasto-plastic relocation checks:

    Figure 11  Through-stack membrane stress relocation (LPBF)
    Figure 12  Thickness independence (substrate + film sweeps)
    + single-element elasto-plastic confirmation (LPBF 1300 C and anneal 400 C)

Requires:  CalculiX `ccx` on PATH, plus the analytical engine fgm_model.py
(for material constants and the optimised stack), NumPy and Matplotlib.

    $ python fea_validation.py
"""
from __future__ import annotations
import os, subprocess, pickle
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from matplotlib.lines import Line2D
import fgm_model as M

DT_LPBF, DT_ANNEAL = 1200.0, 380.0
H_SUB_BASE = 0.15e-3          # baseline substrate thickness (m)

# ---- stack from the analytical engine ---------------------------------------
_stack = M.build_stack(5, optimised=True)
_films = [(c, m, t * 1e-9) for c, m, t in _stack if c != "Metal"]
FILM_LAYERS = _films[1:]
SUB_MAT = _films[0][1]
T0 = np.array([t for _, _, t in FILM_LAYERS])
CODES = [c for c, _, _ in FILM_LAYERS]


def _biax(mat):
    return mat.E / (1.0 - mat.nu)


# ---- generic 2D plane-strain stack model (graded mesh) ----------------------
def run_stack(h_sub, dT, tag, t_film=T0, W=30e-6, edge_refine=False, keep=False):
    yi = np.concatenate([[0.0], np.cumsum(t_film)])
    yn, ns, r = [], 18, 1.35
    sz = np.array([r ** i for i in range(ns)]); sz = sz / sz.sum() * h_sub
    y = -h_sub; yn.append(y)
    for s in sz[::-1]:
        y += s; yn.append(y)
    yn[-1] = 0.0
    for k in range(len(FILM_LAYERS)):
        n = max(4, int(round(t_film[k] / 40e-9)))
        for i in range(1, n + 1):
            yn.append(yi[k] + t_film[k] * i / n)
    yn = np.array(sorted(set(np.round(yn, 13))))
    if edge_refine:
        xc = np.linspace(0, W - 1e-6, 34)
        xf = W - 1e-6 + (np.linspace(0, 1, 18) ** 2) * 1e-6
        xn = np.unique(np.concatenate([xc, xf]))
    else:
        xn = (np.linspace(0, 1, 41) ** 1.3) * W
    NX, NY = len(xn), len(yn)
    nid = lambda i, j: j * NX + i + 1
    nodes = [(nid(i, j), xn[i], yn[j]) for j in range(NY) for i in range(NX)]

    def mat_at(yc):
        if yc < 0: return "SUB"
        for k in range(len(FILM_LAYERS)):
            if yi[k] - 1e-15 <= yc < yi[k + 1] + 1e-15: return FILM_LAYERS[k][0]
        return FILM_LAYERS[-1][0]

    els, esets, ecent, e = [], {}, {}, 1
    for j in range(NY - 1):
        yc = 0.5 * (yn[j] + yn[j + 1]); ms = mat_at(yc)
        for i in range(NX - 1):
            els.append((e, nid(i, j), nid(i + 1, j), nid(i + 1, j + 1), nid(i, j + 1)))
            ecent[e] = (0.5 * (xn[i] + xn[i + 1]), yc); esets.setdefault(ms, []).append(e); e += 1
    mm = {"SUB": SUB_MAT}
    for c, m, t in FILM_LAYERS: mm[c] = m

    fn = f"_fea_{tag}"
    with open(fn + ".inp", "w") as f:
        f.write("*NODE\n")
        for n, x, yv in nodes: f.write(f"{n},{x:.10e},{yv:.10e}\n")
        f.write("*ELEMENT,TYPE=CPE4,ELSET=ALL\n")
        for el in els: f.write(",".join(map(str, el)) + "\n")
        for ms, e_ in esets.items():
            f.write(f"*ELSET,ELSET=E_{ms}\n")
            for k in range(0, len(e_), 8): f.write(",".join(map(str, e_[k:k + 8])) + "\n")
        for ms, m in mm.items():
            f.write(f"*MATERIAL,NAME=M_{ms}\n*ELASTIC\n{m.E*1e9:.6e},{m.nu:.3f}\n"
                    f"*EXPANSION,ZERO=20.\n{m.alpha*1e-6:.6e}\n"
                    f"*SOLID SECTION,ELSET=E_{ms},MATERIAL=M_{ms}\n")
        left = [nid(0, j) for j in range(NY)]; bot = [nid(i, 0) for i in range(NX)]
        f.write("*NSET,NSET=NALL,GENERATE\n1,%d\n" % len(nodes))
        f.write("*NSET,NSET=LEFT\n")
        for k in range(0, len(left), 8): f.write(",".join(map(str, left[k:k + 8])) + "\n")
        f.write("*NSET,NSET=BOT\n")
        for k in range(0, len(bot), 8): f.write(",".join(map(str, bot[k:k + 8])) + "\n")
        f.write("*INITIAL CONDITIONS,TYPE=TEMPERATURE\nNALL,20.\n"
                "*STEP\n*STATIC\n*BOUNDARY\nLEFT,1,1,0.\nBOT,2,2,0.\n")
        f.write("*TEMPERATURE\nNALL,%.4f\n*EL PRINT,ELSET=ALL\nS\n*END STEP\n" % (20.0 - dT))
    subprocess.run(["ccx", fn], capture_output=True, text=True)

    per, grab = {}, False
    for ln in open(fn + ".dat"):
        s = ln.strip()
        if "stresses" in s.lower(): grab = True; continue
        if grab:
            p = s.split()
            if len(p) >= 8 and p[0].isdigit():
                per.setdefault(int(p[0]), []).append(list(map(float, p[2:8])))
            elif s and not p[0].isdigit(): grab = False
    st = {el: np.mean(v, 0) for el, v in per.items()}
    if not keep:
        for ext in (".inp", ".dat", ".frd", ".cvg", ".sta", ".12d"):
            try: os.remove(fn + ext)
            except OSError: pass
    return ecent, st, yi, W


def _central_layer_means(ecent, st, yi):
    els = sorted(ecent)
    xc = np.array([ecent[e][0] for e in els]); yc = np.array([ecent[e][1] for e in els])
    sxx = np.array([st[e][0] for e in els]) / 1e9
    xmin = np.unique(xc)[:3].max()
    out = []
    for k in range(1, len(yi)):
        sel = (xc <= xmin) & (yc >= yi[k - 1]) & (yc < yi[k])
        out.append((CODES[k - 1], float(np.mean(sxx[sel])), 0.5 * (yi[k - 1] + yi[k]) * 1e6))
    return out


# ---- Figure 11: relocation profile ------------------------------------------
def fig_11_relocation(outfile="fig_09_relocation.png"):
    ecent, st, yi, W = run_stack(H_SUB_BASE, DT_LPBF, "f11")
    prof = _central_layer_means(ecent, st, yi)
    labs = [p[0] for p in prof]; vals = np.array([p[1] for p in prof]); ymid = np.array([p[2] for p in prof])

    fig, ax = plt.subplots(figsize=(7.2, 5.4))
    band = {"F1": "#dfe7f0", "F2": "#dfe7f0", "F3": "#dfe7f0", "F4": "#f4ead6", "F5": "#ecd7d0",
            "F6": "#e3efe7", "F7": "#d2e7da", "F8": "#bcdcc7", "F9": "#9fceaf", "F10": "#7dbd92", "F11": "#57a877"}
    for k in range(1, len(yi)):
        ax.axhspan(yi[k - 1] * 1e6, yi[k] * 1e6, color=band.get(CODES[k - 1], "#eee"), alpha=.7, zorder=0)
    yRT, yHOT = 0.35, 0.02
    ax.axvspan(-yRT, yRT, color="0.70", alpha=.28, zorder=0)
    ax.axvline(yHOT, color="#b03020", lw=1.0, ls=":", zorder=1)
    ax.axvline(-yHOT, color="#b03020", lw=1.0, ls=":", zorder=1)
    ax.axvline(0, color="k", lw=.8)
    ax.plot(vals, ymid, "o-", color="#0b3d2e", lw=2, ms=6, zorder=5)
    for v, y in zip(vals, ymid):
        ax.annotate(f"{v:+.2f}", (v, y), textcoords="offset points",
                    xytext=(10 if v >= 0 else -10, 0), ha="left" if v >= 0 else "right",
                    va="center", fontsize=7.5, color="#0b3d2e", fontweight="bold")
    ax.annotate("ceramic interface\nin compression", xy=(-0.32, 0.22),
                xytext=(0.5, 0.35), fontsize=8, color="#7a1f17",
                ha="left", va="center",
                arrowprops=dict(arrowstyle="->", color="#7a1f17", lw=0.9))
    ax.annotate("residual mismatch in\nductile Ni-Cr ($\\approx$100$\\times$ hot yield)\n"
                "$\\rightarrow$ relieved by plastic flow",
                xy=(vals[-1], ymid[-1]), xytext=(0.9, 5.55),
                fontsize=8, color="#0b3d2e", ha="center", va="center",
                arrowprops=dict(arrowstyle="->", color="#0b3d2e", lw=0.9,
                                connectionstyle="arc3,rad=0.0"))
    ax.set_xlabel("In-plane membrane stress  $\\sigma_{xx}$  (GPa)   ---   LPBF, $\\Delta T=1200$ K")
    ax.set_ylabel("Height in stack above sapphire  ($\\mu$m)")
    ax.set_xlim(-3, 3.2); ax.set_ylim(0, yi[-1] * 1e6)
    ax.set_title("Through-stack stress relocation (LPBF)", fontsize=10, fontweight="bold", pad=10)
    leg = [Line2D([0], [0], color="#0b3d2e", lw=2, marker="o", label="FEA membrane $\\sigma_{xx}$ (centre)"),
           Patch(facecolor="0.70", alpha=.5, label="within Ni-Cr yield at 25 $^\\circ$C ($\\pm$350 MPa)"),
           Line2D([0], [0], color="#b03020", lw=1.0, ls=":", label="Ni-Cr yield at ~1300 $^\\circ$C ($\\pm$20 MPa)")]
    ax.legend(handles=leg, loc="upper left", fontsize=7.5, framealpha=.95)
    ax2 = ax.twinx(); ax2.set_ylim(ax.get_ylim())
    ax2.set_yticks([0.5 * (yi[k - 1] + yi[k]) * 1e6 for k in range(1, len(yi))])
    ax2.set_yticklabels(CODES, fontsize=8, color="#333"); ax2.set_ylabel("FGM layer", fontsize=9)


    fig.tight_layout(); fig.savefig(outfile, dpi=300); plt.close(fig)
    print(f"  fig 11: F11 membrane = {vals[-1]:+.3f} GPa (laminate cross-check 2.29 GPa)")


# ---- Figure 12: thickness sweeps --------------------------------------------
def _f11_central(ecent, st, yi):
    els = sorted(ecent)
    xc = np.array([ecent[e][0] for e in els]); yc = np.array([ecent[e][1] for e in els])
    sxx = np.array([st[e][0] for e in els]) / 1e9
    xmin = np.unique(xc)[:3].max()
    sel = (xc <= xmin) & (yc >= yi[-2]) & (yc < yi[-1])
    return float(np.mean(sxx[sel]))


def fig_12_sweeps(outfile="fig_S4_sweeps.png"):
    subs = [0.006, 0.012, 0.025, 0.05, 0.1, 0.5, 1.0]
    sub = [(h, _f11_central(*run_stack(h * 1e-3, DT_LPBF, f"s{int(h*1e6)}")[:3])) for h in subs]
    films = [0.5, 1.0, 2.0, 3.0]
    flm = [(sc, _f11_central(*run_stack(1.0e-3, DT_LPBF, f"f{int(sc*10)}", t_film=T0 * sc)[:3])) for sc in films]

    hs = np.array([h * 1e3 for h, _ in sub]); ss = np.array([v for _, v in sub])
    fs = np.array([s for s, _ in flm]); sf = np.array([v for _, v in flm])

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10.4, 4.7))
    fig.suptitle("Thickness independence of the membrane stress", fontsize=10, fontweight="bold", y=0.99)
    ax1.semilogx(hs, ss, "o-", color="#0b3d2e", lw=2, ms=7)
    ax1.axhline(ss[-1], color="0.6", ls="--", lw=1)
    ax1.fill_between([hs.min(), hs.max()], ss[-1] * .98, ss[-1] * 1.02, color="#0b3d2e", alpha=.10, zorder=0)
    ax1.axvline(6.15, color="#b03020", ls=":", lw=1)
    ax1.text(6.6, 0.25, "$h_{sub}=h_{film}$", color="#b03020", fontsize=8, rotation=90, va="bottom")
    ax1.set_xlabel("Substrate thickness  $h_{sub}$  ($\\mu$m)")
    ax1.set_ylabel("Peak F11 membrane stress  (GPa)")
    ax1.set_title("(a)  Substrate-thickness independence", fontsize=10)
    ax1.set_ylim(0, 2.6); ax1.grid(alpha=.25, which="both")
    ax1.annotate("flat to $\\pm$2%  (6 $\\mu$m $\\to$ 1 mm)\nconstraint set by sapphire modulus",
                 (hs[2], ss[2]), xytext=(0.30, 0.30), textcoords="axes fraction", fontsize=8.5, color="#0b3d2e")
    ax2.plot(fs, sf, "s-", color="#1c5d8c", lw=2, ms=7)
    for x, y in zip(fs, sf):
        ax2.annotate(f"{y:.2f}", (x, y), textcoords="offset points", xytext=(0, 9), ha="center", fontsize=8, color="#1c5d8c")
    ax2.set_xlabel("Film thickness scale factor  ($\\times$ baseline)")
    ax2.set_ylabel("Peak F11 membrane stress  (GPa)")
    ax2.set_title("(b)  Film-thickness dependence", fontsize=10)
    ax2.set_ylim(0, 2.6); ax2.grid(alpha=.25)
    ax2.annotate("interface-step stress: exactly\nthickness-independent (by construction)\n"
                 "membrane: $-$33% over 0.5$\\to$3$\\times$\n(curvature relief, favourable)",
                 (fs[1], sf[1]), xytext=(0.30, 0.18), textcoords="axes fraction", fontsize=8.5, color="#1c5d8c")
    fig.tight_layout(rect=[0, 0, 1, 0.95]); fig.savefig(outfile, dpi=300); plt.close(fig)
    print(f"  fig 12: substrate flat {ss.min():.2f}-{ss.max():.2f} GPa; film {sf[0]:.2f}->{sf[-1]:.2f} GPa")


# ---- single-element elasto-plastic relocation check -------------------------
def plastic_check(temp_label, sigma_y, elastic_equiv, tag):
    E, nu, L = 220e9, 0.31, 1e-6
    eps = elastic_equiv * (1 - nu) / E
    with open(f"_pl_{tag}.inp", "w") as f:
        for n, x, y in [(1, 0, 0), (2, L, 0), (3, L, L), (4, 0, L)]:
            f.write(("*NODE\n" if n == 1 else "") + f"{n},{x:.6e},{y:.6e}\n")
        f.write("*ELEMENT,TYPE=CPE4,ELSET=ALL\n1,1,2,3,4\n")
        f.write(f"*MATERIAL,NAME=NICR\n*ELASTIC\n{E:.6e},{nu:.3f}\n*PLASTIC\n{sigma_y:.6e},0.0\n")
        f.write("*SOLID SECTION,ELSET=ALL,MATERIAL=NICR\n")
        f.write("*STEP,NLGEOM\n*STATIC\n0.05,1.0\n*BOUNDARY\n1,1,2,0.\n4,1,1,0.\n2,2,2,0.\n")
        f.write(f"*BOUNDARY\n2,1,1,{eps*L:.6e}\n3,1,1,{eps*L:.6e}\n*EL PRINT,ELSET=ALL\nS,PEEQ\n*END STEP\n")
    subprocess.run(["ccx", f"_pl_{tag}"], capture_output=True, text=True)
    S, P, mode = [], [], None
    for ln in open(f"_pl_{tag}.dat"):
        s = ln.strip()
        if s.lower().startswith("stresses"): mode = "S"; continue
        if "equivalent plastic strain" in s.lower(): mode = "P"; continue
        p = s.split()
        if mode == "S" and len(p) >= 8 and p[0].isdigit() and p[1] == "1": S.append(list(map(float, p[2:8])))
        if mode == "P" and len(p) >= 3 and p[0].isdigit() and p[1] == "1":
            try: P.append(float(p[2]))
            except ValueError: pass
    sxx, syy, szz, sxy, sxz, syz = S[-1]
    vm = np.sqrt(0.5 * ((sxx - syy) ** 2 + (syy - szz) ** 2 + (szz - sxx) ** 2) + 3 * (sxy ** 2 + sxz ** 2 + syz ** 2))
    for ext in (".inp", ".dat", ".frd", ".cvg", ".sta", ".12d"):
        try: os.remove(f"_pl_{tag}" + ext)
        except OSError: pass
    print(f"  plastic ({temp_label}): imposed {elastic_equiv/1e9:.2f} GPa -> "
          f"von Mises {vm/1e6:.1f} MPa (yield {sigma_y/1e6:.0f}), PEEQ {P[-1]:.2e}")


def main():
    if subprocess.run(["which", "ccx"], capture_output=True).returncode != 0:
        raise SystemExit("CalculiX `ccx` not found on PATH. Install ccx (v2.21) to run the FEA validation.")
    print("FEA validation (CalculiX):")
    fig_11_relocation()
    fig_12_sweeps()
    plastic_check("LPBF 1300 C", 20e6, 2.2e9, "lpbf")
    plastic_check("anneal 400 C", 350e6, 0.73e9, "anneal")
    print("Done. Figures 11-12 generated; plastic relocation confirmed.")


if __name__ == "__main__":
    main()
