# from pxr import Usd, UsdGeom
# stage = Usd.Stage.Open('workflows/simbox/example_assets/task/sort_the_rubbish/recyclable_garbage/bottle_0/Aligned_obj.usd')
# print('metersPerUnit:', UsdGeom.GetStageMetersPerUnit(stage))
# print('upAxis:', UsdGeom.GetStageUpAxis(stage))
# dp = stage.GetDefaultPrim()
# print('defaultPrim:', dp.GetPath() if dp else 'None')
# xf = UsdGeom.Xformable(dp)
# for op in xf.GetOrderedXformOps():
#     print(f'  xformOp: {op.GetName()} = {op.Get()}')
# for child in dp.GetChildren():
#     cxf = UsdGeom.Xformable(child)
#     ops = cxf.GetOrderedXformOps()
#     if ops:
#         for op in ops:
#             print(f'  child {child.GetPath()} xformOp: {op.GetName()} = {op.Get()}')

# from pxr import Usd, UsdGeom                                                                                                                                        
# stage = Usd.Stage.Open('workflows/simbox/example_assets/task/sort_the_rubbish/recyclable_garbage/bottle_0/Aligned_obj.usd')
# print('metersPerUnit:', UsdGeom.GetStageMetersPerUnit(stage))                                                                                                       
# print('upAxis:', UsdGeom.GetStageUpAxis(stage))                                                                                                                     
# for prim in stage.Traverse():                                                                                                                                       
#     print(f'  {prim.GetPath()} type={prim.GetTypeName()}')                                                                                                          
#     xf = UsdGeom.Xformable(prim)                                                                                                                                    
#     if xf:                                                                                                                                                          
#         for op in xf.GetOrderedXformOps():                                                                                                                          
#             print(f'    {op.GetName()} = {op.Get()}')

from pxr import Usd, UsdGeom                                                                                                                                        
import yaml, os                                                                                                                                                     
                                                                                                                                                                    
with open('workflows/simbox/core/configs/tasks/example/sort_the_rubbish.yaml') as f:                                                                                
    cfg = yaml.safe_load(f)                                                                                                                                         
                                                                                                                                                                    
tasks = cfg.get('tasks', [cfg])                                                                                                                                     
for task in tasks:                                                                                                                                                  
    asset_root = task.get('asset_root', '')                                                                                                                         
    for obj in task.get('objects', []):                                                                                                                             
        name = obj['name']             
        path = obj.get('path', '')                                                                                                                                  
        full = os.path.join(asset_root, path)                                                                                                                       
        print(f'=== {name}: {full} ===')                                                                                                                            
        try:                                                                                                                                                        
            stage = Usd.Stage.Open(full)                                                                                                                            
            print(f'  metersPerUnit: {UsdGeom.GetStageMetersPerUnit(stage)}')                                                                                       
            print(f'  upAxis: {UsdGeom.GetStageUpAxis(stage)}')                                                                                                     
        except Exception as e:                                                                                                                                      
            print(f'  ERROR: {e}')

# ===== Check mesh vertex ranges =====
print("\n===== Mesh vertex ranges =====")
import numpy as np
usd_files = {
    'pick_object_right': 'workflows/simbox/example_assets/task/sort_the_rubbish/recyclable_garbage/bottle_0/Aligned_obj.usd',
    'gso_box_right': 'workflows/simbox/example_assets/task/sort_the_rubbish/garbage_can/recyclable_can/Aligned_obj.usd',
}
for name, path in usd_files.items():
    stage = Usd.Stage.Open(path)
    for prim in stage.Traverse():
        if prim.GetTypeName() == 'Mesh':
            mesh = UsdGeom.Mesh(prim)
            pts = mesh.GetPointsAttr().Get()
            if pts:
                pts = np.array(pts)
                print(f'{name} mesh={prim.GetPath()} npts={len(pts)}')
                print(f'  X: [{pts[:,0].min():.4f}, {pts[:,0].max():.4f}]')
                print(f'  Y: [{pts[:,1].min():.4f}, {pts[:,1].max():.4f}]  (upAxis=Y, this is height)')
                print(f'  Z: [{pts[:,2].min():.4f}, {pts[:,2].max():.4f}]')
                print(f'  Y_max * 0.01 = {pts[:,1].max()*0.01:.6f} m')
                print(f'  Y_extent = {pts[:,1].max()-pts[:,1].min():.4f}')
                print(f'  Y_max (raw cm) = {pts[:,1].max():.4f}')
                # 109.75 / Y_max ?
                if pts[:,1].max() > 0:
                    print(f'  109.75 / Y_max = {109.75 / pts[:,1].max():.4f}')
                    print(f'  109.75 / (Y_max*0.01) = {109.75 / (pts[:,1].max()*0.01):.4f}')

# ===== Fix USD metadata: set metersPerUnit=1.0 and upAxis=Z =====
print("\n===== Fixing USD metadata =====")
import shutil

fix_files = [
    'workflows/simbox/example_assets/task/sort_the_rubbish/recyclable_garbage/bottle_0/Aligned_obj.usd',
    'workflows/simbox/example_assets/task/sort_the_rubbish/garbage_can/recyclable_can/Aligned_obj.usd',
    'workflows/simbox/example_assets/task/sort_the_rubbish/garbage_can/nonrecyclable_can/Aligned_obj.usd',
]
# Also try to find pick_object_left USD
pick_left_path = 'workflows/simbox/example_assets/task/sort_the_rubbish/non_recyclable_garbage/obj_0/Aligned_obj.usd'
if os.path.exists(pick_left_path):
    fix_files.append(pick_left_path)

for fpath in fix_files:
    if not os.path.exists(fpath):
        print(f'  SKIP (not found): {fpath}')
        continue
    # Backup
    bak = fpath + '.bak'
    if not os.path.exists(bak):
        shutil.copy2(fpath, bak)
        print(f'  Backed up: {fpath} -> {bak}')
    else:
        print(f'  Backup already exists: {bak}')

    stage = Usd.Stage.Open(fpath)
    old_mpu = UsdGeom.GetStageMetersPerUnit(stage)
    old_up = UsdGeom.GetStageUpAxis(stage)
    print(f'  {fpath}: old metersPerUnit={old_mpu}, upAxis={old_up}')

    UsdGeom.SetStageMetersPerUnit(stage, 1.0)
    UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
    stage.GetRootLayer().Save()

    # Verify
    stage2 = Usd.Stage.Open(fpath)
    print(f'  -> new metersPerUnit={UsdGeom.GetStageMetersPerUnit(stage2)}, upAxis={UsdGeom.GetStageUpAxis(stage2)}')

print("\nDone! Now run launcher to test if Z is normal.")
print("To restore: for each .bak file, copy it back over the .usd file")