"""`infer/objects.py` on synthetic class + height maps whose answer is known.

No torch, no checkpoint: `pip install numpy scipy scikit-image pillow pytest`
and run `pytest -q` from the repo root.
"""

from __future__ import annotations

import json

import numpy as np
import pytest

from config import CLASS_NAMES
from infer.objects import (
    BUILDING_ID, TREE_ID, WATER_ID, ObjectParams, extract_objects, summarise,
)

GSD = 0.5
GROUND_ID = CLASS_NAMES.index("ground")


def blank(h=120, w=160):
    return np.full((h, w), GROUND_ID, np.uint8), np.zeros((h, w), np.float32)


def add_cone(seg, height, cy, cx, radius_px, top_m):
    """A conical crown: `top_m` at the centre, 0.5 * top at the rim."""
    yy, xx = np.mgrid[:seg.shape[0], :seg.shape[1]]
    d = np.hypot(yy - cy, xx - cx)
    inside = d <= radius_px
    seg[inside] = TREE_ID
    cone = top_m * (1.0 - 0.5 * d / radius_px)
    height[inside] = np.maximum(height[inside], cone[inside])


def test_three_separate_trees():
    seg, h = blank()
    truth = [(30, 30, 8, 12.0), (30, 110, 6, 7.0), (90, 70, 10, 18.0)]
    for cy, cx, r, top in truth:
        add_cone(seg, h, cy, cx, r, top)

    obj = extract_objects(seg, h, GSD)
    trees = obj["trees"]
    assert len(trees) == 3
    # sorted tallest first
    assert [t["h"] for t in trees] == sorted((t["h"] for t in trees), reverse=True)
    for (cy, cx, r, top), t in zip(sorted(truth, key=lambda t: -t[3]), trees):
        assert t["h"] == pytest.approx(top, abs=0.6)
        # corner-origin pixels: the centre of pixel (cy, cx) is (cx + 0.5, cy + 0.5)
        assert t["x"] == pytest.approx(cx + 0.5, abs=1.0)
        assert t["y"] == pytest.approx(cy + 0.5, abs=1.0)
        assert t["r"] == pytest.approx(r * GSD, rel=0.2)


def test_touching_crowns_are_split():
    seg, h = blank()
    add_cone(seg, h, 60, 60, 12, 10.0)
    add_cone(seg, h, 60, 80, 12, 11.0)          # overlaps the first
    trees = extract_objects(seg, h, GSD)["trees"]
    assert len(trees) == 2
    xs = sorted(t["x"] for t in trees)
    assert xs[0] < 70.5 < xs[1]


def test_flat_topped_blob_is_still_one_tree():
    seg, h = blank()
    seg[40:60, 40:60] = TREE_ID
    h[40:60, 40:60] = 6.0
    trees = extract_objects(seg, h, GSD)["trees"]
    assert len(trees) == 1
    assert trees[0]["h"] == pytest.approx(6.0, abs=0.01)


def test_shrubs_below_min_height_are_not_trees():
    seg, h = blank()
    add_cone(seg, h, 60, 60, 10, 1.5)
    assert extract_objects(seg, h, GSD)["trees"] == []


def test_building_footprint_and_height():
    seg, h = blank()
    seg[30:50, 10:40] = BUILDING_ID
    h[30:50, 10:40] = 12.0
    h[30, 10:40] = 3.0                          # blurred roof edge, eroded away
    obj = extract_objects(seg, h, GSD)
    (b,) = obj["buildings"]
    assert b["h"] == pytest.approx(12.0)
    assert b["area_m2"] == pytest.approx(20 * 30 * GSD**2)
    poly = np.array(b["poly"])
    assert 4 <= len(poly) <= 8                  # a rectangle, simplified
    assert poly[:, 0].min() == pytest.approx(10, abs=0.6)
    assert poly[:, 0].max() == pytest.approx(40, abs=0.6)
    assert poly[:, 1].min() == pytest.approx(30, abs=0.6)
    assert poly[:, 1].max() == pytest.approx(50, abs=0.6)
    # clockwise on screen (y down) == positive shoelace
    x, y = poly[:, 0], poly[:, 1]
    assert np.sum(x * np.roll(y, -1) - np.roll(x, -1) * y) > 0


def test_flat_building_is_rejected():
    seg, h = blank()
    seg[30:50, 10:40] = BUILDING_ID
    h[30:50, 10:40] = 0.8
    assert extract_objects(seg, h, GSD)["buildings"] == []


def test_water_keeps_lakes_and_drops_puddles():
    seg, h = blank()
    seg[60:100, 60:120] = WATER_ID              # 600 m^2
    seg[10:14, 10:14] = WATER_ID                # 4 m^2
    water = extract_objects(seg, h, GSD)["water"]
    assert len(water) == 1
    assert water[0]["area_m2"] == pytest.approx(40 * 60 * GSD**2)


def test_caps_keep_the_tallest_and_flag_truncation():
    seg, h = blank()
    for cx, top in ((25, 5.0), (80, 9.0), (135, 7.0)):
        add_cone(seg, h, 60, cx, 8, top)
    obj = extract_objects(seg, h, GSD, ObjectParams(max_trees=2))
    assert [round(t["h"]) for t in obj["trees"]] == [9, 7]
    assert obj["truncated"]["trees"] is True
    assert obj["counts"]["trees"] == 2


def test_contract_is_json_and_self_describing():
    seg, h = blank()
    add_cone(seg, h, 60, 60, 8, 9.0)
    seg[10:30, 100:140] = BUILDING_ID
    h[10:30, 100:140] = 7.0
    obj = json.loads(json.dumps(extract_objects(seg, h, GSD)))
    assert obj["version"] == 1
    assert obj["grid"] == {**obj["grid"], "width": 160, "height": 120, "gsd_m": GSD}
    assert obj["counts"] == {"trees": 1, "buildings": 1, "water": 0}
    assert "1 trees" in summarise(obj)


def test_empty_scene_and_bad_input():
    seg, h = blank()
    obj = extract_objects(seg, h, GSD)
    assert obj["counts"] == {"trees": 0, "buildings": 0, "water": 0}
    with pytest.raises(ValueError):
        extract_objects(seg, h[:10], GSD)
    with pytest.raises(ValueError):
        extract_objects(seg, h, 0)


def test_previews_render():
    from viz.classes import class_preview, objects_overlay

    seg, h = blank()
    add_cone(seg, h, 60, 60, 8, 9.0)
    rgb = np.full(seg.shape + (3,), 128, np.uint8)
    assert class_preview(rgb, seg, CLASS_NAMES).size == (160, 120)
    assert objects_overlay(rgb, extract_objects(seg, h, GSD)).size == (160, 120)