"""Extras: current company stock €, sensitivity sweep on floor."""
import pandas as pd
import numpy as np

CURRENT_YEAR, CURRENT_WEEK = 2026, 20
HORIZON_END = (2026, 40)
PO_WINDOW_YW = {2026 * 100 + w for w in range(22, 28)}

inc = pd.read_csv("data/incoming_supply.csv")
sup = pd.read_csv("data/supply_master.csv")
sup["is_abc"] = sup["supplier"].fillna("").str.contains("ABC NUTRITIONAL", case=False)
stock_wh = pd.read_csv("data/stock.csv")
stores = pd.concat([pd.read_csv(f"data/{f}") for f in
                    ["stock_stores.csv", "stock_stores_at.csv", "stock_stores_slo.csv"]],
                   ignore_index=True)
stores.columns = [c.lower() for c in stores.columns]
stores = stores.groupby("sku", as_index=False)["on_hand"].sum()
opening = pd.concat([stock_wh[["sku", "on_hand"]], stores], ignore_index=True)
opening = opening.groupby("sku", as_index=False)["on_hand"].sum()
costs = pd.read_csv("data/sku_costs.csv")
cost_map = dict(zip(costs["sku"], costs["cost_price"]))
fc = pd.read_csv("data/forecast_for_supply.csv")
fc["yw"] = fc["year"] * 100 + fc["week"]
fc_lookup = {(r.sku, r.yw): r.demand for r in fc.itertuples()}

inc["yw"] = inc["year"] * 100 + inc["week"]
inc_m = inc.merge(sup[["sku", "is_abc"]], on="sku", how="left")
inc_m["is_abc"] = inc_m["is_abc"].fillna(False).astype(bool)

# total stock value now
opening["val"] = opening["on_hand"] * opening["sku"].map(cost_map).fillna(0)
total_now_eur = float(opening["val"].sum())
total_now_units = int(opening["on_hand"].sum())

# total incoming € by week (all suppliers + ABC only)
inc_m["unit_cost"] = inc_m["sku"].map(cost_map).fillna(0)
inc_m["eur"] = inc_m["qty"] * inc_m["unit_cost"]
all_by_w = inc_m.groupby(["year", "week"]).agg(units=("qty", "sum"), eur=("eur", "sum"))
abc_by_w = inc_m[inc_m["is_abc"]].groupby(["year", "week"]).agg(units=("qty", "sum"), eur=("eur", "sum"))

print("=" * 60)
print(f"CURRENT COMPANY STOCK (WH + stores HR/AT/SLO):")
print(f"  Units: {total_now_units:,}")
print(f"  EUR  : {total_now_eur:,.0f}")
print()
print("INCOMING BY WEEK (all suppliers):")
print(all_by_w.round(0).to_string())
print()
print("INCOMING BY WEEK (ABC only):")
print(abc_by_w.round(0).to_string())
print()

# ---------- sensitivity sweep ----------
# Try floors of 0w, 1w, 2w, 3w, 4w forward demand
def yw_(y, w): return y * 100 + w
def week_iter(sy, sw, ey, ew):
    y, w = sy, sw
    while (y, w) <= (ey, ew):
        yield y, w
        w += 1
        if w > 52: w, y = 1, y + 1

on_hand_map = dict(zip(opening["sku"], opening["on_hand"]))
inc_lookup_full = {}
for r in inc.itertuples():
    inc_lookup_full[(r.sku, r.yw)] = inc_lookup_full.get((r.sku, r.yw), 0) + r.qty

def project(sku, cancelled):
    s = float(on_hand_map.get(sku, 0))
    out = {}
    for y, w in week_iter(CURRENT_YEAR, CURRENT_WEEK, *HORIZON_END):
        key = (sku, yw_(y, w))
        d = float(fc_lookup.get(key, 0))
        i = float(inc_lookup_full.get(key, 0))
        if key in cancelled:
            i = 0
        s = max(0.0, s + i - d)
        out[yw_(y, w)] = s
    return out

def fwd_avg(sku, y, w, n=8):
    vals = []
    for _ in range(n):
        vals.append(fc_lookup.get((sku, yw_(y, w)), 0))
        w += 1
        if w > 52: w, y = 1, y + 1
    return float(np.mean(vals)) if vals else 0.0

abc_window = inc_m[inc_m["is_abc"] & inc_m["yw"].isin(PO_WINDOW_YW)].copy()

print("SENSITIVITY (cumulative greedy cancel, per SKU):")
print(f"{'Floor':<10} {'POs cancel':>12} {'Qty':>12} {'EUR':>14}")
for floor_w in [0, 1, 2, 3, 4]:
    total_eur = 0
    total_qty = 0
    total_pos = 0
    for sku, g in abc_window.groupby("sku"):
        g_sorted = g.sort_values(["year", "week"])
        cancelled = set()
        for _, po in g_sorted.iterrows():
            pyw = int(po["yw"])
            trial = cancelled | {(sku, pyw)}
            proj = project(sku, trial)
            weeks_from = [w for w in proj if w >= pyw]
            min_after = min(proj[w] for w in weeks_from) if weeks_from else 0
            fa = fwd_avg(sku, po["year"], po["week"])
            if (fa == 0) or (min_after >= floor_w * fa):
                cancelled = trial
                total_qty += int(po["qty"])
                total_eur += float(po["eur"])
                total_pos += 1
    print(f"{floor_w}w {'':<7} {total_pos:>12} {total_qty:>12,} {total_eur:>14,.0f}")
