"""Three levels, separated regions and an explicit fluid-region remap.

Run on an approved idle GPU (logical device 0):
    PYTHONPATH=src CUDA_VISIBLE_DEVICES=0 python examples/multigrid_remap.py --verify-jet
Host-only geometry check, without CUDA:
    PYTHONPATH=src python examples/multigrid_remap.py --check-geometry

A short public-API smoke example, not an accuracy/stability benchmark. No files
are written. Edit the rectangles below; this is manual remapping, not auto AMR.
"""
import argparse
from contextlib import ExitStack

import numpy as np

from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.jets import Jet, Parallel, Tube
from CelerisLab.body.objects import SimObject
from CelerisLab.config import LBMConfig
from CelerisLab.lbm.equilibrium import compute_feq_d2q9
from CelerisLab.lbm.refinement import FixedPatchSet2D
from CelerisLab.lbm.synchronized_remap import RemapTopology

CFG = LBMConfig(nx=48, ny=32, viscosity=.1)  # D2Q9/SRT/FP32, periodic
PARENTS = ((4, 4, 20, 28), (24, 4, 44, 28))  # L0 cell-edge rectangles
PIN = (16, 16, 32, 24)  # L1 allocated-grid indices; whole body region stays fixed
BEFORE = (PIN, (56, 32, 72, 40))
AFTER = (PIN, (60, 32, 76, 40))  # move only the separated fluid rectangle
LIVE_AMPLITUDE = -.004


def geometry():
    # Retain these original L0 edge-frame definitions for every remap; never
    # substitute the stepper's projected fine-local boundary objects.
    objects = (SimObject(0, geometry=CircleGeometry(12., 10., .8),
                         center=(12., 10.), radius=.8),)
    jets = (Jet(Tube((10.5, 10.), 0., 3., .3), Parallel(0.), amplitude=.002),)
    parent = FixedPatchSet2D.from_config(CFG, PARENTS)
    meshes = tuple(RemapTopology.from_config(
        CFG, parent, rectangles, child_objects=objects, child_jets=jets)
        for rectangles in (BEFORE, AFTER))
    if any(mesh.pinned_rectangles != (PIN,) for mesh in meshes):
        raise RuntimeError("the body-containing rectangle must remain pinned")
    return objects, jets, meshes


def diagnostics(run, mesh):
    """Masks come from geometry, never from whether downloaded values are finite."""
    mass = 0.
    for level, (g, leaf) in enumerate(zip(run.download(), mesh.leaf_masks)):
        fluid = g[:, leaf].astype(np.float64)
        rho = fluid.sum(axis=0)
        if not np.isfinite(fluid).all() or np.any(rho <= 0):
            raise RuntimeError(f"invalid L{level} fluid populations/density")
        if not np.isnan(g[:, ~leaf]).all():
            raise RuntimeError(f"unexpected L{level} nonleaf values")
        mass += rho.sum() / 4**level  # area weights: 1, 1/4, 1/16
    print(f"absolute base step {run.step_count}: finite fluid leaves, mass={mass:.9f}")
    # Driven boundary injects/extracts mass; mass need not stay constant in time.


def verify_live_jet(current, objects):
    """Optional public-API check; normal continuation needs no control solver."""
    # Identity remap copies the same state. Explicitly setting the known live
    # amplitude on the control detects a stale amplitude on the continuation.
    # No private fields or constructor-time Jet snapshots inspect the solver.
    with current.remap_child(AFTER, child_objects=objects) as control:
        control.set_jet_amplitude(0, LIVE_AMPLITUDE, level=2)
        current.step()
        control.step()
        for actual, expected in zip(current.download(), control.download()):
            if not np.array_equal(actual, expected, equal_nan=True):
                raise RuntimeError("live Jet continuation differs from explicitly set control")
    print("live Jet amplitude preserved: next step equals explicitly set control")


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--check-geometry", action="store_true")
    parser.add_argument("--verify-jet", action="store_true",
                        help="compare one step with an explicitly driven control")
    args = parser.parse_args()
    objects, jets, (old_mesh, new_mesh) = geometry()
    if args.check_geometry:
        print("geometry OK: three levels, two L1 regions, fixed body pin, moved fluid region")
        return

    # Lazy imports keep --check-geometry useful on a CPU-only machine.
    from CelerisLab.cuda.context import CudaContext
    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper

    # ExitStack owns every old/new stepper, including on failure. All are closed
    # before the shared context. remap_child never consumes/closes its source.
    with CudaContext(0) as context, ExitStack() as owned:
        old = owned.enter_context(FixedPatchStepper(
            CFG, PARENTS, context=context, base_block=(4, 4), fine_block=(4, 4),
            child_block=(4, 4), child_rectangles_in_parent=BEFORE,
            child_objects=objects, child_jets=jets))
        rest = np.asarray(compute_feq_d2q9(1., 0., 0.), np.float32)[:, None, None]
        old.initialize(*(np.broadcast_to(rest, (9, *mask.shape))
                         for mask in old_mesh.leaf_masks))
        old.step(2)  # one base step advances L1 twice and L2 four times
        old.set_jet_amplitude(0, LIVE_AMPLITUDE, level=2)
        diagnostics(old, old_mesh)
        count = old.step_count
        current = owned.enter_context(old.remap_child(AFTER, child_objects=objects))
        if current is old or current.step_count != count:
            raise RuntimeError("remap must return a distinct continuation at the same time")
        diagnostics(current, new_mesh)

        if args.verify_jet:
            verify_live_jet(current, objects)  # optional control; advances one step
        else:
            current.step()
        if old.step_count != count or current.step_count != count + 1:
            raise RuntimeError("old/new absolute step ownership changed")
        old.close()  # retire the old state only after successful replacement
        current.step(2)
        if current.step_count != count + 3:
            raise RuntimeError("continuation lost its absolute step count")
        diagnostics(current, new_mesh)
        print("smoke example only: no accuracy or long-time stability claim")


if __name__ == "__main__":
    main()