"""v4 bottom-up assembly with top-down reconciliation (spec §1, §5).

The full-engine backtest showed pure bottom-up loses to a top-down aggregate on
high-volume (Gold) SKUs — aggregation cancels noise the leaf sum can't. The
diagnostic confirmed it's genuine aggregation error (cleaning even helps), so we
reconcile: forecast the SKU total TOP-DOWN (accurate) and the leaves BOTTOM-UP
(the per-region / per-buyer detail), then scale the leaves to sum to the
top-down total. Reconciled total == top-down accuracy; leaves keep their shares.

Each leaf is forecast on its CLEANED baseline (regime-aware), and the top-down
total is forecast on the summed cleaned series — consistent basis. The promo
uplift and per-buyer one-off layers are reported separately (the planner/KAM
adjustments), never baked into the baseline.
"""
from __future__ import annotations

import numpy as np

from forecast_v4.data import load_sku, next_iso_week, load_store_retail, horizon_weeks
from forecast_v4.stores import allocate_retail_to_stores
from forecast_v4.cleaning_retail import clean_retail
from forecast_v4.cleaning_wholesale import strip_one_offs
from forecast_v4.retail import _backtest_pick, _ses, _wma, _seasonal, _apply_bounds
from forecast_v4.wholesale import forecast_buyer
from forecast_v4.regime import discount_pct

_EPS = 1e-9


def _promo_mask(weeks, price, promo_yws, covered_yws):
    """ERP-calendar promo flag where the SKU is covered; regime-local price
    discount (>10%) as fallback for weeks outside calendar coverage."""
    cov = np.array([w in covered_yws for w in weeks])
    cal = np.array([w in promo_yws for w in weeks])
    pricep = discount_pct(np.asarray(price, float)) > 0.10
    return np.where(cov, cal, pricep)


def _retail_model(series, iso, start_cw, h):
    m = _backtest_pick(series, iso, start_cw, test=4)
    fc = (_seasonal(series, iso, start_cw, h) if m == "seasonal"
          else _ses(series, h) if m == "ses" else _wma(series, h))
    return _apply_bounds(fc, series), m


def forecast_sku(pid: int, h: int = 13, max_yw: int = 999999,
                 reconcile: bool = True, start_yw: int | None = None) -> dict | None:
    d = load_sku(pid, max_yw=max_yw)
    if d is None:
        return None
    # Anchor the forecast horizon to a GLOBAL week (latest data week across all
    # SKUs) so every SKU forecasts the same W+1..W+h, not its own last-sale week.
    anchor = start_yw if start_yw is not None else d["weeks"][-1]
    iso, start_cw = d["iso"], next_iso_week(anchor)
    T = len(d["weeks"])
    leaves = []                       # (kind, key, fc[h], cleaned_series[T], meta)
    raw_total = np.zeros(T)           # total DEMAND (incl. promos/one-offs) -> top-down basis

    # ---- retail / web leaves (cleaned baseline for shares) ----
    for kind, store in (("retail", d["retail"]), ("web", d["web"])):
        for reg, (qa, price) in store.items():
            if qa.sum() <= 0:
                continue
            mask = _promo_mask(d["weeks"], price, d["promo_yws"], d["covered_yws"])
            base_series, up = clean_retail(qa, promo_mask=mask)
            fc, model = _retail_model(base_series, iso, start_cw, h)
            leaves.append((kind, reg, fc, base_series,
                           {"model": model, "uplift": round(up, 2)}))
            raw_total += np.asarray(qa, float)

    # ---- wholesale leaves: per-buyer cleaned rhythm ----
    oneoff_total = 0.0
    for bpid, q in d["wholesale"].items():
        q = np.asarray(q, float)
        cleaned, excess = strip_one_offs(q)
        oneoff_total += float(excess.sum())
        fc = forecast_buyer(cleaned, h)
        leaves.append(("wholesale", bpid, fc, cleaned, {}))
        raw_total += q

    bottom_up_sum = float(sum(lf[2].sum() for lf in leaves))

    # ---- top-down total demand on the RAW summed series (the accurate aggregate,
    #      matches the backtest "top-down" winner; leaves below are scaled to it) ----
    td_fc, td_model = _retail_model(raw_total, iso, start_cw, h)
    top_down_total = float(td_fc.sum())

    # ---- reconcile leaves to the top-down total ----
    scale = (top_down_total / bottom_up_sum) if (reconcile and bottom_up_sum > _EPS) else 1.0

    out = {"pid": pid, "retail": {}, "web": {}, "wholesale": {},
           "td_model": td_model, "scale": round(scale, 3),
           "weeks": horizon_weeks(anchor, h)}
    ws_total = 0.0
    ws_buyers = 0
    retail_week = np.zeros(h)      # forecast_retail incl. web, per week (reconciled)
    ws_week = np.zeros(h)          # forecast_wholesale per week (reconciled)
    for kind, key, fc, _series, meta in leaves:
        sfc = np.asarray(fc, float) * scale
        if kind == "wholesale":
            ws_total += float(sfc.sum())
            ws_buyers += 1
            ws_week += sfc
        else:
            out[kind][key] = {**meta, "total13": float(sfc.sum()), "week": sfc}
            retail_week += sfc
    out["wholesale"] = {"total13": ws_total, "n_buyers": ws_buyers,
                        "historical_oneoff_units": round(oneoff_total), "week": ws_week}
    out["retail_week"] = retail_week           # retail + web, reconciled, per week
    out["wholesale_week"] = ws_week
    out["total_week"] = retail_week + ws_week  # = top-down total in aggregate, sums to total13
    out["total13"] = float((retail_week + ws_week).sum())
    out["bottom_up_total13"] = bottom_up_sum
    return out


def forecast_sku_stores(pid: int, h: int = 13, max_yw: int = 999999) -> dict | None:
    """forecast_sku plus per-store allocation of the retail forecast: each
    region's retail total13 is split across that region's stores by recency-
    weighted share. Adds out['stores'] = {store_id: {name, region, total13, share}}."""
    r = forecast_sku(pid, h=h, max_yw=max_yw)
    if r is None:
        return None
    retail_by_region = {reg: v["total13"] for reg, v in r["retail"].items()}
    _, stores = load_store_retail(pid, max_yw=max_yw)
    r["stores"] = allocate_retail_to_stores(retail_by_region, stores)
    return r