"""Research-only field capture helpers for the Dütsch driver (numpy only).

Snapshots are taken at the first step of the scored phase bins in declared
cycles: a near-body crop at full resolution and a fluid-weighted block-mean
downsample of the whole domain; plus scalar trends at every bin read.
"""
from __future__ import annotations

import numpy as np

CROP_X, CROP_Y, DOWN = 8.0, 6.0, 8      # crop half-widths in D; downsample factor
STRIP = 2.0                             # inlet/outlet strip width in D


def crop_box(xm, ym, D, n, body_extent):
    """Inclusive integer crop (x0, x1, y0, y1); fail closed if it leaves the domain
    interior or does not contain the body's swept extent (x_lo, x_hi, y_lo, y_hi)."""
    x0, x1 = int(round(xm - CROP_X * D)), int(round(xm + CROP_X * D))
    y0, y1 = int(round(ym - CROP_Y * D)), int(round(ym + CROP_Y * D))
    if x0 < 1 or y0 < 1 or x1 > n - 2 or y1 > n - 2:
        raise ValueError("capture crop leaves the domain interior")
    bx0, bx1, by0, by1 = body_extent
    if not (x0 <= bx0 and bx1 <= x1 and y0 <= by0 and by1 <= y1):
        raise ValueError("capture crop clips the body's swept extent")
    return x0, x1, y0, y1


def downsample(field, fluid, f=DOWN):
    """Fluid-weighted f x f block means; NaN where a block has no fluid cell."""
    ny, nx = field.shape
    h, w = ny // f, nx // f
    v = np.where(fluid, field, 0.0)[:h * f, :w * f].reshape(h, f, w, f).sum(axis=(1, 3))
    c = fluid[:h * f, :w * f].reshape(h, f, w, f).sum(axis=(1, 3))
    with np.errstate(invalid="ignore", divide="ignore"):
        return np.where(c > 0, v / np.maximum(c, 1), np.nan).astype(np.float32)


def fluid_mean(field, fluid, xs=slice(None), ys=slice(None)):
    m = fluid[ys, xs]
    return float(field[ys, xs][m].mean()) if m.any() else float("nan")


def trend_row(s, k, ph, rho, ux, uy, fluid, D, n, crop):
    """(s, cycle, phase, <rho>, <ux>, <uy>, <rho>_crop, inlet <rho>, <ux>, outlet <rho>, <ux>)."""
    w = max(1, int(round(STRIP * D)))
    x0, x1, y0, y1 = crop
    cx, cy = slice(x0, x1 + 1), slice(y0, y1 + 1)
    inlet, outlet, inner = slice(1, 1 + w), slice(n - 1 - w, n - 1), slice(1, n - 1)
    return (s, k, ph, fluid_mean(rho, fluid), fluid_mean(ux, fluid), fluid_mean(uy, fluid),
            fluid_mean(rho, fluid, cx, cy),
            fluid_mean(rho, fluid, inlet, inner), fluid_mean(ux, fluid, inlet, inner),
            fluid_mean(rho, fluid, outlet, inner), fluid_mean(ux, fluid, outlet, inner))