# SPDX-License-Identifier: MIT
"""Convert the W1–W5 tables in HomaModule util/dist.cc to size/CDF text files.

Source (BSD-2-Clause OR GPL-2.0+, Homa Developers), pinned commit:
https://raw.githubusercontent.com/PlatformLab/HomaModule/40891f01b290872d4fe781c4941deb984785b4b2/util/dist.cc

Messages are truncated at 1 MB, as in the Homa/Linux paper (ATC 2021,
Table 3), whose means this reproduces.
"""
import re
import sys
import urllib.request
from pathlib import Path

URL = ("https://raw.githubusercontent.com/PlatformLab/HomaModule/"
       "40891f01b290872d4fe781c4941deb984785b4b2/util/dist.cc")
CAP = 1_000_000
HERE = Path(__file__).resolve().parent

src = urllib.request.urlopen(URL).read().decode()
for w in ("w1", "w2", "w3", "w4", "w5"):
    body = re.search(rf"dist_point_gen::{w}\[\] = \{{(.*?)\n\}};", src, re.S).group(1)
    pts = [(int(a), float(b)) for a, b in re.findall(r"\{(\d+),\s*([0-9.eE+-]+)\}", body) if int(a) > 0]
    total = sum(f for _, f in pts)
    lines, cum, mean = [], 0.0, 0.0
    merged = {}
    for length, f in pts:
        L = min(length, CAP)
        merged[L] = merged.get(L, 0.0) + f / total
    for L in sorted(merged):
        cum += merged[L]
        mean += L * merged[L]
        lines.append(f"{L} {min(cum, 1.0):.9f}")
    head = f"# {w} from HomaModule util/dist.cc (pinned commit), truncated at {CAP} B\n# mean {mean:.0f} B\n"
    (HERE / "workloads" / f"{w}.txt").write_text(head + "\n".join(lines) + "\n")
    print(w, len(lines), "points, mean", round(mean))