import json, os, sys, time, numpy as np, trimesh
from multiprocessing import Pool
sys.path.insert(0, '/root')
from caption_filter import is_single_object

root = '/root/lowpoly'
OUT = '/root/data'
os.makedirs(OUT, exist_ok=True)

paths = {}
for dirpath, _, names in os.walk(root):
    for n in names:
        if n.endswith('.obj'):
            paths[n[:-4]] = os.path.join(dirpath, n)
print('obj files on disk:', len(paths), flush=True)

caps = json.load(open('/root/captions.json'))
uids = sorted(u for u in caps if u in paths)
print('captioned meshes:', len(uids), flush=True)

def vox(u):
    try:
        m = trimesh.load(paths[u], force='mesh')
        v = m.vertices - m.vertices.mean(0)
        s = np.abs(v).max()
        m.vertices = v / (s if s > 0 else 1)
        g = m.voxelized(pitch=2.0/31).matrix.astype(np.uint8)
        pad = np.zeros((32,32,32), np.uint8)
        sh = [min(32,d) for d in g.shape]
        pad[:sh[0],:sh[1],:sh[2]] = g[:sh[0],:sh[1],:sh[2]]
        if pad.sum() == 0:
            return u, None
        return u, np.packbits(pad.reshape(-1))
    except Exception:
        return u, None

t0 = time.time()
ok_uids, rows, failed = [], [], 0
with Pool(32) as p:
    for i, (u, packed) in enumerate(p.imap(vox, uids, chunksize=64)):
        if packed is None:
            failed += 1
        else:
            ok_uids.append(u); rows.append(packed)
        if (i+1) % 5000 == 0:
            print(f'  voxelized {i+1}/{len(uids)}  failed {failed}  {time.time()-t0:.0f}s', flush=True)

vox_arr = np.stack(rows)
np.save(f'{OUT}/voxels_packed.npy', vox_arr)
captions = [caps[u] for u in ok_uids]
keep = [bool(is_single_object(c)) for c in captions]
json.dump({'uids': ok_uids, 'captions': captions, 'keep': keep},
          open(f'{OUT}/meta.json', 'w'))

occ = np.unpackbits(vox_arr, axis=1).mean(1) * 100
print(f'\ndone in {time.time()-t0:.0f}s')
print(f'meshes voxelized {len(ok_uids)}, failed/empty {failed}')
print(f'array {vox_arr.shape} {vox_arr.nbytes/2**20:.0f} MiB packed')
print(f'filter survivors {sum(keep)} = {sum(keep)/len(keep)*100:.1f}%')
print(f'surface occupancy mean {occ.mean():.2f}%  median {np.median(occ):.2f}%')
print('BUILDDONE', flush=True)