"""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.

Amendments logged in research/decisions.md (estimand unchanged):
  A-001 complete-block design: the banner set is the B banners with the widest basket coverage and the item set is the
        items all B banners carry (B = 5 primary; 4, 6, 7 sensitivity), so that K = 1 is feasible without imputation.
  A-002 candidate screen: listings whose unit price is outside [1/4, 4] x the item's cross-banner median unit price are
        dropped (mis-parsed sizes / wrong matches produce absurd unit prices).
"""
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
        med = d.unit_p.median(); d = d[(d.unit_p >= med / 4) & (d.unit_p <= med * 4)]  # A-002
        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 B in (4, 5, 6, 7):
        for w, d in cand.groupby("week"):
            tbl = d.groupby(["item", "banner"]).item_cost.min().unstack()
            top = tbl.notna().sum().sort_values(ascending=False).index[:B]      # A-001: widest-coverage banners
            full = tbl[list(top)].dropna()                                        # items carried by all B banners
            if len(full) < 5: continue
            for K in (1, 2, 3, 4, 5, B):
                if K > B: continue
                tot, used, n = solve(full, K)
                results.append(dict(B=B, week=w, K=K if K < B else 99, cost=tot, banners=";".join(used), n_items=n, n_banners_avail=B))
            for b in full.columns:
                results.append(dict(B=B, week=w, K=1, cost=float(full[b].sum()), banners=b, n_items=len(full), n_banners_avail=B, single=True))
            if B == 5:
                for fc in (5.0, 10.0, 20.0):
                    tot, used, n = solve(full, B, fixed_cost=fc)
                    results.append(dict(B=B, week=w, K=99, cost=tot, banners=";".join(used), n_items=n, n_banners_avail=B, fixed_cost=fc, n_used=len(used), cost_incl_travel=tot + fc * len(used)))
    # Fixed block over time (for QGPAI component 4): B=5 widest-coverage banners over the whole period and the items they
    # all carry in EVERY week, so that the weekly basket is constant.
    cov_all = cand.groupby(["banner", "week"]).item.nunique().groupby("banner").mean().sort_values(ascending=False); top5 = list(cov_all.index[:5])
    tbls = {w: d[d.banner.isin(top5)].groupby(["item", "banner"]).item_cost.min().unstack().reindex(columns=top5) for w, d in cand.groupby("week")}
    common_items = set.intersection(*[set(t.dropna().index) for t in tbls.values()])
    fixed = []
    for w, t in tbls.items():
        full = t.loc[sorted(common_items)]
        for K in (1, 2, 5):
            tot, used, n = solve(full, K); fixed.append(dict(week=w, K=K if K < 5 else 99, cost=tot, banners=";".join(used), n_items=n))
    fx = pd.DataFrame(fixed); fx.to_csv(P / "modelE_fixed_block.csv", index=False)
    out["fixed_block"] = dict(banners=top5, n_items=len(common_items), mean_cost_K1=float(fx[fx.K == 1].cost.mean()), mean_cost_K2=float(fx[fx.K == 2].cost.mean()), mean_cost_unc=float(fx[fx.K == 99].cost.mean()))
    log.info("fixed block: banners %s, %d items constant over weeks; K1 %.2f K2 %.2f unc %.2f", top5, len(common_items), out["fixed_block"]["mean_cost_K1"], out["fixed_block"]["mean_cost_K2"], out["fixed_block"]["mean_cost_unc"])
    r = pd.DataFrame(results); r.to_csv(P / "modelE_results.csv", index=False)
    r5 = r[r.B == 5]
    base = r5[(r5.single.isna()) & (r5.fixed_cost.isna())]
    front = base.groupby("K").cost.mean(); singles = r5[r5.single == True].groupby("week").cost.agg(["min", "max", "mean"])
    sens = []
    for B, rb in r.groupby("B"):
        bb = rb[(rb.single.isna()) & (rb.fixed_cost.isna())].groupby("K").cost.mean()
        sens.append(dict(B=B, n_items=int(rb.n_items.iloc[0]), cost_K1=bb.get(1), cost_K2=bb.get(2), cost_unconstrained=bb.get(99), savings_share=(bb.get(1) - bb.get(99)) / bb.get(1), share_at_K2=(bb.get(1) - bb.get(2)) / (bb.get(1) - bb.get(99))))
    pd.DataFrame(sens).to_csv(P / "modelE_block_sensitivity.csv", index=False); out["block_sensitivity"] = sens
    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 = r5[r5.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()