PuGa/puga/simulate.py

"""Replica of PRUNplanner's plan simulator (usePlanCalculation.ts): workforce, efficiency, material I/O, cost, profit.
Input is a PRUNplanner plan JSON (as stored by the API, or built by tools/plan_push.build_payload). Pure given the data passed in.
Validated against a screenshot of the real simulator (tests/test_simulate.py)."""
from . import econ

TIERS = econ.TIERS
FARMS = {"FRM", "ORC"}
DEGRADATION_DAYS = 180          # usePlanCalculation.ts: degradation/day = construction cost / 180


def simulate(plan: dict, recipes: list[dict], buildings: list[dict], resources: list[dict], fertility: float,
             price, faction: str | None = None, permits: tuple[float, float] = (1, 2), mcg_per_area: float = 4.0,
             built: dict[str, int] | None = None) -> dict:
    """price(ticker, side) -> AIC per unit or None; side is "buy" (net consumed) or "sell" (net produced), like PRUNplanner BUY/SELL exchange preferences; construction uses "buy". resources: planet resources [{material_ticker, resource_type, daily_extraction}]."""
    d = plan["plan_data"]
    bl = {b["building_ticker"]: b for b in buildings}
    rec = {r["recipe_id"]: r for r in recipes}
    cogc = plan.get("plan_cogc")
    cogc = None if cogc in (None, "---") else cogc
    experts = {e["type"].upper(): e["amount"] for e in d["experts"]}
    lux = {w["type"]: (w["lux1"], w["lux2"]) for w in d["workforce"]}

    req = {t: sum(bl[b["name"]][t + "s"] * b["amount"] for b in d["buildings"]) for t in TIERS}
    cap = dict.fromkeys(TIERS, 0)
    for i in d["infrastructure"]:
        hab = bl[i["building"]].get("habitations")
        if hab:
            for t in TIERS:
                cap[t] += hab[t + "s"] * i["amount"]
    tier_eff = {t: econ.tier_efficiency(cap[t], req[t], *lux[t]) for t in TIERS}
    workforce = {t: dict(need=req[t], supply=cap[t], open=cap[t] - req[t], eff=tier_eff[t]) for t in TIERS}

    flows: dict[str, dict[str, float]] = {}

    def add(side, tk, amt):
        flows.setdefault(tk, {"in": 0.0, "out": 0.0})[side] += amt

    lines = []
    for b in d["buildings"]:
        info = bl[b["name"]]
        heads = {t + "s": info[t + "s"] for t in TIERS}
        eff, elements = econ.building_efficiency(heads, tier_eff, expertise=info["expertise"], cogc=cogc, hq=plan.get("plan_corphq", False),
                                                 experts=experts, faction=faction, permits_used=permits[0], permits_total=permits[1],
                                                 fertility=fertility, is_farm=b["name"] in FARMS)
        rs = []
        for ar in b["active_recipes"]:
            if ar["amount"] == 0:            # switched off in the UI (quantity 0): no production, workforce still required
                continue
            rid = ar["recipeid"]
            if "#" in rid and "=>" not in rid:                       # extraction, e.g. EXT#ALO: daily rate = factor*70 (60 gas)
                tk = rid.split("#")[1]
                res = next(r for r in resources if r["material_ticker"] == tk)
                rs.append(dict(time_ms=econ.TOTAL_MS_DAY, inputs={}, outputs={tk: res["daily_extraction"]}, amount=ar["amount"]))
            else:
                r = rec[rid]
                rs.append(dict(time_ms=r["time_ms"], inputs={i["material_ticker"]: i["material_amount"] for i in r["inputs"]},
                               outputs={o["material_ticker"]: o["material_amount"] for o in r["outputs"]}, amount=ar["amount"]))
        if rs and eff > 0:
            io = econ.production_io(rs, eff, b["amount"])
            for tk, v in io["in"].items():
                add("in", tk, v)
            for tk, v in io["out"].items():
                add("out", tk, v)
        lines.append(dict(building=b["name"], amount=b["amount"], efficiency=eff, elements=elements, recipes=[a["recipeid"] for a in b["active_recipes"]]))

    for t in TIERS:
        for tk, v in econ.workforce_consumption(t, req[t], cap[t], *lux[t]).items():
            add("in", tk, v)

    def unit_cost(tk):
        info = bl[tk]
        return sum(c["material_amount"] * (price(c["material_ticker"], "buy") or 0) for c in info["costs"]) + mcg_per_area * info["area_cost"] * (price("MCG", "buy") or 0)

    # Verified vs simulator screenshots: degradation = production buildings' construction cost / 180 (infrastructure excluded);
    # plan cost and area also include one core module (CM), added automatically.
    prod_cost = sum(unit_cost(b["name"]) * b["amount"] for b in d["buildings"])
    infra_cost = sum(unit_cost(i["building"]) * i["amount"] for i in d["infrastructure"])
    cm_cost = unit_cost("CM")
    plan_cost = prod_cost + infra_cost + cm_cost
    area = (sum(bl[b["name"]]["area_cost"] * b["amount"] for b in d["buildings"])
            + sum(bl[i["building"]]["area_cost"] * i["amount"] for i in d["infrastructure"]) + bl["CM"]["area_cost"])
    degradation = prod_cost / DEGRADATION_DAYS
    rows = {}
    missing = []
    for tk, f in flows.items():
        delta = f["out"] - f["in"]
        p = price(tk, "sell" if delta > 0 else "buy")
        if p is None:
            missing.append(tk)
            p = 0.0
        rows[tk] = dict(inp=f["in"], out=f["out"], delta=delta, value=delta * p, price=p)
    gross = sum(r["value"] for r in rows.values())
    profit = gross - degradation
    # construction still to do = planned minus already built (incl. CM), like PRUNplanner's Construction Cart
    built = built or {}
    planned = {}
    for tk, n in [(b["name"], b["amount"]) for b in d["buildings"]] + [(i["building"], i["amount"]) for i in d["infrastructure"]]:
        planned[tk] = planned.get(tk, 0) + n
    new_capex = sum(unit_cost(tk) * max(n - built.get(tk, 0), 0) for tk, n in planned.items()) + (0 if built.get("CM") else cm_cost)
    return dict(cm_cost=cm_cost, new_capex=new_capex, planned=planned, area=area, prod_cost=prod_cost, infra_cost=infra_cost, workforce=workforce, buildings=lines, flows=rows, plan_cost=plan_cost, degradation=degradation, gross=gross,
                profit=profit, roi_days=(plan_cost / profit if profit > 0 else None), missing_prices=missing)