import _ from 'lodash';
import '../../../../../assert/js/assert.js';
import '../../../../../globals.js';
// Copyright 2023-2024, University of Colorado Boulder
/**
 * A template that performs a reduce operation on a single workgroup. The value will be valid at local_id.x === 0.
 *
 * @author Jonathan Olson <jonathan.olson@colorado.edu>
 */
import { optionize3 } from '../../../../../phet-core/js/optionize.js';
import { alpenglow, binaryExpressionStatementWGSL, commentWGSL, LOCAL_INDEXABLE_DEFAULTS, logValueWGSL, u32S, unrollWGSL, wgsl } from '../../../imports.js';
export const REDUCE_DEFAULTS = {
    mapScratchIndex: _.identity,
    convergent: false,
    scratchPreloaded: false,
    valuePreloaded: true,
    ...LOCAL_INDEXABLE_DEFAULTS // eslint-disable-line phet/no-object-spread-on-non-literals
};
const reduceWGSL = (providedOptions) => {
    const options = optionize3()({}, REDUCE_DEFAULTS, providedOptions);
    const value = options.value;
    const scratch = options.scratch;
    const workgroupSize = options.workgroupSize;
    const binaryOp = options.binaryOp;
    const localIndex = options.localIndex;
    const mapScratchIndex = options.mapScratchIndex;
    const convergent = options.convergent;
    const scratchPreloaded = options.scratchPreloaded;
    const valuePreloaded = options.valuePreloaded;
    assert && assert(!convergent || Number.isInteger(Math.log2(workgroupSize)));
    assert && assert(scratchPreloaded || valuePreloaded);
    const start = convergent ? Math.log2(workgroupSize) : 0;
    const end = convergent ? 0 : Math.log2(workgroupSize);
    const condition = (i) => convergent
        ? wgsl `${localIndex} < ${u32S(1 << (i - 1))}`
        : wgsl `${localIndex} % ${u32S(1 << (i + 1))} == 0u`;
    const accessIndex = (i) => convergent
        ? wgsl `${localIndex} + ${u32S(1 << (i - 1))}`
        : wgsl `${localIndex} + ${u32S(1 << i)}`;
    const combineToValue = (varName, a, b) => {
        return binaryExpressionStatementWGSL(varName, binaryOp.combineExpression || null, binaryOp.combineStatements || null, a, b);
    };
    return wgsl `
    ${commentWGSL(`begin reduce convergent:${convergent}`)}
    ${!scratchPreloaded ? wgsl `
      ${scratch}[ ${mapScratchIndex(localIndex)} ] = ${value};
    ` : wgsl ``}
    ${!valuePreloaded ? wgsl `
      ${value} = ${scratch}[ ${mapScratchIndex(localIndex)} ];
    ` : wgsl ``}
    
    ${logValueWGSL({
        name: `before reduce convergent:${convergent}`,
        value: 'value',
        type: binaryOp.type
    })}
    
    ${unrollWGSL(start, end, (i, isFirst, isLast) => wgsl `
      // We don't need the first workgroupBarrier() if scratchPreloaded is true
      ${!scratchPreloaded || !isFirst ? wgsl `
        workgroupBarrier();
      ` : wgsl ``}

      // TODO: check performance differences with a select/combine?
      if ( ${condition(i)} ) {
        ${combineToValue(value, value, wgsl `${scratch}[ ${mapScratchIndex(accessIndex(i))} ]`)}

        ${!isLast ? wgsl `
          ${scratch}[ ${mapScratchIndex(localIndex)} ] = ${value};
        ` : wgsl ``}
      }
    `)}
    ${commentWGSL('end reduce')}
  `;
};
export default reduceWGSL;
alpenglow.register('reduceWGSL', reduceWGSL);
//# sourceMappingURL=reduceWGSL.js.map