core/sop_rules.py

"""
SOP Rule Engine — Tapi Tag Ontology Phase II
QC checks and error messages derived directly from the official SOP (page 12).
"""

from __future__ import annotations
from dataclasses import dataclass, field
from enum import Enum

from .tapi_ontology import (
    WEATHER_LABELS, TIME_OF_DAY_LABELS, VISIBILITY_LABELS,
    SETTING_AREA_LABELS, SETTING_INFRASTRUCTURE_LABELS, SETTING_NATURE_LABELS,
    ROAD_PROPERTY_LABELS, EVENT_LABELS,
    EGO_MOVEMENT_LABELS, VEHICLE_MOVEMENT_LABELS,
    PEDESTRIAN_MOVEMENT_LABELS, ANIMAL_MOVEMENT_LABELS,
    ACTOR_VEHICLE_LABELS, ACTOR_PROPERTY_LABELS,
    ACTOR_ANIMAL_LABELS, ACTOR_PEDESTRIAN_LABELS,
    TRAFFIC_SIGNAL_LABELS, TRAFFIC_SIGN_LABELS,
    ALL_VALID_LABELS, URBAN_RESIDENTIAL, EVENT_MAX_SECONDS,
)


class Severity(str, Enum):
    ERROR   = "error"
    WARNING = "warning"


@dataclass
class RuleViolation:
    rule_id:  str
    severity: Severity
    message:  str
    frame:    int
    track_id: str | None = None
    details:  dict       = field(default_factory=dict)


# ── Rule registry ────────────────────────────
RULES: dict[str, dict] = {}

def rule(rule_id: str, description: str, severity: Severity = Severity.ERROR):
    def decorator(fn):
        RULES[rule_id] = {"fn": fn, "description": description, "severity": severity}
        return fn
    return decorator


# ════════════════════════════════════════════════
#  ENVIRONMENT — SOP page 12
# ════════════════════════════════════════════════

@rule("ENV-01", "Missing Environment label for any frames")
def check_env_missing(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    env = WEATHER_LABELS | TIME_OF_DAY_LABELS | VISIBILITY_LABELS
    if not any(l in env for l in labels):
        return [_v("ENV-01", f"Missing Environment label: frame {fi}", fi)]
    return []


@rule("ENV-02", "Contradictory annotation within Environment labels based on Time of Day")
def check_tod_change(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    tod = [l for l in labels if l in TIME_OF_DAY_LABELS]
    if len(tod) > 1:
        return [_v("ENV-02",
            f"Found multiple values for Time of Day: Environment Label Time of Day cannot change within a task "
            f"(frame {fi}): {tod}", fi)]
    saved = ctx.get("task_tod")
    if tod:
        if saved is None:
            ctx["task_tod"] = tod[0]
        elif tod[0] != saved:
            return [_v("ENV-02",
                f"Environment label Time of Day cannot change within a task "
                f"(was '{saved}', now '{tod[0]}' at frame {fi})", fi)]
    return []


@rule("ENV-03", "Contradictory annotation within Environment labels based on Weather")
def check_weather_conflict(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    weather = [l for l in labels if l in WEATHER_LABELS]
    vis     = [l for l in labels if l in VISIBILITY_LABELS]
    if len(weather) > 1:
        return [_v("ENV-03",
            f"Environment label cannot have multiple Weather values (frame {fi}): {weather}", fi)]
    if weather and vis:
        w = weather[0]
        v = vis[0]
        if w == "clear" and v == "poor visibility from weather":
            return [_v("ENV-03",
                f"Environment label '{v}' cannot be set if Weather is '{w}' (frame {fi})", fi)]
        if w == "foggy" and v == "clear":
            return [_v("ENV-03",
                f"Environment label Visibility 'clear' cannot be set if Weather is 'foggy' (frame {fi})", fi)]
    return []


@rule("ENV-04", "Environment labels overlapping for any frames")
def check_env_overlap(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    vis = [l for l in labels if l in VISIBILITY_LABELS]
    if len(vis) > 1:
        return [_v("ENV-04",
            f"Environment labels overlapping: frame {fi} — multiple Visibility values: {vis}", fi)]
    return []


# ════════════════════════════════════════════════
#  SETTING — SOP page 12
# ════════════════════════════════════════════════

@rule("SET-01", "Labels for Setting is missing for any frames")
def check_setting_missing(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    all_set = SETTING_AREA_LABELS | SETTING_INFRASTRUCTURE_LABELS | SETTING_NATURE_LABELS
    if not any(l in all_set for l in labels):
        return [_v("SET-01", f"Missing Setting label: frame {fi}", fi)]
    return []


@rule("SET-02", "Setting labels are overlapping for any frames")
def check_setting_overlap(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    area = [l for l in labels if l in SETTING_AREA_LABELS]
    if len(area) > 1:
        return [_v("SET-02",
            f"Setting labels overlapping: frame {fi} — multiple Area values: {area}", fi)]
    return []


@rule("SET-03", "Setting Area Urban/Residential must also have 'buildings' infrastructure")
def check_urban_buildings(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    area = next((l for l in labels if l in SETTING_AREA_LABELS), None)
    if area in URBAN_RESIDENTIAL and "buildings" not in labels:
        return [_v("SET-03",
            f"Setting error: {area.title()}/Residential area must have 'buildings' selected (frame {fi})", fi)]
    return []


# ════════════════════════════════════════════════
#  ROAD TOPOLOGY — SOP page 12
# ════════════════════════════════════════════════

@rule("ROAD-01", "Labels for Road Topology is missing for any frames")
def check_road_missing(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    has_road = any(l in ROAD_PROPERTY_LABELS or l == "total lanes" for l in labels)
    if not has_road:
        return [_v("ROAD-01", f"Missing Road Topology label: frame {fi}", fi)]
    return []


@rule("ROAD-02", "Road Topology labels are overlapping for any frames")
def check_road_overlap(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    props = [l for l in labels if l in ROAD_PROPERTY_LABELS]
    # median-separated is structural, should not mix with wet/snow etc.
    if "median-separated" in props and len(props) > 3:
        return [_v("ROAD-02",
            f"Road Topology labels overlapping: frame {fi} — check combinations {props}",
            fi, severity=Severity.WARNING)]
    return []


# ════════════════════════════════════════════════
#  EVENT — SOP page 12
# ════════════════════════════════════════════════

@rule("EVT-01", "Length of an Event label is greater than 5 seconds")
def check_event_duration(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    fps = ctx.get("fps", 25)
    max_frames = int(EVENT_MAX_SECONDS * fps)
    violations = []
    event_start: dict = ctx.setdefault("event_start", {})
    for obj in ann.get("objects", []):
        label = _n(obj.get("label", ""))
        tid   = obj.get("track_id", label)
        if label in EVENT_LABELS:
            if tid not in event_start:
                event_start[tid] = fi
            elif (fi - event_start[tid]) > max_frames:
                violations.append(_v("EVT-01",
                    f"Length of the event tag (frame {event_start[tid]} - {fi}) is longer than 5 seconds. "
                    f"Please make sure timestamps are accurate.", fi, tid))
    return violations


@rule("EVT-02", "Missing traffic signal static object during event frames")
def check_event_signal(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    has_event  = any(l in EVENT_LABELS for l in labels)
    has_signal = any(l in TRAFFIC_SIGNAL_LABELS for l in labels)
    if has_event and not has_signal:
        return [_v("EVT-02",
            f"Missing traffic signal when event happens (frame {fi}) — "
            f"static object 'traffic signal' must be present during event frames", fi)]
    return []


# ════════════════════════════════════════════════
#  MOVEMENT — EGO — SOP page 12
# ════════════════════════════════════════════════

@rule("EGO-01", "Labels for Movements with the Ego-Car is missing for more than 10 consecutive frames")
def check_ego_missing(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    has_ego = any(l in EGO_MOVEMENT_LABELS for l in labels)
    streak: int = ctx.get("ego_miss_streak", 0)
    if not has_ego:
        ctx["ego_miss_streak"] = streak + 1
        if ctx["ego_miss_streak"] >= 10:
            return [_v("EGO-01",
                f"Missing labels for ego car: frame {fi - ctx['ego_miss_streak'] + 1} - {fi}", fi)]
    else:
        ctx["ego_miss_streak"] = 0
    return []


@rule("EGO-02", "Ego parked overlapping with other Ego-Movement labels")
def check_ego_parked_overlap(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    ego = [l for l in labels if l in EGO_MOVEMENT_LABELS]
    parked    = [l for l in ego if "ego parked" in l or l == "ego parking"]
    non_parked = [l for l in ego if "ego parked" not in l and l != "ego parking"]
    if parked and non_parked:
        return [_v("EGO-02",
            f"'{parked[0]}' overlapping with other Ego-Movement label '{non_parked[0]}' (frame {fi})", fi)]
    return []


@rule("EGO-03", "Ego waits overlapping with other Ego-Movement labels")
def check_ego_waits_overlap(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    ego = [l for l in labels if l in EGO_MOVEMENT_LABELS]
    waits     = [l for l in ego if l.startswith("ego waits")]
    non_waits = [l for l in ego if not l.startswith("ego waits")]
    if waits and non_waits:
        return [_v("EGO-03",
            f"'Ego waits...' overlapping with other Ego-Movement label (frame {fi}): "
            f"'{waits[0]}' + '{non_waits[0]}'", fi)]
    return []


@rule("EGO-04", "Ego merges conflicting labels (e.g. merges into + merges out of same place)")
def check_ego_merges_conflict(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    merges = [l for l in labels if l in EGO_MOVEMENT_LABELS and "ego merges" in l]
    into_l  = [l for l in merges if "merges into" in l]
    out_l   = [l for l in merges if "merges out of" in l]
    if into_l and out_l:
        return [_v("EGO-04",
            f"Ego merges overlapping with conflicting label (frame {fi}): "
            f"'{into_l[0]}' + '{out_l[0]}'", fi)]
    return []


@rule("EGO-05", "Ego turns conflicting labels (left + right at same time)")
def check_ego_turns_conflict(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    turns  = [l for l in labels if l in EGO_MOVEMENT_LABELS and "ego turns" in l]
    lefts  = [l for l in turns if "left" in l]
    rights = [l for l in turns if "right" in l]
    if lefts and rights:
        return [_v("EGO-05",
            f"'Ego turns...' overlapping with conflicting label (frame {fi}): "
            f"'{lefts[0]}' + '{rights[0]}'", fi)]
    return []


@rule("EGO-06", "Ego changes lane right and left simultaneously")
def check_ego_lane_change(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    right = "ego changes lane right out of ego lane" in labels
    left  = "ego changes lane left out of ego lane" in labels
    if right and left:
        return [_v("EGO-06",
            f"'Ego changes lane right' overlapping with 'Ego changes lane left' (frame {fi})", fi)]
    return []


@rule("EGO-07", "Ego merges into Rotary and Ego merges out of Rotary cannot overlap")
def check_ego_rotary(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    labels = _labels(ann)
    if "ego merges into rotary" in labels and "ego merges out of rotary" in labels:
        return [_v("EGO-07",
            f"'Ego merges into Rotary' overlapping with 'Ego merges out of Rotary' (frame {fi})", fi)]
    return []


# ════════════════════════════════════════════════
#  GENERAL MOVEMENT — SOP page 12
# ════════════════════════════════════════════════

@rule("MOV-01", "Label used for multiple annotations within same timeline")
def check_duplicate_timeline(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    count: dict[str, int] = {}
    for obj in ann.get("objects", []):
        l = _n(obj.get("label", ""))
        count[l] = count.get(l, 0) + 1
    violations = []
    for label, n in count.items():
        if n > 1:
            violations.append(_v("MOV-01",
                f"'{label}' at frame {fi} used for multiple annotations within the same timeline. "
                f"Create a separate annotation for each movement.", fi))
    return violations


# ════════════════════════════════════════════════
#  BBOX / LABEL VALIDATION
# ════════════════════════════════════════════════

@rule("BOX-01", "Bounding box too large or too small", Severity.WARNING)
def check_box_size(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    W = ctx.get("frame_width", 1920)
    H = ctx.get("frame_height", 1080)
    area = W * H
    violations = []
    for obj in [o for o in ann.get("objects", []) if o.get("track_id")]:
        bb = obj.get("bbox")
        if not bb or len(bb) < 4:
            continue
        x, y, w, h = bb
        ratio = (w * h) / area
        tid = obj.get("track_id")
        if ratio > 0.90:
            violations.append(_v("BOX-01",
                f"Bounding box too large — covers {ratio*100:.1f}% of frame (frame {fi})",
                fi, tid, Severity.WARNING))
        elif 0 < w * h < area * 0.0001:
            violations.append(_v("BOX-01",
                f"Bounding box extremely small — possible noise annotation (frame {fi})",
                fi, tid, Severity.WARNING))
    return violations


@rule("BOX-02", "Bounding box outside frame boundary")
def check_out_of_bounds(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    W = ctx.get("frame_width", 1920)
    H = ctx.get("frame_height", 1080)
    violations = []
    for obj in ann.get("objects", []):
        bb = obj.get("bbox")
        if not bb or len(bb) < 4:
            continue
        x, y, w, h = bb
        if x < 0 or y < 0 or (x + w) > W or (y + h) > H:
            violations.append(_v("BOX-02",
                f"Bounding box extends outside frame boundary (frame {fi})",
                fi, obj.get("track_id")))
    return violations


@rule("LBL-01", "Label not in Tapi Tag Ontology Phase II")
def check_valid_label(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    violations = []
    for obj in ann.get("objects", []):
        label = _n(obj.get("label", ""))
        if label and label not in ALL_VALID_LABELS:
            violations.append(_v("LBL-01",
                f"Unknown label '{obj.get('label')}' — not in Tapi Tag Ontology Phase II (frame {fi})",
                fi, obj.get("track_id")))
    return violations


@rule("TRK-01", "Track ID label must remain consistent across frames")
def check_track_consistency(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    history: dict = ctx.setdefault("track_history", {})
    violations = []
    for obj in ann.get("objects", []):
        tid   = obj.get("track_id")
        label = _n(obj.get("label", ""))
        if not tid:
            continue
        if tid not in history:
            history[tid] = label
        elif history[tid] != label:
            violations.append(_v("TRK-01",
                f"ID {tid} label changed: '{history[tid]}' → '{label}' (frame {fi})",
                fi, tid))
    return violations


@rule("TRK-02", "Duplicate bounding boxes for same object in same frame")
def check_duplicates(ann: dict, fi: int, ctx: dict) -> list[RuleViolation]:
    # Only check actor objects (those with track_id)
    objects = [o for o in ann.get("objects", []) if o.get("track_id")]
    seen: list = []
    violations = []
    for obj in objects:
        bb = obj.get("bbox")
        if not bb or len(bb) < 4:
            continue
        for s in seen:
            if _iou(bb, s) > 0.70:
                violations.append(_v("TRK-02",
                    f"Duplicate annotation detected (IoU > 0.70) (frame {fi})",
                    fi, obj.get("track_id")))
        seen.append(bb)
    return violations


# ════════════════════════════════════════════════
#  Helpers
# ════════════════════════════════════════════════

def _n(s: str) -> str:
    return s.strip().lower()

def _labels(ann: dict) -> list[str]:
    return [_n(o.get("label", "")) for o in ann.get("objects", [])]

def _v(rule_id: str, message: str, frame: int,
       track_id: str | None = None,
       severity: Severity = Severity.ERROR) -> RuleViolation:
    return RuleViolation(rule_id=rule_id, severity=severity,
                         message=message, frame=frame, track_id=track_id)

def _iou(a: list, b: list) -> float:
    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


def run_all_rules(ann: dict, frame_idx: int, ctx: dict) -> list[RuleViolation]:
    violations: list[RuleViolation] = []
    for meta in RULES.values():
        try:
            violations.extend(meta["fn"](ann, frame_idx, ctx))
        except Exception:
            pass
    return violations