# /// 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
    import torch
    from PIL import Image

    return Image, np, pl, torch


@app.cell
def project_imports():
    """Project module imports using new architecture."""
    from compressors import HashPipeline, hamming_distance
    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,
        save_object_image,
        save_room_view,
    )
    from utils.image import numpy_to_pil

    return (
        HashPipeline,
        HabitatSimulatorConfig,
        ObjectNode,
        RoomNode,
        SimpleSceneGraph,
        TopDownSceneElements,
        cfg_manager,
        collect_room_views_by_room,
        create_habitat_simulator,
        hamming_distance,
        numpy_to_pil,
        render_topdown_scene_map,
        save_object_image,
        save_room_view,
    )


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

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

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

    # Flatten room views into (room_id, view_idx, PIL.Image) tuples.
    room_view_dataset = [
        (_room_id, _view_idx, numpy_to_pil(_rgb))
        for _room_id, _views in all_room_views.items()
        for _view_idx, _rgb in enumerate(_views)
    ]

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


@app.cell
def build_scene_graph(
    ObjectNode,
    SimpleSceneGraph,
    cfg_manager,
    np,
    pipeline,
    room_nodes,
    room_view_dataset,
    save_object_image,
    save_room_view,
):
    scene_graph = SimpleSceneGraph(
        rooms={_room.room_id: _room for _room in room_nodes},
        objects={},
    )

    # Storage for cropped object images (for visualization).
    object_images = {}

    output_dir = cfg_manager.get().output.directory / "verification"
    output_dir.mkdir(parents=True, exist_ok=True)

    _images = [item[2] for item in room_view_dataset]
    _metadata = [(item[0], item[1]) for item in room_view_dataset]

    _text_labels = [
        "a chair",
        "a table",
        "a sofa",
        "a cabinet",
        "a shelf",
        "a lamp",
        "a picture",
        "a window",
        "a door",
        "a plant",
    ]
    _output = pipeline.process_batch(
        _images, _text_labels, batch_size=32, return_debug_details=True
    )
    _cropped_images = _output.cropped_images
    hash_tensor = _output.hash_bits

    from collections import Counter

    _reasons = Counter(m["fallback_reason"] or "ok" for m in _output.debug_meta)
    print(f"Fallback breakdown: {_reasons}")

    # Save original room views.
    for _room_id, _view_idx, _image in room_view_dataset:
        save_room_view(output_dir, _room_id, _view_idx, _image)

    # Prefix sum: map flat crop index to (input_image_idx, mask_idx).
    _num_selected = [_m["num_selected"] for _m in _output.debug_meta]
    assert sum(_num_selected) == len(_cropped_images), (
        f"Sum of num_selected ({sum(_num_selected)}) != cropped_images count ({len(_cropped_images)})"
    )
    _prefix_sums = [0]
    for _n in _num_selected:
        _prefix_sums.append(_prefix_sums[-1] + _n)

    _obj_counter = 0
    for _img_idx, _n_crops in enumerate(_num_selected):
        _room_id, _view_idx = _metadata[_img_idx]
        for _mask_idx in range(_n_crops):
            _crop_flat_idx = _prefix_sums[_img_idx] + _mask_idx
            _cropped = _cropped_images[_crop_flat_idx]
            _hash_bits = hash_tensor[_crop_flat_idx]

            _obj_id = f"{_room_id}_v{_view_idx:03d}_m{_mask_idx:02d}"

            _bits_array = _hash_bits.detach().cpu().numpy().reshape(-1)
            _bits_binary = (_bits_array > 0).astype(np.uint8)
            _hash_bytes = np.packbits(_bits_binary).tobytes()

            object_images[_obj_id] = _cropped
            save_object_image(
                output_dir, _room_id, _obj_id, _view_idx, _mask_idx, _cropped
            )

            scene_graph.objects[_obj_id] = ObjectNode(
                obj_id=_obj_id,
                room_id=_room_id,
                position=np.array([0.0, 0.0, 0.0], dtype=np.float32),
                visual_hash=_hash_bytes,
                semantic_hash=_hash_bytes,
                hit_count=1,
                last_seen_frame=_view_idx,
            )
            _obj_counter += 1

    _fallback_count = sum(
        1 for _meta in _output.debug_meta if _meta["fallback_reason"] is not None
    )

    print(f"Created {_obj_counter} objects")
    print(f"Saved cropped images to: {output_dir}")
    print(f"Fallback frames: {_fallback_count}/{len(_output.debug_meta)}")

    return hash_tensor, object_images, output_dir, scene_graph


@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] + "...",
        }
        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 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,
    hamming_distance,
    np,
    mo,
    object_images,
    pipeline,
    scene_graph,
    torch,
):
    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")

    _text_labels = [
        "a chair",
        "a table",
        "a sofa",
        "a cabinet",
        "a shelf",
        "a lamp",
        "a picture",
        "a window",
        "a door",
        "a plant",
    ]
    _output = pipeline.process_batch([_query_image], _text_labels, batch_size=1)
    _query_bits = (_output.hash_bits > 0).to(dtype=torch.int32)

    if _query_bits.numel() > 0 and scene_graph.objects:
        _obj_ids = list(scene_graph.objects.keys())
        _obj_hashes = []
        for _obj_id in _obj_ids:
            _obj = scene_graph.objects[_obj_id]
            _bits = np.unpackbits(np.frombuffer(_obj.visual_hash, dtype=np.uint8))[
                : pipeline.hash_bits
            ].astype(np.int32)
            _obj_hashes.append(_bits)

        _db_tensor = torch.tensor(np.array(_obj_hashes), dtype=torch.int32).to(
            _query_bits.device
        )

        _distances = hamming_distance(_query_bits, _db_tensor)
        _best_query_idx = int(_distances.min(dim=1).values.argmin().item())

        _query_tensor = _query_bits[_best_query_idx]
        query_cropped = _output.cropped_images[_best_query_idx]
        _query_distances = _distances[_best_query_idx].cpu().numpy()
        _query_hash_hex = (
            np.packbits(_query_tensor.cpu().numpy().astype(np.uint8)).tobytes().hex()
        )

        _top_k = min(5, len(_obj_ids))
        _top_indices = np.argsort(_query_distances)[:_top_k]

        top_matches = [
            {
                "obj_id": _obj_ids[_i],
                "distance": int(_query_distances[_i]),
                "similarity": 1.0 - _query_distances[_i] / float(pipeline.hash_bits),
                "hash_hex": scene_graph.objects[_obj_ids[_i]].visual_hash.hex(),
            }
            for _i in _top_indices
        ]

        query_result = {
            "query_cropped": query_cropped,
            "query_hash_hex": _query_hash_hex,
            "top_matches": top_matches,
        }

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