"""Regime / level-shift primitive — the shared cleaning building block (spec §3.1).

Given a weekly series, distinguish a TRANSIENT deviation (a 1-3 week dip/spike
that reverts) from a SUSTAINED level shift (a new level that persists). Both
channels need this:

  - Retail price: a transient dip below the local normal = a PROMO; a sustained
    shift = a REPRICE (re-baseline the normal, don't call it a promo). This is
    what `regime_normal_price` produces — the fix for the "reprice trap" where a
    global / last-12-week normal mislabels a permanent reprice as a discount.
  - Wholesale per-buyer qty: a transient spike above the buyer's level = a
    ONE-OFF (KAM on-top / buyer promo, to strip); a sustained rise = a new
    baseline (to keep). This is what `classify_weeks` labels.

Pure functions over 1-D weekly arrays — no DB, no I/O — so they unit-test cleanly.
"""
from __future__ import annotations

import numpy as np
import pandas as pd


def local_baseline(
    y, window: int = 13, center: bool = True, exclude_mask=None,
    min_periods: int | None = None,
) -> np.ndarray:
    """Regime-tracking baseline: a rolling median that FOLLOWS sustained level
    shifts but IGNORES transient 1-3 week deviations (a short dip/spike barely
    moves a 13-week median; a persistent shift moves it).

    y            : 1-D weekly values.
    window       : rolling window in weeks (13 ~= one quarter).
    center       : True for historical cleaning (sees both sides); False for the
                   live forecast edge (trailing only — no future leakage).
    exclude_mask : optional bool array; True weeks are excluded from the median
                   (e.g. known promo weeks, so the normal reflects shelf price).
                   Excluded weeks still receive a baseline (from neighbours).
    Returns an np.ndarray aligned to y (no NaNs).
    """
    s = pd.Series(np.asarray(y, dtype=float))
    inc = s.mask(np.asarray(exclude_mask, dtype=bool)) if exclude_mask is not None else s
    mp = min_periods if min_periods is not None else max(3, window // 3)
    base = inc.rolling(window, center=center, min_periods=mp).median()
    # Fill edges / windows that were entirely excluded.
    base = base.bfill().ffill()
    if base.isna().all():  # series was all-excluded / empty -> fall back to overall
        fill = float(np.nanmedian(s.to_numpy())) if len(s) else 0.0
        base = pd.Series(np.full(len(s), fill))
    return base.to_numpy()


def classify_weeks(
    y, baseline=None, window: int = 13, rel_thresh: float = 0.5, min_run: int = 3,
) -> np.ndarray:
    """Label each week relative to the regime baseline:

      'normal' : within rel_thresh of the baseline
      'event'  : a transient deviation (a same-direction run shorter than min_run)
      'shift'  : part of a sustained run (>= min_run weeks) at a new level

    rel_thresh : relative deviation that counts as a deviation. Lumpy wholesale
                 wants a large value (e.g. 0.5+); retail price wants small (~0.08).
    min_run    : consecutive same-direction deviating weeks to call it a shift.

    Returns an np.ndarray[str] of labels aligned to y.
    """
    y = np.asarray(y, dtype=float)
    n = len(y)
    if baseline is None:
        baseline = local_baseline(y, window=window)
    baseline = np.asarray(baseline, dtype=float)

    out = np.array(["normal"] * n, dtype=object)
    if n == 0:
        return out

    with np.errstate(divide="ignore", invalid="ignore"):
        dev = np.where(baseline > 0, (y - baseline) / baseline, 0.0)
    deviating = np.abs(dev) > rel_thresh
    direction = np.sign(dev)

    i = 0
    while i < n:
        if not deviating[i]:
            i += 1
            continue
        j = i
        while j + 1 < n and deviating[j + 1] and direction[j + 1] == direction[i]:
            j += 1
        run_len = j - i + 1
        label = "shift" if run_len >= min_run else "event"
        out[i:j + 1] = label
        i = j + 1
    return out


def regime_normal_price(
    price, promo_mask=None, window: int = 13, center: bool = True,
) -> np.ndarray:
    """Regime-local NORMAL (non-promo) shelf price per week.

    Computed as a rolling median of the NON-promo weeks, so it tracks permanent
    reprices (15.5 -> 21.6 -> 24.6) while ignoring transient promo dips. Use with
    promo weeks from the ERP calendar (`promo_mask`); when no mask is given,
    promo weeks are inferred as transient downward deviations.

    Returns an np.ndarray of normal price aligned to `price` (no NaNs).
    """
    price = np.asarray(price, dtype=float)
    if promo_mask is None:
        labels = classify_weeks(price, window=window, rel_thresh=0.08, min_run=3)
        # promo = transient downward event; a reprice is a 'shift' (kept as normal)
        base0 = local_baseline(price, window=window, center=center)
        promo_mask = (labels == "event") & (price < base0)
    return local_baseline(price, window=window, center=center, exclude_mask=promo_mask)


def discount_pct(price, normal=None, **kw) -> np.ndarray:
    """Per-week discount vs the regime-local normal price (fraction, 0..1).
    Positive => below normal (a real promo). Near-zero across a reprice (the
    trap avoided)."""
    price = np.asarray(price, dtype=float)
    if normal is None:
        normal = regime_normal_price(price, **kw)
    normal = np.asarray(normal, dtype=float)
    with np.errstate(divide="ignore", invalid="ignore"):
        d = np.where(normal > 0, 1.0 - price / normal, 0.0)
    return d