import marimo

__generated_with = "0.21.1"
app = marimo.App()


@app.cell
def _():
    import marimo as mo
    import numpy as np
    from PIL import Image, ImageDraw, ImageFont

    return Image, ImageDraw, ImageFont, mo, np


@app.cell
def _(np):
    from scenegraph import RoomNode
    from simulator import (
        HabitatSimulatorConfig,
        TopDownSceneElements,
        create_habitat_simulator,
        render_topdown_scene_map,
    )
    from habitat.utils.visualizations import maps

    scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
    image_size = 768
    num_rooms = 3
    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),
            )
        )

    render_topdown_scene_map(
        pathfinder=sim.pathfinder,
        elements=TopDownSceneElements(room_nodes=room_nodes),
        meters_per_pixel=meters_per_pixel,
        maps_module=maps,
    )
    return agent, room_nodes, sim, views_per_room


@app.cell
def _(agent, room_nodes, sim, views_per_room):
    from simulator import collect_room_views_by_room

    all_room_views = collect_room_views_by_room(
        agent=agent,
        sim=sim,
        room_nodes=room_nodes,
        views_per_room=views_per_room,
    )
    return (all_room_views,)


@app.cell
def _(ImageDraw, ImageFont, all_room_views, mo):
    from compressors import HashPipeline
    from utils.common import get_device
    from utils.image import numpy_to_pil

    device = get_device()

    score_threshold = 0.25
    postprocess_threshold = 0.1
    text_labels = [
        "a cat",
        "a dog",
        "a TV remote control",
        "a chair",
        "a table",
        "a vase",
        "a painting",
        "a window",
    ]

    pipeline = HashPipeline(
        owlv2_model="google/owlv2-base-patch16-ensemble",
        score_threshold=score_threshold,
        postprocess_threshold=postprocess_threshold,
    )

    image = numpy_to_pil(all_room_views["room_00"][1])

    output = pipeline.process_batch(
        images=[image],
        text_labels=text_labels,
        batch_size=1,
        return_debug_details=True,
    )
    meta = output.debug_meta[0] if output.debug_meta else {}
    boxes = meta.get("boxes_xyxy", [])
    scores = meta.get("scores", [])
    labels = meta.get("labels", [])
    filtered_items = list(zip(boxes, scores, labels, strict=False))

    _vis_image = image.copy()
    _draw = ImageDraw.Draw(_vis_image)

    try:
        font = ImageFont.truetype("arial.ttf", 18)
    except Exception:
        font = ImageFont.load_default()

    for _box, _score, _text_label in filtered_items:
        _x1, _y1, _x2, _y2 = [float(v) for v in _box]
        _label = f"{_text_label}: {_score:.3f}"

        _draw.rectangle([_x1, _y1, _x2, _y2], outline="red", width=3)

        try:
            _tx1, _ty1, _tx2, _ty2 = _draw.textbbox((_x1, _y1), _label, font=font)
        except Exception:
            _w = _draw.textlength(_label, font=font)
            _h = 20
            _tx1, _ty1, _tx2, _ty2 = _x1, _y1, _x1 + _w + 6, _y1 + _h

        text_bg = [_tx1, max(0, _ty1 - 2), _tx2 + 4, _ty2 + 2]
        _draw.rectangle(text_bg, fill="red")
        _draw.text((_x1 + 2, max(0, _y1)), _label, fill="white", font=font)

    # 8. 结果文本
    detection_lines = []
    for _box, _score, _text_label in filtered_items:
        box_rounded = [round(_v, 2) for _v in _box]
        detection_lines.append(
            f"- {_text_label}: score={_score:.3f}, box={box_rounded}"
        )

    if not detection_lines:
        detection_text = f"没有检测到 score >= {score_threshold:.2f} 的目标"
    else:
        detection_text = "\n".join(detection_lines)

    mo.vstack(
        [
            mo.md(
                "## OWLv2 检测可视化结果"
                f"\n\ndevice: `{device}`"
                f"\n\n过滤阈值：`score >= {score_threshold:.2f}`"
            ),
            mo.image(_vis_image, width=700),
            mo.md(detection_text),
        ]
    )
    return device, filtered_items, image, meta


@app.cell
def _(meta):
    proposals = meta.get("masks", [])
    return (proposals,)


@app.cell
def _(Image, ImageDraw, filtered_items, image, mo, np, proposals):
    from compressors.filter import (
        MaskScoringConfig,
        compute_mask_features,
        score_mask,
        should_reject_mask,
    )

    image_shape = (image.height, image.width)
    config = MaskScoringConfig()

    _kept = []
    _rejected = []
    for _idx, proposal in enumerate(proposals):
        _feat = compute_mask_features(proposal, image_shape)
        _is_rejected = should_reject_mask(_feat, config)
        _score = score_mask(proposal, image_shape, config)
        _owl_label = (
            filtered_items[_idx][2] if _idx < len(filtered_items) else f"obj_{_idx}"
        )
        _owl_score = filtered_items[_idx][1] if _idx < len(filtered_items) else 0.0
        _owl_bbox = (
            filtered_items[_idx][0] if _idx < len(filtered_items) else [0, 0, 0, 0]
        )

        _entry = {
            "idx": _idx,
            "proposal": proposal,
            "features": _feat,
            "mask_score": _score,
            "owl_label": _owl_label,
            "owl_score": _owl_score,
            "owl_bbox": _owl_bbox,
        }
        if _is_rejected:
            _rejected.append(_entry)
        else:
            _kept.append(_entry)

    _colors = [
        (255, 0, 0, 90),
        (0, 255, 0, 90),
        (0, 0, 255, 90),
        (255, 255, 0, 90),
        (255, 0, 255, 90),
        (0, 255, 255, 90),
        (255, 128, 0, 90),
        (128, 0, 255, 90),
    ]

    def _overlay_masks(base_img, entries, border_color, show_score=False):
        """Overlay masks on image and draw bounding boxes with labels."""
        _rgba = base_img.convert("RGBA").copy()
        for _e in entries:
            _p = _e["proposal"]
            _c = _colors[_e["idx"] % len(_colors)]
            _mask_rgba = np.zeros((base_img.height, base_img.width, 4), dtype=np.uint8)
            _mask_rgba[_p["segment"]] = _c
            _rgba = Image.alpha_composite(
                _rgba, Image.fromarray(_mask_rgba, mode="RGBA")
            )

        _draw = ImageDraw.Draw(_rgba)
        for _e in entries:
            _x1, _y1, _x2, _y2 = [float(v) for v in _e["owl_bbox"]]
            _mask_area = int(_e["proposal"]["area"])
            if show_score:
                _label = f"{_e['owl_label']} | score={_e['mask_score']:.2f} | area={_mask_area}"
            else:
                _label = f"{_e['owl_label']} | area={_mask_area}"
            _draw.rectangle([_x1, _y1, _x2, _y2], outline=border_color, width=3)

            try:
                _tb = _draw.textbbox((_x1, _y1), _label)
            except Exception:
                _tb = (_x1, _y1, _x1 + 200, _y1 + 20)
            _draw.rectangle(
                [_tb[0], max(0, _tb[1] - 2), _tb[2] + 4, _tb[3] + 2],
                fill=border_color,
            )
            _draw.text((_x1 + 2, max(0, _y1)), _label, fill="white")
        return _rgba

    _all_entries = _kept + _rejected
    _all_entries.sort(key=lambda e: e["idx"])

    _before_img = _overlay_masks(image, _all_entries, "red", show_score=False)
    _after_img = _overlay_masks(image, _kept, (0, 180, 0), show_score=True)

    _draw_after = ImageDraw.Draw(_after_img)
    for _e in _rejected:
        _x1, _y1, _x2, _y2 = [float(v) for v in _e["owl_bbox"]]
        _feat = _e["features"]
        _reason_parts = []
        if _feat.area_ratio < config.min_area_ratio:
            _reason_parts.append(f"area_ratio={_feat.area_ratio:.4f}<min")
        if _feat.area_ratio > config.max_area_ratio:
            _reason_parts.append(f"area_ratio={_feat.area_ratio:.4f}>max")
        _aspect = max(_feat.aspect_ratio, 1.0 / max(_feat.aspect_ratio, 1e-6))
        if _aspect > config.max_aspect_ratio:
            _reason_parts.append(f"aspect={_aspect:.1f}>max")
        if _feat.fill_ratio < config.min_fill_ratio_hard:
            _reason_parts.append(f"fill={_feat.fill_ratio:.3f}<min")
        if (
            _feat.num_components > config.max_components
            and _feat.largest_component_ratio < config.min_largest_component_ratio
        ):
            _reason_parts.append(f"fragments={_feat.num_components}")
        if _feat.touch_edge_count >= config.reject_edge_touch_count:
            _reason_parts.append(f"edge_touch={_feat.touch_edge_count}")
        if (
            _feat.touch_edge_count >= config.reject_large_edge_touch_count
            and _feat.area_ratio > config.reject_large_edge_area_ratio
        ):
            _reason_parts.append("edge+large")

        _reason = ", ".join(_reason_parts) if _reason_parts else "unknown"
        _label = f"X {_e['owl_label']}: {_reason}"

        _dash_len = 8
        _gap_len = 4
        for _seg_start in range(int(_x1), int(_x2), _dash_len + _gap_len):
            _seg_end = min(_seg_start + _dash_len, int(_x2))
            _draw_after.line([(_seg_start, _y1), (_seg_end, _y1)], fill="red", width=2)
            _draw_after.line([(_seg_start, _y2), (_seg_end, _y2)], fill="red", width=2)
        for _seg_start in range(int(_y1), int(_y2), _dash_len + _gap_len):
            _seg_end = min(_seg_start + _dash_len, int(_y2))
            _draw_after.line([(_x1, _seg_start), (_x1, _seg_end)], fill="red", width=2)
            _draw_after.line([(_x2, _seg_start), (_x2, _seg_end)], fill="red", width=2)

        _draw_after.text((_x1 + 2, max(0, _y1 - 18)), _label, fill="red")

    _total = len(proposals)
    _kept_count = len(_kept)
    _rej_count = len(_rejected)

    _rej_detail_lines = []
    for _e in _rejected:
        _f = _e["features"]
        _rej_detail_lines.append(
            f"  - **{_e['owl_label']}** (idx={_e['idx']}): "
            f"area_ratio={_f.area_ratio:.4f}, fill_ratio={_f.fill_ratio:.3f}, "
            f"aspect_ratio={_f.aspect_ratio:.1f}, touch_edge={_f.touch_edge_count}, "
            f"components={_f.num_components}, mask_score={_e['mask_score']:.3f}"
        )

    _kept_detail_lines = []
    for _e in _kept:
        _f = _e["features"]
        _kept_detail_lines.append(
            f"  - **{_e['owl_label']}** (idx={_e['idx']}): "
            f"area_ratio={_f.area_ratio:.4f}, mask_score={_e['mask_score']:.3f}"
        )

    _detail_parts = []
    if _kept_detail_lines:
        _detail_parts.append("**保留的 mask：**\n" + "\n".join(_kept_detail_lines))
    if _rej_detail_lines:
        _detail_parts.append("**过滤掉的 mask：**\n" + "\n".join(_rej_detail_lines))

    _detail_text = "\n\n".join(_detail_parts) if _detail_parts else "无 mask 数据"

    mo.vstack(
        [
            mo.md(
                "## Mask 过滤对比"
                f"\n\n共 {_total} 个 mask → 保留 **{_kept_count}** 个，过滤掉 **{_rej_count}** 个"
            ),
            mo.hstack(
                [
                    mo.vstack(
                        [
                            mo.md(f"### 过滤前（{_total} 个）"),
                            mo.image(_before_img, width=480),
                        ]
                    ),
                    mo.vstack(
                        [
                            mo.md(
                                f"### 过滤后（{_kept_count} 个保留，{_rej_count} 个过滤）"
                            ),
                            mo.image(_after_img, width=480),
                        ]
                    ),
                ]
            ),
            mo.md(_detail_text),
        ]
    )
    return


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