"""
Generate test tube grasp annotations by adapting working bottle grasps.

Strategy: Load verified working bottle grasps, scale positions to match
test tube dimensions. The rotation conventions are already correct.

Bottle: ~5cm radius, ~10cm height
Test tube: ~1.5cm radius, ~10.6cm height (similar height, 3x smaller radius)

Usage:
    python migrate/gen_tube_grasp.py
"""
import numpy as np
import os

BOTTLE_GRASP = "workflows/simbox/example_assets/task/sort_the_rubbish/recyclable_garbage/bottle_0/Aligned_grasp_sparse.npy"
OUTPUT = "workflows/simbox/example_assets/task/pick_test_tube/test_tube/Aligned_grasp_sparse.npy"


def main():
    # Load working bottle grasps
    bottle = np.load(BOTTLE_GRASP)
    print(f"Loaded {len(bottle)} bottle grasps")
    print(f"Bottle pos range: X[{bottle[:,13].min():.4f},{bottle[:,13].max():.4f}] "
          f"Y[{bottle[:,14].min():.4f},{bottle[:,14].max():.4f}] "
          f"Z[{bottle[:,15].min():.4f},{bottle[:,15].max():.4f}]")

    tube = bottle.copy()

    # Scale XY positions (radius: bottle ~3cm -> tube ~1.5cm, factor ~0.5)
    # Keep Z positions similar (both ~10cm height)
    xy_scale = 0.5
    tube[:, 13] *= xy_scale  # X position
    tube[:, 14] *= xy_scale  # Y position
    # Z positions stay the same (similar height)

    # Reduce gripper width for smaller object
    # Bottle width: 0.044-0.10, Tube needs: 0.02-0.06
    tube[:, 1] = np.clip(tube[:, 1] * 0.6, 0.02, 0.06)

    print(f"\nTube pos range: X[{tube[:,13].min():.4f},{tube[:,13].max():.4f}] "
          f"Y[{tube[:,14].min():.4f},{tube[:,14].max():.4f}] "
          f"Z[{tube[:,15].min():.4f},{tube[:,15].max():.4f}]")
    print(f"Tube width range: [{tube[:,1].min():.3f},{tube[:,1].max():.3f}]")

    np.save(OUTPUT, tube)
    print(f"\nSaved {len(tube)} grasps to: {OUTPUT}")


if __name__ == "__main__":
    main()