model.py

from typing import Literal, TypedDict

EnergyClass = Literal["very_low", "low", "medium", "high"]

class EnergyEstimate(TypedDict):
    energy_class: EnergyClass
    estimated_drop_pct: int

def estimate_energy(duration_min: int, distance_m: float, task_type: str) -> EnergyEstimate:
    """Rough energy estimate for a single robot task.

    This is a heuristic demo only.
    """
    base = duration_min * 0.4 + distance_m * 0.02

    task_type = task_type.lower()
    if task_type == "vacuum":
        base *= 1.3
    elif task_type == "mop":
        base *= 1.2
    elif task_type == "patrol":
        base *= 1.1
    elif task_type == "idle":
        base *= 0.2

    drop = max(1, min(int(round(base)), 40))

    if drop <= 5:
        cls: EnergyClass = "very_low"
    elif drop <= 12:
        cls = "low"
    elif drop <= 22:
        cls = "medium"
    else:
        cls = "high"

    return {
        "energy_class": cls,
        "estimated_drop_pct": drop,
    }

if __name__ == "__main__":
    tests = [
        (20, 200.0, "vacuum"),
        (10, 80.0, "mop"),
        (5, 20.0, "idle"),
        (30, 400.0, "patrol"),
    ]
    for d, dist, t in tests:
        print(d, dist, t, "->", estimate_energy(d, dist, t))