#!/usr/bin/env python3
"""The working set of an aggregation is the number of DISTINCT keys, not the
input size -- DDIA 2e Ch. 11, "Sorting Versus In-Memory Aggregation".

Two ways to count URL hits in a log file:

  1. In-memory hash aggregation:  counts[url] += 1  over a Python dict.
  2. The Unix pipeline:           sort | uniq -c | sort -rn | head.

Both compute the same top-N. But their MEMORY profiles differ in a way a reader
who expects "memory grows with the input" gets wrong. This script holds the
LINE COUNT fixed (2,000,000 lines every time) and varies only the number of
DISTINCT keys: 1, then 1,000, then 100,000, then 1,000,000. For each dataset it
runs the hash aggregation in a SUBPROCESS (so ru_maxrss reflects just that run)
and measures the BSD `sort` step under /usr/bin/time -l.

The result: the hash table's peak RSS is flat in the line count but RISES steeply
with the distinct-key count -- one dict slot per distinct group must stay resident.
A file of 2,000,000 identical lines costs the dict exactly ONE entry. The `sort`
step's RSS is essentially INSENSITIVE to cardinality (it barely moves as distinct
keys span a millionfold range); what it tracks instead is the input's byte size.
That is the external-sort memory model: RAM bounded by the data's size and a fixed
buffer, sorted runs spilled to disk -- never one slot per distinct group.

Measurement note: macOS BSD `sort` mmaps its input, so its RSS reflects the mapped
file pages (bounded by file bytes, reclaimable), not the distinct-key count. The
reproducible result is the SCALING -- hash grows with cardinality, sort does not.

Pure standard library + BSD `sort`/`uniq` via subprocess. Seeded -> reproducible
datasets (exact RSS bytes vary slightly by allocator; the SCALING is the result).
"""
import os
import random
import re
import resource
import subprocess
import sys
import tempfile

LINE_COUNT = 2_000_000           # held CONSTANT across every dataset
DISTINCTS = [1, 1_000, 100_000, 1_000_000]
SEED = 42
URL = "http://example.com/page/{}".format


# --------------------------------------------------------------------------- #
# Dataset generation: exactly LINE_COUNT lines, exactly `distinct` unique URLs.
# Every one of the `distinct` keys is emitted at least once (so the distinct
# count is exact), then the remaining lines are filled with a skew toward low
# indices so a few "hot" URLs clearly top the ranking. Seeded -> deterministic.
# --------------------------------------------------------------------------- #

def generate(path, distinct, line_count=LINE_COUNT, seed=SEED):
    rng = random.Random(seed)
    with open(path, "w") as f:
        # Emit each distinct key once so the cardinality is exactly `distinct`.
        seeded = list(range(distinct))
        rng.shuffle(seeded)
        for i in seeded:
            f.write(URL(i))
            f.write("\n")
        # Fill the rest with a cube-law skew toward index 0 -> clear hot keys.
        for _ in range(line_count - distinct):
            idx = int(distinct * rng.random() ** 3)
            if idx >= distinct:
                idx = distinct - 1
            f.write(URL(idx))
            f.write("\n")


# --------------------------------------------------------------------------- #
# Hash aggregation worker. Run as a subprocess so ru_maxrss isolates THIS run:
# baseline interpreter RSS + the dict of one slot per distinct key. Reads the
# file line by line (never loads it whole) so only the dict grows.
# --------------------------------------------------------------------------- #

def hash_aggregate(path):
    counts = {}
    with open(path) as f:
        for line in f:
            url = line.rstrip("\n")
            counts[url] = counts.get(url, 0) + 1
    peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss  # bytes on macOS
    top = sorted(counts.items(), key=lambda kv: (-kv[1], kv[0]))[:5]
    # Internal protocol, captured by the orchestrator (not shown to the reader).
    print(f"RSS_BYTES {peak}")
    print(f"DISTINCT {len(counts)}")
    for url, c in top:
        print(f"TOP {c} {url}")


# --------------------------------------------------------------------------- #
# Orchestration helpers.
# --------------------------------------------------------------------------- #

def run_hash_worker(path):
    """Run the hash aggregation in a fresh process; return (peak_bytes, distinct, top)."""
    out = subprocess.run(
        [sys.executable, os.path.abspath(__file__), "--hash", path],
        capture_output=True, text=True, check=True,
    ).stdout
    peak = distinct = None
    top = []
    for line in out.splitlines():
        if line.startswith("RSS_BYTES "):
            peak = int(line.split()[1])
        elif line.startswith("DISTINCT "):
            distinct = int(line.split()[1])
        elif line.startswith("TOP "):
            _, c, url = line.split(" ", 2)
            top.append((int(c), url))
    return peak, distinct, top


def measure_sort_rss(path):
    """Run BSD `sort` under /usr/bin/time -l; return its peak RSS in bytes (or None)."""
    proc = subprocess.run(
        ["/usr/bin/time", "-l", "sort", path],
        stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True,
    )
    m = re.search(r"(\d+)\s+maximum resident set size", proc.stderr)
    return int(m.group(1)) if m else None


def sort_pipeline_topn(path, n=5):
    """The book's pipeline: sort | uniq -c | sort -rn | head -n. Returns [(count, url)]."""
    p1 = subprocess.Popen(["sort", path], stdout=subprocess.PIPE)
    p2 = subprocess.Popen(["uniq", "-c"], stdin=p1.stdout, stdout=subprocess.PIPE)
    p1.stdout.close()
    p3 = subprocess.Popen(["sort", "-rn"], stdin=p2.stdout, stdout=subprocess.PIPE)
    p2.stdout.close()
    p4 = subprocess.Popen(["head", "-n", str(n)], stdin=p3.stdout,
                          stdout=subprocess.PIPE, text=True)
    p3.stdout.close()
    out, _ = p4.communicate()
    result = []
    for line in out.splitlines():
        count, url = line.split(None, 1)
        result.append((int(count), url))
    return result


def mib(nbytes):
    return nbytes / (1024 * 1024)


# --------------------------------------------------------------------------- #
# Main.
# --------------------------------------------------------------------------- #

def main():
    print("Counting URL hits two ways: an in-memory hash table vs. `sort | uniq -c`.")
    print(f"Line count is held CONSTANT at {LINE_COUNT:,} lines for every dataset;")
    print("only the number of DISTINCT keys changes.\n")

    with tempfile.TemporaryDirectory(prefix="ddia_ch11_") as tmp:
        rows = []
        hot_topn_file = None
        for distinct in DISTINCTS:
            path = os.path.join(tmp, f"log_{distinct}.txt")
            generate(path, distinct)
            if distinct == 1_000:
                hot_topn_file = path
            peak, counted, top = run_hash_worker(path)
            sort_rss = measure_sort_rss(path)
            fsize = os.path.getsize(path)
            rows.append((distinct, counted, peak, sort_rss, fsize, top))
            print(f"  built + measured dataset: {distinct:>9,} distinct keys")
        print()

        # ---- The table: distinct keys vs peak RSS, line count held constant ----
        print(f"Peak memory vs. distinct-key count  (input fixed at {LINE_COUNT:,} lines):")
        print()
        print(f"    {'distinct keys':>13}  {'lines':>11}  {'input':>9}  "
              f"{'hash-table RSS':>16}  {'sort RSS':>12}")
        print(f"    {'-'*13}  {'-'*11}  {'-'*9}  {'-'*16}  {'-'*12}")
        for distinct, counted, peak, sort_rss, fsize, _ in rows:
            hash_col = f"{mib(peak):8.1f} MiB"
            sort_col = f"{mib(sort_rss):6.1f} MiB" if sort_rss else "     n/a"
            print(f"    {counted:>13,}  {LINE_COUNT:>11,}  {mib(fsize):6.1f} MiB  "
                  f"{hash_col:>16}  {sort_col:>12}")
        print()
        base = rows[0][2]
        top_row = rows[-1]
        print(f"    hash table: {mib(base):.1f} MiB for 1 distinct key -> "
              f"{mib(top_row[2]):.1f} MiB for {top_row[1]:,} distinct "
              f"({top_row[2] / base:.1f}x) -- tracks DISTINCT KEYS, same line count.")
        print(f"    sort:       {mib(rows[0][3]):.1f} -> {mib(top_row[3]):.1f} MiB over the same "
              f"millionfold cardinality swing ({top_row[3] / rows[0][3]:.2f}x) -- ")
        print(f"                insensitive to distinct-key count; its RSS tracks input bytes "
              f"(BSD `sort` mmaps the file).\n")

        # ---- Same top-N from both methods on the skewed 1,000-key dataset ----
        print("Both methods agree on the answer. Top 5 URLs of the 1,000-key dataset:")
        print()
        hash_top = run_hash_worker(hot_topn_file)[2]
        sort_top = sort_pipeline_topn(hot_topn_file)
        print(f"    {'hash: counts[url] += 1':<34}  {'sort | uniq -c | sort -rn | head':<34}")
        print(f"    {'-'*32}    {'-'*34}")
        for (hc, hu), (sc, su) in zip(hash_top, sort_top):
            left = f"{hc:>7,}  {hu.rsplit('/', 1)[-1]}"
            right = f"{sc:>7,}  {su.rsplit('/', 1)[-1]}"
            print(f"    {left:<34}  {right:<34}")
        agree = [c for c, _ in hash_top] == [c for c, _ in sort_top]
        print()
        print(f"    identical top-5 counts: {agree}")


if __name__ == "__main__":
    if len(sys.argv) == 3 and sys.argv[1] == "--hash":
        hash_aggregate(sys.argv[2])
    else:
        main()