"""Job-scoped instrumentation of installed TokenRig; no shared code modified.

Run inside the existing DGXA container. One prediction produces raw + voxel
postprocessed outputs, avoiding stochastic differences between the two trials.
"""
import hashlib
import os
from pathlib import Path
import runpy
import sys

job=Path(__file__).resolve().parent
source=Path('/workspace/SkinTokens/demo.py').read_text()
needle='        if use_postprocess:\n'
assert source.count(needle)==1
instrument='''        import numpy as np
        import os
        pair_folder = os.environ["TOKENRIG_PAIR_FOLDER"]
        np.savez_compressed(pair_folder + "/raw_prediction.npz", vertices=asset.vertices, faces=asset.faces, joints=asset.joints, parents=asset.parents, skin=asset.skin)
        post_bpy_payload("transfer", dict(source_asset=asset, target_path=asset.path, export_path=pair_folder + "/rigged_raw.glb", group_per_vertex=4))
        print("PAIR_RAW_EXPORTED", flush=True)
'''
source=source.replace(needle,instrument+needle)
needle2='            asset.normalize_skin()\n'
assert source.count(needle2)==1
source=source.replace(needle2,needle2+'''            np.savez_compressed(pair_folder + "/post_prediction.npz", vertices=asset.vertices, faces=asset.faces, joints=asset.joints, parents=asset.parents, skin=asset.skin)
            print("PAIR_POSTPROCESSED", flush=True)
''')
source='import random\nimport numpy as np\nimport torch\nrandom.seed(20261023)\nnp.random.seed(20261023)\ntorch.manual_seed(20261023)\n'+source
script=job/'instrumented_demo.py'
script.write_text(source)
(job/'installed_demo.sha256').write_text(hashlib.sha256(Path('/workspace/SkinTokens/demo.py').read_bytes()).hexdigest())
os.environ['TOKENRIG_PAIR_FOLDER']=str(job)
sys.argv=[str(script),'--input',str(job/'source.glb'),'--output',str(job/'rigged_post.glb'),
          '--use_skeleton','--use_transfer','--use_postprocess','--skeleton_template','original','--rig_cls','vroid']
runpy.run_path(str(script),run_name='__main__')