"""
TorsoCAE Validation -- cfd/stokes case 1: planar Poiseuille flow.

Classic creeping-flow (Stokes) benchmark: pressure-driven flow between two
parallel no-slip plates develops the exact parabolic velocity profile

    u(y) = U_max * (1 - (2y/H)^2),   y in [-H/2, H/2]
    U_avg = (2/3) * U_max
    dp/dx = -8*mu*U_max / H^2   (pressure drop per unit length)

dolfinx_backend.py's _solve_cfd_stokes solves the LINEAR Stokes equations
(no advection term at all -- Re never enters the solved PDE), so the
"entrance length" for the profile to settle to the exact parabola from a
uniform (plug-flow) inlet is a purely diffusive decay length of order H,
not an inertial one -- a channel a few H past the inlet is already exact.
The channel here is 10H long and the profile is sampled at mid-channel
(5H downstream), well clear of both the inlet's plug-profile mismatch and
any outlet artifact.

The domain is a thin 3D slab (LZ << H) with a `symmetry` (slip, z=0)
condition on the front/back faces to force 2D-equivalent flow rather than
building an actual 2D mesh -- same technique used for the pendulum's plate
thickness in structural/geometric_nonlinear_dynamics.

Velocity uses Taylor-Hood P2 elements, which represent an exact parabola
exactly (a single quadratic element reproduces a quadratic field) -- so
mesh convergence here is expected to be extremely fast, unlike the P1
bending-locking issue found in thermal_stress/case_02.

Two independent checks:
  1. Velocity profile at mid-channel vs. the exact parabola (RMS/max, all
     sampled points).
  2. Streamwise pressure gradient (finite difference between two interior
     x-stations) vs. the exact -8*mu*U_max/H^2.

Run: HWLOC_COMPONENTS=-gl python3 validation/cfd/stokes/case_01_poiseuille_flow/run_validation.py
"""
import json
import sys
import time
from pathlib import Path

import numpy as np

CASE_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(CASE_DIR.parents[3] / "server"))
from pytorsocae import TorsoCAESession

H, LX, LZ = 0.02, 0.2, 0.005      # m -- channel gap, length (10H), thin slab depth
MU, RHO = 1.0e-3, 1000.0          # Pa.s, kg/m3 (rho unused by the linear Stokes driver)
U_MAX = 1.0                       # m/s, exact centerline velocity (Stokes is linear -- magnitude is arbitrary)
U_AVG = (2.0 / 3.0) * U_MAX
DPDX_EXACT = -8.0 * MU * U_MAX / H ** 2   # Pa/m

# Plain (non-boolean) csg box: 1=x=0, 2=x=dx, 3=y=cy, 4=y=cy+dy, 5=z=cz, 6=z=cz+dz
TAG_INLET, TAG_OUTLET, TAG_ZLO, TAG_ZHI = 1, 2, 5, 6


def u_exact(y: np.ndarray) -> np.ndarray:
    return U_MAX * (1.0 - (2.0 * y / H) ** 2)


def run_case(size_factor: float) -> dict:
    session = TorsoCAESession()
    session.csg_make_shape("box", {"dx": LX, "dy": H, "dz": LZ, "cx": 0.0, "cy": -0.5 * H, "cz": -0.5 * LZ},
                            name="Channel", body_id="body_0")
    session.csg_select("body_0")
    session.mesh(mesh_id="mesh_0", algo_id="hex_transfinite", size_factor=size_factor,
                 size_min=0, size_max=0, order=2, dim=3)
    session.solid("Channel").material(mu=MU, rho=RHO)
    session.surface(TAG_INLET, scope_id="body_0").bc("inlet_velocity", values=[U_AVG, 0, 0])
    session.surface(TAG_OUTLET, scope_id="body_0").bc("outlet_pressure", values=[0.0])
    session.surface(TAG_ZLO, scope_id="body_0").bc("symmetry", values=[0.0], components=["z"])
    session.surface(TAG_ZHI, scope_id="body_0").bc("symmetry", values=[0.0], components=["z"])
    # y=-H/2, y=+H/2 (tags 3,4) left unassigned -> CFD driver default-applies no-slip 'wall'.
    session.set_physics("cfd", submodel="stokes", backend="dolfinx")
    session.set_solver_options(algo="direct", precond="none", tol=1e-10, max_iter=500, n_cores=2, device="cpu")
    t0 = time.time()
    result = session.compute(mesh_ids=["mesh_0"])
    elapsed = time.time() - t0

    npz = np.load(result["npz_path"])
    u_coords = npz["u_dof_coords"]
    velocity = npz["velocity"]
    p_coords = npz["p_dof_coords"]
    pressure = npz["pressure"]

    # Velocity profile at three x-stations (all z -- flow is z-invariant by symmetry BC),
    # to distinguish genuine mesh-discretization error from residual inlet/outlet
    # entrance-decay effects (error should be flat across stations if discretization,
    # decreasing away from both ends if entrance-decay).
    denom = U_MAX
    unique_x = np.unique(u_coords[:, 0])
    station_errors = {}
    for frac in (0.3, 0.5, 0.7):
        x_station = unique_x[np.argmin(np.abs(unique_x - frac * LX))]
        mask = np.abs(u_coords[:, 0] - x_station) < 1e-9
        y_s = u_coords[mask, 1]
        ux_s = velocity[mask, 0]
        ref_s = u_exact(y_s)
        station_errors[frac] = float(np.abs(ux_s - ref_s).max() / denom * 100.0)

    x_mid = unique_x[np.argmin(np.abs(unique_x - 0.5 * LX))]
    mid_mask = np.abs(u_coords[:, 0] - x_mid) < 1e-9
    y_mid = u_coords[mid_mask, 1]
    ux_fem = velocity[mid_mask, 0]
    ux_ref = u_exact(y_mid)
    profile_rms_pct = float(np.sqrt(np.mean((ux_fem - ux_ref) ** 2)) / denom * 100.0)
    profile_max_pct = station_errors[0.5]

    # Streamwise pressure gradient: linear fit over interior P1 pressure nodes (avoids
    # inlet/outlet end effects and doesn't depend on any x-station landing exactly on a
    # mesh vertex plane -- pressure is exactly linear in x for this exact solution).
    interior = (p_coords[:, 0] > 0.2 * LX) & (p_coords[:, 0] < 0.8 * LX)
    dpdx_fem = float(np.polyfit(p_coords[interior, 0], pressure[interior], 1)[0])
    dpdx_error_pct = abs(dpdx_fem - DPDX_EXACT) / abs(DPDX_EXACT) * 100.0

    return {
        "size_factor": size_factor,
        "n_velocity_dofs": int(u_coords.shape[0]),
        "station_max_error_pct": station_errors,
        "profile_rms_error_pct": profile_rms_pct,
        "profile_max_error_pct": profile_max_pct,
        "dpdx_exact_pa_per_m": DPDX_EXACT,
        "dpdx_fem_pa_per_m": dpdx_fem,
        "dpdx_error_pct": dpdx_error_pct,
        "elapsed_s": elapsed,
    }


def main():
    runs = [run_case(sf) for sf in (1.0, 0.5)]
    finest = runs[-1]

    summary = {
        "case": "cfd/stokes / planar Poiseuille flow between parallel plates",
        "citation": "Standard exact solution of the 2D Stokes equations for pressure-driven channel flow "
                     "(e.g. White, Viscous Fluid Flow); u(y)=U_max*(1-(2y/H)^2), dp/dx=-8*mu*U_max/H^2",
        "geometry": {"channel_gap_H_m": H, "length_LX_m": LX, "slab_depth_LZ_m": LZ},
        "material": {"mu_Pa_s": MU, "rho_kg_m3": RHO},
        "U_max_m_s": U_MAX,
        "U_avg_m_s": U_AVG,
        "runs": runs,
        "profile_rms_error_pct": finest["profile_rms_error_pct"],
        "profile_max_error_pct": finest["profile_max_error_pct"],
        "dpdx_error_pct": finest["dpdx_error_pct"],
        "pass": bool(finest["profile_max_error_pct"] < 3.0 and finest["dpdx_error_pct"] < 3.0),
    }
    (CASE_DIR / "results.json").write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2), flush=True)


if __name__ == "__main__":
    main()