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

import marimo

__generated_with = "0.20.4"
app = marimo.App(width="medium", app_title="Pipeline Verification")


@app.cell
def import_packages():
    from io import BytesIO

    import habitat_sim
    import marimo as mo
    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 (
        BytesIO,
        HashPipeline,
        Image,
        ObjectNode,
        RoomNode,
        SimpleSceneGraph,
        extract_masked_region,
        habitat_sim,
        maps,
        mo,
        np,
        pl,
        plt,
        segment_image,
    )


@app.cell
def setup_habitat_simulator(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 sample_room_nodes(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 render_topdown_room_map(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 collect_room_views(
    agent,
    habitat_sim,
    mo,
    plt,
    room_nodes,
    sim,
    views_per_room,
):
    all_room_views = {}

    for _node in mo.status.progress_bar(
        room_nodes,
        title="Collecting room views",
        subtitle="Sampling observations from Habitat",
        show_eta=True,
        show_rate=True,
    ):
        _agent_state = habitat_sim.AgentState()
        _agent_state.position = _node.center.copy()
        agent.set_state(_agent_state)

        _room_views = []
        for _ in mo.status.progress_bar(
            range(views_per_room),
            title=f"Capturing {_node.room_id}",
            subtitle="Rotating agent viewpoints",
            show_eta=True,
            show_rate=True,
        ):
            _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 build_hash_pipeline(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 build_scene_graph_from_views(
    Image,
    ObjectNode,
    SimpleSceneGraph,
    all_room_views,
    extract_masked_region,
    hash_pipeline,
    mo,
    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

    _view_jobs = [
        (_room_id, _view_idx, _rgb)
        for _room_id, _views in all_room_views.items()
        for _view_idx, _rgb in enumerate(_views)
    ]

    for _room_id, _view_idx, _rgb in mo.status.progress_bar(
        _view_jobs,
        title="Extracting masks and hashes",
        subtitle="Running SAM + HashPipeline",
        show_eta=True,
        show_rate=True,
    ):
        _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,
                image_bytes=np.array(_masked_image).tobytes(),
                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 build_room_and_object_tables(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(disabled=True)
def display_rooms_table(rooms_table):
    rooms_table
    return


@app.cell(disabled=True)
def display_objects_table(objects_table):
    objects_table
    return


@app.cell
def create_file_upload(mo):
    file_upload = mo.ui.file(
        filetypes=["image/*"], kind="area", label="Upload a query image"
    )
    file_upload
    return (file_upload,)


@app.cell
def load_uploaded_image(BytesIO, Image, file_upload):
    uploaded_image = None
    if file_upload.value:
        _contents = file_upload.contents()
        if _contents:
            uploaded_image = Image.open(BytesIO(_contents))
    return (uploaded_image,)


@app.cell
def display_uploaded_image(mo, np, uploaded_image):
    mo.image(np.array(uploaded_image), alt="Uploaded query image")
    return


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