"""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 --mode slot_reservoir
Both levels use blocked storage (--base-block, --fine-block tiles).

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, SOLID
from CelerisLab.lbm.equilibrium import compute_feq_d2q9
from CelerisLab.lbm.refinement import ACTIVE
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_masks(run, mode):
    """Known fluid leaves: uncovered coarse cells (minus wall rows in slot mode)
    and fine ACTIVE cells outside the body."""
    p = run.hierarchy.pairs[0]
    x0, y0, x1, y1 = p.bounds
    coarse = np.ones(p.coarse_shape, bool)
    coarse[y0:y1, x0:x1] = False
    if mode == "slot_reservoir":
        coarse[[0, -1]] = False
    solid = (run.boundaries[1].flags.reshape(p.fine_shape) & SOLID) != 0
    return coarse, (p.fine_kind == ACTIVE) & ~solid


def leaf_diagnostics(coarse, fine, masks):
    """Fluid leaves only; a non-finite population or rho <= 0 there is an error.

    Mass counts a fine cell as 1/4 of a coarse cell (2D, 2:1)."""
    out = {}
    for name, g, leaf, area in (("coarse", coarse, masks[0], 1.0), ("fine", fine, masks[1], .25)):
        f = g[:, leaf].astype(np.float64)
        rho = f.sum(axis=0)
        if not np.isfinite(f).all() or (rho <= 0).any():
            raise RuntimeError(f"non-finite population or non-positive density on a {name} fluid leaf")
        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("--base-block", type=int, nargs=2, metavar=("BY", "BX"), default=(12, 8),
                        help="root tile; must divide the 24 x 32 root")
    parser.add_argument("--fine-block", type=int, nargs=2, metavar=("BY", "BX"), default=(16, 20),
                        help="fine tile; must divide the whole fine allocation (32 x 40 here)")
    parser.add_argument("--steps", type=int, default=128)
    args = parser.parse_args()
    if args.steps < 0:
        parser.error("--steps must be nonnegative")
    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)
    # 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, context=context, coarse_boundary=args.mode, refinements=[BOUNDS],
            level_blocks=[tuple(args.base_block), tuple(args.fine_block)],
            level_objects=[(), (body,)]) as run:
        p = run.hierarchy.pairs[0]
        # 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 root {tuple(args.base_block)} fine {tuple(args.fine_block)}"
        print(f"mode={args.mode} storage={storage} speed={args.speed:+g} "
              f"(fine allocation {p.fine_shape[0]}x{p.fine_shape[1]}, fine origin {p.fine_origin})")
        masks = leaf_masks(run, args.mode)
        start = leaf_diagnostics(*run.download(), masks)
        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, masks)
            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()