"""
Problem definition for the DVRPTW-TA (Dynamic Vehicle Routing Problem
with Time Windows and Traffic Awareness).

Defines Route, Solution, and ProblemContext classes aligned with
the paper's formulation (§3.1).
"""

import copy
import hashlib
from typing import Dict, List, Optional, Tuple
import numpy as np


class Route:
    """A single vehicle's delivery route: sequence of customer nodes.

    Format: [0, c1, c2, ..., ck, 0] where 0 is the depot.
    """

    def __init__(self, vehicle_id: int, nodes: Optional[List[int]] = None):
        self.vehicle_id = vehicle_id
        self.nodes: List[int] = nodes if nodes is not None else [0, 0]

    @property
    def customers(self) -> List[int]:
        """Return customer nodes only (excluding depot(s))."""
        return [n for n in self.nodes if n != 0]

    @property
    def n_customers(self) -> int:
        return len(self.customers)

    def total_demand(self, customers: Dict[int, "Customer"]) -> float:
        """Sum of demand for all customers on this route."""
        return sum(customers[n].demand_kg for n in self.customers if n in customers)

    def is_capacity_feasible(self, capacity: float, customers: Dict[int, "Customer"]) -> bool:
        return self.total_demand(customers) <= capacity

    def insert(self, customer_id: int, position: int):
        """Insert a customer at the given position (after depot start)."""
        # position is 1-based index in the customer sequence
        assert 1 <= position <= len(self.customers) + 1
        # Find insertion point in self.nodes (skip leading depot)
        insert_at = position
        self.nodes.insert(insert_at, customer_id)

    def remove(self, customer_id: int) -> bool:
        """Remove a customer from the route. Returns True if found."""
        if customer_id in self.nodes:
            self.nodes.remove(customer_id)
            return True
        return False

    def copy(self) -> "Route":
        return Route(self.vehicle_id, list(self.nodes))

    def __repr__(self) -> str:
        return f"Route(v{self.vehicle_id}: {self.nodes})"


class Solution:
    """A complete routing solution with m routes (one per vehicle)."""

    def __init__(self, n_vehicles: int):
        self.n_vehicles = n_vehicles
        self.routes: List[Route] = [Route(k) for k in range(n_vehicles)]

    def copy(self) -> "Solution":
        s = Solution(self.n_vehicles)
        s.routes = [r.copy() for r in self.routes]
        return s

    def find_route(self, customer_id: int) -> Optional[int]:
        """Find which vehicle (route index) serves a customer."""
        for k, route in enumerate(self.routes):
            if customer_id in route.nodes:
                return k
        return None

    def find_position(self, customer_id: int) -> Optional[Tuple[int, int]]:
        """Find (vehicle_idx, position_in_nodes) for a customer."""
        for k, route in enumerate(self.routes):
            try:
                pos = route.nodes.index(customer_id)
                return (k, pos)
            except ValueError:
                continue
        return None

    def unassigned_customers(self, all_customer_ids: set) -> set:
        """Return set of customer IDs not assigned to any route."""
        assigned = set()
        for route in self.routes:
            assigned.update(route.customers)
        return all_customer_ids - assigned

    def get_assignments(self) -> Dict[int, int]:
        """Return {customer_id: vehicle_id} for all assigned customers."""
        result = {}
        for k, route in enumerate(self.routes):
            for c in route.customers:
                result[c] = k
        return result

    def total_customers(self) -> int:
        return sum(len(r.customers) for r in self.routes)

    def __repr__(self) -> str:
        parts = [f"Solution({self.n_vehicles} vehicles, {self.total_customers()} customers)"]
        for r in self.routes:
            parts.append(f"  {r}")
        return "\n".join(parts)


class ProblemContext:
    """Bundles all problem instance data for efficient access."""

    def __init__(
        self,
        customers: Dict[int, "Customer"],
        depot: "Depot",
        traffic: "TrafficData",
        n_vehicles: int,
        vehicle_capacity: float,
        op_start: float = 360.0,
        op_end: float = 1080.0,
        n_intervals: int = 12,
    ):
        self.customers = customers
        self.depot = depot
        self.traffic = traffic
        self.n_vehicles = n_vehicles
        self.vehicle_capacity = vehicle_capacity
        self.op_start = op_start
        self.op_end = op_end
        self.n_intervals = n_intervals
        self.n_nodes = 1 + len(customers)

    @property
    def customer_ids(self) -> set:
        return set(self.customers.keys())

    @property
    def interval_duration(self) -> float:
        return (self.op_end - self.op_start) / self.n_intervals

    def time_to_interval(self, minutes: float) -> int:
        """Map time to interval h.

        The paper discretizes traffic feeds for 6:00-18:00. Customer windows may
        extend to 20:00, so departures after the final traffic bucket use
        hold-last-value semantics instead of extrapolating unobserved traffic.
        """
        minutes = max(self.op_start, min(minutes, self.op_end - 1))
        return int((minutes - self.op_start) / self.interval_duration)

    def get_travel_time(self, i: int, j: int, depart_minutes: float) -> float:
        """Get time-dependent travel time for arc (i,j) at departure time."""
        if i == j:
            return 0.0
        h = self.time_to_interval(depart_minutes)
        val = self.traffic.travel_time[i, j, h]
        return val if not np.isinf(val) else np.inf

    def get_congestion_penalty(self, i: int, j: int, depart_minutes: float) -> float:
        """Get congestion penalty ρ_ij(T_i) from the traffic tensor."""
        if i == j:
            return 0.0
        h = self.time_to_interval(depart_minutes)
        return self.traffic.congestion_penalty[i, j, h]

    def get_risk_adjusted_time(self, i: int, j: int, depart_minutes: float,
                                 beta: float = 0.3) -> float:
        """Get risk-adjusted travel time: t'_ij = t_ij + β·η_ij (Eq.10)."""
        t = self.get_travel_time(i, j, depart_minutes)
        h = self.time_to_interval(depart_minutes)
        eta = self.traffic.uncertainty[i, j, h]
        return t + beta * eta