import face_recognition
import json
import annoy
from typing import Tuple

EMBEDDING_DIMENSION=128
ANNOY_INDEX_FILE = "models/actors_annoy_index.ann"
ANNOY_METADATA_FILE = "models/actors_annoy_metadata.json"
ANNOY_MAPPING_FILE = "models/actors_mapping.json"

def load_annoy_index(
    index_file = ANNOY_INDEX_FILE, 
    metadata_file = ANNOY_METADATA_FILE, 
    mapping_file = ANNOY_MAPPING_FILE
) -> Tuple[annoy.AnnoyIndex, dict]:
    """Load annoy index and associated mapping file"""
    with open(metadata_file) as f: 
        annoy_index_metadata = json.load(f)

    annoy_index = annoy.AnnoyIndex(f=EMBEDDING_DIMENSION, **annoy_index_metadata)
    annoy_index.load(index_file)

    with open(mapping_file) as f: 
        mapping = json.load(f)
        mapping = {int(k): v for k, v in mapping.items()}
    return annoy_index, mapping

def analyze_image(image, annoy_index, n_matches: int = 1, num_jitters: int = 1, model: str = "large"):
    """Extract face location, embeddings, and top n_matches matches"""
    face_locations = face_recognition.face_locations(image)
    embeddings = face_recognition.face_encodings(image, num_jitters=num_jitters, model=model, known_face_locations=face_locations)
    matches = []
    distances = []
    for emb in embeddings:
        m, d = annoy_index.get_nns_by_vector(emb, n_matches, include_distances=True)
        matches.append(m)
        distances.append(d)
    return [dict(embeddings=e, matches=m, distances=d, face_locations=f) for e,m,d,f in zip(embeddings, matches, distances, face_locations)]