"""Wholesale per-buyer forecaster (spec §4.2).

For each (SKU, buyer): strip one-offs (cleaning_wholesale) to get the recurring
rhythm, then forecast that rhythm routed by the buyer's ORDERING ARCHETYPE — the
structure the validation backtest showed is essential (Croston on every buyer,
incl. one-order ones, over-forecasts; tiering fixes it).

The archetype is keyed on EFFECTIVE FREQUENCY = active weeks / weeks-since-first-
order (not an absolute active-week count). Absolute counts don't scale: over a
74-week history a buyer active 6 weeks is sporadic (freq .08), but the old
`>=6 weeks` rule handed it full Croston rhythm and over-forecast it. Effective
frequency is invariant to history length and fair to newly-listed buyers.
Cutoffs calibrated to the real buyer profile (real sporadic export buyers sit at
<=.05; monthly buyers like AHEMA/Tifon/Kaufland are >=.16):

  freq >= 0.33  rhythm        -> Croston rate (size/interval) + churn decay
  0.06..0.33    intermittent  -> damped run-rate
  freq < 0.06   sporadic      -> 0 baseline (no cadence; KAM on-top only)
  < 3 orders                  -> 0 (1-2 orders can't establish a cadence, and two
                                 equal large buys masquerade as a perfect rhythm —
                                 e.g. SBI's two ~75k export deals; on-top only)

Forecast the top-N buyers individually + a single "Other" bucket for the long
tail, then sum to the SKU wholesale baseline. The stripped one-offs are returned
separately — they re-enter additively as KAM on-tops at assembly.
"""
from __future__ import annotations

import numpy as np

from forecast_v4.cleaning_wholesale import strip_one_offs

RHYTHM_MIN_FREQ = 0.33      # >= ~biweekly: genuine cadence -> Croston
SPORADIC_MAX_FREQ = 0.06    # < this (or <2 orders): no recurring baseline
SPARSE_MIN_WEEKS = 3        # need >=3 orders to infer any cadence (2 equal buys fake a rhythm)
CHURN_WEEKS = 8


def _croston_rate(qty) -> float:
    """Per-week expected demand via Croston: EWMA of order size / EWMA of
    inter-order interval."""
    qty = np.asarray(qty, dtype=float)
    nz = np.where(qty > 0)[0]
    if len(nz) == 0:
        return 0.0
    if len(nz) == 1:
        return float(qty[nz[0]] / len(qty))
    d = qty[nz]
    iv = [nz[0] + 1] + [nz[k] - nz[k - 1] for k in range(1, len(nz))]
    a = 0.3
    z, p = float(d[0]), float(iv[0])
    for k in range(1, len(d)):
        z = a * d[k] + (1 - a) * z
        p = a * iv[k] + (1 - a) * p
    return float(z / max(p, 1.0))


def active_weeks(qty, window: int | None = None) -> int:
    qty = np.asarray(qty, dtype=float)
    if window:
        qty = qty[-window:]
    return int((qty > 0).sum())


def effective_freq(qty) -> float:
    """Active weeks / weeks-since-first-order. Invariant to total history length
    and to when a buyer was first listed (a buyer active in all 8 weeks since
    onset reads 1.0, not 8/74)."""
    qty = np.asarray(qty, dtype=float)
    nz = np.where(qty > 0)[0]
    if len(nz) == 0:
        return 0.0
    eff_span = len(qty) - int(nz[0])           # weeks from first order to now
    return len(nz) / max(eff_span, 1)


def classify_buyer(qty) -> str:
    """'rhythm' | 'intermittent' | 'sporadic' from effective frequency + order count."""
    aw = active_weeks(qty)
    freq = effective_freq(qty)
    if aw < SPARSE_MIN_WEEKS or freq < SPORADIC_MAX_FREQ:
        return "sporadic"
    if freq >= RHYTHM_MIN_FREQ:
        return "rhythm"
    return "intermittent"


def forecast_buyer(qty, h: int = 13, churn_weeks: int = CHURN_WEEKS,
                   sparse_window: int = 26) -> np.ndarray:
    """Archetype-routed per-buyer rhythm forecast -> per-week expected qty, length h."""
    qty = np.asarray(qty, dtype=float)
    n = len(qty)
    arch = classify_buyer(qty)
    recent = bool((qty[-churn_weeks:] > 0).any()) if n else False

    if arch == "rhythm":
        rate = _croston_rate(qty)
        if not recent:
            rate *= 0.25                       # churn decay
    elif arch == "intermittent":
        w = qty[-sparse_window:] if n > sparse_window else qty
        rate = float(w.mean()) * (1.0 if recent else 0.5)   # damped run-rate
    else:                                      # sporadic / single order
        rate = 0.0                             # no recurring baseline -> KAM on-top only
    return np.full(h, max(rate, 0.0))


def forecast_wholesale_sku(buyer_qty: dict, h: int = 13, top_n: int = 25):
    """buyer_qty: {buyer_id: weekly array}. Returns
        (total[h], per_buyer_forecast{buyer: array[h]}, oneoff_excess{buyer: array}).
    Forecasts the top_n buyers by recurring volume individually + one "__other__"
    bucket for the rest."""
    cleaned, excess = {}, {}
    for b, q in buyer_qty.items():
        c, e = strip_one_offs(np.asarray(q, dtype=float))
        cleaned[b], excess[b] = c, e

    ranked = sorted(cleaned, key=lambda b: cleaned[b].sum(), reverse=True)
    top, tail = ranked[:top_n], ranked[top_n:]

    per_buyer = {b: forecast_buyer(cleaned[b], h) for b in top}
    if tail:
        other = np.sum([cleaned[b] for b in tail], axis=0)
        per_buyer["__other__"] = forecast_buyer(other, h)

    total = np.sum(list(per_buyer.values()), axis=0) if per_buyer else np.zeros(h)
    return total, per_buyer, excess