"""
Scientific ML Studio -- reference code (written by hand, not generated by the Studio)
 
Licence : Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0)
          https://creativecommons.org/licenses/by-nc/4.0/
(c) 2026 Scientific ML Studio. All rights are reserved to the Scientific ML Studio platform
except the permissions granted by the licence above.
 
You may share and adapt this code for NON-COMMERCIAL purposes, provided you credit
"Scientific ML Studio" with a link to https://scimlstudio.com/, and indicate any changes you make.
Commercial use, paid client work and production deployment need prior written permission,
which you can ask for through the site's Contact page.
-------------------------------------------------------------------------------

check the lid-driven cavity numbers without PyTorch and without training.

Reads the weights of a trained network (weights.npz, written by lid_driven_cavity_final.py, or the
file cavity_weights_run1.npz that comes with the lesson), evaluates the network on a 201 x 201 grid,
and compares it with the finite element reference (cavity_reference_N128.npz, or solved by
cavity_reference.py if the file is missing).

    python cavity_compare.py                                   the lesson's run
    python cavity_compare.py --weights weights.npz             your own run

The network is the one in the lesson: inputs (x, y) scaled to [-1, 1], four hidden layers of 50 tanh
neurons, outputs (u, v, p). The vorticity v_x - u_y is computed exactly by carrying the derivatives
through the layers. Needs numpy (and scipy only if the reference has to be solved).

Licence: CC BY-NC 4.0, (c) 2026 Scientific ML Studio. Credit "Scientific ML Studio", https://scimlstudio.com/
"""
import argparse
from pathlib import Path

import numpy as np

_trapz = getattr(np, "trapezoid", None) or np.trapz


def load_network(path):
    w = np.load(path)
    layers = [(w[f"net.{k}.weight"].astype(np.float64), w[f"net.{k}.bias"].astype(np.float64))
              for k in sorted({int(n.split(".")[1]) for n in w.files})]
    return layers


def evaluate(layers, x, y):
    """u, v, p and vorticity at the points (x, y). Inputs are scaled from [0, 1] to [-1, 1]."""
    z = 2.0 * np.stack([x, y], axis=1) - 1.0
    dx = np.zeros_like(z)
    dy = np.zeros_like(z)
    dx[:, 0] = 2.0                                    # d z / d x
    dy[:, 1] = 2.0                                    # d z / d y
    for i, (W, b) in enumerate(layers):
        a = z @ W.T + b
        dax, day = dx @ W.T, dy @ W.T
        if i < len(layers) - 1:
            z = np.tanh(a)
            s = 1.0 - z ** 2
            dx, dy = dax * s, day * s
        else:
            z, dx, dy = a, dax, day
    u, v, p = z[:, 0], z[:, 1], z[:, 2]
    return u, v, p, dx[:, 1] - dy[:, 0]               # vorticity = v_x - u_y


def domain_mean(f, g):
    return float(_trapz(_trapz(f, g, axis=1), g, axis=0) / (g[-1] - g[0]) ** 2)


def stream_function(u, g):
    inc = 0.5 * (u[:, 1:] + u[:, :-1]) * np.diff(g)[None, :]
    return np.concatenate([np.zeros((u.shape[0], 1)), np.cumsum(inc, axis=1)], axis=1)


def vortex_centre(psi, g):
    """Minimum of the stream function, refined with a parabola through three grid values."""
    i, j = np.unravel_index(int(np.argmin(psi)), psi.shape)

    def refine(fm, f0, fp, h, c):
        d = fm - 2 * f0 + fp
        return c + 0.5 * h * (fm - fp) / d if d > 0 else c
    h = g[1] - g[0]
    return (refine(psi[i - 1, j], psi[i, j], psi[i + 1, j], h, g[i]),
            refine(psi[i, j - 1], psi[i, j], psi[i, j + 1], h, g[j]))


def main():
    here = Path(__file__).resolve().parent
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[1])
    ap.add_argument("--weights", default=str(here / "cavity_weights_run1.npz"))
    ap.add_argument("--reference", default=str(here / "cavity_reference_N128.npz"))
    a = ap.parse_args()

    g = np.linspace(0.0, 1.0, 201)
    X, Y = np.meshgrid(g, g, indexing="ij")
    u, v, p, om = (f.reshape(201, 201) for f in evaluate(load_network(a.weights), X.ravel(), Y.ravel()))
    p = p - domain_mean(p, g)                         # the pressure is only defined up to a constant

    if Path(a.reference).exists():
        ref = dict(np.load(a.reference))
    else:
        import cavity_reference as cr
        print("reference file not found: solving on 128 x 128 elements (about two minutes)")
        ref = cr.reference_fields(cr.Cavity(n_el=128, re=100.0).solve())
    ref["p"] = ref["p"] - domain_mean(ref["p"], g)

    rel = lambda f, r: np.linalg.norm(f - r) / np.linalg.norm(r)
    print("Relative L2 error against the finite element reference, 201 x 201 grid")
    for name, f, r in (("u", u, ref["u"]), ("v", v, ref["v"]), ("p (mean removed)", p, ref["p"]),
                       ("speed", np.hypot(u, v), np.hypot(ref["u"], ref["v"])), ("vorticity", om, ref["omega"])):
        print(f"  {name:<18s}{rel(f, r):.3e}   largest abs error {np.abs(f - r).max():.3e}")
    vec = np.sqrt(np.sum((u - ref["u"]) ** 2 + (v - ref["v"]) ** 2) / np.sum(ref["u"] ** 2 + ref["v"] ** 2))
    print(f"  {'velocity (u, v)':<18s}{vec:.3e}")
    cx, cy = vortex_centre(stream_function(u, g), g)
    rx, ry = vortex_centre(stream_function(ref["u"], g), g)
    print(f"Primary vortex centre: network ({cx:.4f}, {cy:.4f}), reference ({rx:.4f}, {ry:.4f})")


if __name__ == "__main__":
    main()