# /// 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 base_dependencies():
    """Basic dependencies for data processing."""
    import numpy as np
    import polars as pl
    from PIL import Image

    return Image, np, pl


@app.cell
def project_imports():
    """Project module imports using new architecture."""
    from compressors import HashPipeline
    from configs import cfg_manager
    from scenegraph import (
        RoomNode,
        SceneGraphBuildConfig,
        SceneGraphBuilder,
        load_scene_graph,
        query_image_against_scene_graph,
        save_scene_graph,
    )
    from simulator import (
        HabitatSimulatorConfig,
        TopDownRenderStyle,
        TopDownSceneElements,
        collect_room_views_by_room,
        create_habitat_simulator,
        flatten_room_views,
        render_topdown_scene_map,
        save_object_image,
        save_room_view,
    )
    from utils.image import numpy_to_pil

    return (
        HashPipeline,
        HabitatSimulatorConfig,
        RoomNode,
        SceneGraphBuildConfig,
        SceneGraphBuilder,
        TopDownRenderStyle,
        TopDownSceneElements,
        cfg_manager,
        collect_room_views_by_room,
        create_habitat_simulator,
        flatten_room_views,
        load_scene_graph,
        numpy_to_pil,
        query_image_against_scene_graph,
        render_topdown_scene_map,
        save_object_image,
        save_room_view,
        save_scene_graph,
    )


@app.cell
def text_labels():
    """Shared text labels for detection during graph build and query."""
    text_labels = [
        "a chair",
        "a table",
        "a sofa",
        "a cabinet",
        "a shelf",
        "a lamp",
        "a picture",
        "a window",
        "a door",
        "a plant",
    ]
    return (text_labels,)


@app.cell
def habitat_setup(HabitatSimulatorConfig, RoomNode, create_habitat_simulator, np):
    """Initialize Habitat simulator and sample room nodes."""
    _scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
    _image_size = 768
    _num_rooms = 5
    views_per_room = 12
    meters_per_pixel = 0.05

    habitat_config = 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,
    )
    sim, agent = create_habitat_simulator(habitat_config)

    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(f"  {_node.room_id}: {_node.center}")

    return agent, habitat_config, meters_per_pixel, room_nodes, sim, views_per_room


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


@app.cell
def pipeline_init(HashPipeline):
    pipeline = HashPipeline(
        dino_model="facebook/dinov2-large",
        sam_model="facebook/sam2.1-hiera-large",
        hash_bits=512,
        score_threshold=0.10,
        postprocess_threshold=0.05,
    )
    print(f"Pipeline initialized: {pipeline.hash_bits} bits")
    return (pipeline,)


@app.cell
def collect_views(
    agent,
    collect_room_views_by_room,
    flatten_room_views,
    habitat_config,
    numpy_to_pil,
    room_nodes,
    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,
        depth_sensor_uuid=habitat_config.depth_sensor_uuid,
    )

    room_views = flatten_room_views(all_room_views)

    # Build dataset of (room_id, view_idx, PIL.Image) tuples.
    room_view_dataset = [
        (_view.room_id, _view.view_idx, numpy_to_pil(_view.rgb))
        for _view in room_views
    ]

    print(f"Collected {len(room_view_dataset)} room views")
    return all_room_views, room_view_dataset, room_views


@app.cell
def build_scene_graph(
    SceneGraphBuildConfig,
    SceneGraphBuilder,
    cfg_manager,
    habitat_config,
    load_scene_graph,
    mo,
    numpy_to_pil,
    pipeline,
    room_nodes,
    room_views,
    save_object_image,
    save_room_view,
    save_scene_graph,
    text_labels,
):
    output_dir = cfg_manager.get().output.directory / "verification"
    output_dir.mkdir(parents=True, exist_ok=True)
    cache_path = output_dir / "scene_graph.json"

    # Try loading cached scene graph first.
    if cache_path.exists():
        scene_graph = load_scene_graph(cache_path)
        print(f"Loaded cached scene graph from {cache_path}")
        print(f"  {len(scene_graph.rooms)} rooms, {len(scene_graph.objects)} objects")
        object_images = {}
        build_artifacts = None
    else:
        builder = SceneGraphBuilder(
            pipeline=pipeline,
            config=SceneGraphBuildConfig(
                inference_batch_size=4,
                position_strategy="bbox_depth_center",
                camera_hfov_degrees=habitat_config.hfov_degrees,
            ),
        )

        pil_room_views = [
            type(_view)(
                room_id=_view.room_id,
                view_idx=_view.view_idx,
                rgb=numpy_to_pil(_view.rgb),
                depth=_view.depth,
                agent_position=_view.agent_position,
                agent_rotation=_view.agent_rotation,
                camera_position=_view.camera_position,
                camera_rotation=_view.camera_rotation,
            )
            for _view in room_views
        ]

        with mo.status.spinner(title="Building scene graph from room views"):
            scene_graph, build_artifacts = builder.build_from_room_views(
                room_nodes=room_nodes,
                room_views=pil_room_views,
                text_labels=text_labels,
            )

        object_images = build_artifacts.object_images
        debug_meta = build_artifacts.debug_meta

        # Save scene graph to cache.
        save_scene_graph(cache_path, scene_graph)
        print(f"Saved scene graph to {cache_path}")

        # Save original room views.
        for _room_view in mo.status.progress_bar(
            pil_room_views,
            title="Saving room-view snapshots",
            subtitle="Writing original room images to disk",
            completion_title="Room-view snapshots saved",
            completion_subtitle=f"Saved {len(pil_room_views)} room views",
            show_eta=True,
            show_rate=True,
            remove_on_exit=False,
        ):
            save_room_view(output_dir, _room_view.room_id, _room_view.view_idx, _room_view.rgb)

        # Save object crops.
        for _obj_id, _cropped in mo.status.progress_bar(
            object_images.items(),
            title="Saving object crops",
            subtitle="Writing cropped object images to disk",
            completion_title="Object crops saved",
            completion_subtitle=f"Saved {len(object_images)} object crops",
            show_eta=True,
            show_rate=True,
            remove_on_exit=False,
        ):
            _node = scene_graph.objects[_obj_id]
            save_object_image(
                output_dir,
                _node.room_id,
                _obj_id,
                _node.last_seen_frame,
                0,  # mask idx 0 ok for M0
                _cropped,
            )

        from collections import Counter

        _meta_fallbacks = [_meta.get("fallback_reason") for _meta in debug_meta]
        fallback_count = sum(1 for f in _meta_fallbacks if f is not None)
        _reasons = Counter(f or "ok" for f in _meta_fallbacks)
        print(f"Fallback breakdown: {_reasons}")
        print(f"Fallback frames: {fallback_count}/{len(debug_meta)}")

    print(f"Created {len(scene_graph.objects)} objects")
    print(f"Output directory: {output_dir}")

    return build_artifacts, object_images, output_dir, scene_graph


@app.cell
def render_scene_graph_birdseye(
    TopDownRenderStyle,
    TopDownSceneElements,
    meters_per_pixel,
    mo,
    render_topdown_scene_map,
    room_nodes,
    scene_graph,
    sim,
):
    """Bird's-eye view of the full scene graph: rooms, objects, and edges."""
    _object_nodes = list(scene_graph.objects.values())
    _edges = [
        (_obj.room_id, _obj.obj_id) for _obj in _object_nodes
    ]

    _style = TopDownRenderStyle(
        title="Scene Graph — Bird's-Eye View",
        figure_size=(10, 10),
    )

    _elements = TopDownSceneElements(
        room_nodes=room_nodes,
        object_nodes=_object_nodes,
        edges=_edges,
    )

    _birdseye_image = render_topdown_scene_map(
        pathfinder=sim.pathfinder,
        elements=_elements,
        meters_per_pixel=meters_per_pixel,
        style=_style,
    )

    mo.md("## Scene Graph Bird's-Eye View"), mo.image(_birdseye_image)
    return (_birdseye_image,)


@app.cell
def build_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]),
        }
        for _room in scene_graph.rooms.values()
    ]

    object_rows = [
        {
            "obj_id": obj.obj_id,
            "room_id": obj.room_id,
            "visual_hash": obj.visual_hash.hex()[:16] + "...",
            "label": obj.label,
            "confidence": obj.confidence,
            "source_view_id": obj.source_view_id,
            "position_x": float(obj.position[0]) if obj.position is not None else None,
            "position_y": float(obj.position[1]) if obj.position is not None else None,
            "position_z": float(obj.position[2]) if obj.position is not None else None,
            "bbox_xyxy": str(obj.bbox_xyxy) if obj.bbox_xyxy is not None else None,
            "position_confidence": obj.position_confidence,
        }
        for obj in scene_graph.objects.values()
    ]

    rooms_df = pl.DataFrame(room_rows)
    objects_df = pl.DataFrame(object_rows)

    return objects_df, rooms_df


@app.cell
def display_tables(mo, objects_df, rooms_df):
    mo.vstack(
        [
            mo.md("## Rooms"),
            mo.ui.table(rooms_df),
            mo.md("## Objects"),
            mo.ui.table(objects_df),
        ]
    )
    return


@app.cell
def display_objects_by_room(mo, objects_df, pl):
    """Table of all objects with their room assignments."""
    _obj_room_df = objects_df.select(
        ["obj_id", "room_id", "label", "confidence", "position_confidence"]
    ).sort(["room_id", "label"])

    mo.vstack([
        mo.md("## Objects by Room"),
        mo.ui.table(_obj_room_df),
    ])
    return


@app.cell
def display_room_summary(mo, objects_df, pl, rooms_df):
    """Table of all rooms with their object counts."""
    _counts = (
        objects_df.group_by("room_id")
        .agg(pl.col("obj_id").count().alias("object_count"))
    )
    _room_summary = rooms_df.select(["room_id"]).join(
        _counts, on="room_id", how="left"
    ).with_columns(
        pl.col("object_count").fill_null(0).cast(pl.Int64)
    ).sort("room_id")

    mo.vstack([
        mo.md("## Room Summary"),
        mo.ui.table(_room_summary),
    ])
    return


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


@app.cell
def query_matching(
    Image,
    file_upload,
    mo,
    pipeline,
    query_image_against_scene_graph,
    scene_graph,
    text_labels,
):
    from io import BytesIO

    query_result = None
    query_cropped = None
    top_matches = []

    _file_contents = file_upload.contents()
    mo.stop(not _file_contents, mo.md("请先上传文件"))

    _query_image = Image.open(BytesIO(_file_contents)).convert("RGB")

    _query_results = query_image_against_scene_graph(
        image=_query_image,
        pipeline=pipeline,
        scene_graph=scene_graph,
        text_labels=text_labels,
        top_k=5,
        batch_size=1,
    )

    if _query_results:
        _best_result = max(
            _query_results,
            key=lambda result: result.matches[0].score if result.matches else -1,
        )
        query_cropped = _best_result.query_crop
        top_matches = [
            {
                "obj_id": match.obj_id,
                "distance": int(pipeline.hash_bits - match.score),
                "similarity": match.similarity,
                "hash_hex": match.hash_bytes.hex(),
            }
            for match in _best_result.matches
        ]

        query_result = {
            "query_cropped": query_cropped,
            "query_hash_hex": _best_result.query_hash.hex(),
            "top_matches": top_matches,
            "num_query_results": len(_query_results),
        }

    return query_cropped, query_result, top_matches


@app.cell
def display_results(mo, object_images, query_cropped, query_result, top_matches):
    mo.stop(not query_result, mo.md("No query results yet. Upload an image above."))

    _result_items = [
        mo.vstack(
            [
                mo.md("**Query (cropped)**"),
                mo.image(query_cropped),
            ],
            align="center",
        )
    ]

    for _match in top_matches:
        _obj_id = _match["obj_id"]
        _obj_img = object_images.get(_obj_id)

        if _obj_img is not None:
            _result_items.append(
                mo.vstack(
                    [
                        mo.md(f"**{_obj_id}**"),
                        mo.image(_obj_img),
                        mo.md(f"Distance: {_match['distance']}"),
                        mo.md(f"Similarity: {_match['similarity']:.2%}"),
                    ],
                    align="center",
                )
            )

    mo.vstack(_result_items, justify="center", gap=2)

    return


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