"""
nd_choice.py — lookup-based nd selector for adaptive chunk processing.

Tags are (nd: int, single_region: bool) tuples:
  nd=-k, True  → 2^k × 2^k super-cell (coarser grid). k=1..9 is benchmarked
                 data; beyond that, best_nd_tag falls back to the exact
                 aggregation-depth formula (get_exact_L) rather than
                 extrapolating the cost model — see best_nd_tag docstring.
  nd>=0, True  → normal nd with SINGLE_REGION forced on
  nd>=0, False → normal nd, multi-region chunking

Usage::
    from .nd_choice import best_nd_tag
    tag = best_nd_tag(sr=2.0, ppc=150)   # -> (nd, single_region)
"""
import math as _math
import os as _os
import numpy as _np
from aabpl import config as _cfg
from .nd_bench_data import BENCH_DATA as _RAW_BENCH

# ---------------------------------------------------------------------------
# Benchmark data: imported from nd_bench_data.py (a plain Python literal,
# generated by temp/regen_nd_tables.py from the raw benchmark checkpoint --
# was a JSON file loaded at import time; moved to a .py module so production
# use needs no file I/O/JSON parsing, just a normal import).
#   BENCH_DATA : sr -> tag_key -> {ppc: ms_per_N}   (raw, no processing)
# tag_key is "{nd},{single_region}" e.g. "-4,True" — parsed straight back
# into the (nd, single_region) tuple used as the tag everywhere else.
# ---------------------------------------------------------------------------

def _parse_tag_key(tag_key):
    _nd_str, _sr_str = tag_key.split(',')
    return (int(_nd_str), _sr_str == 'True')


def _tag_dict(raw_sr_dict):
    return {_parse_tag_key(_tag_key): {float(_ppc): _v for _ppc, _v in _ppc_map.items()}
            for _tag_key, _ppc_map in raw_sr_dict.items()}


_BENCH = {float(_sr): _tag_dict(_tags) for _sr, _tags in _RAW_BENCH.items()}

_SR_VALS = _np.array(sorted(_BENCH.keys()), dtype=float)

# All tags observed across the benchmark data, ordered from deepest
# super-cell to highest subdivision. Built from the data itself (rather
# than hardcoded) so newly benchmarked tags are picked up automatically.
_ALL_TAGS = sorted(
    {tag for _tags in _BENCH.values() for tag in _tags},
    key=lambda t: (t[0], t[1]),
)

# ---------------------------------------------------------------------------
# Interpolation
# ---------------------------------------------------------------------------

def _interp_ppc(timing_dict, ppc):
    """Log-linear interpolation/extrapolation of ms/N at ppc.

    Clamps at the low end; extrapolates log-linearly beyond the high end,
    with the extrapolation slope floored at 1.0 in log-log space (cost can't
    grow slower than linearly in ppc — more points per cell is irreducibly
    more work). This also enforces monotonicity (result >= last measured
    value). Without the floor, tags benchmarked only over a narrow low-ppc
    range (e.g. deep super-cells, deliberately measured only at ppc<6) can
    have a sub-linear empirical slope from just 2-3 sample points, and
    extrapolating that slope 100-1000x beyond the data produces wildly
    understated costs — this is what let e.g. sm6T "win" at ppc=1000 despite
    never being benchmarked past ppc=1.
    """
    if not timing_dict:
        return float('inf')
    ppcs = sorted(timing_dict)
    if ppc <= ppcs[0]:
        return timing_dict[ppcs[0]]
    if ppc >= ppcs[-1]:
        if len(ppcs) < 2 or ppcs[-1] == ppcs[-2]:
            return timing_dict[ppcs[-1]]
        # extrapolate log-linear trend from last two breakpoints, slope floored at 1.0
        lo, hi = ppcs[-2], ppcs[-1]
        log_lo, log_hi = _math.log(lo), _math.log(hi)
        val_lo, val_hi = timing_dict[lo], timing_dict[hi]
        if val_lo <= 0 or val_hi <= 0 or log_hi == log_lo:
            return max(val_hi, 0.0)
        slope = max(1.0, (_math.log(val_hi) - _math.log(val_lo)) / (log_hi - log_lo))
        return val_hi * _math.exp(slope * (_math.log(ppc) - log_hi))
    i = 0
    while ppcs[i + 1] < ppc:
        i += 1
    lo, hi = ppcs[i], ppcs[i + 1]
    t = (_math.log(ppc) - _math.log(lo)) / (_math.log(hi) - _math.log(lo))
    return timing_dict[lo] * (1.0 - t) + timing_dict[hi] * t


def _cost_at(sr_key, tag, ppc):
    return _interp_ppc(_BENCH.get(sr_key, {}).get(tag, {}), ppc)


# ---------------------------------------------------------------------------
# Precomputed per-(sr, tag) time-per-point functions, for the chunk-merge
# cost comparison (estimate separate vs. merged processing time from ppc
# alone). _interp_ppc re-sorts its timing_dict and linear-scans for the
# bracketing breakpoints on every call -- fine for a one-off nd decision,
# wasteful if called repeatedly for the same (sr, tag) with different ppc
# values, which is exactly the merge-decision use case (every candidate
# pair, every round). get_time_fn sorts once and returns a closure over
# plain Python lists/floats (not full numpy arrays -- these timing_dicts
# are tiny, ~10-25 points, so a pure-Python bisect is faster than paying
# numpy's per-call overhead here, same lesson as this session's earlier
# array-vs-dict experiment). Same log-linear interpolation and slope-floored
# extrapolation as _interp_ppc -- this is a faster restatement of that
# function, not a different model.
# ---------------------------------------------------------------------------
_time_fn_cache = {}


def get_time_fn(sr, tag):
    """Return a cached ``f(ppc) -> ms_per_point`` closure for (sr, tag).

    Falls back to a constant-inf function if (sr, tag) has no benchmark
    data at all, so callers can treat every tag uniformly without a
    presence check. Uses exact benchmarked sr only (like _build_nd_curve --
    no cross-sr blending); pass whichever sr key matches cfg.FIXED_SPACING_RATIO.
    """
    key = (sr, tag)
    fn = _time_fn_cache.get(key)
    if fn is not None:
        return fn

    timing_dict = _BENCH.get(sr, {}).get(tag, {})
    if not timing_dict:
        fn = lambda ppc: float('inf')
        _time_fn_cache[key] = fn
        return fn

    ppcs = sorted(timing_dict)
    vals = [timing_dict[p] for p in ppcs]
    log_ppcs = [_math.log(p) for p in ppcs]
    n = len(ppcs)

    # Precompute the high-end extrapolation slope once (same floor-at-1.0
    # rule as _interp_ppc).
    if n >= 2 and vals[-1] > 0 and vals[-2] > 0 and log_ppcs[-1] != log_ppcs[-2]:
        _hi_slope = max(1.0, (_math.log(vals[-1]) - _math.log(vals[-2]))
                        / (log_ppcs[-1] - log_ppcs[-2]))
    else:
        _hi_slope = None  # signals "just return vals[-1]" (flat/degenerate case)

    def _fn(ppc, _ppcs=ppcs, _vals=vals, _logs=log_ppcs, _n=n, _hi_slope=_hi_slope):
        if ppc <= _ppcs[0]:
            return _vals[0]
        if ppc >= _ppcs[-1]:
            if _hi_slope is None:
                return _vals[-1]
            return _vals[-1] * _math.exp(_hi_slope * (_math.log(ppc) - _logs[-1]))
        # bisect for the bracketing breakpoints (timing_dicts are tiny --
        # linear scan is fine and avoids importing bisect for ~20 items).
        i = 0
        while _ppcs[i + 1] < ppc:
            i += 1
        lo, hi = _logs[i], _logs[i + 1]
        t = (_math.log(ppc) - lo) / (hi - lo)
        return _vals[i] * (1.0 - t) + _vals[i + 1] * t

    _time_fn_cache[key] = _fn
    return _fn


def estimate_time_ms(sr, tag, ppc, n_points):
    """Estimated total processing time (ms) for n_points at density ppc,
    using tag's benchmarked cost curve. Thin convenience wrapper around
    get_time_fn for one-off calls; prefer get_time_fn directly when calling
    repeatedly for the same (sr, tag) (e.g. across many candidate merges)."""
    return get_time_fn(sr, tag)(ppc) * n_points


# ---------------------------------------------------------------------------
# Precomputed crossover thresholds per sr
# ---------------------------------------------------------------------------
# _CROSSOVERS[sr] = sorted list of (ppc_lower_bound, best_tag) transitions.
# First entry always starts at ppc=1.  Lookup: find last entry with
# ppc_lower_bound <= query_ppc.

def get_exact_L(ppc, sr):
    """Aggregation depth L (nd=-L) at which expected points/supercell crosses 1.

    ppc is defined relative to the search circle (ppc = density * pi * r^2),
    not the grid cell -- a fine cell has area (r/sr)^2 = r^2/sr^2, so the
    fine-cell density is ppc/(pi*sr^2), not ppc itself. Each smNT level
    doubles cell spacing -> 4x cell area -> 4x expected points, so
    pts_per_supercell_at_L = ppc/(pi*sr^2) * 4^L. Solve for L where that
    equals 1: L = log_4(pi*sr^2/ppc). (Previously this dropped the pi*sr^2
    factor entirely and just used log_4(1/ppc), which systematically
    underestimated the true crossover depth by log_4(pi*sr^2) levels --
    e.g. ~1.8 levels at sr=2, worse for larger sr -- silently leaving
    singleton-group overhead on the table for very sparse outlier data.)

    Used by best_nd_tag beyond the benchmarked range: the old supercell
    geometry builder (_expand_coarse) was combinatorial in 4^k and
    infeasible past k~10-12, but the coarse-cell-batching fast path in
    disk_aggregation_chunk_adaptive_nd.py builds no per-position geometry at
    all (just a block-radius formula + point grouping by row>>k), so there's
    no longer a depth ceiling forcing us to stay within benchmarked data.
    """
    return max(0.0, _math.log(_math.pi * sr * sr / max(ppc, 1e-300), 4))


# Deepest nd actually covered by benchmark data (most negative tag). Beyond
# this, get_exact_L takes over directly rather than extrapolating the cost
# model further — see best_nd_tag.
_DEEPEST_BENCHED_L = -min((t[0] for t in _ALL_TAGS), default=0)

# Ceiling on L. Not driven by any real limit: numpy int right-shift by >=64
# safely returns 0 (correctly meaning "the whole dataset is one supercell"),
# and coarse_spacing = grid_spacing*2**k is a plain float, fine to well past
# k=1000. The only genuine bound is get_exact_L's own floor on ppc
# (1e-300), which caps its own output at get_exact_L(0) ~= 498 — deeper than
# that is not a finer answer, just noise on an already-degenerate input.
_MAX_SUPERCELL_L = 500


def _build_crossovers():
    ppc_grid = list(_np.logspace(-4, 4, 400))  # 1e-4..10000, 400 log-spaced points
    result = {}
    for sr in _SR_VALS:
        sr = float(sr)
        sr_tags = list(_BENCH.get(sr, {}).keys())
        zones = []
        prev_tag = None
        for ppc in ppc_grid:
            best_tag, best_cost = None, float('inf')
            for tag in _ALL_TAGS:
                if tag not in sr_tags:
                    continue
                c = _cost_at(sr, tag, ppc)
                if c < best_cost:
                    best_cost = c
                    best_tag = tag
            if best_tag != prev_tag:
                zones.append((float(ppc), best_tag))
                prev_tag = best_tag
        result[sr] = zones
    return result

_CROSSOVERS = _build_crossovers()


def _smooth_crossovers(crossovers):
    """Enforce nd non-decreasing as ppc increases within each sr's zone list.

    nd should never decrease as ppc rises -- denser data never needs *less*
    subdivision/regions than sparser data. Real crossover tables sometimes
    have a local dip (e.g. sr=2.0: ...,(4159,3),(5003,2),(8313,3),... at
    ppc~5000) from measurement noise in the underlying benchmark rather than
    a genuine reversal. A running-max pass over nd removes these dips at the
    source (every caller benefits, not just documentation) without needing
    to hand-identify which specific zones are noisy.
    """
    smoothed = {}
    for sr, zones in crossovers.items():
        out = []
        running_tag = None
        running_max_nd = -_math.inf
        for ppc, tag in zones:
            nd, _single_region = tag
            if nd >= running_max_nd:
                running_max_nd = nd
                running_tag = tag
                out.append((ppc, tag))
            else:
                # This zone would dip below a depth already reached at a
                # lower ppc -- reuse the tag from whichever prior zone set
                # that running max, not this (noisy) zone's own tag, so
                # single_region stays paired with the nd it actually came
                # from.
                out.append((ppc, running_tag))
        smoothed[sr] = out
    return smoothed


_CROSSOVERS = _smooth_crossovers(_CROSSOVERS)


def _crossover_lookup(zones, ppc):
    """Binary search: return tag for this ppc in precomputed crossover list."""
    lo, hi = 0, len(zones) - 1
    while lo < hi:
        mid = (lo + hi + 1) // 2
        if zones[mid][0] <= ppc:
            lo = mid
        else:
            hi = mid - 1
    return zones[lo][1]


_curve_cache = {}

# Standard ppc grid, mirrored from bench_adaptive.py's PPC_LIST -- observations
# at off-grid ppc (e.g. the anchor*1.4-shifted checks, ratio-variant sweep
# points) aren't shown in the HTML table and shouldn't silently influence the
# curve either. Keep in sync with bench_adaptive.py's PPC_LIST by hand; there's
# no live import across the temp/ <-> aabpl/ boundary.
_STANDARD_PPC_GRID = frozenset([
    0.01, 0.1, 1, 6, 10, 20, 63, 100, 200, 300, 400, 630, 800,
    1000, 1500, 2000, 3000, 4000, 6300, 8000, 10000, 20000, 30000, 100000, 200000,
])


def _build_nd_curve(sr):
    """Single source of truth for "what nd wins at each ppc", shared by
    best_nd_tag and plot_nd_choice_curve (previously duplicated between the
    two -- the plot's construction was more careful than what best_nd_tag
    actually used at runtime, so the documentation plot didn't reflect real
    decisions; this unifies them). Cached per sr.

    Only applies when sr exactly matches a benchmarked sr (sr in _BENCH) --
    interpolating between two different sr's raw observations to build one
    combined keyframe curve isn't well-defined the same way, so best_nd_tag
    keeps the older crossover-zone blending logic for that case.

    Returns (kf_ppc, kf_nd, kf_sreg, low_offset, high_a, high_b):
      kf_*        -- sorted keyframe arrays spanning the tested ppc range.
                     Within this range, nd = kf_nd[j] where j is the last
                     index with kf_ppc[j] <= ppc (a step function).
      low_offset  -- anchors get_exact_L's formula to the curve's own value
                     at its lowest keyframe, for ppc below the tested range.
      high_a/b    -- linear trend (nd ~ a + b*log10(ppc)) fit through the
                     top 2-3 highest-ppc real observations, for ppc above
                     the tested range.

    Keyframe construction: for each tested ppc, the empirically-best tag
    among tags actually measured there. A single-point spike (nd higher
    than both immediate neighbors, reversing right back down) is treated as
    a noisy one-off rather than a real local optimum -- excluded from the
    keyframe sequence, with the resulting gap filled by synthetic
    unit-integer intermediate steps between the two flanking trusted points
    (front-loaded to reach the next confirmed value by the excluded point's
    own ppc, the highest value plausible under monotonicity, not spread
    evenly across the whole gap or left as one big leap).
    """
    if sr in _curve_cache:
        return _curve_cache[sr]

    obs_ppc, obs_nd, obs_sreg = [], [], []
    bd = _BENCH.get(sr, {})
    all_ppc = sorted({p for d in bd.values() for p in d if p in _STANDARD_PPC_GRID})
    for p in all_ppc:
        best_t, best_v = None, float('inf')
        for tag, d in bd.items():
            v = d.get(p)
            if v is not None and v < best_v:
                best_v, best_t = v, tag
        if best_t is not None:
            obs_ppc.append(p)
            obs_nd.append(best_t[0])
            obs_sreg.append(best_t[1])
    obs_ppc_arr = _np.array(obs_ppc)
    obs_nd_arr = _np.array(obs_nd)
    obs_sreg_arr = _np.array(obs_sreg, dtype=bool)

    n_obs = len(obs_nd_arr)
    is_outlier = _np.zeros(n_obs, dtype=bool)
    for i in range(1, n_obs - 1):
        if obs_nd_arr[i] > obs_nd_arr[i - 1] and obs_nd_arr[i] > obs_nd_arr[i + 1]:
            is_outlier[i] = True
    trusted_idx = [i for i in range(n_obs) if not is_outlier[i]]

    keyframes = []
    for k, i in enumerate(trusted_idx):
        keyframes.append((obs_ppc_arr[i], obs_nd_arr[i], obs_sreg_arr[i]))
        if k + 1 < len(trusted_idx):
            j = trusted_idx[k + 1]
            nd_i, nd_j = obs_nd_arr[i], obs_nd_arr[j]
            if nd_i != nd_j:
                # Never leave a >1 nd gap between consecutive trusted
                # keyframes, regardless of whether an outlier was removed
                # between them or the gap is just inherent to sparse
                # testing -- always step through every intermediate integer.
                # Arrival point: front-loaded to the last removed outlier's
                # own ppc if there was one (see docstring), else the next
                # trusted point's own ppc (nothing to front-load to).
                _outliers_between = [obs_ppc_arr[m] for m in range(i + 1, j) if is_outlier[m]]
                arrival_ppc = max(_outliers_between) if _outliers_between else obs_ppc_arr[j]
                step = 1 if nd_j > nd_i else -1
                n_steps = abs(int(nd_j - nd_i)) - 1
                if n_steps > 0:
                    log_positions = _np.linspace(
                        _math.log10(obs_ppc_arr[i]), _math.log10(arrival_ppc), n_steps + 2
                    )[1:-1]
                    for s in range(n_steps):
                        keyframes.append((10.0 ** log_positions[s], nd_i + step * (s + 1), obs_sreg_arr[i]))
                if arrival_ppc != obs_ppc_arr[j]:
                    keyframes.append((arrival_ppc, nd_j, obs_sreg_arr[i]))
    kf_ppc = _np.array([kf[0] for kf in keyframes])
    kf_nd = _np.array([kf[1] for kf in keyframes])
    kf_sreg = _np.array([kf[2] for kf in keyframes], dtype=bool)

    # low_offset below anchors to the TRUE first observation's own ppc/nd
    # (not the shifted boundary computed next), so compute it before
    # shifting kf_ppc.
    if len(kf_ppc):
        floor_L_table = float(-kf_nd[0]) if kf_nd[0] < 0 else 0.0
        floor_L_exact = get_exact_L(float(kf_ppc[0]), sr)
        low_offset = floor_L_table - floor_L_exact
    else:
        low_offset = 0.0

    if n_obs >= 2:
        top = list(zip(obs_ppc, obs_nd))[-3:] if n_obs >= 3 else list(zip(obs_ppc, obs_nd))[-2:]
        xs = _np.log10([p for p, _nd in top])
        ys = _np.array([nd for _p, nd in top], dtype=float)
        high_b, high_a = _np.polyfit(xs, ys, 1)
    else:
        high_a = high_b = 0.0

    # Shift each transition to the geometric midpoint between it and the
    # previous keyframe, so the step happens BETWEEN two competing answers
    # rather than glued exactly to one specific measurement's own ppc (a
    # step function that holds the old value right up until the next
    # keyframe's ppc, then jumps, ties the breakpoint to whichever
    # observation happens to sit on the right -- not a principled
    # crossover location). kf_ppc[0] is left alone: it's the start of the
    # trusted range itself, not a transition between two answers.
    if len(kf_ppc) > 1:
        kf_ppc = kf_ppc.copy()
        kf_ppc[1:] = _np.sqrt(kf_ppc[:-1] * kf_ppc[1:])

    result = (kf_ppc, kf_nd, kf_sreg, low_offset, float(high_a), float(high_b))
    _curve_cache[sr] = result
    return result


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def best_nd_tag(sr, ppc):
    """Return the (nd, single_region) tag with lowest estimated cost at (sr, ppc).

    Uses precomputed crossover thresholds for O(log n) lookup.
    Interpolates between the two nearest sr table entries.

    Beyond the deepest benchmarked super-cell depth, switches to the exact
    aggregation-depth formula (get_exact_L) rather than extrapolating the
    cost model further — at that point there's no ranking data to trust, but
    the analytically "right" depth (≈1 point/cell) is unambiguous, and the
    coarse-batch fast path can build it in O(1) regardless of how deep it is.
    """
    _fixed_nd = getattr(_cfg, 'FIXED_NEST_DEPTH', None)
    if _fixed_nd is not None:
        # Single source of truth for "fixed nd" testing: short-circuit the cost
        # model entirely and force every chunk to this tag. single_region must
        # come from cfg.SINGLE_REGION here (not hardcoded False) -- callers like
        # bench_adaptive.py convey a single_region=True tag (e.g. 's1T') purely
        # via FIXED_NEST_DEPTH + SINGLE_REGION, with no monkeypatch, so dropping
        # this makes _chunk_tag's single_region bit disagree with the region
        # structure the grid was actually built with -> KeyError on
        # cntd_offset_by_region[region_trgl_id] (single-region grids only
        # populate region 0).
        return (_fixed_nd, bool(getattr(_cfg, 'SINGLE_REGION', False)))

    ppc = float(ppc)

    if sr in _BENCH:
        # Exact match to a benchmarked sr: single source of truth shared
        # with plot_nd_choice_curve (_build_nd_curve) -- both step through
        # the same outlier-filtered, intermediate-step-filled keyframes, so
        # the documentation plot always reflects what this function
        # actually decides, not a separately-maintained approximation of it.
        kf_ppc, kf_nd, kf_sreg, low_offset, high_a, high_b = _build_nd_curve(sr)
        if len(kf_ppc):
            lo_obs, hi_obs = float(kf_ppc[0]), float(kf_ppc[-1])
            if ppc > hi_obs:
                # Beyond the highest ppc ever actually benchmarked at this sr:
                # follow the measured trend from the top 2-3 real
                # observations instead of flatlining at the last keyframe,
                # floored at that keyframe's own nd so it never steps back
                # shallower than the last real measurement.
                _nd_pred = high_a + high_b * _math.log10(ppc)
                _nd = max(int(kf_nd[-1]), int(round(_nd_pred)))
                return (_nd, bool(kf_sreg[-1]))
            if ppc < lo_obs:
                # Below the lowest ppc ever tested at this sr: trust
                # get_exact_L's *shape*, anchored so it agrees exactly with
                # the curve's own value at its lowest keyframe (low_offset)
                # -- otherwise the formula's independent zero-point creates
                # a discontinuous jump right at the boundary, which then
                # partially reverses further out (a visible hump).
                _l_exact = get_exact_L(ppc, sr)
                _l = max(0, int(round(_l_exact + low_offset)))
                return (-min(_MAX_SUPERCELL_L, _l), True)
            j = int(_np.searchsorted(kf_ppc, ppc, side='right')) - 1
            j = max(0, j)
            return (int(kf_nd[j]), bool(kf_sreg[j]))

    # sr doesn't exactly match a benchmarked value (or has no observations):
    # cross-sr interpolation between the two nearest benched sr's, using the
    # older crossover-zone tables (unaffected by the above -- this path is
    # only reached for the rarer non-exact-sr case).
    _l_exact = get_exact_L(ppc, sr)
    if _l_exact > _DEEPEST_BENCHED_L:
        # Floor at _DEEPEST_BENCHED_L: this fallback only ever triggers once
        # we're already past the benchmarked depth, so the returned L should
        # never round back to something shallower than that boundary --
        # continuation from where the benchmark data leaves off, not a step
        # backward.
        _l = max(_DEEPEST_BENCHED_L, int(round(_l_exact)))
        return (-min(_MAX_SUPERCELL_L, _l), True)
    ppc = max(1e-6, ppc)
    idx = int(_np.searchsorted(_SR_VALS, sr, side='right')) - 1
    idx = max(0, min(idx, len(_SR_VALS) - 2))
    lo_sr, hi_sr = float(_SR_VALS[idx]), float(_SR_VALS[idx + 1])

    tag_lo = _crossover_lookup(_CROSSOVERS[lo_sr], ppc)
    tag_hi = _crossover_lookup(_CROSSOVERS[hi_sr], ppc)
    if tag_lo == tag_hi:
        return tag_lo

    # Near a crossover boundary: compute actual blended costs for both candidates
    t_sr = 0.0 if hi_sr == lo_sr else max(0.0, min(1.0, (sr - lo_sr) / (hi_sr - lo_sr)))
    candidates = {tag_lo, tag_hi}
    best_tag, best_cost = None, float('inf')
    for tag in candidates:
        cost = ((1.0 - t_sr) * _cost_at(lo_sr, tag, ppc)
                + t_sr        * _cost_at(hi_sr, tag, ppc))
        if cost < best_cost:
            best_cost = cost
            best_tag = tag
    return best_tag


# ---------------------------------------------------------------------------
# Point-count-weighted per-block decision (opt-in via cfg.USE_WEIGHTED_ND_DECISION)
# ---------------------------------------------------------------------------
# best_nd_tag above collapses a chunk's whole density distribution into one
# percentile (p75) before ever consulting the cost model -- correct for a
# roughly-uniform chunk, systematically misleading for one with a real
# density mixture (e.g. a dense hotspot next to a sparse background): p75 of
# occupied-block DENSITIES also can't tell "1 of 100 blocks occupied" from
# "99 of 100 blocks occupied" at the same per-block density, since it never
# sees occupancy count at all. best_nd_tag_weighted instead evaluates the
# real cost model on every occupied block's own ppc (weighted by that
# block's actual point count, so total-workload/occupancy is correctly
# reflected in the sum itself) and blends 50/50 against a total-chunk-AREA
# ppc (points / full chunk footprint including empty blocks -- occupancy-
# aware by construction, robust to any single noisy block).

_sorted_bp_cache = {}

def _get_sorted_breakpoints(sr_key, tag):
    """Cached (ppcs_sorted_array, vals_array) for one (sr_key, tag) -- computed
    once, reused across every block/chunk that queries this tag at this sr.
    _interp_ppc re-sorts timing_dict.keys() on every call; doing that once per
    (sr, tag) instead of once per (block, tag) is what makes per-block
    evaluation viable at all -- see the session discussion on this."""
    ck = (sr_key, tag)
    cached = _sorted_bp_cache.get(ck)
    if cached is not None:
        return cached
    timing_dict = _BENCH.get(sr_key, {}).get(tag, {})
    if not timing_dict:
        result = (None, None)
    else:
        ppcs = sorted(timing_dict)
        result = (_np.array(ppcs, dtype=float), _np.array([timing_dict[p] for p in ppcs], dtype=float))
    _sorted_bp_cache[ck] = result
    return result


def cost_at_array(sr_key, tag, ppc_array):
    """Vectorized _interp_ppc over an array of ppc values, for one exact
    sr_key (must be a value already present in _BENCH -- this is a single-sr
    lookup, no sr-interpolation; callers needing that blend across two nearby
    sr_key calls themselves, same as best_nd_tag does for its own crossover
    lookup). Mirrors _interp_ppc's clamp/extrapolation logic exactly.
    """
    ppc_array = _np.asarray(ppc_array, dtype=float)
    ppcs, vals = _get_sorted_breakpoints(sr_key, tag)
    if ppcs is None:
        return _np.full(ppc_array.shape, _np.inf)
    out = _np.empty(ppc_array.shape, dtype=float)

    lo_mask = ppc_array <= ppcs[0]
    out[lo_mask] = vals[0]

    hi_mask = ppc_array >= ppcs[-1]
    if hi_mask.any():
        if len(ppcs) < 2 or ppcs[-1] == ppcs[-2]:
            out[hi_mask] = vals[-1]
        else:
            lo, hi = ppcs[-2], ppcs[-1]
            log_lo, log_hi = _math.log(lo), _math.log(hi)
            val_lo, val_hi = vals[-2], vals[-1]
            if val_lo <= 0 or val_hi <= 0 or log_hi == log_lo:
                out[hi_mask] = max(val_hi, 0.0)
            else:
                slope = max(1.0, (_math.log(val_hi) - _math.log(val_lo)) / (log_hi - log_lo))
                out[hi_mask] = val_hi * _np.exp(slope * (_np.log(ppc_array[hi_mask]) - log_hi))

    mid_mask = ~lo_mask & ~hi_mask
    if mid_mask.any():
        idx = _np.searchsorted(ppcs, ppc_array[mid_mask])
        idx = _np.clip(idx, 1, len(ppcs) - 1)
        lo_ppc, hi_ppc = ppcs[idx - 1], ppcs[idx]
        t = (_np.log(ppc_array[mid_mask]) - _np.log(lo_ppc)) / (_np.log(hi_ppc) - _np.log(lo_ppc))
        out[mid_mask] = vals[idx - 1] * (1.0 - t) + vals[idx] * t

    return out


def _nearest_sr_key(sr):
    idx = int(_np.searchsorted(_SR_VALS, sr, side='right')) - 1
    idx = max(0, min(idx, len(_SR_VALS) - 1))
    return float(_SR_VALS[idx])


def best_nd_tag_weighted(sr, block_ppcs, block_counts, total_area_ppc,
                          candidate_tags=None, share_occupied=None):
    """Point-count-weighted, occupancy-aware alternative to best_nd_tag.

    block_ppcs / block_counts: per-occupied-coarse-block ppc and point count
        arrays (same length, same order) for the chunk being decided.
    total_area_ppc: points_in_chunk / full_chunk_area (fine cells, including
        empty blocks) * pi*sr^2 -- the occupancy-aware whole-chunk density.
    share_occupied: fraction of coarse blocks in the chunk's bounding box
        that are actually occupied (n_occupied / n_blocks_possible, all
        blocks counted whether populated or not). Drives how much to trust
        total_area_ppc -- see alpha below. Falls back to a raw-count-based
        proxy (old behaviour) if not supplied, e.g. by callers that haven't
        computed the full bounding-box block count.
    candidate_tags: which (nd, single_region) tags to compare; defaults to
        every tag with benchmark data (_ALL_TAGS). Narrowing this is the
        obvious next lever if the full sweep proves too expensive.

    Respects FIXED_NEST_DEPTH the same way best_nd_tag does (short-circuits
    before touching any of the above), so fixed-nd testing is unaffected.
    """
    _fixed_nd = getattr(_cfg, 'FIXED_NEST_DEPTH', None)
    if _fixed_nd is not None:
        return (_fixed_nd, bool(getattr(_cfg, 'SINGLE_REGION', False)))

    if len(block_ppcs) == 0:
        return best_nd_tag(sr, max(total_area_ppc, 1e-6))

    sr_key = _nearest_sr_key(sr)
    block_ppcs = _np.asarray(block_ppcs, dtype=float)
    block_counts = _np.asarray(block_counts, dtype=float)
    total_pts = float(block_counts.sum())
    total_area_ppc = max(total_area_ppc, 1e-6)
    tags = candidate_tags if candidate_tags is not None else _ALL_TAGS

    # alpha = weight on the occupancy-aware total-area-ppc term (the
    # "average" side); (1-alpha) weights the point-count-weighted per-block
    # sum (the "mixture" side). total_area_ppc divides total points by the
    # FULL bounding-box area, empty blocks included -- when most blocks are
    # occupied that's close to an unbiased local-density estimate (trust it);
    # when points cluster in a few blocks out of many, the same average is
    # diluted by all that empty space and badly understates density where
    # points actually are (defer to the per-block mixture sum instead).
    # share_occupied (not raw n_occupied) is what actually captures this: 42
    # occupied blocks out of 50 (dense coverage) and 42 out of 5000 (sparse,
    # clustered) have the same n_occupied but call for opposite trust in the
    # area average.
    _n_occupied = len(block_ppcs)
    if share_occupied is None:
        # Fallback for callers that haven't computed the full bounding-box
        # block count (old count-only heuristic).
        alpha = 0.1 + 0.9 / (1.0 + _n_occupied ** 0.5)
    else:
        alpha = 0.9 * max(0.0, min(1.0, share_occupied))

    best_tag, best_cost = None, float('inf')
    for tag in tags:
        cost = blended_cost_ms(sr_key, tag, block_ppcs, block_counts, total_area_ppc, alpha,
                                total_pts=total_pts)
        if cost < best_cost:
            best_cost = cost
            best_tag = tag
    return best_tag


def blended_cost_ms(sr_key, tag, block_ppcs, block_counts, total_area_ppc, alpha, total_pts=None):
    """Estimated total processing ms for one specific `tag` on this chunk (or
    chunk-piece), using the same point-count-weighted-sum / occupancy-aware-
    average blend as best_nd_tag_weighted's own internal tag search -- shared
    so any other cost comparison (e.g. the row-strip/column-band merge
    decision in disk_aggregation_chunk_adaptive_nd.py) reasons about cost
    with the same model actually used to pick nd, instead of a cheaper but
    inconsistent single-ppc estimate.

    sr_key: an sr value already resolved via _nearest_sr_key (or an exact
        _BENCH key) -- not a raw runtime sr.
    block_ppcs / block_counts / total_area_ppc: see best_nd_tag_weighted.
    alpha: already resolved (0..1), e.g. 0.9 * share_occupied -- see
        best_nd_tag_weighted for how this is normally derived.
    total_pts: optional -- pass to skip recomputing sum(block_counts) when
        the caller already has it.
    """
    block_ppcs = _np.asarray(block_ppcs, dtype=float)
    block_counts = _np.asarray(block_counts, dtype=float)
    if total_pts is None:
        total_pts = float(block_counts.sum())
    total_area_ppc = max(total_area_ppc, 1e-6)
    per_block_cost = cost_at_array(sr_key, tag, block_ppcs)
    sum_cost = float(_np.dot(block_counts, per_block_cost))
    avg_cost = total_pts * float(cost_at_array(sr_key, tag, _np.array([total_area_ppc]))[0])
    return (1.0 - alpha) * sum_cost + alpha * avg_cost


# ---------------------------------------------------------------------------
# Documentation plot: chosen nd across the full ppc range for a given sr.
# ---------------------------------------------------------------------------
def plot_nd_choice_curve(sr=2.0, ppc_lo=0.001, ppc_hi=1_000_000, n_points=600,
                          out_path=None, show=False):
    """Plot best_nd_tag(sr, ppc)'s chosen nd across ppc_lo..ppc_hi (log scale).

    Line color encodes single_region (True/False). Markers overlay the real
    benchmark observations backing this decision (from _BENCH[sr]) at
    whichever ppc values were actually tested -- the empirically-best tag at
    each tested ppc, so markers sit on the curve exactly where it reflects
    real data rather than the get_exact_L extrapolation (see best_nd_tag).

    Saves a PNG to out_path (default: docs/img/nd_choice_by_ppc_sr{sr}.png
    relative to the repo root) unless out_path is False. Returns the
    matplotlib Figure. Requires matplotlib (imported lazily -- not a hard
    dependency of the library, only of this documentation helper).
    """
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D

    # The curve is just best_nd_tag itself, sampled densely -- best_nd_tag
    # uses _build_nd_curve internally for any sr that exactly matches a
    # benchmarked value, so this plot always reflects real decisions rather
    # than a separately-maintained approximation of them.
    ppcs = _np.geomspace(ppc_lo, ppc_hi, n_points)
    nds = _np.empty(n_points)
    sregs = _np.empty(n_points, dtype=bool)
    for i, ppc in enumerate(ppcs):
        nd, sreg = best_nd_tag(sr, float(ppc))
        nds[i], sregs[i] = nd, sreg

    # Real benchmark observations, for the scatter markers only (the curve
    # itself already incorporates these via _build_nd_curve).
    obs_ppc, obs_nd, obs_sreg = [], [], []
    if sr in _BENCH:
        all_ppc = sorted({p for d in _BENCH[sr].values() for p in d if p in _STANDARD_PPC_GRID})
        for p in all_ppc:
            best_t, best_v = None, float('inf')
            for tag, d in _BENCH[sr].items():
                v = d.get(p)
                if v is not None and v < best_v:
                    best_v, best_t = v, tag
            if best_t is not None:
                obs_ppc.append(p)
                obs_nd.append(best_t[0])
                obs_sreg.append(best_t[1])
    obs_ppc_arr = _np.array(obs_ppc)
    obs_nd_arr = _np.array(obs_nd)
    obs_sreg_arr = _np.array(obs_sreg, dtype=bool)

    color_true, color_false = '#2563eb', '#dc2626'   # blue = single_region, red = multi-region

    fig, ax = plt.subplots(figsize=(9, 5), dpi=150)

    # Draw the line in constant-color segments so the color genuinely
    # switches at each single_region transition instead of blending.
    seg_start = 0
    for i in range(1, n_points + 1):
        if i == n_points or sregs[i] != sregs[seg_start]:
            c = color_true if sregs[seg_start] else color_false
            ax.plot(ppcs[seg_start:i], nds[seg_start:i], color=c, linewidth=2.2,
                     solid_capstyle='round', zorder=2)
            seg_start = i

    if len(obs_ppc_arr):
        for mask, c in [(obs_sreg_arr, color_true), (~obs_sreg_arr, color_false)]:
            if mask.any():
                ax.scatter(obs_ppc_arr[mask], obs_nd_arr[mask], color=c, edgecolor='white',
                           linewidth=1.0, s=55, zorder=3)

    from matplotlib.ticker import MultipleLocator, LogLocator

    ax.set_xscale('log')
    ax.set_xlabel('ppc (points per search circle)', fontsize=11)
    ax.set_ylabel('chosen nd', fontsize=11)
    ax.set_title(f'Adaptive nd choice vs. point density (sr={sr:g})', fontsize=13, fontweight='bold')

    # y: integer nd steps only, no fractional gridlines.
    ax.yaxis.set_major_locator(MultipleLocator(1))

    # x: explicit minor ticks at 2..9 within each decade (log scale doesn't
    # always show these by default depending on backend/range), so the grid
    # includes the smaller sub-decade steps, not just each power of 10.
    ax.xaxis.set_minor_locator(LogLocator(base=10.0, subs=tuple(range(2, 10)), numticks=100))

    ax.grid(True, which='major', linestyle='-', linewidth=0.5, alpha=0.35)
    ax.grid(True, which='minor', linestyle=':', linewidth=0.3, alpha=0.2)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    legend_elems = [
        Line2D([0], [0], color=color_true, lw=2.2, label='single_region = True'),
        Line2D([0], [0], color=color_false, lw=2.2, label='single_region = False'),
        Line2D([0], [0], marker='o', color='none', markerfacecolor='#6b7280',
               markeredgecolor='white', markersize=8, label='benchmark observation'),
    ]
    ax.legend(handles=legend_elems, loc='upper left', frameon=True, fontsize=9.5)

    fig.tight_layout()

    if out_path is not False:
        if out_path is None:
            _repo_root = _os.path.abspath(_os.path.join(_os.path.dirname(__file__), '..', '..', '..'))
            out_path = _os.path.join(_repo_root, 'docs', 'img', f'nd_choice_by_ppc_sr{sr:g}.png')
        _os.makedirs(_os.path.dirname(out_path), exist_ok=True)
        fig.savefig(out_path, bbox_inches='tight')
        print(f'saved: {out_path}')

    if show:
        plt.show()
    return fig
