"""Given code for Task 2: a discrete-event simulation of RL rollout scheduling.

You do not need to edit this file, but DO read `run_simulation` once: it is short,
and it is exactly the model behind the chapter's rollout-timeline demo.

The model (disaggregated placement: generation and training on different GPUs)
------------------------------------------------------------------------------
* `n_lanes` generation lanes. A lane produces one response at a time. Think of a
  lane as one sequence slot in an inference engine's running batch.
* Response j takes `lengths[j]` time units to generate (time is proportional to
  the number of tokens; a median-length response takes about 1 unit).
* One trainer. A training step consumes `batch_size` finished responses, takes
  `train_time` units, and then publishes a new policy version (0, 1, 2, ...).
* A response is tagged with the policy version that was current when it STARTED.
  Its *lag* (staleness) is how many versions old that tag is when the trainer uses
  it: lag = trainer_version - tag.
* Admission control (your `may_start`): a free lane may start a new response only
  if that keeps the generators at most `max_staleness` batches ahead of the trainer.
  max_staleness = 0 is ordinary synchronous RL.
* Batching (your `take_batch`): the trainer takes finished responses in the order
  they finished (first in, first out), optionally discarding ones that are too stale.
"""

import numpy as np


def lognormal_lengths(n, sigma, cap=8.0, seed=0):
    """n response lengths with median 1 and a long right tail.

    length = min(cap, exp(sigma * z)), z ~ N(0, 1). sigma = 0 gives all lengths = 1;
    sigma = 1 gives a heavy tail (the mean is about 1.6, a few responses hit the cap).
    `cap` plays the role of the maximum generation length.
    """
    z = np.random.default_rng(seed).standard_normal(n)
    return np.minimum(cap, np.exp(sigma * z))


def run_simulation(lengths, n_lanes, batch_size, max_staleness, train_time, n_steps, drop_stale,
                   may_start, take_batch):
    """Simulate `n_steps` training steps. Returns a dict of statistics (see bottom)."""
    lengths = list(lengths)
    t = 0.0
    version = 0                 # number of finished training steps = current policy version
    lanes = [None] * n_lanes    # None (idle) or the sample dict being generated
    queue = []                  # finished samples waiting for the trainer, in finish order
    trainer_end = None          # time the current training step ends (None = trainer idle)
    n_started = 0
    n_dropped = 0
    trace = []                  # every sample ever started (for plotting)
    trained = []                # samples the trainer actually used
    steps = []                  # (start, end, version_before) for each training step

    while True:
        # 1) Is the trainer free and is there a full batch to train on?
        if trainer_end is None:
            batch, dropped, rest = take_batch(queue, version, batch_size, max_staleness, drop_stale)
            for s in dropped:
                s["dropped"] = True
            n_dropped += len(dropped)
            queue = rest
            if batch is not None:
                trained.extend(batch)
                trainer_end = t + train_time
                steps.append((t, trainer_end, version))

        # 2) Fill idle lanes, as long as admission control allows.
        for i in range(n_lanes):
            if lanes[i] is None and may_start(n_started, n_dropped, version, max_staleness, batch_size):
                if n_started >= len(lengths):
                    raise RuntimeError("ran out of response lengths: pass a longer `lengths` array")
                s = {"id": n_started, "lane": i, "tag": version, "start": t,
                     "end": t + float(lengths[n_started]), "lag": None, "dropped": False}
                lanes[i] = s
                trace.append(s)
                n_started += 1

        # 3) Jump to the next event: a lane finishing or the trainer finishing.
        ends = [s["end"] for s in lanes if s is not None]
        if trainer_end is not None:
            ends.append(trainer_end)
        if not ends:
            raise RuntimeError("deadlock: nothing is running and nothing may start "
                               "(check may_start and take_batch)")
        t = min(ends)

        # 4) Process everything that finishes at time t.
        done_now = [s for s in lanes if s is not None and s["end"] <= t + 1e-12]
        for s in sorted(done_now, key=lambda s: (s["end"], s["id"])):
            queue.append(s)
            lanes[s["lane"]] = None
        if trainer_end is not None and trainer_end <= t + 1e-12:
            trainer_end = None
            version += 1
            if version == n_steps:
                break

    wall = t
    busy = sum(min(s["end"], wall) - s["start"] for s in trace if s["start"] < wall)
    lags = np.array([s["lag"] for s in trained], dtype=float)
    dropped_time = sum(s["end"] - s["start"] for s in trace if s["dropped"])
    return {
        "wall_time": wall,
        "lane_util": busy / (n_lanes * wall),          # fraction of lane-time spent generating
        "trainer_util": n_steps * train_time / wall,   # fraction of time the trainer is busy
        "throughput": n_steps * batch_size / wall,     # trained samples per time unit
        "mean_lag": float(lags.mean()),
        "max_lag": int(lags.max()),
        "frac_stale": float((lags > 0).mean()),        # fraction of trained samples with lag >= 1
        "n_dropped": n_dropped,
        "wasted_frac": dropped_time / busy if busy > 0 else 0.0,
        "trace": trace,
        "steps": steps,
        "trained": trained,
    }