"""Model E — fixed-basket shopping optimisation (H6). Preregistration §2 H6 and config/basket.yaml (frozen).

For each complete week: map each basket item to candidate listings per banner (category + name regex, parsable size,
unit price), take the cheapest unit price per banner, and solve min-cost assignment with at most K banners (K=1..5) as a
MIP (PuLP/CBC). Travel costs are unobserved: an illustrative fixed per-banner cost grid is added as a labelled assumption.
"""
import json, re, numpy as np, pandas as pd, pulp, yaml
from src.common import path, ROOT, get_logger, ledger
log = get_logger("modelE"); P = path("processed"); E = ROOT / "experiments"

def candidates(pw, items):
    rows = []
    for it in items:
        d = pw[(pw.category == it["category"]) & pw.name.fillna("").str.contains(it["name_regex"], case=False, regex=True)]
        if it.get("exclude_regex"): d = d[~d.name.fillna("").str.contains(it["exclude_regex"], case=False, regex=True)]
        d = d[(d.unit_base == ("un" if it["base"] == "un" else f"100 {it['base']}")) & d.unit_p.notna() & (d.unit_p > 0)]
        if len(d) == 0: continue
        d = d.assign(item=it["id"], item_cost=lambda x: x.unit_p * (it["qty"] if it["base"] == "un" else it["qty"] / 100.0))
        rows.append(d[["item", "week", "banner", "company", "uid", "name", "size_label", "p_mode", "unit_p", "item_cost"]])
    return pd.concat(rows, ignore_index=True)

def solve(cost_tbl, K, fixed_cost=0.0):
    """cost_tbl: item x banner matrix of min item cost (NaN = unavailable). Returns (total cost, banners used, items covered)."""
    items = list(cost_tbl.index); banners = list(cost_tbl.columns)
    prob = pulp.LpProblem("basket", pulp.LpMinimize)
    x = {(i, b): pulp.LpVariable(f"x_{n}_{m}", cat="Binary") for n, i in enumerate(items) for m, b in enumerate(banners) if not np.isnan(cost_tbl.at[i, b])}
    y = {b: pulp.LpVariable(f"y_{m}", cat="Binary") for m, b in enumerate(banners)}
    prob += pulp.lpSum(cost_tbl.at[i, b] * v for (i, b), v in x.items()) + fixed_cost * pulp.lpSum(y.values())
    for i in items:
        prob += pulp.lpSum(v for (ii, b), v in x.items() if ii == i) == 1
    for (i, b), v in x.items():
        prob += v <= y[b]
    prob += pulp.lpSum(y.values()) <= K
    prob.solve(pulp.PULP_CBC_CMD(msg=False))
    if pulp.LpStatus[prob.status] != "Optimal": return np.nan, [], 0
    used = [b for b, v in y.items() if v.value() > 0.5]
    tot = sum(cost_tbl.at[i, b] * v.value() for (i, b), v in x.items())
    return float(tot), used, len(items)

def main():
    out = {}
    items = yaml.safe_load((ROOT / "config" / "basket.yaml").read_text())["items"]
    pw = pd.read_parquet(P / "panel_weekly.parquet"); sf = pw[pw.channel == "online_storefront"]
    cand = candidates(sf, items); cand.to_parquet(P / "modelE_candidates.parquet", index=False)
    cov = cand.groupby("item").agg(listings=("uid", "nunique"), banners=("banner", "nunique")).reset_index(); cov.to_csv(P / "modelE_item_coverage.csv", index=False)
    log.info("basket items with candidates: %d / %d ; median banners per item %.0f", cov.item.nunique(), len(items), cov.banners.median())
    results = []
    for w, d in cand.groupby("week"):
        tbl = d.groupby(["item", "banner"]).item_cost.min().unstack()
        # keep banners carrying >= 60 % of covered items, and items available at >= 2 banners (so that K=1 is feasible on a common set)
        tbl = tbl.loc[:, tbl.notna().mean() >= 0.6]
        tbl = tbl[tbl.notna().sum(axis=1) >= 2]
        # for K=1 feasibility, restrict to items available at every retained banner
        full = tbl.dropna()
        for K in (1, 2, 3, 4, 5, len(tbl.columns)):
            tot, used, n = solve(full, K)
            results.append(dict(week=w, K=K if K < len(tbl.columns) else 99, cost=tot, banners=";".join(used), n_items=n, n_banners_avail=len(tbl.columns)))
        # single-banner cost for each banner (full item set)
        for b in full.columns:
            results.append(dict(week=w, K=1, cost=float(full[b].sum()), banners=b, n_items=len(full), n_banners_avail=len(tbl.columns), single=True))
        for fc in (5.0, 10.0, 20.0):
            tot, used, n = solve(full, len(full.columns), fixed_cost=fc)
            results.append(dict(week=w, K=99, cost=tot - fc * len(used), banners=";".join(used), n_items=n, n_banners_avail=len(tbl.columns), fixed_cost=fc, n_used=len(used), cost_incl_travel=tot))
    r = pd.DataFrame(results); r.to_csv(P / "modelE_results.csv", index=False)
    base = r[(r.get("single").isna()) & (r.fixed_cost.isna())] if "fixed_cost" in r else r[r.get("single").isna()]
    front = base.groupby("K").cost.mean(); singles = r[r.get("single") == True].groupby("week").cost.agg(["min", "max", "mean"])
    out["items_in_common_set_mean"] = float(base.n_items.mean()); out["banners_available_mean"] = float(base.n_banners_avail.mean())
    out["frontier_mean_cost_by_K"] = {int(k): float(v) for k, v in front.items()}
    c1, c2, cinf = front.get(1), front.get(2), front.get(99)
    out["savings_unconstrained_vs_best_single"] = float(c1 - cinf); out["savings_share_unconstrained"] = float((c1 - cinf) / c1)
    out["share_of_unconstrained_savings_at_K2"] = float((c1 - c2) / (c1 - cinf)) if c1 - cinf > 0 else None
    out["single_banner_cost_spread_mean"] = float((singles["max"] - singles["min"]).mean()); out["single_banner_mean_cost"] = float(singles["mean"].mean())
    out["most_expensive_vs_cheapest_single_share"] = float(((singles["max"] - singles["min"]) / singles["min"]).mean())
    if "fixed_cost" in r:
        tc = r[r.fixed_cost.notna()].groupby("fixed_cost").agg(n_used=("n_used", "mean"), cost_goods=("cost", "mean"), cost_incl=("cost_incl_travel", "mean")).reset_index()
        tc.to_csv(P / "modelE_travel_sensitivity.csv", index=False); out["travel_cost_sensitivity"] = tc.to_dict(orient="records")
    (E / "modelE_summary.json").write_text(json.dumps(out, indent=2, default=float))
    log.info("E: common items %.1f, banners %.1f ; cost K=1 %.2f, K=2 %.2f, unconstrained %.2f ; unconstrained savings %.1f%% ; share captured at K=2 %.2f ; single-banner spread %.2f (%.1f%%)",
             out["items_in_common_set_mean"], out["banners_available_mean"], c1, c2, cinf, 100 * out["savings_share_unconstrained"], out["share_of_unconstrained_savings_at_K2"] or float("nan"), out["single_banner_cost_spread_mean"], 100 * out["most_expensive_vs_cheapest_single_share"])
    ledger("4", "model_e", f"K1 {c1:.2f} K2 {c2:.2f} Kinf {cinf:.2f}")

if __name__ == "__main__":
    main()