import sys, time, resource
import numpy as np
sys.path.insert(0, "scripts")
from common import UP, fcp, load_mesh
sys.modules["fcp.encode"].MAX_ATOMS = 1 << 24
m = load_mesh(sys.argv[1]); L = int(sys.argv[2])
if sys.argv[3] == "lowmem":
    import lowmem  # noqa
t = time.perf_counter()
d = fcp.encode(m, atoms=1 << 24, layers=L, N=64, M=64, up=UP, device="cuda", parts="tracked-fit")
np.save(f"cache/lm_{sys.argv[3]}.npy", np.column_stack([d.atoms[0].astype(float), d.atoms[1]]))
print(sys.argv[3], f"{time.perf_counter()-t:.1f} s, peak {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1e6:.2f} GB, atoms {d.n_atoms}, parts {len(d.parts)}")