cores/correlation/graph.py

"""Correlation graph builder — deterministic relationship graph.

Builds a CorrelationGraph (nodes + edges) from evidence collected
across multiple providers.  No AI — only deterministic matching.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import List

from models.reports import (
    CorrelationGraph,
    CorrelationNode,
    CorrelationEdge,
)


class EdgeType(str, Enum):
    SAME_FACE = "same_face"
    SAME_OBJECT = "same_object"
    SAME_LOCATION = "same_location"
    SAME_CAMERA = "same_camera"
    SAME_HASH = "same_hash"
    SAME_EMBEDDING = "same_embedding"
    SAME_METADATA = "same_metadata"
    SAME_TIMESTAMP = "same_timestamp"


@dataclass
class Node:
    id: str
    node_type: str          # "face" | "object" | "location" | "metadata" | "image" | "embedding"
    label: str = ""
    properties: dict = field(default_factory=dict)


@dataclass
class Edge:
    source: str
    target: str
    edge_type: EdgeType
    confidence: float = 1.0
    evidence: str = ""


class CorrelationGraphBuilder:
    """Builds a CorrelationGraph from nodes + edges."""

    def __init__(self) -> None:
        self._nodes: list[Node] = []
        self._edges: list[Edge] = []

    def add_node(self, node: Node) -> None:
        self._nodes.append(node)

    def add_nodes(self, nodes: List[Node]) -> None:
        self._nodes.extend(nodes)

    def add_edge(self, edge: Edge) -> None:
        self._edges.append(edge)

    def add_edges(self, edges: List[Edge]) -> None:
        self._edges.extend(edges)

    def add_matches(
        self,
        edge_type: EdgeType,
        matches: List[tuple],
    ) -> None:
        """Add matched pairs as edges.

        Args:
            edge_type: the type of relationship.
            matches: list of (source_id, target_id, confidence, evidence_str) tuples.
        """
        for source, target, confidence, evidence in matches:
            self._edges.append(Edge(
                source=source,
                target=target,
                edge_type=edge_type,
                confidence=confidence,
                evidence=evidence,
            ))

    def build(self, elapsed_ms: float = 0.0) -> CorrelationGraph:
        """Build the final CorrelationGraph."""
        # Deduplicate nodes by ID
        seen_ids: set[str] = set()
        unique_nodes: list[Node] = []
        for n in self._nodes:
            if n.id not in seen_ids:
                seen_ids.add(n.id)
                unique_nodes.append(n)

        # Deduplicate edges (same source+target+type)
        seen_edges: set[tuple[str, str, str]] = set()
        unique_edges: list[Edge] = []
        for e in self._edges:
            key = (e.source, e.target, e.edge_type.value)
            reverse_key = (e.target, e.source, e.edge_type.value)
            if key not in seen_edges and reverse_key not in seen_edges:
                seen_edges.add(key)
                unique_edges.append(e)

        return CorrelationGraph(
            nodes=[
                CorrelationNode(
                    id=n.id,
                    node_type=n.node_type,
                    label=n.label,
                    properties=n.properties,
                )
                for n in unique_nodes
            ],
            edges=[
                CorrelationEdge(
                    source=e.source,
                    target=e.target,
                    edge_type=e.edge_type.value,
                    confidence=e.confidence,
                    evidence=e.evidence,
                )
                for e in unique_edges
            ],
            num_nodes=len(unique_nodes),
            num_edges=len(unique_edges),
            elapsed_ms=round(elapsed_ms, 3),
        )