face_model/analyzer.py

"""
face_model/analyzer.py — InsightFace wrapper với RAM cache và fallback matching

Chức năng:
    - Khởi tạo InsightFace model (singleton)
    - Detect khuôn mặt & trích xuất embedding
    - RAM cache: known_embeddings, names, person_ids, etc.
    - find_match_ram: cosine similarity fallback khi Supabase RPC thất bại
"""
import logging
import os

import cv2
import numpy as np
from insightface.app import FaceAnalysis

import config

logger = logging.getLogger("face_analyzer")


class FaceAnalyzer:
    """Singleton quản lý InsightFace model và RAM embedding cache."""

    def __init__(self):
        self._analyzer: FaceAnalysis | None = None

        # RAM cache — parallel arrays (index-aligned)
        self.known_embeddings:  list[np.ndarray] = []
        self.known_names:       list[str]         = []
        self.known_person_ids:  list[str]         = []
        self.known_embedding_ids: list[str]       = []
        self.known_mongo_ids:   list[str]         = []

    # ── Model Init ────────────────────────────────────────────────────────

    def initialize(self) -> None:
        """Khởi tạo InsightFace model. Gọi một lần khi khởi động server."""
        model_root = os.getenv("INSIGHTFACE_ROOT", "~/.insightface")
        logger.info(
            f"[Model] Loading InsightFace '{config.MODEL_NAME}' from root '{model_root}' "
            f"with provider '{config.MODEL_PROVIDER}'..."
        )
        self._analyzer = FaceAnalysis(
            name=config.MODEL_NAME,
            root=model_root,
            providers=[config.MODEL_PROVIDER],
        )
        self._analyzer.prepare(ctx_id=0, det_size=config.DET_SIZE)
        logger.info("[Model] InsightFace initialized successfully.")

    @property
    def is_ready(self) -> bool:
        return self._analyzer is not None

    # ── Face Detection & Embedding ────────────────────────────────────────

    def get_faces(self, img: np.ndarray) -> list:
        """Detect tất cả khuôn mặt trong ảnh. Returns list of face objects."""
        if not self.is_ready:
            raise RuntimeError("Face analyzer not initialized.")
        return self._analyzer.get(img)

    # ── RAM Cache ─────────────────────────────────────────────────────────

    def reload_cache(self, records: list[dict]) -> int:
        """
        Tải lại RAM cache từ danh sách records.
        records: list[{person_id, full_name, embedding_id, embedding, face_crop_mongo_id}]
        Returns: số lượng embeddings đã load
        """
        temp_embeddings:    list[np.ndarray] = []
        temp_names:         list[str]         = []
        temp_person_ids:    list[str]         = []
        temp_embedding_ids: list[str]         = []
        temp_mongo_ids:     list[str]         = []

        for r in records:
            emb = r.get("embedding", [])
            if not emb:
                continue
            temp_embeddings.append(np.array(emb, dtype=np.float32))
            temp_names.append(r.get("full_name", ""))
            temp_person_ids.append(r.get("person_id", ""))
            temp_embedding_ids.append(r.get("embedding_id", ""))
            temp_mongo_ids.append(r.get("face_crop_mongo_id", ""))

        self.known_embeddings    = temp_embeddings
        self.known_names         = temp_names
        self.known_person_ids    = temp_person_ids
        self.known_embedding_ids = temp_embedding_ids
        self.known_mongo_ids     = temp_mongo_ids

        logger.info(f"[Cache] Reloaded {len(self.known_embeddings)} embeddings into RAM.")
        return len(self.known_embeddings)

    def reload_from_local_folder(self, folder: str) -> int:
        """
        Fallback: Load embeddings từ thư mục ảnh cục bộ (khi Supabase offline).
        """
        if not os.path.exists(folder):
            return 0

        self.known_embeddings    = []
        self.known_names         = []
        self.known_person_ids    = []
        self.known_embedding_ids = []
        self.known_mongo_ids     = []

        files = sorted([
            f for f in os.listdir(folder)
            if f.lower().endswith((".png", ".jpg", ".jpeg"))
        ])
        for filename in files:
            img_path = os.path.join(folder, filename)
            img = cv2.imread(img_path)
            if img is None:
                continue
            faces = self._analyzer.get(img)
            if faces:
                self.known_embeddings.append(faces[0].normed_embedding)
                name = os.path.splitext(filename)[0]
                self.known_names.append(name)
                self.known_person_ids.append("")
                self.known_embedding_ids.append("")
                self.known_mongo_ids.append("")

        logger.info(f"[Cache Fallback] Loaded {len(self.known_embeddings)} from '{folder}'.")
        return len(self.known_embeddings)

    def add_to_cache(
        self,
        embedding:    np.ndarray,
        name:         str,
        person_id:    str,
        embedding_id: str,
        mongo_id:     str,
    ) -> None:
        """Thêm một embedding mới vào RAM cache."""
        self.known_embeddings.append(embedding)
        self.known_names.append(name)
        self.known_person_ids.append(person_id)
        self.known_embedding_ids.append(embedding_id)
        self.known_mongo_ids.append(mongo_id)

    def update_name_in_cache(self, person_id: str, new_name: str) -> None:
        """Cập nhật tên trong RAM cache sau khi update trên DB."""
        for i, pid in enumerate(self.known_person_ids):
            if pid == person_id:
                self.known_names[i] = new_name

    def remove_from_cache(self, person_id: str) -> None:
        """Xóa tất cả entries của person khỏi RAM cache."""
        indices = [i for i, pid in enumerate(self.known_person_ids) if pid == person_id]
        for i in reversed(indices):
            self.known_embeddings.pop(i)
            self.known_names.pop(i)
            self.known_person_ids.pop(i)
            self.known_embedding_ids.pop(i)
            self.known_mongo_ids.pop(i)

    @property
    def total(self) -> int:
        return len(self.known_embeddings)

    # ── RAM-based Matching (fallback) ─────────────────────────────────────

    def find_match_ram(
        self,
        embedding: np.ndarray,
        threshold: float,
    ) -> tuple[float, int]:
        """
        Tìm khuôn mặt khớp nhất trong RAM cache bằng cosine similarity.
        Returns:
            (max_similarity, best_index)  — best_index = -1 nếu không khớp
        """
        if not self.known_embeddings:
            return 0.0, -1

        similarities = [float(np.dot(embedding, e)) for e in self.known_embeddings]
        max_sim = max(similarities)
        best_idx = similarities.index(max_sim)

        if max_sim >= threshold:
            return max_sim, best_idx
        return max_sim, -1


# ── Singleton ─────────────────────────────────────────────────────────────
face_analyzer = FaceAnalyzer()