# ndarray_factory.js

    // --- factory functions ---
    import { NDArray, DTYPE_MAP } from './ndarray_core.js';

    /**
     * Creates an NDArray from a regular array or TypedArray.
     * @param {Array|TypedArray} source - The source data.
     * @param {string} [dtype='float64'] - The desired data type.
     * @return {NDArray}
     */
    export function array(source, dtype = 'float64') {
        const Ctor = DTYPE_MAP[dtype] || Float64Array;
        const isTyped = (v) => ArrayBuffer.isView(v) && !(v instanceof DataView);

        // 1. Shape Inference & Strict Validation
        const getShape = (v) => {
            if (isTyped(v)) return [v.length];
            if (!Array.isArray(v)) return []; // Scalar
            if (v.length === 0) throw new Error("Input array cannot be empty.");

            const sub = getShape(v[0]);
            const subStr = sub.join(',');

            for (let i = 1; i < v.length; i++) {
                if (getShape(v[i]).join(',') !== subStr) {
                    throw new Error(`Jagged array detected at index ${i}.`);
                }
            }
            return [v.length, ...sub];
        };

        const shape = getShape(source);
        const size = shape.reduce((a, b) => a * b, 1);
        const data = new Ctor(size);

        // 2. Fill Data
        // Special handling for scalar (0-dim) to avoid logic complexity in loop
        if (shape.length === 0) {
            data[0] = source;
        } else {
            let offset = 0;
            const lastDim = shape.length - 1; // The index of the innermost dimension

            const fill = (v, dim) => {
                if (dim === lastDim) {
                    // Optimization: We are at the bottom level.
                    // v is guaranteed to be an Array-like (JS Array or TypedArray)
                    // data.set is the fastest way to copy a block of memory.
                    data.set(v, offset);
                    offset += v.length;
                } else {
                    // We are not at the bottom. v is guaranteed to be an Array.
                    // Recurse deeper.
                    for (let i = 0; i < v.length; i++) {
                        fill(v[i], dim + 1);
                    }
                }
            };
            fill(source, 0);
        }

        return new NDArray(data, { shape, dtype });
    }

    /**
     * Creates an NDArray from a regular array or TypedArray.
     * @param {Array|TypedArray} source - The source data.
     * @return {NDArray}
     */
    export function float64(source){
        return array(source,"float64");
    }
    /**
     * Creates an NDArray from a regular array or TypedArray.
     * @param {Array|TypedArray} source - The source data.
     * @return {NDArray}
     */
    export function float32(source){
        return array(source,"float32");
    }
    /**
     * Creates an NDArray from a regular array or TypedArray.
     * @param {Array|TypedArray} source - The source data.
     * @return {NDArray}
     */
    export function uint32(source){
        return array(source,"uint32");
    }
    /**
     * Creates an NDArray from a regular array or TypedArray.
     * @param {Array|TypedArray} source - The source data.
     * @return {NDArray}
     */
    export function int32(source){
        return array(source,"int32");
    }
    /**
     * Creates an NDArray from a regular array or TypedArray.
     * @param {Array|TypedArray} source - The source data.
     * @return {NDArray}
     */
    export function int16(source){
        return array(source,"int16");
    }
    /**
     * Creates an NDArray from a regular array or TypedArray.
     * @param {Array|TypedArray} source - The source data.
     * @return {NDArray}
     */
    export function uint16(source){
        return array(source,"uint16");
    }
    /**
     * Creates an NDArray from a regular array or TypedArray.
     * @param {Array|TypedArray} source - The source data.
     * @return {NDArray}
     */
    export function int8(source){
        return array(source,"int8");
    }
    /**
     * Creates an NDArray from a regular array or TypedArray.
     * @param {Array|TypedArray} source - The source data.
     * @return {NDArray}
     */
    export function uint8(source){
        return array(source,"uint8");
    }
    /**
     * Creates an NDArray from a regular array or TypedArray.
     * @param {Array|TypedArray} source - The source data.
     * @return {NDArray}
     */
    export function uint8c(source){
        return array(source,"uint8c");
    }

    /**
     * Creates a new NDArray of the given shape, filled with zeros.
     * @param {Array<number>} shape - The shape of the new array.
     * @param {string} [dtype='float64'] - The data type of the new array.
     * @returns {NDArray}
     */
    export function zeros(shape, dtype = 'float64') {
        const size = shape.reduce((a, b) => a * b, 1);
        const Constructor = DTYPE_MAP[dtype];
        return new NDArray(new Constructor(size), { shape, dtype });
    }

    /**
     * Creates a new NDArray of the given shape, filled with ones.
     * @param {Array<number>} shape - The shape of the new array.
     * @param {string} [dtype='float64'] - The data type of the new array.
     * @returns {NDArray}
     */
    export function ones(shape, dtype = 'float64') {
        const arr = zeros(shape, dtype);
        arr.data.fill(1);
        return arr;
    }

    /**
     * Creates a new NDArray of the given shape, filled with a specified value.
     * @param {Array<number>} shape - The shape of the new array.
     * @param {*} value - The value to fill the array with.
     * @param {string} [dtype='float64'] - The data type of the new array.
     * @returns {NDArray}
     */
    export function full(shape, value, dtype = 'float64') {
        const arr = zeros(shape, dtype);
        arr.data.fill(value);
        return arr;
    }

    /**
     * Like Python's arange: [start, stop)
     * @param {number} start the starting value of the sequence.
     * @param {number | null} stop the end value of the sequence.
     * @param {number | null} [step=1] the spacing between values.
     * @param {string} [dtype='float64'] the data type of the resulting array.
     * @returns {NDArray}
     */
    export function arange(start, stop = null, step = 1, dtype = 'float64') {
        if (stop === null) {
            stop = start;
            start = 0;
        }
        const size = Math.ceil((stop - start) / step);
        const Constructor = DTYPE_MAP[dtype];
        const data = new Constructor(size);
        for (let i = 0; i < size; i++) {
            data[i] = start + i * step;
        }
        return new NDArray(data, { shape: [size], dtype });
    }

    /**
     * Linearly spaced points.
     * @param {number} start the starting value of the sequence.
     * @param {number | null} stop the end value of the sequence.
     * @param {number} [num=50] the number of samples to generate.
     * @param {string} [dtype='float64'] the data type of the resulting array.
     * @returns {NDArray}
     */
    export function linspace(start, stop, num = 50, dtype = 'float64') {
        const step = (stop - start) / (num - 1);
        const Constructor = DTYPE_MAP[dtype];
        const data = new Constructor(num);
        for (let i = 0; i < num; i++) {
            data[i] = start + i * step;
        }
        return new NDArray(data, { shape: [num], dtype });
    }

    const toNDArray = (val, defaultDtype = 'float64') => {
        return (val instanceof NDArray) ? val : array(val, defaultDtype);
    };

    /**
     * argwhere(condition)
     *
     * @param {NDArray|Array|number} condition - Input array to check for non-zero elements.
     * @returns {NDArray} A new contiguous NDArray with shape[nz, condition.ndim], where nz is the number of non-zero elements in condition.
     * Each row contains the indices of a non-zero element in condition.
     */
    export function argwhere(condition) {
        const cond = toNDArray(condition, 'int8').asContiguous();
        let nz = 0;

        for (let i = 0; i < cond.size; i++) {
            if (cond.data[cond.offset + i]) {
                nz++;
            }
        }

        let ret = new NDArray(new Int32Array(nz * cond.ndim), {
            shape: [nz, cond.ndim],
            dtype: 'int32'
        });

        let idx = 0;
        for (let i = 0; i < cond.size; i++) {
            if (cond.data[cond.offset + i]) {
                if (cond.ndim > 0) {
                    ret.data[idx * cond.ndim] = Math.floor(i / cond.strides[0]);
                    for (let j = 1; j < cond.ndim; j++) {
                        ret.data[idx * cond.ndim + j] = Math.floor((i % cond.strides[j - 1]) / cond.strides[j]);
                    }
                }
                idx++;
            }
        }

        return ret;
    }

    /**
     * where(condition, x, y)
     * Returns a new array with elements chosen from x or y depending on condition.
     * Supports NumPy-style broadcasting across all three arguments, including 0-sized dimensions.
     *
     * @param {NDArray|Array|number} condition - Where True, yield x, otherwise yield y.
     * @param {NDArray|Array|number} x - Values from which to choose if condition is True.
     * @param {NDArray|Array|number} y - Values from which to choose if condition is False.
     * @returns {NDArray} A new contiguous NDArray.
     */
    export function where(condition, x, y) {
        // 1. Standardize inputs to NDArray
        const toNDArray = (val, defaultDtype = 'float64') => {
            return (val instanceof NDArray) ? val : array(val, defaultDtype);
        };

        const cond = toNDArray(condition, 'int8');

        if( x===undefined || y===undefined){
            return argwhere(cond).transpose(); // shape [cond.ndim, nz]
        }

        const dtypes = [(x instanceof NDArray)?x.dtype:null , (y instanceof NDArray)?y.dtype:null];
        let dtype = null;
        for(let i of dtypes){
            if(dtype === null){
                dtype = i;
            }else if(dtype !== i && i !== null){
                throw new Error(`dtype is different ${dtype} ${i}`);
            }
        }
        if(dtype===null){
            dtype="float64";
        }
        const xArr = toNDArray(x, dtype);
        const yArr = toNDArray(y, dtype);

        // 2. Determine output shape using strict NumPy broadcasting rules
        const outNdim = Math.max(cond.ndim, xArr.ndim, yArr.ndim);
        const outShape = new Int32Array(outNdim);

        for (let i = 1; i <= outNdim; i++) {
            // Dimension sizes from the right (defaulting to 1 for missing dimensions)
            const dC = (cond.ndim - i >= 0) ? cond.shape[cond.ndim - i] : 1;
            const dX = (xArr.ndim - i >= 0) ? xArr.shape[xArr.ndim - i] : 1;
            const dY = (yArr.ndim - i >= 0) ? yArr.shape[yArr.ndim - i] : 1;

            // Broadcasting logic:
            // A dimension is compatible if all values are equal, or if some are 1.
            let maxDim = 1;
            const dims = [dC, dX, dY];
            for (let j = 0; j < 3; j++) {
                const d = dims[j];
                if (d !== 1) {
                    if (maxDim !== 1 && maxDim !== d) {
                        throw new Error(`Incompatible shapes for where: condition [${cond.shape}], x [${xArr.shape}], y [${yArr.shape}]`);
                    }
                    maxDim = d;
                }
            }
            outShape[outNdim - i] = maxDim;
        }

        // 3. Prepare traversal metadata
        // reduce returns 1 for outNdim == 0 (scalars) and 0 if any dim is 0.
        const outSize = outShape.reduce((a, b) => a * b, 1);

        const sC = new Int32Array(outNdim);
        const sX = new Int32Array(outNdim);
        const sY = new Int32Array(outNdim);
        const wrapC = new Float64Array(outNdim);
        const wrapX = new Float64Array(outNdim);
        const wrapY = new Float64Array(outNdim);

        for (let i = 0; i < outNdim; i++) {
            const dim = outShape[i];
            const axisFromRight = outNdim - i;

            const getStride = (arr, axisRight) => {
                const axis = arr.ndim - axisRight;
                // Broadcast (stride 0) if dimension is missing (axis < 0) or size is 1
                if (axis < 0 || arr.shape[axis] === 1) return 0;
                return arr.strides[axis];
            };

            sC[i] = getStride(cond, axisFromRight);
            sX[i] = getStride(xArr, axisFromRight);
            sY[i] = getStride(yArr, axisFromRight);

            // wrapBack is the distance to jump from the last element back to the first in a dimension
            wrapC[i] = (dim - 1) * sC[i];
            wrapX[i] = (dim - 1) * sX[i];
            wrapY[i] = (dim - 1) * sY[i];
        }

        // 4. Execution Engine
        const Constructor = DTYPE_MAP[xArr.dtype] || Float64Array;
        const newData = new Constructor(outSize);
        const currentIdx = new Int32Array(outNdim);

        let pC = cond.offset;
        let pX = xArr.offset;
        let pY = yArr.offset;

        const dataC = cond.data;
        const dataX = xArr.data;
        const dataY = yArr.data;

        // Loop runs outSize times. If outSize is 0, it doesn't run at all.
        // If outSize is 1 (0D scalar), it runs exactly once.
        for (let i = 0; i < outSize; i++) {
            // Core ternary selection
            newData[i] = dataC[pC] ? dataX[pX] : dataY[pY];

            // Odometer increment with manual pointer updates
            for (let d = outNdim - 1; d >= 0; d--) {
                if (++currentIdx[d] < outShape[d]) {
                    pC += sC[d];
                    pX += sX[d];
                    pY += sY[d];
                    break;
                } else {
                    currentIdx[d] = 0;
                    pC -= wrapC[d];
                    pX -= wrapX[d];
                    pY -= wrapY[d];
                }
            }
        }

        return new NDArray(newData, {
            shape: outShape,
            dtype: xArr.dtype
        });
    };

    /**
     * Joins a sequence of arrays along an existing axis.
     * @param {Array<NDArray>} arrays The arrays must have the same shape, except in the dimension corresponding to `axis`.
     * @param {number} [axis=0] The axis along which the arrays will be joined.
     * @returns {NDArray} A new NDArray.
     */
    export function concat(arrays, axis = 0) {
        if (!arrays || arrays.length === 0) {
            throw new Error('Array list cannot be empty.');
        }

        if (arrays.length === 1) {
            return arrays[0].copy(); // .copy() is available in this file.
        }

        const firstArr = arrays[0];
        const { dtype, ndim } = firstArr;

        // --- Validation and Metadata ---
        if (ndim === 0) {
            if (!arrays.every(arr => arr.ndim === 0)) {
                throw new Error('All arrays must be scalars to concatenate with a scalar.');
            }
            const newData = new (DTYPE_MAP[dtype])(arrays.length);
            for (let i = 0; i < arrays.length; i++) {
                newData[i] = arrays[i].data[arrays[i].offset];
            }
            return new NDArray(newData, { shape: [arrays.length], dtype });
        }

        const finalAxis = axis < 0 ? ndim + axis : axis;
        if (finalAxis < 0 || finalAxis >= ndim) {
            throw new Error(`Axis ${axis} is out of bounds for array of dimension ${ndim}.`);
        }

        let newShape = [...firstArr.shape];
        let concatDimSize = 0;

        for (const arr of arrays) {
            if (arr.ndim !== ndim) throw new Error('All arrays must have same number of dimensions.');
            if (arr.dtype !== dtype) throw new Error('All arrays must have same dtype.');
            for (let j = 0; j < ndim; j++) {
                if (j !== finalAxis && arr.shape[j] !== newShape[j]) {
                    throw new Error(`Dimension mismatch on axis ${j}: expected ${newShape[j]} but got ${arr.shape[j]}.`);
                }
            }
            concatDimSize += arr.shape[finalAxis];
        }
        newShape[finalAxis] = concatDimSize;
        // --- End Validation ---

        const result = zeros(newShape, dtype);

        let destAxisOffset = 0;
        for (const arr of arrays) {
            if (arr.size === 0) continue;

            const dimSize = arr.shape[finalAxis];

            const sliceSpec = Array(ndim).fill(null);
            sliceSpec[finalAxis] = [destAxisOffset, destAxisOffset + dimSize];

            const view = result.slice(...sliceSpec);

            view.set(arr);

            destAxisOffset += dimSize;
        }

        return result;
    }

    /**
     * Joins a sequence of arrays along a new axis.
     * The `stack` function creates a new dimension, whereas `concat` joins along an existing one.
     * All input arrays must have the same shape and dtype.
     *
     * @param {Array<NDArray>} arrays - The list of arrays to stack.
     * @param {number} [axis=0] - The axis in the result array along which the input arrays are stacked.
     * @returns {NDArray} A new NDArray.
     */
    export function stack(arrays, axis = 0) {
        if (!arrays || arrays.length === 0) {
            throw new Error('Array list cannot be empty.');
        }

        const firstArr = arrays[0];
        const { shape: oldShape, ndim: oldNdim, dtype } = firstArr;

        // The result of stack will have one more dimension than the input arrays
        const targetNdim = oldNdim + 1;

        // Normalize negative axis (e.g., -1 refers to the last possible axis)
        let finalAxis = axis < 0 ? targetNdim + axis : axis;
        if (finalAxis < 0 || finalAxis >= targetNdim) {
            throw new Error(`Axis ${axis} is out of bounds for stack into ${targetNdim} dimensions.`);
        }

        // Step 1: Validate shapes and expand dimensions
        const expandedArrays = arrays.map((arr, index) => {
            // Ensure dimensionality matches
            if (arr.ndim !== oldNdim) {
                throw new Error(`Array at index ${index} has ${arr.ndim} dimensions, expected ${oldNdim}.`);
            }

            // Ensure dtypes match
            if (arr.dtype !== dtype) {
                throw new Error(`Dtype mismatch at index ${index}. Expected ${dtype}.`);
            }

            // Ensure shapes match exactly
            for (let i = 0; i < oldNdim; i++) {
                if (arr.shape[i] !== oldShape[i]) {
                    throw new Error(
                        `Shape mismatch at index ${index}: expected [${oldShape}] but got [${arr.shape}].`
                    );
                }
            }

            /**
             * Step 2: Reshape the array to add a singleton dimension at the target axis.
             * Example: If shape is [3, 4] and axis is 1, new shape becomes [3, 1, 4].
             */
            const newShape = [...arr.shape];
            newShape.splice(finalAxis, 0, 1);

            // return a new view with the expanded shape
            return arr.reshape(...newShape);
        });

        /**
         * Step 3: Concatenate the expanded arrays along the newly created axis.
         * Since every array now has a size of 1 at 'finalAxis',
         * concat will join them to a size of 'arrays.length' at that axis.
         */
        return concat(expandedArrays, finalAxis);
    };

    /**
     * Creates a 2D identity matrix.
     * @param {number} n - Number of rows.
     * @param {number} [m] - Number of columns. Defaults to n if not provided.
     * @param {string} [dtype='float64'] - Data type of the array.
     * @returns {NDArray} A 2D NDArray with ones on the main diagonal.
     */
    export function eye(n, m, dtype = 'float64') {
        const rows = n;
        const cols = m === undefined ? n : m;
        const shape = [rows, cols];

        // 1. Initialize an all-zero NDArray.
        // This handles data allocation and dtype mapping via the existing zeros method.
        const res = zeros(shape, dtype);

        const offset = 0;
        // 2. Fill the main diagonal with 1.
        // For a contiguous 2D array, the element at (i, i) is at index: i * cols + i.
        const minDim = Math.min(rows, cols);
        for (let i = 0; i < minDim; i++) {
            res.data[offset + i * cols + i] = 1;
        }

        return res;
    }