embodied intelligence1

from ai2thor.controller import Controller
import os
import cv2
import matplotlib
import matplotlib.pyplot as plt
from PIL import Image
import json

sav_path = './tmp_data/'
if not os.path.exists(sav_path):
    os.makedirs(sav_path, exist_ok = True)

# controller = Controller(
#     agentMode="default",
#     visibilityDistance=1.5,
#     scene="FloorPlan1",
#
#     # step sizes
#     gridSize=0.25,
#     snapToGrid=True,
#     rotateStepDegrees=90,
#
#     # image modalities
#     renderDepthImage=False,
#     renderInstanceSegmentation=False,
#
#     # camera properties
#     width=2000,
#     height=1000,
#     fieldOfView=100
# )
# controller.step("PausePhysicsAutoSim")
# controller.step(
#     action="AdvancePhysicsStep",
#     timeStep=0.02
# )
# controller.step(
#     action="PickupObject",
#     objectId="Apple|1|1|1",
#     forceAction=False,
#     manualInteract=False
# )
# controller.step(
#     action="MoveAhead",
#     moveMagnitude=5
# )
#
# controller.step(
#     action="MoveAhead",
#     moveMagnitude=5
# )
# controller.step(
#     action="MoveAhead",
#     moveMagnitude=5
# )
# controller.step(
#     action="RotateRight",
#     degrees=30
# )
# controller.step(
#     action="LookUp",
#     degrees=30
# )
# controller.step("UnpausePhysicsAutoSim")
# Other supported directions
# controller.step("MoveBack")
# controller.step("MoveLeft")
# controller.step("MoveRight")


# print(event.third_party_camera_frames)

controller = Controller(
        scene = "FloorPlan1",
        gridSize = 0.5,
        rotateStepDegrees = 90,
        renderDepthImage = True,
        renderInstanceSegmentation = True,
        width=300,
        height=300,
        fieldOfView = 90
    )
print("Start the controller.")

num_step = 0
event = controller.step(action = "GetReachablePositions")
reachable_positions = event.metadata["actionReturn"]
print(f"Reachable Points: {reachable_positions}")
for i in event.metadata["objects"]:
    # if i["pickupable"]==True:
    #     print(i['objectId'])
    # if i['objectId']=='Shelf|+01.75|+00.55|-02.56':
    #     print(i)
    print(i['objectId'])
# print(event.metadata["objects"])
plt.imsave(os.path.join(sav_path, f'frame_{num_step}.png'), event.frame)
num_step += 1

event = controller.step(
    action = "Teleport",
    position = dict(x = -1.0, y = 0.900999128818512, z = 1.0),
    rotation = dict(x = 0, y = 0, z = 0),
    horizon = 0,
    standing = True
)
print("Teleport {}".format(event.metadata["lastActionSuccess"]))
plt.imsave(os.path.join(sav_path, f'frame_{num_step}.png'), event.frame)
num_step += 1



nav_actions = ['MoveAhead', 'RotateRight']

for act in nav_actions:
    event = controller.step(action = act)
    print("action:{} {}".format(act, event.metadata["lastActionSuccess"]))
    plt.imsave(os.path.join(sav_path, f'frame_{num_step}.png'), event.frame)
    num_step += 1

event = controller.step(
    action="PickupObject",
    objectId="Book|+00.15|+01.10|+00.62"
)
print("PickUp {}".format(event.metadata["lastActionSuccess"]))
plt.imsave(os.path.join(sav_path, f'frame_{num_step}.png'), event.frame)
num_step += 1

nav_actions2 = ['MoveAhead', 'MoveAhead', 'MoveAhead', 'MoveAhead', 'MoveAhead']
for act in nav_actions2:
    event = controller.step(action = act)
    print("action:{} {}".format(act, event.metadata["lastActionSuccess"]))
    plt.imsave(os.path.join(sav_path, f'frame_{num_step}.png'), event.frame)
    num_step += 1

event = controller.step(
    action = "PutObject",
    objectId = "Floor|+00.00|+00.00|+00.00",
    forceAction = True
)
print("Put Book {}".format(event.metadata["lastActionSuccess"]))
plt.imsave(os.path.join(sav_path, f'frame_{num_step}.png'), event.frame)
num_step += 1

event = controller.step(action = "GetMapViewCameraProperties", raise_for_failure = True)
event = controller.step(
    action = "AddThirdPartyCamera",
    **event.metadata["actionReturn"]
)
top_down_frame = event.third_party_camera_frames[-1]
print("Draw Top-Down View")
plt.imsave(os.path.join(sav_path, f'top_down.png'), top_down_frame)