"""Model-free allocation/cache receipts from the installed pinned MTPLX runtime.

No replacement allocator, timings, model loading or downloads. stdout is JSONL.
Physical addresses are replaced with run-local identities before emission.
"""
import gc
import json
import os

import mlx.core as mx
import numpy as np


def main():
    assert mx.__version__ == "0.32.2"
    arrays = {}
    identities = {}
    page = os.sysconf("SC_PAGE_SIZE")
    mx.set_cache_limit(4 * 1024 * 1024)
    mx.clear_cache()
    assert mx.get_active_memory() == 0

    def emit(op, **fields):
        print(json.dumps(dict(op=op, active=mx.get_active_memory(),
                              cached=mx.get_cache_memory(), **fields)), flush=True)

    def alloc(name, size):
        before = mx.get_active_memory()
        arrays[name] = mx.array(np.zeros(size, dtype=np.uint8))
        mx.eval(arrays[name])
        mx.synchronize()
        pointer = np.asarray(arrays[name]).__array_interface__["data"][0] if size else None
        identity = identities.setdefault(pointer, len(identities)) if size else None
        emit("alloc", name=name, size=size, buffer=identity,
             storage=mx.get_active_memory() - before)

    def free(name):
        del arrays[name]
        gc.collect()
        mx.synchronize()
        emit("free", name=name)

    def clear():
        assert not arrays
        mx.clear_cache()
        emit("clear")

    def limit(size):
        mx.set_cache_limit(size)
        emit("limit", size=size)

    emit("init", page=page, limit=4 * 1024 * 1024, runtime=mx.__version__)
    for i, size in enumerate([0, 1, 255, 256, 257, page - 1, page, page + 1,
                              2 * page - 1, 2 * page, 2 * page + 1]):
        alloc(str(i), size)
    for name in list(arrays):
        free(name)
    for i, size in enumerate([256, 1, page, page + 1, 2 * page + 1]):
        alloc(str(i), size)
    for name in list(arrays):
        free(name)
    clear()

    # Equal-size multimap entries reuse the oldest insertion, not the newest.
    for name in "abc":
        alloc(name, 257)
    for name in "bac":
        free(name)
    for name in "xyz":
        alloc(name, 256)
    for name in "xyz":
        free(name)
    clear()

    # Strict upper bounds: min(2 * request, request + 2 * page).
    for stored, requested in [(512, 256), (511, 256),
                              (4 * page, 2 * page), (3 * page, 2 * page),
                              (6 * page, 4 * page), (5 * page, 4 * page)]:
        alloc("old", stored)
        free("old")
        alloc("new", requested)
        free("new")
        clear()

    for name, size in zip("abcd", [64, 128, 256, 512]):
        alloc(name, size)
    for name in "cadb":
        free(name)
    limit(600)
    alloc("e", 2048)  # trim from oldest, overshooting by complete buffers
    free("e")       # free checks the pre-insertion size, allowing overshoot
    alloc("f", 10000)
    free("f")
    limit(1)
    alloc("g", 20000)  # >=90% requested release clears the whole pool
    free("g")
    clear()
    limit(4096)
    alloc("old", 64)
    free("old")
    limit(0)  # changing the limit does not clear; reuse is still attempted first
    alloc("new", 64)
    free("new")
    clear()


if __name__ == "__main__":
    main()