"""
The shelf as a plane, not a stream of boxes.

Every measurement so far treats a frame as the world: a facing is "a track id",
and a track id dies the moment the camera pans past and comes back. Label
inheritance, drift checks and restart counters were each correct fixes for that,
and all of them are symptoms. The shelf does not move; the camera does.

This module keeps a per-session **shelf map**. Each frame is registered to the map
by its content shift, every detection is placed in shelf coordinates, and a
detection that lands where a facing already is *merges* into it. Identity votes and
tag readings accumulate on the physical facing rather than on a transient track.

That is what the current literature does for retail: RetailGlue (CVPR 2026W) builds
shelf panoramas by matching products across frames and reports 88.4 % stitching F1,
and the 7-Eleven Taiwan system (Sci. Rep. 2025, 7,000+ stores) stitches "virtual
shelves" before checking planograms.

**What is measured, and what is not.** The synthetic walk registers to within
0.01 px of its ground truth - but it is a pure translation of a single photograph,
so that validates the harness, not handheld video. The number that matters (real
registration error, and the duplicate rate that follows from it) needs real
footage, which `record.py` exists to collect.
"""

from __future__ import annotations

import threading
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence, Tuple

# A frame whose content shift cannot be estimated this confidently is treated as
# discontinuous: the map starts a new section rather than guessing a position.
MIN_REGISTRATION_RESPONSE = 0.20
# Two detections are the same physical facing when their shelf boxes overlap this
# much. Measured: adjacent facings of one product overlap ~0 by construction.
MERGE_IOU = 0.45


def _iou(a: Sequence[int], b: Sequence[int]) -> float:
    x1, y1 = max(a[0], b[0]), max(a[1], b[1])
    x2, y2 = min(a[2], b[2]), min(a[3], b[3])
    intersection = max(0, x2 - x1) * max(0, y2 - y1)
    area_a = max(0, a[2] - a[0]) * max(0, a[3] - a[1])
    area_b = max(0, b[2] - b[0]) * max(0, b[3] - b[1])
    return intersection / max(1.0, float(area_a + area_b - intersection))


@dataclass
class Facing:
    """One physical position on the shelf, with everything observed there."""

    facing_id: str
    shelf_box: List[int]
    row: int = 0
    observations: int = 0
    first_seq: int = 0
    last_seq: int = 0
    identity_votes: Dict[str, float] = field(default_factory=dict)
    tag_votes: Dict[str, float] = field(default_factory=dict)
    tag_readings: List[Dict[str, Any]] = field(default_factory=list)
    seen_track_ids: set = field(default_factory=set)

    def vote_identity(self, sku_id: Optional[str], weight: float = 1.0) -> None:
        if sku_id:
            self.identity_votes[sku_id] = self.identity_votes.get(sku_id, 0.0) + weight

    def vote_tag(self, sku_id: Optional[str], reading: Optional[Dict[str, Any]] = None) -> None:
        if sku_id:
            self.tag_votes[sku_id] = self.tag_votes.get(sku_id, 0.0) + 1.0
        if reading:
            self.tag_readings.append(reading)

    def best_identity(self) -> Tuple[Optional[str], float]:
        """The identity the facing's pooled evidence supports, and its agreement."""
        if not self.identity_votes:
            return None, 0.0
        sku_id, weight = max(self.identity_votes.items(), key=lambda item: item[1])
        total = sum(self.identity_votes.values())
        return sku_id, weight / max(1e-6, total)

    def best_tag_identity(self) -> Tuple[Optional[str], float]:
        if not self.tag_votes:
            return None, 0.0
        sku_id, weight = max(self.tag_votes.items(), key=lambda item: item[1])
        total = sum(self.tag_votes.values())
        return sku_id, weight / max(1e-6, total)

    def as_payload(self) -> Dict[str, Any]:
        identity, agreement = self.best_identity()
        tag_identity, tag_agreement = self.best_tag_identity()
        return {
            "facing_id": self.facing_id,
            "box": list(self.shelf_box),
            "row": self.row,
            "observations": self.observations,
            "track_ids": len(self.seen_track_ids),
            "sku_id": identity,
            "identity_agreement": round(agreement, 3),
            "tag_sku_id": tag_identity,
            "tag_agreement": round(tag_agreement, 3),
            "tag_readings": self.tag_readings[-3:],
        }


class ShelfMap:
    """Registers frames and accumulates facings in shelf coordinates."""

    def __init__(self, settings: Optional[Any] = None) -> None:
        self._lock = threading.Lock()
        self._prev_gray: Optional[Any] = None
        self._offset: List[float] = [0.0, 0.0]      # shelf position of the frame origin
        self._facings: List[Facing] = []
        self._orb = None
        self.registrations = 0
        self.registration_failures = 0
        self.anchors = 0
        self.continuations = 0
        self.detections = 0
        self.merges = 0
        self.responses: List[float] = []

    # ------------------------------------------------------------- registration
    def _gray(self, frame: Any, scale: float = 0.5) -> Any:
        import cv2

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        if scale != 1.0:
            gray = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
        return gray

    def register(self, frame: Any) -> Optional[Tuple[float, float]]:
        """Content shift from the previous frame, in full-resolution pixels.

        Phase correlation first: it is the right tool for a translating scene and
        costs a couple of milliseconds. ORB + RANSAC only when the correlation is
        weak (blur, a large jump, a scene with no texture), because that is the case
        where a fast estimator quietly returns nonsense.
        """
        import cv2
        import numpy as np

        gray = self._gray(frame)
        with self._lock:
            previous = self._prev_gray
            self._prev_gray = gray
        if previous is None or previous.shape != gray.shape:
            return (0.0, 0.0)
        window = cv2.createHanningWindow((gray.shape[1], gray.shape[0]), cv2.CV_32F)
        (dx, dy), response = cv2.phaseCorrelate(previous.astype(np.float32), gray.astype(np.float32), window)
        self.responses.append(float(response))
        if response >= MIN_REGISTRATION_RESPONSE:
            self.registrations += 1
            return (dx / 0.5, dy / 0.5)          # back to full resolution

        shift = self._orb_shift(previous, gray)
        if shift is None:
            self.registration_failures += 1
            return None
        self.registrations += 1
        return (shift[0] / 0.5, shift[1] / 0.5)

    def _orb_shift(self, previous: Any, current: Any) -> Optional[Tuple[float, float]]:
        import cv2
        import numpy as np

        if self._orb is None:
            self._orb = cv2.ORB_create(nfeatures=1500, fastThreshold=10)
        key_a, desc_a = self._orb.detectAndCompute(previous, None)
        key_b, desc_b = self._orb.detectAndCompute(current, None)
        if desc_a is None or desc_b is None or len(key_a) < 20 or len(key_b) < 20:
            return None
        matches = cv2.BFMatcher(cv2.NORM_HAMMING).knnMatch(desc_a, desc_b, k=2)
        good = [m for m, n in matches if m.distance < 0.75 * n.distance]
        if len(good) < 12:
            return None
        src = np.float32([key_a[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
        dst = np.float32([key_b[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
        matrix, _ = cv2.estimateAffinePartial2D(src, dst, method=cv2.RANSAC, ransacReprojThreshold=3.0)
        if matrix is None:
            return None
        return (float(matrix[0, 2]), float(matrix[1, 2]))

    # ------------------------------------------------------------------ facings
    def _anchor_to_map(
        self, detections: Sequence[Sequence[int]], offset_x: float, offset_y: float
    ) -> Tuple[float, float]:
        """Correct the chained offset by matching detections against the map.

        Chaining pairwise shifts drifts: measured over a 79-step walk with *perfect*
        pairwise registration, the accumulated x offset reached 2047 px on a 1500 px
        shelf, the same facing landed in two places, and a shelf of ~60 facings
        became 334 with 113 of them seen once. Drift is inherent to chaining - the
        fix is to anchor the frame to the structure already on the map, which is
        what the retail-stitching literature does with product correspondences.

        One ICP-style step is enough: for every detection that has a facing nearby,
        the correction it implies is a vote, and the median vote wins.
        """
        if not self._facings or len(detections) < 4:
            return offset_x, offset_y
        corrections = []
        search = 40.0
        for box in detections:
            centre_x = (box[0] + box[2]) / 2.0 + offset_x
            centre_y = (box[1] + box[3]) / 2.0 + offset_y
            best, best_distance = None, search
            for facing in self._facings:
                facing_x = (facing.shelf_box[0] + facing.shelf_box[2]) / 2.0
                facing_y = (facing.shelf_box[1] + facing.shelf_box[3]) / 2.0
                distance = max(abs(facing_x - centre_x), abs(facing_y - centre_y))
                if distance < best_distance:
                    best, best_distance = facing, distance
            if best is not None:
                corrections.append(
                    (
                        (best.shelf_box[0] + best.shelf_box[2]) / 2.0 - ((box[0] + box[2]) / 2.0 + offset_x),
                        (best.shelf_box[1] + best.shelf_box[3]) / 2.0 - ((box[1] + box[3]) / 2.0 + offset_y),
                    )
                )
        if len(corrections) < 4:
            return offset_x, offset_y
        dx = sorted(c[0] for c in corrections)[len(corrections) // 2]
        dy = sorted(c[1] for c in corrections)[len(corrections) // 2]
        self.anchors += 1
        return offset_x + dx, offset_y + dy

    def update(self, seq: int, frame: Any, tracks: Sequence[Dict[str, Any]]) -> None:
        """Register the frame, then place or merge every detection on the shelf."""
        shift = self.register(frame)
        boxes = [
            [int(v) for v in (track.get("box") or track.get("raw_box"))]
            for track in tracks
            if (track.get("box") or track.get("raw_box")) and len(track.get("box") or track.get("raw_box")) == 4
        ]
        with self._lock:
            if shift is None:
                # A discontinuous frame: keep the last known offset rather than
                # inventing one, so its detections land where the shelf is.
                shift = (0.0, 0.0)
            # `shift` is how far the CONTENT moved in the frame, so the frame's
            # origin moves the opposite way along the shelf. Adding it instead
            # (measured) walks the map backwards: the same facing lands 50 px away
            # in the next frame and a shelf of ~60 facings becomes 857.
            self._offset[0] -= shift[0]
            self._offset[1] -= shift[1]
            if self._facings:
                self._offset[0], self._offset[1] = self._anchor_to_map(
                    boxes, self._offset[0], self._offset[1]
                )
            offset_x, offset_y = self._offset
            for track in tracks:
                box = track.get("box") or track.get("raw_box")
                if not box or len(box) != 4:
                    continue
                self.detections += 1
                shelf_box = [
                    int(round(box[0] + offset_x)), int(round(box[1] + offset_y)),
                    int(round(box[2] + offset_x)), int(round(box[3] + offset_y)),
                ]
                track_id = track.get("track_id")
                facing = self._find(shelf_box, int(track_id) if track_id is not None else None)
                if facing is None:
                    facing = Facing(
                        facing_id=f"face_{uuid.uuid4().hex[:8]}",
                        shelf_box=shelf_box,
                        first_seq=seq,
                    )
                    self._facings.append(facing)
                else:
                    self.merges += 1
                    # Average the box: repeated observations of a static object
                    # should converge, and the average is what the eye expects.
                    for index in range(4):
                        facing.shelf_box[index] = int(round((facing.shelf_box[index] * facing.observations + shelf_box[index]) / (facing.observations + 1)))
                facing.observations += 1
                facing.last_seq = seq
                identity = track.get("identity") or {}
                if identity.get("state") == "recognized":
                    facing.vote_identity(identity.get("sku_id"), float(identity.get("score") or 1.0))
                tag_identity = track.get("tag_identity") or {}
                tag = track.get("price_tag") or {}
                if tag_identity.get("sku_id") or tag.get("state") == "readable":
                    facing.vote_tag(tag_identity.get("sku_id"), tag if tag else None)
                if track.get("track_id") is not None:
                    facing.seen_track_ids.add(int(track["track_id"]))

    def _find(self, shelf_box: Sequence[int], track_id: Optional[int] = None) -> Optional[Facing]:
        """The facing a detection belongs to.

        Continuity first: if this track id was already placed nearby, it is the same
        facing, and no overlap test is needed - the tracker already decided that.
        Geometry only when the id is new (a pan that outran the tracker, or the
        tracker renumbering). Measured: with geometry alone, box jitter split the
        same facing often enough that the map held 210 facings where the tracker
        held 156 ids - it was adding duplicates, not removing them.
        """
        width = max(1, shelf_box[2] - shelf_box[0])
        height = max(1, shelf_box[3] - shelf_box[1])
        if track_id is not None:
            radius = 0.75 * max(width, height)
            best, best_distance = None, radius
            for facing in self._facings:
                if track_id not in facing.seen_track_ids:
                    continue
                distance = max(
                    abs((facing.shelf_box[0] + facing.shelf_box[2]) / 2.0 - (shelf_box[0] + shelf_box[2]) / 2.0),
                    abs((facing.shelf_box[1] + facing.shelf_box[3]) / 2.0 - (shelf_box[1] + shelf_box[3]) / 2.0),
                )
                if distance < best_distance:
                    best, best_distance = facing, distance
            if best is not None:
                self.continuations += 1
                return best
        best, best_overlap = None, MERGE_IOU
        for facing in self._facings:
            overlap = _iou(shelf_box, facing.shelf_box)
            if overlap > best_overlap:
                best, best_overlap = facing, overlap
        return best

    # ------------------------------------------------------------------ reporting
    def facings(self) -> List[Facing]:
        with self._lock:
            return list(self._facings)

    def stats(self) -> Dict[str, Any]:
        with self._lock:
            facings = list(self._facings)
            responses = list(self.responses[-200:])
        consistent = 0
        for facing in facings:
            sku_id, agreement = facing.best_identity()
            if sku_id and agreement >= 0.8:
                consistent += 1
        return {
            "facings": len(facings),
            "detections": self.detections,
            "merges": self.merges,
            "continuations": self.continuations,
            "registrations": self.registrations,
            "anchored": self.anchors,
            "registration_failures": self.registration_failures,
            "median_response": round(sorted(responses)[len(responses) // 2], 3) if responses else None,
            "identity_consistent": consistent,
            "offset": [round(value, 1) for value in self._offset],
        }