diff --git a/tests/test_whatif.py b/tests/test_whatif.py new file mode 100644 index 0000000..b58f9c2 --- /dev/null +++ b/tests/test_whatif.py @@ -0,0 +1,29 @@ +from pathlib import Path +import sys +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "tools")) +import whatif + + +def test_apply_add_reuses_existing_recipe(): + spec = {"buildings": [{"building": "POL", "amount": 1, "recipes": ["EPO,PG=>DEC"]}]} + whatif.apply_add(spec, "POL:1") + assert spec["buildings"] == [{"building": "POL", "amount": 2, "recipes": ["EPO,PG=>DEC"]}] + + +def test_apply_add_new_building_needs_explicit_recipe(): + spec = {"buildings": []} + import pytest + with pytest.raises(SystemExit): + whatif.apply_add(spec, "CHP:1") + + +def test_apply_add_appends_new_building_with_recipe(): + spec = {"buildings": []} + whatif.apply_add(spec, "CHP:1:MG,S,CU=>IND") + assert spec["buildings"] == [{"building": "CHP", "amount": 1, "recipes": ["MG,S,CU=>IND"]}] + + +def test_apply_infra_adds_delta(): + spec = {"infrastructure": {"HB1": 4}} + whatif.apply_infra(spec, "HB1:+1,HB2:+2") + assert spec["infrastructure"] == {"HB1": 5, "HB2": 2} diff --git a/tools/whatif.py b/tools/whatif.py new file mode 100644 index 0000000..788ce84 --- /dev/null +++ b/tools/whatif.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python3 +"""Marginal ROI of adding to a base you already run: whole-plan-with-addition minus whole-plan- +without, per docs/library.md's diffing rule (NOT the plan-level ROI simulate.py/PRUNplanner print, +which divides whole-plan profit by only the new capex and badly overstates once a base exists). +Also flags any workforce tier the addition pushes over capacity (the recurring surprise this +session - e.g. PP2 overloading pioneer housing after POL+CHP already maxed it). + + tools/whatif.py plans/examples/deimos_base.yaml --add CHP:1:MG,S,CU=>IND --infra HB1:+1 + tools/whatif.py before.yaml --after after.yaml # two explicit full-base specs instead + tools/whatif.py before.yaml --add POL:1 # reuses POL's existing recipe in before.yaml +""" +import argparse, copy, sys +from pathlib import Path +import yaml +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) +from puga import config, market, prunplanner as pp +from puga.simulate import simulate + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +import plan_push + + +def apply_add(spec: dict, add: str): + """BUILDING:AMOUNT[:RECIPE[|RECIPE...]] - amount is a delta; recipe reuses an existing + entry of the same building ticker in spec if omitted.""" + parts = add.split(":", 2) + if len(parts) < 2: + sys.exit(f"--add needs BUILDING:AMOUNT[:RECIPE], got {add!r}") + ticker, amount = parts[0], int(parts[1]) + recipes = parts[2].split("|") if len(parts) == 3 else None + for b in spec.setdefault("buildings", []): + if b["building"] == ticker and (recipes is None or b["recipes"] == recipes): + b["amount"] += amount + return + if recipes is None: + sys.exit(f"no existing {ticker} in the plan to infer a recipe from - give one: {ticker}:{amount}:RECIPE") + spec["buildings"].append({"building": ticker, "amount": amount, "recipes": recipes}) + + +def apply_infra(spec: dict, infra: str): + """HB1:+1,HB2:+2 - deltas added to spec['infrastructure'].""" + inf = spec.setdefault("infrastructure", {}) + for part in infra.split(","): + ticker, amount = part.split(":") + inf[ticker] = inf.get(ticker, 0) + int(amount) + + +def run(spec, recipes, blds, price, faction, perm, built): + plan = plan_push.build_payload(spec, recipes, {b["building_ticker"] for b in blds}) + planet = pp._g(f"/data/planet/{plan['planet_natural_id']}/", 3600) + return plan, simulate(plan, recipes, blds, planet["resources"], planet["fertility"], price, faction, perm, built=built) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("before", help="full-base plan spec as it stands (or as a baseline you want to compare from)") + ap.add_argument("--after", help="full-base plan spec with the addition already in it; skips --add/--infra") + ap.add_argument("--add", action="append", default=[], help="BUILDING:AMOUNT[:RECIPE[|RECIPE]], repeatable") + ap.add_argument("--infra", action="append", default=[], help="HB1:+1,HB2:+2, repeatable") + ap.add_argument("--cx", default=config.DEFAULT_CX) + ap.add_argument("--basis", default="uni30", choices=["real", "uni30", "vwap30", "vwap7", "ask", "bid", "mid"]) + ap.add_argument("--no-faction", action="store_true") + a = ap.parse_args() + + before_spec = yaml.safe_load(Path(a.before).read_text()) + if a.after: + after_spec = yaml.safe_load(Path(a.after).read_text()) + else: + if not a.add and not a.infra: + sys.exit("give --after, or at least one --add/--infra") + after_spec = copy.deepcopy(before_spec) + for add in a.add: + apply_add(after_spec, add) + for infra in a.infra: + apply_infra(after_spec, infra) + + recipes, blds = pp.recipes(), pp.buildings() + snap = market.snapshot() + + def price(t, side="buy"): + if a.basis == "uni30": + v = market.uni30(snap, t) + if v: + return v + q = snap.get((t, a.cx)) + return q.ask if q else None + q = snap.get((t, a.cx)) + if not q: + return None + return {"real": q.ask if side == "buy" else (q.vwap7 or q.vwap30 or q.bid), + "vwap30": q.vwap30 or q.vwap7 or q.ask, "vwap7": q.vwap7 or q.vwap30 or q.ask, + "ask": q.ask, "bid": q.bid, "mid": (q.ask + q.bid) / 2 if q.ask and q.bid else None}[a.basis] + + st = yaml.safe_load(config.state_path().read_text()) + faction = None if a.no_faction else (st.get("company") or {}).get("faction") + perm = (st.get("permits", {}).get("used", 1), st.get("permits", {}).get("total", 2)) + planet_id = before_spec["planet"] + built = next((b.get("buildings", {}) for b in st.get("bases", []) if b.get("planet") == planet_id), {}) + + # hq comes from each spec file's own "hq:" field (plan_push.build_payload reads it directly) - + # keep both before/after specs' hq set explicitly rather than relying on state here. + plan_b, r_b = run(before_spec, recipes, blds, price, faction, perm, built) + plan_a, r_a = run(after_spec, recipes, blds, price, faction, perm, built) + + d_profit = r_a["profit"] - r_b["profit"] + d_capex = r_a["new_capex"] - r_b["new_capex"] + + print(f"BEFORE profit/day {r_b['profit']:,.0f} new_capex {r_b['new_capex']:,.0f}") + print(f"AFTER profit/day {r_a['profit']:,.0f} new_capex {r_a['new_capex']:,.0f}") + print(f"\nMARGINAL profit/day {d_profit:,.0f} new_capex {d_capex:,.0f}", end="") + if d_capex > 0 and d_profit > 0: + print(f" ROI/day {100*d_profit/d_capex:.2f}% payback {d_capex/d_profit:.1f} d") + elif d_capex <= 0: + print(" (no new capex)") + else: + print(" (marginal profit <= 0 - not worth it)") + + print("\nWORKFORCE (after) need supply open eff%") + newly_overloaded = [] + for t, w in r_a["workforce"].items(): + if w["need"] or w["supply"]: + was_already_overloaded = r_b["workforce"].get(t, {}).get("open", 0) < 0 + flag = "" + if w["open"] < 0 and not was_already_overloaded: + flag = " <- NEWLY OVERLOADED by this addition, drags every building's efficiency down" + newly_overloaded.append(t) + elif w["open"] < 0: + flag = " (already overloaded before this addition too - not new)" + print(f" {t:10} {w['need']:5.0f} {w['supply']:6.0f} {w['open']:5.0f} {w['eff']*100:7.2f}{flag}") + if newly_overloaded: + print(f"\nFix before building: more housing for {', '.join(newly_overloaded)} (see which HB covers each tier).") + + +if __name__ == "__main__": + main()