services/correlation_engine_service.py

"""
Correlation Engine service — correlates evidence across all providers
into a deterministic relationship graph.

Takes a UnifiedFaceReport (the output of any analysis job) and:
  1. Extracts nodes: faces, objects, metadata, images, embeddings.
  2. Runs deterministic matchers to find relationships.
  3. Builds a CorrelationGraph with typed edges.

No AI — only deterministic matching via cores.correlation.matchers.

Usage:
  POST /analysis/correlate with an image → runs full pipeline → correlates
  POST /analysis/correlate with a pre-computed report → correlates only
"""

from __future__ import annotations

import time
from typing import Optional

import numpy as np

from cores.correlation import (
    CorrelationGraphBuilder,
    Node,
    EdgeType,
    match_faces,
    match_objects,
    match_locations,
    match_cameras,
    match_hashes,
    match_embeddings,
    match_metadata,
    match_timestamps,
)
from cores.vision import sha256_bytes, phash
from models.reports import UnifiedFaceReport
from pipeline import InputValidator, ImagePreprocessor, ImageHasher
from utils.logging import execution_context, new_execution_id


class CorrelationEngineService:
    """Correlates evidence across providers into a relationship graph."""

    def __init__(
        self,
        validator: InputValidator,
        preprocessor: ImagePreprocessor,
        hasher: ImageHasher,
    ) -> None:
        self._validator = validator
        self._preprocessor = preprocessor
        self._hasher = hasher

    async def correlate(self, report: dict) -> dict:
        """Correlate a pre-computed report.

        Args:
            report: a UnifiedFaceReport as a dict (or the report sub-object).

        Returns:
            {"success": True, "correlation_graph": {...}, "elapsed_ms": float}
        """
        eid = new_execution_id()
        with execution_context(execution_id=eid, provider_id="correlation_engine"):
            t0 = time.perf_counter()

            builder = CorrelationGraphBuilder()

            # --- Extract nodes ---

            # Face nodes (from detections)
            face_embeddings: list[tuple[str, np.ndarray]] = []
            for i, det in enumerate(report.get("detections", [])):
                node_id = f"face_{i}"
                builder.add_node(Node(
                    id=node_id,
                    node_type="face",
                    label=f"Face {i}",
                    properties={"box": det.get("box", {})},
                ))
                if det.get("embedding"):
                    face_embeddings.append((node_id, np.array(det["embedding"])))

            # Object nodes (from object_detections + object_intelligence)
            objects_for_matching: list[tuple[str, str, dict]] = []
            # From object_intelligence (richer)
            obj_intel = report.get("object_intelligence") or {}
            for i, obj in enumerate(obj_intel.get("objects", [])):
                node_id = f"object_{i}"
                builder.add_node(Node(
                    id=node_id,
                    node_type="object",
                    label=obj.get("class_label", "object"),
                    properties={"box": obj.get("box", {}), "confidence": obj.get("confidence", 0)},
                ))
                objects_for_matching.append((node_id, obj.get("class_label", ""), obj.get("box", {})))
            # Fallback: from object_detections (provider-level)
            if not objects_for_matching:
                for od in report.get("object_detections", []):
                    for i, obj in enumerate(od.get("objects", [])):
                        node_id = f"object_{od.get('provider', 'x')}_{i}"
                        builder.add_node(Node(
                            id=node_id,
                            node_type="object",
                            label=obj.get("label", "object"),
                            properties={"box": obj.get("box", {})},
                        ))
                        objects_for_matching.append((node_id, obj.get("label", ""), obj.get("box", {})))

            # Metadata nodes (from metadata_extractions + forensic_metadata)
            metadata_for_matching: list[tuple[str, dict]] = []
            for i, meta in enumerate(report.get("metadata_extractions", [])):
                node_id = f"metadata_{i}"
                builder.add_node(Node(
                    id=node_id,
                    node_type="metadata",
                    label=meta.get("format", "metadata"),
                    properties={"provider": meta.get("provider", "")},
                ))
                metadata_for_matching.append((node_id, meta.get("exif", {})))
            # Forensic metadata
            forensic_meta = report.get("forensic_metadata")
            if forensic_meta:
                node_id = "forensic_metadata"
                builder.add_node(Node(
                    id=node_id,
                    node_type="metadata",
                    label="forensic_metadata",
                    properties={"camera_make": forensic_meta.get("camera_make"),
                                "camera_model": forensic_meta.get("camera_model")},
                ))
                metadata_for_matching.append((node_id, forensic_meta.get("exif", {})))

            # Image nodes (hash)
            image_hashes: list[tuple[str, str, str]] = []
            image_hash = report.get("metadata", {}).get("image_hash")
            if image_hash:
                node_id = "image_0"
                builder.add_node(Node(
                    id=node_id,
                    node_type="image",
                    label="source_image",
                    properties={"sha256": image_hash},
                ))
                image_hashes.append((node_id, image_hash, ""))  # phash empty if not available

            # Embedding nodes
            image_embeddings: list[tuple[str, np.ndarray]] = []
            for i, emb in enumerate(report.get("embedding_results", [])):
                node_id = f"embedding_{i}"
                builder.add_node(Node(
                    id=node_id,
                    node_type="embedding",
                    label=emb.get("model", "embedding"),
                    properties={"dimensions": emb.get("dimensions", 0)},
                ))
                if emb.get("embedding"):
                    image_embeddings.append((node_id, np.array(emb["embedding"])))

            # Location nodes
            location_items: list[tuple[str, dict]] = []
            loc_est = report.get("location_estimate")
            if loc_est and loc_est.get("gps"):
                node_id = "location_0"
                builder.add_node(Node(
                    id=node_id,
                    node_type="location",
                    label="gps_location",
                    properties=loc_est["gps"],
                ))
                location_items.append((node_id, loc_est["gps"]))

            # --- Run matchers ---

            if face_embeddings:
                builder.add_matches(EdgeType.SAME_FACE, match_faces(face_embeddings))

            if objects_for_matching:
                builder.add_matches(EdgeType.SAME_OBJECT, match_objects(objects_for_matching))

            if location_items:
                builder.add_matches(EdgeType.SAME_LOCATION, match_locations(location_items))

            if image_hashes:
                builder.add_matches(EdgeType.SAME_HASH, match_hashes(image_hashes))

            if image_embeddings:
                builder.add_matches(EdgeType.SAME_EMBEDDING, match_embeddings(image_embeddings))

            if metadata_for_matching:
                builder.add_matches(EdgeType.SAME_METADATA, match_metadata(metadata_for_matching))
                # Camera matching
                cameras: list[tuple[str, str, str, str]] = []
                for item_id, exif in metadata_for_matching:
                    make = exif.get("Make", "")
                    model = exif.get("Model", "")
                    fp = ""  # fingerprint would come from forensic_metadata
                    cameras.append((item_id, make, model, fp))
                if forensic_meta and forensic_meta.get("camera_fingerprint"):
                    cameras.append(("forensic_metadata",
                                    forensic_meta.get("camera_make", ""),
                                    forensic_meta.get("camera_model", ""),
                                    forensic_meta["camera_fingerprint"]))
                builder.add_matches(EdgeType.SAME_CAMERA, match_cameras(cameras))

                # Timestamp matching
                timestamps: list[tuple[str, str]] = []
                for item_id, exif in metadata_for_matching:
                    ts = exif.get("DateTimeOriginal") or exif.get("DateTime")
                    if ts:
                        timestamps.append((item_id, str(ts)))
                if timestamps:
                    builder.add_matches(EdgeType.SAME_TIMESTAMP, match_timestamps(timestamps))

            elapsed = (time.perf_counter() - t0) * 1000.0
            graph = builder.build(elapsed_ms=elapsed)

            return {
                "success": True,
                "correlation_graph": graph.model_dump(),
                "elapsed_ms": round(elapsed, 3),
            }