#!/usr/bin/env python3
"""Hands-on P06 — a MapReduce framework, assembled from six lego blocks."""
import os, random, statistics, tempfile
from collections import defaultdict
from _harness import block, run_all
rng = random.Random(0)
TEXT = ("the quick brown fox the lazy dog the fox runs " * 400).split()

@block(1, "Input splitting", "a split must never cut a record in half")
def b1(s, show):
    def split(words, m):
        per = max(1, len(words) // m)
        return [words[i:i+per] for i in range(0, len(words), per)]
    parts = split(TEXT, 8)
    assert sum(len(p) for p in parts) == len(TEXT), "lost or duplicated records"
    if show:
        print(f"  {len(TEXT)} words -> {len(parts)} splits of ~{len(parts[0])}")
        print(f"  conservation: sum(len) == len(input): "
              f"{sum(len(p) for p in parts) == len(TEXT)}")
        print("  on real files the boundary lands mid-line; each split reads to the")
        print("  next delimiter and skips a leading partial. Test a record that")
        print("  spans a boundary or you will silently lose or double it.")
    return {"split": split}

@block(2, "map + partition", "hash(key) % R decides the whole shuffle topology")
def b2(s, show):
    def do_map(words): return [(w, 1) for w in words]
    def partition(pairs, R):
        out = defaultdict(list)
        for k, v in pairs: out[hash(k) % R].append((k, v))
        return out
    parts = s["split"](TEXT, 8)
    R = 4
    m0 = do_map(parts[0]); p0 = partition(m0, R)
    if show:
        print(f"  map split 0: {len(parts[0])} words -> {len(m0)} pairs")
        print(f"  partitioned into R={R}: sizes {[len(p0[i]) for i in range(R)]}")
        print("  the SAME key always lands in the same partition, from every mapper.")
        print("  That is what makes the reduce side a merge instead of a join.")
    return {"do_map": do_map, "partition": partition, "R": R}

@block(3, "Shuffle: the all-to-all", "M x R transfers -- the pattern that scales worst")
def b3(s, show):
    def shuffle(map_outputs, R):
        bytes_moved = 0
        red_in = {r: [] for r in range(R)}
        for parts in map_outputs:
            for r in range(R):
                chunk = parts.get(r, [])
                red_in[r].extend(chunk)
                bytes_moved += sum(len(k) + 8 for k, _ in chunk)
        return red_in, bytes_moved
    splits = s["split"](TEXT, 8)
    mo = [s["partition"](s["do_map"](sp), s["R"]) for sp in splits]
    red_in, moved = shuffle(mo, s["R"])
    if show:
        print(f"  M={len(splits)} mappers x R={s['R']} reducers = "
              f"{len(splits)*s['R']} transfers")
        print(f"  {moved/1024:.1f} KB moved across the shuffle")
        print(f"  reducer input sizes: {[len(v) for v in red_in.values()]}")
        print("  all-to-all is why the shuttle dominates: it grows as M*R, and the")
        print("  network is the scarcest resource in the cluster.")
    return {"shuffle": shuffle}

@block(4, "Combiner", "pre-aggregate on the map side -- only valid if reduce is associative")
def b4(s, show):
    def combine(pairs):
        agg = defaultdict(int)
        for k, v in pairs: agg[k] += v
        return list(agg.items())
    splits = s["split"](TEXT, 8)
    raw = [s["partition"](s["do_map"](sp), s["R"]) for sp in splits]
    comb = [s["partition"](combine(s["do_map"](sp)), s["R"]) for sp in splits]
    _, b_raw = s["shuffle"](raw, s["R"])
    _, b_com = s["shuffle"](comb, s["R"])
    if show:
        print(f"  shuffle bytes without combiner: {b_raw/1024:>7.1f} KB")
        print(f"  shuffle bytes with combiner:    {b_com/1024:>7.1f} KB  "
              f"({b_raw/b_com:.0f}x less)")
        print("  valid ONLY because + is associative and commutative. A combiner")
        print("  applied to a non-associative reduce is silently wrong -- the")
        print("  framework must REJECT it, not trust you.")
    return {"combine": combine}

@block(5, "Reduce + atomic commit", "at-least-once execution plus an atomic rename = exactly-once effect")
def b5(s, show):
    D = tempfile.mkdtemp(prefix="h06-")
    def do_reduce(pairs):
        agg = defaultdict(int)
        for k, v in pairs: agg[k] += v
        return sorted(agg.items())
    def commit(result, rid, attempt):
        tmp = os.path.join(D, f"part-{rid}.attempt{attempt}.tmp")
        with open(tmp, "w") as f:
            for k, v in result: f.write(f"{k}\t{v}\n")
        os.replace(tmp, os.path.join(D, f"part-{rid}"))   # ATOMIC
        return os.path.join(D, f"part-{rid}")
    if show:
        splits = s["split"](TEXT, 8)
        mo = [s["partition"](s["combine"](s["do_map"](sp)), s["R"]) for sp in splits]
        red_in, _ = s["shuffle"](mo, s["R"])
        r0 = do_reduce(red_in[0])
        p1 = commit(r0, 0, attempt=1)
        p2 = commit(r0, 0, attempt=2)        # a duplicate/speculative task
        print(f"  reducer 0 produced {len(r0)} keys, e.g. {r0[:2]}")
        print(f"  attempt 1 and attempt 2 both committed -> one file: "
              f"{p1 == p2 and os.path.exists(p1)}")
        print(f"  no .tmp files left behind: "
              f"{not any(f.endswith('.tmp') for f in os.listdir(D))}")
        print("  write-temp-then-rename is why a duplicated task is harmless. Same")
        print("  trick as P03's segment flush and P07's checkpoint: make the state")
        print("  transition and the position advance ATOMIC.")
    return {"do_reduce": do_reduce, "commit": commit, "D": D}

@block(6, "Stragglers and backup tasks", "job time is a MAXIMUM, and maxima behave badly")
def b6(s, show):
    def job(ntasks=200, nworkers=20, frac=0.0, mult=1, backup=False, trials=400):
        out = []
        for _ in range(trials):
            t = [10.0 * (mult if rng.random() < frac else 1.0) for _ in range(ntasks)]
            if backup: t = [min(x, 20.0) for x in t]
            w = [0.0] * nworkers
            for x in sorted(t, reverse=True):
                i = w.index(min(w)); w[i] += x
            out.append(max(w))
        return statistics.fmean(out)
    if show:
        base = job()
        print(f"  {'scenario':<22}{'completion':>12}{'vs ideal':>10}{'w/ backup':>12}")
        for frac, mult, lbl in ((0.0,1,"no stragglers"),(0.01,10,"1% at 10x"),
                                (0.05,10,"5% at 10x"),(0.01,50,"1% at 50x")):
            a = job(frac=frac, mult=mult); b = job(frac=frac, mult=mult, backup=True)
            print(f"  {lbl:<22}{a:>10.1f}s{a/base:>9.2f}x{b:>10.1f}s")
        print("  two tasks in two hundred inflate the job 4.5x. Backup tasks recover")
        print("  nearly all of it. That is MapReduce section 3.6, generated not quoted.")
    return {"job": job}

def assembly(s):
    print("\nSix blocks = a batch framework. Run word count, then kill workers.\n")
    def run_job(text, M=8, R=4, kill_workers=0, use_combiner=True):
        splits = s["split"](text, M)
        map_out, retries = [], 0
        for i, sp in enumerate(splits):
            attempt = 0
            while True:
                attempt += 1
                # a "worker failure" loses the task's output; the coordinator retries
                if i < kill_workers and attempt == 1:
                    retries += 1; continue
                pairs = s["do_map"](sp)
                if use_combiner: pairs = s["combine"](pairs)
                map_out.append(s["partition"](pairs, R)); break
        red_in, moved = s["shuffle"](map_out, R)
        result = {}
        for r in range(R):
            for k, v in s["do_reduce"](red_in[r]): result[k] = v
        return result, retries, moved

    truth = {}
    for w in TEXT: truth[w] = truth.get(w, 0) + 1
    print(f"  {'run':<28}{'keys':>7}{'matches oracle':>17}{'retries':>9}")
    for lbl, kw, comb in (("clean", 0, True), ("3 workers killed", 3, True),
                          ("6 workers killed", 6, True), ("no combiner", 0, False)):
        res, ret, moved = run_job(TEXT, kill_workers=kw, use_combiner=comb)
        print(f"  {lbl:<28}{len(res):>7}{str(res == truth):>17}{ret:>9}")
    print("\n  Identical output under every fault schedule. That is the ONLY")
    print("  correctness criterion that matters, and it is possible because map and")
    print("  reduce are pure functions of their input -- the framework may re-run")
    print("  any task, anywhere, at any time.")
    print("\n  The restriction IS the feature. Let map read shared mutable state and")
    print("  retry, speculation and rescheduling all become unsound at once.")
    print("\n  Built: splitting -> map/partition -> shuffle -> combiner -> reduce +")
    print("  atomic commit -> straggler simulation.")
    print("  Missing, on the project page: real processes and sockets (m3-m5),")
    print("  coordinator checkpointing (m10), data locality (m11), the LATE")
    print("  scheduler (E5, where naive speculation makes things WORSE), and E12 --")
    print("  the hand-written comparison that is the report's thesis.")

if __name__ == "__main__":
    run_all(assembly, "HANDS-ON P06 — MapReduce, block by block")