"""Optional exact-header GPU helper probe. Not run in this audit.
Run ONLY through the authorized Horsea executor. Uses existing PyCUDA/NumPy.
This is helper parity, not full collide_dispatch or turbulence validation.
"""

import argparse, json, pathlib
from audit import DEFAULT_ROOT, populations, wrapper


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=pathlib.Path, default=DEFAULT_ROOT)
    ap.add_argument("--device", type=int, default=0)
    ap.add_argument("--run-gpu", action="store_true")
    ap.add_argument("--output", type=pathlib.Path, required=True)
    args = ap.parse_args()
    if not args.run_gpu:
        ap.error("Explicit --run-gpu required; use authorized Horsea executor only")
    import numpy as np
    import pycuda.driver as cuda
    from pycuda.compiler import SourceModule

    cuda.init()
    ctx = cuda.Device(args.device).make_context()
    rows = []
    try:
        for dim in (2, 3):
            for cs in (0.16, 0.0):
                for closed, iters in ((1, 3), (0, 3), (0, 200)):
                    module = SourceModule(
                        wrapper(dim, closed, cs, iters, gpu=True),
                        include_dirs=[
                            str(args.root.resolve() / "src/CelerisLab/lbm/kernels")
                        ],
                        options=["-std=c++11", "--use_fast_math"],
                        no_extern_c=True,
                    )
                    fn = module.get_function("probe")
                    for case, kind, p, tau in [
                        ("equilibrium", "equilibrium", 0.0, 0.8),
                        ("pure_shear", "shear", 0.02, 0.8),
                        ("isotropic", "isotropic", 0.02, 0.8),
                        ("cap_active", "shear", 0.2, 0.503),
                    ]:
                        f, eq = populations(dim, kind, p)
                        out = np.zeros(1, np.float32)
                        fn(
                            cuda.In(np.asarray(f, np.float32)),
                            cuda.In(np.asarray(eq, np.float32)),
                            np.float32(1),
                            np.float32(1 / tau),
                            cuda.Out(out),
                            block=(1, 1, 1),
                            grid=(1, 1),
                        )
                        ctx.synchronize()
                        rows.append(
                            dict(
                                dim=dim,
                                cs=cs,
                                closed=closed,
                                iters=iters,
                                case=case,
                                omega=float(out[0]),
                            )
                        )
        args.output.write_text(
            json.dumps(
                dict(
                    device=cuda.Device(args.device).name(),
                    diagnostic_only=True,
                    results=rows,
                ),
                indent=2,
            )
            + "\n"
        )
    finally:
        ctx.pop()
        ctx.detach()


if __name__ == "__main__":
    main()