/**
 * @file Helper module for `Tensor` processing.
 * 
 * These functions and classes are only used internally, 
 * meaning an end-user shouldn't need to access anything here.
 * 
 * @module utils/tensor
 */

import { ONNX } from '../backends/onnx.js';

import {
    interpolate_data,
    permute_data
} from './maths.js';


const DataTypeMap = Object.freeze({
    float32: Float32Array,
    float64: Float64Array,
    string: Array, // string[]
    int8: Int8Array,
    uint8: Uint8Array,
    int16: Int16Array,
    uint16: Uint16Array,
    int32: Int32Array,
    uint32: Uint32Array,
    int64: BigInt64Array,
    uint64: BigUint64Array,
    bool: Uint8Array,
});

/**
 * @typedef {keyof typeof DataTypeMap} DataType
 * @typedef {import('./maths.js').AnyTypedArray | any[]} DataArray
 */

const ONNXTensor = ONNX.Tensor;

export class Tensor {
    /** @type {number[]} Dimensions of the tensor. */
    dims;

    /** @type {DataType} Type of the tensor. */
    type;

    /** @type {DataArray} The data stored in the tensor. */
    data;

    /** @type {number} The number of elements in the tensor. */
    size;

    /**
     * Create a new Tensor or copy an existing Tensor.
     * @param {[DataType, DataArray, number[]]|[import('onnxruntime-common').Tensor]} args
     */
    constructor(...args) {
        if (args[0] instanceof ONNXTensor) {
            // Create shallow copy
            Object.assign(this, args[0]);

        } else {
            // Create new tensor
            Object.assign(this, new ONNXTensor(
                /** @type {DataType} */(args[0]),
                /** @type {Exclude<import('./maths.js').AnyTypedArray, Uint8ClampedArray>} */(args[1]),
                args[2]
            ));
        }

        return new Proxy(this, {
            get: (obj, key) => {
                if (typeof key === 'string') {
                    let index = Number(key);
                    if (Number.isInteger(index)) {
                        // key is an integer (i.e., index)
                        return obj._getitem(index);
                    }
                }
                // @ts-ignore
                return obj[key];
            },
            set: (obj, key, value) => {
                // TODO allow setting of data

                // @ts-ignore
                return obj[key] = value;
            }
        });
    }

    /**
     * Returns an iterator object for iterating over the tensor data in row-major order.
     * If the tensor has more than one dimension, the iterator will yield subarrays.
     * @returns {Iterator} An iterator object for iterating over the tensor data in row-major order.
     */
    *[Symbol.iterator]() {
        const [iterLength, ...iterDims] = this.dims;

        if (iterDims.length > 0) {
            const iterSize = iterDims.reduce((a, b) => a * b);
            for (let i = 0; i < iterLength; ++i) {
                yield this._subarray(i, iterSize, iterDims);
            }
        } else {
            yield* this.data
        }

    }

    /**
     * Index into a Tensor object.
     * @param {number} index The index to access.
     * @returns {Tensor} The data at the specified index.
     */
    _getitem(index) {
        const [iterLength, ...iterDims] = this.dims;

        index = safeIndex(index, iterLength);

        if (iterDims.length > 0) {
            const iterSize = iterDims.reduce((a, b) => a * b);
            return this._subarray(index, iterSize, iterDims);
        } else {
            return new Tensor(this.type, [this.data[index]], iterDims);
        }
    }

    /**
     * @param {number|bigint} item The item to search for in the tensor
     * @returns {number} The index of the first occurrence of item in the tensor data.
     */
    indexOf(item) {
        for (let index = 0; index < this.data.length; ++index) {
            // Note: == instead of === so we can match Ints with BigInts
            if (this.data[index] == item) {
                return index;
            }
        }
        return -1;
    }

    /**
     * @param {number} index 
     * @param {number} iterSize 
     * @param {any} iterDims 
     * @returns {Tensor}
     */
    _subarray(index, iterSize, iterDims) {
        const o1 = index * iterSize;
        const o2 = (index + 1) * iterSize;

        // We use subarray if available (typed array), otherwise we use slice (normal array)
        const data =
            ('subarray' in this.data)
                ? this.data.subarray(o1, o2)
                : this.data.slice(o1, o2);
        return new Tensor(this.type, data, iterDims);
    }

    /**
     * Returns the value of this tensor as a standard JavaScript Number. This only works
     * for tensors with one element. For other cases, see `Tensor.tolist()`.
     * @returns {number|bigint} The value of this tensor as a standard JavaScript Number.
     * @throws {Error} If the tensor has more than one element.
     */
    item() {
        if (this.data.length !== 1) {
            throw new Error(`a Tensor with ${this.data.length} elements cannot be converted to Scalar`);
        }
        return this.data[0];
    }

    /**
     * Convert tensor data to a n-dimensional JS list
     * @returns {Array}
     */
    tolist() {
        return reshape(this.data, this.dims)
    }

    /**
     * Return a new Tensor with the sigmoid function applied to each element.
     * @returns {Tensor} The tensor with the sigmoid function applied.
     */
    sigmoid() {
        return this.clone().sigmoid_();
    }

    /**
     * Applies the sigmoid function to the tensor in place.
     * @returns {Tensor} Returns `this`.
     */
    sigmoid_() {
        for (let i = 0; i < this.data.length; ++i) {
            this.data[i] = 1 / (1 + Math.exp(-this.data[i]));
        }
        return this;
    }

    /**
     * Return a new Tensor with every element multiplied by a constant.
     * @param {number} val The value to multiply by.
     * @returns {Tensor} The new tensor.
     */
    mul(val) {
        return this.clone().mul_(val);
    }

    /**
     * Multiply the tensor by a constant in place.
     * @param {number} val The value to multiply by.
     * @returns {Tensor} Returns `this`.
     */
    mul_(val) {
        for (let i = 0; i < this.data.length; ++i) {
            this.data[i] *= val;
        }
        return this;
    }


    /**
     * Return a new Tensor with every element added by a constant.
     * @param {number} val The value to add by.
     * @returns {Tensor} The new tensor.
     */
    add(val) {
        return this.clone().add_(val);
    }

    /**
     * Add the tensor by a constant in place.
     * @param {number} val The value to add by.
     * @returns {Tensor} Returns `this`.
     */
    add_(val) {
        for (let i = 0; i < this.data.length; ++i) {
            this.data[i] += val;
        }
        return this;
    }
    clone() {
        return new Tensor(this.type, this.data.slice(), this.dims.slice());
    }

    slice(...slices) {
        // This allows for slicing with ranges and numbers
        let newTensorDims = [];
        let newOffsets = [];

        // slices is an array of numbers or arrays of numbers
        // e.g., slices = [0, [1, 3], null, [0, 3]]
        for (let sliceIndex = 0; sliceIndex < this.dims.length; ++sliceIndex) {
            let slice = slices[sliceIndex];

            if (slice === null || slice === undefined) {
                // null or undefined means take the whole dimension
                newOffsets.push([0, this.dims[sliceIndex]]);
                newTensorDims.push(this.dims[sliceIndex]);

            } else if (typeof slice === 'number') {
                slice = safeIndex(slice, this.dims[sliceIndex], sliceIndex);

                // A number means take a single element
                newOffsets.push([slice, slice + 1]);

            } else if (Array.isArray(slice) && slice.length === 2) {
                // An array of length 2 means take a range of elements

                if (slice[0] > slice[1]) {
                    throw new Error(`Invalid slice: ${slice}`);
                }

                let offsets = [
                    Math.max(slice[0], 0),
                    Math.min(slice[1], this.dims[sliceIndex])
                ];

                newOffsets.push(offsets);
                newTensorDims.push(offsets[1] - offsets[0]);

            } else {
                throw new Error(`Invalid slice: ${slice}`);
            }
        }

        let newDims = newOffsets.map(([start, end]) => end - start);
        let newBufferSize = newDims.reduce((a, b) => a * b);

        // Allocate memory
        // @ts-ignore
        let data = new this.data.constructor(newBufferSize);

        // Precompute strides
        const stride = this.stride();

        for (let i = 0; i < newBufferSize; ++i) {
            let originalIndex = 0;
            for (let j = newDims.length - 1, num = i; j >= 0; --j) {
                const size = newDims[j];
                originalIndex += ((num % size) + newOffsets[j][0]) * stride[j];
                num = Math.floor(num / size);
            }
            data[i] = this.data[originalIndex];
        }
        return new Tensor(this.type, data, newTensorDims);

    }

    /**
     * Return a permuted version of this Tensor, according to the provided dimensions.
     * @param  {...number} dims Dimensions to permute.
     * @returns {Tensor} The permuted tensor.
     */
    permute(...dims) {
        return permute(this, dims);
    }

    // TODO: implement transpose. For now (backwards compatibility), it's just an alias for permute()
    transpose(...dims) {
        return this.permute(...dims);
    }

    // TODO add .max() and .min() methods

    /**
     * Returns the sum of each row of the input tensor in the given dimension dim.
     * 
     * @param {number} [dim=null] The dimension or dimensions to reduce. If `null`, all dimensions are reduced.
     * @param {boolean} keepdim Whether the output tensor has `dim` retained or not.
     * @returns The summed tensor
     */
    sum(dim = null, keepdim = false) {
        return this.norm(1, dim, keepdim);
    }

    /**
     * Returns the matrix norm or vector norm of a given tensor.
     * @param {number|string} [p='fro'] The order of norm
     * @param {number} [dim=null] Specifies which dimension of the tensor to calculate the norm across.
     * If dim is None, the norm will be calculated across all dimensions of input.
     * @param {boolean} [keepdim=false] Whether the output tensors have dim retained or not.
     * @returns {Tensor} The norm of the tensor.
     */
    norm(p = 'fro', dim = null, keepdim = false) {
        if (p === 'fro') {
            // NOTE: Since we only support integer dims, Frobenius norm produces the same result as p=2.
            p = 2;
        } else if (typeof p === 'string') {
            throw Error(`Unsupported norm: ${p}`);
        }

        if (dim === null) {
            // @ts-ignore
            let val = this.data.reduce((a, b) => a + (b ** p), 0) ** (1 / p);
            return new Tensor(this.type, [val], []);
        }

        // Negative indexing
        dim = safeIndex(dim, this.dims.length);

        // Calculate the shape of the resulting array after summation
        const resultDims = this.dims.slice(); // Copy the original dimensions
        resultDims[dim] = 1; // Remove the specified axis

        // Create a new array to store the accumulated values
        // @ts-ignore
        const result = new this.data.constructor(this.data.length / this.dims[dim]);

        // Iterate over the data array
        for (let i = 0; i < this.data.length; ++i) {

            // Calculate the index in the resulting array
            let resultIndex = 0;

            for (let j = this.dims.length - 1, num = i, resultMultiplier = 1; j >= 0; --j) {
                const size = this.dims[j];
                if (j !== dim) {
                    const index = num % size;
                    resultIndex += index * resultMultiplier;
                    resultMultiplier *= resultDims[j];
                }
                num = Math.floor(num / size);
            }

            // Accumulate the value at the current index
            result[resultIndex] += (this.data[i]) ** p;
        }

        if (p !== 1) {
            for (let i = 0; i < result.length; ++i) {
                result[i] = result[i] ** (1 / p);
            }
        }

        if (!keepdim) {
            resultDims.splice(dim, 1);
        }

        return new Tensor(this.type, result, resultDims);
    }

    /**
     * Performs `L_p` normalization of inputs over specified dimension. Operates in place.
     * @param {number} [p=2] The exponent value in the norm formulation
     * @param {number} [dim=1] The dimension to reduce
     * @returns {Tensor} `this` for operation chaining.
     */
    normalize_(p = 2.0, dim = 1) {
        dim = safeIndex(dim, this.dims.length);

        const norm = this.norm(p, dim, true);

        for (let i = 0; i < this.data.length; ++i) {

            // Calculate the index in the resulting array
            let resultIndex = 0;

            for (let j = this.dims.length - 1, num = i, resultMultiplier = 1; j >= 0; --j) {
                const size = this.dims[j];
                if (j !== dim) {
                    const index = num % size;
                    resultIndex += index * resultMultiplier;
                    resultMultiplier *= this.dims[j];
                }
                num = Math.floor(num / size);
            }

            // Divide by normalized value
            this.data[i] /= norm.data[resultIndex];
        }

        return this;
    }

    /**
     * Performs `L_p` normalization of inputs over specified dimension.
     * @param {number} [p=2] The exponent value in the norm formulation
     * @param {number} [dim=1] The dimension to reduce
     * @returns {Tensor} The normalized tensor.
     */
    normalize(p = 2.0, dim = 1) {
        return this.clone().normalize_(p, dim);
    }

    /**
     * Compute and return the stride of this tensor.
     * Stride is the jump necessary to go from one element to the next one in the specified dimension dim.
     * @returns {number[]} The stride of this tensor.
     */
    stride() {
        return dimsToStride(this.dims);
    }

    /**
     * Returns a tensor with all specified dimensions of input of size 1 removed.
     * 
     * NOTE: The returned tensor shares the storage with the input tensor, so changing the contents of one will change the contents of the other.
     * If you would like a copy, use `tensor.clone()` before squeezing.
     * 
     * @param {number} [dim=null] If given, the input will be squeezed only in the specified dimensions.
     * @returns The squeezed tensor
     */
    squeeze(dim = null) {
        return new Tensor(
            this.type,
            this.data,
            calc_squeeze_dims(this.dims, dim)
        )
    }

    /**
     * In-place version of @see {@link Tensor.squeeze}
     */
    squeeze_(dim = null) {
        this.dims = calc_squeeze_dims(this.dims, dim);
        return this;
    }

    /**
     * Returns a new tensor with a dimension of size one inserted at the specified position.
     * 
     * NOTE: The returned tensor shares the same underlying data with this tensor.
     * 
     * @param {number} dim The index at which to insert the singleton dimension
     * @returns The unsqueezed tensor
     */
    unsqueeze(dim = null) {
        return new Tensor(
            this.type,
            this.data,
            calc_unsqueeze_dims(this.dims, dim)
        );
    }

    /**
     * In-place version of @see {@link Tensor.unsqueeze}
     */
    unsqueeze_(dim = null) {
        this.dims = calc_unsqueeze_dims(this.dims, dim);
        return this;
    }

    /**
     * In-place version of @see {@link Tensor.flatten}
     */
    flatten_(start_dim = 0, end_dim = -1) {
        // TODO validate inputs
        end_dim = (end_dim + this.dims.length) % this.dims.length;

        let dimsToKeepBefore = this.dims.slice(0, start_dim);
        let dimsToFlatten = this.dims.slice(start_dim, end_dim + 1);
        let dimsToKeepAfter = this.dims.slice(end_dim + 1);

        this.dims = [...dimsToKeepBefore, dimsToFlatten.reduce((a, b) => a * b, 1), ...dimsToKeepAfter]
        return this;
    }

    /**
     * Flattens input by reshaping it into a one-dimensional tensor.
     * If `start_dim` or `end_dim` are passed, only dimensions starting with `start_dim`
     * and ending with `end_dim` are flattened. The order of elements in input is unchanged.
     * @param {number} start_dim the first dim to flatten
     * @param {number} end_dim the last dim to flatten
     * @returns The flattened tensor.
     */
    flatten(start_dim = 0, end_dim = -1) {
        return this.clone().flatten_(start_dim, end_dim);
    }

    /**
     * Returns a new tensor with the same data as the `self` tensor but of a different `shape`.
     * @param  {...number} dims the desired size
     * @returns {Tensor} The tensor with the same data but different shape
     */
    view(...dims) {
        // TODO: validate dims
        let inferredIndex = -1;
        for (let i = 0; i < dims.length; ++i) {
            if (dims[i] === -1) {
                if (inferredIndex !== -1) {
                    throw new Error("Only one dimension can be inferred");
                }
                inferredIndex = i;
            }
        }

        if (inferredIndex !== -1) {
            // Some dimension must be inferred
            const productOther = dims.reduce((product, curr, index) => {
                return index !== inferredIndex ? product * curr : product
            }, 1);

            dims[inferredIndex] = this.data.length / productOther;
        }
        return new Tensor(this.type, this.data, dims); // NOTE: uses same underlying storage
    }

    neg_() {
        for (let i = 0; i < this.data.length; ++i) {
            this.data[i] = -this.data[i];
        }
        return this;
    }
    neg() {
        return this.clone().neg_();
    }

    /**
     * In-place version of @see {@link Tensor.clamp}
     */
    clamp_(min, max) {
        for (let i = 0; i < this.data.length; ++i) {
            this.data[i] = Math.min(Math.max(this.data[i], min), max);
        }
        return this;
    }

    /**
     * Clamps all elements in input into the range [ min, max ]
     * @param {number} min lower-bound of the range to be clamped to
     * @param {number} max upper-bound of the range to be clamped to
     * @returns the output tensor.
     */
    clamp(min, max) {
        return this.clone().clamp_(min, max);
    }

    /**
     * In-place version of @see {@link Tensor.round}
     */
    round_() {
        for (let i = 0; i < this.data.length; ++i) {
            this.data[i] = Math.round(this.data[i]);
        }
        return this;
    }

    /**
     * Rounds elements of input to the nearest integer.
     * @returns the output tensor.
     */
    round() {
        return this.clone().round_();
    }

    /**
     * Performs Tensor dtype conversion.
     * @param {DataType} type The desired data type.
     * @returns {Tensor} The converted tensor.
     */
    to(type) {
        // If the self Tensor already has the correct dtype, then self is returned.
        if (this.type === type) return this;

        // Otherwise, the returned tensor is a copy of self with the desired dtype.
        if (!DataTypeMap.hasOwnProperty(type)) {
            throw new Error(`Unsupported type: ${type}`);
        }
        // @ts-ignore
        return new Tensor(type, DataTypeMap[type].from(this.data), this.dims);
    }
}

/**
 * This creates a nested array of a given type and depth (see examples).
 * 
 * @example
 *   NestArray<string, 1>; // string[]
 * @example
 *   NestArray<number, 2>; // number[][]
 * @example
 *   NestArray<string, 3>; // string[][][] etc.
 * @template T
 * @template {number} Depth
 * @template {never[]} [Acc=[]]
 * @typedef {Acc['length'] extends Depth ? T : NestArray<T[], Depth, [...Acc, never]>} NestArray
 */

/**
 * Reshapes a 1-dimensional array into an n-dimensional array, according to the provided dimensions.
 *
 * @example
 *   reshape([10                    ], [1      ]); // Type: number[]      Value: [10]
 *   reshape([1, 2, 3, 4            ], [2, 2   ]); // Type: number[][]    Value: [[1, 2], [3, 4]]
 *   reshape([1, 2, 3, 4, 5, 6, 7, 8], [2, 2, 2]); // Type: number[][][]  Value: [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]
 *   reshape([1, 2, 3, 4, 5, 6, 7, 8], [4, 2   ]); // Type: number[][]    Value: [[1, 2], [3, 4], [5, 6], [7, 8]]
 * @param {T[]|DataArray} data The input array to reshape.
 * @param {DIM} dimensions The target shape/dimensions.
 * @template T
 * @template {[number]|number[]} DIM
 * @returns {NestArray<T, DIM["length"]>} The reshaped array.
 */
function reshape(data, dimensions) {

    const totalElements = data.length;
    const dimensionSize = dimensions.reduce((a, b) => a * b);

    if (totalElements !== dimensionSize) {
        throw Error(`cannot reshape array of size ${totalElements} into shape (${dimensions})`);
    }

    /** @type {any} */
    let reshapedArray = data;

    for (let i = dimensions.length - 1; i >= 0; i--) {
        reshapedArray = reshapedArray.reduce((acc, val) => {
            let lastArray = acc[acc.length - 1];

            if (lastArray.length < dimensions[i]) {
                lastArray.push(val);
            } else {
                acc.push([val]);
            }

            return acc;
        }, [[]]);
    }

    return reshapedArray[0];
}

/**
 * Permutes a tensor according to the provided axes.
 * @param {any} tensor The input tensor to permute.
 * @param {Array} axes The axes to permute the tensor along.
 * @returns {Tensor} The permuted tensor.
 */
export function permute(tensor, axes) {
    const [permutedData, shape] = permute_data(tensor.data, tensor.dims, axes);
    return new Tensor(tensor.type, permutedData, shape);
}


/**
 * Interpolates an Tensor to the given size.
 * @param {Tensor} input The input tensor to interpolate. Data must be channel-first (i.e., [c, h, w])
 * @param {number[]} size The output size of the image
 * @param {string} mode The interpolation mode
 * @param {boolean} align_corners Whether to align corners.
 * @returns {Tensor} The interpolated tensor.
 */
export function interpolate(input, [out_height, out_width], mode = 'bilinear', align_corners = false) {

    // Input image dimensions
    const in_channels = input.dims.at(-3) ?? 1;
    const in_height = input.dims.at(-2);
    const in_width = input.dims.at(-1);

    let output = interpolate_data(
        /** @type {import('./maths.js').TypedArray}*/(input.data),
        [in_channels, in_height, in_width],
        [out_height, out_width],
        mode,
        align_corners
    );
    return new Tensor(input.type, output, [in_channels, out_height, out_width]);
}

/**
 * Perform mean pooling of the last hidden state followed by a normalization step.
 * @param {Tensor} last_hidden_state Tensor of shape [batchSize, seqLength, embedDim]
 * @param {Tensor} attention_mask Tensor of shape [batchSize, seqLength]
 * @returns {Tensor} Returns a new Tensor of shape [batchSize, embedDim].
 */
export function mean_pooling(last_hidden_state, attention_mask) {
    // last_hidden_state: [batchSize, seqLength, embedDim]
    // attention_mask:    [batchSize, seqLength]

    let shape = [last_hidden_state.dims[0], last_hidden_state.dims[2]];
    // @ts-ignore
    let returnedData = new last_hidden_state.data.constructor(shape[0] * shape[1]);
    let [batchSize, seqLength, embedDim] = last_hidden_state.dims;

    let outIndex = 0;
    for (let i = 0; i < batchSize; ++i) {
        let offset = i * embedDim * seqLength;

        for (let k = 0; k < embedDim; ++k) {
            let sum = 0;
            let count = 0;

            let attnMaskOffset = i * seqLength;
            let offset2 = offset + k;
            // Pool over all words in sequence
            for (let j = 0; j < seqLength; ++j) {
                // index into attention mask
                let attn = Number(attention_mask.data[attnMaskOffset + j]);

                count += attn;
                sum += last_hidden_state.data[offset2 + j * embedDim] * attn;
            }

            let avg = sum / count;
            returnedData[outIndex++] = avg;
        }
    }

    return new Tensor(
        last_hidden_state.type,
        returnedData,
        shape
    )
}

/**
 * Apply Layer Normalization for last certain number of dimensions.
 * @param {Tensor} input The input tensor
 * @param {number[]} normalized_shape input shape from an expected input of size
 * @param {Object} options The options for the layer normalization
 * @param {number} [options.eps=1e-5] A value added to the denominator for numerical stability.
 * @returns {Tensor} The normalized tensor.
 */
export function layer_norm(input, normalized_shape, {
    eps = 1e-5,
} = {}) {
    if (input.dims.length !== 2) {
        throw new Error('`layer_norm` currently only supports 2D input.');
    }

    const [batchSize, featureDim] = input.dims;

    if (normalized_shape.length !== 1 && normalized_shape[0] !== featureDim) {
        throw new Error('`normalized_shape` must be a 1D array with shape `[input.dims[1]]`.');
    }

    const [std, mean] = std_mean(input, 1, 0, true);

    // @ts-ignore
    const returnedData = new input.data.constructor(input.data.length);

    for (let i = 0; i < batchSize; ++i) {
        const offset = i * featureDim;
        for (let j = 0; j < featureDim; ++j) {
            const offset2 = offset + j;
            returnedData[offset2] = (input.data[offset2] - mean.data[i]) / (std.data[i] + eps);
        }
    }
    return new Tensor(input.type, returnedData, input.dims);
}

/**
 * Helper function to calculate new dimensions when performing a squeeze operation.
 * @param {number[]} dims The dimensions of the tensor.
 * @param {number|number[]|null} dim The dimension(s) to squeeze.
 * @returns The new dimensions.
 * @private
 */
function calc_squeeze_dims(dims, dim) {
    dims = dims.slice();
    if (dim === null) {
        dims = dims.filter((d) => d !== 1);
    } else if (typeof dim === 'number') {
        if (dims[dim] === 1) {
            dims.splice(dim, 1);
        }
    } else if (Array.isArray(dim)) {
        dims = dims.filter((x, i) => {
            return x !== 1 || !dim.includes(i);
        });
    }
    return dims;
}

/**
 * Helper function to calculate new dimensions when performing an unsqueeze operation.
 * @param {number[]} dims The dimensions of the tensor.
 * @param {number} dim The dimension to unsqueeze.
 * @returns The new dimensions.
 * @private
 */
function calc_unsqueeze_dims(dims, dim) {
    // Dimension out of range (e.g., "expected to be in range of [-4, 3], but got 4")
    // + 1 since we allow inserting at the end (i.e. dim = -1)
    dim = safeIndex(dim, dims.length + 1);
    dims = dims.slice();
    // Insert 1 into specified dimension
    dims.splice(dim, 0, 1);
    return dims;
}

/**
 * Safely calculate the index for an array of a given size, allowing negative indexing.
 * @param {number} index The index that will be used.
 * @param {number} size The size of the array.
 * @param {number} [dimension=null] The dimension that the index is for (optional).
 * @returns {number} The index, guaranteed to be non-negative and less than `arrayLength`.
 * 
 * @throws {Error} If the index is out of range.
 * @private
 */
function safeIndex(index, size, dimension = null) {
    if (index < -size || index >= size) {
        throw new Error(`IndexError: index ${index} is out of bounds for dimension${dimension === null ? '' : ' ' + dimension} with size ${size}`);
    }

    if (index < 0) {
        // Negative indexing, ensuring positive index
        index = ((index % size) + size) % size;
    }
    return index;
}

/**
 * Concatenates an array of tensors along a specified dimension.
 * @param {Tensor[]} tensors The array of tensors to concatenate.
 * @param {number} dim The dimension to concatenate along.
 * @returns {Tensor} The concatenated tensor.
 */
export function cat(tensors, dim = 0) {
    dim = safeIndex(dim, tensors[0].dims.length);

    // TODO do validation of shapes

    const resultDims = tensors[0].dims.slice();
    resultDims[dim] = tensors.reduce((a, b) => a + b.dims[dim], 0);

    // Create a new array to store the accumulated values
    const resultSize = resultDims.reduce((a, b) => a * b, 1);
    // @ts-ignore
    const result = new tensors[0].data.constructor(resultSize);

    // Create output tensor of same type as first
    const resultType = tensors[0].type;

    if (dim === 0) {
        // Handle special case for performance reasons

        let offset = 0;
        for (let t of tensors) {
            result.set(t.data, offset);
            offset += t.data.length;
        }

    } else {

        let currentDim = 0;

        for (let t = 0; t < tensors.length; ++t) {
            let tensor = tensors[t];

            // Iterate over the data array
            for (let i = 0; i < tensor.data.length; ++i) {
                // Calculate the index in the resulting array
                let resultIndex = 0;

                for (let j = tensor.dims.length - 1, num = i, resultMultiplier = 1; j >= 0; --j) {
                    const size = tensor.dims[j];
                    let index = num % size;
                    if (j === dim) {
                        index += currentDim;
                    }
                    resultIndex += index * resultMultiplier;
                    resultMultiplier *= resultDims[j];
                    num = Math.floor(num / size);
                }
                // Accumulate the value at the current index
                result[resultIndex] = tensor.data[i];
            }

            currentDim += tensor.dims[dim];
        }
    }
    return new Tensor(resultType, result, resultDims);
}

/**
 * Stack an array of tensors along a specified dimension.
 * @param {Tensor[]} tensors The array of tensors to stack.
 * @param {number} dim The dimension to stack along.
 * @returns {Tensor} The stacked tensor.
 */
export function stack(tensors, dim = 0) {
    // TODO do validation of shapes
    // NOTE: stack expects each tensor to be equal size
    return cat(tensors.map(t => t.unsqueeze(dim)), dim);
}


/**
 * Calculates the standard deviation and mean over the dimensions specified by dim. dim can be a single dimension or `null` to reduce over all dimensions.
 * @param {Tensor} input the input tenso
 * @param {number|null} dim the dimension to reduce. If None, all dimensions are reduced.
 * @param {number} correction difference between the sample size and sample degrees of freedom. Defaults to Bessel's correction, correction=1.
 * @param {boolean} keepdim whether the output tensor has dim retained or not.
 * @returns {Tensor[]} A tuple of (std, mean) tensors.
 */
export function std_mean(input, dim = null, correction = 1, keepdim = false) {

    if (dim === null) {
        // None to reduce over all dimensions.
        // @ts-ignore
        const sum = input.data.reduce((a, b) => a + b, 0);
        const mean = sum / input.data.length;
        // @ts-ignore
        const std = Math.sqrt(input.data.reduce((a, b) => a + (b - mean) ** 2, 0) / (input.data.length - correction));

        const meanTensor = new Tensor(input.type, [mean], [/* scalar */]);
        const stdTensor = new Tensor(input.type, [std], [/* scalar */]);

        return [stdTensor, meanTensor];
    }

    // Negative indexing
    dim = safeIndex(dim, input.dims.length);

    const meanTensor = mean(input, dim, keepdim);

    // Calculate the shape of the resulting array after summation
    const resultDims = input.dims.slice(); // Copy the original dimensions
    resultDims[dim] = 1; // Remove the specified axis

    // Create a new array to store the accumulated values
    // @ts-ignore
    const result = new input.data.constructor(input.data.length / input.dims[dim]);

    // Iterate over the data array
    for (let i = 0; i < input.data.length; ++i) {

        // Calculate the index in the resulting array
        let resultIndex = 0;

        for (let j = input.dims.length - 1, num = i, resultMultiplier = 1; j >= 0; --j) {
            const size = input.dims[j];
            if (j !== dim) {
                const index = num % size;
                resultIndex += index * resultMultiplier;
                resultMultiplier *= resultDims[j];
            }
            num = Math.floor(num / size);
        }

        // Accumulate the value at the current index
        result[resultIndex] += (input.data[i] - meanTensor.data[resultIndex]) ** 2;
    }

    for (let i = 0; i < result.length; ++i) {
        result[i] = Math.sqrt(result[i] / (input.dims[dim] - correction));
    }

    if (!keepdim) {
        resultDims.splice(dim, 1);
    }

    const stdTensor = new Tensor(input.type, result, resultDims);

    return [stdTensor, meanTensor];
}


/**
 * Returns the mean value of each row of the input tensor in the given dimension dim.
 * @param {Tensor} input the input tensor.
 * @param {number|null} dim the dimension to reduce.
 * @param {boolean} keepdim whether the output tensor has dim retained or not.
 * @returns A new tensor with means taken along the specified dimension.
 */
export function mean(input, dim = null, keepdim = false) {

    if (dim === null) {
        // None to reduce over all dimensions.
        // @ts-ignore
        let val = input.data.reduce((a, b) => a + b, 0);
        return new Tensor(input.type, [val / input.data.length], [/* scalar */]);
    }

    // Negative indexing
    dim = safeIndex(dim, input.dims.length);

    // Calculate the shape of the resulting array after summation
    const resultDims = input.dims.slice(); // Copy the original dimensions
    resultDims[dim] = 1; // Remove the specified axis

    // Create a new array to store the accumulated values
    // @ts-ignore
    const result = new input.data.constructor(input.data.length / input.dims[dim]);

    // Iterate over the data array
    for (let i = 0; i < input.data.length; ++i) {

        // Calculate the index in the resulting array
        let resultIndex = 0;

        for (let j = input.dims.length - 1, num = i, resultMultiplier = 1; j >= 0; --j) {
            const size = input.dims[j];
            if (j !== dim) {
                const index = num % size;
                resultIndex += index * resultMultiplier;
                resultMultiplier *= resultDims[j];
            }
            num = Math.floor(num / size);
        }

        // Accumulate the value at the current index
        result[resultIndex] += input.data[i];
    }

    if (input.dims[dim] !== 1) {
        for (let i = 0; i < result.length; ++i) {
            result[i] = result[i] / input.dims[dim];
        }
    }

    if (!keepdim) {
        resultDims.splice(dim, 1);
    }

    return new Tensor(input.type, result, resultDims);
}


/**
 *
 * Measures similarity between two temporal sequences (e.g., input audio and output tokens
 * to generate token-level timestamps).
 * @param {Tensor} matrix 
 * @returns {number[][]}
 */
export function dynamicTimeWarping(matrix) {
    const [output_length, input_length] = matrix.dims;

    const outputShape = [output_length + 1, input_length + 1];

    const cost = new Tensor(
        'float32',
        new Float32Array(outputShape[0] * outputShape[1]).fill(Infinity),
        outputShape
    );

    const trace = new Tensor(
        'float32',
        new Float32Array(outputShape[0] * outputShape[1]).fill(-1),
        outputShape
    )

    // same as `cost[0][0] = 0`;
    cost[0].data[0] = 0;

    for (let j = 1; j < input_length + 1; ++j) {
        for (let i = 1; i < output_length + 1; ++i) {

            const c0 = cost[i - 1][j - 1].item();
            const c1 = cost[i - 1][j].item();
            const c2 = cost[i][j - 1].item();

            let c, t;
            if (c0 < c1 && c0 < c2) {
                c = c0;
                t = 0;
            } else if (c1 < c0 && c1 < c2) {
                c = c1;
                t = 1;
            } else {
                c = c2;
                t = 2;
            }

            cost[i].data[j] = matrix[i - 1][j - 1].item() + c;
            trace[i].data[j] = t;
        }
    }

    // backtrace
    let i = output_length;
    let j = input_length;

    // @ts-ignore
    trace.data.fill(2, 0, outputShape[1]) // trace[0, :] = 2
    for (let i = 0; i < outputShape[0]; ++i) { // trace[:, 0] = 1
        trace[i].data[0] = 1;
    }

    let text_indices = [];
    let time_indices = [];

    while (i > 0 || j > 0) {
        text_indices.push(i - 1);
        time_indices.push(j - 1);

        const t = trace[i][j].item();
        switch (t) {
            case 0:
                --i; --j;
                break;
            case 1:
                --i;
                break;
            case 2:
                --j;
                break;
            default:
                throw new Error(
                    `Internal error in dynamic time warping. Unexpected trace[${i}, ${j}]. Please file a bug report.`
                )
        }
    }

    text_indices.reverse();
    time_indices.reverse();

    return [text_indices, time_indices];

}

function dimsToStride(dims) {
    const stride = new Array(dims.length);
    for (let i = dims.length - 1, s2 = 1; i >= 0; --i) {
        stride[i] = s2;
        s2 *= dims[i];
    }
    return stride;
}

/**
 * Returns a tensor filled with the scalar value 1, with the shape defined by the variable argument size.
 * @param {number[]} size A sequence of integers defining the shape of the output tensor.
 */
export function ones(size) {
    const numElements = size.reduce((a, b) => a * b, 1);
    return new Tensor(
        'int64',
        new BigInt64Array(numElements).fill(1n),
        size
    )
}

/**
 * Returns a tensor filled with the scalar value 1, with the same size as input.
 * @param {Tensor} tensor The size of input will determine size of the output tensor.
 * @returns The ones tensor.
 */
export function ones_like(tensor) {
    return ones(tensor.dims);
}

/**
 * Quantizes the embeddings tensor to binary or unsigned binary precision.
 * @param {Tensor} tensor The tensor to quantize.
 * @param {'binary'|'ubinary'} precision The precision to use for quantization.
 * @returns {Tensor} The quantized tensor.
 */
export 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'");
    }

    const signed = precision === 'binary';
    const dtype = signed ? 'int8' : 'uint8';

    // Create a typed array to store the packed bits
    const cls = signed ? Int8Array : Uint8Array;
    const inputData = tensor.data;
    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]);
}
