src/data/gallery_store.py

from __future__ import annotations

import json
import re
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np
from PIL import Image

from src import config


@dataclass
class GalleryRecord:
    person: str
    gallery_filename: str
    query_filename: str
    gallery_path: str
    query_path: str


class DemoGallery:
    def __init__(self) -> None:
        gallery_npz = Path(config.LFW_FEATURES_PATH)
        manifest_path = Path(config.LFW_VISIBLE_MANIFEST)
        example_manifest_path = Path(config.LFW_EXAMPLES_MANIFEST)

        data = np.load(gallery_npz, allow_pickle=True)
        self.base_features = data["features"].astype(np.float32)
        self.base_people = data["persons"].astype(str)
        self.base_filenames = data["filenames"].astype(str)
        self.base_identity_count = int(len(set(self.base_people.tolist())))
        self.feature_dim = int(self.base_features.shape[1])

        calibration_npz = Path(config.LFW_CALIBRATION_PATH)
        calibration = np.load(calibration_npz, allow_pickle=True)
        self.calibration_features = calibration["features"].astype(np.float32)
        self.calibration_people = calibration["persons"].astype(str)
        self.calibration_filenames = calibration["filenames"].astype(str)

        with open(manifest_path) as fh:
            self.base_records = [GalleryRecord(**row) for row in json.load(fh)]
        with open(example_manifest_path) as fh:
            self.examples = json.load(fh)

        self.features = self.base_features
        self.people = self.base_people
        self.filenames = self.base_filenames
        self._reload_enrollments()

    def _empty_features(self) -> np.ndarray:
        return np.empty((0, self.feature_dim), dtype=np.float32)

    def _person_index(self, person: str) -> int | None:
        normalized = person.strip().lower()
        for idx, existing in enumerate(self.enrolled_people.tolist()):
            if str(existing).strip().lower() == normalized:
                return idx
        return None

    @staticmethod
    def _slugify(person: str) -> str:
        slug = re.sub(r"[^a-zA-Z0-9]+", "_", person.strip()).strip("_")
        return slug or "user"

    def _reload_enrollments(self) -> None:
        if config.ENROLLMENT_FEATURES_PATH.exists():
            data = np.load(config.ENROLLMENT_FEATURES_PATH, allow_pickle=True)
            self.enrolled_features = data["features"].astype(np.float32)
            self.enrolled_people = data["persons"].astype(str)
            self.enrolled_filenames = data["filenames"].astype(str)
        else:
            self.enrolled_features = self._empty_features()
            self.enrolled_people = np.array([], dtype=str)
            self.enrolled_filenames = np.array([], dtype=str)

        if config.ENROLLMENT_MANIFEST_PATH.exists():
            with open(config.ENROLLMENT_MANIFEST_PATH) as fh:
                self.enrolled_records = [GalleryRecord(**row) for row in json.load(fh)]
        else:
            self.enrolled_records = []

        self.records = self.enrolled_records + self.base_records
        combined_people = self.base_people.tolist() + self.enrolled_people.tolist()
        self.identity_count = int(len(set(combined_people)))
        self.enrolled_identity_count = int(len(set(self.enrolled_people.tolist())))

    def _save_enrollments(self) -> None:
        np.savez(
            config.ENROLLMENT_FEATURES_PATH,
            features=self.enrolled_features.astype(np.float32),
            persons=self.enrolled_people.astype(str),
            filenames=self.enrolled_filenames.astype(str),
        )
        with open(config.ENROLLMENT_MANIFEST_PATH, "w") as fh:
            json.dump([asdict(record) for record in self.enrolled_records], fh, indent=2)

    def has_enrollments(self) -> bool:
        return len(self.enrolled_people) > 0

    def enroll(self, person: str, image: Image.Image, embedding: np.ndarray) -> tuple[GalleryRecord, bool]:
        person = " ".join(person.split())
        if not person:
            raise ValueError("Person name cannot be empty")

        filename = f"enrolled_{self._slugify(person)}.jpg"
        image_path = config.ENROLLMENT_IMAGES_DIR / filename
        image.convert("RGB").save(image_path, format="JPEG", quality=95)

        embedding_row = np.asarray(embedding, dtype=np.float32).reshape(1, -1)
        if embedding_row.shape[1] != self.feature_dim:
            raise ValueError(f"Embedding dimension mismatch: expected {self.feature_dim}, got {embedding_row.shape[1]}")

        record = GalleryRecord(
            person=person,
            gallery_filename=filename,
            query_filename=filename,
            gallery_path=str(image_path),
            query_path=str(image_path),
        )

        existing_idx = self._person_index(person)
        replaced = existing_idx is not None
        if replaced:
            assert existing_idx is not None
            self.enrolled_features[existing_idx : existing_idx + 1] = embedding_row
            self.enrolled_people[existing_idx] = person
            self.enrolled_filenames[existing_idx] = filename
            self.enrolled_records[existing_idx] = record
        else:
            self.enrolled_features = np.vstack([self.enrolled_features, embedding_row])
            self.enrolled_people = np.concatenate([self.enrolled_people, np.array([person], dtype=str)])
            self.enrolled_filenames = np.concatenate([self.enrolled_filenames, np.array([filename], dtype=str)])
            self.enrolled_records.append(record)

        self._save_enrollments()
        self._reload_enrollments()
        return record, replaced

    def summary(self) -> dict[str, object]:
        return {
            "gallery_size": int(len(self.base_people) + len(self.enrolled_people)),
            "identity_count": self.identity_count,
            "visible_gallery_size": int(len(self.records)),
            "calibration_size": int(len(self.calibration_people)),
            "people": [str(x.person) for x in self.records],
            "dimensions": self.feature_dim,
            "enrolled_gallery_size": int(len(self.enrolled_people)),
            "enrolled_identity_count": self.enrolled_identity_count,
            "demo_gallery_size": int(len(self.base_people)),
            "demo_identity_count": self.base_identity_count,
        }

    def example_choices(self) -> list[tuple[str, str]]:
        return [(item["person"], str(config.EXAMPLES_DIR / item["filename"])) for item in self.examples]

    def load_example_image(self, filename: str) -> Image.Image:
        return Image.open(config.EXAMPLES_DIR / filename).convert("RGB")

    def gallery_image_path(self, gallery_filename: str, person: str | None = None) -> Path:
        direct = Path(gallery_filename)
        if direct.is_absolute() and direct.exists():
            return direct

        enrolled_direct = config.ENROLLMENT_IMAGES_DIR / gallery_filename
        if enrolled_direct.exists():
            return enrolled_direct

        base_direct = config.GALLERY_DIR / gallery_filename
        if base_direct.exists():
            return base_direct

        def resolve_record_path(record: GalleryRecord) -> Path | None:
            for candidate in (
                Path(record.gallery_path),
                config.GALLERY_DIR / record.gallery_path,
                config.EXAMPLES_DIR / record.gallery_path,
                config.GALLERY_DIR / record.gallery_filename,
                config.EXAMPLES_DIR / record.gallery_filename,
            ):
                if candidate.exists():
                    return candidate
            return None

        if person is not None:
            for record in self.enrolled_records + self.base_records:
                if record.person == person:
                    candidate = resolve_record_path(record)
                    if candidate is not None:
                        return candidate

        for record in self.enrolled_records + self.base_records:
            if record.gallery_filename == gallery_filename:
                candidate = resolve_record_path(record)
                if candidate is not None:
                    return candidate

        raise FileNotFoundError(f"No visible gallery image found for {gallery_filename}")

    def gallery_items(self, page: int = 1, per_page: int = 24) -> list[tuple[str, str]]:
        page = max(1, int(page or 1))
        start = (page - 1) * per_page
        end = start + per_page
        items: list[tuple[str, str]] = []
        for record in self.records[start:end]:
            items.append((str(self.gallery_image_path(record.gallery_filename, record.person)), record.person))
        return items

    def total_pages(self, per_page: int = 24) -> int:
        return max(1, (len(self.records) + per_page - 1) // per_page)

    def database_info(self, page: int = 1, per_page: int = 24) -> str:
        summary = self.summary()
        total_pages = self.total_pages(per_page)
        page = min(max(1, int(page or 1)), total_pages)
        start = (page - 1) * per_page + 1
        end = min(page * per_page, len(self.records))
        return (
            f"{summary['identity_count']} identities | {summary['gallery_size']} images\n"
            f"Enrolled: {summary['enrolled_identity_count']} identities / {summary['enrolled_gallery_size']} images\n"
            f"Visible: {start}-{end} of {summary['visible_gallery_size']} | Page {page}/{total_pages}"
        )