"""Tests for retail/web cleaning — uplift estimation, promo deflation, OOS
imputation — calibrated to the whey (~5x) vs creatine (~1.6x) patterns."""
import numpy as np

from forecast_v4.cleaning_retail import estimate_uplift, impute_oos, clean_retail


def _series(normal, promo, n_normal=40, n_promo=8):
    qty = np.full(n_normal + n_promo, float(normal))
    qty[n_normal:] = float(promo)
    mask = np.zeros_like(qty, dtype=bool)
    mask[n_normal:] = True
    return qty, mask


def test_whey_uplift_high():
    qty, mask = _series(normal=450, promo=2500, n_promo=8)
    up = estimate_uplift(qty, mask)
    assert 4.0 < up < 7.0, f"whey uplift off: {up:.2f}"


def test_creatine_uplift_low():
    qty, mask = _series(normal=700, promo=1150, n_promo=8)
    up = estimate_uplift(qty, mask)
    assert 1.4 < up < 1.9, f"creatine uplift off: {up:.2f}"


def test_promo_weeks_deflated_to_baseline():
    qty, mask = _series(normal=450, promo=2500, n_promo=8)
    baseline, up = clean_retail(qty, promo_mask=mask)
    # promo weeks should come back down near the normal level
    assert np.all(np.abs(baseline[mask] - 450) < 120), baseline[mask]
    assert np.all(baseline[~mask] == 450)


def test_sparse_promo_shrinks_toward_prior():
    # Only ONE promo week — don't trust its raw 5.5x ratio; shrink toward prior.
    qty, mask = _series(normal=450, promo=2500, n_promo=1)
    up = estimate_uplift(qty, mask, prior=2.0)
    assert up < 4.0, f"sparse uplift not shrunk: {up:.2f}"


def test_no_promo_uplift_is_one():
    qty = np.full(30, 300.0)
    assert estimate_uplift(qty, np.zeros(30, bool)) == 1.0


def test_oos_gap_imputed():
    qty = np.full(20, 100.0)
    qty[10] = 0.0          # out of stock one week
    filled = impute_oos(qty)
    assert abs(filled[10] - 100.0) < 1e-6, f"OOS not imputed: {filled[10]}"


def test_clean_retail_infers_promo_from_price():
    # No mask: a deep price dip should be inferred as a promo and deflated.
    qty = np.full(40, 400.0); qty[20:23] = 1600.0
    price = np.full(40, 20.0); price[20:23] = 14.0   # ~30% off
    baseline, up = clean_retail(qty, price=price)
    assert up > 2.0
    assert baseline[21] < 1000.0, "inferred promo week not deflated"