"""A work-stealing scheduler: P workers, one deque each, over a virtual clock.

Every worker owns a deque. It pushes and pops its **own end**. When its deque
runs dry it becomes a thief: it picks a victim and takes a task from the
victim's **far end** -- the end the owner does not touch. That end split is
the famous part of the design, and this toy exists to measure what it is
actually worth.

Time is an integer count of ticks. The clock moves for exactly one reason: a
worker spent a tick. Each tick, each worker does exactly one of:

  1. execute one tick of the task in its hands,
  2. pop a task off its own deque (free -- the deque is local), or
  3. spend one tick of a `steal_cost`-tick steal attempt, which resolves at
     the end of that window and may come back empty-handed.

There are no threads, no locks and no `time.monotonic()` anywhere, so a run
is a pure function of its parameters and every number on the commentary page
is byte-identical on every machine. The price is stated in the commentary's
section 8: with no caches and no atomics, this models the *scheduling* half
of the end split and none of the memory half.

Three workloads, and the difference between them is the whole result:

  preload()    a bag of independent, indivisible tasks, dealt unevenly
  spawn_flat() one root that spawns n equal tasks at once
  spawn_tree() divide-and-conquer: a task of size n splits into n/2 + n/2
               until it is small enough to run

Only the third gives the deque a *size gradient* -- oldest task at the far
end is the largest undivided subtree -- and only the third can tell the two
ends apart.
"""

from collections import deque

FAR, OWN = "far", "own"     # which end a thief takes from
SPLIT_TICKS = 1             # cost of subdividing a task into two halves


class Task:
    """A unit of work. Divisible tasks cost `SPLIT_TICKS` and produce two
    children; indivisible ones cost their whole size and produce nothing.

    `leaf_size` is the granularity cutoff. `Task(c, c)` is therefore an
    always-indivisible task of cost c, which is what `preload` deals out.
    """

    __slots__ = ("size", "divisible")

    def __init__(self, size, leaf_size):
        self.size = size
        self.divisible = size > leaf_size

    def cost(self):
        return SPLIT_TICKS if self.divisible else self.size


class Worker:
    """A deque plus the task currently in hand.

    `dq[-1]` is the own end (push/pop, LIFO for the owner); `dq[0]` is the far
    end (oldest task, what a thief takes when `steal_end == FAR`).
    """

    __slots__ = ("wid", "dq", "cur", "left", "stealing", "busy", "overhead")

    def __init__(self, wid):
        self.wid = wid
        self.dq = deque()
        self.cur = None        # task in hand
        self.left = 0          # ticks remaining on it
        self.stealing = 0      # ticks remaining in the current steal attempt
        self.busy = 0          # ticks spent executing tasks
        self.overhead = 0      # ticks spent inside steal attempts


class Sim:
    """The scheduler. Step every worker once per tick, in worker-id order.

    `steal_end` is the counterfactual this toy is built around: FAR is the
    Cilk/Chase-Lev rule, OWN is the same scheduler stealing from the end the
    owner uses. `steal_cost` is the declared price of one attempt, successful
    or not; at 1 the whole effect vanishes, which the demo prints rather than
    hides.
    """

    def __init__(self, nworkers, leaf=8, steal_end=FAR, victim="rr",
                 steal_cost=8, seed=12345):
        self.w = [Worker(i) for i in range(nworkers)]
        self.leaf = leaf
        self.steal_end = steal_end
        self.victim = victim
        self.steal_cost = steal_cost
        self.clock = 0
        self.steals = 0        # attempts that came back with a task
        self.failed = 0        # attempts that found an empty deque
        self.moved = 0         # ticks of subtree work carried by those steals
        self.log = []          # (tick, victim, thief, size, subtree_ticks)
        self.rng = seed
        self.probe = [(i + 1) % nworkers for i in range(nworkers)]

    def _rand(self, n):
        """A seeded LCG, used only by victim='random'. The default policy is
        round-robin precisely so the headline numbers owe nothing to a seed.
        """
        self.rng = (self.rng * 1103515245 + 12345) % (1 << 31)
        return (self.rng >> 16) % n

    def subtree(self, t):
        """Total ticks of work hiding under a task, children included.

        A divisible task of size n eventually becomes n/leaf leaves costing n
        ticks in total, plus n/leaf - 1 internal splits at SPLIT_TICKS each.
        This is the number that makes the far end worth taking: it is what one
        steal actually moves, which is not the same as the task's own cost.
        """
        if not t.divisible:
            return t.size
        return t.size + SPLIT_TICKS * (t.size // self.leaf - 1)

    def pick_victim(self, w):
        n = len(self.w)
        if self.victim == "rr":
            v = self.probe[w.wid]
            self.probe[w.wid] = (v + 1) % n
            if v == w.wid:                       # never rob yourself
                v = self.probe[w.wid]
                self.probe[w.wid] = (v + 1) % n
            return self.w[v]
        if self.victim == "random":
            v = self._rand(n)
            while v == w.wid:
                v = self._rand(n)
            return self.w[v]
        if self.victim == "richest":
            return max((v for v in self.w if v.wid != w.wid),
                       key=lambda v: len(v.dq))
        raise ValueError(f"unknown victim policy {self.victim!r}")

    def run(self, stealing=True, limit=10_000_000):
        """Advance until every deque is empty and no worker holds a task.

        Returns the makespan: the tick at which the last worker finished.

        The termination test sits at the *top* of the pass, before any worker
        steps. Put it at the bottom instead and the run bills one extra tick
        for the pass in which everyone discovers there is nothing left --
        enough to break `makespan == the hoarder's own load` in section 6.
        """
        while self.clock < limit:
            if not any(w.dq or w.cur for w in self.w):
                break
            for w in self.w:
                if w.cur is None and w.stealing == 0 and w.dq:
                    w.cur = w.dq.pop()           # own end, free, LIFO
                    w.left = w.cur.cost()
                if w.cur is not None:
                    w.busy += 1
                    w.left -= 1
                    if w.left == 0:
                        t, w.cur = w.cur, None
                        if t.divisible:
                            half = t.size // 2
                            w.dq.append(Task(half, self.leaf))
                            w.dq.append(Task(t.size - half, self.leaf))
                elif stealing:
                    self._steal_tick(w)
            self.clock += 1
        return self.clock

    def _steal_tick(self, w):
        """Spend one tick inside a steal attempt; resolve it when the window
        closes. A failed attempt costs exactly what a successful one costs --
        the thief cannot know the deque was empty without going to look.
        """
        if w.stealing == 0:
            w.stealing = self.steal_cost
        w.stealing -= 1
        w.overhead += 1
        if w.stealing != 0:
            return
        v = self.pick_victim(w)
        if not v.dq:
            self.failed += 1
            return
        t = v.dq.popleft() if self.steal_end == FAR else v.dq.pop()
        self.steals += 1
        self.moved += self.subtree(t)
        self.log.append((self.clock, v.wid, w.wid, t.size, self.subtree(t)))
        w.cur, w.left = t, t.cost()


def preload(sim, loads):
    """Deal a bag of indivisible tasks: loads[i] is worker i's list of costs."""
    for wid, costs in enumerate(loads):
        for c in costs:
            sim.w[wid].dq.append(Task(c, c))


def spawn_flat(sim, n, each, wid=0):
    """One worker holding n equal indivisible tasks: spawning without a
    gradient. The control that isolates task *size* from LIFO-vs-FIFO.
    """
    for _ in range(n):
        sim.w[wid].dq.append(Task(each, each))


def spawn_tree(sim, root, wid=0):
    """One divisible root. Splitting it builds the size gradient: the far end
    keeps the shallowest, largest subtree; the own end holds the smallest.
    """
    sim.w[wid].dq.append(Task(root, sim.leaf))


def total_work(sim):
    """Ticks of real work executed (steal overhead excluded)."""
    return sum(w.busy for w in sim.w)


def uneven_bag(seed, nworkers=4, hoarded=24, lo=20, hi=200):
    """A deterministic uneven workload: worker 0 hoards `hoarded` tasks, the
    others get one each. The LCG is here so the demo can sweep 500 orderings
    without pulling in `random`, whose stream is not pinned across versions.
    """
    x, loads = seed, [[] for _ in range(nworkers)]
    def nxt():
        nonlocal x
        x = (x * 1103515245 + 12345) % (1 << 31)
        return lo + (x >> 16) % (hi - lo + 1)
    for _ in range(hoarded):
        loads[0].append(nxt())
    for wid in range(1, nworkers):
        loads[wid].append(nxt())
    return loads