"""Public API demo: fixed 2:1 patch with a stationary, optionally actuated circle.

From this checkout, after checking the selected physical GPU is idle:
    PYTHONPATH="$PWD/src" CUDA_VISIBLE_DEVICES=0 python examples/fixed_patch_circle.py
    ... --speed -0.002 --fine-block 16 20 --mode slot_reservoir

A short smoke run (default 128 coarse steps): it shows setup, stepping and
fluid-only diagnostics, not a developed flow, accuracy or long-time stability.
Normal blowing/suction injects/extracts mass, so leaf mass is expected to change.
No files, plots or checkpoints are written.
"""
import argparse
from pathlib import Path

import numpy as np

import CelerisLab
from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.objects import SimObject, SurfaceActuation
from CelerisLab.config import LBMConfig
from CelerisLab.cuda.context import CudaContext
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY
from CelerisLab.lbm.equilibrium import compute_feq_d2q9
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper

NX, NY, BOUNDS = 32, 24, (8, 6, 24, 18)  # coarse cells; patch is half-open (x0, y0, x1, y1)
CENTER, RADIUS = (16.0, 12.0), 2.15       # coarse CELL-EDGE frame: cell I spans [I, I+1)


def config(mode):
    if mode == "periodic":  # x and y periodic coarse level
        return LBMConfig(nx=NX, ny=NY, viscosity=.1)
    # Bottom on-site velocity slot, top fixed-density reservoir, x periodic.
    return LBMConfig(nx=NX, ny=NY, viscosity=.1, velocity=0, wall_normal_velocity=.001,
                     wall_slot_width=8, y_wall_bc="on_site_velocity", top_wall_bc="reservoir")


def leaf_diagnostics(coarse, fine):
    """Fluid leaves only (downloads are NaN on covered/ghost/solid cells).

    Mass counts a fine cell as 1/4 of a coarse cell (2D, 2:1)."""
    out = {}
    for name, g, area in (("coarse", coarse, 1.0), ("fine", fine, .25)):
        leaf = np.isfinite(g).all(axis=0)
        f = g[:, leaf].astype(np.float64)
        rho = f.sum(axis=0)
        ux = (np.array(D2Q9_EX)[:, None] * f).sum(axis=0) / rho
        uy = (np.array(D2Q9_EY)[:, None] * f).sum(axis=0) / rho
        out[name] = (area * rho.sum(), rho.min(), rho.max(), np.hypot(ux, uy).max())
    return out


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--mode", choices=("periodic", "slot_reservoir"), default="periodic")
    parser.add_argument("--speed", type=float, default=.002,
                        help="normal surface velocity, coarse lattice units (+ blows, - sucks, 0 passive)")
    parser.add_argument("--fine-block", type=int, nargs=2, metavar=("BY", "BX"),
                        help="blocked fine storage; must divide the whole fine allocation (32 x 40 here)")
    parser.add_argument("--steps", type=int, default=128)
    args = parser.parse_args()
    root = Path(__file__).resolve().parents[1]
    if Path(CelerisLab.__file__).resolve() != root / "src/CelerisLab/__init__.py":
        raise RuntimeError(f"Set PYTHONPATH to {root / 'src'}")

    # Aperture centred at angle 0.7 rad, half-width 0.85 rad (the tested fixture).
    spec = SurfaceActuation(0.7, 0.85, args.speed) if args.speed else None
    body = SimObject(0, geometry=CircleGeometry(*CENTER, RADIUS), center=CENTER, radius=RADIUS,
                     surface_actuation=spec)
    cfg = config(args.mode)
    fine_block = tuple(args.fine_block) if args.fine_block else None
    # The caller owns the CUDA context; the stepper must be closed first (the
    # nested with-statement does that).
    with CudaContext(0) as context, FixedPatchStepper(
            cfg, BOUNDS, context=context, coarse_boundary=args.mode,
            fine_objects=[body], fine_block=fine_block) as run:
        p = run.patch
        # Rest equilibrium is already postcollision G, so no initial collision is
        # needed. Non-leaf entries (covered, ghost, solid) are discarded by initialize.
        rest = np.asarray(compute_feq_d2q9(1.0, 0.0, 0.0), np.float32)[:, None, None]
        run.initialize(np.broadcast_to(rest, (9, *p.coarse_shape)),
                       np.broadcast_to(rest, (9, *p.fine_shape)))
        storage = f"blocked {fine_block}" if fine_block else "dense"
        print(f"mode={args.mode} fine storage={storage} speed={args.speed:+g} "
              f"(fine allocation {p.fine_shape[0]}x{p.fine_shape[1]}, fine origin {p.fine_origin})")
        start = leaf_diagnostics(*run.download())
        for _ in range(0, args.steps, 32):
            run.step(min(32, args.steps - run.step_count))
            coarse, fine = run.download()  # synchronized postcollision G, NaN on non-leaves
            d = leaf_diagnostics(coarse, fine)
            mass = d["coarse"][0] + d["fine"][0]
            rho_lo = min(d["coarse"][1], d["fine"][1])
            rho_hi = max(d["coarse"][2], d["fine"][2])
            speed = max(d["coarse"][3], d["fine"][3])
            print(f"coarse step {run.step_count:4d}: leaf mass {mass:.9f} "
                  f"(change {mass - start['coarse'][0] - start['fine'][0]:+.3e}) "
                  f"rho [{rho_lo:.6f}, {rho_hi:.6f}] max |u| {speed:.3e}")
        print("smoke run only: storage/stepping demonstration, not accuracy or stability evidence")


if __name__ == "__main__":
    main()