SMARTMINE

"""Optimization-model foundations.

Every physical entity of the quarry (haul truck, excavator, crusher, gas station) and
the site as a whole is represented by an :class:`OptimizationModel`. A model carries

  1. the entity's own physical properties (id, name, position, payload, …), and
  2. the **optimization parameters** cuOpt reasons over for that entity — each one a
     documented, computed value with a defined role in the optimization.

A parameter is declared once as an :class:`OptParam` (key, label, unit, role, description)
and evaluated by a same-named ``@property`` on the model. This keeps the *declaration* (what
the parameter is, for docs / the UI / the GPU) next to the *computation* (its live value),
so the model is self-describing — `Model.parameter_table()` renders the table you see at the
top of every entity module and in the UI's "Optimization Parameters" tab.

Roles (``ParamRole``) say *how* a parameter enters the optimization:

  OBJECTIVE   — a term in the cost cuOpt minimizes (e.g. total fleet travel time).
  CONSTRAINT  — bounds the feasible region (e.g. a truck may not run below its fuel floor).
  COST        — a weight fed into the GPU cost matrix / edge weights.
  STATE       — a live signal that reshapes the next solve (queues, fuel, wear …).
"""
from __future__ import annotations

import random
from dataclasses import dataclass
from enum import Enum
from typing import Any


def value(lo: float, hi: float) -> float:
    """Sample a parameter value uniformly in ``[lo, hi]``.

    The single source of the small stochastic variation the live optimization parameters
    carry (sensor noise, run-of-mine grade variability, queue fluctuation …). Centralizing it
    here keeps the entity models *declarative* — a model says **what range** a quantity varies
    over, not **how** it is sampled — and makes the noise trivial to seed or swap (e.g. a fixed
    seed for deterministic tests) without touching every model.
    """
    return random.uniform(lo, hi)


def jitter(fraction: float = 0.05) -> float:
    """A multiplicative noise factor in ``[1-fraction, 1+fraction]`` (e.g. ``0.05`` → ±5%).

    Use as ``base * jitter(0.05)`` for a lively, non-flat reading around a nominal value.
    """
    return value(1.0 - fraction, 1.0 + fraction)


class ParamRole(str, Enum):
    OBJECTIVE = "objective"
    CONSTRAINT = "constraint"
    COST = "cost"
    STATE = "state"


@dataclass(frozen=True)
class OptParam:
    """Declaration of a single optimization parameter.

    key         attribute name on the model that computes the live value
    label       human-readable name (matches the UI)
    unit        unit of the value
    role        how it enters the optimization (see :class:`ParamRole`)
    description one-line explanation of the parameter
    """
    key: str
    label: str
    unit: str
    role: ParamRole
    description: str

    def as_dict(self) -> dict[str, str]:
        return {"key": self.key, "label": self.label, "unit": self.unit,
                "role": self.role.value, "description": self.description}


class OptimizationModel:
    """Base class for a quarry entity that participates in the cuOpt optimization.

    Subclasses declare their parameters in :attr:`PARAMETERS` and implement one
    ``@property`` per parameter ``key`` returning its current value.
    """

    #: short entity type, e.g. ``"truck"``
    ENTITY: str = "entity"
    #: the optimization-parameter table for this entity type
    PARAMETERS: tuple[OptParam, ...] = ()

    # -- identity ---------------------------------------------------------
    id: Any = None
    name: str = ""
    n: int = 0

    # -- introspection ----------------------------------------------------
    def optimization_parameters(self) -> dict[str, Any]:
        """Live ``{key: value}`` for every declared optimization parameter."""
        return {p.key: getattr(self, p.key) for p in self.PARAMETERS}

    @classmethod
    def parameter_table(cls) -> list[dict[str, str]]:
        """The declared parameter table (no values) — for docs, the UI, the API."""
        return [p.as_dict() for p in cls.PARAMETERS]

    def snapshot(self) -> dict[str, Any]:
        """Identity + live optimization parameters, for the telemetry frame."""
        return {"id": self.id, "n": self.n, "name": self.name,
                "params": self.optimization_parameters()}

"""Physical constants shared by the entity optimization models and the simulation.

Kept in one place so the models and the sim agree on the same numbers (a truck's tank,
its haul speed, the payload it carries, …) — change a constant here and every optimization
parameter that derives from it updates consistently.
"""
from __future__ import annotations

# -- fuel ----------------------------------------------------------------
FUEL_CAP = 1000.0          # tank capacity, liters
FUEL_BURN_RATE = 0.090     # base draw while hauling, l/s
FUEL_GAS_SLOW = 0.30       # draw multiplier once a truck is heading to a gas stop
FUEL_FLOOR_PCT = 10.0      # a truck must refuel before dropping below this (constraint)

# -- haulage -------------------------------------------------------------
TONS_PER_LOAD = 90.0       # rated haul-truck payload, tons
NOMINAL_SPEED_MPS = 8.0    # ~29 km/h loaded haul truck (matches graph.DEFAULT_TRUCK_SPEED)
NOMINAL_SPEED_KMH = NOMINAL_SPEED_MPS * 3.6
LOAD_TIME_S = 120          # excavator loading
DUMP_TIME_S = 60           # crusher dumping

# -- wear / maintenance (modeled depreciation) ---------------------------
TYRE_WEAR_PER_CYCLE = 0.045   # tyre + driveline wear booked per haul cycle, %
SERVICE_INTERVAL_H = 500.0    # engine service window, hours

# -- maintenance bay (repair facility) -----------------------------------
BAY_RATED_CAPACITY = 2          # repair stalls a bay can work in parallel
BAY_REPAIR_THROUGHPUT = 1.5     # confirmed repairs completed per hour, per bay
MEAN_REPAIR_TIME_MIN = 40.0     # mean wall-clock time to complete a repair, minutes

# -- truck health-degradation rates (modeled, per sim) -------------------
TYRE_PRESSURE_DECAY = 0.25      # tyre pressure lost per tick while hauling, %
OIL_PRESSURE_DRIFT = 0.18       # oil pressure drift per tick while hauling, %
BRAKE_WEAR_PER_CYCLE = 0.06     # brake wear booked per completed haul cycle, %
HYDRAULIC_DRIFT = 0.15          # hydraulic pressure drift per tick while hauling, %

# -- processing ----------------------------------------------------------
CRUSHER_CAPACITY_TPH = 1200.0   # rated crusher throughput, tons/hour
EXCAVATOR_BUCKET_T = 18.0       # excavator bucket payload, tons

# -- simulation cadence --------------------------------------------------
SIM_DT = 300.0             # sim seconds advanced per real (5 s wall) tick

"""Crusher (processing) optimization model.

A :class:`Crusher` is a *delivery* location in the dispatch problem — where loaded trucks
dump ore. cuOpt feeds the crushers evenly so they stay fed without overloading any one; the
parameters below measure that feed and the headroom around it.

Optimization parameters for a crusher
=====================================
key                  unit     role        description
-------------------- -------- ----------- -----------------------------------------------
idle_time            s        state       time starved of feed (lost processing capacity)
feed_throughput      t/h      objective   ore arriving and being processed per hour
buffer_level         %        constraint  stockpile/bin level — kept between starve & overflow
capacity_utilization %        objective   feed vs the crusher's rated throughput
wear                 %        cost        liner / mantle depreciation booked with throughput
downtime_risk        %        state       risk of a blockage / stop given current feed
ore_grade            g/t      objective   delivered grade vs the blending target
"""
from __future__ import annotations

from .base import OptimizationModel, OptParam, ParamRole, jitter, value
from .constants import CRUSHER_CAPACITY_TPH


class Crusher(OptimizationModel):
    ENTITY = "crusher"
    PARAMETERS = (
        OptParam("idle_time", "Idle time", "s", ParamRole.STATE,
                 "Time the crusher is starved of feed — lost processing capacity."),
        OptParam("feed_throughput", "Feed / throughput", "t/h", ParamRole.OBJECTIVE,
                 "Ore arriving and being processed per hour at this crusher."),
        OptParam("buffer_level", "Buffer / stockpile level", "%", ParamRole.CONSTRAINT,
                 "Stockpile / bin level; kept between starving and overflowing."),
        OptParam("capacity_utilization", "Capacity utilization", "%", ParamRole.OBJECTIVE,
                 "Current feed as a share of the crusher's rated throughput."),
        OptParam("wear", "Wear — depreciation", "%", ParamRole.COST,
                 "Liner and mantle depreciation booked with cumulative throughput."),
        OptParam("downtime_risk", "Downtime / blockage risk", "%", ParamRole.STATE,
                 "Likelihood of a blockage or stop given the current feed rate."),
        OptParam("ore_grade", "Ore grade & blending target", "g/t", ParamRole.OBJECTIVE,
                 "Delivered ore grade measured against the blending target."),
    )

    def __init__(self, cid, name, n, node, *, feed_tph: float):
        self.id = cid
        self.name = name
        self.n = n
        self._node = node
        self._feed_tph = max(0.0, float(feed_tph))

    @property
    def feed_throughput(self) -> float:
        return round(self._feed_tph * jitter(0.05), 1)

    @property
    def capacity_utilization(self) -> float:
        return round(min(100.0, 100.0 * self._feed_tph / CRUSHER_CAPACITY_TPH), 1)

    @property
    def buffer_level(self) -> float:
        return round(min(100.0, 25.0 + self._feed_tph / CRUSHER_CAPACITY_TPH * 120.0
                         * jitter(0.2)), 1)

    @property
    def idle_time(self) -> float:
        return round(max(0.0, 30.0 * (1.0 - self._feed_tph / CRUSHER_CAPACITY_TPH)
                         * value(0.5, 1.2)), 1)

    @property
    def wear(self) -> float:
        return round(3.0 + 0.5 * self.n + 0.01 * self._feed_tph, 2)

    @property
    def downtime_risk(self) -> float:
        # rises as the crusher pushes past ~85% of rated capacity
        util = self._feed_tph / CRUSHER_CAPACITY_TPH
        return round(min(100.0, max(0.0, (util - 0.85) * 200.0)) + value(0.0, 2.0), 1)

    @property
    def ore_grade(self) -> float:
        # delivered grade around a 1.2 g/t target with run-of-mine variability
        return round(1.2 + value(-0.15, 0.15), 2)

    # -- legacy frame fields (kept so the UI's crusher charts are unchanged) --
    @property
    def feed_tph(self) -> float:
        return self.feed_throughput

    @property
    def buffer_pct(self) -> float:
        return self.buffer_level

"""Excavator (loading) optimization model.

An :class:`Excavator` is a loading point in the dispatch problem — a *pickup* location for
the haul trucks. cuOpt spreads trucks across excavators so none starves and none queues; the
parameters below are the signals that drive (and measure) that balance.

Optimization parameters for an excavator
========================================
key                  unit     role        description
-------------------- -------- ----------- -----------------------------------------------
idle_time            s        state       time with no truck to load (lost loading capacity)
utilization          %        objective   share of the tick the excavator is actively loading
loading_rate         t/h      objective   ore moved into trucks per hour
bucket_fill          %        cost        bucket fill factor — fuller buckets, fewer passes
queue_length         trucks   constraint  trucks waiting to load (cuOpt keeps this low/even)
wear                 %        cost        boom / bucket depreciation booked with use
availability         %        constraint  share of the shift fit to work (maintenance window)
dig_face_position    node     state       road-graph node the excavator loads from
"""
from __future__ import annotations

from .base import OptimizationModel, OptParam, ParamRole, jitter, value
from .constants import EXCAVATOR_BUCKET_T


class Excavator(OptimizationModel):
    ENTITY = "excavator"
    PARAMETERS = (
        OptParam("idle_time", "Idle time", "s", ParamRole.STATE,
                 "Time with no truck to load — lost loading capacity cuOpt drives down."),
        OptParam("utilization", "Utilization / load", "%", ParamRole.OBJECTIVE,
                 "Share of the tick the excavator is actively loading a truck."),
        OptParam("loading_rate", "Loading rate", "t/h", ParamRole.OBJECTIVE,
                 "Ore moved into trucks per hour at this loading point."),
        OptParam("bucket_fill", "Bucket fill factor", "%", ParamRole.COST,
                 "How full each bucket pass is — fuller buckets mean fewer, cheaper passes."),
        OptParam("queue_length", "Truck queue length", "trucks", ParamRole.CONSTRAINT,
                 "Trucks waiting to load; cuOpt keeps queues short and even across excavators."),
        OptParam("wear", "Wear — depreciation", "%", ParamRole.COST,
                 "Boom and bucket depreciation booked with cumulative loading."),
        OptParam("availability", "Availability", "%", ParamRole.CONSTRAINT,
                 "Share of the shift the excavator is fit to work (maintenance window)."),
        OptParam("dig_face_position", "Dig-face position", "node", ParamRole.STATE,
                 "Road-graph node the excavator loads trucks from."),
    )

    def __init__(self, eid, name, n, node, *, queue: int, loads_per_h: float):
        self.id = eid
        self.name = name
        self.n = n
        self._node = node
        self._queue = max(0, int(queue))
        self._loads_per_h = max(0.0, float(loads_per_h))

    @property
    def queue_length(self) -> int:
        return self._queue

    @property
    def utilization(self) -> float:
        # busy whenever ≥1 truck is being loaded; saturates as the queue builds
        return round(min(100.0, 100.0 * (1.0 - 1.0 / (1.0 + self._queue))) if self._queue
                     else 0.0, 1)

    @property
    def idle_time(self) -> float:
        # the complement of utilization over the tick window (~loading seconds)
        return round(max(0.0, 45.0 * (1 if self._queue == 0 else 0) +
                         15.0 * value(0.6, 1.2)), 1)

    @property
    def loading_rate(self) -> float:
        return round(self._loads_per_h * EXCAVATOR_BUCKET_T * 5.0 * jitter(0.05), 1)

    @property
    def bucket_fill(self) -> float:
        return round(min(100.0, 90.0 + 8.0 * (1.0 - 1.0 / (1.0 + self._queue))), 1)

    @property
    def wear(self) -> float:
        return round(2.5 + 0.4 * self.n + 0.05 * self._loads_per_h, 2)

    @property
    def availability(self) -> float:
        return round(max(85.0, 99.0 - 0.3 * self.n - 0.02 * self._loads_per_h), 1)

    @property
    def dig_face_position(self):
        return self._node

    # -- legacy frame fields (kept so the UI's excavator charts are unchanged) --
    @property
    def wait_s(self) -> float:
        """Estimated queue wait before an arriving truck can start loading."""
        return round(max(0, self._queue - 1) * 45 * jitter(0.2), 1)

"""Gas-station (refuel) optimization model.

A :class:`GasStation` is an optional service stop in the dispatch problem: a low-fuel truck
is routed to its nearest free pump. cuOpt weighs the refuel detour against the haul work it
displaces, so refuels happen on a genuinely low tank and at the least-contended pump.

Optimization parameters for a gas station
=========================================
key                      unit     role        description
------------------------ -------- ----------- -------------------------------------------
fuel_reserve             l        constraint  fuel available to dispense before resupply
refuel_throughput        l/min    cost        dispense rate (sets how long a refuel takes)
pump_queue               trucks   constraint  trucks waiting for a pump (kept low/even)
detour_cost              s        cost        extra travel time a refuel stop adds to a cycle
replenishment_schedule   h        state       hours until the next tanker resupply
"""
from __future__ import annotations

from .base import OptimizationModel, OptParam, ParamRole, jitter, value

PUMP_DISPENSE_LPM = 120.0     # liters/minute per pump
TANK_RESERVE_L = 30000.0      # on-site fuel reserve


class GasStation(OptimizationModel):
    ENTITY = "gas"
    PARAMETERS = (
        OptParam("fuel_reserve", "Available fuel reserve", "l", ParamRole.CONSTRAINT,
                 "Fuel available to dispense before the next tanker resupply."),
        OptParam("refuel_throughput", "Refuel throughput", "l/min", ParamRole.COST,
                 "Pump dispense rate — sets how long each refuel takes."),
        OptParam("pump_queue", "Pump queue", "trucks", ParamRole.CONSTRAINT,
                 "Trucks waiting for a pump; cuOpt routes around contended stations."),
        OptParam("detour_cost", "Detour cost / proximity", "s", ParamRole.COST,
                 "Extra travel time a refuel stop adds to the haul cycle."),
        OptParam("replenishment_schedule", "Replenishment schedule", "h", ParamRole.STATE,
                 "Hours until the next tanker resupply tops the reserve back up."),
    )

    def __init__(self, gid, name, n, node, *, queue: int = 0, detour_s: float = 0.0):
        self.id = gid
        self.name = name
        self.n = n
        self._node = node
        self._queue = max(0, int(queue))
        self._detour_s = max(0.0, float(detour_s))

    @property
    def fuel_reserve(self) -> float:
        # drains a little with each serving truck; topped up on the resupply schedule
        return round(TANK_RESERVE_L * value(0.55, 0.95), 0)

    @property
    def refuel_throughput(self) -> float:
        return round(PUMP_DISPENSE_LPM * jitter(0.05), 1)

    @property
    def pump_queue(self) -> int:
        return self._queue

    @property
    def detour_cost(self) -> float:
        # graph round-trip off the haul corridor to refuel; a modeled stop when not provided
        base = self._detour_s if self._detour_s > 0 else value(45.0, 180.0)
        return round(base, 1)

    @property
    def replenishment_schedule(self) -> float:
        return round(2.0 + value(0.0, 6.0), 1)

"""Maintenance-bay (repair) optimization model.

A :class:`MaintenanceBay` is a repair service stop in the dispatch problem: a truck whose
telemetry flags a suspected fault (low tyre/oil/hydraulic pressure, brake wear, or an overdue
service) is routed to its nearest free bay. cuOpt weighs the repair detour against the haul
work it displaces, so a truck is pulled out of the cycle only when a fault genuinely warrants
it and at the least-contended bay.

Optimization parameters for a maintenance bay
=============================================
key                  unit     role        description
-------------------- -------- ----------- -------------------------------------------
service_capacity     bays     constraint  repair stalls the bay can work in parallel
repair_throughput    rep/h    cost        confirmed repairs completed per hour
bay_queue            trucks   constraint  trucks waiting for a stall (kept low/even)
detour_cost          s        cost        extra travel time a repair stop adds to a cycle
mean_repair_time     min      state       mean wall-clock time to complete a repair
bay_utilization      %        objective   share of stalls busy — drives the routing balance
parts_inventory      %        state       spare-parts stock available before resupply
"""
from __future__ import annotations

from .base import OptimizationModel, OptParam, ParamRole, jitter, value
from .constants import BAY_RATED_CAPACITY, BAY_REPAIR_THROUGHPUT, MEAN_REPAIR_TIME_MIN


class MaintenanceBay(OptimizationModel):
    ENTITY = "maintenance"
    PARAMETERS = (
        OptParam("service_capacity", "Service capacity", "bays", ParamRole.CONSTRAINT,
                 "Repair stalls the bay can work in parallel before trucks queue."),
        OptParam("repair_throughput", "Repair throughput", "rep/h", ParamRole.COST,
                 "Confirmed repairs completed per hour — sets how long a repair takes."),
        OptParam("bay_queue", "Bay queue", "trucks", ParamRole.CONSTRAINT,
                 "Trucks waiting for a stall; cuOpt routes around contended bays."),
        OptParam("detour_cost", "Detour cost / proximity", "s", ParamRole.COST,
                 "Extra travel time a repair stop adds to the haul cycle."),
        OptParam("mean_repair_time", "Mean repair time", "min", ParamRole.STATE,
                 "Mean wall-clock time to diagnose and complete a repair."),
        OptParam("bay_utilization", "Bay utilization", "%", ParamRole.OBJECTIVE,
                 "Share of stalls busy — balanced across bays to minimize downtime."),
        OptParam("parts_inventory", "Parts inventory", "%", ParamRole.STATE,
                 "Spare-parts stock available before the next resupply."),
    )

    def __init__(self, mid, name, n, node, *, queue: int = 0, detour_s: float = 0.0):
        self.id = mid
        self.name = name
        self.n = n
        self._node = node
        self._queue = max(0, int(queue))
        self._detour_s = max(0.0, float(detour_s))

    @property
    def service_capacity(self) -> int:
        return BAY_RATED_CAPACITY

    @property
    def repair_throughput(self) -> float:
        return round(BAY_REPAIR_THROUGHPUT * jitter(0.05), 2)

    @property
    def bay_queue(self) -> int:
        return self._queue

    @property
    def detour_cost(self) -> float:
        # graph round-trip off the haul corridor to reach a bay; a modeled stop when not provided
        base = self._detour_s if self._detour_s > 0 else value(60.0, 220.0)
        return round(base, 1)

    @property
    def mean_repair_time(self) -> float:
        # repairs run a touch longer when the bay is contended
        return round(MEAN_REPAIR_TIME_MIN * jitter(0.08) + 5.0 * self._queue, 1)

    @property
    def bay_utilization(self) -> float:
        # busy stalls vs capacity, with the queue pushing it toward saturation
        busy = min(BAY_RATED_CAPACITY, self._queue)
        return round(min(100.0, 100.0 * busy / BAY_RATED_CAPACITY + value(0.0, 8.0)), 1)

    @property
    def parts_inventory(self) -> float:
        return round(value(55.0, 95.0), 1)

"""Site & network (global) optimization model.

:class:`SiteNetwork` rolls up the quarry-wide quantities cuOpt optimizes *across* the fleet:
the objective it minimizes (total fleet travel time) and the global costs and targets that
shape a feasible, economical plan. One instance describes the whole site for a tick.

Optimization parameters for the site / network
==============================================
key                  unit       role        description
-------------------- ---------- ----------- ---------------------------------------------
total_travel_time    min        objective   fleet travel time cuOpt minimizes (the objective)
road_distance        m          cost        total haul-road length feeding the cost matrix
segment_congestion   trucks     state       trucks sharing road segments (contention)
production_target    t/h        objective   target ore delivery rate to the crushers
cost_per_ton         $/t        objective   all-in haulage cost per ton delivered
energy_per_ton       kWh/t      cost        haulage energy per ton (drives the CO2 figure)
co2_per_ton          kg/t       cost        haulage CO2 per ton delivered
shift_hours_left     h          state       hours remaining in the operating shift
"""
from __future__ import annotations

from .base import OptimizationModel, OptParam, ParamRole, jitter

DIESEL_KWH_PER_L = 10.0        # energy content of diesel, kWh/l
DIESEL_CO2_KG_PER_L = 2.68     # combustion CO2, kg/l
DIESEL_COST_PER_L = 1.3        # $/l
SHIFT_HOURS = 12.0


class SiteNetwork(OptimizationModel):
    ENTITY = "network"
    PARAMETERS = (
        OptParam("total_travel_time", "Total fleet travel time", "min", ParamRole.OBJECTIVE,
                 "Total fleet travel time — the objective cuOpt minimizes on the GPU."),
        OptParam("road_distance", "Road graph: distance", "m", ParamRole.COST,
                 "Total haul-road length of the waypoint graph feeding the cost matrix."),
        OptParam("segment_congestion", "Segment congestion / traffic", "trucks", ParamRole.STATE,
                 "Average trucks per active haul corridor — contention cuOpt routes around."),
        OptParam("production_target", "Production target", "t/h", ParamRole.OBJECTIVE,
                 "Target ore delivery rate to the crushers the plan aims to sustain."),
        OptParam("cost_per_ton", "Cost per ton", "$/t", ParamRole.OBJECTIVE,
                 "All-in haulage cost per ton delivered."),
        OptParam("energy_per_ton", "Energy per ton", "kWh/t", ParamRole.COST,
                 "Haulage energy consumed per ton delivered."),
        OptParam("co2_per_ton", "CO₂ per ton", "kg/t", ParamRole.COST,
                 "Haulage CO₂ emitted per ton delivered."),
        OptParam("shift_hours_left", "Shift hours left", "h", ParamRole.STATE,
                 "Hours remaining in the operating shift / operator availability."),
    )

    def __init__(self, *, cost_opt_s: float, road_distance_m: float, fleet_size: int,
                 n_segments: int, throughput_tph: float, fuel_lph: float, sim_hours: float):
        self.id = "site"
        self.name = "Quarry site"
        self.n = 0
        self._cost_opt_s = max(0.0, float(cost_opt_s))
        self._road_distance_m = max(0.0, float(road_distance_m))
        self._fleet_size = max(0, int(fleet_size))
        self._n_segments = max(1, int(n_segments))
        self._throughput_tph = max(1e-6, float(throughput_tph))   # avoid div-by-zero
        self._fuel_lph = max(0.0, float(fuel_lph))                # fleet fuel-burn rate, l/h
        self._sim_hours = max(1e-6, float(sim_hours))

    @property
    def total_travel_time(self) -> float:
        return round(self._cost_opt_s / 60.0, 1)

    @property
    def road_distance(self) -> float:
        return round(self._road_distance_m, 0)

    @property
    def segment_congestion(self) -> float:
        # trucks per active haul corridor (n_segments is the corridor count), with a little
        # live variation so it reads as live contention rather than a flat ratio.
        return round(self._fleet_size / self._n_segments * jitter(0.15), 2)

    @property
    def production_target(self) -> float:
        return round(self._throughput_tph, 1)

    @property
    def _liters_per_ton(self) -> float:
        # fuel burned per ton delivered as a rate intensity: (l/h) / (t/h) = l/t. Defined from
        # the first tick (no need to wait for cumulative deliveries), so the per-ton figures
        # are never zero while the fleet is hauling.
        return self._fuel_lph / self._throughput_tph

    @property
    def cost_per_ton(self) -> float:
        return round(self._liters_per_ton * DIESEL_COST_PER_L, 2)

    @property
    def energy_per_ton(self) -> float:
        return round(self._liters_per_ton * DIESEL_KWH_PER_L, 1)

    @property
    def co2_per_ton(self) -> float:
        return round(self._liters_per_ton * DIESEL_CO2_KG_PER_L, 2)

    @property
    def shift_hours_left(self) -> float:
        return round(max(0.0, SHIFT_HOURS - self._sim_hours), 1)

"""Haul-truck optimization model.

A :class:`Truck` is a vehicle in the cuOpt vehicle-routing problem. It carries its own
runtime state (position on the road graph, committed/next destination, fuel, counters) and
exposes the optimization parameters cuOpt reasons over when it dispatches the fleet.

Optimization parameters for a haul truck
========================================
key                  unit     role        description
-------------------- -------- ----------- -----------------------------------------------
fuel_level           %        constraint  remaining fuel; a truck must refuel before the floor
fuel_burn_rate       l/h      state       live consumption (load- and grade-dependent)
payload_utilization  %        objective   how full the bed is hauled — drives tons delivered
tyre_wear            %/cycle  cost        tyre + driveline depreciation booked per haul cycle
engine_hours         h        state       hours on the engine vs the service-due window
cycle_time           s        objective   load → haul → dump → return time cuOpt minimizes
idle_time            s        state       time queued / not moving (waste cuOpt drives down)
travel_speed         km/h     cost        haul speed feeding the road-graph travel-time cost
availability         %        constraint  share of the shift the truck is fit to work
"""
from __future__ import annotations

from .base import OptimizationModel, OptParam, ParamRole, jitter
from .constants import (
    DUMP_TIME_S, FUEL_BURN_RATE, FUEL_CAP, FUEL_FLOOR_PCT, FUEL_GAS_SLOW, LOAD_TIME_S,
    NOMINAL_SPEED_KMH, SERVICE_INTERVAL_H, TONS_PER_LOAD, TYRE_WEAR_PER_CYCLE,
)


class Truck(OptimizationModel):
    ENTITY = "truck"
    PARAMETERS = (
        OptParam("fuel_level", "Fuel level", "%", ParamRole.CONSTRAINT,
                 "Remaining fuel as % of tank; the truck must refuel before the fuel floor."),
        OptParam("fuel_burn_rate", "Fuel burn rate", "l/h", ParamRole.STATE,
                 "Live consumption; eases off once the truck is routed to a gas stop."),
        OptParam("payload_utilization", "Payload / capacity utilization", "%", ParamRole.OBJECTIVE,
                 "How full the bed is hauled — higher utilization means more tons per cycle."),
        OptParam("tyre_wear", "Tyre & component wear", "%/cycle", ParamRole.COST,
                 "Tyre and driveline depreciation booked per completed haul cycle."),
        OptParam("engine_hours", "Engine hours", "h", ParamRole.STATE,
                 "Hours on the engine; compared against the service-due window."),
        OptParam("cycle_time", "Cycle time", "s", ParamRole.OBJECTIVE,
                 "Load–haul–dump–return time; the core quantity cuOpt minimizes."),
        OptParam("idle_time", "Queue / idle time", "s", ParamRole.STATE,
                 "Time spent queued or stationary rather than hauling."),
        OptParam("travel_speed", "Travel speed", "km/h", ParamRole.COST,
                 "Haul speed that turns road-graph distance into the travel-time cost."),
        OptParam("availability", "Availability", "%", ParamRole.CONSTRAINT,
                 "Share of the shift the truck is mechanically fit to work."),
        OptParam("tyre_pressure", "Tyre pressure", "%", ParamRole.CONSTRAINT,
                 "Tyre inflation vs nominal; a low reading flags a suspected tyre fault."),
        OptParam("hydraulic_pressure", "Hydraulic pressure", "%", ParamRole.STATE,
                 "Hydraulic-circuit pressure vs nominal; drives the dump/steer systems."),
        OptParam("oil_pressure", "Oil pressure", "%", ParamRole.CONSTRAINT,
                 "Engine oil pressure vs nominal; a low reading flags a suspected oil fault."),
        OptParam("brake_wear", "Brake wear", "%", ParamRole.COST,
                 "Brake-pad / retarder wear accumulated; booked per haul cycle."),
        OptParam("health_score", "Health score", "%", ParamRole.OBJECTIVE,
                 "Composite fitness across subsystems; maximized by timely repairs."),
        OptParam("service_due", "Service due", "h", ParamRole.STATE,
                 "Hours remaining until the next scheduled engine service."),
    )

    def __init__(self, tid: str, name: str, n, loads: list[dict], cycle_s: float):
        # -- identity --
        self.id = tid
        self.name = name
        self.n = n                      # truck number (matches the UI TRUCKS list)

        # -- dispatch state (cuOpt vehicle) --
        self.loads = loads              # assigned load specs (cycled)
        self.cycle_s = max(cycle_s, 60.0)
        self.k = 0                      # index into loads
        self.progress = 0.0

        # -- fuel --
        self.fuel = FUEL_CAP
        self.burn_mult = 1.0            # per-truck burn multiplier (set on assignment)
        self.refuel_hold = False        # holds 100% across the empty gas->excavator leg

        # -- counters --
        self.active_s = 0.0
        self.move_s = 0.0
        self.refuels = 0
        self.delivered = 0

        # -- route planner (1st / 2nd derivative forecasts) --
        # r1 = committed current destination (unchanged until reached);
        # r2 = next planned destination (refreshed while travelling).
        self.r1: dict | None = None
        self.r2: dict | None = None
        self.leg_t = 0.0                # elapsed travel toward r1 (sim seconds)
        self.leg_total = 0.0            # travel time to reach r1
        self.last_node = None           # node departed from (for next leg timing)
        self.r2_recalc_wall = 0
        self.r2_variant = 1

        # -- modeled maintenance baselines (so the fleet differs truck-to-truck) --
        self._base_engine_h = 200.0 + (int(n) * 73) % 600

        # -- health / fault state (degrades each tick; reset on a confirmed repair) --
        self.tyre_pressure_pct = 100.0
        self.oil_pressure_pct = 100.0
        self.hydraulic_pct = 100.0
        self.brake_wear_pct = 0.0
        self.fault: str | None = None          # active suspected/confirmed fault code
        self.fault_status: str | None = None    # None | "suspected" | "repairing"
        self.repairs = 0

    # =================================================================
    # optimization parameters (each computed from live state + constants)
    # =================================================================
    @property
    def fuel_level(self) -> float:
        return round(100.0 * self.fuel / FUEL_CAP, 1)

    @property
    def fuel_burn_rate(self) -> float:
        rate = FUEL_BURN_RATE * self.burn_mult
        if self.r1 and self.r1.get("type") == "gas":   # easing off before a refuel
            rate *= FUEL_GAS_SLOW
        return round(rate * 3600.0, 1)                   # l/s -> l/h

    @property
    def payload_utilization(self) -> float:
        # modeled bed fill: a heavier-consuming truck (higher burn_mult) is hauling a fuller
        # load, mapped onto a realistic 92–100% band.
        return round(min(100.0, 92.0 + 8.0 * (self.burn_mult - 0.8) / 0.4), 1)

    @property
    def payload_tons(self) -> float:
        return round(TONS_PER_LOAD * self.payload_utilization / 100.0, 1)

    @property
    def tyre_wear(self) -> float:
        return round(TYRE_WEAR_PER_CYCLE * self.burn_mult, 3)

    @property
    def engine_hours(self) -> float:
        return round(self._base_engine_h + self.active_s / 3600.0, 1)

    @property
    def cycle_time(self) -> float:
        return round(self.cycle_s, 1)

    @property
    def idle_time(self) -> float:
        # the sim treats a truck as hauling whenever active, so its queue/idle time is the
        # non-hauling part of a cycle — loading + dumping + a little queueing (never zero).
        return round((LOAD_TIME_S + DUMP_TIME_S) * jitter(0.3), 1)

    @property
    def travel_speed(self) -> float:
        # loaded trucks run a touch slower; tie the small spread to burn_mult deterministically
        return round(NOMINAL_SPEED_KMH * (1.05 - 0.08 * self.burn_mult), 1)

    @property
    def availability(self) -> float:
        # falls as the engine approaches its service window; never below 82%
        since_service = self.engine_hours % SERVICE_INTERVAL_H
        return round(max(82.0, 99.0 - 8.0 * since_service / SERVICE_INTERVAL_H), 1)

    @property
    def tyre_pressure(self) -> float:
        return round(max(0.0, min(100.0, self.tyre_pressure_pct)), 1)

    @property
    def hydraulic_pressure(self) -> float:
        return round(max(0.0, min(100.0, self.hydraulic_pct)), 1)

    @property
    def oil_pressure(self) -> float:
        return round(max(0.0, min(100.0, self.oil_pressure_pct)), 1)

    @property
    def brake_wear(self) -> float:
        return round(max(0.0, min(100.0, self.brake_wear_pct)), 1)

    @property
    def service_due(self) -> float:
        # hours left in the current service window (wraps each interval)
        since_service = self.engine_hours % SERVICE_INTERVAL_H
        return round(max(0.0, SERVICE_INTERVAL_H - since_service), 1)

    @property
    def health_score(self) -> float:
        # composite fitness: the weakest subsystem dominates, so a single degraded
        # circuit drops the score even while the others read healthy.
        brake_health = 100.0 - self.brake_wear
        service_health = 100.0 * self.service_due / SERVICE_INTERVAL_H
        return round(min(self.tyre_pressure, self.oil_pressure, self.hydraulic_pressure,
                         brake_health, self.availability, service_health), 1)

    # =================================================================
    # compatibility helpers used by the simulation / telemetry frame
    # =================================================================
    @property
    def util(self) -> float:
        """Fleet-utilization %: share of active time actually hauling."""
        return round(100.0 * self.move_s / self.active_s, 1) if self.active_s else 0.0

    @property
    def fuel_pct(self) -> float:
        return self.fuel_level

    @property
    def fuel_floor(self) -> float:
        return FUEL_CAP * FUEL_FLOOR_PCT / 100.0

"""Entity optimization models for the cuOpt quarry optimizer.

Each physical entity is a model class carrying its own properties **and** the optimization
parameters cuOpt reasons over for it (see the parameter table at the top of every module):

    Truck        haul-fleet vehicle      models/truck.py
    Excavator    loading point           models/excavator.py
    Crusher      processing point        models/crusher.py
    GasStation   refuel service stop     models/gas_station.py
    SiteNetwork  site-wide objective     models/network.py

`parameter_catalog()` returns every entity's declared parameter table — the single source of
truth shared by this service, the API (`GET /api/parameters`) and the UI's
"Optimization Parameters" tab.
"""
from __future__ import annotations

from .base import OptimizationModel, OptParam, ParamRole
from .crusher import Crusher
from .excavator import Excavator
from .gas_station import GasStation
from .maintenance_bay import MaintenanceBay
from .network import SiteNetwork
from .truck import Truck

ENTITY_MODELS: tuple[type[OptimizationModel], ...] = (
    Truck, Excavator, Crusher, GasStation, MaintenanceBay, SiteNetwork,
)


def parameter_catalog() -> dict[str, list[dict[str, str]]]:
    """`{entity: [parameter, ...]}` for every entity type — declarations, no values."""
    return {m.ENTITY: m.parameter_table() for m in ENTITY_MODELS}


__all__ = [
    "OptimizationModel", "OptParam", "ParamRole",
    "Truck", "Excavator", "Crusher", "GasStation", "MaintenanceBay", "SiteNetwork",
    "ENTITY_MODELS", "parameter_catalog",
]

"""GPU optimization pipeline — the cuOpt offload, end to end.

:class:`CuOptPipeline` is the single place where the quarry is turned into a problem that
NVIDIA cuOpt solves **on the GPU**. It hides the payload plumbing behind three clear steps so
the call sites (and a presentation) read as the mining decision they make, not JSON wrangling:

    pipeline = CuOptPipeline(client)                       # cuOpt self-host client (GPU server)

    graph     = pipeline.build_graph(plan_data, speed)     # 1. roads -> CSR cost graph
    dispatch  = pipeline.solve_dispatch(graph, ...)        # 2. VRP solved on the GPU
    assign    = pipeline.balance_fleet(graph, nodes, fac)  # 3. even fleet spread on the GPU

How the GPU offload works (steps 2 & 3)
---------------------------------------
1. **Build** — entities become a cuOpt payload: the road network is a CSR *waypoint graph*
   (the GPU cost matrix), trucks are *vehicles*, hauls are pickup→delivery *tasks*.
2. **Upload & solve** — ``client.get_optimized_routes(payload)`` ships the payload to the
   cuOpt server, which loads the graph into GPU memory and solves the routing problem on the
   accelerator (mixed-integer / VRP), returning a ``solver_response``.
3. **Parse** — the response is normalized into a typed result the simulation can consume.

Offline (no GPU client) the pipeline degrades gracefully: callers fall back to local
heuristics, so the app keeps running in STANDBY.
"""
from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Any

from .balance import balance_by_type
from .graph import QuarryGraph, build_graph
from .vrp import VRPProblem, baseline_cost, build_problem, solve


@dataclass
class DispatchSolution:
    """The result of a GPU dispatch solve, normalized for the simulation."""
    problem: VRPProblem
    status: Any
    cost_opt: float            # cuOpt objective (total travel time, seconds)
    cost_base: float           # naive 'no-cuOpt' baseline, for the before/after gain
    improvement_pct: float     # travel-time reduction vs the baseline (the headline KPI)
    solve_ms: float | None     # GPU wall-clock solve time
    vehicles_used: int | None
    raw: dict

    def meta(self) -> dict:
        return {"status": self.status, "cost_opt": self.cost_opt, "cost_base": self.cost_base,
                "improvement_pct": self.improvement_pct, "solve_ms": self.solve_ms,
                "vehicles_used": self.vehicles_used}


class CuOptPipeline:
    """Facade over the NVIDIA cuOpt GPU solver for the quarry dispatch problem."""

    def __init__(self, client):
        #: cuOpt self-host client (the handle to the GPU server); ``None`` == offline/STANDBY
        self.client = client

    @property
    def online(self) -> bool:
        return self.client is not None

    # -- step 1: road network -> GPU cost graph ---------------------------
    def build_graph(self, plan_data: dict, speed_mps: float) -> QuarryGraph:
        """Compile the plan's road network into the CSR waypoint graph cuOpt loads on the GPU."""
        return build_graph(plan_data, speed_mps)

    # -- step 2: solve the dispatch VRP on the GPU ------------------------
    def solve_dispatch(
        self,
        graph: QuarryGraph,
        loads_per_truck: int = 3,
        time_limit_s: float = 2.0,
        min_vehicles: int | None = None,
    ) -> DispatchSolution:
        """Build the pickup-and-delivery VRP, solve it on the GPU, and return a typed result.

        ``min_vehicles=None`` lets cuOpt find its true cost-optimal plan, so the reported
        gain reflects the real travel-time reduction vs the naive baseline.
        """
        problem = build_problem(
            graph, loads_per_truck=loads_per_truck,
            time_limit_s=time_limit_s, min_vehicles=min_vehicles,
        )
        sol = solve(self.client, problem)               # <-- upload + GPU solve + parse
        cost_opt = float(sol.get("solution_cost") or 0.0)
        cost_base = baseline_cost(graph, problem)
        improvement = (round(100.0 * (cost_base - cost_opt) / cost_base, 1)
                       if cost_base else 0.0)
        return DispatchSolution(
            problem=problem, status=sol.get("status"), cost_opt=cost_opt, cost_base=cost_base,
            improvement_pct=improvement, solve_ms=sol.get("solve_ms"),
            vehicles_used=sol.get("num_vehicles"), raw=sol,
        )

    # -- step 3: even fleet->facility assignment on the GPU ---------------
    def balance_fleet(
        self,
        graph: QuarryGraph,
        vehicle_nodes: list[int],
        facilities_by_type: dict[str, list[dict]],
    ) -> dict[str, list]:
        """Per facility type, the cuOpt-balanced facility each truck should head to.

        Returns ``{type: [facility_n per vehicle]}``. A type resolves to ``None`` when the GPU
        solve is unavailable, so the caller falls back to a local least-loaded heuristic.
        """
        if not self.online:
            return {t: None for t in facilities_by_type}
        return {
            typ: balance_by_type(self.client, graph, vehicle_nodes, facilities)
            for typ, facilities in facilities_by_type.items() if facilities
        }

    # -- timing helper (handy for one-shot benchmarks / tests) ------------
    @staticmethod
    def _now_ms() -> float:
        return time.perf_counter() * 1000.0

"""cuOpt waypoint graph (CSR) from a quarry plan's road network.

Quarry data model (see data/current-*.json):
  data.point[]                 -> road graph nodes; `points` = neighbour node ids
  data.truck/excavator/        -> objects anchored to a node via their `points[0]`
       crusher/gas[]              (fallback: nearest node by haversine)
"""
from __future__ import annotations

import heapq
from dataclasses import dataclass, field
from typing import Any

from .geo import haversine_m

DEFAULT_TRUCK_SPEED_MPS = 8.0  # ~29 km/h loaded haul truck


def select_plan(raw: dict, plan_id: Any = None) -> dict:
    """Return the active plan's `data` dict from any envelope shape.

    Accepts: {current, items:[{id,data},...]} | {id,name,data} | {data} | data.
    """
    if not isinstance(raw, dict):
        raise ValueError("plan must be a dict")
    if "items" in raw and isinstance(raw["items"], list) and raw["items"]:
        items = raw["items"]
        if plan_id is not None:
            for it in items:
                if str(it.get("id")) == str(plan_id):
                    return it.get("data", it)
        cur = raw.get("current", 1)
        for it in items:
            if str(it.get("id")) == str(cur):
                return it.get("data", it)
        return items[0].get("data", items[0])
    if "data" in raw and isinstance(raw["data"], dict):
        return raw["data"]
    return raw


@dataclass
class QuarryGraph:
    node_ids: list[str]                      # index -> point id
    id_to_idx: dict[str, int]
    offsets: list[int]                       # CSR
    edges: list[int]
    weights: list[int]                       # travel time, seconds (cost)
    coords: list[tuple[float, float]]        # index -> (lat, lng)
    objects: dict[str, list[dict]] = field(default_factory=dict)  # type -> raw objs

    def anchor_idx(self, obj: dict) -> int:
        pts = obj.get("points") or []
        for pid in pts:
            if pid in self.id_to_idx:
                return self.id_to_idx[pid]
        # fallback: nearest node by haversine
        best, bi = float("inf"), 0
        for i, (la, ln) in enumerate(self.coords):
            d = haversine_m(obj["lat"], obj["lng"], la, ln)
            if d < best:
                best, bi = d, i
        return bi

    @property
    def n_nodes(self) -> int:
        return len(self.node_ids)

    def dijkstra(self, src: int) -> list[float]:
        """Shortest travel time (seconds) from src to every node."""
        INF = float("inf")
        dist = [INF] * self.n_nodes
        dist[src] = 0.0
        pq = [(0.0, src)]
        while pq:
            d, u = heapq.heappop(pq)
            if d > dist[u]:
                continue
            for k in range(self.offsets[u], self.offsets[u + 1]):
                v, w = self.edges[k], self.weights[k]
                nd = d + w
                if nd < dist[v]:
                    dist[v] = nd
                    heapq.heappush(pq, (nd, v))
        return dist

    def spt(self, a: int, b: int) -> float:
        """Shortest travel time a->b (seconds), cached per source."""
        cache = getattr(self, "_spt_cache", None)
        if cache is None:
            cache = {}
            object.__setattr__(self, "_spt_cache", cache)
        if a not in cache:
            cache[a] = self.dijkstra(a)
        return cache[a][b]


def build_graph(data: dict, speed_mps: float = DEFAULT_TRUCK_SPEED_MPS) -> QuarryGraph:
    points = data.get("point", [])
    node_ids = [p["id"] for p in points]
    id_to_idx = {pid: i for i, pid in enumerate(node_ids)}
    coords = [(p["lat"], p["lng"]) for p in points]

    # bidirectional adjacency -> dedup
    adj: list[set[int]] = [set() for _ in points]
    for i, p in enumerate(points):
        for nb in p.get("points", []) or []:
            j = id_to_idx.get(nb)
            if j is None or j == i:
                continue
            adj[i].add(j)
            adj[j].add(i)  # treat haul roads as two-way

    offsets = [0]
    edges: list[int] = []
    weights: list[int] = []
    for i in range(len(points)):
        for j in sorted(adj[i]):
            dist = haversine_m(*coords[i], *coords[j])
            edges.append(j)
            weights.append(max(1, round(dist / speed_mps)))  # seconds, >=1
        offsets.append(len(edges))

    objects = {
        t: list(data.get(t, []))
        for t in ("truck", "excavator", "crusher", "gas", "maintenance")
    }
    return QuarryGraph(node_ids, id_to_idx, offsets, edges, weights, coords, objects)

"""Solve the quarry dispatch VRP with cuOpt.

Pickup-and-Delivery VRP on the road waypoint graph.
  - vehicles  = trucks (start/end at their anchor node)
  - each task = one load: pickup at an excavator -> delivery at a crusher
  - capacity  = loads carried at once (default 1)
  - cost      = travel time (seconds) along haul roads; cuOpt routes shortest paths
"""
from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Any

from .graph import QuarryGraph

LOAD_TIME_S = 120   # excavator loading
DUMP_TIME_S = 60    # crusher dumping


@dataclass
class VRPProblem:
    payload: dict
    truck_ids: list[str]
    load_specs: list[dict]   # per load: excavator_id, crusher_id, nodes


def _assign_loads(graph: QuarryGraph, loads_per_truck: int) -> list[dict]:
    """Round-robin excavator -> nearest crusher loads to haul."""
    exc = graph.objects.get("excavator", [])
    cru = graph.objects.get("crusher", [])
    trucks = graph.objects.get("truck", [])
    if not exc or not cru or not trucks:
        return []

    # precompute each excavator's anchor and the nearest crusher by graph weight proxy
    specs = []
    n_loads = max(1, loads_per_truck) * len(trucks)
    for k in range(n_loads):
        e = exc[k % len(exc)]
        e_idx = graph.anchor_idx(e)
        # round-robin crushers so deliveries spread across all of them (not just #1)
        c = cru[k % len(cru)]
        c_idx = graph.anchor_idx(c)
        specs.append({
            "excavator_id": e["id"], "crusher_id": c["id"],
            "pickup_node": e_idx, "delivery_node": c_idx,
        })
    return specs


def build_problem(
    graph: QuarryGraph,
    loads_per_truck: int = 3,
    capacity: int = 1,
    time_limit_s: float = 2.0,
    min_vehicles: int | None = None,
) -> VRPProblem:
    trucks = graph.objects.get("truck", [])
    truck_ids = [t["id"] for t in trucks]
    specs = _assign_loads(graph, loads_per_truck)

    # ---- task data (interleaved pickup, delivery per load) ----
    task_locations: list[int] = []
    demand: list[int] = []
    service_times: list[int] = []
    pd_pairs: list[list[int]] = []
    for s in specs:
        pi = len(task_locations)
        task_locations.append(s["pickup_node"]);   demand.append(+capacity); service_times.append(LOAD_TIME_S)
        di = len(task_locations)
        task_locations.append(s["delivery_node"]);  demand.append(-capacity); service_times.append(DUMP_TIME_S)
        pd_pairs.append([pi, di])

    # ---- fleet data ----
    # capacities/demand are per-dimension: [[v0,v1,...]] for vehicles, [[t0,t1,...]] for tasks
    vehicle_locations = [[graph.anchor_idx(t), graph.anchor_idx(t)] for t in trucks]
    capacities = [[capacity for _ in trucks]]
    fleet_data: dict = {
        "vehicle_locations": vehicle_locations,
        "capacities": capacities,
    }
    if min_vehicles:
        fleet_data["min_vehicles"] = int(min_vehicles)

    payload = {
        "cost_waypoint_graph_data": {
            "waypoint_graph": {
                0: {"offsets": graph.offsets, "edges": graph.edges, "weights": graph.weights}
            }
        },
        "fleet_data": fleet_data,
        "task_data": {
            "task_locations": task_locations,
            "demand": [demand],
            "pickup_and_delivery_pairs": pd_pairs,
            "service_times": service_times,
        },
        "solver_config": {"time_limit": time_limit_s},
    }
    return VRPProblem(payload, truck_ids, specs)


def parse_solution(resp: Any) -> dict:
    """Normalize a cuOpt routing response into a plain dict."""
    if hasattr(resp, "json"):
        resp = resp.json()
    if isinstance(resp, dict) and "response" in resp:
        resp = resp["response"]
    sr = resp.get("solver_response", resp) if isinstance(resp, dict) else {}
    return {
        "status": sr.get("status"),
        "solution_cost": sr.get("solution_cost"),
        "num_vehicles": sr.get("num_vehicles"),
        "vehicle_data": sr.get("vehicle_data", {}),
        "dropped_tasks": sr.get("num_infeasible_nodes") or sr.get("dropped_tasks"),
        "msg": sr.get("msg") or sr.get("error"),
    }


def baseline_cost(graph: QuarryGraph, problem: VRPProblem) -> float:
    """Naive 'no-cuOpt' cost: each load hauled as an independent round trip
    (pickup -> delivery -> back to pickup), summed. Used for before/after."""
    total = 0.0
    for s in problem.load_specs:
        leg = graph.spt(s["pickup_node"], s["delivery_node"])
        total += 2.0 * leg + LOAD_TIME_S + DUMP_TIME_S
    return round(total, 1)


def solve(client, problem: VRPProblem) -> dict:
    t0 = time.perf_counter()
    resp = client.get_optimized_routes(problem.payload)
    solve_ms = round((time.perf_counter() - t0) * 1000.0, 1)
    sol = parse_solution(resp)
    sol["solve_ms"] = solve_ms
    return sol

"""GPU load-balancer (cuOpt plugin).

Spreads the fleet evenly across the facilities of a type so the route forecasts —
the 1st derivative (where each truck heads now) and the 2nd derivative (its next
stop) — distribute over ALL excavators / crushers / gas stations instead of
clustering on one. The balancing is done on the GPU by cuOpt, not by a Python
heuristic:

  - vehicles      = trucks, starting from their current graph node
  - tasks         = one "slot" per truck, slots placed round-robin across the
                    facilities (so the slot pool itself is evenly spread)
  - capacity 1    = each truck serves exactly one slot
  - min_vehicles  = fleet size, so every truck is placed
  - cost          = travel time on the road graph

cuOpt then matches each truck to its nearest free slot, yielding an even,
proximity-aware assignment. The chosen slot's facility is the recommendation.
"""
from __future__ import annotations

from .vrp import parse_solution


def _served_slot(vehicle_data) -> int | None:
    """The slot (task index) a vehicle served. cuOpt tags depot stops as 'Depot'
    in `task_id`; the remaining numeric entry is the served task."""
    tid = vehicle_data.get("task_id")
    if isinstance(tid, list):
        for t in tid:
            s = str(t)
            if s.lstrip("-").isdigit():
                return int(s)
    return None


def balance_by_type(client, graph, vehicle_nodes, facilities) -> list | None:
    """Assign each vehicle to one facility of a single type via cuOpt.

    vehicle_nodes : list[int]   graph node each truck starts from
    facilities    : list[dict]  [{"n":.., "node":..}, ...] of ONE type
    returns       : list        facility "n" chosen per vehicle (len == vehicles),
                                 or None on failure (caller uses a local fallback).
    """
    V = len(vehicle_nodes)
    F = len(facilities)
    if V == 0 or F == 0:
        return None
    slot_fac = [i % F for i in range(V)]                       # slot -> facility index
    task_locations = [facilities[slot_fac[i]]["node"] for i in range(V)]
    if any(n is None for n in vehicle_nodes) or any(n is None for n in task_locations):
        return None

    payload = {
        "cost_waypoint_graph_data": {"waypoint_graph": {
            0: {"offsets": graph.offsets, "edges": graph.edges, "weights": graph.weights}}},
        "fleet_data": {
            "vehicle_locations": [[n, n] for n in vehicle_nodes],
            "capacities": [[1] * V],
            "min_vehicles": V,
        },
        "task_data": {
            "task_locations": task_locations,
            "demand": [[1] * V],
        },
        "solver_config": {"time_limit": 1.0},
    }
    try:
        sol = parse_solution(client.get_optimized_routes(payload))
    except Exception:
        return None
    vd = sol.get("vehicle_data") or {}
    if not vd:
        return None

    out = [None] * V
    for vid, data in vd.items():
        try:
            vi = int(vid)
        except (TypeError, ValueError):
            continue
        slot = _served_slot(data)
        if slot is not None and 0 <= slot < V and 0 <= vi < V:
            out[vi] = facilities[slot_fac[slot]]["n"]
    # never return a partial plan — fill gaps round-robin so every truck has a target
    for i in range(V):
        if out[i] is None:
            out[i] = facilities[i % F]["n"]
    return out