"""Wholesale per-buyer cleaning (spec §3.3).

Wholesale demand is order-driven and lumpy. Per (SKU, buyer):

  - A ONE-OFF is an order much larger than that buyer's own typical order size,
    AND isolated (not part of a sustained run of big orders). These are the
    buyer's own promo sell-ins / KAM on-tops — strip them from the recurring
    rhythm baseline; the excess re-enters additively (the engine forecasts the
    rhythm, KAMs supply the one-offs).
  - A SUSTAINED rise (>= min_run consecutive big orders) is a genuine new
    baseline — KEEP it.
  - NO OOS imputation: a zero week means the buyer didn't order, not a stock-out.

Detection is per-buyer and variance-aware (median + k*MAD of that buyer's own
order sizes) with an absolute floor, so a global threshold doesn't over-flag a
big buyer or spuriously flag a tiny one. Sparse buyers (< min_orders) are left
alone — the sparse tier damps them elsewhere.
"""
from __future__ import annotations

import numpy as np

_MAD_TO_STD = 1.4826


def detect_one_offs(
    qty, k: float = 3.0, floor_units: float = 300.0,
    min_run: int = 3, min_orders: int = 4,
) -> np.ndarray:
    """Boolean mask of one-off weeks in a weekly qty series (zeros = no order).

    A week is a one-off when its order size exceeds `max(median + k*MAD,
    floor_units)` over the buyer's own orders AND it is isolated (a same-level
    run shorter than `min_run` orders). Buyers with fewer than `min_orders`
    total orders are not judged (returns all-False).
    """
    qty = np.asarray(qty, dtype=float)
    mask = np.zeros(len(qty), dtype=bool)
    order_idx = np.where(qty > 0)[0]
    if len(order_idx) < min_orders:
        return mask

    sizes = qty[order_idx]
    med = float(np.median(sizes))
    mad = float(np.median(np.abs(sizes - med))) * _MAD_TO_STD
    bar = max(med + k * mad, floor_units)
    big = sizes > bar

    # Flag isolated big orders (run < min_run); keep sustained runs (a step-up).
    i = 0
    while i < len(big):
        if not big[i]:
            i += 1
            continue
        j = i
        while j + 1 < len(big) and big[j + 1]:
            j += 1
        if (j - i + 1) < min_run:
            for t in range(i, j + 1):
                mask[order_idx[t]] = True
        i = j + 1
    return mask


def strip_one_offs(qty, **kw):
    """Split a weekly qty series into (rhythm_baseline, one_off_excess).

    One-off weeks are reduced to the buyer's typical order size; the removed
    excess is returned separately as the additive (KAM-on-top) layer. Their sum
    reconstructs the original series exactly.
    """
    qty = np.asarray(qty, dtype=float)
    mask = detect_one_offs(qty, **kw)
    nz = qty[qty > 0]
    typical = float(np.median(nz)) if len(nz) else 0.0
    cleaned = qty.copy()
    excess = np.zeros_like(qty)
    excess[mask] = np.maximum(qty[mask] - typical, 0.0)
    cleaned[mask] = qty[mask] - excess[mask]
    return cleaned, excess