import marimo

__generated_with = "0.21.1"
app = marimo.App()


@app.cell
def _():
    import habitat_sim
    import numpy as np
    from habitat.utils.visualizations import maps
    from matplotlib import pyplot as plt
    from scenegraph import RoomNode

    return RoomNode, habitat_sim, maps, np, plt


@app.cell
def _(habitat_sim):
    scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
    num_rooms = 4
    views_per_room = 6
    image_size = 256
    meters_per_pixel = 0.05

    sim_cfg = habitat_sim.SimulatorConfiguration()
    sim_cfg.scene_id = scene_path
    sim_cfg.enable_physics = False

    agent_cfg = habitat_sim.agent.AgentConfiguration()
    rgb_sensor_spec = habitat_sim.CameraSensorSpec()
    rgb_sensor_spec.uuid = "color_sensor"
    rgb_sensor_spec.sensor_type = habitat_sim.SensorType.COLOR
    rgb_sensor_spec.resolution = [image_size, image_size]
    rgb_sensor_spec.position = [0.0, 1.5, 0.0]
    agent_cfg.sensor_specifications = [rgb_sensor_spec]

    turn_angle = 360.0 / views_per_room
    agent_cfg.action_space = {
        "move_forward": habitat_sim.agent.ActionSpec(
            "move_forward", habitat_sim.agent.ActuationSpec(amount=0.25)
        ),
        "turn_left": habitat_sim.agent.ActionSpec(
            "turn_left", habitat_sim.agent.ActuationSpec(amount=turn_angle)
        ),
        "turn_right": habitat_sim.agent.ActionSpec(
            "turn_right", habitat_sim.agent.ActuationSpec(amount=turn_angle)
        ),
    }

    cfg = habitat_sim.Configuration(sim_cfg, [agent_cfg])
    sim = habitat_sim.Simulator(cfg)
    agent = sim.initialize_agent(0)
    return agent, meters_per_pixel, num_rooms, sim, views_per_room


@app.cell
def _(RoomNode, np, num_rooms, sim):
    room_nodes = []
    for _idx in range(num_rooms):
        _point = sim.pathfinder.get_random_navigable_point()
        _room_node = RoomNode(
            room_id=f"room_{_idx:02d}",
            center=np.asarray(_point, dtype=np.float32),
            bbox_extent=np.asarray([1.5, 2.0, 1.5], dtype=np.float32),
        )
        room_nodes.append(_room_node)

    room_points = np.vstack([_node.center for _node in room_nodes])
    print("Sampled room centers:")
    for _node in room_nodes:
        print(_node.room_id, _node.center)
    return (room_nodes,)


@app.cell
def _(maps, meters_per_pixel, plt, room_nodes, sim):
    top_down_map = maps.get_topdown_map(
        sim.pathfinder,
        height=float(room_nodes[0].center[1]),
        meters_per_pixel=meters_per_pixel,
    )

    plt.figure(figsize=(8, 8))
    plt.imshow(top_down_map, cmap="gray")

    for _node in room_nodes:
        _gy, _gx = maps.to_grid(
            float(_node.center[2]),
            float(_node.center[0]),
            top_down_map.shape,
            pathfinder=sim.pathfinder,
        )
        plt.scatter(_gx, _gy, c="red", s=50)
        plt.text(_gx + 2, _gy + 2, _node.room_id, color="yellow", fontsize=8)

    plt.title("RoomNode Top-Down Map")
    plt.axis("off")
    plt.show()
    return


@app.cell
def _(agent, habitat_sim, plt, room_nodes, sim, views_per_room):
    all_room_views = {}

    for _node in room_nodes:
        _agent_state = habitat_sim.AgentState()
        _agent_state.position = _node.center.copy()
        agent.set_state(_agent_state)

        _room_views = []
        for _ in range(views_per_room):
            _observations = sim.get_sensor_observations()
            _rgb = _observations["color_sensor"]
            _room_views.append(_rgb)
            sim.step("turn_left")

        all_room_views[_node.room_id] = _room_views

        _fig, _axes = plt.subplots(2, 3, figsize=(10, 6))
        for _view_idx, _ax in enumerate(_axes.flatten()):
            _ax.imshow(_room_views[_view_idx])
            _ax.set_title(f"{_node.room_id} - view {_view_idx + 1}")
            _ax.axis("off")
        plt.tight_layout()
        plt.show()
    return


if __name__ == "__main__":
    app.run()