"""
Cost functions for the DVRPTW-TA problem.

Implements the composite objective (Eq.1) and insertion cost evaluation
(Eq.16-17) from the paper.
"""

from typing import Dict, List, Tuple
import numpy as np


class CostCalculator:
    """Computes costs for routing solutions under traffic-aware conditions."""

    def __init__(
        self,
        lambda_lateness: float = 1.0,
        lambda_congestion: float = 1.0,
        lambda_stability: float = 0.3,
    ):
        self.lambda_lateness = lambda_lateness
        self.lambda_congestion = lambda_congestion
        self.lambda_stability = lambda_stability

    # ------------------------------------------------------------------
    # Route-level time propagation
    # ------------------------------------------------------------------

    def propagate_route(
        self,
        route_nodes: List[int],
        problem_ctx: "ProblemContext",
        depot_start_time: float = 360.0,
    ) -> dict:
        """Compute all time values along a route.

        Implements Eq.5:
          A_j = T_i + t_ij(T_i)
          S_j = max(A_j, e_j)
          T_j = S_j + s_j
          δ_j = max(0, S_j - l_j)

        Returns dict with arrival_times, service_starts, departures,
        delays, wait_times, and totals.
        """
        customers = problem_ctx.customers
        n = len(route_nodes)
        arrivals = [0.0] * n
        service_starts = [0.0] * n
        departures = [0.0] * n
        delays = [0.0] * n
        waits = [0.0] * n
        arc_travel_time = 0.0
        congestion_exposure = 0.0

        # First node (depot)
        departures[0] = depot_start_time
        service_starts[0] = depot_start_time
        arrivals[0] = depot_start_time

        for idx in range(1, n):
            i = route_nodes[idx - 1]
            j = route_nodes[idx]

            # Arrival at j (Eq.5). Eq.1's travel term is the sum of arc
            # traversal times only; waiting and service duration are temporal
            # propagation state, not travel cost.
            travel_time = problem_ctx.get_travel_time(i, j, departures[idx - 1])
            arc_travel_time += travel_time
            arrivals[idx] = departures[idx - 1] + travel_time

            # Accumulate congestion exposure
            congestion_exposure += problem_ctx.get_congestion_penalty(
                i, j, departures[idx - 1]
            )

            # Service start (wait if early)
            if j == 0:
                # Return to depot: no service time, no time window
                service_starts[idx] = arrivals[idx]
                delays[idx] = 0.0
                waits[idx] = 0.0
            else:
                cust = customers.get(j)
                if cust is None:
                    # Should not happen in valid solutions
                    service_starts[idx] = arrivals[idx]
                    delays[idx] = 0.0
                    waits[idx] = 0.0
                else:
                    # Soft time window: wait if early, record delay if late
                    service_starts[idx] = max(arrivals[idx], cust.earliest_time_min)
                    waits[idx] = max(0.0, cust.earliest_time_min - arrivals[idx])
                    delays[idx] = max(0.0, service_starts[idx] - cust.latest_time_min)

            # Departure (Eq.5)
            if j == 0:
                departures[idx] = arrivals[idx]
            else:
                service_time = customers[j].service_time_min if j in customers else 0.0
                departures[idx] = service_starts[idx] + service_time

        return {
            "arrivals": arrivals,
            "service_starts": service_starts,
            "departures": departures,
            "delays": delays,
            "waits": waits,
            "total_travel_time": arc_travel_time,
            "route_duration": departures[-1] - depot_start_time,
            "total_delay": sum(delays),
            "total_wait": sum(waits),
            "congestion_exposure": congestion_exposure,
        }

    # ------------------------------------------------------------------
    # Composite cost (Eq.1)
    # ------------------------------------------------------------------

    def compute_total_cost(
        self,
        solution: "Solution",
        problem_ctx: "ProblemContext",
        previous_solution: "Solution" = None,
    ) -> float:
        """Compute composite objective value per Eq.1.

        Cost = Σ_k Σ_(i,j) [t_ij(T_i) + λ₂·ρ_ij(T_i)] + λ₁·Σ_j δ_j

        For RRD: also includes λ₃ × stability penalty.
        """
        total_travel_time = 0.0
        total_delay = 0.0
        total_congestion = 0.0

        for route in solution.routes:
            if len(route.nodes) < 2:
                continue
            result = self.propagate_route(route.nodes, problem_ctx)

            # Travel time: sum of all arc travel times
            total_travel_time += result["total_travel_time"]

            # Delay penalty: sum of lateness at all customers
            total_delay += result["total_delay"]

            # Congestion: sum of ρ_ij on traversed arcs
            total_congestion += result["congestion_exposure"]

        cost = (
            total_travel_time
            + self.lambda_lateness * total_delay
            + self.lambda_congestion * total_congestion
        )

        # Stability penalty (RRD only, when comparing to previous solution)
        if previous_solution is not None:
            stability = self._compute_stability(solution, previous_solution)
            cost += self.lambda_stability * stability

        return cost

    def compute_travel_time_cost(
        self, solution: "Solution", problem_ctx: "ProblemContext"
    ) -> float:
        total = 0.0
        for route in solution.routes:
            if len(route.nodes) < 2:
                continue
            for idx in range(1, len(route.nodes)):
                i, j = route.nodes[idx - 1], route.nodes[idx]
                depart = self._get_departure_at(route.nodes, idx - 1, problem_ctx)
                total += problem_ctx.get_travel_time(i, j, depart)
        return total

    def compute_lateness_penalty(
        self, solution: "Solution", problem_ctx: "ProblemContext"
    ) -> float:
        total = 0.0
        for route in solution.routes:
            result = self.propagate_route(route.nodes, problem_ctx)
            total += result["total_delay"]
        return self.lambda_lateness * total

    def compute_congestion_exposure(
        self, solution: "Solution", problem_ctx: "ProblemContext"
    ) -> float:
        total = 0.0
        for route in solution.routes:
            result = self.propagate_route(route.nodes, problem_ctx)
            total += result["congestion_exposure"]
        return total

    # ------------------------------------------------------------------
    # Insertion cost (Eq.16-17)
    # ------------------------------------------------------------------

    def compute_insertion_cost(
        self,
        route: "Route",
        customer_id: int,
        position: int,
        problem_ctx: "ProblemContext",
    ) -> float:
        """Compute marginal cost of inserting a customer at a position.

        Eq.16: Δf_ijk = t_ji(T_j) + s_i + t_ik(T_j + t_ji + s_i) - t_jk(T_j)
                       + λ₁·δ_i(T_i) + λ₂·ρ_ji(T_j)

        This computes only the local change; a full suffix propagation
        (Eq.17) would recompute downstream times, which is done in
        compute_insertion_cost_full().
        """
        nodes = route.nodes
        customers = problem_ctx.customers
        cust = customers[customer_id]

        # Find predecessor (j) and successor (k) in the route
        # Position is 1-based among customers
        if position <= len(route.customers):
            # Insert between existing nodes
            insert_idx = position  # position in nodes list
            j = nodes[insert_idx - 1]
            k = nodes[insert_idx]
        else:
            # Insert before returning to depot
            j = nodes[-2]  # last customer before depot
            k = nodes[-1]  # depot (0)

        # Get departure time from j
        depart_j = self._get_departure_at(nodes, nodes.index(j), problem_ctx)

        # Old cost: t_jk(T_j)
        old_cost = problem_ctx.get_travel_time(j, k, depart_j)

        # New cost: t_ji(T_j) + s_i + t_ik(T_j + t_ji + s_i)
        t_ji = problem_ctx.get_travel_time(j, customer_id, depart_j)
        arrival_i = depart_j + t_ji
        service_start_i = max(arrival_i, cust.earliest_time_min)
        t_ik = problem_ctx.get_travel_time(
            customer_id, k, service_start_i + cust.service_time_min
        )
        travel_component = t_ji + cust.service_time_min + t_ik - old_cost

        # Delay penalty at i (λ₁·δ_i)
        lateness_i = max(0.0, service_start_i - cust.latest_time_min)
        delay_penalty = self.lambda_lateness * lateness_i

        # Congestion penalty for new arc (j,i) (λ₂·ρ_ji)
        congestion_penalty = self.lambda_congestion * problem_ctx.get_congestion_penalty(
            j, customer_id, depart_j
        )

        return travel_component + delay_penalty + congestion_penalty

    def compute_insertion_cost_full(
        self,
        route: "Route",
        customer_id: int,
        position: int,
        problem_ctx: "ProblemContext",
    ) -> float:
        """Full insertion cost with suffix propagation (Eq.17).

        Evaluates the complete impact of insertion including all
        downstream time changes due to FIFO shift.
        """
        # Build candidate route with customer inserted
        candidate_nodes = list(route.nodes)
        # position is 1-based among customers
        insert_at = position
        candidate_nodes.insert(insert_at, customer_id)

        # Evaluate old route cost
        old_result = self.propagate_route(route.nodes, problem_ctx)
        old_cost = (
            old_result["total_travel_time"]
            + self.lambda_lateness * old_result["total_delay"]
            + self.lambda_congestion * old_result["congestion_exposure"]
        )

        # Evaluate new route cost (full propagation)
        new_result = self.propagate_route(candidate_nodes, problem_ctx)
        new_cost = (
            new_result["total_travel_time"]
            + self.lambda_lateness * new_result["total_delay"]
            + self.lambda_congestion * new_result["congestion_exposure"]
        )

        return new_cost - old_cost

    def find_best_insertion(
        self,
        route: "Route",
        customer_id: int,
        problem_ctx: "ProblemContext",
        use_full: bool = True,
    ) -> Tuple[int, float]:
        """Find best position to insert a customer into a route.

        Returns (best_position, best_cost).
        position is 1-based among customers (1 = first customer, n+1 = last).
        """
        n_cust = len(route.customers)
        best_pos = 1
        best_cost = float("inf")
        cost_fn = self.compute_insertion_cost_full if use_full else self.compute_insertion_cost

        for pos in range(1, n_cust + 2):
            cost = cost_fn(route, customer_id, pos, problem_ctx)
            if cost < best_cost:
                best_cost = cost
                best_pos = pos

        return best_pos, best_cost

    # ------------------------------------------------------------------
    # Stability and helpers
    # ------------------------------------------------------------------

    def _compute_stability(
        self, new_sol: "Solution", old_sol: "Solution"
    ) -> float:
        """Compute route stability penalty (Eq.41).

        Stability = Σ_k |R^a_k ∩ R^s_k| / |R^a_k ∪ R^s_k|
        Higher = more stable. We invert for penalty.
        """
        stability = 0.0
        for k in range(min(new_sol.n_vehicles, old_sol.n_vehicles)):
            new_set = set(new_sol.routes[k].customers)
            old_set = set(old_sol.routes[k].customers)
            union = new_set | old_set
            inter = new_set & old_set
            if len(union) > 0:
                stability += len(inter) / len(union)
        # Penalize route changes (invert: more stable = lower penalty)
        return (new_sol.n_vehicles - stability) * 100.0

    def _get_departure_at(
        self,
        nodes: List[int],
        idx: int,
        problem_ctx: "ProblemContext",
        start_time: float = 360.0,
    ) -> float:
        """Compute the departure time at a given position along a route."""
        result = self.propagate_route(nodes, problem_ctx, start_time)
        return result["departures"][idx]

    # ------------------------------------------------------------------
    # Solution quality metrics
    # ------------------------------------------------------------------

    def evaluate_solution(
        self,
        solution: "Solution",
        problem_ctx: "ProblemContext",
        include_penalties: bool = True,
    ) -> dict:
        """Comprehensive solution evaluation returning all metrics.

        Args:
            include_penalties: If False, total_cost = pure travel time only
                              (used for Static-VRPTW baseline which doesn't
                               account for congestion or time-dependent costs).
        """
        total_travel = 0.0
        total_delay = 0.0
        total_congestion = 0.0
        n_ontime = 0
        n_late = 0
        delays_list = []
        max_delay = 0.0
        avg_route_duration = 0.0

        all_cust_ids = problem_ctx.customer_ids
        assigned = solution.get_assignments()

        for route in solution.routes:
            if len(route.nodes) < 2:
                continue
            result = self.propagate_route(route.nodes, problem_ctx)

            total_travel += result["total_travel_time"]
            total_delay += result["total_delay"]
            total_congestion += result["congestion_exposure"]
            avg_route_duration += result["route_duration"]

            # Count on-time deliveries
            for idx, node in enumerate(route.nodes):
                if node == 0:
                    continue
                delay = result["delays"][idx]
                if delay <= 0:
                    n_ontime += 1
                else:
                    n_late += 1
                    delays_list.append(delay)
                    max_delay = max(max_delay, delay)

        n_total = len(all_cust_ids)
        otdr = n_ontime / n_total if n_total > 0 else 0.0
        avg_delay = np.mean(delays_list) if delays_list else 0.0
        avg_route_dur = (
            avg_route_duration / solution.n_vehicles if solution.n_vehicles > 0 else 0.0
        )

        if include_penalties:
            total_cost = (
                total_travel
                + self.lambda_lateness * total_delay
                + self.lambda_congestion * total_congestion
            )
        else:
            total_cost = total_travel

        return {
            "total_cost": total_cost,
            "travel_time_cost": total_travel,
            "delay_penalty": self.lambda_lateness * total_delay,
            "congestion_cost": total_congestion,
            "otdr": otdr,
            "avg_delay": avg_delay,
            "max_delay": max_delay,
            "late_customers": n_late,
            "ces": total_congestion,
            "avg_route_duration": avg_route_dur,
            "n_assigned": len(assigned),
            "n_total": n_total,
        }