# /// script
# requires-python = ">=3.10"
# dependencies = ["lz4", "zstandard"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Compress or not: the rate comparison.

compress + transfer(compressed) + decompress against transfer(uncompressed),
for a JSON file, with LZ4, Zstd levels 1, 3 and 10, and gzip. Codecs are
timed on this machine; transfers are computed for 10 Mb/s, 1 Gb/s and
100 Gb/s links.
"""
import argparse, gzip, json, random, time
import lz4.frame, zstandard as zstd

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--mb", type=int, default=500, help="Size of the generated JSON text, in MB.")
p.add_argument("--links", type=lambda s: [float(v) for v in s.split(",")], default=[0.01, 1, 100], help="Link speeds in Gb/s.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(0)
line = lambda i: json.dumps({"id": i, "user": f"u{random.randint(1, 50_000)}", "event": random.choice(["view", "click", "buy"]),
                             "path": f"/p/{random.randint(1, 2000)}", "ms": random.randint(1, 900)}) + "\n"
chunk = "".join(line(i) for i in range(200_000)).encode()
data = (chunk * (args.mb * 1_000_000 // len(chunk) + 1))[: args.mb * 1_000_000]

codecs = {
    "none": (lambda b: b, lambda b: b),
    "LZ4": (lz4.frame.compress, lz4.frame.decompress),
    "Zstd 1": (zstd.ZstdCompressor(level=1).compress, zstd.ZstdDecompressor().decompress),
    "Zstd 3": (zstd.ZstdCompressor(level=3).compress, zstd.ZstdDecompressor().decompress),
    "Zstd 10": (zstd.ZstdCompressor(level=10).compress, zstd.ZstdDecompressor().decompress),
    "gzip 6": (lambda b: gzip.compress(b, 6), gzip.decompress),
}
rows, totals = [], []
for name, (comp, decomp) in codecs.items():
    t = time.perf_counter(); z = comp(data); tc = time.perf_counter() - t
    t = time.perf_counter(); decomp(z); td = time.perf_counter() - t
    rows.append({"codec": name, "ratio": len(data) / len(z), "compress_s": tc, "decompress_s": td})
    for gbps in args.links:
        total = tc + len(z) * 8 / (gbps * 1e9) + td
        totals.append({"codec": name, "gbps": gbps, "total_s": total})
    print(f"{name:<8} ratio {len(data) / len(z):5.2f}x  compress {tc:6.2f} s  decompress {td:5.2f} s   " +
          "  ".join(f"{gbps:g} Gb/s: {next(x['total_s'] for x in totals if x['codec'] == name and x['gbps'] == gbps):8.2f} s" for gbps in args.links))
best = {gbps: min((x for x in totals if x["gbps"] == gbps), key=lambda x: x["total_s"])["codec"] for gbps in args.links}
print("fastest end to end: " + ", ".join(f"{g:g} Gb/s -> {c}" for g, c in best.items()))

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"codecs": rows, "totals": totals, "best": [{"gbps": g, "codec": c} for g, c in best.items()]}}, f)