# /// 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 configs import cfg_manager
    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_dataset

    return (
        HabitatSimulatorConfig,
        HashPipeline,
        Image,
        ObjectNode,
        RoomNode,
        SimpleSceneGraph,
        TopDownSceneElements,
        collect_room_views_by_room,
        create_habitat_simulator,
        cfg_manager,
        extract_masked_region,
        maps,
        mo,
        np,
        pl,
        plt,
        render_topdown_scene_map,
        segment_image_dataset,
    )


@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,
    cfg_manager,
    extract_masked_region,
    hash_bits,
    mo,
    np,
    pipeline_batch_size,
    room_nodes,
    sam_max_masks,
    sam_min_area,
    segment_image_dataset,
    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={},
    )
    verification_output_dir = cfg_manager.get().output.directory / "verification"
    verification_output_dir.mkdir(parents=True, exist_ok=True)

    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 = []

    room_view_images = []
    for _, _, rgb in room_view_dataset:
        rgb3 = rgb[..., :3] if rgb.shape[-1] > 3 else rgb
        room_view_images.append(Image.fromarray(rgb3.astype(np.uint8)))

    masks_dataset = segment_image_dataset(
        hash_pipeline.mask_generator,
        room_view_images,
        min_area=hash_pipeline.sam_min_mask_area,
        max_masks=hash_pipeline.sam_max_masks,
        points_per_batch=hash_pipeline.sam_points_per_batch,
    )
    if len(masks_dataset) != len(room_view_dataset):
        raise RuntimeError("SAM dataset output size mismatch with room_view_dataset.")

    dataset_jobs = list(zip(room_view_dataset, room_view_images, masks_dataset))
    for (room_id, view_idx, _), image, masks in mo.status.progress_bar(
        dataset_jobs,
        title="Building object dataset",
        subtitle="Running SAM segmentation",
        show_eta=True,
        show_rate=True,
    ):
        room_output_dir = verification_output_dir / room_id
        room_output_dir.mkdir(parents=True, exist_ok=True)
        room_view_path = room_output_dir / f"view_{view_idx:03d}.png"
        image.convert("RGB").save(room_view_path, format="PNG")

        total_masks += len(masks)

        for mask_idx, mask in enumerate(masks):
            masked_image = extract_masked_region(image, mask["segment"])
            object_dataset.append(
                (room_id, view_idx, mask_idx, mask["bbox"], masked_image)
            )

    if object_dataset:
        masked_images = [item[4] for item in object_dataset]
        if any(not isinstance(img, Image.Image) for img in masked_images):
            raise TypeError(
                "object_dataset contains non-image entries for batch inference."
            )
        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, view_idx, mask_idx, bbox, masked_image) 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

        room_output_dir = verification_output_dir / room_id
        room_output_dir.mkdir(parents=True, exist_ok=True)
        object_image_path = (
            room_output_dir / f"{obj_id}_view{view_idx:03d}_mask{mask_idx:02d}.png"
        )
        masked_image.convert("RGB").save(object_image_path, format="PNG")

        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}")
    print(f"Saved object images to: {verification_output_dir}")
    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()