"""Pedestrian movement physics and capacity-constrained admission.

Two ideas do all the work here:

1. **Speed depends on density.** Walking speed collapses as a corridor fills.
   This is what turns excess demand into a visible, measurable queue instead of
   an ever-faster stream of dots.

2. **Throughput is bounded twice.** A person moving from one link to the next
   must pass a *node* budget (how many people per minute the gate/exit can
   process) and an *edge* budget (how many people per minute the next corridor
   accepts), and the next corridor must have physical room. Everything that
   cannot pass waits, in arrival order.

Both are vectorised over all agents; there is no per-agent Python loop.
"""

from __future__ import annotations

import numpy as np

from ..config import MovementConfig


def weidmann_speed(density: np.ndarray, cfg: MovementConfig) -> np.ndarray:
    """Free walking speed as a function of local density (Weidmann 1993).

        v(rho) = v_free * (1 - exp(-gamma * (1/rho - 1/rho_jam)))

    Below `free_flow_density` the relation is clamped to free speed, which
    avoids the 1/rho singularity for an almost empty corridor.
    """
    rho = np.asarray(density, dtype=np.float64)
    safe = np.maximum(rho, cfg.free_flow_density)
    exponent = -cfg.weidmann_gamma * (1.0 / safe - 1.0 / cfg.jam_density)
    v = cfg.free_speed_mps * (1.0 - np.exp(exponent))
    v = np.where(rho <= cfg.free_flow_density, cfg.free_speed_mps, v)
    return np.clip(v, cfg.min_speed_mps, cfg.free_speed_mps)


def group_rank(sorted_keys: np.ndarray) -> np.ndarray:
    """Rank of each element within its run of equal keys (keys must be sorted).

    Used to implement "the first N in this queue may pass" without a loop.
    """
    n = sorted_keys.shape[0]
    if n == 0:
        return np.empty(0, dtype=np.int64)
    idx = np.arange(n, dtype=np.int64)
    new_run = np.empty(n, dtype=bool)
    new_run[0] = True
    if n > 1:
        new_run[1:] = sorted_keys[1:] != sorted_keys[:-1]
    starts = np.maximum.accumulate(np.where(new_run, idx, np.int64(0)))
    return idx - starts


class CapacityBudget:
    """Integer-per-step budget derived from a per-minute rate.

    Fractional capacity is carried across steps so that, for example, a rate of
    90 people/minute with a 1-second step really admits 90 people per minute
    rather than silently rounding down to 60.
    """

    _UNBOUNDED = np.int64(1 << 40)

    def __init__(self, rate_ppm: np.ndarray) -> None:
        self.base_rate = np.asarray(rate_ppm, dtype=np.float64).copy()
        self.multiplier = np.ones_like(self.base_rate)
        self.carry = np.zeros_like(self.base_rate)

    @property
    def effective_rate(self) -> np.ndarray:
        return self.base_rate * self.multiplier

    def accrue(self, dt_s: float) -> np.ndarray:
        """Advance the budget by `dt_s` and return the integer allowance."""
        rate = self.effective_rate
        finite = np.isfinite(rate)
        self.carry[finite] += rate[finite] * dt_s / 60.0
        allowance = np.where(finite, np.floor(self.carry), self._UNBOUNDED)
        return allowance.astype(np.int64)

    def consume(self, used: np.ndarray) -> None:
        finite = np.isfinite(self.base_rate * self.multiplier)
        self.carry[finite] -= used[finite]
        np.maximum(self.carry, 0.0, out=self.carry)

    def clamp_carry(self, max_seconds: float, dt_s: float) -> None:
        """Stop unused capacity accumulating without bound while a link is idle."""
        rate = self.effective_rate
        finite = np.isfinite(rate)
        cap = rate[finite] * max_seconds / 60.0
        self.carry[finite] = np.minimum(self.carry[finite], np.maximum(cap, dt_s))

    def state(self) -> dict[str, np.ndarray]:
        return {"base_rate": self.base_rate.copy(),
                "multiplier": self.multiplier.copy(),
                "carry": self.carry.copy()}

    def restore(self, state: dict[str, np.ndarray]) -> None:
        self.base_rate = state["base_rate"].copy()
        self.multiplier = state["multiplier"].copy()
        self.carry = state["carry"].copy()


def admit(
    candidate_group: np.ndarray,
    priority: np.ndarray,
    allowance: np.ndarray,
) -> np.ndarray:
    """First-come-first-served admission within each group.

    Parameters
    ----------
    candidate_group
        Group index (node index or edge index) each candidate is queueing for.
    priority
        Lower goes first. In practice the time the agent joined the queue.
    allowance
        Integer allowance per group, indexed by group id.

    Returns
    -------
    Boolean mask over the candidates, True where the candidate may pass.
    """
    n = candidate_group.shape[0]
    if n == 0:
        return np.zeros(0, dtype=bool)
    order = np.lexsort((priority, candidate_group))
    ranked_groups = candidate_group[order]
    rank = group_rank(ranked_groups)
    permitted_sorted = rank < allowance[ranked_groups]
    permitted = np.zeros(n, dtype=bool)
    permitted[order] = permitted_sorted
    return permitted