"""ABC 3-scenario postpone/cancel analysis using REAL weekly demand (13w).

Implements the spec exactly:
- Real weeks-of-cover (walk-forward, no smoothing)
- Real post-delivery cover (stock_now + PO_qty, demand from delivery week)
- 3 scenarios via cancel threshold ∈ [None, 26, 36] weeks
- Multi-PO sanity check per SKU
- Company-wide weekly € projection with real weekly outflows
- Movers diff vs 8w-avg classifier
"""
import json
import pandas as pd
import numpy as np

# ---------------- constants ----------------
CURRENT_YEAR = 2026
CURRENT_WEEK = 20
HORIZON_WEEKS = 13                                    # CW20..CW32
PO_WINDOW = [(2026, w) for w in range(22, 28)]        # CW22..CW27
POSTPONE_DELAY = 4
CANCEL_THRESHOLDS = [None, 26, 36]
POSTPONE_COVER_TRIGGER = 4
CURRENT_STOCK_EUR = 4_616_194
PROJECTION_WEEKS = 14                                 # CW20..CW33 for output

# ---------------- load ----------------
inc = pd.read_csv("data/incoming_supply.csv")
sup = pd.read_csv("data/supply_master.csv")
stock_wh = pd.read_csv("data/stock.csv")
stores_frames = []
for f in ["stock_stores.csv", "stock_stores_at.csv", "stock_stores_slo.csv"]:
    s = pd.read_csv(f"data/{f}")
    s.columns = [c.lower() for c in s.columns]
    stores_frames.append(s[["sku", "on_hand"]])
stores = pd.concat(stores_frames, ignore_index=True).groupby("sku", as_index=False)["on_hand"].sum()
fc = pd.read_csv("data/forecast_for_supply.csv")
costs = pd.read_csv("data/sku_costs.csv")
plan = pd.read_csv("data/sku_plan_list.csv")
plan.columns = [c.lower() for c in plan.columns]
plan = plan.rename(columns={"cat": "category", "oznaka": "tier"})

# normalise
sup["is_abc"] = sup["supplier"].fillna("").str.contains("ABC NUTRITIONAL", case=False)
cost_map = dict(zip(costs["sku"], costs["cost_price"]))
plan_min = plan[["sku", "category", "tier"]].copy()
wh_map = dict(zip(stock_wh["sku"], stock_wh["on_hand"]))
stores_map = dict(zip(stores["sku"], stores["on_hand"]))
all_skus = set(stock_wh["sku"]) | set(stores["sku"]) | set(fc["sku"])
def combined_stock(sku):
    return float(wh_map.get(sku, 0)) + float(stores_map.get(sku, 0))

fc["yw"] = fc["year"] * 100 + fc["week"]
fc_lookup = {(r.sku, r.yw): float(r.demand) for r in fc.itertuples()}

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

# ---------------- helpers ----------------
def week_step(y, w, n=1):
    """Step (y, w) forward by n weeks, naive 52-week year."""
    w_new = w + n
    while w_new > 52:
        w_new -= 52
        y += 1
    while w_new < 1:
        w_new += 52
        y -= 1
    return (y, w_new)

def week_series(start_y, start_w, n_weeks):
    out = []
    y, w = start_y, start_w
    for _ in range(n_weeks):
        out.append((y, w))
        y, w = week_step(y, w, 1)
    return out

def yw_int(y, w):
    return y * 100 + w

def demand_series_for(sku, start_y, start_w, n=HORIZON_WEEKS):
    return [fc_lookup.get((sku, yw_int(y, w)), 0.0) for (y, w) in week_series(start_y, start_w, n)]

def real_weeks_cover(stock_now, demand_series):
    """Walk-forward cover. Returns fractional weeks. Extrapolates trailing with series avg."""
    remaining = float(stock_now)
    weeks = 0.0
    for d in demand_series:
        d = float(d)
        if d <= 0:
            weeks += 1
            continue
        if remaining >= d:
            remaining -= d
            weeks += 1
        else:
            weeks += remaining / d
            return weeks
    avg = sum(demand_series) / max(len(demand_series), 1)
    if avg > 0:
        return weeks + remaining / avg
    return weeks + 999  # no demand at all → effectively infinite

def classify(real_cover_now, real_post_delivery_cover, cancel_threshold):
    if cancel_threshold is not None and real_post_delivery_cover > cancel_threshold:
        return "CANCEL"
    if real_cover_now >= POSTPONE_COVER_TRIGGER:
        return "POSTPONE"
    if real_cover_now >= 2:
        return "REVIEW"
    return "PRODUCE"

# ---------------- per-PO core analysis ----------------
abc_window_yw = {yw_int(y, w) for (y, w) in PO_WINDOW}
abc_window = inc_m[inc_m["is_abc"] & inc_m["yw"].isin(abc_window_yw)].copy()

records = []
for _, po in abc_window.iterrows():
    sku = po["sku"]
    po_y = int(po["year"])
    po_w = int(po["week"])
    qty = float(po["qty"])
    cost = float(cost_map.get(sku, 0.0))
    eur_value = qty * cost

    wh_now = float(wh_map.get(sku, 0))
    stores_now = float(stores_map.get(sku, 0))
    combined_now = wh_now + stores_now

    series_now = demand_series_for(sku, CURRENT_YEAR, CURRENT_WEEK, HORIZON_WEEKS)
    real_cover_wh = real_weeks_cover(wh_now, series_now)
    real_cover_combined = real_weeks_cover(combined_now, series_now)

    series_post = demand_series_for(sku, po_y, po_w, HORIZON_WEEKS)
    real_post_delivery = real_weeks_cover(wh_now + qty, series_post)

    avg_13w = float(np.mean(series_now))
    max_13w = float(np.max(series_now)) if series_now else 0.0
    peak_idx = int(np.argmax(series_now)) if series_now else 0
    peak_week_label = f"CW{week_series(CURRENT_YEAR, CURRENT_WEEK, HORIZON_WEEKS)[peak_idx][1]}"

    sa = classify(real_cover_wh, real_post_delivery, None)
    sb = classify(real_cover_wh, real_post_delivery, 26)
    sc = classify(real_cover_wh, real_post_delivery, 36)

    def new_week_for(action):
        if action == "CANCEL":
            return ""
        if action == "POSTPONE":
            ny, nw = week_step(po_y, po_w, POSTPONE_DELAY)
            return f"CW{nw}" + (f"/{ny}" if ny != po_y else "")
        return f"CW{po_w}"

    plan_row = plan_min[plan_min["sku"] == sku]
    tier = plan_row["tier"].iloc[0] if len(plan_row) else ""
    category = plan_row["category"].iloc[0] if len(plan_row) else ""

    records.append({
        "sku": sku, "tier": tier, "category": category,
        "po_year": po_y, "po_week": po_w, "po_label": f"CW{po_w}",
        "qty": int(qty), "cost_price": round(cost, 2),
        "eur_value": round(eur_value, 0),
        "wh_stock_now": int(round(wh_now)),
        "stores_stock_now": int(round(stores_now)),
        "combined_stock_now": int(round(combined_now)),
        "real_weeks_cover_wh": round(real_cover_wh, 2),
        "real_weeks_cover_combined": round(real_cover_combined, 2),
        "real_post_delivery_cover": round(real_post_delivery, 2),
        "avg_weekly_demand_13w": round(avg_13w, 1),
        "max_weekly_demand_13w": round(max_13w, 1),
        "peak_demand_week": peak_week_label,
        "scenario_a": sa, "scenario_b": sb, "scenario_c": sc,
        "new_delivery_week_a": new_week_for(sa),
        "new_delivery_week_b": new_week_for(sb),
        "new_delivery_week_c": new_week_for(sc),
    })

df = pd.DataFrame(records)

# ---------------- multi-PO sanity check per SKU per scenario ----------------
def project_sku_stock_walk(sku, scenario_actions, walk_weeks):
    """Walk WH stock for one SKU through `walk_weeks` with scenario adjustments applied.
    scenario_actions: dict {(sku, year, week): action} — only window POs
    Returns list of (yw_label, stock_end_of_week, demand_that_week).
    """
    # Build adjusted inflows for this SKU
    sku_inflows = {}  # {yw_int: qty}
    sku_pos = inc_m[inc_m["sku"] == sku]
    for _, p in sku_pos.iterrows():
        py, pw = int(p["year"]), int(p["week"])
        pq = float(p["qty"])
        key = (sku, py, pw)
        if key in scenario_actions:
            act = scenario_actions[key]
            if act == "CANCEL":
                continue
            if act == "POSTPONE":
                ny, nw = week_step(py, pw, POSTPONE_DELAY)
                sku_inflows[yw_int(ny, nw)] = sku_inflows.get(yw_int(ny, nw), 0) + pq
                continue
        # PRODUCE / REVIEW / non-window PO → leave on original week
        sku_inflows[yw_int(py, pw)] = sku_inflows.get(yw_int(py, pw), 0) + pq

    s = float(wh_map.get(sku, 0))
    out = []
    for (y, w) in walk_weeks:
        key = yw_int(y, w)
        inflow = sku_inflows.get(key, 0)
        d = fc_lookup.get((sku, key), 0.0)
        s = max(0.0, s + inflow - d)
        out.append((key, s, d))
    return out

walk_weeks = week_series(CURRENT_YEAR, CURRENT_WEEK, HORIZON_WEEKS)  # CW20..CW32

def sanity_check_and_adjust(scenario_col):
    """For each SKU with 2+ POSTPONE POs in window, simulate stacked postpones.
    If stock projection drops below 1 week real cover at any week through CW32,
    un-postpone latest-week PO first, re-check."""
    adjusted = df[scenario_col].copy()

    sku_groups = df.groupby("sku")
    for sku, g in sku_groups:
        window_idx = g.index.tolist()
        # Build actions for this SKU
        actions = {}
        for i in window_idx:
            row = df.loc[i]
            actions[(sku, int(row["po_year"]), int(row["po_week"]))] = adjusted.loc[i]

        postpones = [i for i in window_idx if adjusted.loc[i] == "POSTPONE"]
        if len(postpones) < 2:
            continue

        # iterate flipping latest-week POSTPONE if violated
        while True:
            walk = project_sku_stock_walk(sku, actions, walk_weeks)
            # check: stock(W) >= demand(W) at each week (1w real cover)
            violation = False
            for (key, s, d) in walk:
                if d > 0 and s < d:
                    violation = True
                    break
            if not violation:
                break
            # find latest-week POSTPONE PO still active
            postpone_pos = sorted(
                [i for i in window_idx if actions[(sku, int(df.loc[i, "po_year"]), int(df.loc[i, "po_week"]))] == "POSTPONE"],
                key=lambda i: (df.loc[i, "po_year"], df.loc[i, "po_week"]),
                reverse=True,
            )
            if not postpone_pos:
                break  # can't fix further
            flip = postpone_pos[0]
            actions[(sku, int(df.loc[flip, "po_year"]), int(df.loc[flip, "po_week"]))] = "PRODUCE"
            adjusted.loc[flip] = "PRODUCE"

    return adjusted

df["scenario_a_final"] = sanity_check_and_adjust("scenario_a")
df["scenario_b_final"] = sanity_check_and_adjust("scenario_b")
df["scenario_c_final"] = sanity_check_and_adjust("scenario_c")

# Overwrite the scenario columns with sanity-checked values; keep raw for diagnostics
for s in ["a", "b", "c"]:
    diff_count = int((df[f"scenario_{s}"] != df[f"scenario_{s}_final"]).sum())
    print(f"  Sanity adjustments scenario_{s}: {diff_count} POs un-postponed")
    df[f"scenario_{s}"] = df[f"scenario_{s}_final"]

df = df.drop(columns=["scenario_a_final", "scenario_b_final", "scenario_c_final"])

# Recompute new_delivery_week columns after sanity check
for s in ["a", "b", "c"]:
    def _newweek(row, s=s):
        action = row[f"scenario_{s}"]
        if action == "CANCEL":
            return ""
        if action == "POSTPONE":
            ny, nw = week_step(int(row["po_year"]), int(row["po_week"]), POSTPONE_DELAY)
            return f"CW{nw}" + (f"/{ny}" if ny != int(row["po_year"]) else "")
        return f"CW{int(row['po_week'])}"
    df[f"new_delivery_week_{s}"] = df.apply(_newweek, axis=1)

df.to_csv("three_scenarios_per_po.csv", index=False)
print(f"\nWrote three_scenarios_per_po.csv ({len(df)} rows)")

# ---------------- cash projection ----------------
projection_weeks = week_series(CURRENT_YEAR, CURRENT_WEEK, PROJECTION_WEEKS)  # CW20..CW33
walk_extra = week_series(CURRENT_YEAR, CURRENT_WEEK, PROJECTION_WEEKS + 5)    # walk a bit further for safety

def build_adjusted_inflows(scenario_col):
    """Return list of (sku, year, week, qty) with scenario actions applied to ABC window POs."""
    out = []
    abc_window_keys = {(r["sku"], int(r["po_year"]), int(r["po_week"])): r[scenario_col]
                       for _, r in df.iterrows()}
    for _, p in inc_m.iterrows():
        sku, y, w, q = p["sku"], int(p["year"]), int(p["week"]), float(p["qty"])
        key = (sku, y, w)
        if key in abc_window_keys:
            action = abc_window_keys[key]
            if action == "CANCEL":
                continue
            if action == "POSTPONE":
                ny, nw = week_step(y, w, POSTPONE_DELAY)
                out.append((sku, ny, nw, q))
                continue
        out.append((sku, y, w, q))
    return out

def project_company_eur(adjusted_inflows):
    weekly = {}
    s = float(CURRENT_STOCK_EUR)
    inflow_by_week = {}
    for (sku, y, w, q) in adjusted_inflows:
        cost = cost_map.get(sku, 0.0)
        key = yw_int(y, w)
        inflow_by_week[key] = inflow_by_week.get(key, 0.0) + q * cost

    # CW20 = anchor (start-of-week). Walk from CW21 onward.
    weekly[(CURRENT_YEAR, CURRENT_WEEK)] = s
    for i, (y, w) in enumerate(walk_extra):
        if i == 0:
            continue  # CW20 = anchor, already set
        inflow = inflow_by_week.get(yw_int(y, w), 0.0)
        outflow = sum(fc_lookup.get((sku, yw_int(y, w)), 0.0) * cost_map.get(sku, 0.0) for sku in all_skus)
        s = max(0.0, s + inflow - outflow)
        weekly[(y, w)] = s
    return weekly

# Baseline = all POs as-scheduled (no scenario action applied)
baseline_inflows = [(p["sku"], int(p["year"]), int(p["week"]), float(p["qty"])) for _, p in inc_m.iterrows()]
baseline_proj = project_company_eur(baseline_inflows)
proj_a = project_company_eur(build_adjusted_inflows("scenario_a"))
proj_b = project_company_eur(build_adjusted_inflows("scenario_b"))
proj_c = project_company_eur(build_adjusted_inflows("scenario_c"))

proj_rows = []
for (y, w) in projection_weeks:
    proj_rows.append({
        "week_label": f"CW{w}",
        "week_iso": f"{y}-W{w:02d}",
        "baseline_eur": round(baseline_proj.get((y, w), 0)),
        "scenario_a_eur": round(proj_a.get((y, w), 0)),
        "scenario_b_eur": round(proj_b.get((y, w), 0)),
        "scenario_c_eur": round(proj_c.get((y, w), 0)),
    })
proj_df = pd.DataFrame(proj_rows)
proj_df.to_csv("weekly_inventory_projection.csv", index=False)
print(f"Wrote weekly_inventory_projection.csv ({len(proj_df)} rows)")

# ---------------- 8w-avg classifier for movers diff ----------------
def avg_classify(scenario_col_letter):
    """Run an 8w-avg version of classification to find SKUs that moved."""
    threshold_map = {"a": None, "b": 26, "c": 36}
    cancel_threshold = threshold_map[scenario_col_letter]
    out = {}
    for _, po in df.iterrows():
        sku = po["sku"]
        py, pw = int(po["po_year"]), int(po["po_week"])
        wh_now = float(wh_map.get(sku, 0))

        # 8w avg from now
        series_now_8 = demand_series_for(sku, CURRENT_YEAR, CURRENT_WEEK, 8)
        avg_now = float(np.mean(series_now_8))
        cover_now_avg = wh_now / avg_now if avg_now > 0 else 999

        # post-delivery 8w avg
        series_post_8 = demand_series_for(sku, py, pw, 8)
        avg_post = float(np.mean(series_post_8))
        cover_post_avg = (wh_now + float(po["qty"])) / avg_post if avg_post > 0 else 999

        out[(sku, py, pw)] = classify(cover_now_avg, cover_post_avg, cancel_threshold)
    return out

avg_b = avg_classify("b")
# Diff: compare real-demand scenario_b vs 8w-avg scenario_b
movers_to_postpone = []
movers_to_produce = []
movers_to_cancel = []
movers_off_cancel = []
for _, po in df.iterrows():
    key = (po["sku"], int(po["po_year"]), int(po["po_week"]))
    real = po["scenario_b"]
    avg = avg_b[key]
    if real != avg:
        if real == "POSTPONE" and avg != "POSTPONE":
            movers_to_postpone.append(po["sku"])
        if real == "PRODUCE" and avg != "PRODUCE":
            movers_to_produce.append(po["sku"])
        if real == "CANCEL" and avg != "CANCEL":
            movers_to_cancel.append(po["sku"])
        if avg == "CANCEL" and real != "CANCEL":
            movers_off_cancel.append(po["sku"])

# ---------------- summary JSON ----------------
def scenario_totals(scenario_col):
    sub = df[df[scenario_col].isin(["CANCEL", "POSTPONE"])]
    cancel_eur = float(df[df[scenario_col] == "CANCEL"]["eur_value"].sum())
    postpone_eur = float(df[df[scenario_col] == "POSTPONE"]["eur_value"].sum())
    return {
        "cancel_eur": round(cancel_eur),
        "postpone_eur": round(postpone_eur),
        "total_impact_eur": round(cancel_eur + postpone_eur),
    }

def scenario_counts(scenario_col):
    return {
        "cancel_pos": int((df[scenario_col] == "CANCEL").sum()),
        "postpone_pos": int((df[scenario_col] == "POSTPONE").sum()),
        "cancel_skus": int(df[df[scenario_col] == "CANCEL"]["sku"].nunique()),
        "postpone_skus": int(df[df[scenario_col] == "POSTPONE"]["sku"].nunique()),
    }

def peak_info(proj_map):
    week_eur = {(y, w): proj_map.get((y, w), 0) for (y, w) in projection_weeks}
    peak_yw = max(week_eur, key=week_eur.get)
    peak_eur = week_eur[peak_yw]
    return {"eur": round(peak_eur), "week": f"CW{peak_yw[1]}", "under_5m": bool(peak_eur < 5_000_000)}

summary = {
    "weeks": [f"CW{w}" for (_, w) in projection_weeks],
    "scenarios": {
        "baseline":   [round(baseline_proj.get((y, w), 0)) for (y, w) in projection_weeks],
        "scenario_a": [round(proj_a.get((y, w), 0)) for (y, w) in projection_weeks],
        "scenario_b": [round(proj_b.get((y, w), 0)) for (y, w) in projection_weeks],
        "scenario_c": [round(proj_c.get((y, w), 0)) for (y, w) in projection_weeks],
    },
    "totals": {
        "scenario_a": scenario_totals("scenario_a"),
        "scenario_b": scenario_totals("scenario_b"),
        "scenario_c": scenario_totals("scenario_c"),
    },
    "counts": {
        "scenario_a": scenario_counts("scenario_a"),
        "scenario_b": scenario_counts("scenario_b"),
        "scenario_c": scenario_counts("scenario_c"),
    },
    "peak": {
        "baseline":   peak_info(baseline_proj),
        "scenario_a": peak_info(proj_a),
        "scenario_b": peak_info(proj_b),
        "scenario_c": peak_info(proj_c),
    },
    "movers": {
        "to_postpone_after_real_demand": sorted(set(movers_to_postpone)),
        "to_produce_after_real_demand": sorted(set(movers_to_produce)),
        "to_cancel_after_real_demand": sorted(set(movers_to_cancel)),
        "off_cancel_after_real_demand": sorted(set(movers_off_cancel)),
        "total_pos_reclassified_vs_8w_avg": int(sum(
            df.apply(lambda r: avg_b[(r["sku"], int(r["po_year"]), int(r["po_week"]))] != r["scenario_b"], axis=1)
        )),
    },
}

with open("scenario_summary.json", "w", encoding="utf-8") as f:
    json.dump(summary, f, indent=2)
print("Wrote scenario_summary.json")

# ---------------- supplier-ready files for scenario B ----------------
scen_b_actions = df[df["scenario_b"].isin(["CANCEL", "POSTPONE"])].copy()
scen_b_actions["original_delivery_week"] = scen_b_actions["po_label"]
scen_b_actions["action"] = scen_b_actions["scenario_b"]
scen_b_actions["new_delivery_week_or_blank"] = scen_b_actions["new_delivery_week_b"].where(
    scen_b_actions["scenario_b"] == "POSTPONE", ""
)
scen_b_actions["our_notes"] = ""

cancel_b = scen_b_actions[scen_b_actions["action"] == "CANCEL"][[
    "sku", "tier", "original_delivery_week", "qty", "cost_price", "eur_value",
    "action", "new_delivery_week_or_blank", "our_notes"
]].sort_values(["tier", "eur_value"], ascending=[True, False])
postpone_b = scen_b_actions[scen_b_actions["action"] == "POSTPONE"][[
    "sku", "tier", "original_delivery_week", "qty", "cost_price", "eur_value",
    "action", "new_delivery_week_or_blank", "our_notes"
]].sort_values(["tier", "eur_value"], ascending=[True, False])

cancel_b.to_csv("abc_cancel_list_scenario_b.csv", index=False)
postpone_b.to_csv("abc_postpone_list_scenario_b.csv", index=False)
print(f"Wrote abc_cancel_list_scenario_b.csv ({len(cancel_b)} rows)")
print(f"Wrote abc_postpone_list_scenario_b.csv ({len(postpone_b)} rows)")

# ---------------- console summary ----------------
print("\n" + "=" * 70)
print("HEADLINE NUMBERS")
print("=" * 70)
print(f"Current company stock (anchor): €{CURRENT_STOCK_EUR:,}")
print(f"Baseline peak: €{summary['peak']['baseline']['eur']:,} at {summary['peak']['baseline']['week']}")
print(f"  Under €5M? {summary['peak']['baseline']['under_5m']}")
print()
for s, name in [("a", "A: postpone-only"), ("b", "B: postpone + cancel @ >26w cover"),
                ("c", "C: postpone + cancel @ >36w cover")]:
    p = summary["peak"][f"scenario_{s}"]
    t = summary["totals"][f"scenario_{s}"]
    c = summary["counts"][f"scenario_{s}"]
    print(f"Scenario {name}")
    print(f"  Peak: €{p['eur']:,} at {p['week']} | under €5M? {p['under_5m']}")
    print(f"  Cancel: {c['cancel_pos']} POs / {c['cancel_skus']} SKUs / €{t['cancel_eur']:,}")
    print(f"  Postpone: {c['postpone_pos']} POs / {c['postpone_skus']} SKUs / €{t['postpone_eur']:,}")
    print(f"  Total impact: €{t['total_impact_eur']:,}")
    print()
print(f"Movers vs 8w-avg classification (scenario B): {summary['movers']['total_pos_reclassified_vs_8w_avg']} POs reclassified")
