core/qc_agent.py

"""
QC Agent — stateful cross-frame quality control engine
"""

from __future__ import annotations
from collections import defaultdict
from dataclasses import dataclass, field

from .sop_rules import run_all_rules, RuleViolation, Severity


@dataclass
class FrameAnnotation:
    frame_idx: int
    objects: list[dict]   # [{"track_id","label","bbox":[x,y,w,h],"confidence"}]


@dataclass
class QCIssue:
    check:    str
    severity: str
    message:  str
    frame:    int
    track_id: str | None = None

    def user_message(self) -> str:
        icon = "🔴" if self.severity == "error" else "🟡"
        tid  = f" · Object {self.track_id}" if self.track_id else ""
        return f"{icon} Frame {self.frame}{tid} — {self.message}"


@dataclass
class QCReport:
    total_frames:    int
    total_issues:    int
    issues_by_type:  dict[str, int]
    frames_affected: list[int]
    issues:          list[QCIssue]
    summary:         str = ""

    def to_display(self) -> str:
        lines = [
            f"QC Report  ({self.total_frames} frames)",
            f"Total issues : {self.total_issues}",
            f"Frames affected: {len(self.frames_affected)}",
            "", "Issues by type:",
        ]
        for k, v in sorted(self.issues_by_type.items(), key=lambda x: -x[1]):
            lines.append(f"  {k}: {v}")
        lines.append("\nDetail:")
        for i in self.issues:
            lines.append("  " + i.user_message())
        return "\n".join(lines)


class QCAgent:
    def __init__(self, ctx: dict | None = None):
        self.ctx: dict = ctx or {}
        self.ctx.setdefault("track_history", {})
        self.ctx.setdefault("track_frame_count", {})
        self.ctx.setdefault("track_last_seen", {})
        self.ctx.setdefault("detection_index", {})
        self._issues: list[QCIssue] = []
        self._frames: list[FrameAnnotation] = []
        self._FLICKER_GAP    = self.ctx.get("flicker_gap", 5)
        self._MISSING_THRESH = 0.30

    # ── public ────────────────────────────────

    def process_frame(self, fa: FrameAnnotation):
        self._frames.append(fa)
        ann = {"objects": fa.objects}

        # SOP rule engine
        for v in run_all_rules(ann, fa.frame_idx, self.ctx):
            self._issues.append(QCIssue(
                check=v.rule_id, severity=v.severity.value,
                message=v.message, frame=v.frame, track_id=v.track_id,
            ))

        # cross-frame checks
        self._check_flicker(fa)
        self._check_missed(fa)
        self._update_state(fa)

    def generate_report(self) -> QCReport:
        by_type: dict[str, int] = defaultdict(int)
        affected: set[int] = set()
        for i in self._issues:
            by_type[i.check] += 1
            affected.add(i.frame)

        report = QCReport(
            total_frames=len(self._frames),
            total_issues=len(self._issues),
            issues_by_type=dict(by_type),
            frames_affected=sorted(affected),
            issues=self._issues,
        )
        if self._issues:
            top = max(by_type, key=by_type.get)
            report.summary = (
                f"⚠️ {report.total_issues} issue(s) across "
                f"{len(affected)} frame(s). Most common: {top} ({by_type[top]})."
            )
        else:
            report.summary = "✅ All annotations passed QC checks."
        return report

    def get_frame_warnings(self, frame_idx: int) -> list[str]:
        return [i.user_message() for i in self._issues if i.frame == frame_idx]

    def load_ai_detections(self, index: dict[int, list[dict]]):
        self.ctx["detection_index"] = index

    # ── cross-frame checks ────────────────────

    def _check_flicker(self, fa: FrameAnnotation):
        for obj in fa.objects:
            tid = obj.get("track_id")
            if not tid:
                continue
            last = self.ctx["track_last_seen"].get(tid)
            if last is not None:
                gap = fa.frame_idx - last
                if 1 < gap <= self._FLICKER_GAP:
                    self._issues.append(QCIssue(
                        check="FLICKER", severity="warning",
                        message=f"Object disappeared for {gap-1} frame(s) then reappeared — possible flickering tag",
                        frame=fa.frame_idx, track_id=tid,
                    ))

    def _check_missed(self, fa: FrameAnnotation):
        dets = self.ctx["detection_index"].get(fa.frame_idx, [])
        if not dets:
            return
        human_boxes = [o["bbox"] for o in fa.objects if o.get("bbox")]
        for det in dets:
            db = det.get("bbox", [])
            if not any(_iou(db, hb) > self._MISSING_THRESH for hb in human_boxes):
                self._issues.append(QCIssue(
                    check="MISSED", severity="error",
                    message=(
                        f"AI detected '{det.get('label','object')}' "
                        f"(conf {det.get('confidence',0):.0%}) but no annotation found"
                    ),
                    frame=fa.frame_idx,
                ))

    def _update_state(self, fa: FrameAnnotation):
        for obj in fa.objects:
            tid = obj.get("track_id")
            if not tid:
                continue
            self.ctx["track_frame_count"][tid] = self.ctx["track_frame_count"].get(tid, 0) + 1
            self.ctx["track_last_seen"][tid] = fa.frame_idx


def _iou(a: list, b: list) -> float:
    if len(a) < 4 or len(b) < 4:
        return 0.0
    x1,y1,w1,h1 = a[:4]; x2,y2,w2,h2 = b[:4]
    ix = max(0, min(x1+w1,x2+w2)-max(x1,x2))
    iy = max(0, min(y1+h1,y2+h2)-max(y1,y2))
    inter = ix*iy
    union = w1*h1+w2*h2-inter
    return inter/union if union > 0 else 0.0