"""Tests for per-store allocation."""
import numpy as np

from forecast_v4.stores import store_shares, allocate, allocate_retail_to_stores


def test_shares_sum_to_one():
    sq = {1: np.full(30, 100.0), 2: np.full(30, 50.0), 3: np.full(30, 25.0)}
    sh = store_shares(sq)
    assert abs(sum(sh.values()) - 1.0) < 1e-9
    assert sh[1] > sh[2] > sh[3]          # bigger store -> bigger share


def test_closed_store_gets_near_zero():
    # Store 2 stopped selling 26+ weeks ago -> recent share ~0.
    sq = {1: np.full(40, 100.0),
          2: np.array([100.0] * 8 + [0.0] * 32)}
    sh = store_shares(sq, window=26)
    assert sh[2] < 0.02, f"closed store share too high: {sh[2]:.3f}"
    assert sh[1] > 0.98


def test_new_store_gets_share():
    # Store 2 only started recently but sells well now -> meaningful share.
    sq = {1: np.full(40, 100.0),
          2: np.array([0.0] * 32 + [100.0] * 8)}
    sh = store_shares(sq, window=26)
    assert sh[2] > 0.2, f"new store under-weighted: {sh[2]:.3f}"


def test_allocate_sums_to_total():
    sh = {1: 0.5, 2: 0.3, 3: 0.2}
    alloc = allocate(1000.0, sh)
    assert abs(sum(float(v) for v in alloc.values()) - 1000.0) < 1e-6


def test_allocate_retail_splits_within_region():
    stores = {
        1: {"name": "Arena", "region": "HR", "qty": np.full(30, 200.0)},
        2: {"name": "Split", "region": "HR", "qty": np.full(30, 100.0)},
        3: {"name": "Koper", "region": "SI", "qty": np.full(30, 50.0)},
    }
    out = allocate_retail_to_stores({"HR": 900.0, "SI": 300.0}, stores)
    # HR total 900 split 2:1 between Arena/Split; SI 300 all to Koper.
    assert abs(out[1]["total13"] - 600.0) < 1.0
    assert abs(out[2]["total13"] - 300.0) < 1.0
    assert abs(out[3]["total13"] - 300.0) < 1.0