import simulate as sm 
from simulate.assets.action_mapping import ActionMapping
import time

CAMERA_HEIGHT = 40
CAMERA_WIDTH = 64

scene = sm.Scene(engine="unity")
# scene = sm.Scene()
scene += sm.LightSun()
scene += sm.Box(
    name="table_surface",
    position=[1, 0, 1],
    bounds=[0, 7, 0, 1, 0, 11],
    material=sm.Material.GREEN,
    with_collider=True,
    with_rigid_body=True,
    physics_component=sm.RigidBodyComponent(
        use_gravity=True,
        kinematic=True,
        collision_detection="continuous"
    ),
)

scene += sm.Box(
        name="border_left",
        position=[0, 0, 0],
        bounds=[0, 1, 0, 2, 0, 13],
        material=sm.Material.GREEN,
        with_collider=True,
        with_rigid_body=True,
        physics_component=sm.RigidBodyComponent(
            use_gravity=True,
            kinematic=True,
            collision_detection="continuous"
        ),
)

scene += sm.Box(
        name="border_right",
        position=[7, 0, 0],
        bounds=[0, 1, 0, 2, 0, 13],
        material=sm.Material.GREEN,
        with_collider=True,
        with_rigid_body=True,
        physics_component=sm.RigidBodyComponent(
            use_gravity=True,
            mass=1,
            kinematic=True,
            collision_detection="continuous"
        ),
)

scene += sm.Box(
        name="border_top",
        position=[0, 0, 0],
        bounds=[0, 8, 0, 2, 0, 1],
        material=sm.Material.GREEN,
        with_collider=True,
        with_rigid_body=True,
        physics_component=sm.RigidBodyComponent(
            use_gravity=True,
            kinematic=True,
            collision_detection="continuous"
        ),
)

scene += sm.Box(
        name="border_bottom",
        position=[0, 0, 12],
        bounds=[0, 8, 0, 2, 0, 1],
        material=sm.Material.GREEN,
        with_collider=True,
        with_rigid_body=True,
        physics_component=sm.RigidBodyComponent(
            use_gravity=True,
            kinematic=True,
            collision_detection="continuous"
        ),
)

cue_ball = sm.Sphere(
    name="cue_ball",
    position=[4, 1.2, 9],
    radius=0.2,
    material=sm.Material.WHITE,
    with_collider=True,
    with_rigid_body=True,
    physics_component=sm.RigidBodyComponent(
        use_gravity=True,
        constraints=["freeze_position_y"],
        mass=1,
        collision_detection="continuous"
    ),
)


cue_ball.is_actor = True 

mapping = [
    ActionMapping("add_force", axis=[0, 0, -1], is_impulse=True),
]

cue_ball.actuator = sm.Actuator(
    mapping=mapping,
    high=[10],
    low=[-10],
    shape=(1,)
)

cue_ball += sm.StateSensor(target_entity=cue_ball, properties=["position"])

scene += cue_ball

target_ball = sm.Sphere(
    name="target_ball",
    position=[4, 1.2, 6],
    radius=0.2,
    material=sm.Material.RED,
    with_collider=True,
    with_rigid_body=True,
    physics_component=sm.RigidBodyComponent(
        use_gravity=True,
        constraints=["freeze_position_y"],
        mass=0.5,
        collision_detection="continuous"
    ),
)

scene += target_ball

scene += sm.Camera(name="camera", position=[0,0,0], rotation=[35, 0, 0])

scene.show()

# # uncomment here
env = sm.RLEnv(scene)
env.reset()

for i in range(500):
    action = scene.action_space.sample()
    
    print("###############################")
    obs, reward, done, info = env.step(action)
    print("###############################")
    print(action)
    print(f"Ball position {obs.get('StateSensor')}")
    print(f"step {i}, reward {reward[0]}, observation")
    time.sleep(0.1)

env.close()



# input("Press enter to continue...")