milp_solver.py

"""Exact MILP reference implementation (PuLP + HiGHS).

Implements the formulation documented in Section 3.3 of the technical
documentation: arc variables x_ijk, serve variables y_i, service-start
times T_i, big-M time propagation (which also eliminates subtours), and
priority-scaled drop penalties. Solved to proven optimality with HiGHS.

Two details are added so the MILP's solution space matches solver.py's
OR-Tools model exactly (required for a fair objective comparison):
  * waiting at a stop is capped at config.MAX_WAIT_MINUTES (OR-Tools
    dimension slack has the same cap);
  * a claim's latest service start is max(a_i, b_i - s_i), matching the
    tightened window in solver.py.

This scales to roughly 15-25 claims; it exists as a correctness oracle
for the heuristic, not as a production solver.
"""

from __future__ import annotations

import itertools
import time
from dataclasses import dataclass, field

import pulp

import config
from data_gen import Adjuster, Claim


@dataclass
class ExactResult:
    status: str                 # "Optimal" = proven optimum
    engine: str = "highs"       # which solver produced this
    objective: float | None = None
    best_bound: float | None = None   # rigorous lower bound (dual bound)
    served_ids: set[str] = field(default_factory=set)
    routes: dict[str, list[str]] = field(default_factory=dict)
    solve_seconds: float = 0.0
    arrivals: dict[str, int] = field(default_factory=dict)
    spans: dict[str, tuple[int, int]] = field(default_factory=dict)


def _warm_values(adjusters, claims, x, y, T, Ts, Te, heur) -> dict:
    """Map MILP variables to the heuristic solution's values (MIP start)."""
    n_adj = len(adjusters)
    node_of = {c.claim_id: n_adj + i for i, c in enumerate(claims)}
    warm: dict = {}
    for r in heur.routes:
        k = next(kk for kk, a in enumerate(adjusters)
                 if a.adjuster_id == r.adjuster.adjuster_id)
        if r.stops:
            path = [k] + [node_of[s.claim.claim_id] for s in r.stops] + [k]
            for a_, b_ in zip(path, path[1:]):
                warm[x[(a_, b_, k)]] = 1.0
            for s in r.stops:
                warm[T[node_of[s.claim.claim_id]]] = float(s.arrival_min)
                warm[y[node_of[s.claim.claim_id]]] = 1.0
            warm[Ts[k]] = float(r.start_min)
            warm[Te[k]] = float(r.end_min)
    return warm


def _solve_with_mip_start(prob, solver, warm) -> None:
    """Replicate pulp.HiGHS.actualSolve with a MIP start injected between
    model build and run (PuLP's HiGHS wrapper has no warmStart param)."""
    import highspy

    solver.createAndConfigureSolver(prob)
    solver.buildSolverModel(prob)  # assigns var.index column ids
    values = [0.0] * len(prob.variables())
    for var in prob.variables():
        v = warm.get(var)
        if v is None:
            v = var.lowBound if var.lowBound is not None else 0.0
        values[var.index] = float(v)
    hsol = highspy.HighsSolution()
    hsol.col_value = values
    prob.solverModel.setSolution(hsol)

    solver.callSolver(prob)
    status, sol_status = solver.findSolutionValues(prob)
    for var in prob.variables():
        var.modified = False
    for constraint in prob.constraints.values():
        constraint.modifier = False
    prob.assignStatus(status, sol_status)


def _pick_engine(engine: str) -> str:
    """'auto' prefers Gurobi when a usable license is present."""
    if engine == "auto":
        try:
            if pulp.GUROBI(msg=False).available():
                return "gurobi"
        except Exception:
            pass
        return "highs"
    return engine


def solve_exact(adjusters: list[Adjuster], claims: list[Claim],
                travel_min: list[list[int]],
                time_limit_s: int = 300,
                warm_start=None,
                miles: list[list[float]] | None = None,
                engine: str = "highs",
                lunch_break: bool = False,
                balance: bool = False) -> ExactResult:
    """Solve to proven optimality.

    warm_start: a solver.Solution from the OR-Tools heuristic. Injected as
    a MIP start, so HiGHS begins with the heuristic incumbent and spends
    the whole budget proving (or improving) it. Without a warm start, hard
    instances can time out before HiGHS even finds a competitive incumbent.
    """
    n_adj = len(adjusters)
    H = list(range(n_adj))                       # home node of vehicle k is k
    C = list(range(n_adj, n_adj + len(claims)))  # claim nodes

    svc = {n: 0 for n in H}
    svc.update({n_adj + i: c.service_minutes for i, c in enumerate(claims)})
    # Eligibility = skills + service territory (radius needs the mileage
    # matrix; without it, radius-limited adjusters fall back to skill-only
    # - pass `miles` whenever any adjuster has max_radius_miles set).
    qual = {n_adj + i: [k for k, a in enumerate(adjusters)
                        if (config.is_eligible(a, c, miles[k][n_adj + i])
                            if miles is not None else c.peril in a.skills)]
            for i, c in enumerate(claims)}
    pen = {n_adj + i: config.effective_penalty(c.priority, c.age_days)
           for i, c in enumerate(claims)}
    # Service-start bounds. For split-availability claims these are
    # the HULL of all windows (sound for arc pruning and big-M
    # tightening); the window-choice disjunction below pins the start
    # inside one actual window.
    from data_gen import arrival_ranges
    t_lo = {n_adj + i: arrival_ranges(c)[0][0]
            for i, c in enumerate(claims)}
    t_hi = {n_adj + i: arrival_ranges(c)[-1][1]
            for i, c in enumerate(claims)}

    W = config.MAX_WAIT_MINUTES

    prob = pulp.LpProblem("vrptw_exact", pulp.LpMinimize)

    # Arc variables x[(i, j, k)], created only where feasible:
    # home(k)->claim, claim->claim, claim->home(k), all skill-filtered.
    # Arcs that can never satisfy the window/shift/wait-cap timing are
    # pruned outright - fewer binaries and much tighter big-Ms below.
    x = {}
    for k in range(n_adj):
        e_k, l_k = adjusters[k].shift_start, adjusters[k].shift_end
        for j in C:
            if k not in qual[j]:
                continue
            # home -> claim: arrival window reachable from the shift?
            if (e_k + travel_min[k][j] <= t_hi[j]
                    and l_k + travel_min[k][j] + W >= t_lo[j]):
                x[(k, j, k)] = pulp.LpVariable(f"x_h{k}_{j}_{k}", cat="Binary")
            # claim -> home: possible to be back before shift end?
            if t_lo[j] + svc[j] + travel_min[j][k] <= l_k:
                x[(j, k, k)] = pulp.LpVariable(f"x_{j}_h{k}_{k}", cat="Binary")
    for i, j in itertools.permutations(C, 2):
        if (t_lo[i] + svc[i] + travel_min[i][j] <= t_hi[j]
                and t_hi[i] + svc[i] + travel_min[i][j] + W >= t_lo[j]):
            for k in set(qual[i]) & set(qual[j]):
                x[(i, j, k)] = pulp.LpVariable(f"x_{i}_{j}_{k}", cat="Binary")

    y = {i: pulp.LpVariable(f"y_{i}", cat="Binary") for i in C}
    T = {i: pulp.LpVariable(f"T_{i}", lowBound=t_lo[i], upBound=t_hi[i])
         for i in C}

    # Split availability (multiple time windows): t_lo/t_hi above are
    # the HULL, sound for arc pruning and big-M tightening. For claims
    # with more than one window, a served claim must additionally pick
    # exactly one allowed service-start range; the linking big-M is the
    # per-claim hull width - the tightest constant available.
    for i, c in enumerate(claims):
        ranges = arrival_ranges(c)
        if len(ranges) <= 1:
            continue
        node_i = n_adj + i
        Mw = t_hi[node_i] - t_lo[node_i]
        zs = []
        for r, (lo, hi) in enumerate(ranges):
            z = pulp.LpVariable(f"win_{node_i}_{r}", cat="Binary")
            zs.append(z)
            prob += T[node_i] >= lo - Mw * (1 - z), f"wlo_{node_i}_{r}"
            prob += T[node_i] <= hi + Mw * (1 - z), f"whi_{node_i}_{r}"
        prob += pulp.lpSum(zs) == y[node_i], f"wpick_{node_i}"
    Ts = {k: pulp.LpVariable(f"Tstart_{k}", lowBound=a.shift_start,
                             upBound=a.shift_end)
          for k, a in enumerate(adjusters)}
    Te = {k: pulp.LpVariable(f"Tend_{k}", lowBound=a.shift_start,
                             upBound=a.shift_end)
          for k, a in enumerate(adjusters)}

    # Objective: travel minutes + drop penalties (+ optional balance:
    # the span term is linear in existing variables - no big-M needed).
    obj = (pulp.lpSum(travel_min[i][j] * v for (i, j, k), v in x.items())
           + pulp.lpSum(pen[i] * (1 - y[i]) for i in C))
    if balance:
        obj += config.BALANCE_COEFFICIENT * pulp.lpSum(
            Te[k] - Ts[k] for k in range(n_adj))
    prob += obj

    # Coverage: each served claim entered exactly once (qualified k only,
    # by variable construction).
    for i in C:
        prob += (pulp.lpSum(v for (a_, b_, k), v in x.items() if b_ == i)
                 == y[i])

    # Flow conservation at each claim, per vehicle.
    for i in C:
        for k in qual[i]:
            inflow = pulp.lpSum(v for (a_, b_, kk), v in x.items()
                                if b_ == i and kk == k)
            outflow = pulp.lpSum(v for (a_, b_, kk), v in x.items()
                                 if a_ == i and kk == k)
            prob += inflow == outflow

    # Each vehicle leaves home at most once and returns as often as it
    # leaves (0 or 1 times).
    for k in range(n_adj):
        dep = pulp.lpSum(v for (a_, b_, kk), v in x.items()
                         if a_ == k and kk == k)
        ret = pulp.lpSum(v for (a_, b_, kk), v in x.items()
                         if b_ == k and kk == k)
        prob += dep <= 1
        prob += dep == ret
        prob += Te[k] >= Ts[k]

    # Time propagation (big-M) with the wait cap; eliminates subtours.
    # Big-Ms are tightened per arc from the variable bounds - the weakest
    # value that still deactivates the constraint when the arc is unused.
    # Claim-to-claim constraints are aggregated over vehicles (at most one
    # vehicle uses arc i->j, so sum_k x_ijk is 0/1): one constraint pair
    # per arc instead of per vehicle, and a tighter LP relaxation.
    cc_arcs: dict[tuple[int, int], list] = {}
    for (i, j, k), v in x.items():
        if i in T and j in T:
            cc_arcs.setdefault((i, j), []).append(v)
    for (i, j), vs in cc_arcs.items():
        t_ij = travel_min[i][j]
        used = pulp.lpSum(vs)
        m_lo = max(0, t_hi[i] + svc[i] + t_ij - t_lo[j])
        m_up = max(0, t_hi[j] - (t_lo[i] + svc[i] + t_ij + W))
        prob += T[j] >= T[i] + svc[i] + t_ij - m_lo * (1 - used)
        prob += T[j] <= T[i] + svc[i] + t_ij + W + m_up * (1 - used)
    for (i, j, k), v in x.items():
        t_ij = travel_min[i][j]
        if i == k and j in T:            # home -> first claim
            e_k, l_k = adjusters[k].shift_start, adjusters[k].shift_end
            m_lo = max(0, l_k + t_ij - t_lo[j])
            m_up = max(0, t_hi[j] - (e_k + t_ij + W))
            prob += T[j] >= Ts[k] + t_ij - m_lo * (1 - v)
            prob += T[j] <= Ts[k] + t_ij + W + m_up * (1 - v)
        elif i in T and j == k:          # last claim -> home
            m_lo = max(0, t_hi[i] + svc[i] + t_ij - adjusters[k].shift_start)
            prob += Te[k] >= T[i] + svc[i] + t_ij - m_lo * (1 - v)

    # Optional mandatory lunch: mirrors the OR-Tools break interval.
    # Each adjuster takes one 30-min break starting inside the window;
    # the break may not overlap any visit and, when it falls in a
    # driven gap, it consumes 30 minutes of that gap's slack (the
    # classic 'or' -> indicator-binary + big-M pattern per placement).
    if lunch_break:
        L_lo = config.LUNCH_BREAK["earliest_start"]
        L_hi = config.LUNCH_BREAK["latest_start"]
        dur = config.LUNCH_BREAK["duration"]
        for k in range(n_adj):
            e_k, l_k = adjusters[k].shift_start, adjusters[k].shift_end
            Lk = pulp.LpVariable(f"lunch_{k}", lowBound=L_lo,
                                 upBound=L_hi)
            placements = []
            u_pre = pulp.LpVariable(f"lun_pre_{k}", cat="Binary")
            m = max(0, L_hi + dur - e_k)
            prob += Lk + dur <= Ts[k] + m * (1 - u_pre)
            placements.append(u_pre)
            u_post = pulp.LpVariable(f"lun_post_{k}", cat="Binary")
            m = max(0, l_k - L_lo)
            prob += Lk >= Te[k] - m * (1 - u_post)
            placements.append(u_post)
            for (i, j, kk), v in x.items():
                if kk != k:
                    continue
                dep_lo = e_k if i == k else t_lo[i] + svc[i]
                dep_hi = l_k if i == k else t_hi[i] + svc[i]
                arr_lo = e_k if j == k else t_lo[j]
                arr_hi = l_k if j == k else t_hi[j]
                dep = Ts[k] if i == k else T[i] + svc[i]
                arr = Te[k] if j == k else T[j]
                w = pulp.LpVariable(f"lun_{i}_{j}_{k}", cat="Binary")
                prob += w <= v
                m = max(0, dep_hi - L_lo)
                prob += Lk >= dep - m * (1 - w)
                m = max(0, L_hi + dur - arr_lo)
                prob += Lk + dur <= arr + m * (1 - w)
                m = max(0, dep_hi + travel_min[i][j] + dur - arr_lo)
                prob += arr >= dep + travel_min[i][j] + dur - m * (1 - w)
                placements.append(w)
            prob += pulp.lpSum(placements) == 1

    # 2-cycle elimination cuts: implied by time propagation in the integer
    # model, but they strengthen the LP relaxation considerably.
    for (i, j) in cc_arcs:
        if i < j and (j, i) in cc_arcs:
            prob += pulp.lpSum(cc_arcs[(i, j)] + cc_arcs[(j, i)]) <= 1

    engine = _pick_engine(engine)
    warm = (None if warm_start is None else
            _warm_values(adjusters, claims, x, y, T, Ts, Te, warm_start))
    t0 = time.time()
    if engine == "gurobi":
        # PuLP's Gurobi interface supports MIP starts natively.
        if warm:
            for var, val in warm.items():
                var.setInitialValue(val)
        # Gurobi has twice been observed returning a WRONG 'Optimal'
        # certificate on this formulation, via two distinct mechanisms:
        # (1) lunch-break disjunctions: an unsound presolve reduction
        #     certified a value 9 above the true optimum (seed 3003) -
        #     fixed by Presolve=0;
        # (2) base models: default integer/feasibility tolerances let
        #     the big-M rows leak, pruning the true optimum (seed 404:
        #     certified 2,326 where 2,314 is feasible and optimal -
        #     HiGHS, OR-Tools at 60 s, CP-SAT, and warm-started Gurobi
        #     itself all agree) - fixed by the 1e-9 tolerances.
        # The full hardening set therefore applies to EVERY Gurobi
        # solve. Measured cost on the hardest referee proof (seed 505):
        # 77.8 s vs 75.3 s default - about 3%.
        params = dict(msg=False, timeLimit=time_limit_s, gapRel=0.0,
                      warmStart=bool(warm), NumericFocus=3,
                      IntFeasTol=1e-9, FeasibilityTol=1e-9, Presolve=0)
        solver = pulp.GUROBI(**params)
        try:
            prob.solve(solver)
        except Exception as e:
            # Typical cause: size-limited trial license on big models.
            print(f"Gurobi failed ({str(e)[:80]}); falling back to HiGHS.")
            engine = "highs"
    if engine == "highs":
        solver = pulp.HiGHS(msg=False, timeLimit=time_limit_s, gapRel=0.0)
        if warm is None:
            prob.solve(solver)
        else:
            _solve_with_mip_start(prob, solver, warm)
    wall_seconds = time.time() - t0

    # prob.status says "Optimal" even when HiGHS merely hit its time limit
    # with a feasible incumbent; prob.sol_status distinguishes a *proven*
    # optimum (LpSolutionOptimal) from a plain feasible solution.
    if prob.sol_status == pulp.LpSolutionOptimal:
        status = "Optimal"
    elif prob.sol_status == pulp.LpSolutionIntegerFeasible:
        status = "Feasible (limit)"
    else:
        status = pulp.LpStatus[prob.status]

    # PuLP's Gurobi wrapper does not copy variable values back when
    # the time limit fires, even though the solver may hold a feasible
    # incumbent - recover it so a time-limited run still returns real
    # routes instead of an empty (all-dropped) reconstruction.
    if (engine == "gurobi"
            and prob.sol_status != pulp.LpSolutionOptimal
            and getattr(prob, "solverModel", None) is not None):
        try:
            if prob.solverModel.SolCount > 0:
                for v in prob.variables():
                    if getattr(v, "solverVar", None) is not None:
                        v.varValue = v.solverVar.X
                status = "Feasible (limit)"
        except Exception:
            pass

    best_bound = None
    if getattr(prob, "solverModel", None) is not None:
        try:
            best_bound = prob.solverModel.getInfo().mip_dual_bound
        except Exception:
            try:
                best_bound = prob.solverModel.ObjBound
            except Exception:
                pass

    result = ExactResult(status=status, engine=engine,
                         objective=pulp.value(prob.objective),
                         best_bound=best_bound,
                         solve_seconds=wall_seconds)
    if status not in ("Optimal", "Feasible (limit)"):
        return result

    result.served_ids = {claims[i - n_adj].claim_id
                         for i in C if y[i].value() > 0.5}
    for i in C:
        if y[i].value() > 0.5 and T[i].value() is not None:
            result.arrivals[claims[i - n_adj].claim_id] = int(
                round(T[i].value()))
    for k, adj in enumerate(adjusters):
        if Ts[k].value() is not None and Te[k].value() is not None:
            result.spans[adj.adjuster_id] = (int(round(Ts[k].value())),
                                             int(round(Te[k].value())))
    # Reconstruct routes by following selected arcs.
    chosen = {(i, j): k for (i, j, k), v in x.items() if v.value() > 0.5}
    for k, adj in enumerate(adjusters):
        route, node = [], k
        while True:
            nxt = next((j for (i, j) in chosen if i == node
                        and chosen[(i, j)] == k), None)
            if nxt is None or nxt == k:
                break
            route.append(claims[nxt - n_adj].claim_id)
            node = nxt
        result.routes[adj.adjuster_id] = route
    return result


def solve_as_backend(adjusters, claims, miles, travel_min,
                     time_limit_s: int = 60, engine: str = "auto",
                     lunch_break: bool = False, balance: bool = False):
    """Use the exact MILP as a schedule-producing backend (--solver milp).

    Warm-started by a quick OR-Tools solve, then solved exactly. Returns
    (Solution, ExactResult) - the Solution carries earliest-feasible stop
    times reconstructed from the proven visit sequences, in the same
    shape every report/map/app consumer expects. Sensible up to ~15-25
    claims; beyond that expect a time-limited incumbent, not a proof.
    """
    import solver as ortools_solver
    from solver import Route, Solution, Stop

    heur = ortools_solver.solve(adjusters, claims, miles, travel_min,
                                time_limit_s=min(5, time_limit_s),
                                lunch_break=lunch_break, balance=balance)
    exact = solve_exact(adjusters, claims, travel_min,
                        time_limit_s=time_limit_s, warm_start=heur,
                        miles=miles, engine=engine,
                        lunch_break=lunch_break, balance=balance)
    if exact.objective is None:
        return heur, exact
    if heur is not None and exact.objective > heur.objective:
        return heur, exact   # time-limited MILP lost to its own start

    n_adj = len(adjusters)
    node = {c.claim_id: n_adj + i for i, c in enumerate(claims)}
    by_id = {c.claim_id: c for c in claims}
    routes, served = [], set()
    for k, a in enumerate(adjusters):
        seq = exact.routes.get(a.adjuster_id, [])
        route = Route(adjuster=a)
        prev, t = k, a.shift_start
        for cid in seq:
            c = by_id[cid]
            leg = travel_min[prev][node[cid]]
            arrival = exact.arrivals.get(
                cid, max(c.window_start, t + leg))
            route.stops.append(Stop(
                claim=c, arrival_min=arrival,
                departure_min=arrival + c.service_minutes,
                travel_miles_from_prev=miles[prev][node[cid]],
                travel_min_from_prev=leg))
            route.total_miles += miles[prev][node[cid]]
            route.total_travel_min += leg
            route.total_service_min += c.service_minutes
            served.add(cid)
            prev, t = node[cid], arrival + c.service_minutes
        span = (exact.spans.get(a.adjuster_id)
                if (lunch_break or balance) else None)
        if seq:
            route.start_min = (span[0] if span else
                               route.stops[0].arrival_min
                               - route.stops[0].travel_min_from_prev)
            route.total_miles += miles[prev][k]
            route.total_travel_min += travel_min[prev][k]
            route.end_min = (span[1] if span
                             else t + travel_min[prev][k])
        else:
            route.start_min = route.end_min = a.shift_start
        routes.append(route)

    dropped = [c for c in claims if c.claim_id not in served]
    unservable = [c for i, c in enumerate(claims)
                  if not any(config.is_eligible(a, c,
                                                miles[k][n_adj + i])
                             for k, a in enumerate(adjusters))]
    travel_total = sum(r.total_travel_min for r in routes)
    objective = travel_total + sum(
        config.effective_penalty(c.priority, c.age_days)
        for c in dropped)
    if balance:
        objective += config.BALANCE_COEFFICIENT * sum(
            r.end_min - r.start_min for r in routes if r.stops)
    sol = Solution(
        routes=routes, dropped=dropped, unservable=unservable,
        objective=objective,
        total_miles=sum(r.total_miles for r in routes),
        total_travel_min=travel_total)
    if heur is not None and sol.objective > heur.objective:
        return heur, exact   # reconstruction sanity: never worse than start
    return sol, exact