import numpy as np

from compressors.object_score import (
    MaskScoringConfig,
    compute_mask_features,
    rank_masks,
    score_mask,
    select_best_mask,
)


def _rect_mask(height: int, width: int, x: int, y: int, w: int, h: int) -> np.ndarray:
    mask = np.zeros((height, width), dtype=bool)
    mask[y : y + h, x : x + w] = True
    return mask


def test_compute_mask_features_core_metrics() -> None:
    mask = _rect_mask(height=20, width=20, x=5, y=4, w=6, h=5)
    mask_dict = {
        "segment": mask,
        "area": int(mask.sum()),
        "bbox": [5, 4, 6, 5],
        "predicted_iou": 0.8,
        "stability_score": 0.9,
    }

    features = compute_mask_features(mask_dict, image_shape=(20, 20))

    assert features.area_ratio == 30 / 400
    assert features.fill_ratio == 1.0
    assert features.aspect_ratio == 6 / 5
    assert features.touch_top is False
    assert features.touch_left is False
    assert features.num_components == 1
    assert features.largest_component_ratio == 1.0
    assert features.num_holes == 0


def test_rank_masks_rejects_extreme_small_and_fragmented_masks() -> None:
    cfg = MaskScoringConfig(min_area_ratio=0.02)
    good_mask = _rect_mask(height=30, width=30, x=6, y=6, w=10, h=10)

    fragmented = np.zeros((30, 30), dtype=bool)
    fragmented[2, 2] = True
    fragmented[4, 7] = True
    fragmented[8, 12] = True
    fragmented[12, 16] = True
    fragmented[16, 20] = True
    fragmented[20, 24] = True
    fragmented[24, 26] = True

    masks = [
        {"segment": np.zeros((30, 30), dtype=bool), "area": 1, "bbox": [0, 0, 1, 1]},
        {
            "segment": fragmented,
            "area": int(fragmented.sum()),
            "bbox": [2, 2, 25, 25],
        },
        {
            "segment": good_mask,
            "area": int(good_mask.sum()),
            "bbox": [6, 6, 10, 10],
            "predicted_iou": 0.9,
            "stability_score": 0.9,
        },
    ]

    ranked = rank_masks(masks=masks, image_shape=(30, 30), config=cfg, max_masks=3)

    assert len(ranked) == 1
    assert ranked[0]["area"] == int(good_mask.sum())
    assert "mask_score" in ranked[0]


def test_score_mask_prefers_stable_reasonable_object() -> None:
    cfg = MaskScoringConfig()

    candidate = {
        "segment": _rect_mask(height=100, width=100, x=30, y=20, w=24, h=25),
        "area": 24 * 25,
        "bbox": [30, 20, 24, 25],
        "predicted_iou": 0.92,
        "stability_score": 0.91,
    }
    weak = {
        "segment": _rect_mask(height=100, width=100, x=0, y=0, w=4, h=60),
        "area": 4 * 60,
        "bbox": [0, 0, 4, 60],
        "predicted_iou": 0.4,
        "stability_score": 0.3,
    }

    score_candidate = score_mask(candidate, image_shape=(100, 100), config=cfg)
    score_weak = score_mask(weak, image_shape=(100, 100), config=cfg)

    assert score_candidate > score_weak


def test_select_best_mask_falls_back_to_largest_area_when_all_rejected() -> None:
    cfg = MaskScoringConfig(min_area_ratio=0.2)
    tiny = _rect_mask(height=20, width=20, x=1, y=1, w=2, h=2)
    larger = _rect_mask(height=20, width=20, x=5, y=5, w=4, h=4)

    masks = [
        {"segment": tiny, "area": int(tiny.sum()), "bbox": [1, 1, 2, 2]},
        {"segment": larger, "area": int(larger.sum()), "bbox": [5, 5, 4, 4]},
    ]

    best = select_best_mask(masks=masks, image_shape=(20, 20), config=cfg)

    assert best is not None
    assert best["area"] == int(larger.sum())