// Keywords for a file from its CLIP embedding.
//
// Ente stores no keywords. A file's ML data holds a CLIP embedding, the 512
// numbers Ente's image model gave the photo, and Ente's photos app sorts files
// into the categories of its Discover section with it: a file falls in a
// category when the cosine similarity of its CLIP embedding and the text
// embedding of the category's query is at least the category's `minScore`,
// unless the file is a video and the category leaves out videos. A file's
// keywords are the titles of the categories it falls in. The categories and
// their text embeddings are in `clip-keyword-vectors.ts`.

import { DISCOVER_CATEGORIES } from "./clip-keyword-vectors.js";
import type { MLData } from "./mldata-fetch.js";

// How many numbers a CLIP embedding from Ente's image model has.
const EMBEDDING_LENGTH = 512;

const length = (vector: number[]): number => Math.hypot(...vector);

const cosineSimilarity = (a: number[], b: number[]): number => {
    let dot = 0;
    a.forEach((value, i) => {
        dot += value * (b[i] ?? 0);
    });
    return dot / (length(a) * length(b));
};

// The titles of the Discover categories a file falls in, highest score first,
// or undefined when its ML data has no CLIP embedding of 512 finite numbers
// that are not all zero. Tests pass their own `categories`.
export const clipKeywords = (
    mlData: MLData,
    isVideo: boolean,
    categories = DISCOVER_CATEGORIES,
): string[] | undefined => {
    const clip = mlData.clip;
    if (typeof clip !== "object" || clip === null) return undefined;
    const embedding = (clip as { embedding?: unknown }).embedding;
    if (
        !Array.isArray(embedding) ||
        embedding.length !== EMBEDDING_LENGTH ||
        !embedding.every((v) => typeof v === "number" && Number.isFinite(v)) ||
        length(embedding) === 0
    ) {
        return undefined;
    }

    const matches: { title: string; score: number }[] = [];
    for (const category of categories) {
        if (isVideo && !category.showVideo) continue;
        const score = cosineSimilarity(embedding, category.textEmbedding);
        if (score >= category.minScore) {
            matches.push({ title: category.title, score });
        }
    }
    return matches
        .sort((a, b) => b.score - a.score)
        .map((match) => match.title);
};