solver.py

"""OR-Tools routing model for daily claim assignment.

Model summary
-------------
- Each adjuster is a "vehicle" that starts and ends its route at their home.
- Objective: minimize total travel minutes + penalties for dropped claims.
- Time dimension tracks the clock (minutes since midnight). The transit
  between two stops = service time at the first stop + drive time. Slack in
  the dimension lets an adjuster wait for a policyholder's window to open.
- Constraints:
    * policyholder availability window on each claim's arrival time
    * adjuster shift window on each route's start and end
    * skill matching via SetAllowedVehiclesForIndex
    * every claim is in a disjunction, so it can be dropped (rescheduled to
      a later day) at a priority-scaled penalty; MUST-TODAY claims carry a
      penalty so large they are only dropped when physically infeasible.
"""

from __future__ import annotations

from dataclasses import dataclass, field

from ortools.constraint_solver import pywrapcp, routing_enums_pb2

import config
import data_gen
from data_gen import Adjuster, Claim


@dataclass
class Stop:
    claim: Claim
    arrival_min: int          # service starts on arrival (within the window)
    departure_min: int
    travel_miles_from_prev: float
    travel_min_from_prev: int


@dataclass
class Route:
    adjuster: Adjuster
    stops: list[Stop] = field(default_factory=list)
    start_min: int = 0
    end_min: int = 0
    total_miles: float = 0.0
    total_travel_min: int = 0
    total_service_min: int = 0


@dataclass
class Solution:
    routes: list[Route]
    dropped: list[Claim]
    unservable: list[Claim]   # no adjuster has the required skill at all
    objective: int
    total_miles: float = 0.0
    total_travel_min: int = 0

    @property
    def dropped_must_today(self) -> list[Claim]:
        return [c for c in self.dropped
                if c.priority == config.PRIORITY_MUST_TODAY]


def solve(adjusters: list[Adjuster], claims: list[Claim],
          miles: list[list[float]], travel_min: list[list[int]],
          time_limit_s: int = config.SOLVER_TIME_LIMIT_SECONDS,
          lunch_break: bool = False, balance: bool = False,
          ) -> Solution | None:
    n_adj = len(adjusters)
    n_nodes = n_adj + len(claims)
    # Service time by node (0 at adjuster homes).
    service_min = [0] * n_adj + [c.service_minutes for c in claims]

    starts = list(range(n_adj))
    ends = list(range(n_adj))
    manager = pywrapcp.RoutingIndexManager(n_nodes, n_adj, starts, ends)
    routing = pywrapcp.RoutingModel(manager)

    # --- Objective: pure travel time on arcs (service time is a constant
    # for every served claim, so keeping it out of the cost makes the
    # objective easier to interpret).
    def travel_cb(from_index: int, to_index: int) -> int:
        f = manager.IndexToNode(from_index)
        t = manager.IndexToNode(to_index)
        return travel_min[f][t]

    travel_cb_idx = routing.RegisterTransitCallback(travel_cb)
    routing.SetArcCostEvaluatorOfAllVehicles(travel_cb_idx)

    # --- Time dimension: clock advances by service-at-origin + drive time.
    def time_cb(from_index: int, to_index: int) -> int:
        f = manager.IndexToNode(from_index)
        t = manager.IndexToNode(to_index)
        return service_min[f] + travel_min[f][t]

    time_cb_idx = routing.RegisterTransitCallback(time_cb)
    routing.AddDimension(
        time_cb_idx,
        config.MAX_WAIT_MINUTES,       # slack: allowed waiting at a stop
        config.TIME_HORIZON_MINUTES,   # dimension upper bound
        False,                         # don't force start cumul to zero
        "Time",
    )
    time_dim = routing.GetDimensionOrDie("Time")

    # Policyholder availability: arrival (= service start) inside the window.
    # The window is tightened so the visit also *finishes* by window end.
    for i, claim in enumerate(claims):
        idx = manager.NodeToIndex(n_adj + i)
        ranges = data_gen.arrival_ranges(claim)
        cumul = time_dim.CumulVar(idx)
        cumul.SetRange(ranges[0][0], ranges[-1][1])
        for (_, hi), (lo2, _) in zip(ranges, ranges[1:]):
            cumul.RemoveInterval(hi + 1, lo2 - 1)

    # Adjuster shifts bound each route's start and end times.
    for v, adj in enumerate(adjusters):
        time_dim.CumulVar(routing.Start(v)).SetRange(adj.shift_start, adj.shift_end)
        time_dim.CumulVar(routing.End(v)).SetRange(adj.shift_start, adj.shift_end)

    # Optional 30-min lunch break inside a configurable window.
    if lunch_break:
        cp = routing.solver()
        visit_transits = [service_min[manager.IndexToNode(i)]
                          for i in range(routing.Size())]
        for v in range(n_adj):
            interval = cp.FixedDurationIntervalVar(
                config.LUNCH_BREAK["earliest_start"],
                config.LUNCH_BREAK["latest_start"],
                config.LUNCH_BREAK["duration"], False, f"lunch_{v}")
            time_dim.SetBreakIntervalsOfVehicle([interval], v,
                                                visit_transits)

    # Optional workload balancing: charge each minute of route span so
    # stops spread across adjusters instead of piling onto one.
    if balance:
        time_dim.SetSpanCostCoefficientForAllVehicles(
            config.BALANCE_COEFFICIENT)

    # Prefer schedules that start late / end early when travel cost ties.
    for v in range(n_adj):
        routing.AddVariableMinimizedByFinalizer(time_dim.CumulVar(routing.End(v)))
        routing.AddVariableMaximizedByFinalizer(time_dim.CumulVar(routing.Start(v)))

    # --- Skill matching + droppable claims.
    # Skill matching constrains each claim's VehicleVar to qualified
    # adjusters; -1 stays in the domain so the claim can still be dropped
    # (the disjunction below decides whether that's worth the penalty).
    unservable: list[Claim] = []
    for i, claim in enumerate(claims):
        idx = manager.NodeToIndex(n_adj + i)
        allowed = [v for v, adj in enumerate(adjusters)
                   if config.is_eligible(adj, claim, miles[v][n_adj + i])]
        if not allowed:
            unservable.append(claim)
        routing.VehicleVar(idx).SetValues([-1] + allowed)
        routing.AddDisjunction(
            [idx], config.effective_penalty(claim.priority, claim.age_days))

    # --- Search strategy.
    params = pywrapcp.DefaultRoutingSearchParameters()
    params.first_solution_strategy = (
        routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC)
    params.local_search_metaheuristic = (
        routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH)
    params.time_limit.FromSeconds(time_limit_s)

    assignment = routing.SolveWithParameters(params)
    if assignment is None:
        return None

    # --- Extract solution.
    routes: list[Route] = []
    served_nodes: set[int] = set()
    for v, adj in enumerate(adjusters):
        route = Route(adjuster=adj)
        index = routing.Start(v)
        route.start_min = assignment.Value(time_dim.CumulVar(index))
        prev_node = manager.IndexToNode(index)
        while not routing.IsEnd(index):
            index = assignment.Value(routing.NextVar(index))
            node = manager.IndexToNode(index)
            leg_miles = miles[prev_node][node]
            leg_min = travel_min[prev_node][node]
            if routing.IsEnd(index):
                route.end_min = assignment.Value(time_dim.CumulVar(index))
                route.total_miles += leg_miles
                route.total_travel_min += leg_min
                break
            claim = claims[node - n_adj]
            arrival = assignment.Value(time_dim.CumulVar(index))
            route.stops.append(Stop(
                claim=claim,
                arrival_min=arrival,
                departure_min=arrival + claim.service_minutes,
                travel_miles_from_prev=leg_miles,
                travel_min_from_prev=leg_min,
            ))
            route.total_miles += leg_miles
            route.total_travel_min += leg_min
            route.total_service_min += claim.service_minutes
            served_nodes.add(node)
            prev_node = node
        routes.append(route)

    dropped = [c for i, c in enumerate(claims) if (n_adj + i) not in served_nodes]
    sol = Solution(
        routes=routes,
        dropped=dropped,
        unservable=unservable,
        objective=assignment.ObjectiveValue(),
        total_miles=sum(r.total_miles for r in routes),
        total_travel_min=sum(r.total_travel_min for r in routes),
    )
    return sol