from __future__ import annotations

from decimal import ROUND_HALF_UP, Decimal
from typing import List, Sequence, Tuple

from env.models import Order
from env.utils import manhattan


def clamp(value: float, low: float, high: float) -> float:
    return max(low, min(high, value))


def round_score(value: float, decimals: int = 4) -> float:
    """Round with HALF_UP behavior for stable scoring outputs."""
    safe = clamp(value, 0.0, 1.0)

    if decimals > 0:
        epsilon = float(Decimal("1") / (Decimal(10) ** decimals))
        if safe <= 0.0:
            safe = epsilon
        elif safe >= 1.0:
            safe = 1.0 - epsilon

    quantizer = Decimal("1") if decimals <= 0 else Decimal(f"1.{'0' * decimals}")
    return float(Decimal(str(safe)).quantize(quantizer, rounding=ROUND_HALF_UP))


def completion_rate(delivered_weight: float, total_weight: float) -> float:
    if total_weight <= 0:
        # Neutral when no work was available.
        return 0.5
    return clamp(delivered_weight / total_weight, 0.0, 1.0)


def efficiency_score(steps_taken: int, optimal_steps: int) -> float:
    """Compute bounded efficiency ratio from optimal-step lower bound."""
    if optimal_steps <= 0:
        return 1.0 if steps_taken <= 0 else 0.0
    if steps_taken <= 0:
        return 0.0
    return clamp(optimal_steps / steps_taken, 0.0, 1.0)


def penalty_score(
    invalid_actions: int,
    delay_events: int,
    battery_depletion_events: int,
    no_progress_events: int,
    total_steps: int,
) -> float:
    """Map penalty events to a [0, 1] quality score where 1 is best."""
    denominator = max(total_steps, 1)

    weighted_penalties = (
        invalid_actions * 1.0
        + delay_events * 0.6
        + battery_depletion_events * 2.0
        + no_progress_events * 0.25
    )
    base_rate = weighted_penalties / denominator

    # Add deterministic duration pressure for sustained non-progress behavior so
    # longer stalled episodes are graded worse than short stalls.
    stalled_fraction = no_progress_events / denominator
    stall_duration_pressure = stalled_fraction * min(0.35, denominator * 0.02)

    rate = base_rate + stall_duration_pressure
    return clamp(1.0 - rate, 0.0, 1.0)


def optimal_steps_single_agent(
    orders: Sequence[Order],
    start_location: Tuple[int, int] = (0, 0),
    max_exact_orders: int = 10,
) -> int:
    """
    Estimate minimal steps for one-agent execution with one active order at a time.

    For small order counts this is exact DP over completion order and chosen drop location.
    For larger sets it uses a deterministic nearest-first fallback to stay efficient.
    """
    order_list: List[Order] = list(orders)
    count = len(order_list)
    if count == 0:
        return 0

    if count > max_exact_orders:
        return _heuristic_optimal_steps(order_list, start_location)

    pickups = [(o.pickup.x, o.pickup.y) for o in order_list]
    drop_choices = [
        [(d.x, d.y) for d in (o.delivery_locations or [o.dropoff])]
        for o in order_list
    ]

    max_mask = 1 << count
    inf = float("inf")

    dp = [
        [
            [inf for _ in drop_choices[i]]
            for i in range(count)
        ]
        for _ in range(max_mask)
    ]

    for i in range(count):
        for k, drop in enumerate(drop_choices[i]):
            travel = manhattan(start_location, pickups[i]) + manhattan(pickups[i], drop)
            dp[1 << i][i][k] = float(travel)

    for mask in range(max_mask):
        for last in range(count):
            for last_k, current in enumerate(dp[mask][last]):
                if current == inf:
                    continue
                last_drop = drop_choices[last][last_k]

                for nxt in range(count):
                    if mask & (1 << nxt):
                        continue

                    transition_to_pickup = manhattan(last_drop, pickups[nxt])
                    next_mask = mask | (1 << nxt)

                    for nxt_k, nxt_drop in enumerate(drop_choices[nxt]):
                        candidate = (
                            current
                            + transition_to_pickup
                            + manhattan(pickups[nxt], nxt_drop)
                        )
                        if candidate < dp[next_mask][nxt][nxt_k]:
                            dp[next_mask][nxt][nxt_k] = candidate

    full_mask = max_mask - 1
    optimal_travel = min(
        dp[full_mask][last][last_k]
        for last in range(count)
        for last_k in range(len(drop_choices[last]))
    )

    if optimal_travel == inf:
        return _heuristic_optimal_steps(order_list, start_location)

    # Minimal action overhead per order: accept + pickup-action + deliver-action.
    action_overhead = 3 * count
    return int(optimal_travel + action_overhead)


def _heuristic_optimal_steps(
    orders: Sequence[Order],
    start_location: Tuple[int, int],
) -> int:
    remaining = list(orders)
    current = start_location
    travel = 0

    while remaining:
        best_index = 0
        best_cost = float("inf")
        best_drop = current

        for idx, order in enumerate(remaining):
            pickup = (order.pickup.x, order.pickup.y)
            drop_candidates = [(d.x, d.y) for d in (order.delivery_locations or [order.dropoff])]
            nearest_drop = min(drop_candidates, key=lambda d: manhattan(pickup, d))
            cost = manhattan(current, pickup) + manhattan(pickup, nearest_drop)
            if cost < best_cost:
                best_cost = cost
                best_index = idx
                best_drop = nearest_drop

        travel += int(best_cost)
        current = best_drop
        remaining.pop(best_index)

    return int(travel + 3 * len(orders))