from __future__ import annotations

import random
from collections import deque
from dataclasses import dataclass, field
from typing import Deque, Dict, List, Optional, Tuple


@dataclass(slots=True)
class Vehicle:
    vid: int
    desired_exit: Dict[int, int] = field(default_factory=dict)
    birth_step: Optional[int] = None


@dataclass(frozen=True, slots=True)
class Exit:
    exit_cell: int
    dst_roundabout: Optional[int]
    dst_entry_cell: int = 0


@dataclass(frozen=True, slots=True)
class BehaviorProfile:
    name: str
    sink_bias: float
    use_rotation: bool
    random_ties: bool
    slowdown_prob: float
    injection_prob: float
    reroute_prob: float
    swap_prob: float
    inject_window: int
    max_steps: int


PROFILES: Dict[str, BehaviorProfile] = {
    "balanced": BehaviorProfile(
        name="balanced",
        sink_bias=0.0,
        use_rotation=False,
        random_ties=True,
        slowdown_prob=0.07,
        injection_prob=0.9,
        reroute_prob=0.18,
        swap_prob=0.0,
        inject_window=1,
        max_steps=320,
    ),
    "jammy": BehaviorProfile(
        name="jammy",
        sink_bias=-0.5,
        use_rotation=False,
        random_ties=False,
        slowdown_prob=0.22,
        injection_prob=1.0,
        reroute_prob=0.30,
        swap_prob=0.0,
        inject_window=1,
        max_steps=480,
    ),
    "drain_fast": BehaviorProfile(
        name="drain_fast",
        sink_bias=1.1,
        use_rotation=True,
        random_ties=True,
        slowdown_prob=0.0,
        injection_prob=1.0,
        reroute_prob=0.06,
        swap_prob=0.0,
        inject_window=1,
        max_steps=220,
    ),
    "chaotic": BehaviorProfile(
        name="chaotic",
        sink_bias=0.2,
        use_rotation=False,
        random_ties=True,
        slowdown_prob=0.35,
        injection_prob=0.85,
        reroute_prob=0.45,
        swap_prob=0.0,
        inject_window=1,
        max_steps=520,
    ),
    "optimized": BehaviorProfile(
        name="optimized",
        sink_bias=-0.8,
        use_rotation=False,
        random_ties=True,
        slowdown_prob=0.12,
        injection_prob=0.75,
        reroute_prob=0.40,
        swap_prob=0.30,
        inject_window=5,
        max_steps=800,
    ),
}


@dataclass(frozen=True, slots=True)
class Snapshot:
    step: int
    queue: List[int]
    sink: List[int]
    rings: Dict[int, List[Optional[int]]]


@dataclass(frozen=True, slots=True)
class Metrics:
    emitted: int
    steps: int
    throughput: float
    mean_travel_time: float
    kendall_inversions: int
    mean_abs_displacement: float
    adjacent_pairs_preserved: float
    bit_bias: float


@dataclass(frozen=True, slots=True)
class SimulationResult:
    shuffled: List[int]
    history: List[Snapshot]
    metrics: Metrics


class Roundabout:
    __slots__ = ("length", "entry_cell", "cells", "exits")

    def __init__(self, length: int, exits: List[Exit], entry_cell: int = 0):
        self.length = length
        self.entry_cell = entry_cell
        self.cells: List[Optional[Vehicle]] = [None] * length
        self.exits: List[Exit] = exits

    def advance_exclusion(
        self, rng: random.Random, slowdown_prob: float, swap_prob: float = 0.0
    ) -> List[Optional[Vehicle]]:
        length = self.length
        old = self.cells if swap_prob <= 0 else list(self.cells)

        if swap_prob > 0:
            k = 0
            while k < length - 1:
                j = k + 1
                if (
                    old[k] is not None
                    and old[j] is not None
                    and rng.random() < swap_prob
                ):
                    old[k], old[j] = old[j], old[k]
                    k += 2
                    continue
                k += 1

        new = [None] * length
        for k, vehicle in enumerate(old):
            if vehicle is None:
                continue
            k_next = (k + 1) % length
            can_move = old[k_next] is None and (rng.random() >= slowdown_prob)
            if can_move:
                new[k_next] = vehicle
            else:
                new[k] = vehicle
        return new

    def advance_rotate(self) -> List[Optional[Vehicle]]:
        length = self.length
        old = self.cells
        new = [None] * length
        for k, vehicle in enumerate(old):
            new[(k + 1) % length] = vehicle
        return new


class RoundaboutNetworkSim:
    __slots__ = (
        "rbs",
        "head",
        "inject_points",
        "_inject_idx",
        "source",
        "profile",
        "sink",
        "t",
        "rng",
        "travel_times",
        "rotor_ptr",
        "_exit_weights",
    )

    def __init__(
        self,
        roundabouts: List[Roundabout],
        head: int,
        source: Deque[Vehicle],
        profile: BehaviorProfile,
        seed: int,
        inject_points: Optional[List[int]] = None,
    ):
        self.rbs = roundabouts
        self.head = head
        self.inject_points = inject_points or [head]
        self._inject_idx = 0
        self.source = source
        self.profile = profile
        self.sink: List[Tuple[int, int]] = []
        self.t = 0
        self.rng = random.Random(seed)
        self.travel_times: List[int] = []
        self.rotor_ptr: Dict[Tuple[int, int], int] = {}
        self._exit_weights = self._build_exit_weights()

    def _build_exit_weights(self) -> List[List[float]]:
        sink_w = max(0.05, 1.0 + self.profile.sink_bias)
        net_w = max(0.05, 1.0 - 0.4 * self.profile.sink_bias)
        all_weights: List[List[float]] = []
        for rb in self.rbs:
            weights: List[float] = []
            for ex in rb.exits:
                weights.append(sink_w if ex.dst_roundabout is None else net_w)
            all_weights.append(weights)
        return all_weights

    def _pick_vehicle_exit(self, vehicle: Vehicle, rb_id: int) -> int:
        if rb_id in vehicle.desired_exit and (
            self.rng.random() > self.profile.reroute_prob
        ):
            return vehicle.desired_exit[rb_id]
        chosen = self.rng.choices(
            range(len(self.rbs[rb_id].exits)), weights=self._exit_weights[rb_id], k=1
        )[0]
        vehicle.desired_exit[rb_id] = chosen
        return chosen

    def _choose_winner(
        self,
        dst_key: Tuple[int, int],
        contenders: List[Tuple[int, int, int, Vehicle]],
    ) -> int:
        if len(contenders) == 1:
            return 0
        if self.profile.random_ties:
            return self.rng.randrange(len(contenders))
        ptr = self.rotor_ptr.get(dst_key, 0)
        self.rotor_ptr[dst_key] = (ptr + 1) % len(contenders)
        return ptr

    def step(self) -> None:
        moved_states: List[List[Optional[Vehicle]]] = []
        use_rotation = self.profile.use_rotation
        slowdown_prob = self.profile.slowdown_prob
        swap_prob = self.profile.swap_prob

        for rb in self.rbs:
            moved = (
                rb.advance_rotate()
                if use_rotation
                else rb.advance_exclusion(self.rng, slowdown_prob, swap_prob)
            )
            moved_states.append(moved)

        next_states = [cells[:] for cells in moved_states]
        requests_by_dst: Dict[Tuple[int, int], List[Tuple[int, int, int, Vehicle]]] = {}

        for src_id, rb in enumerate(self.rbs):
            cells = moved_states[src_id]
            for exit_idx, ex in enumerate(rb.exits):
                src_cell = ex.exit_cell
                vehicle = cells[src_cell]
                if vehicle is None:
                    continue
                if self._pick_vehicle_exit(vehicle, src_id) != exit_idx:
                    continue

                if ex.dst_roundabout is None:
                    self.sink.append((vehicle.vid, self.t))
                    if vehicle.birth_step is not None:
                        self.travel_times.append(self.t - vehicle.birth_step)
                    next_states[src_id][src_cell] = None
                else:
                    key = (ex.dst_roundabout, ex.dst_entry_cell)
                    requests_by_dst.setdefault(key, []).append(
                        (src_id, src_cell, exit_idx, vehicle)
                    )

        for dst_key, contenders in requests_by_dst.items():
            dst_rb_id, dst_entry_cell = dst_key
            if next_states[dst_rb_id][dst_entry_cell] is not None:
                continue
            winner_idx = self._choose_winner(dst_key, contenders)
            src_id, src_cell, _, vehicle = contenders[winner_idx]
            next_states[src_id][src_cell] = None
            next_states[dst_rb_id][dst_entry_cell] = vehicle

        if self.source and self.rng.random() <= self.profile.injection_prob:
            for attempt in range(len(self.inject_points)):
                rb_id = self.inject_points[
                    (self._inject_idx + attempt) % len(self.inject_points)
                ]
                entry = self.rbs[rb_id].entry_cell
                if next_states[rb_id][entry] is None:
                    window = min(self.profile.inject_window, len(self.source))
                    pick = self.rng.randrange(window) if window > 1 else 0
                    vehicle = self.source[pick]
                    del self.source[pick]
                    if vehicle.birth_step is None:
                        vehicle.birth_step = self.t
                    next_states[rb_id][entry] = vehicle
                    self._inject_idx = (self._inject_idx + attempt + 1) % len(
                        self.inject_points
                    )
                    break

        for rb, cells in zip(self.rbs, next_states):
            rb.cells = cells
        self.t += 1

    def is_done(self) -> bool:
        if self.source:
            return False
        return all(all(cell is None for cell in rb.cells) for rb in self.rbs)

    def snapshot(self) -> Snapshot:
        return Snapshot(
            step=self.t,
            queue=[vehicle.vid for vehicle in self.source],
            sink=[vid for vid, _ in self.sink],
            rings={
                idx: [None if v is None else v.vid for v in rb.cells]
                for idx, rb in enumerate(self.rbs)
            },
        )


def build_network() -> Tuple[List[Roundabout], List[int]]:
    rb0 = Roundabout(8, [Exit(3, 1, 0), Exit(6, 2, 0)], entry_cell=0)
    rb1 = Roundabout(6, [Exit(2, 2, 3), Exit(4, 3, 0), Exit(5, 4, 0)], entry_cell=0)
    rb2 = Roundabout(7, [Exit(5, None), Exit(3, 3, 2), Exit(6, 0, 4)], entry_cell=0)
    rb3 = Roundabout(6, [Exit(2, 4, 2), Exit(4, None), Exit(5, 1, 3)], entry_cell=0)
    rb4 = Roundabout(7, [Exit(5, None), Exit(3, 0, 3), Exit(6, 2, 4)], entry_cell=0)
    inject_points = [0, 2]
    return [rb0, rb1, rb2, rb3, rb4], inject_points


def kendall_inversion_count(order: List[int]) -> int:
    if len(order) < 2:
        return 0

    vals = sorted(set(order))
    rank = {v: i + 1 for i, v in enumerate(vals)}
    size = len(vals) + 1
    bit = [0] * (size + 1)

    def bit_add(i: int, delta: int) -> None:
        while i <= size:
            bit[i] += delta
            i += i & -i

    def bit_sum(i: int) -> int:
        s = 0
        while i > 0:
            s += bit[i]
            i -= i & -i
        return s

    inversions = 0
    seen = 0
    for v in order:
        i = rank[v]
        less_or_eq = bit_sum(i)
        inversions += seen - less_or_eq
        bit_add(i, 1)
        seen += 1
    return inversions


def mean_abs_displacement(order: List[int]) -> float:
    if not order:
        return 0.0
    return sum(abs(i - v) for i, v in enumerate(order)) / len(order)


def adjacent_pairs_preserved(order: List[int]) -> float:
    if len(order) < 2:
        return 1.0
    preserved = 0
    for a, b in zip(order, order[1:]):
        if abs(a - b) == 1:
            preserved += 1
    return preserved / (len(order) - 1)


def bit_bias_from_sequence(order: List[int]) -> float:
    if len(order) < 2:
        return 0.0
    ones = 0
    total = 0
    for prev, cur in zip(order, order[1:]):
        ones += 1 if cur > prev else 0
        total += 1
    return (ones / total) - 0.5


def complete_with_residuals(
    sim: RoundaboutNetworkSim, expected_count: int
) -> List[int]:
    shuffled = [vid for vid, _ in sim.sink]
    if len(shuffled) >= expected_count:
        return shuffled

    seen = set(shuffled)
    for vehicle in sim.source:
        if vehicle.vid not in seen:
            shuffled.append(vehicle.vid)
            seen.add(vehicle.vid)
    for rb in sim.rbs:
        for vehicle in rb.cells:
            if vehicle is not None and vehicle.vid not in seen:
                shuffled.append(vehicle.vid)
                seen.add(vehicle.vid)
    return shuffled


def compute_metrics(sim: RoundaboutNetworkSim, shuffled: List[int]) -> Metrics:
    emitted = len(sim.sink)
    steps = max(sim.t, 1)
    mean_travel = (
        sum(sim.travel_times) / len(sim.travel_times) if sim.travel_times else 0.0
    )
    return Metrics(
        emitted=emitted,
        steps=sim.t,
        throughput=emitted / steps,
        mean_travel_time=mean_travel,
        kendall_inversions=kendall_inversion_count(shuffled),
        mean_abs_displacement=mean_abs_displacement(shuffled),
        adjacent_pairs_preserved=adjacent_pairs_preserved(shuffled),
        bit_bias=bit_bias_from_sequence(shuffled),
    )


def run_simulation(
    items: List[int],
    seed: int,
    profile: BehaviorProfile,
    record_history: bool = False,
    max_steps_override: Optional[int] = None,
) -> SimulationResult:
    roundabouts, inject_points = build_network()
    vehicles = [Vehicle(vid=i) for i in items]
    source: Deque[Vehicle] = deque(vehicles)

    sim = RoundaboutNetworkSim(
        roundabouts=roundabouts,
        head=0,
        source=source,
        profile=profile,
        seed=seed,
        inject_points=inject_points,
    )

    history: List[Snapshot] = []
    if record_history:
        history.append(sim.snapshot())

    max_steps = profile.max_steps if max_steps_override is None else max_steps_override
    for _ in range(max_steps):
        sim.step()
        if record_history:
            history.append(sim.snapshot())
        if sim.is_done():
            break

    shuffled = complete_with_residuals(sim, expected_count=len(items))
    metrics = compute_metrics(sim, shuffled)
    return SimulationResult(shuffled=shuffled, history=history, metrics=metrics)


def print_result(profile: BehaviorProfile, result: SimulationResult) -> None:
    m = result.metrics
    print(f"[{profile.name}] shuffled: {result.shuffled}")
    print(
        "  metrics: "
        f"emitted={m.emitted} "
        f"steps={m.steps} "
        f"throughput={m.throughput:.3f} "
        f"mean_travel={m.mean_travel_time:.2f} "
        f"inversions={m.kendall_inversions} "
        f"mad={m.mean_abs_displacement:.2f} "
        f"adj_preserved={m.adjacent_pairs_preserved:.2f} "
        f"bit_bias={m.bit_bias:+.3f}"
    )


def with_overrides(
    base: BehaviorProfile,
    slowdown_prob: Optional[float],
    injection_prob: Optional[float],
) -> BehaviorProfile:
    return BehaviorProfile(
        name=base.name,
        sink_bias=base.sink_bias,
        use_rotation=base.use_rotation,
        random_ties=base.random_ties,
        slowdown_prob=base.slowdown_prob if slowdown_prob is None else slowdown_prob,
        injection_prob=base.injection_prob
        if injection_prob is None
        else injection_prob,
        reroute_prob=base.reroute_prob,
        swap_prob=base.swap_prob,
        inject_window=base.inject_window,
        max_steps=base.max_steps,
    )
