import marimo

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


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

    return Image, ImageDraw, ImageFont, mo, np, torch


@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 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 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, torch):
    from transformers import Owlv2Processor, Owlv2ForObjectDetection
    from utils import get_device, numpy_to_pil

    device = get_device()

    # 你真正想控制的显示阈值
    score_threshold = 0.25

    # 1. 加载模型
    processor = Owlv2Processor.from_pretrained(
        "google/owlv2-base-patch16-ensemble"
    )
    model = Owlv2ForObjectDetection.from_pretrained(
        "google/owlv2-base-patch16-ensemble"
    ).to(device)
    model.eval()

    # 2. 读取图片
    image = numpy_to_pil(all_room_views["room_00"][1])

    # 3. 文本查询
    text_labels = [
        [
            "a cat",
            "a dog",
            "a TV remote control",
            "a chair",
            "a table",
            "a vase",
            "a painting",
            "a window"
        ]
    ]

    # 4. 推理
    inputs = processor(text=text_labels, images=image, return_tensors="pt").to(
        device
    )

    with torch.no_grad():
        outputs = model(**inputs)

    # 5. 后处理
    target_sizes = torch.tensor([(image.height, image.width)])
    results = processor.post_process_grounded_object_detection(
        outputs=outputs,
        target_sizes=target_sizes,
        threshold=0.1,  # 模型后处理的第一层粗筛
        text_labels=text_labels,
    )

    result = results[0]
    boxes = result["boxes"]
    scores = result["scores"]
    pred_labels = result["text_labels"]

    # 6. 二次过滤：只保留高分目标
    filtered_items = [
        (_box, _score, _text_label)
        for _box, _score, _text_label in zip(boxes, scores, pred_labels)
        if _score.item() >= score_threshold
    ]

    # 7. 画框
    _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.tolist()]
        _label = f"{_text_label}: {_score.item():.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.tolist()]
        detection_lines.append(
            f"- {_text_label}: score={_score.item():.3f}, box={box_rounded}"
        )

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

    # 9. marimo 输出
    mo.vstack(
        [
            mo.md(
                f"## OWLv2 检测可视化结果\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, torch):
    from transformers import Sam2Processor, Sam2Model

    # -----------------------------
    # 0. 加载 SAM2
    # -----------------------------
    sam2_model_id = "facebook/sam2.1-hiera-large"

    sam2_processor = Sam2Processor.from_pretrained(sam2_model_id)
    sam2_model = Sam2Model.from_pretrained(sam2_model_id).to(device)
    sam2_model.eval()

    # -----------------------------
    # 1. 从上一格的 OWLv2 结果中取框
    # filtered_items: [(box, score, text_label), ...]
    # SAM2 需要的 box 格式：
    # [[[x1, y1, x2, y2], [x1, y1, x2, y2], ...]]
    # 这里最外层是 batch 维度，因为我们当前只处理 1 张图
    # -----------------------------
    input_boxes = [
        [
            [float(v) for v in box.tolist()]
            for box, score, text_label in filtered_items
        ]
    ]

    # 没有框就直接提前返回
    if len(input_boxes[0]) == 0:
        mo.vstack(
            [
                mo.md("## SAM2 分割结果"),
                mo.md(
                    "没有可用于分割的检测框，请先降低 OWLv2 的 score_threshold 或检查检测结果。"
                ),
            ]
        )
    else:
        # -----------------------------
        # 2. 预处理并推理
        # 官方文档支持 bounding box prompt
        # 并且 multimask_output=False 时，每个框只返回一张最佳 mask
        # -----------------------------
        sam2_inputs = sam2_processor(
            images=image,
            input_boxes=input_boxes,
            return_tensors="pt",
        ).to(device)

        with torch.no_grad():
            sam2_outputs = sam2_model(**sam2_inputs, multimask_output=False)

        # -----------------------------
        # 3. 后处理到原图尺寸
        # 文档说明 pred_masks 需要 post_process_masks
        # -----------------------------
        all_masks = sam2_processor.post_process_masks(
            sam2_outputs.pred_masks.cpu(),
            sam2_inputs["original_sizes"],
        )

        # 当前只有 1 张图，所以取第 0 个
        # 形状通常是 [num_objects, num_masks, H, W]
        # 由于 multimask_output=False，num_masks 通常为 1
        masks_for_image = all_masks[0]

        # 兼容处理：压掉单 mask 维度
        # 目标形状变成 [num_objects, H, W]
        if masks_for_image.ndim == 4 and masks_for_image.shape[1] == 1:
            masks_for_image = masks_for_image[:, 0]

        # IoU 分数，文档中也给出了 outputs.iou_scores
        # 形状通常是 [batch_size, point_batch_size, num_masks]
        iou_scores = sam2_outputs.iou_scores.detach().cpu()[0]
        if iou_scores.ndim == 2 and iou_scores.shape[1] == 1:
            iou_scores = iou_scores[:, 0]

        # -----------------------------
        # 4. 可视化：把 mask 半透明叠加到原图上
        # -----------------------------
        base_rgba = image.convert("RGBA")
        overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
        overlay_draw = ImageDraw.Draw(overlay)

        # 准备一个可重复的颜色列表
        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),
        ]

        # 先画 mask，再画 box
        _vis_image = base_rgba.copy()

        for _idx, ((_box, _score, _text_label), mask) in enumerate(
            zip(filtered_items, masks_for_image)
        ):
            color = colors[_idx % len(colors)]

            # mask -> numpy bool
            mask_np = mask.numpy() > 0

            # 做一层彩色 mask
            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 _idx, ((_box, _score, _text_label), mask) in enumerate(
            zip(filtered_items, masks_for_image)
        ):
            _x1, _y1, _x2, _y2 = [float(v) for v in _box.tolist()]
            iou_val = (
                float(iou_scores[_idx].item())
                if _idx < len(iou_scores)
                else float("nan")
            )
            _label = (
                f"{_text_label} | owl={_score.item():.3f} | sam2_iou={iou_val:.3f}"
            )

            _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")

        # -----------------------------
        # 5. 输出摘要
        # -----------------------------
        summary_lines = []
        for _idx, ((_box, _score, _text_label), mask) in enumerate(
            zip(filtered_items, masks_for_image)
        ):
            mask_area = int((mask.numpy() > 0).sum())
            iou_val = (
                float(iou_scores[_idx].item())
                if _idx < len(iou_scores)
                else float("nan")
            )
            summary_lines.append(
                f"- {_text_label}: owl_score={_score.item():.3f}, sam2_iou={iou_val:.3f}, mask_area={mask_area}"
            )


    mo.vstack(
        [
            mo.md("## SAM2 分割可视化结果"),
            mo.image(_vis_image, width=700),
            mo.md(
                "\n".join(summary_lines) if summary_lines else "没有生成任何 mask"
            ),
        ]
    )
    return


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