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.model_loader import load_owlv2_model
    from compressors.proposal.core import detect_objects_batch
    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",
    ]

    owl_processor, owl_model = load_owlv2_model(
        model_name="google/owlv2-base-patch16-ensemble"
    )

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

    detection_batch = detect_objects_batch(
        model=owl_model,
        processor=owl_processor,
        images=[image],
        text_labels_per_image=[text_labels],
        score_threshold=score_threshold,
        postprocess_threshold=postprocess_threshold,
    )
    detections = detection_batch[0] if detection_batch else []
    filtered_items = [(det["bbox"], det["score"], det["label"]) for det in detections]

    _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


@app.cell
def _(Image, ImageDraw, device, filtered_items, image, mo, np):
    from compressors.model_loader import load_sam_model
    from compressors.proposal.core import generate_proposals_batch

    sam2_processor, sam2_model = load_sam_model(
        model_name="facebook/sam2.1-hiera-large"
    )

    input_boxes = [[box for box, _score, _text_label in filtered_items]]
    proposal_batch = generate_proposals_batch(
        model=sam2_model,
        processor=sam2_processor,
        images=[image],
        bboxes_per_image=input_boxes,
    )
    proposals = proposal_batch[0] if proposal_batch else []

    base_rgba = image.convert("RGBA")
    _vis_image = base_rgba.copy()
    summary_lines = []

    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),
    ]

    for _idx, ((_box, _score, _text_label), proposal) in enumerate(
        zip(filtered_items, proposals)
    ):
        mask_np = proposal["segment"]
        color = colors[_idx % len(colors)]

        mask_rgba = np.zeros((image.height, image.width, 4), dtype=np.uint8)
        mask_rgba[mask_np] = color

        mask_img = Image.fromarray(mask_rgba, mode="RGBA")
        _vis_image = Image.alpha_composite(_vis_image, mask_img)

    _draw = ImageDraw.Draw(_vis_image)
    for (_box, _score, _text_label), proposal in zip(filtered_items, proposals):
        _x1, _y1, _x2, _y2 = [float(v) for v in _box]
        mask_area = int(proposal["area"])
        _label = f"{_text_label} | owl={_score:.3f} | mask_area={mask_area}"

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

        try:
            _tx1, _ty1, _tx2, _ty2 = _draw.textbbox((_x1, _y1), _label)
        except Exception:
            _tx1, _ty1, _tx2, _ty2 = _x1, _y1, _x1 + 220, _y1 + 20

        _draw.rectangle(
            [_tx1, max(0, _ty1 - 2), _tx2 + 4, _ty2 + 2],
            fill=(255, 0, 0, 220),
        )
        _draw.text((_x1 + 2, max(0, _y1)), _label, fill="white")
        summary_lines.append(
            f"- {_text_label}: owl_score={_score:.3f}, mask_area={mask_area}"
        )

    if not filtered_items:
        summary_text = (
            "没有可用于分割的检测框，请先降低 OWLv2 的 score_threshold 或检查检测结果。"
        )
    elif not summary_lines:
        summary_text = "没有生成任何 mask"
    else:
        summary_text = "\n".join(summary_lines)

    mo.vstack(
        [
            mo.md(f"## SAM2 分割可视化结果\n\ndevice: `{device}`"),
            mo.image(_vis_image, width=700),
            mo.md(summary_text),
        ]
    )
    return


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