vector_db.py

"""
FAISS Vector Database for Multi-Modal Profile Search
======================================================
Maintains separate FAISS indices for each biometric modality:
  - gait_index:       256-d (gait embeddings)
  - biomech_index:     64-d (biomechanical features)
  - appearance_index: 512-d (OSNet appearance embeddings)
  - face_index:       512-d (ArcFace embeddings, optional)

Uses IndexFlatIP (inner product) on L2-normalized vectors = cosine similarity.
"""

import os
import json
import numpy as np

# Try FAISS GPU, then CPU, then fallback to brute force
FAISS_AVAILABLE = False
FAISS_GPU = False
try:
    import faiss
    FAISS_AVAILABLE = True
    # Check for GPU support
    try:
        if faiss.get_num_gpus() > 0:
            FAISS_GPU = True
            print("[VectorDB] FAISS-GPU available")
        else:
            print("[VectorDB] FAISS-CPU available")
    except Exception:
        print("[VectorDB] FAISS-CPU available")
except ImportError:
    print("[VectorDB] FAISS not found — using brute-force numpy fallback")


class _NumpyFallbackIndex:
    """Brute-force numpy fallback when FAISS is not available."""

    def __init__(self, dim):
        self.dim = dim
        self.vectors = np.zeros((0, dim), dtype=np.float32)

    @property
    def ntotal(self):
        return self.vectors.shape[0]

    def add(self, vectors):
        self.vectors = np.vstack([self.vectors, vectors.astype(np.float32)])

    def search(self, query, k):
        if self.ntotal == 0:
            return np.array([[-1.0] * k]), np.array([[-1] * k])
        k = min(k, self.ntotal)
        # Inner product (= cosine sim on normalized vectors)
        sims = np.dot(self.vectors, query.T).flatten()
        top_k = np.argsort(sims)[::-1][:k]
        scores = sims[top_k].reshape(1, -1)
        indices = top_k.reshape(1, -1)
        return scores, indices

    def reset(self):
        self.vectors = np.zeros((0, self.dim), dtype=np.float32)


# Modality definitions
MODALITIES = {
    "gait":       {"dim": 256, "weight": 0.35},
    "biomech":    {"dim": 64,  "weight": 0.25},
    "appearance": {"dim": 512, "weight": 0.25},
    "height":     {"dim": 1,   "weight": 0.15},
    "face":       {"dim": 512, "weight": 0.30},  # optional, applied differently
}


class ProfileVectorDB:
    """
    Multi-modal FAISS vector database for person profiles.

    Each profile is stored across multiple indices (one per modality).
    All vectors must be L2-normalized before insertion.
    """

    def __init__(self, db_dir="saved_profiles"):
        self.db_dir = db_dir
        os.makedirs(db_dir, exist_ok=True)

        self.indices = {}     # modality_name -> FAISS index
        self.profile_ids = [] # ordered list of profile IDs (position = FAISS index position)
        self._id_file = os.path.join(db_dir, "faiss_profile_ids.json")

        # Build indices
        for mod_name, mod_info in MODALITIES.items():
            self.indices[mod_name] = self._create_index(mod_info["dim"])

        # Load existing
        self._load()

    def _create_index(self, dim):
        """Create a FAISS IndexFlatIP (inner product = cosine on normalized vecs)."""
        if FAISS_AVAILABLE:
            index = faiss.IndexFlatIP(dim)
            return index
        else:
            return _NumpyFallbackIndex(dim)

    def _load(self):
        """Load saved indices and profile ID mapping."""
        if os.path.exists(self._id_file):
            try:
                with open(self._id_file, "r") as f:
                    self.profile_ids = json.load(f)
            except Exception:
                self.profile_ids = []

        if FAISS_AVAILABLE:
            for mod_name in MODALITIES:
                idx_path = os.path.join(self.db_dir, f"faiss_{mod_name}.index")
                if os.path.exists(idx_path):
                    try:
                        self.indices[mod_name] = faiss.read_index(idx_path)
                    except Exception:
                        self.indices[mod_name] = self._create_index(MODALITIES[mod_name]["dim"])

    def save(self):
        """Persist indices and ID mapping to disk."""
        with open(self._id_file, "w") as f:
            json.dump(self.profile_ids, f)

        if FAISS_AVAILABLE:
            for mod_name in MODALITIES:
                idx_path = os.path.join(self.db_dir, f"faiss_{mod_name}.index")
                try:
                    faiss.write_index(self.indices[mod_name], idx_path)
                except Exception as e:
                    print(f"[VectorDB] Failed to save {mod_name} index: {e}")

    def add_profile(self, profile_id, vectors):
        """
        Add a profile's vectors to all indices.

        Args:
            profile_id: str
            vectors: dict with keys matching MODALITIES, values are np.array or None
        """
        self.profile_ids.append(profile_id)

        for mod_name, mod_info in MODALITIES.items():
            vec = vectors.get(mod_name)
            if vec is not None:
                vec = np.array(vec, dtype=np.float32).reshape(1, -1)
                # L2 normalize
                norm = np.linalg.norm(vec)
                if norm > 0:
                    vec = vec / norm
                self.indices[mod_name].add(vec)
            else:
                # Add zero vector as placeholder to keep indices aligned
                zero = np.zeros((1, mod_info["dim"]), dtype=np.float32)
                self.indices[mod_name].add(zero)

        self.save()

    def remove_profile(self, profile_id):
        """Remove a profile. Requires rebuilding indices (FAISS doesn't support deletion)."""
        if profile_id not in self.profile_ids:
            return

        idx = self.profile_ids.index(profile_id)
        self.profile_ids.pop(idx)

        # Rebuild indices without the removed profile
        # This is expensive but deletion is rare
        for mod_name, mod_info in MODALITIES.items():
            old_index = self.indices[mod_name]
            dim = mod_info["dim"]

            if old_index.ntotal <= 1:
                self.indices[mod_name] = self._create_index(dim)
                continue

            if FAISS_AVAILABLE:
                # Extract all vectors, remove the one at idx, rebuild
                all_vecs = faiss.rev_swig_ptr(
                    old_index.get_xb(), old_index.ntotal * dim
                ).reshape(old_index.ntotal, dim).copy()
                all_vecs = np.delete(all_vecs, idx, axis=0)
                new_index = self._create_index(dim)
                if len(all_vecs) > 0:
                    new_index.add(all_vecs)
                self.indices[mod_name] = new_index
            else:
                # Numpy fallback
                all_vecs = old_index.vectors.copy()
                all_vecs = np.delete(all_vecs, idx, axis=0)
                new_index = _NumpyFallbackIndex(dim)
                if len(all_vecs) > 0:
                    new_index.add(all_vecs)
                self.indices[mod_name] = new_index

        self.save()

    def search(self, query_vectors, k=5):
        """
        Search for closest profiles using multi-modal weighted ensemble.

        Args:
            query_vectors: dict with modality names as keys, np.array as values
                           e.g., {"gait": vec_256, "appearance": vec_512, ...}
            k: number of results to return

        Returns:
            list of (profile_id, weighted_score, per_modality_scores) sorted by score desc.
        """
        if len(self.profile_ids) == 0:
            return []

        k = min(k, len(self.profile_ids))
        profile_scores = {}  # profile_id -> {mod_name: score}

        for mod_name, mod_info in MODALITIES.items():
            query = query_vectors.get(mod_name)
            if query is None:
                continue
            if np.all(query == 0):
                continue

            query = np.array(query, dtype=np.float32).reshape(1, -1)
            norm = np.linalg.norm(query)
            if norm > 0:
                query = query / norm

            index = self.indices[mod_name]
            if index.ntotal == 0:
                continue

            scores, indices = index.search(query, min(k, index.ntotal))

            for i in range(scores.shape[1]):
                idx = int(indices[0][i])
                if 0 <= idx < len(self.profile_ids):
                    pid = self.profile_ids[idx]
                    if pid not in profile_scores:
                        profile_scores[pid] = {}
                    profile_scores[pid][mod_name] = max(0.0, float(scores[0][i]))

        # Compute weighted ensemble score per profile
        results = []
        for pid, mod_scores in profile_scores.items():
            # Standard 4-modal ensemble
            gait_sim = mod_scores.get("gait", 0.0)
            biomech_sim = mod_scores.get("biomech", 0.0)
            appearance_sim = mod_scores.get("appearance", 0.0)
            height_sim = mod_scores.get("height", 0.5)

            weighted = (0.35 * gait_sim +
                        0.25 * biomech_sim +
                        0.25 * appearance_sim +
                        0.15 * height_sim)

            # Optional face modality
            face_sim = mod_scores.get("face", None)
            if face_sim is not None and face_sim > 0:
                weighted = 0.70 * weighted + 0.30 * face_sim

            results.append((pid, weighted, mod_scores))

        results.sort(key=lambda x: x[1], reverse=True)
        return results[:k]

    def get_profile_count(self):
        return len(self.profile_ids)