export var Jit = { debug: false };

/**
 * Generates a JIT-compiled unary kernel.
 * Output is always contiguous, so ptrOut simply increments by 1.
 * @private
 */
export function _createUnaryKernel(cacheKey, shape, sIn, fnOrStr) {
    let kernel = UNARY_KERNEL_CACHE.get(cacheKey);
    if (kernel) {
        return kernel;
    }
    
    const ndim = shape.length;
    const opBody = prepareUnaryOp(fnOrStr);

    let fnSource;

    // Handle 0-dimensional arrays (scalars)
    if (ndim === 0) {
        fnSource = `
            return function(dataIn, dataOut, offIn, offOut) {
                "use strict";
                dataOut[offOut] = ${opBody.replace('ptrIn', 'offIn')};
            }
        `;
    }else{

        // --- Build Nested Loops ---
        
        // 1. Innermost loop: Apply operation and step the input pointer.
        // Since the output of copy/map is a fresh contiguous array, 
        // ptrOut always increments by exactly 1.
        let code = `
            dataOut[ptrOut++] = ${opBody};
            ptrIn += ${sIn[ndim - 1]};
        `;
        code = `for (let i${ndim - 1} = 0; i${ndim - 1} < ${shape[ndim - 1]}; i${ndim - 1}++) { ${code} }`;

        // 2. Outer loops with "Gap" adjustment for the input pointer.
        for (let d = ndim - 2; d >= 0; d--) {
            // Adjustment: Move to next dimension start - distance covered by inner loops
            const adjIn = sIn[d] - (shape[d + 1] * sIn[d + 1]);
            code += ` ptrIn += ${adjIn};`;
            code = `for (let i${d} = 0; i${d} < ${shape[d]}; i${d}++) { ${code} }`;
        }

        // 3. Assemble final function
        fnSource = `
            return function(dataIn, dataOut, offIn, offOut) {
                "use strict";
                let ptrIn = offIn, ptrOut = offOut;
                ${code}
            };
        `;
    }
    if(Jit.debug) {
        console.log(cacheKey, "Unary Kernel Source:\n", fnSource);
    }
    kernel = new Function(fnSource)();
    UNARY_KERNEL_CACHE.set(cacheKey, kernel);
    return kernel;
}


/**
 * Generates a kernel with unrolled nested loops and cumulative pointer increments.
 * @private
 */
export function _createBinKernel(cacheKey, shape, sA, sB, sOut, opStr, isComparison) {
    let kernel = BIN_KERNEL_CACHE.get(cacheKey);
    if (kernel) {
        return kernel;
    }
    
    const ndim = shape.length;
    let body = extractOpBody(opStr);
    
    // Inline the data access into the operation body
    body = body.replace(/\ba|x\b/g, 'dataA[ptrA]').replace(/\bb|y\b/g, 'dataB[ptrB]');
    if (isComparison) body = `(${body}) ? 1 : 0`;

    let fnSource;
    // Handle 0-dimensional arrays (scalars)
    if (ndim === 0) {
        fnSource = `
            return function(dataA, dataB, dataOut, offA, offB, offOut) {
                "use strict";
                let ptrA = offA, ptrB = offB, ptrOut = offOut;
                dataOut[ptrOut] = ${body};
            }
        `;
    }else{

        // --- Build the Nested Loops String ---
        
        // 1. Innermost Loop logic: Apply operation and increment pointers by the last dimension's stride.
        let code = `
            dataOut[ptrOut] = ${body};
            ptrA += ${sA[ndim - 1]};
            ptrB += ${sB[ndim - 1]};
            ptrOut += ${sOut[ndim - 1]};
        `;
        code = `for (let i${ndim - 1} = 0; i${ndim - 1} < ${shape[ndim - 1]}; i${ndim - 1}++) { ${code} }`;

        // 2. Outer Loops logic: Wrap inner loops and apply pointer adjustments (gaps).
        // The adjustment ensures that after an inner loop finishes, the pointer moves 
        // to the correct start of the next index in the parent dimension.
        for (let d = ndim - 2; d >= 0; d--) {
            // Gap = Stride of current dimension - distance already moved by the inner dimension
            const adjA = sA[d] - (shape[d + 1] * sA[d + 1]);
            const adjB = sB[d] - (shape[d + 1] * sB[d + 1]);
            const adjOut = sOut[d] - (shape[d + 1] * sOut[d + 1]);
            
            code += ` ptrA += ${adjA}; ptrB += ${adjB}; ptrOut += ${adjOut};`;
            code = `for (let i${d} = 0; i${d} < ${shape[d]}; i${d}++) { ${code} }`;
        }

        // 3. Construct final function
        fnSource = `
            return function(dataA, dataB, dataOut, offA, offB, offOut) {
                "use strict";
                let ptrA = offA, ptrB = offB, ptrOut = offOut;
                ${code}
            };
        `;
    }

    if(Jit.debug) {
        console.log(cacheKey, "Bin Kernel Source:\n", fnSource);
    }
    kernel = new Function(fnSource)();
    BIN_KERNEL_CACHE.set(cacheKey, kernel);
    return kernel;
}



/**
     * Generates a JIT kernel using nested loops and pointer gaps.
     * @private
     */
export function _createReduceKernel(cacheKey, shape, strides, iterAxes, reduceAxes, reducer, finalFn) {
    let kernel = REDUCE_KERNEL_CACHE.get(cacheKey);
    if (kernel) {
        return kernel;
    }
    const nRed = reduceAxes.length;
    const nIter = iterAxes.length;

    const redExpr = prepareReduceExpr(reducer, 'reducer');
    const finalExpr = finalFn ? prepareReduceExpr(finalFn, 'finalizer') : 'acc';

    // --- 1. Build Reduction Space (Inner Loops) ---
    // The innermost operation: accumulate and step by the last reduction axis stride
    let redCode = `
        acc = ${redExpr};
        pIn += ${strides[reduceAxes[nRed - 1]]};
    `;
    
    // Nest reduction loops if multiple axes are being collapsed
    for (let d = nRed - 1; d >= 0; d--) {
        const ax = reduceAxes[d];
        // Gap to jump to the next index of the parent reduction dimension
        const gap = (d === 0) ? 0 : (strides[reduceAxes[d - 1]] - shape[ax] * strides[ax]);
        redCode = `
            for (let r${d} = 0; r${d} < ${shape[ax]}; r${d}++) {
                ${redCode}
            }
            pIn += ${gap};
        `;
    }

    // --- 2. Build Iteration Space (Outer Loops) ---
    // The "body" of the iteration produces one output element per reduction block
    let fullCode = `
        let acc = initVal;
        ${redCode}
        dataOut[pOut++] = ${finalExpr};
    `;

    if (nIter > 0) {
        // Calculate the total memory distance covered by the reduction block
        // Displacement = (shape of outermost reduction axis) * (stride of outermost reduction axis)
        const innerBlockDisplacement = shape[reduceAxes[0]] * strides[reduceAxes[0]];
        
        for (let d = nIter - 1; d >= 0; d--) {
            const ax = iterAxes[d];
            const stride = strides[ax];
            // The displacement to account for is either the reduction block (if innermost iter)
            // or the inner iteration loop displacement.
            const movedByChild = (d === nIter - 1) ? innerBlockDisplacement : (shape[iterAxes[d+1]] * strides[iterAxes[d+1]]);
            const gap = stride - movedByChild;

            // CRITICAL: pIn += gap must be INSIDE the loop to adjust ptr for next i{d}
            fullCode = `
                for (let i${d} = 0; i${d} < ${shape[ax]}; i${d}++) {
                    ${fullCode}
                    pIn += ${gap}; 
                }
            `;
        }
    }

    // --- 3. Assemble Final Kernel ---
    const fnSource = `
        return function(dataIn, dataOut, offIn, initVal, count) {
            "use strict";
            let pIn = offIn;
            let pOut = 0;
            ${fullCode}
        };
    `;

    if(Jit.debug) {
        console.log(cacheKey, "reduce Kernel Source:\n", fnSource);
    }
    kernel = new Function(fnSource)();
    REDUCE_KERNEL_CACHE.set(cacheKey, kernel);
    return kernel;
}

/**
 * Global cache for set JIT kernels to avoid redundant compilations.
 * @private
 */
const SET_KERNEL_CACHE = new Map();

/**
 * Generates a JIT-compiled kernel for the set operation.
 * Optimized for V8 by using incremental pointer arithmetic and static nesting.
 * 
 * @param {string} cacheKey - Unique key for the specific array structure.
 * @param {number} ndim - Number of dimensions of the target array.
 * @param {Array} targetShape - The logical shape of the selection.
 * @param {Int32Array} tStrides - Strides of the target NDArray.
 * @param {Int32Array} sStrides - Pre-computed broadcasting strides for the source.
 * @param {Uint8Array} hasPSet - Flags indicating if a dimension uses fancy indexing or scalar.
 * @param {Array<boolean>} isDimReduced - Flags indicating if a dimension is collapsed by scalar indexing.
 * @private
 */
export function _createSetKernel(cacheKey, ndim, targetShape, tStrides, sStrides, hasPSet, isDimReduced) {
    let kernel = SET_KERNEL_CACHE.get(cacheKey);
    if (kernel) return kernel;

    let targetDimIdx = 0;

    /**
     * Recursively builds the nested loop string.
     * @param {number} d - Current dimension index.
     * @returns {string} - The generated code block for this dimension.
     */
    function buildLevel(d) {
        if (d === ndim) {
            // Innermost level: Perform the actual data assignment
            return `dataT[pT${d}] = dataS[pS${d}];`;
        }

        const tStride = tStrides[d];
        const sStride = sStrides[d];
        const pT_prev = `pT${d}`;
        const pS_prev = `pS${d}`;
        const pT_next = `pT${d + 1}`;
        const pS_next = `pS${d + 1}`;

        if (isDimReduced[d]) {
            // SCALAR INDEX: Offset the pointer once and move to the next dimension level.
            return `
            const ${pT_next} = ${pT_prev} + ps[${d}][0] * ${tStride};
            const ${pS_next} = ${pS_prev};
            ${buildLevel(d + 1)}`;
        } else {
            const len = targetShape[targetDimIdx++];
            if (!hasPSet[d]) {
                // SLICE INDEX (:): Use incremental pointer arithmetic in the loop header for V8 speed.
                return `
                for (let i${d} = 0, ${pT_next} = ${pT_prev}, ${pS_next} = ${pS_prev}; i${d} < ${len}; i${d}++, ${pT_next} += ${tStride}, ${pS_next} += ${sStride}) {
                    ${buildLevel(d + 1)}
                }`;
            } else {
                // FANCY INDEXING: Map the pointer using the pick-set array at each iteration.
                return `
                for (let i${d} = 0; i${d} < ${len}; i${d}++) {
                    const ${pT_next} = ${pT_prev} + ps[${d}][i${d}] * ${tStride};
                    const ${pS_next} = ${pS_prev} + i${d} * ${sStride};
                    ${buildLevel(d + 1)}
                }`;
            }
        }
    }

    const fnSource = `
        return function(dataT, dataS, offT, offS, ps) {
            "use strict";
            const pT0 = offT;
            const pS0 = offS;
            ${buildLevel(0)}
        };
    `;

    if(Jit.debug) {
        console.log(cacheKey, "Set Kernel Source:\n", fnSource);
    }

    kernel = new Function(fnSource)();
    SET_KERNEL_CACHE.set(cacheKey, kernel);
    return kernel;
}

/**
 * @private
 */
const BIN_KERNEL_CACHE = new Map();

/**
 * Extracts the expression body from a function.
 * Supports both: (a, b) => a + b  AND  function(a, b) { return a + b; }
 * @private
 */
function extractOpBody(fnStr) {
    // Try arrow function first
    let match = fnStr.match(/=>\s*([\s\S]+)/);
    if (match) return match[1].trim().replace(/;$/, '');
    // Try standard function return
    match = fnStr.match(/\{[\s\S]*return\s+([\s\S]+?);?\s*\}/);
    if (match) return match[1].trim();
    return fnStr; // Fallback
}




/**
 * Global cache for unary JIT kernels (copy/map).
 * @private
 */
const UNARY_KERNEL_CACHE = new Map();

/**
 * Extracts function body or processes template strings into executable JS code.
 * @private
 */
function prepareUnaryOp(fnOrStr) {
    if (typeof fnOrStr === 'string') {
        // If it's a template like 'Math.sin(${val})', replace with pointer access
        if (fnOrStr.includes('${val}')) {
            return fnOrStr.replace(/\$\{val\}/g, 'dataIn[ptrIn]');
        }
        return fnOrStr.replace(/\bval\b/g, 'dataIn[ptrIn]');
    }
    // If it's a function, stringify and replace the parameter with pointer access
    const fnStr = fnOrStr.toString();
    const body = extractOpBody(fnStr); // Reuses the helper
    // Matches common parameter names like x, val, or item
    return body.replace(/\b(x|a|val|item)\b/g, 'dataIn[ptrIn]');
}







/**
 * @private
 */
const REDUCE_KERNEL_CACHE = new Map();

/**
 * Normalizes the reducer and finalizer into inline expressions.
 * @private
 */
function prepareReduceExpr(fnOrStr, type = 'reducer') {
    const s = fnOrStr.toString();
    const body = extractOpBody(s);
    if (type === 'reducer') {
        // Map arguments: (acc, val) -> acc + dataIn[pIn]
        return body.replace(/\bacc|a\b/g, 'acc').replace(/\b(b|val|v|item)\b/g, 'dataIn[pIn]');
    }
    // Map arguments: (acc, count) -> acc / count
    return body.replace(/\bacc|a\b/g, 'acc').replace(/\b(n|count|len)\b/g, 'count');
}





/**
 * Global cache for pick JIT kernels.
 * @private
 */
const PICK_KERNEL_CACHE = new Map();

/**
 * Generates a JIT-compiled kernel for the pick operation.
 * Optimized for V8 with incremental pointer arithmetic and contiguous output writes.
 * 
 * @param {string} cacheKey - Unique key for the specific array structure.
 * @param {number} ndim - Number of dimensions of the source array.
 * @param {Int32Array} sStrides - Strides of the source NDArray.
 * @param {Uint8Array} isFullSlice - Flags indicating if a dimension is a full ":" slice.
 * @param {Array<boolean>} isDimReduced - Flags indicating if a dimension is collapsed by scalar indexing.
 * @param {Int32Array} odometerShape - The lengths of the pick-sets for each dimension.
 * @private
 */
export function _createPickKernel(cacheKey, ndim, sStrides, isFullSlice, isDimReduced, odometerShape) {
    let kernel = PICK_KERNEL_CACHE.get(cacheKey);
    if (kernel) return kernel;

    /**
     * Recursively builds the nested loop string.
     * @param {number} d - Current source dimension index.
     */
    function buildLevel(d) {
        if (d === ndim) {
            // Innermost level: Read from calculated source pointer and write to contiguous output
            return `dataOut[pOut++] = dataIn[pIn${d}];`;
        }

        const sStride = sStrides[d];
        const pIn_prev = `pIn${d}`;
        const pIn_next = `pIn${d + 1}`;
        const len = odometerShape[d];

        if (isDimReduced[d]) {
            // SCALAR INDEX: Offset the source pointer once and move to next dimension.
            return `
            const ${pIn_next} = ${pIn_prev} + ps[${d}][0] * ${sStride};
            ${buildLevel(d + 1)}`;
        } else if (isFullSlice[d]) {
            // FULL SLICE (:): Optimization - use incremental pointer arithmetic (no pick-set lookup).
            return `
            for (let i${d} = 0, ${pIn_next} = ${pIn_prev}; i${d} < ${len}; i${d}++, ${pIn_next} += ${sStride}) {
                ${buildLevel(d + 1)}
            }`;
        } else {
            // FANCY INDEXING: Lookup source index from the pick-set (ps) array.
            return `
            for (let i${d} = 0; i${d} < ${len}; i${d}++) {
                const ${pIn_next} = ${pIn_prev} + ps[${d}][i${d}] * ${sStride};
                ${buildLevel(d + 1)}
            }`;
        }
    }

    const fnSource = `
        return function(dataIn, dataOut, offIn, ps) {
            "use strict";
            let pOut = 0;
            const pIn0 = offIn;
            ${buildLevel(0)}
        };
    `;

    if(Jit.debug) { 
        console.log(cacheKey, "Pick Kernel Source:\n", fnSource);
    }

    kernel = new Function(fnSource)();
    PICK_KERNEL_CACHE.set(cacheKey, kernel);
    return kernel;
}