"""
Global constants used across simbox workflows.

- CATEGORIES: per-category preferred Euler angles for data randomization
  (see `dr.get_category_euler`). Keys are dataset/category names
  (e.g. \"omniobject3d-banana\"), values are allowed rotations on axes x/y/z.
- HEARTH_KITCHENS: pre-defined kitchen room scenes (geometry USDs) with
  associated transforms, used to load different kitchen layouts.
- CATEGORIES_SCALE_SUGGESTED: default scaling for container / place targets
  by category name (e.g. basket, bowl, dish) in place-style tasks.
- OBJECT_SCALE_SUGGESTED: per-asset scaling overrides for specific USD assets
  (e.g. particular bowls or boxes that need non-uniform scale).
- HANDOVER_RANGE: list of category names that are considered suitable for
  handover-style skills (thin / graspable objects).
- CUROBO_BATCH_SIZE: default batch size used in cuRobo planners.

These constants are imported by utilities in `dr.py`, place/pick skills, and
scene setup code to keep configuration in a single shared location.
"""

CATEGORIES = {
    # large-scale pick objects
    "google_scan-book": {"y": [0, 90, 180, 270]},
    "google_scan-bottle": {"x": [0, 90]},
    "google_scan-clock": {"x": [0, 90, 180, 270]},
    "google_scan-dinosaur": {"x": [90]},
    "google_scan-garage_kit": {"x": [90]},
    "google_scan-handbag": {"x": [90, 270]},
    "google_scan-hat": {"x": [90]},
    # "google_scan-medicine_bottle": {"x": [0, 90, 180]},
    "google_scan-medicine_bottle": {"x": [0, 180]},
    "google_scan-mug": {"x": [0, 90]},
    "google_scan-shoe": {"x": [0, 90, 180, 270]},
    "google_scan-toy_animals": {"x": [90]},
    "google_scan-toy_boat": {"x": [90]},
    "google_scan-toy_bus": {"x": [90]},
    "google_scan-toy_car": {"x": [90]},
    "google_scan-toy_plane": {"x": [90]},
    "google_scan-toy_train": {"x": [90]},
    "google_scan-toy_truck": {"x": [90]},
    "omniobject3d-banana": {"x": [0, 180]},
    "omniobject3d-belt": {"x": [90, 270]},
    "omniobject3d-book": {"y": [0, 90, 180, 270]},
    # "omniobject3d-bottle": {"x": [0, 90, 180]},
    "omniobject3d-bottle": {"x": [0, 180]},
    "omniobject3d-bowl": {"x": [90]},
    "omniobject3d-boxed_beverage": {"x": [0, 90, 180, 270]},
    "omniobject3d-bread": {"x": [90, 270]},
    "omniobject3d-brush": {"y": [0, 90, 180, 270]},
    "omniobject3d-calculator": {"x": [0]},
    "omniobject3d-can": {"x": [0, 90]},
    "omniobject3d-carrot": {"x": [0, 180]},
    "omniobject3d-chess": {"x": [90]},
    "omniobject3d-chili": {"x": [0, 180]},
    "omniobject3d-chinese_chess": {"x": [90]},
    "omniobject3d-clock": {"x": [90]},
    "omniobject3d-corn": {"x": [0, 180]},
    "omniobject3d-cucumber": {"x": [0, 180]},
    "omniobject3d-cup": {"x": [90]},
    "omniobject3d-dinosaur": {"x": [90]},
    "omniobject3d-dish": {"x": [90]},
    "omniobject3d-egg_tart": {"x": [90]},
    "omniobject3d-flash_light": {"x": [0, 180]},
    "omniobject3d-garage_kit": {"x": [90]},
    "omniobject3d-glasses": {"x": [90]},
    "omniobject3d-glasses_case": {"x": [90, 270]},
    "omniobject3d-hair_dryer": {"x": [0, 180]},
    "omniobject3d-hamburger": {"x": [90]},
    "omniobject3d-hammer": {"x": [0, 180]},
    "omniobject3d-hot_dog": {"x": [90]},
    "omniobject3d-knife": {"y": [90, 270]},
    "omniobject3d-medicine_bottle": {"x": [0, 90, 180]},
    "omniobject3d-mooncake": {"x": [90]},
    "omniobject3d-mouse": {"x": [90]},
    "omniobject3d-mug": {"x": [90]},
    "omniobject3d-pen": {"x": [0, 180]},
    "omniobject3d-pineapple": {"x": [90]},
    "omniobject3d-power_strip": {"x": [90]},
    "omniobject3d-razor": {"y": [0, 90, 180, 270]},
    "omniobject3d-red_wine_glass": {"x": [90]},
    "omniobject3d-remote_control": {"y": [270]},
    "omniobject3d-sausage": {"x": [0, 180]},
    "omniobject3d-scissor": {"y": [90, 270]},
    "omniobject3d-shampoo": {"x": [0, 90, 180]},
    "omniobject3d-shoe": {"x": [90]},
    "omniobject3d-soap": {"x": [0, 90, 180, 270]},
    "omniobject3d-spanner": {"x": [90, 270]},
    "omniobject3d-starfish": {"x": [90]},
    "omniobject3d-sweet_potato": {"x": [0, 180]},
    "omniobject3d-table_tennis_bat": {"x": [90, 270]},
    "omniobject3d-teapot": {"x": [90]},
    "omniobject3d-timer": {"y": [270]},
    "omniobject3d-tissue": {"x": [0, 90, 180, 270]},
    "omniobject3d-tooth_brush": {"y": [0, 90, 180, 270]},
    "omniobject3d-tooth_paste": {"y": [0, 90, 180, 270]},
    "omniobject3d-toy_animals": {"x": [90]},
    "omniobject3d-toy_boat": {"x": [90]},
    "omniobject3d-toy_bus": {"x": [90]},
    "omniobject3d-toy_car": {"x": [90]},
    "omniobject3d-toy_plane": {"x": [90]},
    "omniobject3d-toy_plant": {"x": [90]},
    "omniobject3d-toy_train": {"x": [90]},
    "omniobject3d-toy_truck": {"x": [90]},
    "omniobject3d-umbrella": {"x": [0, 180]},
    "omniobject3d-watch": {"x": [0, 180]},
    "phocal-bottle": {"x": [0, 90, 180]},
    "phocal-cup": {"x": [90]},
    "phocal-mug": {"x": [90]},
    "phocal-red_wine_glass": {"x": [90]},
    "phocal-remote_control": {"y": [270]},
    "phocal-shampoo": {"x": [0, 90, 180]},
    "phocal-spoon": {"y": [0, 90, 180, 270]},
    "phocal-teapot": {"x": [90]},
    "real-camera": {"x": [0, 90, 180, 270]},
    # Supp
    "omniobject3d-asparagus": {"x": [0]},
    "omniobject3d-bamboo_shoots": {"x": [0]},
    "omniobject3d-battery": {"x": [0, 90]},
    "omniobject3d-cake": {"x": [0, 90]},
    "omniobject3d-candle": {"x": [0, 90]},
    "omniobject3d-chicken_leg": {"x": [0]},
    "omniobject3d-conch": {"x": [0]},
    "omniobject3d-donut": {"x": [90]},
    "omniobject3d-dumpling": {"x": [90]},
    "omniobject3d-egg": {"x": [0]},
    "omniobject3d-eraser": {"x": [90]},
    "omniobject3d-facial_cream": {"x": [90]},
    "omniobject3d-garlic": {"x": [90]},
    "omniobject3d-hand_cream": {"x": [0]},
    # "discarded": {"x": [0]},
    # "google_scan-box": {"x": [0]},
    # "omniobject3d-ball": {"x": [0]},
    # "omniobject3d-box": {"x": [0]},
    "omniobject3d-lemon": {"x": [0]},
    "omniobject3d-lipstick": {"x": [0]},
    "omniobject3d-lotus_root": {"x": [0]},
    "omniobject3d-mango": {"x": [0]},
    "omniobject3d-mangosteen": {"x": [90]},
    "omniobject3d-onion": {"x": [90]},
    "omniobject3d-orange": {"x": [90]},
    "omniobject3d-pastry": {"x": [90]},
    "omniobject3d-peach": {"x": [0]},
    "omniobject3d-pear": {"x": [0]},
    "omniobject3d-pitaya": {"x": [90]},
    "omniobject3d-pomegranate": {"x": [90]},
    "omniobject3d-rubik_cube": {"x": [0]},
    "omniobject3d-rubik_cube-old": {"x": [0]},
    "omniobject3d-shrimp": {"x": [0]},
    "omniobject3d-tape_measure": {"x": [0]},
    "omniobject3d-tomato": {"x": [90]},
    "omniobject3d-toy_motorcycle": {"x": [0]},
    "omniobject3d-waffle": {"x": [90]},
    # "omniobject3d-whistle": {"x": [0]},
    "phocal-can": {"x": [0, 90]},
    # Basic
    "book": {"y": [0, 90, 180, 270]},
    "plate": {"x": [90], "y": [0], "z": [0]},
    "basket": {"x": [90]},
    "pnp_basket": {"x": [90], "y": [0], "z": [90]},
    "dirt": {"x": [90]},
    # Place
    # "basket": {"x": [90]},
    "bowl": {"x": [90]},
    "dish": {"x": [90]},
    "flower_pot": {"x": [90]},
    "pan": {"x": [90]},
    "box": {"x": [90]},
    "microwaveoven": {"x": [90]},
    "drawing": {"x": [90]},
    "projector": {"x": [90]},
    "pad": {"x": [90]},
    # Long horizon
    "bolt": {"x": [90]},
}

HEARTH_KITCHENS = {
    "001": {
        "name": "scene",
        "path": "home_scenes/rooms/kitchen/MVUHLWYKTKJ5EAABAAAAAAY8_usd_kitchen_0.usd",
        "target_class": "GeometryObject",
        "scale": [0.015, 0.015, 0.015],
        "translation": [-9.8, -3.9, -0.3],
        "euler": [0, 0, 180],
    },
    "002": {
        "name": "scene",
        "path": "home_scenes/rooms/kitchen/MVUHLWYKTKJ5EAABAAAAABA8_usd_kitchen_0.usd",
        "target_class": "GeometryObject",
        "scale": [0.013, 0.013, 0.013],
        "translation": [-8, 1.55, -0.163],
        "euler": [0, 0, -90],
    },
    "003": {
        "name": "scene",
        "path": "home_scenes/rooms/kitchen/MVUHLWYKTKJ5EAABAAAAABI8_usd_kitchen_0.usd",
        "target_class": "GeometryObject",
        "scale": [0.013, 0.013, 0.013],
        "translation": [5.5, -6.2, -0.133],
        "euler": [0, 0, 90],
    },
    "004": {
        "name": "scene",
        "path": "home_scenes/rooms/kitchen/MVUHLWYKTKJ5EAABAAAAABQ8_usd_kitchen_0.usd",
        "target_class": "GeometryObject",
        "scale": [0.013, 0.013, 0.013],
        "translation": [-7.15, -7.2, -0.133],
        "euler": [0, 0, 180],
    },
    "005": {
        "name": "scene",
        "path": "home_scenes/rooms/kitchen/MVUHLWYKTKJ5EAABAAAAABY8_usd_kitchen_0.usd",
        "target_class": "GeometryObject",
        "scale": [0.013, 0.013, 0.013],
        "translation": [7.2, -2.6, -0.1],
        "euler": [0, 0, 90],
    },
    "006": {
        "name": "scene",
        "path": "home_scenes/rooms/kitchen/MVUHLWYKTKJ5EAABAAAAACA8_usd_kitchen_0.usd",
        "target_class": "GeometryObject",
        "scale": [0.015, 0.015, 0.015],
        "translation": [6.27, -8.4, -0.33],
        "euler": [0.0, 0.0, 90.0],
    },
    "007": {
        "name": "scene",
        "path": "home_scenes/rooms/kitchen/MVUHLWYKTKJ5EAABAAAAACI8_usd_kitchen_0.usd",
        "target_class": "GeometryObject",
        "scale": [0.015, 0.015, 0.015],
        "translation": [7.25, -5.1, -0.3],
        "euler": [0, 0, 180],
    },
    "008": {
        "name": "scene",
        "path": "home_scenes/rooms/kitchen/MVUHLWYKTKJ5EAABAAAAACQ8_usd_kitchen_0.usd",
        "target_class": "GeometryObject",
        "scale": [0.015, 0.015, 0.015],
        "translation": [5.84, 0.67, -0.3],
        "euler": [0, 0, 90],
    },
    "009": {
        "name": "scene",
        "path": "home_scenes/rooms/kitchen/MVUHLWYKTKJ5EAABAAAAACY8_usd_kitchen_0.usd",
        "target_class": "GeometryObject",
        "scale": [0.015, 0.015, 0.015],
        "translation": [-7.95, -3, -0.37],
        "euler": [0, 0, 0],
    },
    "010": {
        "name": "scene",
        "path": "home_scenes/rooms/kitchen/MVUHLWYKTKJ5EAABAAAAADA8_usd_kitchen_0.usd",
        "target_class": "GeometryObject",
        "scale": [0.015, 0.015, 0.015],
        "translation": [-5.89, -3.58, -0.3],
        "euler": [0, 0, 180],
    },
}

CATEGORIES_SCALE_SUGGESTED = {
    # Default per-category scale for place tasks (containers / targets)
    "basket": [0.001, 0.001, 0.001],
    "bowl": [0.002, 0.002, 0.002],
    "dish": [0.001, 0.001, 0.001],
    "flower_pot": [0.001, 0.001, 0.001],
    "pan": [0.001, 0.001, 0.001],
    # "box": [0.0007, 0.0007, 0.0007],
    "box": [0.001, 0.001, 0.001],
    "microwaveoven": [0.0005, 0.0005, 0.0005],
    "drawing": [0.001, 0.001, 0.001],
    "projector": [0.001, 0.001, 0.001],
    "pad": [0.0005, 0.0005, 0.0005],
}

OBJECT_SCALE_SUGGESTED = {
    # Per-asset scale overrides for place tasks (special cases)
    "google_scan-picnic_basket_0908": [0.0012, 0.0012, 0.0012],
    "omniobject3d-picnic_basket_006": [0.0012, 0.0012, 0.0012],
    "google_scan-bowl_0070": [0.001, 0.001, 0.001],
    "google_scan-bowl_0071": [0.0008, 0.0008, 0.0008],
    "google_scan-bowl_0495": [0.0008, 0.0008, 0.0008],
    "google_scan-bowl_0498": [0.0015, 0.0015, 0.0015],
    "google_scan-bowl_0539": [0.0015, 0.0015, 0.0015],
    "google_scan-bowl_0746": [0.001, 0.001, 0.001],
    "google_scan-bowl_0752": [0.0008, 0.0008, 0.0008],
    "google_scan-bowl_0839": [0.0015, 0.0015, 0.0015],
    "google_scan-bowl_0891": [0.0015, 0.0015, 0.0015],
    "google_scan-dish_0085": [0.0008, 0.0008, 0.0008],
    "google_scan-dish_0096": [0.002, 0.002, 0.002],
    "google_scan-dish_0099": [0.002, 0.002, 0.002],
    "google_scan-dish_0510": [0.0015, 0.0015, 0.0015],
    "omniobject3d-flower_pot_001": [0.001, 0.001, 0.001],
    "omniobject3d-flower_pot_002": [0.001, 0.001, 0.001],
    "omniobject3d-flower_pot_003": [0.002, 0.002, 0.002],
    "omniobject3d-flower_pot_007": [0.002, 0.002, 0.002],
    "omniobject3d-pan_002": [0.0015, 0.0015, 0.0015],
    "omniobject3d-pan_004": [0.002, 0.002, 0.002],
    "omniobject3d-picnic_basket_003": [0.0009, 0.0009, 0.0009],
    "omniobject3d-picnic_basket_004": [0.0009, 0.0009, 0.0009],
    "omniobject3d-picnic_basket_011": [0.0009, 0.0009, 0.0009],
    "omniobject3d-box_003": [0.0011, 0.0011, 0.0011],
    "omniobject3d-box_022": [0.0006, 0.0006, 0.0006],
    "omniobject3d-box_024": [0.0021, 0.0021, 0.0021],
    "omniobject3d-box_040": [0.0011, 0.0011, 0.0011],
    "omniobject3d-box_043": [0.0016, 0.0016, 0.0016],
    "omniobject3d-box_044": [0.0021, 0.0021, 0.0021],
    "omniobject3d-drawing_001": [0.0003, 0.0003, 0.0003],
    "omniobject3d-drawing_017": [0.0004, 0.0004, 0.0004],
    "omniobject3d-projector_008": [0.0005, 0.0005, 0.0005],
    "omniobject3d-pad_022": [0.0003, 0.0003, 0.0003],
    "omniobject3d-pad_010": [0.0007, 0.0007, 0.0007],
}

HANDOVER_RANGE = [
    "omniobject3d-bamboo_shoots",
    "omniobject3d-banana",
    "omniobject3d-bottle",
    "omniobject3d-carrot",
    "omniobject3d-chili",
    "omniobject3d-corn",
    "omniobject3d-cucumber",
    "omniobject3d-pen",
    "omniobject3d-shampoo",
    "omniobject3d-scissor",
    "omniobject3d-sweet_potato",
    "omniobject3d-tooth_brush",
    "omniobject3d-tooth_paste",
    "phocal-bottle",
    "phocal-spoon",
]

CUROBO_BATCH_SIZE = 20