"""Lead-block reductions of a daily EHD series into the forecast targets.

Two targets, defined on the continuous 46-day daily EHD series (day 0..45), then reduced
onto the six fixed non-overlapping lead blocks W1..W6:

  Primary   : block flag = (>=1 EHD day within the block's day-range).
  Secondary : heatwave = a run of >= `min_run` consecutive EHD days. A heatwave is
              attributed to EVERY block its run OVERLAPS. Runs are detected on the
              CONTINUOUS 0..45 series first, so a run straddling a block boundary is
              detected once and then flags each block it touches. This is the explicit
              straddle rule (documented in CLAUDE.md / results report).

All functions accept boolean arrays with lead_day as axis 0: shape (46,) or (46, ...).
NaNs in the underlying Tx propagate as NaN block flags (a block with any missing day is
returned NaN so it is excluded from scoring rather than silently counted as 0).
"""
from __future__ import annotations
import numpy as np


def _run_membership(ehd: np.ndarray, min_run: int) -> np.ndarray:
    """Boolean: is day i part of a run of >= min_run consecutive True along axis 0?

    A day is a run-member iff it lies inside at least one fully-True window of length
    `min_run`. Implemented as a dilation of the all-True-window mask.
    """
    n = ehd.shape[0]
    e = ehd.astype(bool)
    member = np.zeros_like(e, dtype=bool)
    for j in range(0, n - min_run + 1):
        win_all = np.ones(e.shape[1:], dtype=bool) if e.ndim > 1 else True
        block = e[j:j + min_run]
        win = np.all(block, axis=0)
        for k in range(min_run):
            member[j + k] |= win
    return member


def block_targets(ehd_daily: np.ndarray, blocks: dict, min_run: int = 3):
    """Compute primary & secondary block flags from a daily EHD boolean array.

    ehd_daily: shape (46, ...) boolean, or float with NaN for missing days.
    blocks: {name: (day_start, day_end)} inclusive.
    Returns (primary, secondary): dicts name -> array over the trailing (...) dims,
    with values in {0.0, 1.0, NaN}.
    """
    is_nan = np.isnan(ehd_daily) if np.issubdtype(ehd_daily.dtype, np.floating) else np.zeros(ehd_daily.shape, bool)
    ehd_bool = np.where(is_nan, False, ehd_daily).astype(bool)
    member = _run_membership(ehd_bool, min_run)

    primary, secondary = {}, {}
    for name, (d0, d1) in blocks.items():
        sl = slice(d0, d1 + 1)
        any_ehd = np.any(ehd_bool[sl], axis=0).astype(float)
        any_hw = np.any(member[sl], axis=0).astype(float)
        # mark NaN where any day in the block is missing (works for scalar or grid)
        blk_missing = np.any(is_nan[sl], axis=0)
        primary[name] = np.where(blk_missing, np.nan, any_ehd)
        secondary[name] = np.where(blk_missing, np.nan, any_hw)
    return primary, secondary