"""Chapter 20 lab, Task 2: synchronous vs asynchronous rollout scheduling.

The event loop is given (scheduler_core.run_simulation: read it once, it is short).
You write the two decisions that turn synchronous RL into asynchronous RL:

    Task 2a  may_start    admission control: may a free lane start a new response?
    Task 2b  take_batch   the trainer's batching: which finished responses to train on,
                          how stale each one is, and which ones to throw away

Run `python -m pytest -q -k task2`.

Vocabulary (also in scheduler_core.py):
  * version:  the number of training steps finished so far. Version v is the policy
              the generators currently sample from.
  * tag:      the version that was current when a response STARTED generating.
  * lag:      version - tag, measured when the trainer uses the response. Lag 0 means
              perfectly on-policy; lag 2 means "sampled by the policy from two updates ago".
  * max_staleness k:  how many batches the generators may run ahead of the trainer.
              k = 0 is ordinary synchronous RL.
"""

from scheduler_core import run_simulation


# ---------------------------------------------------------------------------
# Task 2a
# ---------------------------------------------------------------------------
def may_start(n_started, n_dropped, version, max_staleness, batch_size):
    """May a free generation lane start a new response right now? Returns a bool.

    n_started:     responses started so far (including running and dropped ones)
    n_dropped:     responses thrown away so far because they were too stale
    version:       the trainer's current policy version (= training steps finished)
    max_staleness: k
    batch_size:    B, responses per training step

    The rule (it is AReaL's): the number of useful responses ever started,
    n_started - n_dropped, must stay BELOW (version + k + 1) * B.

    Check it with k = 0: at version 0 you may start B responses and then must wait
    until the trainer finishes step 0. That is synchronous RL. With k = 1 you may
    also start the NEXT batch (sampled by the same, soon-to-be-old policy) while
    the trainer works: "one-step off-policy".
    """
    # TODO: one line
    raise NotImplementedError("Task 2a: may_start")


# ---------------------------------------------------------------------------
# Task 2b
# ---------------------------------------------------------------------------
def take_batch(queue, version, batch_size, max_staleness, drop_stale):
    """Pick the trainer's next batch from the finished responses.

    queue:       list of finished samples in the order they finished (oldest first).
                 Each sample is a dict with at least the key 'tag'.
    version:     the trainer's current version (the step about to run updates
                 version -> version + 1)
    batch_size:  B
    max_staleness, drop_stale: if drop_stale is True, first remove from the queue
                 every sample whose lag (version - tag) is greater than max_staleness.

    Returns a tuple (batch, dropped, remaining):
      batch:     None if fewer than B (non-dropped) samples are available. Otherwise
                 the FIRST B samples of the (filtered) queue, and for each of them set
                 sample['lag'] = version - sample['tag'].
      dropped:   list of the samples you removed as too stale (empty if drop_stale is False)
      remaining: the rest of the (filtered) queue, in order, that was not put in the batch.

    Example: queue tags [0, 0, 1], version 2, B = 2, k = 1, drop_stale=True
        -> both tag-0 samples have lag 2 > 1 and are dropped; only one sample is
           left, so batch is None, dropped = the two tag-0 samples, remaining = [tag 1].
    """
    # TODO
    raise NotImplementedError("Task 2b: take_batch")


# ---------------------------------------------------------------------------
# Given: the full simulation, using your two functions
# ---------------------------------------------------------------------------
def simulate(lengths, n_lanes, batch_size, max_staleness, train_time, n_steps, drop_stale=False):
    """Run the discrete-event simulation. See scheduler_core.run_simulation for the outputs:
    wall_time, lane_util, trainer_util, throughput, mean_lag, max_lag, frac_stale,
    n_dropped, wasted_frac, trace, steps, trained."""
    return run_simulation(lengths, n_lanes, batch_size, max_staleness, train_time, n_steps, drop_stale,
                          may_start=may_start, take_batch=take_batch)