import _ from 'lodash';
import '../../../../../assert/js/assert.js';
import '../../../../../globals.js';
// Copyright 2024, University of Colorado Boulder
/**
 * A single level of standalone reduction.
 *
 * @author Jonathan Olson <jonathan.olson@colorado.edu>
 */
import { alpenglow, BufferArraySlot, ceilDivideConstantDivisorWGSL, CompositeModule, getArrayType, MainReduceModule, MainReduceNonCommutativeModule, MainScanModule } from '../../../imports.js';
import { combineOptions, optionize3 } from '../../../../../phet-core/js/optionize.js';
export const SCAN_MODULE_DEFAULTS = {
    name: 'scan',
    log: false, // TODO: how to deduplicate this? - We don't really need all of the defaults, right?
    inputOrder: 'blocked',
    inputAccessOrder: 'striped',
    internalStriping: false,
    areScannedReductionsExclusive: false,
    mainScanModuleOptions: {},
    mainReduceModuleOptions: {},
    mainReduceNonCommutativeModuleOptions: {}
};
// stageInputSize: number
export default class ScanModule extends CompositeModule {
    constructor(providedOptions) {
        const options = optionize3()({}, SCAN_MODULE_DEFAULTS, providedOptions);
        const perStageReduction = options.workgroupSize * options.grainSize;
        const initialStageInputSize = options.input.length;
        const numStages = Math.ceil(Math.log2(initialStageInputSize) / Math.log2(perStageReduction));
        assert && assert(numStages > 0, 'Do not pass in a single-element input');
        let modules;
        let execute;
        let slots;
        let internalSlots;
        const inPlace = options.input === options.output;
        if (numStages === 1) {
            const module = new MainScanModule(combineOptions({
                name: `${options.name} 0`,
                log: options.log,
                binaryOp: options.binaryOp,
                workgroupSize: options.workgroupSize,
                grainSize: options.grainSize,
                lengthExpression: options.lengthExpression,
                exclusive: options.exclusive,
                inputOrder: options.inputOrder,
                inputAccessOrder: options.inputAccessOrder
            }, inPlace ? {
                inPlace: true,
                data: options.input
            } : {
                inPlace: false,
                input: options.input,
                output: options.output
            }, options.mainScanModuleOptions));
            modules = [module];
            execute = (context, inputSize) => {
                module.execute(context, inputSize);
            };
            slots = _.uniq([options.input, options.output]);
            internalSlots = [];
        }
        else if (numStages === 2) {
            const reductionSlot = new BufferArraySlot(getArrayType(options.input.concreteArrayType.elementType, Math.ceil(initialStageInputSize / perStageReduction)));
            const reduceOptions = {
                name: `${options.name} reduce`,
                log: options.log,
                input: options.input,
                output: reductionSlot,
                binaryOp: options.binaryOp,
                workgroupSize: options.workgroupSize,
                grainSize: options.grainSize
            };
            const reduceModule = (options.inputOrder === 'blocked' && options.inputAccessOrder === 'striped' && !options.binaryOp.isCommutative) ? new MainReduceNonCommutativeModule(combineOptions({
                lengthExpression: options.lengthExpression
            }, reduceOptions, options.mainReduceNonCommutativeModuleOptions)) : new MainReduceModule(combineOptions({
                loadReducedOptions: {
                    lengthExpression: options.lengthExpression,
                    inputOrder: options.inputOrder,
                    inputAccessOrder: options.inputAccessOrder
                },
                reduceOptions: {
                    convergent: options.binaryOp.isCommutative
                },
                stripeOutput: options.internalStriping
            }, reduceOptions, options.mainReduceModuleOptions));
            const lowerScanModule = new MainScanModule(combineOptions({
                name: `${options.name} lower scan`,
                log: options.log,
                binaryOp: options.binaryOp,
                workgroupSize: options.workgroupSize,
                grainSize: options.grainSize,
                lengthExpression: options.lengthExpression ? ceilDivideConstantDivisorWGSL(options.lengthExpression, perStageReduction) : null,
                inputOrder: options.internalStriping ? 'striped' : 'blocked',
                inputAccessOrder: options.inputAccessOrder,
                exclusive: options.areScannedReductionsExclusive,
                inPlace: true,
                data: reductionSlot
            }, options.mainScanModuleOptions));
            const upperScanModule = new MainScanModule(combineOptions({
                name: `${options.name} upper scan`,
                log: options.log,
                binaryOp: options.binaryOp,
                workgroupSize: options.workgroupSize,
                grainSize: options.grainSize,
                lengthExpression: options.lengthExpression,
                inputOrder: options.inputOrder,
                inputAccessOrder: options.inputAccessOrder,
                exclusive: options.exclusive,
                areScannedReductionsExclusive: options.areScannedReductionsExclusive,
                addScannedReduction: true,
                scannedReduction: reductionSlot
            }, inPlace ? {
                inPlace: true,
                data: options.input
            } : {
                inPlace: false,
                input: options.input,
                output: options.output
            }, options.mainScanModuleOptions));
            modules = [reduceModule, lowerScanModule, upperScanModule];
            execute = (context, inputSize) => {
                reduceModule.execute(context, inputSize);
                lowerScanModule.execute(context, Math.ceil(inputSize / perStageReduction));
                upperScanModule.execute(context, inputSize);
            };
            slots = _.uniq([options.input, options.output, reductionSlot]);
            internalSlots = [reductionSlot];
        }
        else if (numStages === 3) {
            const reductionSlot = new BufferArraySlot(getArrayType(options.input.concreteArrayType.elementType, Math.ceil(initialStageInputSize / perStageReduction)));
            const doubleReductionSlot = new BufferArraySlot(getArrayType(options.input.concreteArrayType.elementType, Math.ceil(initialStageInputSize / (perStageReduction * perStageReduction))));
            const reduceOptions = {
                name: `${options.name} reduce`,
                log: options.log,
                input: options.input,
                output: reductionSlot,
                binaryOp: options.binaryOp,
                workgroupSize: options.workgroupSize,
                grainSize: options.grainSize
            };
            const reduceModule = (options.inputOrder === 'blocked' && options.inputAccessOrder === 'striped' && !options.binaryOp.isCommutative) ? new MainReduceNonCommutativeModule(combineOptions({
                lengthExpression: options.lengthExpression
            }, reduceOptions, options.mainReduceNonCommutativeModuleOptions)) : new MainReduceModule(combineOptions({
                loadReducedOptions: {
                    lengthExpression: options.lengthExpression,
                    inputOrder: options.inputOrder,
                    inputAccessOrder: options.inputAccessOrder
                },
                reduceOptions: {
                    convergent: options.binaryOp.isCommutative
                },
                stripeOutput: options.internalStriping
            }, reduceOptions, options.mainReduceModuleOptions));
            const middleScanModule = new MainScanModule(combineOptions({
                name: `${options.name} middle scan`,
                log: options.log,
                binaryOp: options.binaryOp,
                workgroupSize: options.workgroupSize,
                grainSize: options.grainSize,
                lengthExpression: options.lengthExpression ? ceilDivideConstantDivisorWGSL(options.lengthExpression, perStageReduction) : null,
                inputOrder: options.internalStriping ? 'striped' : 'blocked',
                inputAccessOrder: options.inputAccessOrder,
                exclusive: options.areScannedReductionsExclusive,
                inPlace: true,
                data: reductionSlot,
                storeReduction: true,
                reduction: doubleReductionSlot,
                stripeReducedOutput: options.internalStriping
            }, options.mainScanModuleOptions));
            const lowerScanModule = new MainScanModule(combineOptions({
                name: `${options.name} lower scan`,
                log: options.log,
                binaryOp: options.binaryOp,
                workgroupSize: options.workgroupSize,
                grainSize: options.grainSize,
                lengthExpression: options.lengthExpression ? ceilDivideConstantDivisorWGSL(options.lengthExpression, perStageReduction * perStageReduction) : null,
                inputOrder: options.internalStriping ? 'striped' : 'blocked',
                inputAccessOrder: options.inputAccessOrder,
                exclusive: options.areScannedReductionsExclusive,
                inPlace: true,
                data: doubleReductionSlot,
                stripeReducedOutput: options.internalStriping
            }, options.mainScanModuleOptions));
            const upperScanModule = new MainScanModule(combineOptions({
                name: `${options.name} upper scan`,
                log: options.log,
                binaryOp: options.binaryOp,
                workgroupSize: options.workgroupSize,
                grainSize: options.grainSize,
                lengthExpression: options.lengthExpression,
                inputOrder: options.inputOrder,
                inputAccessOrder: options.inputAccessOrder,
                exclusive: options.exclusive,
                areScannedReductionsExclusive: options.areScannedReductionsExclusive,
                addScannedReduction: true,
                scannedReduction: reductionSlot,
                addScannedDoubleReduction: true,
                scannedDoubleReduction: doubleReductionSlot
            }, inPlace ? {
                inPlace: true,
                data: options.input
            } : {
                inPlace: false,
                input: options.input,
                output: options.output
            }, options.mainScanModuleOptions));
            modules = [reduceModule, middleScanModule, lowerScanModule, upperScanModule];
            execute = (context, inputSize) => {
                // context.u32Numbers( options.input ).then( histogram => console.log( `histogram`, histogram ) ).catch( e => { throw e; } );
                reduceModule.execute(context, inputSize);
                // context.u32Numbers( reductionSlot ).then( histogram => console.log( `reduction`, histogram ) ).catch( e => { throw e; } );
                middleScanModule.execute(context, Math.ceil(inputSize / perStageReduction));
                // context.u32Numbers( reductionSlot ).then( histogram => console.log( `reduction (scanned)`, histogram ) ).catch( e => { throw e; } );
                // context.u32Numbers( doubleReductionSlot ).then( histogram => console.log( `double reduction`, histogram ) ).catch( e => { throw e; } );
                lowerScanModule.execute(context, Math.ceil(inputSize / (perStageReduction * perStageReduction)));
                // context.u32Numbers( doubleReductionSlot ).then( histogram => console.log( `double reduction (scanned)`, histogram ) ).catch( e => { throw e; } );
                upperScanModule.execute(context, inputSize);
                // context.u32Numbers( options.output ).then( histogram => console.log( `scanned histogram`, histogram ) ).catch( e => { throw e; } );
            };
            slots = _.uniq([options.input, options.output, reductionSlot, doubleReductionSlot]);
            internalSlots = [reductionSlot, doubleReductionSlot];
        }
        else {
            throw new Error(`invalid number of stages: ${numStages}`);
        }
        super(modules, execute);
        this.input = providedOptions.input;
        this.output = providedOptions.output;
        this.slots = slots;
        this.internalSlots = internalSlots;
    }
}
alpenglow.register('ScanModule', ScanModule);
//# sourceMappingURL=ScanModule.js.map