"""One measured L1 strain observation, one growth-only L2 remap, no scheduler.

Run from the repository root: PYTHONPATH=src:. python -m research.shear_refinement.example
Requires CUDA; CPU helpers and tests never import CUDA.
"""
from contextlib import ExitStack
from numbers import Integral
import numpy as np

from CelerisLab.config import LBMConfig
from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.objects import SimObject
from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY
from CelerisLab.lbm.refinement import FixedPatchSet2D
from CelerisLab.lbm.synchronized_remap import RemapTopology
from .boundary_shear_tagging import tag_cells

CFG = LBMConfig(nx=48, ny=32, viscosity=.1)
PARENTS = ((4, 4, 20, 28), (24, 4, 44, 28))
CHILDREN = ((16, 16, 32, 40), (56, 16, 72, 40))
CORRIDOR = (72, 20, 76, 36)  # L1 allocated indices; bounded to right of fluid L2


def grow_rectangle(current, request, *, alignment=4, area_cap=480):
    """Bounding-box growth only, in request-array cell-edge coordinates.

    No body or solver knowledge. Reject invalid/budget-exceeding proposals with
    ValueError; never clip a request, shrink old coverage or mutate an input.
    Alignment is edge divisibility, not a substitute for full topology guards.
    """
    request = np.asarray(request)
    if request.dtype != bool or request.ndim != 2:
        raise ValueError('request must be a 2D boolean mask')
    values = (*current, alignment, area_cap)
    if len(current) != 4 or any(isinstance(v, (bool, np.bool_)) or not isinstance(v, Integral) for v in values):
        raise ValueError('rectangle, alignment and area cap must be integers')
    x0, y0, x1, y1 = current
    ny, nx = request.shape
    if alignment <= 0 or area_cap <= 0 or not (0 <= x0 < x1 <= nx and 0 <= y0 < y1 <= ny):
        raise ValueError('invalid rectangle, alignment or area cap')
    if any(v % alignment for v in current):
        raise ValueError('current rectangle must be aligned')
    y, x = np.nonzero(request)
    if x.size:
        x0, y0 = min(x0, int(x.min())), min(y0, int(y.min()))
        x1, y1 = max(x1, int(x.max())+1), max(y1, int(y.max())+1)
    a = alignment
    result = (x0//a*a, y0//a*a, (x1+a-1)//a*a, (y1+a-1)//a*a)
    if result[2] > nx or result[3] > ny:
        raise ValueError('aligned growth exceeds observation frame')
    if (result[2]-result[0])*(result[3]-result[1]) > area_cap:
        raise ValueError('growth exceeds area cap')
    return result


def objects():
    return (SimObject(71, geometry=CircleGeometry(12., 14., .8),
                      center=(12., 14.), radius=.8),)


def topology(rectangles=CHILDREN):
    return RemapTopology.from_config(CFG, FixedPatchSet2D.from_config(CFG, PARENTS),
                                     rectangles, child_objects=objects())


def observe(g, mesh, *, shear_on=.001):
    """Moments of completed-time postcollision G; only real L1 fluid leaves.

    No reconstructed covered values, ghosts or holes are derivative donors.
    solid is independent geometry (empty on L1: the body is pinned at L2).
    The array is cropped/sparse, so neither derivative axis wraps.
    """
    valid = mesh.leaf_masks[1]
    f = g[:, valid].astype(np.float64)
    rho = f.sum(axis=0)
    if not np.isfinite(f).all() or np.any(rho <= 0):
        raise ValueError('invalid completed-time L1 fluid populations')
    ux, uy = (np.full(valid.shape, np.nan) for _ in range(2))
    ux[valid] = (np.asarray(D2Q9_EX)[:, None]*f).sum(axis=0)/rho
    uy[valid] = (np.asarray(D2Q9_EY)[:, None]*f).sum(axis=0)/rho
    tags = tag_cells(ux, uy, np.zeros(valid.shape, bool), valid=valid, owned=valid,
                     dx_level=.5, origin=mesh.parent.fine_origin, periodic=(False, False),
                     U_ref=.05, L_ref=1., shear_on=shear_on, shear_off=0.,
                     band_on=0., band_off=0., n_buffer=0)
    corridor = np.zeros(valid.shape, bool)
    x0, y0, x1, y1 = CORRIDOR
    corridor[y0:y1, x0:x1] = True
    request = tags.tagged & corridor
    proposed = (CHILDREN[0], grow_rectangle(CHILDREN[1], request))
    # Host guards cover nesting, donor footprints, complete body and separation.
    candidate = topology(proposed)
    candidate.child.fine_layout((4, 4))
    report = dict(requested=int(request.sum()), unknown=int(tags.unevaluable.sum()),
                  corridor_unknown=int((tags.unevaluable & corridor).sum()),
                  corridor_unavailable=int((corridor & ~valid).sum()),
                  out_of_scope=int((tags.tagged & ~corridor).sum()), rectangles=proposed)
    return proposed, report


def initial_fields(mesh):
    """Smooth periodic shear wave, sampled in each level's physical frame."""
    from CelerisLab.lbm.equilibrium import compute_feq_d2q9
    frames = ((np.zeros(2), 1.), (np.asarray(mesh.parent.fine_origin), .5),
              (np.asarray(mesh.parent.fine_origin)+.5*np.asarray(mesh.child.fine_origin), .25))
    arrays = []
    for leaf, (origin, dx) in zip(mesh.leaf_masks, frames):
        y = origin[1]+dx*(np.arange(leaf.shape[0])+.5)
        ux = .01*np.sin(2*np.pi*y/CFG.ny)
        rows = np.stack([compute_feq_d2q9(1., u, 0.) for u in ux], axis=1)
        arrays.append(np.broadcast_to(rows[:, :, None], (9, *leaf.shape)).copy())
    return tuple(arrays)


def initialize(run):
    run.initialize(*initial_fields(topology()))


def make(context):
    from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
    return FixedPatchStepper(CFG, PARENTS, context=context, base_block=(4, 4),
                             fine_block=(4, 4), child_block=(4, 4),
                             child_rectangles_in_parent=CHILDREN, child_objects=objects())


def main():
    from CelerisLab.cuda.context import CudaContext
    with CudaContext(0) as ctx, ExitStack() as owned:
        run = owned.enter_context(make(ctx))
        initialize(run)
        run.step(2)  # Returns only at completed synchronized base time.
        proposed, report = observe(run.download()[1], topology())
        print(report)
        if proposed != CHILDREN:
            run = owned.enter_context(run.remap_child(proposed, child_objects=objects()))
        run.step(2)
        print(f'completed base step {run.step_count}; body rectangle {run.child_patch.bounds[0]}')


if __name__ == '__main__':
    main()