PuGa/tools/whatif.py

#!/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 collections import Counter
from pathlib import Path
import yaml
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from puga import config, contracts, fio, market, prunplanner as pp, shipping
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_remove(spec: dict, remove: str):
    """BUILDING[:AMOUNT] - a teardown scenario: removes AMOUNT (default: all) of that building
    from the plan. For "should we demolish this" questions - pass the current full base as
    --after and the --remove result as the plain positional arg (or vice versa; whatif always
    reports AFTER minus BEFORE, so order which one is which by what you want the sign to mean)."""
    parts = remove.split(":", 1)
    ticker = parts[0]
    amount = int(parts[1]) if len(parts) == 2 else None
    for b in spec.get("buildings", []):
        if b["building"] == ticker:
            b["amount"] = max(0, b["amount"] - amount) if amount else 0
    spec["buildings"] = [b for b in spec.get("buildings", []) if b["amount"] > 0]


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 reclaim_value(ticker: str, qty: int, built_count: int, cx: str):
    """Sell value of the materials you'd actually get back demolishing `qty` of a building you
    already own `built_count` of - real per-site data from FIO's own reclaimables endpoint
    (demolition is NOT a dead loss: observed ~50-90%% of original materials back, varies by
    material, not a fixed fraction - don't approximate it, look it up). None if you don't
    actually own one (can't reclaim a hypothetical building)."""
    if built_count <= 0 or qty <= 0:
        return None
    u = fio.me()
    rec = fio.own(f"/rain/userplanetbuildingreclaimables/{u}", ttl=3600)
    totals = Counter()
    for x in rec:
        if x["BuildingId"].startswith(f"{ticker}-"):
            totals[x["Material"]] += x["Amount"]
    if not totals:
        return None
    value = 0.0
    items = {}
    for mat, amt in totals.items():
        q = round(amt / built_count * min(qty, built_count))  # amt is already summed over all built_count units
        if q <= 0:
            continue
        w = market.walk(mat, cx, q, "sell")
        items[mat] = q
        value += w["total"]
    return {"value": value, "items": items}


def trip_tonnage(flows: dict, mats: dict):
    """Total t/m3 the base needs to haul IN (net importers) and ship OUT (net exporters) per day -
    same accounting scan.py already uses for its own trip-cost estimate (max of the four, since
    imports and exports share the same round trips)."""
    w_in = w_out = v_in = v_out = 0.0
    for t, f in flows.items():
        net = f["inp"] - f["out"]
        m = mats.get(t)
        if not m or abs(net) < 1e-6:
            continue
        if net > 0:
            w_in += net * m["Weight"]; v_in += net * m["Volume"]
        else:
            w_out += -net * m["Weight"]; v_out += -net * m["Volume"]
    return w_in, w_out, v_in, v_out


def marginal_freight(r_b, r_a, planet_id, cx, sf_price, ff_price, mats):
    """AIC/day freight for the EXTRA trips this addition requires - not the base's whole freight
    bill. Continuous amortized rate (extra tonnage/cargo-hold x round-trip cost), proportional to
    the addition's own tonnage, not a step function - a small addition gets a small freight charge,
    not exactly zero (that would need a real per-trip ceiling/rounding this doesn't model). Still
    catches a recipe whose ROI looks great only because its freight cost was left out. Returns
    (freight_aic_per_day, note)."""
    origin = shipping.cx_system(cx)
    if origin is None:
        return 0.0, f"no station found for CX {cx!r} - freight not modelled"
    rt = shipping.round_trip_cost(origin, planet_id, sf_price, ff_price)
    if rt is None:
        return 0.0, "no priceable route (gateway without a measured LOCK fee?) - freight not modelled"
    trips_b = max(trip_tonnage(r_b["flows"], mats)) / contracts.SHIP_CARGO
    trips_a = max(trip_tonnage(r_a["flows"], mats)) / contracts.SHIP_CARGO
    extra_trips = trips_a - trips_b
    return extra_trips * rt, f"{extra_trips:+.3f} extra round trips/day x {rt:,.0f} AIC/trip"


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("--remove", action="append", default=[], help="BUILDING[:AMOUNT] - teardown scenario, removes AMOUNT (default: all), 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.remove and not a.infra:
            sys.exit("give --after, or at least one --add/--remove/--infra")
        after_spec = copy.deepcopy(before_spec)
        for add in a.add:
            apply_add(after_spec, add)
        for remove in a.remove:
            apply_remove(after_spec, remove)
        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"]

    mats = {m["Ticker"]: m for m in fio.materials()}
    sf_px = (snap.get(("SF", a.cx)) or market.Quote("SF", a.cx)).ask or 0
    ff_px = (snap.get(("FF", a.cx)) or market.Quote("FF", a.cx)).ask or 0
    freight, freight_note = marginal_freight(r_b, r_a, planet_id, a.cx, sf_px, ff_px, mats)
    d_profit -= freight

    before_counts = Counter(b["building"] for b in before_spec.get("buildings", []) for _ in range(b["amount"]))
    after_counts = Counter(b["building"] for b in after_spec.get("buildings", []) for _ in range(b["amount"]))
    reclaim_total = 0.0
    reclaim_lines = []
    for ticker in before_counts:
        removed_qty = before_counts[ticker] - after_counts.get(ticker, 0)
        if removed_qty > 0:
            r = reclaim_value(ticker, removed_qty, built.get(ticker, 0), a.cx)
            if r:
                reclaim_total += r["value"]
                items = ", ".join(f"{q}{m}" for m, q in r["items"].items())
                reclaim_lines.append(f"  {ticker} x{removed_qty}: {r['value']:,.0f} AIC ({items})")

    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"FREIGHT  {freight:+,.0f} AIC/day marginal ({freight_note})")
    if reclaim_lines:
        print("\nRECLAIMED MATERIALS (demolishing what's removed, sell value, real FIO data):")
        for line in reclaim_lines:
            print(line)
        print(f"  total: {reclaim_total:,.0f} AIC")
    net_capex = d_capex - reclaim_total
    net_label = f"   net cash GAINED {-net_capex:,.0f}" if net_capex < 0 and reclaim_total else (f"   net new capex {net_capex:,.0f}" if reclaim_total else "")
    print(f"\nMARGINAL  profit/day {d_profit:,.0f}   new_capex {d_capex:,.0f}" + net_label, end="")
    if net_capex > 0 and d_profit > 0:
        print(f"   payback {net_capex/d_profit:.1f} d   (simple, non-compounding - a one-time capex paying itself off; not a reinvestable rate, see docs/decisions.md)")
    elif net_capex <= 0 and d_profit > 0:
        print("   (no net new capex after reclaim)" if reclaim_total else "   (no new capex)")
    else:
        print("   (marginal profit <= 0 - not worth it, regardless of reclaim)")

    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()