"""
Smoke tests. Cheap, exact, and run before any physics benchmark.

Each one isolates a single piece of machinery. If a physics benchmark fails,
run these first -- they localise the fault to a specific operator instead of
leaving you staring at a wrong drag number.
"""

import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import numpy as np

from waterfall_core import Grid, Solver, external_flow
from waterfall_core.poisson import PoissonDCT
from waterfall_core import kernels


def test_poisson_manufactured():
    """
    Poisson solver against a manufactured solution.

    phi = cos(pi x / Lx) cos(pi y / Ly) satisfies homogeneous Neumann on all
    faces exactly, so the discrete solve should recover it to truncation
    error, and that error should fall as h^2.
    """
    errors = []
    for n in (32, 64, 128):
        g = Grid(n, n, 1, 1.0, 1.0, 1.0)
        ps = PoissonDCT(g)
        x, y, _ = g.cell_centers()
        X, Y = np.meshgrid(x, y, indexing="ij")
        exact = (np.cos(np.pi * X) * np.cos(np.pi * Y))[:, :, None]

        # Use the *discrete* Laplacian of the exact field as RHS so we measure
        # the solver's error, not the difference between discrete and
        # continuous operators.
        lam_x = 2.0 * (np.cos(np.pi * 1 / n) - 1.0) / g.dx**2
        rhs = exact * (2.0 * lam_x)

        phi = ps.solve(rhs)[1:-1, 1:-1, 1:-1]
        phi -= phi.mean()
        ref = exact - exact.mean()
        errors.append(np.abs(phi - ref).max())

    ok = errors[-1] < 1e-10
    print(f"  poisson manufactured: max|err| = {errors} -> {'PASS' if ok else 'FAIL'}")
    return ok


def test_freestream_preservation():
    """
    Uniform flow must remain exactly uniform.

    This catches sign errors in the convective operator, bad ghost cells, and
    a projection that leaks. Any solver that fails this cannot be trusted on
    anything else, and the failure is usually invisible in a drag number
    because it looks like a plausible small perturbation.
    """
    g = Grid(48, 32, 1, 3.0, 2.0, 1.0)
    bc = external_flow(g, u_inf=(1.0, 0.0, 0.0))
    s = Solver(g, bc, nu=1e-3, rho=1.0)
    for _ in range(50):
        s.step()
    uc, vc, wc, _ = s.cell_fields()
    err_u = np.abs(uc - 1.0).max()
    err_v = np.abs(vc).max()
    div = s.divergence_norm()
    ok = err_u < 1e-9 and err_v < 1e-9 and div < 1e-10
    print(
        f"  freestream preservation: max|u-1|={err_u:.2e} max|v|={err_v:.2e} "
        f"div={div:.2e} -> {'PASS' if ok else 'FAIL'}"
    )
    return ok


def test_divergence_free():
    """After projection the velocity field must be discretely divergence
    free. Tests the Poisson/gradient pair for consistency: if the divergence
    and gradient operators are not exact adjoints, this leaves a residual."""
    g = Grid(40, 40, 1, 2.0, 2.0, 1.0)
    bc = external_flow(g, u_inf=(1.0, 0.0, 0.0))
    s = Solver(g, bc, nu=1e-2, rho=1.0)
    rng = np.random.default_rng(0)
    s.v[1:-1, 2:g.ny + 1, 1:-1] += 0.3 * rng.standard_normal(
        (g.nx, g.ny - 1, g.nz)
    )
    for _ in range(5):
        s.step()
    div = s.divergence_norm()
    ok = div < 1e-11
    print(f"  divergence free after projection: {div:.3e} -> {'PASS' if ok else 'FAIL'}")
    return ok


def test_ibm_force_sign():
    """
    A blocked cell must produce positive drag.

    Trivial, but it is exactly the check the old ray-based solver failed:
    it reported negative drag coefficients because the per-cell drag map and
    the reported force were computed by two different formulas.
    """
    g = Grid(64, 48, 1, 4.0, 3.0, 1.0)
    bc = external_flow(g, u_inf=(1.0, 0.0, 0.0))
    eta_u = g.zeros_u()
    eta_v = g.zeros_v()
    eta_c = g.zeros_p()
    # solid block in the middle
    eta_u[20:30, 20:30, 1] = 1.0
    eta_v[20:30, 20:30, 1] = 1.0
    eta_c[20:30, 20:30, 1] = 1.0
    s = Solver(g, bc, nu=1e-2, rho=1.0, eta_u=eta_u, eta_v=eta_v, eta_c=eta_c)
    for _ in range(200):
        s.step()
    fx = np.mean([f[0] for f in s.force_history[-50:]])
    ok = fx > 0
    print(f"  ibm drag sign: Fx = {fx:+.4f} N -> {'PASS' if ok else 'FAIL'}")
    return ok


def test_2d_collapse():
    """
    nz=1 must reproduce a genuinely 2D flow: w stays zero and z-derivatives
    contribute nothing. This is load-bearing -- there is no separate 2D code
    path, so if the collapse is wrong every 2D benchmark is wrong.
    """
    g = Grid(32, 32, 1, 2.0, 2.0, 1.0)
    bc = external_flow(g, u_inf=(1.0, 0.0, 0.0))
    s = Solver(g, bc, nu=1e-2, rho=1.0)
    rng = np.random.default_rng(1)
    s.v[1:-1, 2:g.ny + 1, 1:-1] += 0.2 * rng.standard_normal((g.nx, g.ny - 1, 1))
    for _ in range(30):
        s.step()
    wmax = np.abs(s.w).max()
    ok = wmax < 1e-14
    print(f"  2D collapse (max|w|): {wmax:.2e} -> {'PASS' if ok else 'FAIL'}")
    return ok


def main():
    print("Smoke tests")
    results = [
        test_poisson_manufactured(),
        test_freestream_preservation(),
        test_divergence_free(),
        test_2d_collapse(),
        test_ibm_force_sign(),
    ]
    print(f"\n{sum(results)}/{len(results)} passed")
    return 0 if all(results) else 1


if __name__ == "__main__":
    sys.exit(main())