# /// script
# requires-python = ">=3.10"
# dependencies = ["numpy"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Write amplification: same-age against mixed-age writes.

A flash translation layer simulator (not a real drive): erase blocks of 128
pages, out-of-place writes, greedy garbage collection that copies the live
pages out of the block with the fewest. The drive is filled to 50-90% of its
physical pages with logical data, then overwritten for several drive-writes
in two patterns: (a) same-age, whole files written and later deleted together
(data that dies together sits together); (b) mixed-age, single pages
overwritten at random. Write amplification = flash pages written ÷ host
pages written.
"""
import argparse, json, random
from collections import deque

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--blocks", type=int, default=512, help="Erase blocks.")
p.add_argument("--pages", type=int, default=128, help="Pages per block.")
p.add_argument("--fills", default="0.5,0.7,0.8,0.9", help="Logical fill levels (share of physical pages).")
p.add_argument("--drive-writes", type=float, default=4, help="Overwrite volume, in multiples of the logical data.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()

class FTL:
    def __init__(s):
        s.B, s.P = args.blocks, args.pages
        s.valid = [0] * s.B; s.owner = [[None] * s.P for _ in range(s.B)]; s.where = {}
        s.free = deque(range(s.B)); s.cur = s.free.popleft(); s.off = 0; s.flash = 0; s.host = 0
    def _append(s, lpn):
        if s.off == s.P:
            if len(s.free) < 2: s._gc()
            s.cur = s.free.popleft(); s.off = 0
        s.owner[s.cur][s.off] = lpn; s.where[lpn] = (s.cur, s.off); s.valid[s.cur] += 1; s.off += 1; s.flash += 1
    def _invalidate(s, lpn):
        if lpn in s.where:
            b, o = s.where.pop(lpn); s.owner[b][o] = None; s.valid[b] -= 1
    def _gc(s):
        while len(s.free) < 2:
            victim = min((b for b in range(s.B) if b != s.cur and b not in s.free), key=lambda b: s.valid[b])
            live = [l for l in s.owner[victim] if l is not None]
            for l in live: s._invalidate(l)
            s.owner[victim] = [None] * s.P; s.valid[victim] = 0; s.free.append(victim)
            for l in live: s._append(l)                               # copy live pages out: the amplification
    def write(s, lpn): s._invalidate(lpn); s._append(lpn); s.host += 1
    def trim(s, lpn): s._invalidate(lpn)

rows = []
for fill in map(float, args.fills.split(",")):
    logical = int(args.blocks * args.pages * fill)
    for pattern in ("same-age (whole files)", "mixed-age (random pages)"):
        rng = random.Random(args.seed); f = FTL()
        file_pages = args.pages * 4; files = deque()
        for start in range(0, logical, file_pages):                    # initial fill, file by file
            pages = list(range(start, min(start + file_pages, logical))); files.append(pages)
            for l in pages: f.write(l)
        f.flash = f.host = 0; total = int(args.drive_writes * logical)
        while f.host < total:
            if pattern.startswith("same"):
                old = files.popleft()                                  # the oldest file is deleted whole...
                for l in old: f.trim(l)
                for l in old: f.write(l)                               # ...and a new one written in its place
                files.append(old)
            else:
                f.write(rng.randrange(logical))
        rows.append({"fill_pct": fill * 100, "pattern": pattern, "write_amplification": f.flash / f.host})
for r in rows: print(f"fill {r['fill_pct']:3.0f}%  {r['pattern']:<26} write amplification {r['write_amplification']:.2f}")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"wa": rows}}, f)