import marimo

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


@app.cell
def _():
    import habitat_sim
    from habitat_sim import ShortestPath
    import numpy as np
    import plotly.express as px
    from matplotlib import pyplot as plt

    return ShortestPath, habitat_sim, np, plt


@app.cell
def _(habitat_sim):
    # 配置场景
    scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"

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


@app.cell
def _(habitat_sim, sim_cfg):
    # 配置 agent
    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 = [256, 256]
    rgb_sensor_spec.position = [0.0, 1.5, 0.0]

    agent_cfg.sensor_specifications = [rgb_sensor_spec]

    # 创建 simulator 实例
    cfg = habitat_sim.Configuration(sim_cfg, [agent_cfg])
    sim = habitat_sim.Simulator(cfg)
    return (sim,)


@app.cell
def _(habitat_sim, plt, sim):
    # 初始化 agent
    agent = sim.initialize_agent(0)

    # 设置 agent 初始位置：使用可导航点，而不是写死 [0,0,0]
    agent_state = habitat_sim.AgentState()
    agent_state.position = sim.pathfinder.get_random_navigable_point()
    agent.set_state(agent_state)

    state = agent.get_state()
    print("位置:", state.position)
    print("旋转四元数:", state.rotation)

    observations = sim.get_sensor_observations()
    rgb_image = observations["color_sensor"]
    print("RGB shape:", rgb_image.shape)
    print("RGB min/max:", rgb_image[..., :3].min(), rgb_image[..., :3].max())


    plt.imshow(rgb_image)
    plt.axis("off")
    plt.show()
    return (agent,)


@app.cell
def _(agent, plt, sim):
    action_names = list(agent.agent_config.action_space.keys())
    print("可用动作:", action_names)

    # 进行一步移动
    sim.step("move_forward")

    import random
    6
    for i in range(5):
        action = random.choice(action_names)
        obs = sim.step(action)
        rgb = obs["color_sensor"]
        plt.imshow(rgb)
        plt.axis("off")
        plt.show()
    return


@app.cell
def _(sim):
    pathfinder = sim.pathfinder
    print("NavMesh是否加载:", pathfinder.is_loaded)
    print("可导航面积:", pathfinder.navigable_area)
    print("场景边界:", pathfinder.get_bounds())

    # 获取随机可导航点
    rand_point = pathfinder.get_random_navigable_point()
    print("随机可导航点:", rand_point)
    return (pathfinder,)


@app.cell
def _(ShortestPath, pathfinder):
    start = pathfinder.get_random_navigable_point()
    end = pathfinder.get_random_navigable_point()

    path = ShortestPath()
    path.requested_start = start
    path.requested_end = end

    found = pathfinder.find_path(path)
    print("是否找到路径:", found)
    print("路径点:", path.points)
    print("路径长度:", path.geodesic_distance)
    return path, start


@app.cell
def _(path, pathfinder, plt, sim, start):
    from habitat.utils.visualizations import maps

    top_down_map = maps.get_topdown_map(
        pathfinder, height=start[1], meters_per_pixel=0.05
    )

    # 绘制路径
    trajectory = [
        maps.to_grid(p[2], p[0], top_down_map.shape, pathfinder=sim.pathfinder)
        for p in path.points
    ]
    maps.draw_path(top_down_map, trajectory)

    plt.imshow(top_down_map)
    plt.axis("off")
    plt.show()
    return


@app.cell
def _(habitat_sim, np):
    # 高级控制：连续动作和滑动碰撞

    vel_control = habitat_sim.physics.VelocityControl()
    vel_control.controlling_lin_vel = True
    vel_control.lin_vel_is_local = True
    vel_control.controlling_ang_vel = True

    # 设置线速度和角速度
    vel_control.linear_velocity = np.array([0, 0, -1.0])  # forward
    vel_control.angular_velocity = np.array([0, 0.1, 0])  # rotation
    return


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