# SPDX-License-Identifier: MIT
"""YS-EN-0001 · Why Homa schedules the shortest message first.

An idealised model of one receiver's downlink: a single link of rate 1 fed by
Poisson message arrivals (an M/G/1 queue). We compare the order in which the
link serves messages:

  fifo   first come, first served (one queue, no preemption)
  ps     processor sharing: every active message gets an equal share of the
         link. This is the fluid ideal of per-flow fair sharing, which is what
         TCP congestion control aims for.
  srpt   shortest remaining processing time first (preemptive). The policy
         Homa approximates.
  prioK  K strict priority levels chosen by remaining size, fair sharing
         within a level. A coarse stand-in for Homa's handful of switch
         priority queues (not Homa's actual priority-assignment rules).

The model ignores propagation delay, packetisation, the sender side and
incast. It isolates one effect: how much the *order* of service matters.

Slowdown = completion time / size (time the message would take alone on an
idle link). 1.0 is perfect.

Usage:  python queue_sim.py            # writes results.json next to this file
Requires numpy. Deterministic (fixed seeds).
"""
from __future__ import annotations

import heapq
import json
import math
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent


# ---------------------------------------------------------------- workloads
def load_cdf(name: str) -> tuple[np.ndarray, np.ndarray]:
    """Read a two-column (size, cumulative fraction) file from workloads/."""
    sizes, cum = [], []
    for line in (HERE / "workloads" / f"{name}.txt").read_text().splitlines():
        line = line.split("#", 1)[0].strip()
        if line:
            s, c = line.split()
            sizes.append(float(s))
            cum.append(float(c))
    return np.asarray(sizes), np.asarray(cum)


def sampler(sizes: np.ndarray, cum: np.ndarray):
    """Inverse-CDF sampling of an empirical step distribution."""
    def draw(rng: np.random.Generator, n: int) -> np.ndarray:
        idx = np.searchsorted(cum, rng.random(n), side="left")
        return sizes[np.minimum(idx, len(sizes) - 1)]
    return draw


def mean_size(sizes: np.ndarray, cum: np.ndarray) -> float:
    p = np.diff(np.concatenate([[0.0], cum]))
    return float(np.sum(p * sizes))


# ---------------------------------------------------------------- policies
def sim_fifo(arr, size):
    done = np.empty_like(arr)
    t = 0.0
    for i in range(len(arr)):
        t = max(t, arr[i]) + size[i]
        done[i] = t
    return done


def sim_ps(arr, size):
    """Egalitarian processor sharing via virtual time.

    Virtual time V advances at 1/n while n messages are active; a message that
    arrives at virtual time v finishes when V reaches v + size.
    """
    n_jobs = len(arr)
    done = np.empty(n_jobs)
    heap: list[tuple[float, int]] = []  # (virtual finish, job)
    t = 0.0
    v = 0.0
    i = 0
    while i < n_jobs or heap:
        n = len(heap)
        next_arr = arr[i] if i < n_jobs else math.inf
        if n:
            t_fin = t + (heap[0][0] - v) * n
            if t_fin <= next_arr:
                v = heap[0][0]
                t = t_fin
                _, j = heapq.heappop(heap)
                done[j] = t
                continue
            v += (next_arr - t) / n
        t = next_arr
        heapq.heappush(heap, (v + size[i], i))
        i += 1
    return done


def sim_srpt(arr, size):
    """Preemptive shortest-remaining-processing-time first.

    Only the job in service shrinks, so re-keying it on each arrival keeps
    the heap exact.
    """
    n_jobs = len(arr)
    done = np.empty(n_jobs)
    heap: list[tuple[float, int]] = []  # (remaining, job)
    t = 0.0
    i = 0
    while i < n_jobs or heap:
        next_arr = arr[i] if i < n_jobs else math.inf
        if heap:
            rem, j = heap[0]
            if t + rem <= next_arr:
                heapq.heappop(heap)
                t += rem
                done[j] = t
                continue
            heapq.heapreplace(heap, (rem - (next_arr - t), j))
        t = next_arr
        heapq.heappush(heap, (size[i], i))
        i += 1
    return done


def sim_levels(arr, size, cutoffs):
    """K strict priority levels by remaining size, fair sharing within a level.

    Level k holds messages whose remaining bytes lie in (cut[k-1], cut[k]].
    The link serves only the lowest non-empty level, splitting its rate
    equally among that level's messages (packets from several senders
    interleave in one switch queue). As a message drains below a cutoff it
    moves up a level. No cutoffs = processor sharing; a cutoff between every
    pair of sizes = SRPT.

    Each level keeps a heap of (remaining + drained_k) so draining every
    member at once is a single counter update.
    """
    cut = [0.0] + list(cutoffs)
    K = len(cut)
    heaps: list[list[tuple[float, int]]] = [[] for _ in range(K)]
    drained = [0.0] * K
    level_of = lambda rem: int(np.searchsorted(cut, rem, side="left")) - 1
    n_jobs = len(arr)
    done = np.empty(n_jobs)
    t = 0.0
    i = 0
    active = 0
    while i < n_jobs or active:
        next_arr = arr[i] if i < n_jobs else math.inf
        if active:
            k = next(k for k in range(K) if heaps[k])
            h, n = heaps[k], len(heaps[k])
            gap = h[0][0] - drained[k] - cut[k]  # until the smallest finishes or moves up
            dt = gap * n
            if t + dt <= next_arr:
                t += dt
                drained[k] += gap
                key, j = heapq.heappop(h)
                if k == 0:
                    done[j] = t
                    active -= 1
                else:
                    rem = key - drained[k]
                    heapq.heappush(heaps[k - 1], (rem + drained[k - 1], j))
                continue
            drained[k] += (next_arr - t) / n
        t = next_arr
        k = max(level_of(size[i]), 0)
        heapq.heappush(heaps[k], (size[i] + drained[k], i))
        active += 1
        i += 1
    return done


# ---------------------------------------------------------------- experiment
def run(workload: str, load: float, n: int, seed: int, policies: dict):
    sizes, cum = load_cdf(workload)
    draw = sampler(sizes, cum)
    rng = np.random.default_rng(seed)
    es = mean_size(sizes, cum)
    lam = load / es
    arr = np.cumsum(rng.exponential(1 / lam, n))
    size = draw(rng, n)
    warm = n // 10  # drop the first 10 % while the queue fills
    out = {}
    for name, fn in policies.items():
        sd = (fn(arr, size) - arr) / size
        out[name] = (size[warm:], sd[warm:])
    return out


def size_deciles(size, k=10):
    """Bucket messages into k equal-count groups by size (ties split by
    arrival order), like the "message size percentile" axis in the Homa paper."""
    order = np.argsort(size, kind="stable")
    b = np.empty(len(size), dtype=int)
    b[order] = np.arange(len(size)) * k // len(size)
    return b


def log_cutoffs(workload: str, k: int) -> list[float]:
    """k-1 size cutoffs spaced evenly in log(size) between the workload's
    smallest and largest message. A deliberately simple, assumption-light
    rule; Homa itself sets cutoffs from measured traffic (SIGCOMM 2018, §3.4)
    and ranks scheduled messages by remaining bytes."""
    sizes, _ = load_cdf(workload)
    return list(np.geomspace(sizes[0], sizes[-1], k + 1)[1:-1])


def policies_for(workload: str, levels=(8,)) -> dict:
    pol = {"fifo": sim_fifo, "ps": sim_ps}
    for k in levels:
        cut = log_cutoffs(workload, k)
        pol[f"prio{k}"] = lambda a, s, cut=cut: sim_levels(a, s, cut)
    pol["srpt"] = sim_srpt
    return pol


def summarise(sd: np.ndarray) -> dict:
    return {"mean": float(sd.mean()), "p50": float(np.median(sd)), "p99": float(np.quantile(sd, 0.99))}


def main():
    n = 400_000
    results = {"n": n, "by_size": {}, "by_load": {}, "levels": {}, "workloads": {}}

    # 1. W4 (Facebook Hadoop): slowdown by message-size decile at 80 % load
    res = run("w4", 0.8, n, seed=1, policies=policies_for("w4"))
    for name, (sz, sd) in res.items():
        b = size_deciles(sz)
        results["by_size"][name] = [
            {"size_max": float(sz[b == k].max()), **summarise(sd[b == k])} for k in range(10)
        ]

    # 2. W4: mean and p99 slowdown versus load, with a theory check for PS
    for load in (0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9):
        r = run("w4", load, n, seed=1, policies=policies_for("w4"))
        results["by_load"][f"{load:.1f}"] = {name: summarise(sd) for name, (_, sd) in r.items()}
        ps = results["by_load"][f"{load:.1f}"]["ps"]["mean"]
        print(f"W4 load {load:.1f}  PS mean slowdown {ps:6.3f}   M/G/1-PS theory 1/(1-load) = {1/(1-load):6.3f}")

    # 3. W4 at 80 %: how many priority levels does it take?
    r = run("w4", 0.8, n, seed=1, policies=policies_for("w4", levels=(2, 4, 8, 16)))
    for name, (sz, sd) in r.items():
        short = sz <= np.median(sz)
        results["levels"][name] = {"all": summarise(sd), "short": summarise(sd[short])}

    # 4. every workload at 80 %: p99 slowdown of the shortest half of messages
    for w in ("w2", "w3", "w4", "w5"):
        r = run(w, 0.8, n, seed=1, policies=policies_for(w))
        out = {}
        for name, (sz, sd) in r.items():
            order = np.argsort(sz, kind="stable")
            short = np.zeros(len(sz), bool)
            short[order[: len(sz) // 2]] = True
            out[name] = {"short": summarise(sd[short]), "all": summarise(sd)}
        sizes, cum = load_cdf(w)
        results["workloads"][w] = {"mean_size": mean_size(sizes, cum), **out}

    (HERE / "results.json").write_text(json.dumps(results, indent=1))
    print("wrote results.json")


if __name__ == "__main__":
    main()