# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "marimo>=0.21.1",
#     "pyzmq>=27.1.0",
# ]
# ///

import marimo

__generated_with = "0.21.1"
app = marimo.App(app_title="Pipeline Verification")


@app.cell
def import_packages():
    import habitat_sim
    import numpy as np
    import polars as pl
    from habitat.utils.visualizations import maps
    from matplotlib import pyplot as plt
    from PIL import Image
    from compressors.pipeline import HashPipeline
    from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph
    from utils.common import get_device
    from utils.image import extract_masked_region, segment_image

    return (
        HashPipeline,
        Image,
        ObjectNode,
        RoomNode,
        SimpleSceneGraph,
        extract_masked_region,
        habitat_sim,
        maps,
        np,
        pl,
        plt,
        segment_image,
    )


@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)
    sam_max_masks = 5
    sam_min_area = 32 * 32
    sam_points_per_batch = 64
    hash_bits = 512
    return (
        agent,
        hash_bits,
        meters_per_pixel,
        num_rooms,
        sam_max_masks,
        sam_min_area,
        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)

    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 (all_room_views,)


@app.cell
def _(HashPipeline, hash_bits, sam_max_masks, sam_min_area):
    hash_pipeline = HashPipeline(
        dino_model="facebook/dinov2-large",
        sam_model="facebook/sam2.1-hiera-large",
        sam_min_mask_area=sam_min_area,
        sam_max_masks=sam_max_masks,
        hash_bits=hash_bits,
    )
    return (hash_pipeline,)


@app.cell
def _(
    Image,
    ObjectNode,
    SimpleSceneGraph,
    all_room_views,
    extract_masked_region,
    hash_pipeline,
    np,
    room_nodes,
    segment_image,
):
    scene_graph = SimpleSceneGraph(
        rooms={_room.room_id: _room for _room in room_nodes}, objects={}
    )
    total_masks = 0
    _obj_index = 0

    for _room_id, _views in all_room_views.items():
        for _view_idx, _rgb in enumerate(_views):
            _rgb3 = _rgb[..., :3] if _rgb.shape[-1] > 3 else _rgb
            _image = Image.fromarray(_rgb3.astype(np.uint8))

            _masks = segment_image(
                hash_pipeline.mask_generator,
                _image,
                min_area=hash_pipeline.sam_min_mask_area,
                max_masks=hash_pipeline.sam_max_masks,
                points_per_batch=hash_pipeline.sam_points_per_batch,
            )
            total_masks += len(_masks)

            for _mask in _masks:
                _masked_image = extract_masked_region(_image, _mask["segment"])
                _bits = hash_pipeline(_masked_image)

                _bbox = _mask["bbox"]
                _obj_center = np.array(
                    [_bbox[0] + _bbox[2] / 2, _bbox[1] + _bbox[3] / 2, 0.0],
                    dtype=np.float32,
                )

                _obj_id = f"obj_{_obj_index:04d}"
                _obj_index += 1
                _bits_np = _bits.squeeze().detach().cpu().numpy()

                _obj_node = ObjectNode(
                    obj_id=_obj_id,
                    room_id=_room_id,
                    position=_obj_center,
                    visual_hash=_bits_np,
                    semantic_hash=_bits_np,
                    hit_count=1,
                    last_seen_frame=0,
                )
                scene_graph.objects[_obj_id] = _obj_node

    print(f"Total objects created: {len(scene_graph.objects)}")
    print(f"Total processed masks: {total_masks}")
    return (scene_graph,)


@app.cell
def _(pl, scene_graph):
    _room_rows = []
    for _room in scene_graph.rooms.values():
        _room_rows.append(
            {
                "room_id": _room.room_id,
                "center_x": float(_room.center[0]),
                "center_y": float(_room.center[1]),
                "center_z": float(_room.center[2]),
                "bbox_dx": float(_room.bbox_extent[0]),
                "bbox_dy": float(_room.bbox_extent[1]),
                "bbox_dz": float(_room.bbox_extent[2]),
            }
        )

    _object_rows = []
    for _obj in scene_graph.objects.values():
        _object_rows.append(
            {
                "obj_id": _obj.obj_id,
                "room_id": _obj.room_id,
                "last_seen_frame": int(_obj.last_seen_frame),
                "hit_count": int(_obj.hit_count),
                "visual_hash": _obj.visual_hash.tolist(),
                "semantic_hash": _obj.semantic_hash.tolist(),
            }
        )

    rooms_table = pl.DataFrame(_room_rows)
    objects_table = pl.DataFrame(_object_rows)
    return objects_table, rooms_table


@app.cell
def _(rooms_table):
    rooms_table
    return


@app.cell
def _(objects_table):
    objects_table
    return


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