"""Budget prefill formulas (design doc: Budgets). Both are editable UI fields;
these produce the defaults and the ceiling safety bound."""
from __future__ import annotations

import math

ACTIONS_PER_ENTRY = 3        # open, check documents, disposition
HEADROOM = 1.5
STEP_BUDGET_FLOOR = 60
# Measured from run-ee6dd2f7 (JE-01 pop 400, deepseek-v4-flash, 2026-08-06) via
# scripts/measure_tokens.py; tokens/step grows with context length, so this is
# a p90 across full episodes.
EST_TOKENS_PER_STEP = 85000
CEILING_SAFETY = 10


def default_step_budget(population: int) -> int:
    return max(STEP_BUDGET_FLOOR,
               math.ceil(ACTIONS_PER_ENTRY * HEADROOM * population))


def default_token_ceiling(step_budget: int) -> int:
    return CEILING_SAFETY * step_budget * EST_TOKENS_PER_STEP


def spend_usd(tokens_in: int, tokens_out: int,
              prompt_price: float, completion_price: float) -> float:
    return tokens_in * prompt_price + tokens_out * completion_price