"""Counting a frame at several resolutions, and noticing when the count runs away.

A detector reads a fixed 640 px input. Hand it a 4,000 px photograph of a paddock
and every animal is downsampled before the network sees it; hand it one ninth of
that photograph and each animal arrives four times larger. So the same detector
returns a different count depending on how the frame is cut up, and **the way it
changes as you cut finer is the measurement that matters**.

Two things fall out of that, and the second is the important one.

**The count gets better.** Slicing a frame into overlapping tiles and detecting
in each one recovers animals that whole-frame inference loses to downsampling.
This is the standard trick for small objects in large images.

**The count says whether it can be trusted.** In a frame the detector can
actually read, the count stops moving: a paddock with a dozen cattle returns
about a dozen at one tile, at four, and at nine, because there was nothing left
to find. In a broiler house it never stops moving — every finer cut finds more
birds, because there are always more birds hidden behind the ones in front.

That is the difference between a count and a sample, measured rather than
guessed. The old guard asked whether the boxes it *had* were small, which is a
question about the animals the detector found and says nothing about the ones it
missed — on a shed of a thousand birds it saw twenty large foreground birds,
concluded the frame was sparse, and published twenty.
"""

from __future__ import annotations

from dataclasses import dataclass

from PIL import Image

from app.detectors.base import Detection, Detector

#: How much neighbouring tiles overlap, as a share of tile size. An animal
#: sitting exactly on a cut would otherwise be two half-animals, each too
#: partial to detect. Cross-tile NMS then removes the duplicates the overlap
#: creates.
TILE_OVERLAP = 0.20

#: IoU above which two boxes are the same animal. Looser than the within-tile
#: NMS threshold, because the same animal seen in two tiles is cropped
#: differently in each and the boxes never align exactly.
MERGE_IOU = 0.55

#: Intersection over the *smaller* box's area, above which the smaller box is a
#: part of the larger one rather than a second animal.
#:
#: This is the threshold that makes tiling safe, and leaving it out is how the
#: first attempt turned one cow into three. A cow filling the frame is cut into
#: quarters by a 2x2 grid, and the detector obligingly finds a cow in each
#: quarter; those four quarter-boxes barely overlap *each other*, so IoU keeps
#: all four. Each is almost entirely inside the whole-frame box, so containment
#: removes them.
#:
#: 0.85 rather than something lower because two animals standing one behind the
#: other genuinely overlap: on the evaluation set the near animal's box covered
#: up to three quarters of the far animal's. Merging those would trade a
#: duplicate for a lost animal.
MERGE_CONTAINMENT = 0.85

#: Tile grids, coarse to fine. 1 is whole-frame. Stopping at 3 is a cost
#: decision: 1 + 4 + 9 inferences already takes seconds on a CPU, and a frame
#: still finding new animals at 3x3 is a shed — the answer there is that no
#: count exists, not that a fourth grid would find it.
LEVELS: tuple[int, ...] = (1, 2, 3)


@dataclass(frozen=True)
class Level:
    """Everything found at this grid **and every coarser one**.

    Accumulating rather than replacing is what makes the comparison between
    levels mean something: each level is a superset, so a level that adds
    nothing new is a level that found nothing new, and the count can only rise.
    It is also what keeps the whole-frame box of a large animal in the pool to
    absorb the fragments a fine grid makes of it.
    """

    grid: int
    detections: list[Detection]

    @property
    def count(self) -> int:
        return len(self.detections)


def _crops(size: tuple[int, int], grid: int) -> list[tuple[int, int, int, int]]:
    width, height = size
    if grid == 1:
        return [(0, 0, width, height)]

    step_x, step_y = width / grid, height / grid
    pad_x, pad_y = step_x * TILE_OVERLAP, step_y * TILE_OVERLAP
    boxes = []
    for row in range(grid):
        for column in range(grid):
            x0 = max(0, int(column * step_x - pad_x))
            y0 = max(0, int(row * step_y - pad_y))
            x1 = min(width, int((column + 1) * step_x + pad_x))
            y1 = min(height, int((row + 1) * step_y + pad_y))
            boxes.append((x0, y0, x1, y1))
    return boxes


def _overlaps(
    a: tuple[float, float, float, float], b: tuple[float, float, float, float]
) -> tuple[float, float]:
    """`(IoU, intersection over the smaller area)` for two boxes."""
    ax0, ay0, ax1, ay1 = a
    bx0, by0, bx1, by1 = b
    x0, y0 = max(ax0, bx0), max(ay0, by0)
    x1, y1 = min(ax1, bx1), min(ay1, by1)
    overlap = max(0.0, x1 - x0) * max(0.0, y1 - y0)
    if overlap <= 0.0:
        return 0.0, 0.0
    area_a = (ax1 - ax0) * (ay1 - ay0)
    area_b = (bx1 - bx0) * (by1 - by0)
    union = area_a + area_b - overlap
    smaller = min(area_a, area_b)
    return (
        overlap / union if union > 0 else 0.0,
        overlap / smaller if smaller > 0 else 0.0,
    )


def _merge(detections: list[Detection], frame_area: float) -> list[Detection]:
    """One animal, one box, whichever tile found it.

    **Largest box first**, which is the ordering the containment rule needs: the
    whole animal has to be in the kept set before its fragments are tested
    against it. Score order — the usual choice for NMS — would let a confident
    fragment claim the animal and leave its siblings unmatched.

    Boxes come back in frame coordinates, and `area_fraction` is recomputed
    against the whole frame: a bird covering a quarter of its tile covers a
    thirty-sixth of the picture, and everything downstream reasons about the
    picture.
    """
    def area(d: Detection) -> float:
        x0, y0, x1, y1 = d.box
        return (x1 - x0) * (y1 - y0)

    kept: list[Detection] = []
    for detection in sorted(detections, key=area, reverse=True):
        duplicate = False
        for other in kept:
            if other.label != detection.label:
                continue
            iou, containment = _overlaps(other.box, detection.box)
            if iou > MERGE_IOU or containment > MERGE_CONTAINMENT:
                duplicate = True
                break
        if duplicate:
            continue
        kept.append(
            Detection(
                label=detection.label,
                score=detection.score,
                box=detection.box,
                area_fraction=area(detection) / frame_area,
            )
        )
    kept.sort(key=lambda d: d.score, reverse=True)
    return kept


def _raw(detector: Detector, image: Image.Image, grid: int) -> list[Detection]:
    """Every box one grid produced, in frame coordinates, unmerged."""
    width, height = image.size
    gathered: list[Detection] = []
    for x0, y0, x1, y1 in _crops((width, height), grid):
        tile = image if grid == 1 else image.crop((x0, y0, x1, y1))
        for detection in detector.detect(tile):
            tx0, ty0, tx1, ty1 = detection.box
            gathered.append(
                Detection(
                    label=detection.label,
                    score=detection.score,
                    box=(tx0 + x0, ty0 + y0, tx1 + x0, ty1 + y0),
                    # Recomputed by `_merge`; a tile-relative fraction here
                    # would be wrong by the square of the grid.
                    area_fraction=detection.area_fraction,
                )
            )
    return gathered


def detect_at(detector: Detector, image: Image.Image, grid: int) -> list[Detection]:
    """Run one grid on its own. Used by the tests and by nothing else."""
    width, height = image.size
    return _merge(_raw(detector, image, grid), float(width * height))


def pyramid(
    detector: Detector,
    image: Image.Image,
    subject_classes: tuple[str, ...],
    growth_tolerance: float,
    levels: tuple[int, ...] = LEVELS,
) -> list[Level]:
    """Count at successively finer grids, stopping as soon as the count settles.

    Returns every level that was run, coarsest first, each holding the merged
    result of every grid up to and including its own. The caller decides what
    the sequence means; this function only refuses to spend inferences it does
    not need — a frame that has settled is not going to unsettle, and the common
    case is a farmer photographing six animals.
    """
    frame_area = float(image.size[0] * image.size[1])
    gathered: list[Detection] = []
    results: list[Level] = []

    for grid in levels:
        gathered.extend(_raw(detector, image, grid))
        results.append(Level(grid=grid, detections=_merge(gathered, frame_area)))
        if len(results) >= 2 and converged(
            results[-2], results[-1], subject_classes, growth_tolerance
        ):
            break
    return results


def subject_count(level: Level, subject_classes: tuple[str, ...]) -> int:
    return sum(1 for d in level.detections if d.label in subject_classes)


def converged(
    coarser: Level, finer: Level, subject_classes: tuple[str, ...], tolerance: float
) -> bool:
    """Whether cutting the frame finer stopped finding new animals.

    Growth is measured against the coarser count, so it is a proportion rather
    than a difference: three more animals out of six means the frame was not
    read, three more out of sixty means it was.
    """
    before = subject_count(coarser, subject_classes)
    after = subject_count(finer, subject_classes)
    if before == 0:
        # Nothing at the coarse grid. Converged only if the finer grid agrees,
        # otherwise the coarse pass simply could not see the animals.
        return after == 0
    return (after - before) / before <= tolerance