"""Agent population: generation, storage and reproducibility.

Agents are stored as a structure of arrays. A 40,000-agent population is
therefore about a dozen numpy arrays, which is what makes a full simulation
step cost single-digit milliseconds and a counterfactual roll-out cheap enough
to run five of them while the operator waits.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from ..config import MovementConfig
from ..venue.models import CompiledVenue
from ..venue.scenario import Scenario

# Agent lifecycle
STATUS_WAITING = np.int8(0)   # at origin, not yet departed
STATUS_ON_EDGE = np.int8(1)   # somewhere in the pedestrian network
STATUS_ARRIVED = np.int8(2)   # reached its destination

# Routing policies (indices into the next-hop table)
POLICY_SHORTEST = 0    # Baseline A: minimise distance
POLICY_STATIC = 1      # Baseline B: fixed capacity-aware assignment, no feedback
POLICY_ADAPTIVE = 2    # FlowTwin: dynamic cost, recomputed from live state
N_POLICIES = 3

POLICY_NAMES = {
    POLICY_SHORTEST: "shortest_path",
    POLICY_STATIC: "static_assignment",
    POLICY_ADAPTIVE: "flowtwin_adaptive",
}
POLICY_BY_NAME = {v: k for k, v in POLICY_NAMES.items()}


@dataclass
class AgentPopulation:
    """Structure-of-arrays agent store."""

    status: np.ndarray          # int8
    origin: np.ndarray          # int32 node index
    dest_node: np.ndarray       # int32 node index
    dest_slot: np.ndarray       # int32 index into the destination list
    edge: np.ndarray            # int32 directed-edge index, -1 when not on one
    node: np.ndarray            # int32 node the agent is currently at/waiting on
    pos_m: np.ndarray           # float32 metres travelled along the current edge
    speed_factor: np.ndarray    # float32 personal free-speed multiplier
    compliance: np.ndarray      # float32 probability of accepting a reroute
    policy: np.ndarray          # int8 routing policy
    release_t: np.ndarray       # float32 sim time at which the agent departs
    enter_t: np.ndarray         # float32 sim time the agent entered the network
    arrive_t: np.ndarray        # float32 sim time the agent reached its sink
    queue_since: np.ndarray     # float32 time the agent joined its current queue
    reroute_count: np.ndarray   # int16 number of accepted route changes
    speed_now: np.ndarray       # float32 current walking speed (m/s)
    blocked: np.ndarray         # bool: standing in the queue at the end of an edge

    @property
    def size(self) -> int:
        return int(self.status.shape[0])

    def copy(self) -> "AgentPopulation":
        return AgentPopulation(**{k: v.copy() for k, v in self.__dict__.items()})


def _release_offsets(rng: np.random.Generator, n: int, ramp_s: float, shape: str) -> np.ndarray:
    """Sample departure times within a release window of width `ramp_s`."""
    if ramp_s <= 0:
        return np.zeros(n, dtype=np.float64)
    if shape == "uniform":
        u = rng.random(n)
    elif shape == "double":
        # Two waves: an early group and a later group.
        pick = rng.random(n) < 0.55
        a = np.clip(rng.normal(0.22, 0.10, n), 0.0, 1.0)
        b = np.clip(rng.normal(0.68, 0.13, n), 0.0, 1.0)
        u = np.where(pick, a, b)
    else:  # "peaked" — most people leave immediately, with a long tail
        u = np.clip(rng.beta(1.35, 3.1, n), 0.0, 1.0)
    return u * ramp_s


def build_population(
    venue: CompiledVenue,
    scenario: Scenario,
    rng: np.random.Generator,
    movement: MovementConfig,
    crowd_size: int | None = None,
    release_ramp_s: float | None = None,
    compliance_scale: float = 1.0,
    initial_policy: int = POLICY_SHORTEST,
) -> tuple[AgentPopulation, list[int], list[str]]:
    """Create the agent population for a scenario.

    Returns the population, the list of destination node indices (the "slots"
    the routing tables are built for) and their node ids.
    """
    total = int(crowd_size if crowd_size is not None else scenario.crowd_size)
    if total <= 0:
        raise ValueError("crowd_size must be positive")

    groups = scenario.normalised_demand()

    # Destination slots: the distinct sinks used by this scenario.
    dest_ids: list[str] = []
    for group, _ in groups:
        for dest_id in group.destinations:
            if dest_id not in dest_ids:
                dest_ids.append(dest_id)
    for dest_id in dest_ids:
        if dest_id not in venue.node_index:
            raise ValueError(f"scenario references unknown destination node {dest_id!r}")
    dest_indices = [venue.node_index[d] for d in dest_ids]
    slot_of_node = {node_idx: slot for slot, node_idx in enumerate(dest_indices)}

    # Integer split of the crowd across demand groups (largest-remainder, so the
    # totals are exact and reproducible).
    raw = np.array([share * total for _, share in groups], dtype=np.float64)
    counts = np.floor(raw).astype(np.int64)
    remainder = total - int(counts.sum())
    if remainder > 0:
        order = np.argsort(-(raw - counts))
        counts[order[:remainder]] += 1

    origin_arr = np.empty(total, dtype=np.int32)
    dest_arr = np.empty(total, dtype=np.int32)
    slot_arr = np.empty(total, dtype=np.int32)
    release = np.empty(total, dtype=np.float64)

    base_ramp = release_ramp_s if release_ramp_s is not None else scenario.release.ramp_s

    cursor = 0
    for (group, _), count in zip(groups, counts):
        if count == 0:
            continue
        sl = slice(cursor, cursor + int(count))
        cursor += int(count)

        if group.origin not in venue.node_index:
            raise ValueError(f"scenario references unknown origin node {group.origin!r}")
        o_idx = venue.node_index[group.origin]
        origin_arr[sl] = o_idx

        d_ids = list(group.destinations.keys())
        d_w = np.array([group.destinations[d] for d in d_ids], dtype=np.float64)
        d_w = d_w / d_w.sum()
        chosen = rng.choice(len(d_ids), size=int(count), p=d_w)
        d_node = np.array([venue.node_index[d] for d in d_ids], dtype=np.int32)
        dest_arr[sl] = d_node[chosen]
        slot_arr[sl] = np.array([slot_of_node[int(n)] for n in d_node], dtype=np.int32)[chosen]

        ramp = group.release_ramp_s if group.release_ramp_s is not None else base_ramp
        offsets = _release_offsets(rng, int(count), float(ramp), scenario.release.shape)
        release[sl] = scenario.release.start_s + group.release_offset_s + offsets

    speed_factor = np.clip(
        rng.normal(1.0, movement.speed_sigma, total),
        movement.speed_factor_min,
        movement.speed_factor_max,
    )
    lo, hi = scenario.compliance_min, scenario.compliance_max
    compliance = np.clip(rng.uniform(lo, hi, total) * compliance_scale, 0.0, 1.0)

    pop = AgentPopulation(
        status=np.full(total, STATUS_WAITING, dtype=np.int8),
        origin=origin_arr,
        dest_node=dest_arr,
        dest_slot=slot_arr,
        edge=np.full(total, -1, dtype=np.int32),
        node=origin_arr.copy(),
        pos_m=np.zeros(total, dtype=np.float32),
        speed_factor=speed_factor.astype(np.float32),
        compliance=compliance.astype(np.float32),
        policy=np.full(total, np.int8(initial_policy), dtype=np.int8),
        release_t=release.astype(np.float32),
        enter_t=np.full(total, np.nan, dtype=np.float32),
        arrive_t=np.full(total, np.nan, dtype=np.float32),
        queue_since=np.full(total, np.inf, dtype=np.float32),
        reroute_count=np.zeros(total, dtype=np.int16),
        speed_now=np.zeros(total, dtype=np.float32),
        blocked=np.zeros(total, dtype=bool),
    )
    return pop, dest_indices, dest_ids