"""Stamp the background moments onto DiSCo's shape-moment layout (dmipy-sim's shape_moments.stamp_background): one
pass over the columnar source, the bg_<g>.npy columns and the manifest's background and RF schedules written into
DIR, which holds the layout's manifest.json and tiles.npy (the rest of the layout is not read and not touched; copy
the new columns and manifest beside it). Measured on a CPU host, from the Hub: 42.4 GB read, 1,016 s.

    python tools/stamp_disco_background.py DIR
"""
import json, os, sys, time
import numpy as np
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import dmipy_sim as d
from dmipy_sim.replay.shape_moments import stamp_background
out = sys.argv[1]
man = json.load(open(os.path.join(out, "manifest.json")))
shapes = {}
for name, rec in man["shapes"].items():
    fam, kw = rec["build_spec"]
    kw = dict(kw); dirs = kw.pop("gradient_directions")
    if fam == "pgse":
        shapes[name] = d.pgse(dirs, kw["delta"], kw["Delta"], gradient_strengths=kw["gradient_strengths"], TE=kw["TE"], n_t=kw["n_t"], slew_rate=np.inf)
    else:
        shapes[name] = d.pgste(dirs, kw["delta"], kw["TM"], gradient_strengths=kw["gradient_strengths"], n_t=kw["n_t"], slew_rate=np.inf,
                               ste_flip_angles=tuple(kw["ste_flip_angles"]))
last = [0.0]
def progress(rows, nbytes, secs):
    if secs - last[0] > 60:
        last[0] = secs
        print(time.strftime("%H:%M:%S"), f"{rows:,} rows, {nbytes / 1e9:.1f} GB, {secs:.0f} s, {nbytes / max(secs, 1) / 1e6:.0f} MB/s", flush=True)
m = stamp_background(out, "hf://SubstrateCommons/disco-replay/disco", shapes, workers=16, progress=progress)
print("DONE", json.dumps(m["background"]), json.dumps(m["stamps"][-1]), flush=True)