Design a Ride-Sharing Service.

# ride.py, part 1: places, drivers, and the trip's states
from __future__ import annotations

import itertools
import math
import threading
from abc import ABC, abstractmethod
from collections import defaultdict, deque
from dataclasses import dataclass, field
from enum import Enum


@dataclass(frozen=True)
class Location:
    x: float  # km east of a reference point; production uses lat/lng and haversine
    y: float  # km north

    def km_to(self, other: Location) -> float:
        return math.hypot(self.x - other.x, self.y - other.y)


class DriverStatus(Enum):
    OFFLINE = "offline"
    AVAILABLE = "available"
    ON_TRIP = "on_trip"


@dataclass
class Driver:
    driver_id: str
    name: str
    location: Location
    status: DriverStatus = DriverStatus.AVAILABLE
    trip_id: str | None = None
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)

    def try_reserve(self, trip_id: str) -> bool:
        """Compare-and-set AVAILABLE -> ON_TRIP: of two riders, exactly one wins this driver."""
        with self._lock:
            if self.status is not DriverStatus.AVAILABLE:
                return False
            self.status, self.trip_id = DriverStatus.ON_TRIP, trip_id
            return True

    def release(self) -> None:
        with self._lock:
            self.status, self.trip_id = DriverStatus.AVAILABLE, None


class TripStatus(Enum):
    REQUESTED = "requested"
    ACCEPTED = "accepted"
    STARTED = "started"
    COMPLETED = "completed"
    CANCELLED = "cancelled"


ALLOWED = {
    TripStatus.REQUESTED: {TripStatus.ACCEPTED, TripStatus.CANCELLED},
    TripStatus.ACCEPTED: {TripStatus.STARTED, TripStatus.CANCELLED},
    TripStatus.STARTED: {TripStatus.COMPLETED},
    TripStatus.COMPLETED: set(),
    TripStatus.CANCELLED: set(),
}


class InvalidTransition(Exception): ...


@dataclass
class Trip:
    trip_id: str
    rider: str
    pickup: Location
    drop: Location
    surge: float = 1.0  # locked in at request time, never recomputed at the drop
    status: TripStatus = TripStatus.REQUESTED
    driver_id: str | None = None
    accepted_at: float | None = None
    fare: int | None = None
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)

    def move(self, to: TripStatus) -> None:
        with self._lock:  # the rider's Cancel and the driver's Start can't both win
            if to not in ALLOWED[self.status]:
                raise InvalidTransition(f"{self.trip_id}: {self.status.value} -> {to.value}")
            self.status = to

# part 2: where the drivers are: a grid of 1 km cells
class DriverIndex:
    def __init__(self, cell_km: float = 1.0) -> None:
        self.cell_km = cell_km
        self._cells: dict[tuple[int, int], set[str]] = defaultdict(set)
        self._drivers: dict[str, Driver] = {}
        self._mu = threading.Lock()

    def _cell(self, loc: Location) -> tuple[int, int]:
        return math.floor(loc.x / self.cell_km), math.floor(loc.y / self.cell_km)

    def add(self, d: Driver) -> None:
        with self._mu:
            self._drivers[d.driver_id] = d
            self._cells[self._cell(d.location)].add(d.driver_id)

    def get(self, driver_id: str) -> Driver:
        return self._drivers[driver_id]

    def update_location(self, driver_id: str, loc: Location) -> None:
        with self._mu:  # the driver app pings every few seconds
            d = self._drivers[driver_id]
            self._cells[self._cell(d.location)].discard(driver_id)
            d.location = loc
            self._cells[self._cell(loc)].add(driver_id)

    def nearby(self, loc: Location, radius_km: float) -> list[tuple[float, Driver]]:
        """Available drivers within the radius. A snapshot: try_reserve is the real check."""
        r = math.ceil(radius_km / self.cell_km)
        cx, cy = self._cell(loc)
        found = []
        with self._mu:
            for dx in range(-r, r + 1):
                for dy in range(-r, r + 1):
                    for did in self._cells.get((cx + dx, cy + dy), ()):
                        d = self._drivers[did]
                        km = d.location.km_to(loc)
                        if km <= radius_km and d.status is DriverStatus.AVAILABLE:
                            found.append((km, d))
        return found

# part 3: the two strategies: who gets the ride, and what it costs
class MatchingStrategy(ABC):
    @abstractmethod
    def rank(self, pickup: Location, candidates: list[tuple[float, Driver]]) -> list[Driver]: ...


class NearestFirst(MatchingStrategy):
    def rank(self, pickup: Location, candidates: list[tuple[float, Driver]]) -> list[Driver]:
        return [d for _, d in sorted(candidates, key=lambda c: (c[0], c[1].driver_id))]


class FareStrategy(ABC):
    @abstractmethod
    def fare(self, km: float, minutes: float) -> int: ...


@dataclass(frozen=True)
class StandardFare(FareStrategy):
    base: int = 50  # rupees
    per_km: int = 12
    per_min: int = 2
    minimum: int = 80

    def fare(self, km: float, minutes: float) -> int:
        return max(self.minimum, round(self.base + self.per_km * km + self.per_min * minutes))


class SurgeFare(FareStrategy):
    """Decorator: any fare strategy, times the multiplier the rider agreed to."""

    def __init__(self, inner: FareStrategy, multiplier: float) -> None:
        self.inner, self.multiplier = inner, multiplier

    def fare(self, km: float, minutes: float) -> int:
        return round(self.inner.fare(km, minutes) * self.multiplier)


@dataclass(frozen=True)
class SurgePolicy:
    cap: float = 2.0

    def multiplier(self, demand: int, supply: int) -> float:
        return min(self.cap, max(1.0, round(demand / max(supply, 1), 1)))

# part 4: the service: request, start, complete, cancel
class NoDriverNearby(Exception): ...


@dataclass
class RideService:
    index: DriverIndex
    matcher: MatchingStrategy
    pricing: FareStrategy
    surge: SurgePolicy = field(default_factory=SurgePolicy)
    radius_km: float = 3.0
    window_s: float = 600.0  # demand = requests nearby in the last 10 minutes
    free_cancel_s: float = 120.0
    cancel_fee: int = 50
    trips: dict[str, Trip] = field(default_factory=dict)
    _active: dict[str, str] = field(default_factory=dict)  # rider -> open trip
    _recent: deque = field(default_factory=deque)  # (time, pickup) of recent requests
    _ids: itertools.count = field(default_factory=lambda: itertools.count(1))
    _mu: threading.Lock = field(default_factory=threading.Lock)

    def request(self, rider: str, pickup: Location, drop: Location, now: float) -> Trip:
        with self._mu:
            if rider in self._active:
                raise ValueError(f"{rider} already has trip {self._active[rider]}")
            self._recent.append((now, pickup))
            while self._recent[0][0] < now - self.window_s:
                self._recent.popleft()
            demand = sum(1 for _, p in self._recent if p.km_to(pickup) <= self.radius_km)
            trip = Trip(f"T{next(self._ids):04d}", rider, pickup, drop)
            self.trips[trip.trip_id] = trip
            self._active[rider] = trip.trip_id
        candidates = self.index.nearby(pickup, self.radius_km)
        trip.surge = self.surge.multiplier(demand, len(candidates))
        for d in self.matcher.rank(pickup, candidates):
            if d.try_reserve(trip.trip_id):  # lost the race for this driver? try the next one
                trip.driver_id, trip.accepted_at = d.driver_id, now
                trip.move(TripStatus.ACCEPTED)
                return trip
        trip.move(TripStatus.CANCELLED)
        self._close(trip)
        raise NoDriverNearby(f"no driver within {self.radius_km} km")

    def start(self, trip_id: str) -> None:
        self.trips[trip_id].move(TripStatus.STARTED)

    def complete(self, trip_id: str, km: float, minutes: float) -> Trip:
        trip = self.trips[trip_id]
        trip.move(TripStatus.COMPLETED)
        trip.fare = SurgeFare(self.pricing, trip.surge).fare(km, minutes)
        self.index.update_location(trip.driver_id, trip.drop)  # the driver is where the rider got off
        self.index.get(trip.driver_id).release()
        self._close(trip)
        return trip

    def cancel(self, trip_id: str, now: float) -> int:
        trip = self.trips[trip_id]
        trip.move(TripStatus.CANCELLED)  # raises once the trip has started
        late = trip.accepted_at is not None and now - trip.accepted_at > self.free_cancel_s
        trip.fare = self.cancel_fee if late else 0
        if trip.driver_id:
            self.index.get(trip.driver_id).release()
        self._close(trip)
        return trip.fare

    def _close(self, trip: Trip) -> None:
        with self._mu:
            self._active.pop(trip.rider, None)

# part 5: one ordinary ride, then twenty riders and three drivers outside a stadium
if __name__ == "__main__":
    index = DriverIndex()
    for did, name, x, y in [("D1", "Ravi", 0.5, 0.2), ("D2", "Sana", 2.5, 1.0), ("D3", "Imran", 10.2, 10.1),
                            ("D4", "Meera", 10.0, 10.6), ("D5", "Arjun", 11.1, 9.8)]:
        index.add(Driver(did, name, Location(x, y)))
    svc = RideService(index, NearestFirst(), StandardFare())

    t = svc.request("asha", Location(0, 0), Location(6, 8), now=0)
    print(t.trip_id, t.status.value, "driver", t.driver_id, "surge", t.surge)  # D1 is 0.54 km away
    svc.start(t.trip_id)
    svc.complete(t.trip_id, km=10, minutes=25)
    print(t.trip_id, t.status.value, "fare", t.fare)  # 50 + 12*10 + 2*25 = 220
    try:
        svc.cancel(t.trip_id, now=1600)
    except InvalidTransition as e:
        print("rejected:", e)

    stadium, won, refused = Location(10, 10), [], []

    def ride(i: int) -> None:
        try:
            won.append(svc.request(f"fan{i}", stadium, Location(10.5, 10.5), now=2000))
        except NoDriverNearby:
            refused.append(i)

    threads = [threading.Thread(target=ride, args=(i,)) for i in range(20)]
    for th in threads:
        th.start()
    for th in threads:
        th.join()
    drivers = sorted(tr.driver_id for tr in won)
    print("matched", len(won), "| no driver", len(refused), "| drivers", drivers, "| double-booked", len(drivers) - len(set(drivers)))

    by_driver = {tr.driver_id: tr for tr in won}
    print("cancel after 200 s, fee", svc.cancel(by_driver["D4"].trip_id, now=2200), "| D4", index.get("D4").status.value)
    svc.start(by_driver["D3"].trip_id)
    svc.complete(by_driver["D3"].trip_id, km=1, minutes=4)  # D3 drops its fan 0.7 km from the gate

    late = svc.request("neha", stadium, Location(13, 14), now=2400)  # 21 requests nearby, 2 free drivers
    print(late.trip_id, "driver", late.driver_id, "surge", late.surge)
    svc.start(late.trip_id)
    print("fare", svc.complete(late.trip_id, km=5, minutes=15).fare)  # (50 + 60 + 30) * 2.0 = 280

T0001 accepted driver D1 surge 1.0
T0001 completed fare 220
rejected: T0001: completed -> cancelled
matched 3 | no driver 17 | drivers ['D3', 'D4', 'D5'] | double-booked 0
cancel after 200 s, fee 50 | D4 available
T0022 driver D4 surge 2.0
fare 280