# /// 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 marimo as mo

    return (mo,)


@app.cell
def global_base_deps():
    import numpy as np
    import polars as pl
    from PIL import Image

    return (Image, np, pl)


@app.cell
def project_setup():
    from configs import cfg_manager
    from scenegraph import ObjectNode, RoomNode, SimpleSceneGraph
    from simulator import (
        HabitatSimulatorConfig,
        TopDownSceneElements,
        collect_room_views_by_room,
        create_habitat_simulator,
        render_topdown_scene_map,
    )

    return (
        HabitatSimulatorConfig,
        ObjectNode,
        RoomNode,
        SimpleSceneGraph,
        TopDownSceneElements,
        cfg_manager,
        collect_room_views_by_room,
        create_habitat_simulator,
        render_topdown_scene_map,
    )


@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 = 512
    num_rooms = 4
    views_per_room = 6
    meters_per_pixel = 0.05

    sam_max_masks = 5
    sam_candidate_masks = 24
    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_candidate_masks,
        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,
):
    from habitat.utils.visualizations import maps
    from matplotlib import pyplot as plt

    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(
    Image,
    ObjectNode,
    SimpleSceneGraph,
    agent,
    cfg_manager,
    collect_room_views_by_room,
    hash_bits,
    np,
    pipeline_batch_size,
    room_nodes,
    sam_candidate_masks,
    sam_max_masks,
    sam_min_area,
    sim,
    views_per_room,
):
    from rich.progress import track

    from compressors import MaskScoringConfig, rank_masks
    from compressors.pipeline import HashPipeline
    from compressors.proposal import extract_masked_region, generate_proposals_batch
    from simulator import save_object_image, save_room_view
    from utils import numpy_to_pil

    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)

    mask_scoring_config = MaskScoringConfig(
        max_area_ratio=0.45,
        reject_edge_touch_count=3,
        reject_large_edge_touch_count=2,
        reject_large_edge_area_ratio=0.12,
        max_components=4,
        min_largest_component_ratio=0.70,
    )

    total_masks = 0
    total_raw_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 = [numpy_to_pil(rgb) for _, _, rgb in room_view_dataset]

    masks_dataset = generate_proposals_batch(
        hash_pipeline.mask_generator,
        room_view_images,
        min_area=hash_pipeline.sam_min_mask_area,
        max_masks=sam_candidate_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 track(
        dataset_jobs,
        description="Building object dataset...",
    ):
        save_room_view(verification_output_dir, room_id, view_idx, image)
        total_raw_masks += len(masks)

        ranked_masks = rank_masks(
            masks,
            image_shape=(image.height, image.width),
            config=mask_scoring_config,
            max_masks=sam_max_masks,
        )
        total_masks += len(ranked_masks)

        for mask_idx, mask in enumerate(ranked_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

        save_object_image(
            verification_output_dir, room_id, obj_id, view_idx, mask_idx, masked_image
        )

        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 raw masks from SAM: {total_raw_masks}")
    print(f"Total kept masks after ranking: {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()