Uber.

# geospatial/h3_location_index.py

"""
H3-Based Location Index for Uber

Uses Uber's open-source H3 hexagonal grid system for
efficient geospatial queries.
"""

from dataclasses import dataclass
from typing import List, Set, Optional
import h3
import redis


@dataclass
class DriverLocation:
    driver_id: str
    latitude: float
    longitude: float
    timestamp: float
    ride_types: List[str]
    rating: float
    is_available: bool


class H3LocationIndex:
    """
    Location index using H3 hexagonal grid.
    
    Key insight: Convert 2D lat/lng to 1D H3 cell ID
    for O(1) lookups instead of O(n) distance calculations.
    
    Applies concepts:
    - Week 1: Partitioning strategies
    - Week 4: Caching for low-latency lookups
    """
    
    # Resolution 9: ~0.1 km² hexagons (edge ~200 m), good for urban areas
    DEFAULT_RESOLUTION = 9
    
    def __init__(self, redis_client: redis.Redis):
        self.redis = redis_client
        self.resolution = self.DEFAULT_RESOLUTION
    
    def update_driver_location(self, driver: DriverLocation) -> None:
        """
        Update driver location in the index.
        
        Called every 4 seconds per driver.
        Must be extremely fast (< 10ms).
        """
        # Convert lat/lng to H3 cell
        cell_id = h3.latlng_to_cell(
            driver.latitude, 
            driver.longitude, 
            self.resolution
        )
        
        # Remove from old cell (if moved)
        old_cell = self.redis.hget(f"driver:{driver.driver_id}", "cell")
        if old_cell and old_cell != cell_id:
            self.redis.srem(f"cell:{old_cell}", driver.driver_id)
        
        # Add to new cell
        self.redis.sadd(f"cell:{cell_id}", driver.driver_id)
        
        # Store driver details
        self.redis.hset(f"driver:{driver.driver_id}", mapping={
            "cell": cell_id,
            "lat": driver.latitude,
            "lng": driver.longitude,
            "ts": driver.timestamp,
            "available": int(driver.is_available),
            "rating": driver.rating
        })
        
        # Set TTL for automatic cleanup if driver goes offline
        self.redis.expire(f"driver:{driver.driver_id}", 30)
    
    def find_nearby_drivers(
        self, 
        latitude: float, 
        longitude: float,
        radius_km: float = 2.0,
        ride_type: str = "UberX",
        limit: int = 10
    ) -> List[DriverLocation]:
        """
        Find available drivers near a location.
        
        Uses H3 k-ring to get nearby cells, then fetches
        drivers from those cells.
        """
        # Get center cell
        center_cell = h3.latlng_to_cell(
            latitude, longitude, self.resolution
        )
        
        # Widen ring by ring (k = 1, then 2, ...) until there are enough
        # candidates, up to the radius. At resolution 9 the edge is ~200 m,
        # so adjacent cell centres are ~0.35 km apart (sqrt(3) x edge):
        # 2 km allows up to k = 6
        max_k = int(radius_km / 0.35) + 1
        
        candidates = []
        searched = set()
        for k in range(1, max_k + 1):
            # Only the cells this ring adds (grid_disk includes the centre)
            ring_cells = set(h3.grid_disk(center_cell, k)) - searched
            searched |= ring_cells
            for cell_id in ring_cells:
                driver_ids = self.redis.smembers(f"cell:{cell_id}")
                for driver_id in driver_ids:
                    driver_data = self.redis.hgetall(f"driver:{driver_id}")
                    # Hash values come back as strings: "0" is truthy, so
                    # compare explicitly (decode_responses=True assumed)
                    if driver_data and driver_data.get("available") == "1":
                        candidates.append(self._to_driver_location(
                            driver_id, driver_data
                        ))
            if len(candidates) >= limit:
                break   # enough drivers: stop widening
        
        # Sort by distance and return top matches
        candidates.sort(
            key=lambda d: self._haversine_distance(
                latitude, longitude, d.latitude, d.longitude
            )
        )
        
        return candidates[:limit]
    
    def get_drivers_in_cell(self, cell_id: str) -> Set[str]:
        """Get all driver IDs in a specific H3 cell."""
        return self.redis.smembers(f"cell:{cell_id}")
    
    def _haversine_distance(
        self, lat1: float, lng1: float, 
        lat2: float, lng2: float
    ) -> float:
        """Calculate distance between two points in km."""
        from math import radians, sin, cos, sqrt, atan2
        
        R = 6371  # Earth's radius in km
        
        lat1, lng1, lat2, lng2 = map(
            radians, [lat1, lng1, lat2, lng2]
        )
        
        dlat = lat2 - lat1
        dlng = lng2 - lng1
        
        a = sin(dlat/2)**2 + cos(lat1) * cos(lat2) * sin(dlng/2)**2
        c = 2 * atan2(sqrt(a), sqrt(1-a))
        
        return R * c

# pricing/surge_service.py

"""
Dynamic Surge Pricing Service

Calculates real-time surge multipliers based on
supply-demand imbalance in each H3 cell.
"""

from dataclasses import dataclass
from typing import Dict, Optional
from datetime import datetime, timedelta
import h3


@dataclass
class SurgeConfig:
    min_surge: float = 1.0
    max_surge: float = 5.0
    update_interval_seconds: int = 30
    smoothing_factor: float = 0.3  # For exponential moving average
    demand_window_minutes: int = 5


@dataclass
class CellSurge:
    cell_id: str
    surge_multiplier: float
    demand_count: int
    supply_count: int
    updated_at: datetime


class SurgePricingService:
    """
    Calculate surge pricing per H3 cell.
    
    Key principles:
    - Higher demand → higher surge
    - More drivers → lower surge
    - Smooth transitions (no sudden jumps)
    - Respect max cap (emergencies, PR concerns)
    
    Applies concepts:
    - Week 5: Eventual consistency (surge propagates)
    - Week 8: Real-time analytics pipeline
    """
    
    def __init__(
        self, 
        config: SurgeConfig,
        location_index,
        demand_tracker,
        cache
    ):
        self.config = config
        self.location_index = location_index
        self.demand_tracker = demand_tracker
        self.cache = cache
        
        # Surge curve parameters (tuned via ML)
        self.surge_curve = [
            (1.0, 1.0),   # ratio 1.0 → surge 1.0
            (1.5, 1.25),  # ratio 1.5 → surge 1.25
            (2.0, 1.5),   # ratio 2.0 → surge 1.5
            (3.0, 2.0),   # ratio 3.0 → surge 2.0
            (5.0, 3.0),   # ratio 5.0 → surge 3.0
            (10.0, 5.0),  # ratio 10.0 → surge 5.0 (cap)
        ]
    
    def get_surge_for_location(
        self, 
        latitude: float, 
        longitude: float
    ) -> float:
        """Get current surge multiplier for a location."""
        cell_id = h3.latlng_to_cell(latitude, longitude, 9)
        
        # Check cache first
        cached = self.cache.get(f"surge:{cell_id}")
        if cached:
            return float(cached)
        
        # Calculate fresh surge
        surge = self._calculate_cell_surge(cell_id)
        
        # Cache for 30 seconds
        self.cache.setex(f"surge:{cell_id}", 30, surge)
        
        return surge
    
    def _calculate_cell_surge(self, cell_id: str) -> float:
        """Calculate surge for a specific H3 cell."""
        # Get demand (requests in last 5 minutes)
        demand = self.demand_tracker.get_request_count(
            cell_id, 
            window_minutes=self.config.demand_window_minutes
        )
        
        # Get supply (available drivers in cell + neighbors)
        nearby_cells = h3.grid_disk(cell_id, 1)  # Cell + 6 neighbors
        supply = sum(
            len(self.location_index.get_drivers_in_cell(c))
            for c in nearby_cells
        )
        
        # Prevent division by zero
        if supply == 0:
            return self.config.max_surge if demand > 0 else self.config.min_surge
        
        # Calculate demand/supply ratio
        ratio = demand / supply
        
        # Map ratio to surge using curve
        raw_surge = self._ratio_to_surge(ratio)
        
        # Apply smoothing (don't change too fast). The 30 s "surge:"
        # entry has already expired when we get here, so keep the last
        # published value under a longer-lived key.
        previous_surge = self.cache.get(f"surge_last:{cell_id}")
        if previous_surge:
            raw_surge = self._smooth_surge(
                float(previous_surge), 
                raw_surge
            )
        
        # Clamp to bounds
        surge = max(
            self.config.min_surge,
            min(self.config.max_surge, raw_surge)
        )
        self.cache.setex(f"surge_last:{cell_id}", 3600, surge)
        return surge
    
    def _ratio_to_surge(self, ratio: float) -> float:
        """Convert demand/supply ratio to surge multiplier."""
        # Below the first point (ratio < 1.0): no surge
        if ratio <= self.surge_curve[0][0]:
            return self.surge_curve[0][1]
        
        # Linear interpolation between curve points
        for i in range(len(self.surge_curve) - 1):
            r1, s1 = self.surge_curve[i]
            r2, s2 = self.surge_curve[i + 1]
            
            if r1 <= ratio <= r2:
                # Linear interpolation
                t = (ratio - r1) / (r2 - r1)
                return s1 + t * (s2 - s1)
        
        # Above max ratio
        return self.surge_curve[-1][1]
    
    def _smooth_surge(
        self, 
        previous: float, 
        current: float
    ) -> float:
        """Apply exponential moving average for smooth transitions."""
        alpha = self.config.smoothing_factor
        return alpha * current + (1 - alpha) * previous
    
    def calculate_fare(
        self,
        distance_km: float,
        duration_minutes: float,
        surge: float,
        ride_type: str = "UberX"
    ) -> float:
        """Calculate total fare with surge."""
        # Base rates (vary by city/ride type)
        rates = {
            "UberX": {"base": 2.50, "per_km": 1.20, "per_min": 0.25},
            "UberXL": {"base": 4.00, "per_km": 1.80, "per_min": 0.35},
            "UberBlack": {"base": 8.00, "per_km": 3.00, "per_min": 0.50},
        }
        
        rate = rates.get(ride_type, rates["UberX"])
        
        base_fare = (
            rate["base"] + 
            rate["per_km"] * distance_km + 
            rate["per_min"] * duration_minutes
        )
        
        # Apply surge
        surged_fare = base_fare * surge
        
        # Minimum fare
        return max(surged_fare, 5.00)

# matching/dispatch_service.py

"""
Rider-Driver Matching Service

Finds optimal driver for each ride request using
weighted scoring and ETA estimation.
"""

from dataclasses import dataclass
from typing import List, Optional, Tuple
from enum import Enum
import asyncio


class MatchResult(Enum):
    MATCHED = "matched"
    NO_DRIVERS = "no_drivers"
    ALL_DECLINED = "all_declined"
    TIMEOUT = "timeout"


@dataclass
class MatchRequest:
    request_id: str
    rider_id: str
    pickup_lat: float
    pickup_lng: float
    destination_lat: float
    destination_lng: float
    ride_type: str


@dataclass
class DriverCandidate:
    driver_id: str
    latitude: float
    longitude: float
    rating: float
    acceptance_rate: float
    wait_time_seconds: int
    eta_seconds: int
    score: float = 0.0


class MatchingService:
    """
    Match riders with optimal drivers.
    
    Uses weighted scoring to balance multiple factors:
    - ETA (user experience)
    - Driver rating (quality)
    - Acceptance rate (reliability)
    - Wait time (fairness to drivers)
    
    Applies concepts:
    - Week 1: Partitioning (H3-based driver lookup)
    - Week 2: Timeouts and retries
    """
    
    # Scoring weights (tuned via ML/experimentation)
    WEIGHTS = {
        "eta": 0.4,
        "rating": 0.2,
        "acceptance": 0.2,
        "wait_time": 0.2
    }
    
    DRIVER_RESPONSE_TIMEOUT = 15  # seconds
    MAX_CANDIDATES = 10
    
    def __init__(
        self, 
        location_index,
        eta_service,
        driver_service,
        notification_service
    ):
        self.location_index = location_index
        self.eta_service = eta_service
        self.driver_service = driver_service
        self.notifications = notification_service
    
    async def find_match(
        self, 
        request: MatchRequest
    ) -> Tuple[MatchResult, Optional[str]]:
        """
        Find and dispatch to optimal driver.
        
        Returns (result, driver_id or None)
        """
        # Step 1: Find nearby candidates
        candidates = await self._get_candidates(request)
        
        if not candidates:
            return (MatchResult.NO_DRIVERS, None)
        
        # Step 2: Score and rank candidates
        scored = await self._score_candidates(request, candidates)
        scored.sort(key=lambda c: c.score, reverse=True)
        
        # Step 3: Dispatch to drivers in order
        for candidate in scored[:self.MAX_CANDIDATES]:
            accepted = await self._dispatch_to_driver(
                request, candidate
            )
            if accepted:
                return (MatchResult.MATCHED, candidate.driver_id)
        
        return (MatchResult.ALL_DECLINED, None)
    
    async def _get_candidates(
        self, 
        request: MatchRequest
    ) -> List[DriverCandidate]:
        """Get available drivers near pickup location."""
        drivers = self.location_index.find_nearby_drivers(
            request.pickup_lat,
            request.pickup_lng,
            radius_km=3.0,
            ride_type=request.ride_type,
            limit=50
        )
        
        candidates = []
        for driver in drivers:
            if not driver.is_available:
                continue
            
            # Get driver stats
            stats = await self.driver_service.get_stats(driver.driver_id)
            
            # Calculate ETA
            eta = await self.eta_service.calculate(
                driver.latitude, driver.longitude,
                request.pickup_lat, request.pickup_lng
            )
            
            candidates.append(DriverCandidate(
                driver_id=driver.driver_id,
                latitude=driver.latitude,
                longitude=driver.longitude,
                rating=driver.rating,
                acceptance_rate=stats.acceptance_rate,
                wait_time_seconds=stats.wait_time_seconds,
                eta_seconds=eta
            ))
        
        return candidates
    
    async def _score_candidates(
        self,
        request: MatchRequest,
        candidates: List[DriverCandidate]
    ) -> List[DriverCandidate]:
        """Calculate match score for each candidate."""
        max_eta = max(c.eta_seconds for c in candidates) or 1
        
        for candidate in candidates:
            # Normalize scores to 0-1 range
            eta_score = 1 - (candidate.eta_seconds / max_eta)
            rating_score = (candidate.rating - 4.0) / 1.0
            acceptance_score = candidate.acceptance_rate
            wait_score = min(candidate.wait_time_seconds / 1800, 1.0)
            
            # Weighted sum
            candidate.score = (
                self.WEIGHTS["eta"] * eta_score +
                self.WEIGHTS["rating"] * max(0, rating_score) +
                self.WEIGHTS["acceptance"] * acceptance_score +
                self.WEIGHTS["wait_time"] * wait_score
            )
        
        return candidates
    
    async def _dispatch_to_driver(
        self,
        request: MatchRequest,
        candidate: DriverCandidate
    ) -> bool:
        """Send ride request to driver and wait for response."""
        # Mark driver as unavailable (prevent double-dispatch)
        await self.driver_service.set_pending(candidate.driver_id)
        
        try:
            # Send push notification
            await self.notifications.send_ride_request(
                driver_id=candidate.driver_id,
                request_id=request.request_id,
                pickup_address="123 Main St",  # Would resolve address
                eta_seconds=candidate.eta_seconds
            )
            
            # Wait for response with timeout
            response = await asyncio.wait_for(
                self.driver_service.wait_for_response(
                    candidate.driver_id, 
                    request.request_id
                ),
                timeout=self.DRIVER_RESPONSE_TIMEOUT
            )
            
            return response == "accepted"
            
        except asyncio.TimeoutError:
            return False
            
        finally:
            # Release driver if not matched
            await self.driver_service.clear_pending(candidate.driver_id)

# routing/eta_service.py

"""
ETA Prediction Service

Combines routing engine with ML model for accurate
arrival time predictions.
"""

from dataclasses import dataclass
from typing import List, Optional, Tuple
from datetime import datetime


@dataclass
class ETAResult:
    eta_seconds: int
    distance_meters: int
    confidence: float
    route_polyline: str


class ETAService:
    """
    Calculate estimated time of arrival.
    
    Combines:
    - Graph-based shortest path (baseline)
    - Real-time traffic data
    - ML model adjustments (DeepETA)
    - Historical patterns
    
    Applies concepts:
    - Week 8: ML pipeline integration
    - Week 4: Caching for repeated queries
    """
    
    def __init__(
        self, 
        routing_engine,
        traffic_service,
        ml_model,
        cache
    ):
        self.routing = routing_engine
        self.traffic = traffic_service
        self.model = ml_model
        self.cache = cache
    
    async def calculate(
        self,
        origin_lat: float,
        origin_lng: float,
        dest_lat: float,
        dest_lng: float,
        departure_time: Optional[datetime] = None
    ) -> ETAResult:
        """Calculate ETA from origin to destination."""
        departure_time = departure_time or datetime.utcnow()
        
        # Check cache for recent identical query
        cache_key = self._make_cache_key(
            origin_lat, origin_lng, dest_lat, dest_lng
        )
        cached = self.cache.get(cache_key)
        if cached:
            return cached
        
        # Step 1: Get baseline route from routing engine
        route = await self.routing.get_route(
            origin_lat, origin_lng,
            dest_lat, dest_lng
        )
        
        # Step 2: Get current traffic conditions
        traffic_factor = await self.traffic.get_factor(
            route.segments,
            departure_time
        )
        
        # Step 3: ML model adjustment
        features = self._extract_features(
            route, traffic_factor, departure_time
        )
        ml_adjustment = self.model.predict(features)
        
        # Step 4: Combine estimates
        base_eta = route.duration_seconds
        traffic_eta = base_eta * traffic_factor
        final_eta = int(traffic_eta * ml_adjustment)
        
        result = ETAResult(
            eta_seconds=final_eta,
            distance_meters=route.distance_meters,
            confidence=0.95 if traffic_factor < 1.5 else 0.85,
            route_polyline=route.polyline
        )
        
        # Cache for 60 seconds
        self.cache.setex(cache_key, 60, result)
        
        return result
    
    def _extract_features(
        self,
        route,
        traffic_factor: float,
        departure_time: datetime
    ) -> dict:
        """Extract features for ML model."""
        return {
            "distance_km": route.distance_meters / 1000,
            "base_duration_min": route.duration_seconds / 60,
            "traffic_factor": traffic_factor,
            "hour_of_day": departure_time.hour,
            "day_of_week": departure_time.weekday(),
            "is_rush_hour": self._is_rush_hour(departure_time),
            "num_turns": route.num_turns,
            "num_traffic_lights": route.num_traffic_lights,
            "highway_fraction": route.highway_fraction,
        }
    
    def _is_rush_hour(self, dt: datetime) -> bool:
        """Check if time is during typical rush hour."""
        hour = dt.hour
        is_weekday = dt.weekday() < 5
        morning_rush = 7 <= hour <= 9
        evening_rush = 17 <= hour <= 19
        return is_weekday and (morning_rush or evening_rush)