"""Retail / web cleaning (spec §3.2).

Recovers the clean, non-promo baseline for a retail or web weekly series:

  - Promo weeks come from the ERP calendar (our own promos, ground truth). If no
    mask is supplied but a price series is, promos are inferred as regime-local
    discounts > 10% (so a permanent reprice is NOT mislabelled — see regime.py).
  - Per-SKU uplift is estimated from the data (promo-week vs normal-week volume),
    empirically 1.5x (creatine) to 7x (whey) — the old 1.35 default is wrong for
    most. Sparse-promo SKUs are shrunk toward a prior.
  - Promo weeks are deflated by the uplift to recover baseline demand.
  - OOS gaps (near-zero weeks amid normal sales) are imputed — retail/web only.
"""
from __future__ import annotations

import numpy as np

from forecast_v4.regime import discount_pct


def estimate_uplift(
    qty, promo_mask, prior: float = 2.0, full_weight_at: int = 6,
    clamp: tuple = (1.0, 10.0),
) -> float:
    """Per-SKU promo uplift = mean(promo-week volume) / mean(normal-week volume),
    shrunk toward `prior` when promo weeks are sparse, clamped to a sane range.
    Returns 1.0 if no promo weeks are observed (nothing to deflate)."""
    qty = np.asarray(qty, dtype=float)
    m = np.asarray(promo_mask, dtype=bool)
    promo_v = qty[m & (qty > 0)]
    norm_v = qty[(~m) & (qty > 0)]
    if len(promo_v) < 1:
        return 1.0
    if len(norm_v) < 3 or np.mean(norm_v) <= 0:
        return float(prior)
    raw = float(np.mean(promo_v) / np.mean(norm_v))
    w = min(len(promo_v) / float(full_weight_at), 1.0)   # confidence in `raw`
    up = w * raw + (1.0 - w) * prior
    return float(np.clip(up, *clamp))


def impute_oos(qty, thresh_frac: float = 0.05, window: int = 8) -> np.ndarray:
    """Fill out-of-stock gaps: weeks below `thresh_frac` of the series median are
    replaced by the mean of nearby in-stock weeks. Retail/web only (a wholesale
    zero is 'didn't order', not a stock-out)."""
    qty = np.asarray(qty, dtype=float).copy()
    nz = qty[qty > 0]
    if len(nz) < 3:
        return qty
    thr = thresh_frac * float(np.median(nz))
    for t in range(len(qty)):
        if qty[t] < thr:
            lo, hi = max(0, t - window), min(len(qty), t + window + 1)
            nb = [qty[i] for i in range(lo, hi) if i != t and qty[i] >= thr]
            if nb:
                qty[t] = float(np.mean(nb))
    return qty


def clean_retail(qty, promo_mask=None, price=None, uplift=None, impute: bool = True):
    """Return (baseline, uplift): the clean non-promo baseline plus the uplift used.

    promo_mask : from the ERP calendar (preferred). If None and `price` is given,
                 inferred as regime-local discount > 10%.
    uplift     : override; otherwise estimated per-SKU from the data.
    """
    qty = np.asarray(qty, dtype=float).copy()
    if promo_mask is None:
        promo_mask = (discount_pct(price) > 0.10) if price is not None else np.zeros(len(qty), bool)
    promo_mask = np.asarray(promo_mask, dtype=bool)
    if uplift is None:
        uplift = estimate_uplift(qty, promo_mask)
    out = qty.copy()
    if uplift > 1.0:
        out[promo_mask] = qty[promo_mask] / uplift
    if impute:
        out = impute_oos(out)
    return out, float(uplift)