"""AstrAI promo: Paged KV Cache — matching astrai/inference/cache.py & scheduler.py."""

from manim import *

Text.set_default(font="Times New Roman")


class PrefixCache(Scene):
    """Animates PagedCache exact logic: alloc→write→free with real code details."""

    def _small_text(self, text, size=10, color=GRAY, **kwargs):
        return Text(text, font_size=size, color=color, **kwargs)

    def _page_box(self, pos, label, color, sz=0.44):
        s = Square(side_length=sz, color=color, fill_opacity=0.12, stroke_width=1.6)
        s.move_to(pos)
        lbl = self._small_text(label, 10, color).move_to(pos)
        return VGroup(s, lbl)

    def _make_ptable(self, x, y, rows, color, label):
        """Draw OS-style page table. Returns (outline, entry_group, phy_cell_centers)."""
        n = len(rows)
        w = 2.6
        h = 0.40
        vx = w * 0.40

        group = VGroup()

        # Label
        lbl = self._small_text(label, 12, color, weight=BOLD)
        lbl.move_to([x + w / 2, y + 0.15, 0])
        group.add(lbl)

        # Outer rect (header+data rows)
        r = Rectangle(width=w, height=h * (n + 1), color=color, stroke_width=1.8, fill_opacity=0.06)
        r.move_to([x + w / 2, y - h * (n + 1) / 2, 0])
        group.add(r)

        # H dividers
        for ri in range(1, n + 1):
            ly = y - ri * h
            group.add(Line([x, ly, 0], [x + w, ly, 0], color=color, stroke_width=0.8))

        # V divider
        group.add(Line([x + vx, y, 0], [x + vx, y - h * (n + 1), 0], color=color, stroke_width=0.8))

        # Column headers
        lh = self._small_text("logical", 10, color, weight=BOLD)
        lh.move_to([x + vx / 2, y - h / 2, 0])
        ph = self._small_text("physical", 10, color, weight=BOLD)
        ph.move_to([x + vx + (w - vx) / 2, y - h / 2, 0])
        group.add(lh, ph)

        # Row entries
        entries = VGroup()
        phy_centers = []
        for i, (lg, phy) in enumerate(rows):
            cy = y - h / 2 - (i + 1) * h
            lt = self._small_text(lg, 11, color)
            lt.move_to([x + vx / 2, cy, 0])
            pt = self._small_text(phy, 11, color)
            pt.move_to([x + vx + (w - vx) / 2, cy, 0])
            entries.add(VGroup(lt, pt))
            phy_centers.append([x + vx + (w - vx) / 2, cy, 0])

        return group, entries, phy_centers, r

    def construct(self):
        # ── Scene setup ──
        title = self._small_text("Paged KV Cache  —  astrai/inference/cache.py", 26, BLUE)
        title.to_edge(UP, buff=0.15)
        self.play(Write(title))
        self.wait(0.1)

        right_x = 4.5
        step_y = 3.0

        def step_msg(text, color=YELLOW):
            nonlocal step_y
            m = self._small_text(text, 12, color)
            m.move_to([right_x, step_y, 0])
            step_y -= 0.55
            self.play(Write(m))
            return m

        def fade(m):
            self.play(FadeOut(m))

        # ═══════════════════════════════════════════
        #  Phase 0: Initialize PagedCache
        # ═══════════════════════════════════════════
        s0 = step_msg("PagedCache(max_batch=8, page_size=64, head_dim=64, n_kv_heads=4)")

        # Show constructor signature
        ctor = self._small_text(
            'PagedCache(n_layers=24, n_pages=8, page_size=64,\n'
            '          n_kv_heads=4, head_dim=64, device=cuda, dtype=bfloat16)',
            9, GRAY
        )
        ctor.move_to([-5.5, 2.4, 0])
        self.play(Write(ctor))

        # Show the tensor shape
        tensor_shape = self._small_text(
            'k_cache: [24, 8, 64, 4, 64]    v_cache: [24, 8, 64, 4, 64]',
            8, GRAY
        )
        tensor_shape.move_to([-5.5, 2.05, 0])
        self.play(Write(tensor_shape))

        self.wait(0.3)
        fade(s0)
        # Keep ctor & tensor_shape visible

        # ═══════════════════════════════════════════
        #  Physical page pool
        # ═══════════════════════════════════════════
        s1 = step_msg("8 page frames, refs[8] = [0,0,0,0,0,0,0,0], free_mask = (1<<8)-1 = 0xFF")

        pool_y = 1.45
        pool_x0 = -3.8
        pool_sp = 0.68
        pool_pages = []
        pool_pos = []
        for i in range(8):
            x = pool_x0 + i * pool_sp
            pos = np.array([x, pool_y, 0])
            pool_pos.append(pos)
            pb = self._page_box(pos, str(i), GRAY)
            pool_pages.append(pb)
            self.play(FadeIn(pb, scale=0.5), run_time=0.04)
        self.wait(0.1)

        # Bracket
        brack = Brace(VGroup(*[p[0] for p in pool_pages]), DOWN, buff=0.05)
        blbl = Text("page frames [0..7]  each holds 64 KV slots", font_size=10, color=GRAY)
        blbl.next_to(brack, DOWN, buff=0.02)
        self.play(Create(brack), Write(blbl), run_time=0.25)

        # Free bitmask — with bit position labels
        mask_y = 0.55
        bit_labels = VGroup()
        for i in range(8):
            bl = self._small_text(f"bit{i}", 7, DARK_GRAY)
            bl.move_to([pool_x0 + i * pool_sp, mask_y + 0.25, 0])
            bit_labels.add(bl)
        self.play(Write(bit_labels), run_time=0.2)

        mask = self._small_text("11111111  (1 = free, 0 = alloc)", 11, GRAY)
        mask.move_to([-1.5, mask_y, 0])
        self.play(Write(mask))
        self.wait(0.15)

        refs_lbl = self._small_text("refs = [0,0,0,0,0,0,0,0]", 9, GRAY)
        refs_lbl.move_to([-1.5, mask_y - 0.3, 0])
        self.play(Write(refs_lbl))

        def update_mask(bits, desc):
            m2 = self._small_text(f"{bits}  (1 = free, 0 = alloc)", 11, GRAY)
            m2.move_to([-1.5, mask_y, 0])
            self.play(Transform(mask, m2))
            if desc:
                self.play(Transform(refs_lbl, self._small_text(desc, 9, GRAY).move_to([-1.5, mask_y - 0.3, 0])))

        fade(s1)

        # ═══════════════════════════════════════════
        #  Phase 1: Cleanup — nothing to do initially
        # ═══════════════════════════════════════════
        s_phase = self._small_text("Phase 1: Cleanup  (no finished tasks)", 11, GRAY)
        s_phase.move_to([right_x, step_y, 0])
        step_y -= 0.4
        self.play(Write(s_phase))
        self.wait(0.2)

        # ═══════════════════════════════════════════
        #  Phase 2: Refill — Request A arrives
        # ═══════════════════════════════════════════
        step_y2 = step_y
        s2a = step_msg("Phase 2: Refill  —  Request A arrives (prompt_len=120 tokens)", GREEN)
        s2b = step_msg("_n_pages_for(120) = (120 + 64 - 1) // 64 = 183 // 64 = 2", GREEN)

        calc_box = Rectangle(width=4.8, height=1.0, color=GREEN_E, stroke_width=1.2, fill_opacity=0.05)
        calc_box.move_to([right_x - 0.3, step_y - 0.1, 0])
        calc_lines = VGroup(
            self._small_text("def _n_pages_for(n_tokens):", 9, GREEN),
            self._small_text("    return (n_tokens + page_size - 1) // page_size", 9, GREEN),
            self._small_text("_n_pages_for(120) = (120 + 64 - 1) // 64 = 2", 9, WHITE),
        )
        calc_lines.arrange(DOWN, buff=0.08, aligned_edge=LEFT)
        calc_lines.move_to(calc_box.get_center())
        calc_grp = VGroup(calc_box, calc_lines)
        self.play(Create(calc_grp), run_time=0.35)
        step_y -= 0.6

        # alloc() in action
        s2c = step_msg("alloc_n(2)  →  calls alloc() twice", GREEN)

        alloc_code = VGroup(
            self._small_text("def alloc(self) -> int:", 9, GREEN),
            self._small_text("    lsb = self._free_mask & -self._free_mask", 9, GREEN),
            self._small_text("    if lsb == 0: return -1", 9, GREEN),
            self._small_text("    idx = lsb.bit_length() - 1", 9, GREEN),
            self._small_text("    self._free_mask ^= lsb", 9, GREEN),
            self._small_text("    self._refs[idx] = 1", 9, GREEN),
            self._small_text("    return idx", 9, GREEN),
            self._small_text("", 6, GREEN),
            self._small_text("1st alloc: free_mask=11111111, lsb=1, idx=0  →  page 0", 9, WHITE),
            self._small_text("2nd alloc: free_mask=11111110, lsb=2, idx=1  →  page 1", 9, WHITE),
        )
        alloc_code.arrange(DOWN, buff=0.06, aligned_edge=LEFT)
        acbox = Rectangle(width=5.0, height=2.6, color=GREEN_E, stroke_width=1.0, fill_opacity=0.03)
        acbox.move_to([right_x - 0.2, step_y - 0.8, 0])
        alloc_code.move_to(acbox.get_center())
        alloc_grp = VGroup(acbox, alloc_code)
        self.play(Create(alloc_grp), run_time=0.4)
        step_y -= 1.8

        # Now draw page table A (2 rows)
        tblA_x = -4.4
        tblA_y = -0.35
        tblA_outline, tblA_entries, tblA_phys, tblA_rect = self._make_ptable(
            tblA_x, tblA_y,
            [("0", "0"), ("1", "1")],
            GREEN, "Task A  page_table"
        )
        self.play(Create(tblA_outline), run_time=0.35)
        for e in tblA_entries:
            self.play(FadeIn(e), run_time=0.08)

        # Alloc effect on pool + bitmask
        for idx in [0, 1]:
            pg = pool_pages[idx][0]
            self.play(pg.animate.set_fill(GREEN, opacity=0.35), run_time=0.1)
            flash = SurroundingRectangle(pool_pages[idx], color=GREEN, buff=0.04)
            self.play(Create(flash), run_time=0.05)
            self.play(FadeOut(flash), run_time=0.04)

        update_mask("11111100", "refs = [1,1,0,0,0,0,0,0]")
        self.wait(0.2)

        # mapping arrows
        arr_a0 = Arrow(
            [tblA_phys[0][0] + 0.1, tblA_phys[0][1], 0],
            [pool_pos[0][0], pool_pos[0][1] - 0.22, 0],
            color=GREEN, stroke_width=1.5, buff=0.03,
            max_tip_length_to_length_ratio=0.18,
        )
        arr_a1 = Arrow(
            [tblA_phys[1][0] + 0.1, tblA_phys[1][1], 0],
            [pool_pos[1][0], pool_pos[1][1] - 0.22, 0],
            color=GREEN, stroke_width=1.5, buff=0.03,
            max_tip_length_to_length_ratio=0.18,
        )
        self.play(GrowArrow(arr_a0), GrowArrow(arr_a1), run_time=0.35)
        self.wait(0.3)

        # ═══════════════════════════════════════════
        #  Phase 3: Prefill (write KV into allocated pages)
        # ═══════════════════════════════════════════
        s3a = step_msg("Phase 3: Prefill  —  write() KV into pages 0,1", GREEN)

        write_code = VGroup(
            self._small_text("def write(self, layer_id, page_table, start_pos, k, v):", 9, GREEN),
            self._small_text("    first_page = start_pos // page_size     # 0 // 64 = 0", 9, GREEN),
            self._small_text("    last_page  = (start_pos+seq_len-1)//page_size", 9, WHITE),
            self._small_text("    for pi in range(first_page, last_page+1):", 9, WHITE),
            self._small_text("        phys_pages = page_table[:, pi]", 9, WHITE),
            self._small_text("        chunk = min(page_start+page_size, start_pos+seq_len) -", 9, WHITE),
            self._small_text("                max(page_start, start_pos)", 9, WHITE),
            self._small_text("        k_cache[layer_id, phys_pages, offset:offset+chunk] = k", 9, GREEN),
        )
        write_code.arrange(DOWN, buff=0.05, aligned_edge=LEFT)
        wbox = Rectangle(width=5.0, height=2.3, color=GREEN_E, stroke_width=1.0, fill_opacity=0.03)
        wbox.move_to([right_x - 0.2, step_y - 0.4, 0])
        write_code.move_to(wbox.get_center())
        write_grp = VGroup(wbox, write_code)
        self.play(Create(write_grp), run_time=0.35)
        step_y -= 1.6

        # Show a KV block "written" onto the pages
        k_written = self._small_text("KV  written", 8, GREEN)
        k_written.move_to([pool_pos[0][0] + 0.5, pool_pos[0][1] + 0.6, 0])
        self.play(Write(k_written))
        self.wait(0.3)
        fade(k_written)

        s3b = step_msg("CacheView  bundles cache + page_table + total_len for attention", GREEN)
        cv_code = VGroup(
            self._small_text("class CacheView:", 9, GREEN),
            self._small_text("    def __init__(self, cache, page_table, total_len):", 9, GREEN),
            self._small_text('    def write(self, layer_id, start_pos, k, v):', 9, WHITE),
            self._small_text("        self._cache.write(layer_id, self._page_table, ...)", 9, WHITE),
            self._small_text('    def gather(self, layer_id):', 9, WHITE),
            self._small_text("        for pi in range(page_table.size(1)):", 9, WHITE),
            self._small_text("            phys_pages = page_table[:, pi]", 9, WHITE),
            self._small_text("            k_parts.append(k_cache[layer_id, phys_pages])", 9, WHITE),
            self._small_text("        k = torch.cat(k_parts, dim=1)", 9, GREEN),
        )
        cv_code.arrange(DOWN, buff=0.05, aligned_edge=LEFT)
        cvbox = Rectangle(width=5.0, height=2.5, color=GREEN_E, stroke_width=1.0, fill_opacity=0.03)
        cvbox.move_to([right_x - 0.2, step_y - 0.6, 0])
        cv_code.move_to(cvbox.get_center())
        cv_grp = VGroup(cvbox, cv_code)
        self.play(Create(cv_grp), run_time=0.35)
        step_y -= 1.8

        # ═══════════════════════════════════════════
        #  Request B arrives
        # ═══════════════════════════════════════════
        s4 = step_msg("Phase 2: Refill  —  Request B arrives (prompt_len=90)", ORANGE)
        s4b = step_msg("alloc_n(2)  →  pages 2, 3", ORANGE)

        tblB_x = -4.4
        tblB_y = -1.9
        tblB_outline, tblB_entries, tblB_phys, tblB_rect = self._make_ptable(
            tblB_x, tblB_y,
            [("0", "2"), ("1", "3")],
            ORANGE, "Task B  page_table"
        )
        self.play(Create(tblB_outline), run_time=0.3)
        for e in tblB_entries:
            self.play(FadeIn(e), run_time=0.08)

        for idx in [2, 3]:
            pg = pool_pages[idx][0]
            self.play(pg.animate.set_fill(ORANGE, opacity=0.35), run_time=0.1)
            flash = SurroundingRectangle(pool_pages[idx], color=ORANGE, buff=0.04)
            self.play(Create(flash), run_time=0.05)
            self.play(FadeOut(flash), run_time=0.04)

        update_mask("11110000", "refs = [1,1,1,1,0,0,0,0]")

        arr_b0 = Arrow(
            [tblB_phys[0][0] + 0.1, tblB_phys[0][1], 0],
            [pool_pos[2][0], pool_pos[2][1] - 0.22, 0],
            color=ORANGE, stroke_width=1.5, buff=0.03,
            max_tip_length_to_length_ratio=0.18,
        )
        arr_b1 = Arrow(
            [tblB_phys[1][0] + 0.1, tblB_phys[1][1], 0],
            [pool_pos[3][0], pool_pos[3][1] - 0.22, 0],
            color=ORANGE, stroke_width=1.5, buff=0.03,
            max_tip_length_to_length_ratio=0.18,
        )
        self.play(GrowArrow(arr_b0), GrowArrow(arr_b1), run_time=0.35)
        self.wait(0.3)

        # ═══════════════════════════════════════════
        #  Phase 4: Decode — on-demand page growth
        # ═══════════════════════════════════════════
        s5 = step_msg("Phase 4: Decode  —  Task A generates more tokens", PINK)
        s5b = step_msg("total_len grows: 120 → 180 tokens", PINK)
        s5c = step_msg("_n_pages_for(180) = (180+63)//64 = 243//64 = 3", PINK)

        # _maybe_alloc_page logic
        map_code = VGroup(
            self._small_text("def _maybe_alloc_page(self, task, pos):", 9, PINK),
            self._small_text("    needed = _n_pages_for(pos + 1)   # _n_pages_for(181) = 3", 9, PINK),
            self._small_text("    while task.n_pages < needed:          # 2 < 3", 9, PINK),
            self._small_text("        p = self.page_cache.alloc()       # alloc page 4", 9, PINK),
            self._small_text("        task.page_table.append(p)", 9, PINK),
            self._small_text("        task.n_pages += 1                 # 2 → 3", 9, PINK),
            self._small_text("", 5, PINK),
            self._small_text("page_table_A: [0, 1]  →  [0, 1, 4]", 9, WHITE),
        )
        map_code.arrange(DOWN, buff=0.05, aligned_edge=LEFT)
        mbox = Rectangle(width=5.0, height=2.2, color=PINK, stroke_width=1.0, fill_opacity=0.03)
        mbox.move_to([right_x - 0.2, step_y - 0.2, 0])
        map_code.move_to(mbox.get_center())
        map_grp = VGroup(mbox, map_code)
        self.play(Create(map_grp), run_time=0.35)
        step_y -= 1.6

        # Alloc page 4
        idx = 4
        pg = pool_pages[idx][0]
        self.play(pg.animate.set_fill(PINK, opacity=0.35), run_time=0.1)
        flash = SurroundingRectangle(pool_pages[idx], color=PINK, buff=0.04)
        self.play(Create(flash), run_time=0.05)
        self.play(FadeOut(flash), run_time=0.04)

        update_mask("11100000", "refs = [1,1,1,1,1,0,0,0]")

        # Expand Table A: add row 3
        new_bottom = tblA_y - 0.40 * 4  # 3 data rows + 1 header
        line_3 = Line(
            [tblA_x, tblA_y - 0.40 * 3, 0],
            [tblA_x + 2.6, tblA_y - 0.40 * 3, 0],
            color=GREEN, stroke_width=0.8,
        )
        lc3 = [tblA_x + 2.6 * 0.20, tblA_y - 0.40 * 3 - 0.20, 0]
        pc3 = [tblA_x + 2.6 * 0.40 + (2.6 - 2.6 * 0.40) / 2, tblA_y - 0.40 * 3 - 0.20, 0]
        lt3 = self._small_text("2", 11, PINK).move_to(lc3)
        pt3 = self._small_text("4", 11, PINK).move_to(pc3)
        self.play(Create(line_3), run_time=0.08)
        self.play(FadeIn(lt3), FadeIn(pt3), run_time=0.08)

        arr_c = Arrow(
            [pc3[0] + 0.1, pc3[1], 0],
            [pool_pos[4][0], pool_pos[4][1] - 0.22, 0],
            color=PINK, stroke_width=1.5, buff=0.03,
            max_tip_length_to_length_ratio=0.18,
        )
        self.play(GrowArrow(arr_c), run_time=0.2)
        self.wait(0.3)

        # Highlight: page_table list conversion
        pt_list_old = self._small_text("page_table_A: [0, 1]  (2 pages → can hold 128 tokens)", 10, GREEN)
        pt_list_old.move_to([-4.4, tblB_y - 1.35, 0])
        self.play(Write(pt_list_old))

        pt_list_new = self._small_text("page_table_A: [0, 1, 4]  (3 pages → can hold 192 tokens)", 10, GREEN)
        pt_list_new.move_to([-4.4, tblB_y - 1.7, 0])
        self.play(Write(pt_list_new))
        self.wait(0.4)

        # ═══════════════════════════════════════════
        #  Task A finished → free()
        # ═══════════════════════════════════════════
        s6 = step_msg("Task A done  →  free pages 0, 1, 4", YELLOW)
        s6b = step_msg("free()  →  refs[idx] -= 1; if refs[idx]==0: mask |= 1<<idx", YELLOW)

        free_code = VGroup(
            self._small_text("def free(self, idx):", 9, YELLOW),
            self._small_text("    self._refs[idx] -= 1", 9, YELLOW),
            self._small_text("    if self._refs[idx] == 0:", 9, YELLOW),
            self._small_text("        self._free_mask |= 1 << idx", 9, YELLOW),
        )
        free_code.arrange(DOWN, buff=0.06, aligned_edge=LEFT)
        fbox = Rectangle(width=4.2, height=1.3, color=YELLOW, stroke_width=1.0, fill_opacity=0.03)
        fbox.move_to([right_x - 0.2, step_y - 0.2, 0])
        free_code.move_to(fbox.get_center())
        free_grp = VGroup(fbox, free_code)
        self.play(Create(free_grp), run_time=0.3)
        step_y -= 0.9

        for idx in [0, 1]:
            pg = pool_pages[idx][0]
            self.play(pg.animate.set_fill(GRAY, opacity=0.12), run_time=0.08)
            flash = SurroundingRectangle(pool_pages[idx], color=YELLOW, buff=0.04)
            self.play(Create(flash), run_time=0.06)
            self.play(FadeOut(flash), run_time=0.04)

        update_mask("11100011", "refs = [0,0,1,1,1,0,0,0]")

        for idx in [4]:
            pg = pool_pages[idx][0]
            self.play(pg.animate.set_fill(GRAY, opacity=0.12), run_time=0.08)
            flash = SurroundingRectangle(pool_pages[idx], color=YELLOW, buff=0.04)
            self.play(Create(flash), run_time=0.06)
            self.play(FadeOut(flash), run_time=0.04)

        update_mask("11110011", "refs = [0,0,1,1,0,0,0,0]")

        # Show Cleanup phase removes finished tasks
        s_cleanup = self._small_text("Phase 1: Cleanup  —  Task A removed from active list", 11, GRAY)
        s_cleanup.move_to([right_x, step_y, 0])
        step_y -= 0.3
        self.play(Write(s_cleanup))

        pt_list_done = self._small_text("page_table_A cleared.  Frame 0,1,4 returned to pool.", 10, GRAY)
        pt_list_done.move_to([-4.4, tblB_y - 2.1, 0])
        self.play(Write(pt_list_done))
        self.wait(0.5)

        # ═══════════════════════════════════════════
        #  gather() demo — remaining Task B reads KV
        # ═══════════════════════════════════════════
        s7 = step_msg("Task B continues  —  gather() reads KV for attention", BLUE)

        gather_code = VGroup(
            self._small_text("def gather(self, layer_id):", 9, BLUE),
            self._small_text("    for pi in range(page_table.size(1)):", 9, BLUE),
            self._small_text("        phys_pages = page_table[:, pi]", 9, BLUE),
            self._small_text("        k_parts.append(k_cache[layer_id, phys_pages])", 9, BLUE),
            self._small_text("    k = torch.cat(k_parts, dim=1)", 9, BLUE),
            self._small_text("    return k, v", 9, BLUE),
            self._small_text("", 5, BLUE),
            self._small_text("gather reads pages [2, 3] → token positions [128..255]", 9, WHITE),
        )
        gather_code.arrange(DOWN, buff=0.05, aligned_edge=LEFT)
        gbox = Rectangle(width=5.0, height=2.0, color=BLUE, stroke_width=1.0, fill_opacity=0.03)
        gbox.move_to([right_x - 0.2, step_y - 0.4, 0])
        gather_code.move_to(gbox.get_center())
        gather_grp = VGroup(gbox, gather_code)
        self.play(Create(gather_grp), run_time=0.35)
        step_y -= 1.4

        # Highlight B's remaining pages
        for idx in [2, 3]:
            flash = SurroundingRectangle(pool_pages[idx], color=BLUE, buff=0.04)
            self.play(Create(flash), run_time=0.06)
            self.play(FadeOut(flash), run_time=0.04)
        self.wait(0.3)

        # ═══════════════════════════════════════════
        #  Summary
        # ═══════════════════════════════════════════
        self.wait(0.5)
        summary = self._small_text("Paged KV Cache  —  page-table-indirected, O(1) alloc/free, on-demand growth", 20, GREEN)
        summary.to_edge(DOWN, buff=0.4)
        self.play(Write(summary))

        benefits = VGroup(
            self._small_text("✓ No per-request KV pre-allocation — pages allocated on demand", 14, GREEN),
            self._small_text("✓ Page table decouples logical position from physical storage", 14, GREEN),
            self._small_text("✓ Ref-counted free → safe concurrent release across tasks", 14, GREEN),
            self._small_text("✓ Bitmask O(1) alloc/free — no fragmentation", 14, GREEN),
        )
        benefits.arrange(DOWN, buff=0.08, aligned_edge=LEFT)
        benefits.move_to([-5.0, -1.5, 0])
        self.play(Write(benefits), run_time=0.5)

        # Final mask
        summary_mask = self._small_text(
            "Final state: free_mask=11110011  free frames=0,1,4  in-use=2,3  total=8",
            10, GRAY
        )
        summary_mask.next_to(benefits, DOWN, buff=0.25, aligned_edge=LEFT)
        self.play(Write(summary_mask))

        self.wait(2)
        self.play(*[FadeOut(m) for m in self.mobjects])