import { fs } from "@llamaindex/env";
import _ from "lodash";
import { filetypemime } from "magic-bytes.js";
import { DEFAULT_SIMILARITY_TOP_K } from "../constants.js";
export var SimilarityType;
(function(SimilarityType) {
    SimilarityType["DEFAULT"] = "cosine";
    SimilarityType["DOT_PRODUCT"] = "dot_product";
    SimilarityType["EUCLIDEAN"] = "euclidean";
})(SimilarityType || (SimilarityType = {}));
/**
 * The similarity between two embeddings.
 * @param embedding1
 * @param embedding2
 * @param mode
 * @returns similarity score with higher numbers meaning the two embeddings are more similar
 */ export function similarity(embedding1, embedding2, mode = "cosine") {
    if (embedding1.length !== embedding2.length) {
        throw new Error("Embedding length mismatch");
    }
    // NOTE I've taken enough Kahan to know that we should probably leave the
    // numeric programming to numeric programmers. The naive approach here
    // will probably cause some avoidable loss of floating point precision
    // ml-distance is worth watching although they currently also use the naive
    // formulas
    function norm(x) {
        let result = 0;
        for(let i = 0; i < x.length; i++){
            result += x[i] * x[i];
        }
        return Math.sqrt(result);
    }
    switch(mode){
        case "euclidean":
            {
                const difference = embedding1.map((x, i)=>x - embedding2[i]);
                return -norm(difference);
            }
        case "dot_product":
            {
                let result = 0;
                for(let i = 0; i < embedding1.length; i++){
                    result += embedding1[i] * embedding2[i];
                }
                return result;
            }
        case "cosine":
            {
                return similarity(embedding1, embedding2, "dot_product") / (norm(embedding1) * norm(embedding2));
            }
        default:
            throw new Error("Not implemented yet");
    }
}
/**
 * Get the top K embeddings from a list of embeddings ordered by similarity to the query.
 * @param queryEmbedding
 * @param embeddings list of embeddings to consider
 * @param similarityTopK max number of embeddings to return, default 2
 * @param embeddingIds ids of embeddings in the embeddings list
 * @param similarityCutoff minimum similarity score
 * @returns
 */ // eslint-disable-next-line max-params
export function getTopKEmbeddings(queryEmbedding, embeddings, similarityTopK = DEFAULT_SIMILARITY_TOP_K, embeddingIds = null, similarityCutoff = null) {
    if (embeddingIds == null) {
        embeddingIds = Array(embeddings.length).map((_, i)=>i);
    }
    if (embeddingIds.length !== embeddings.length) {
        throw new Error("getTopKEmbeddings: embeddings and embeddingIds length mismatch");
    }
    const similarities = [];
    for(let i = 0; i < embeddings.length; i++){
        const sim = similarity(queryEmbedding, embeddings[i]);
        if (similarityCutoff == null || sim > similarityCutoff) {
            similarities.push({
                similarity: sim,
                id: embeddingIds[i]
            });
        }
    }
    similarities.sort((a, b)=>b.similarity - a.similarity); // Reverse sort
    const resultSimilarities = [];
    const resultIds = [];
    for(let i = 0; i < similarityTopK; i++){
        if (i >= similarities.length) {
            break;
        }
        resultSimilarities.push(similarities[i].similarity);
        resultIds.push(similarities[i].id);
    }
    return [
        resultSimilarities,
        resultIds
    ];
}
// eslint-disable-next-line max-params
export function getTopKEmbeddingsLearner(queryEmbedding, embeddings, similarityTopK, embeddingsIds, queryMode) {
    throw new Error("Not implemented yet");
}
// eslint-disable-next-line max-params
export function getTopKMMREmbeddings(queryEmbedding, embeddings, similarityFn = null, similarityTopK = null, embeddingIds = null, _similarityCutoff = null, mmrThreshold = null) {
    const threshold = mmrThreshold || 0.5;
    similarityFn = similarityFn || similarity;
    if (embeddingIds === null || embeddingIds.length === 0) {
        embeddingIds = Array.from({
            length: embeddings.length
        }, (_, i)=>i);
    }
    const fullEmbedMap = new Map(embeddingIds.map((value, i)=>[
            value,
            i
        ]));
    const embedMap = new Map(fullEmbedMap);
    const embedSimilarity = new Map();
    let score = Number.NEGATIVE_INFINITY;
    let highScoreId = null;
    for(let i = 0; i < embeddings.length; i++){
        const emb = embeddings[i];
        const similarity = similarityFn(queryEmbedding, emb);
        embedSimilarity.set(embeddingIds[i], similarity);
        if (similarity * threshold > score) {
            highScoreId = embeddingIds[i];
            score = similarity * threshold;
        }
    }
    const results = [];
    const embeddingLength = embeddings.length;
    const similarityTopKCount = similarityTopK || embeddingLength;
    while(results.length < Math.min(similarityTopKCount, embeddingLength)){
        results.push([
            score,
            highScoreId
        ]);
        embedMap.delete(highScoreId);
        const recentEmbeddingId = highScoreId;
        score = Number.NEGATIVE_INFINITY;
        for (const embedId of Array.from(embedMap.keys())){
            const overlapWithRecent = similarityFn(embeddings[embedMap.get(embedId)], embeddings[fullEmbedMap.get(recentEmbeddingId)]);
            if (threshold * embedSimilarity.get(embedId) - (1 - threshold) * overlapWithRecent > score) {
                score = threshold * embedSimilarity.get(embedId) - (1 - threshold) * overlapWithRecent;
                highScoreId = embedId;
            }
        }
    }
    const resultSimilarities = results.map(([s, _])=>s);
    const resultIds = results.map(([_, n])=>n);
    return [
        resultSimilarities,
        resultIds
    ];
}
async function blobToDataUrl(input) {
    const buffer = Buffer.from(await input.arrayBuffer());
    const mimes = filetypemime(buffer);
    if (mimes.length < 1) {
        throw new Error("Unsupported image type");
    }
    return "data:" + mimes[0] + ";base64," + buffer.toString("base64");
}
export async function imageToString(input) {
    if (input instanceof Blob) {
        // if the image is a Blob, convert it to a base64 data URL
        return await blobToDataUrl(input);
    } else if (_.isString(input)) {
        return input;
    } else if (input instanceof URL) {
        return input.toString();
    } else {
        throw new Error(`Unsupported input type: ${typeof input}`);
    }
}
export function stringToImage(input) {
    if (input.startsWith("data:")) {
        // if the input is a base64 data URL, convert it back to a Blob
        const base64Data = input.split(",")[1];
        const byteArray = Buffer.from(base64Data, "base64");
        return new Blob([
            byteArray
        ]);
    } else if (input.startsWith("http://") || input.startsWith("https://")) {
        return new URL(input);
    } else if (_.isString(input)) {
        return input;
    } else {
        throw new Error(`Unsupported input type: ${typeof input}`);
    }
}
export async function imageToDataUrl(input) {
    // first ensure, that the input is a Blob
    if (input instanceof URL && input.protocol === "file:" || _.isString(input)) {
        // string or file URL
        const dataBuffer = await fs.readFile(input instanceof URL ? input.pathname : input);
        input = new Blob([
            dataBuffer
        ]);
    } else if (!(input instanceof Blob)) {
        if (input instanceof URL) {
            throw new Error(`Unsupported URL with protocol: ${input.protocol}`);
        } else {
            throw new Error(`Unsupported input type: ${typeof input}`);
        }
    }
    return await blobToDataUrl(input);
}
