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

import marimo

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


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

    import marimo as mo
    import numpy as np
    import polars as pl
    from PIL import Image

    from compressors.pipeline import HashPipeline
    from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph
    from simulator import (
        HabitatSimulatorConfig,
        TopDownSceneElements,
        collect_room_views_by_room,
        create_habitat_simulator,
        render_topdown_scene_map,
    )
    from utils.image import extract_masked_region, segment_image

    return (
        BytesIO,
        HabitatSimulatorConfig,
        HashPipeline,
        Image,
        ObjectNode,
        RoomNode,
        SimpleSceneGraph,
        TopDownSceneElements,
        collect_room_views_by_room,
        create_habitat_simulator,
        extract_masked_region,
        mo,
        np,
        pl,
        render_topdown_scene_map,
        segment_image,
    )


@app.cell
def setup_verification_context(
    HabitatSimulatorConfig, RoomNode, create_habitat_simulator, np
):
    scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
    image_size = 256
    num_rooms = 4
    views_per_room = 6
    meters_per_pixel = 0.05

    sam_max_masks = 5
    sam_min_area = 32 * 32
    hash_bits = 512

    sim, agent = create_habitat_simulator(
        HabitatSimulatorConfig(
            scene_path=scene_path,
            views_per_room=views_per_room,
            image_size=image_size,
            sensor_height=1.5,
            move_forward_step=0.25,
            enable_physics=False,
        )
    )

    room_nodes = []
    for idx in range(num_rooms):
        point = sim.pathfinder.get_random_navigable_point()
        room_nodes.append(
            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),
            )
        )

    print("Sampled room centers:")
    for node in room_nodes:
        print(node.room_id, node.center)

    return (
        agent,
        hash_bits,
        meters_per_pixel,
        room_nodes,
        sam_max_masks,
        sam_min_area,
        sim,
        views_per_room,
    )


@app.cell
def render_topdown_room_map(
    TopDownSceneElements,
    meters_per_pixel,
    render_topdown_scene_map,
    room_nodes,
    sim,
):
    render_topdown_scene_map(
        pathfinder=sim.pathfinder,
        elements=TopDownSceneElements(room_nodes=room_nodes),
        meters_per_pixel=meters_per_pixel,
    )
    return


@app.cell
def build_scene_graph_pipeline(
    agent,
    HashPipeline,
    Image,
    ObjectNode,
    SimpleSceneGraph,
    collect_room_views_by_room,
    extract_masked_region,
    hash_bits,
    mo,
    np,
    room_nodes,
    sam_max_masks,
    sam_min_area,
    segment_image,
    sim,
    views_per_room,
):
    all_room_views = collect_room_views_by_room(
        agent=agent,
        sim=sim,
        room_nodes=room_nodes,
        views_per_room=views_per_room,
    )

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

    scene_graph = SimpleSceneGraph(
        rooms={room.room_id: room for room in room_nodes},
        objects={},
    )
    total_masks = 0
    object_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_{object_index:04d}"
            object_index += 1
            bits_np = bits.squeeze().detach().cpu().numpy()

            scene_graph.objects[obj_id] = 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,
            )

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


@app.cell
def build_room_and_object_tables(pl, scene_graph):
    room_rows = [
        {
            "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]),
        }
        for room in scene_graph.rooms.values()
    ]

    object_rows = [
        {
            "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(),
        }
        for obj in scene_graph.objects.values()
    ]

    rooms_table = pl.DataFrame(room_rows)
    objects_table = pl.DataFrame(object_rows)
    return objects_table, rooms_table


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

    uploaded_image = None
    if file_upload.value:
        contents = file_upload.contents()
        if contents:
            uploaded_image = Image.open(BytesIO(contents))
            mo.image(np.array(uploaded_image), alt="Uploaded query image")

    return file_upload, uploaded_image


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