"""Vehicle routing with time windows with Heusor, on Solomon's c101.

A fleet of vehicles leaves one depot to deliver to 100 customers. Each
customer has a demand, a time window in which delivery must start, and a
service time. A vehicle carries no more than its capacity, waits when it is
early, and must be back before the depot closes. The total distance driven is
as small as possible.

Run it:  python vehicle_routing.py
"""

from __future__ import annotations

import copy
import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Annotated

import heusor
from heusor import (
    CascadingShadow,
    ConstraintFactory,
    HardSoftScore,
    InverseRelationShadow,
    NextElementShadow,
    PlanningListVariable,
    PlanningPin,
    PreviousElementShadow,
    constraint_provider,
    planning_entity,
    planning_solution,
)

# ─── The model ──────────────────────────────────────────────────────────────
#
# Times and distances are in tenths of Solomon's unit: a distance is
# floor(10 × the straight-line distance), and a vehicle drives one unit of
# distance per unit of time. Distances are symmetric, so a visit's drive to the
# depot, ``times[0]``, is also the drive from it.


def arrival_time(visit: "Visit", previous: "Visit | None") -> int:
    """At the head of a route: when the depot opens plus the drive from it.
    Otherwise: when the previous visit is done plus the drive from there."""
    return (
        visit.vehicle.ready + visit.times[0]
        if previous is None
        else previous.departure + previous.times[visit.node]
    )


@planning_entity
@dataclass
class Visit:
    x: int
    y: int
    demand: int
    ready: int  # delivery may start from …
    due: int  # … and must start by
    service: int
    node: int  # this customer's number; the depot is 0
    times: list[int]  # the drive from here to each node, by node number
    vehicle: Annotated["Vehicle | None", InverseRelationShadow] = None
    previous: Annotated["Visit | None", PreviousElementShadow] = None
    next: Annotated["Visit | None", NextElementShadow] = None
    arrival: Annotated[int | None, CascadingShadow(arrival_time)] = None
    departure: Annotated[
        int | None, CascadingShadow(lambda visit, previous: max(visit.arrival, visit.ready) + visit.service)
    ] = None


@planning_entity
@dataclass
class Vehicle:
    id: int
    capacity: int
    ready: int  # the depot opens
    due: int  # the depot closes
    visits: Annotated[list[Visit], PlanningListVariable] = field(default_factory=list)
    pinned: Annotated[bool, PlanningPin] = False


@planning_solution
@dataclass
class Plan:
    vehicles: list[Vehicle]
    visits: list[Visit]
    score: HardSoftScore | None = None


# ─── The constraints ────────────────────────────────────────────────────────


def drive_in(visit: Visit) -> int:
    """The drive to this visit: from the depot or from the previous visit."""
    return visit.times[0] if visit.previous is None else visit.times[visit.previous.node]


@constraint_provider
def define_constraints(f: ConstraintFactory):
    return [
        # Hard: every customer is visited.
        f.for_each_including_unassigned(Visit)
        .filter(lambda visit: visit.vehicle is None)
        .penalize(HardSoftScore.ONE_HARD)
        .as_constraint("Visit assigned"),
        # Hard: a vehicle carries no more than its capacity.
        f.for_each(Vehicle)
        .filter(lambda vehicle: sum(v.demand for v in vehicle.visits) > vehicle.capacity)
        .penalize(HardSoftScore.ONE_HARD, lambda vehicle: sum(v.demand for v in vehicle.visits) - vehicle.capacity)
        .as_constraint("Vehicle capacity"),
        # Hard: delivery starts by the customer's due time — penalized by the lateness.
        f.for_each(Visit)
        .filter(lambda visit: visit.arrival > visit.due)
        .penalize(HardSoftScore.ONE_HARD, lambda visit: visit.arrival - visit.due)
        .as_constraint("Arrival after due time"),
        # Hard: the vehicle is back before the depot closes.
        f.for_each(Visit)
        .filter(lambda visit: visit.next is None and visit.departure + visit.times[0] > visit.vehicle.due)
        .penalize(HardSoftScore.ONE_HARD, lambda visit: visit.departure + visit.times[0] - visit.vehicle.due)
        .as_constraint("Back after depot closes"),
        # Soft: drive as little as possible — each visit's drive in, and the
        # last visit's drive home.
        f.for_each(Visit)
        .penalize(HardSoftScore.ONE_SOFT, lambda visit: drive_in(visit) + (visit.times[0] if visit.next is None else 0))
        .as_constraint("Distance"),
    ]


# ─── The data ───────────────────────────────────────────────────────────────

SCALE = 10


def load_solomon(path: Path = Path(__file__).with_name("c101.txt")) -> Plan:
    """A Solomon VRPTW instance: the vehicle count and capacity, then one row per
    node (number, x, y, demand, ready, due, service); node 0 is the depot."""
    lines = [line.split() for line in path.read_text().splitlines() if line.split()]
    fleet = next(i for i, words in enumerate(lines) if words[0] == "NUMBER") + 1
    count, capacity = int(lines[fleet][0]), int(lines[fleet][1])
    rows = [[int(w) for w in words] for words in lines if words[0].isdigit() and len(words) == 7]
    xy = [(x, y) for _, x, y, *_ in rows]
    times = [[math.floor(SCALE * math.dist(a, b)) for b in xy] for a in xy]
    depot = rows[0]
    visits = [
        Visit(x, y, demand, ready * SCALE, due * SCALE, service * SCALE, node, times[node])
        for node, x, y, demand, ready, due, service in rows[1:]
    ]
    vehicles = [Vehicle(i, capacity, depot[4] * SCALE, depot[5] * SCALE) for i in range(count)]
    return Plan(vehicles, visits)


# ─── Solving and explaining ─────────────────────────────────────────────────


def solve(plan: Plan, time_limit: float | None = 30, step_limit: int | None = None, seed: int = 0) -> Plan:
    return heusor.solve(plan, define_constraints, time_limit=time_limit, step_limit=step_limit, seed=seed)


def explain(plan: Plan) -> None:
    explanation = heusor.score(plan, define_constraints)
    print(f"Score: {explanation.score}")
    for name, score in explanation.constraints.items():
        print(f"  {name:26} {score}")
    used = [v for v in plan.vehicles if v.visits]
    print(f"  {len(used)} vehicles, {-explanation.score.soft / SCALE:.1f} distance")


# ─── Late orders ────────────────────────────────────────────────────────────


def add_late_orders(
    planned: Plan,
    late: list[Visit],
    keep: int | None,
    time_limit: float | None = 10,
    step_limit: int | None = None,
    seed: int = 0,
) -> Plan:
    """Orders ``late`` arrive after the plan was made. The first ``keep``
    vehicles with a route (all of them when ``keep`` is ``None``) are already
    loaded, so their routes are pinned; the late orders go into the other
    routes or into spare vehicles."""
    plan = copy.deepcopy(planned)
    plan.visits += copy.deepcopy(late)
    for vehicle in [v for v in plan.vehicles if v.visits][:keep]:
        vehicle.pinned = True
    return heusor.solve(plan, define_constraints, time_limit=time_limit, step_limit=step_limit, seed=seed)


if __name__ == "__main__":
    full = load_solomon()
    # The plan is made before the last five orders come in.
    late = full.visits[-5:]
    morning = Plan(full.vehicles, full.visits[:-5])
    planned = solve(morning)
    explain(planned)
    updated = add_late_orders(planned, late, keep=5)
    print("\nAfter five late orders, with the first five routes already loaded:")
    explain(updated)
    loaded = [v.id for v in planned.vehicles if v.visits][:5]
    print(f"  Loaded, unchanged: vehicles {loaded}")
    for visit in updated.visits[-5:]:
        print(f"  Order {visit.node} → vehicle {visit.vehicle.id}")