# /// 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 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 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 (
        HabitatSimulatorConfig,
        HashPipeline,
        Image,
        ObjectNode,
        RoomNode,
        SimpleSceneGraph,
        TopDownSceneElements,
        collect_room_views_by_room,
        create_habitat_simulator,
        extract_masked_region,
        maps,
        mo,
        np,
        pl,
        plt,
        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
    pipeline_batch_size = 64

    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,
        pipeline_batch_size,
        room_nodes,
        sam_max_masks,
        sam_min_area,
        sim,
        views_per_room,
    )


@app.cell
def render_topdown_room_map(
    TopDownSceneElements,
    maps,
    meters_per_pixel,
    plt,
    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,
        maps_module=maps,
        plt_module=plt,
    )
    return


@app.cell
def build_scene_graph_pipeline(
    HashPipeline,
    Image,
    ObjectNode,
    SimpleSceneGraph,
    agent,
    collect_room_views_by_room,
    extract_masked_region,
    hash_bits,
    mo,
    np,
    pipeline_batch_size,
    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

    room_view_dataset = [
        (room_id, view_idx, rgb)
        for room_id, views in all_room_views.items()
        for view_idx, rgb in enumerate(views)
    ]
    object_dataset = []

    for room_id, _view_idx, rgb in mo.status.progress_bar(
        room_view_dataset,
        title="Building object dataset",
        subtitle="Running SAM segmentation",
        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"])
            object_dataset.append((room_id, mask["bbox"], masked_image))

    if object_dataset:
        masked_images = [item[2] for item in object_dataset]
        batched_bits = hash_pipeline.forward_dataset(
            masked_images,
            batch_size=pipeline_batch_size,
            apply_sam=False,
        )
        if len(batched_bits) != len(object_dataset):
            raise RuntimeError(
                "Batch output size mismatch between masked images and hash outputs."
            )
    else:
        batched_bits = []

    for ob_idx, (room_id, bbox, _) in enumerate(object_dataset):
        bits = batched_bits[ob_idx]
        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_array = np.asarray(bits.detach().cpu().numpy()).reshape(-1)
        if bits_array.size == 512:
            bits_binary = (bits_array > 0).astype(np.uint8)
            hash_bytes = np.packbits(bits_binary).tobytes()
        elif bits_array.size == 64:
            hash_bytes = bits_array.astype(np.uint8).tobytes()
        else:
            raise ValueError(
                f"Unexpected hash length: {bits_array.size}. Expected 512 bits or 64 bytes."
            )

        scene_graph.objects[obj_id] = ObjectNode(
            obj_id=obj_id,
            room_id=room_id,
            position=obj_center,
            visual_hash=hash_bytes,
            semantic_hash=hash_bytes,
            hit_count=1,
            last_seen_frame=0,
        )

    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 = [
        {
            "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.hex(),
            "semantic_hash": obj.semantic_hash.hex(),
        }
        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(mo):
    file_upload = mo.ui.file(
        filetypes=["image/*"],
        kind="area",
        label="Upload a query image",
    )
    file_upload
    return (file_upload,)


@app.cell
def _(file_upload, mo):
    upload_image = None
    if file_upload.value:
        upload_image = mo.image(file_upload.contents(), alt="Uploaded query image")

    # Build a grid.
    upload_image if upload_image is not None else mo.md("No image uploaded yet.")
    return


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