import { pipeline, Tensor } from "@xenova/transformers";

function hamming_distance(arr1, arr2) {
    if (arr1.length !== arr2.length) {
        throw new Error("Typed arrays must have the same length");
    }

    let distance = 0;

    // Iterate over each byte in the typed arrays
    for (let i = 0; i < arr1.length; ++i) {
        // XOR the bytes to find differing bits
        let xorResult = arr1[i] ^ arr2[i];

        // Count set bits in the XOR result using Brian Kernighan's Algorithm
        while (xorResult) {
            ++distance;
            xorResult &= xorResult - 1;
        }
    }

    return distance;
}

function quantize_embeddings(tensor, precision) {
    if (tensor.dims.length !== 2) {
        throw new Error("The tensor must have 2 dimensions");
    }
    if (tensor.dims.at(-1) % 8 !== 0) {
        throw new Error("The last dimension of the tensor must be a multiple of 8");
    }
    if (!['binary', 'ubinary'].includes(precision)) {
        throw new Error("The precision must be either 'binary' or 'ubinary'");
    }
    // Create a typed array to store the packed bits
    const inputData = tensor.data;

    const signed = precision === 'binary';
    const cls = signed ? Int8Array : Uint8Array;
    const dtype = signed ? 'int8' : 'uint8';
    const outputData = new cls(inputData.length / 8);

    // Iterate over each number in the array
    for (let i = 0; i < inputData.length; ++i) {
        // Determine if the number is greater than 0
        const bit = inputData[i] > 0 ? 1 : 0;

        // Calculate the index in the typed array and the position within the byte
        const arrayIndex = Math.floor(i / 8);
        const bitPosition = i % 8;

        // Pack the bit into the typed array
        outputData[arrayIndex] |= bit << (7 - bitPosition);
        if (signed && bitPosition === 0) {
            outputData[arrayIndex] -= 128;
        }
    };

    return new Tensor(dtype, outputData, [tensor.dims[0], tensor.dims[1] / 8]);
}

const embedder = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2', {
    quantized: false,
});

const texts = ['hello', 'hi', 'banana'];
const output = await embedder(texts, {
    normalize: true,
    pooling: 'mean',
});
const embeddings = quantize_embeddings(output, 'ubinary').tolist();

const pairs = [[0, 1], [0, 2], [1, 2]];
for (const [i, j] of pairs) {
    console.log(`${texts[i]} <-> ${texts[j]}`, '|', hamming_distance(embeddings[i], embeddings[j]));
}