text:
function add_results!(benchmark_results_w_params, benchmark_results)
    benchmark_results_w_params["TrialObj"] = benchmark_results
    benchmark_results_w_params["RawTimes"] = benchmark_results.times
    benchmark_results_w_params["Minimum"] = minimum(benchmark_results).time
    benchmark_results_w_params["Median"] = median(benchmark_results).time
    benchmark_results_w_params["Mean"] = mean(benchmark_results).time
    benchmark_results_w_params["σ"] = std(benchmark_results).time
    benchmark_results_w_params["Maximum"] = maximum(benchmark_results).time
end

function add_metadata!(benchmark_results_w_params, processor, algorithm::KernelAlgorithm{ALGN}) where ALGN
    # Get OS info
    os_name = @static Sys.iswindows() ? "windows" : (@static Sys.isapple() ? "macos" : @static Sys.islinux() ? "linux" : @static Sys.isunix() ? "generic_unix" : throw("Can't determine OS name"))

    # Get CPU info
    cpu_name = Sys.cpu_info()[1].model
    # Workaround for NVIDIA Jetson CPU
    cpu_name == "unknown" ? "NVIDIA ARMv8" : cpu_name

    # Record CUDA version
    cuda_version = string(CUDA.version())

    # Get GPU name
    gpu_name = name(CUDA.CuDevice(0))

    # Add metadata to the results
    benchmark_results_w_params["os"] = os_name
    benchmark_results_w_params["CPU_model"] = cpu_name
    benchmark_results_w_params["GPU_model"] = gpu_name
    benchmark_results_w_params["CUDA"] = cuda_version
    processor == "GPU" ? benchmark_results_w_params["algorithm"] = ALGODICTINV[ALGN] : nothing
end

# CPU Benchmark
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{false},
    num_samples,
    num_ants,
    num_correlators,
    algorithm
)
    system = gnss(use_gpu = enable_gpu)
    start_code_phase = 0.0
    carrier_phase = 0.0
    carrier_frequency = 1500Hz
    prn = 1
    code_frequency = get_code_frequency(system)

    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants=NumAnts(num_ants))
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), num_correlators)
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)

    state = TrackingState(prn, system, carrier_frequency, start_code_phase, num_samples=num_samples, num_ants=NumAnts(num_ants), correlator=correlator)

    code_replica = Tracking.get_code(state)
    Tracking.resize!(code_replica, num_samples + correlator_sample_shifts[end] - correlator_sample_shifts[1])
    carrier_replica = Tracking.get_carrier(state)
    Tracking.resize!(Tracking.choose(carrier_replica, signal), num_samples)
    downconverted_signal_temp = Tracking.get_downconverted_signal(state)
    downconverted_signal = Tracking.resize!(downconverted_signal_temp, size(signal, 1), signal)

    @benchmark Tracking.downconvert_and_correlate!(
        $system,
        $signal,
        $correlator,
        $code_replica,
        $start_code_phase,
        $carrier_replica,
        $carrier_phase,
        $downconverted_signal,
        $code_frequency,
        $correlator_sample_shifts,
        $carrier_frequency,
        $sampling_frequency,
        1,
        $num_samples,
        $prn
    )
end

# GPU Kernel Benchmark for KernelAlgorithm 1_3_cplx_multi
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{true},
    num_samples,
    num_ants,
    num_correlators,
    algorithm::KernelAlgorithm{1330}
)
    # Generate GNSS object and signal information
    system = gnss(use_gpu = enable_gpu)
    codes = system.codes
    code_frequency = get_code_frequency(system)
    code_length = get_code_length(system)
    start_code_phase = 0.0f0
    carrier_phase = 0.0f0
    carrier_frequency = 1500Hz
    prn = 1

    # Generate the signal
    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)

    # Generate correlator
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
    num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]

    # Generate CUDA kernel tuning parameters
    block_dim_z = num_correlators
    block_dim_y = num_ants
    # keep num_corrs and num_ants in seperate dimensions, truncate num_samples accordingly to fit
    block_dim_x = prevpow(2, 512 ÷ block_dim_y ÷ block_dim_z)
    threads_per_block = [(block_dim_x, block_dim_y, block_dim_z), 512]
    blocks_per_grid = cld(num_samples, block_dim_x)
    partial_sum = StructArray{ComplexF32}((CUDA.zeros(Float32, blocks_per_grid, block_dim_y, block_dim_z), CUDA.zeros(Float32, blocks_per_grid, block_dim_y, block_dim_z)))
    shmem_size = [sizeof(ComplexF32) * block_dim_x * block_dim_y * block_dim_z
                sizeof(ComplexF32) * 1024 * num_ants * num_correlators]
    Num_Ants = NumAnts(num_ants)
    @benchmark CUDA.@sync $kernel_algorithm(
        $threads_per_block,
        $blocks_per_grid,
        $shmem_size,
        nothing,
        $codes,
        $code_frequency,
        $sampling_frequency,
        $start_code_phase,
        $prn,
        $num_samples,
        $num_of_shifts,
        $code_length,
        $partial_sum,
        nothing,
        nothing,
        nothing,
        nothing,
        $signal.re,
        $signal.im,
        $correlator_sample_shifts,
        $carrier_frequency,
        $carrier_phase,
        $Num_Ants,
        nothing,
        $algorithm
    )
    # kernel_algorithm(
    #     threads_per_block,
    #     blocks_per_grid,
    #     shmem_size,
    #     nothing,
    #     codes,
    #     code_frequency,
    #     sampling_frequency,
    #     start_code_phase,
    #     prn,
    #     num_samples,
    #     num_of_shifts,
    #     code_length,
    #     partial_sum,
    #     nothing,
    #     nothing,
    #     nothing,
    #     nothing,
    #     signal.re,
    #     signal.im,
    #     correlator_sample_shifts,
    #     carrier_frequency,
    #     carrier_phase,
    #     Num_Ants,
    #     nothing,
    #     algorithm
    # )
end

# GPU Kernel Benchmark for KernelAlgorithm 1_3_cplx_multi_texmem
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{true},
    num_samples,
    num_ants,
    num_correlators,
    algorithm::KernelAlgorithm{1331}
)
    # Generate GNSS object and signal information
    system = gnss(use_gpu = enable_gpu)
    codes = system.codes
    codes = CuTexture(
        CuTextureArray(codes),
        address_mode = CUDA.ADDRESS_MODE_WRAP,
        interpolation = CUDA.NearestNeighbour(),
        normalized_coordinates = true
    )
    code_frequency = get_code_frequency(system)
    code_length = get_code_length(system)
    start_code_phase = 0.0f0
    carrier_phase = 0.0f0
    carrier_frequency = 1500Hz
    prn = 1

    # Generate the signal
    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)

    # Generate correlator
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
    num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]

    # Generate CUDA kernel tuning parameters
    block_dim_z = num_correlators
    block_dim_y = num_ants
    # keep num_corrs and num_ants in seperate dimensions, truncate num_samples accordingly to fit
    block_dim_x = prevpow(2, 512 ÷ block_dim_y ÷ block_dim_z)
    threads_per_block = [(block_dim_x, block_dim_y, block_dim_z), 1024]
    blocks_per_grid = cld(num_samples, block_dim_x)
    partial_sum = StructArray{ComplexF32}((CUDA.zeros(Float32, blocks_per_grid, block_dim_y, block_dim_z), CUDA.zeros(Float32, blocks_per_grid, block_dim_y, block_dim_z)))
    shmem_size = [sizeof(ComplexF32) * block_dim_x * block_dim_y * block_dim_z
                sizeof(ComplexF32) * 1024 * num_ants * num_correlators]
    Num_Ants = NumAnts(num_ants)
    @benchmark CUDA.@sync $kernel_algorithm(
        $threads_per_block,
        $blocks_per_grid,
        $shmem_size,
        nothing,
        $codes,
        $code_frequency,
        $sampling_frequency,
        $start_code_phase,
        $prn,
        $num_samples,
        $num_of_shifts,
        $code_length,
        $partial_sum,
        nothing,
        nothing,
        nothing,
        nothing,
        $signal.re,
        $signal.im,
        $correlator_sample_shifts,
        $carrier_frequency,
        $carrier_phase,
        $Num_Ants,
        nothing,
        $algorithm
    )
end

# GPU Kernel Benchmark for KernelAlgorithm 1_4_cplx_multi_texmem
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{true},
    num_samples,
    num_ants,
    num_correlators,
    algorithm::KernelAlgorithm{1431}
)
    # Generate GNSS object and signal information
    system = gnss(use_gpu = enable_gpu)
    codes = system.codes
    codes = CuTexture(
        CuTextureArray(codes),
        address_mode = CUDA.ADDRESS_MODE_WRAP,
        interpolation = CUDA.NearestNeighbour(),
        normalized_coordinates = true
    )
    code_frequency = get_code_frequency(system)
    code_length = get_code_length(system)
    start_code_phase = 0.0f0
    carrier_phase = 0.0f0
    carrier_frequency = 1500Hz
    prn = 1

    # Generate the signal
    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)

    # Generate correlator
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
    num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]

    # Generate CUDA kernel tuning parameters
    block_dim_z = num_correlators
    block_dim_y = num_ants
    # keep num_corrs and num_ants in seperate dimensions, truncate num_samples accordingly to fit
    num_of_threads = 512
    if num_ants == 1
        num_of_threads = 512
        if num_ants == 4
            num_of_threads = 128
            if num_ants == 16
                num_of_threads = 128
            end
        end
    end
    block_dim_x = prevpow(2, num_of_threads ÷ block_dim_y ÷ block_dim_z)
    threads_per_block = [(block_dim_x, block_dim_y, block_dim_z), num_of_threads]
    blocks_per_grid = cld(cld(num_samples, block_dim_x), 2)
    partial_sum = StructArray{ComplexF32}((CUDA.zeros(Float32, blocks_per_grid, block_dim_y, block_dim_z), CUDA.zeros(Float32, blocks_per_grid, block_dim_y, block_dim_z)))
    shmem_size = [sizeof(ComplexF32) * block_dim_x * block_dim_y * block_dim_z
                sizeof(ComplexF32) * threads_per_block[2] * num_ants * num_correlators]
    Num_Ants = NumAnts(num_ants)
    @benchmark CUDA.@sync $kernel_algorithm(
        $threads_per_block,
        $blocks_per_grid,
        $shmem_size,
        nothing,
        $codes,
        $code_frequency,
        $sampling_frequency,
        $start_code_phase,
        $prn,
        $num_samples,
        $num_of_shifts,
        $code_length,
        $partial_sum,
        nothing,
        nothing,
        nothing,
        nothing,
        $signal.re,
        $signal.im,
        $correlator_sample_shifts,
        $carrier_frequency,
        $carrier_phase,
        $Num_Ants,
        nothing,
        $algorithm
    )
end

# GPU Kernel Benchmark for KernelAlgorithm 2_3_cplx_multi
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{true},
    num_samples,
    num_ants,
    num_correlators,
    algorithm::KernelAlgorithm{2330}
)
    system = gnss(use_gpu = enable_gpu)
    codes = system.codes
    code_frequency = get_code_frequency(system)
    code_length = get_code_length(system)
    start_code_phase = 0.0f0
    carrier_phase = 0.0f0
    carrier_frequency = 1500Hz
    prn = 1
    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
    num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]
    code_replica = CUDA.zeros(Float32, num_samples + num_of_shifts)
    carrier_replica = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples), CUDA.zeros(Float32, num_samples)))
    downconverted_signal = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples, num_ants), CUDA.zeros(Float32, num_samples, num_ants)))
    threads_per_block = zeros(Int, 3)
    blocks_per_grid = zeros(Int, 3)
    accum = StructArray{ComplexF32}((CUDA.zeros(Float32, (num_samples, num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (num_samples, num_ants, length(correlator_sample_shifts)))))
    threads_per_block[3] = 1024
    blocks_per_grid[3] = cld(num_samples, threads_per_block[3])
    phi = StructArray{ComplexF32}((CUDA.zeros(Float32, (blocks_per_grid[3], num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (blocks_per_grid[3], num_ants, length(correlator_sample_shifts)))))
    shmem_size = sizeof(ComplexF32) * 1024 * num_correlators * num_ants
    code_replica_kernel = @cuda launch=false gen_code_replica_strided_kernel!(
            code_replica,
            codes,
            code_frequency,
            sampling_frequency,
            start_code_phase,
            prn,
            num_samples,
            num_of_shifts,
            code_length
        )
    blocks_per_grid[1], threads_per_block[1] = launch_configuration(code_replica_kernel.fun)
    downconvert_and_accumulate_kernel = @cuda launch=false downconvert_and_accumulate_strided_kernel!(
        accum.re,
        accum.im,
        code_replica,
        carrier_replica.re,
        carrier_replica.im,
        downconverted_signal.re,
        downconverted_signal.im,
        signal.re,
        signal.im,
        carrier_frequency,
        sampling_frequency,
        carrier_phase,
        num_samples,
        NumAnts(num_ants),
        correlator_sample_shifts
    )
    blocks_per_grid[2], threads_per_block[2] = launch_configuration(downconvert_and_accumulate_kernel.fun)
    Num_Ants=NumAnts(num_ants)
    @benchmark CUDA.@sync $kernel_algorithm(
        $threads_per_block,
        $blocks_per_grid,
        $shmem_size,
        $code_replica,
        $codes,
        $code_frequency,
        $sampling_frequency,
        $start_code_phase,
        $prn,
        $num_samples,
        $num_of_shifts,
        $code_length,
        $accum.re,
        $accum.im,
        $phi.re,
        $phi.im,
        $carrier_replica.re,
        $carrier_replica.im,
        $downconverted_signal.re,
        $downconverted_signal.im,
        $signal.re,
        $signal.im,
        $correlator_sample_shifts,
        $carrier_frequency,
        $carrier_phase,
        $Num_Ants,
        nothing,
        $algorithm
    )
end

# GPU Kernel Benchmark for KernelAlgorithm 2_3_cplx_multi_textmem
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{true},
    num_samples,
    num_ants,
    num_correlators,
    algorithm::KernelAlgorithm{2331}
)
    system = gnss(use_gpu = enable_gpu)
    codes = system.codes
    #convert to text_mem
    codes = CuTexture(
        CuTextureArray(codes),
        address_mode = CUDA.ADDRESS_MODE_WRAP,
        interpolation = CUDA.NearestNeighbour(),
        normalized_coordinates = true
    )
    code_frequency = get_code_frequency(system)
    code_length = get_code_length(system)
    start_code_phase = 0.0f0
    carrier_phase = 0.0f0
    carrier_frequency = 1500Hz
    prn = 1
    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
    num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]
    code_replica = CUDA.zeros(Float32, num_samples + num_of_shifts)
    carrier_replica = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples), CUDA.zeros(Float32, num_samples)))
    downconverted_signal = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples, num_ants), CUDA.zeros(Float32, num_samples, num_ants)))
    threads_per_block = zeros(Int, 3)
    blocks_per_grid = zeros(Int, 3)
    accum = StructArray{ComplexF32}((CUDA.zeros(Float32, (num_samples, num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (num_samples, num_ants, length(correlator_sample_shifts)))))
    threads_per_block[3] = 1024
    blocks_per_grid[3] = cld(num_samples, threads_per_block[3])
    phi = StructArray{ComplexF32}((CUDA.zeros(Float32, (blocks_per_grid[3], num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (blocks_per_grid[3], num_ants, length(correlator_sample_shifts)))))
    shmem_size = sizeof(ComplexF32) * 1024 * num_correlators * num_ants
    code_replica_kernel = @cuda launch=false gen_code_replica_texture_mem_strided_kernel!(
            code_replica,
            codes,
            code_frequency,
            sampling_frequency,
            start_code_phase,
            prn,
            num_samples,
            num_of_shifts,
            code_length
        )
    blocks_per_grid[1], threads_per_block[1] = launch_configuration(code_replica_kernel.fun)
    downconvert_and_accumulate_kernel = @cuda launch=false downconvert_and_accumulate_strided_kernel!(
        accum.re,
        accum.im,
        code_replica,
        carrier_replica.re,
        carrier_replica.im,
        downconverted_signal.re,
        downconverted_signal.im,
        signal.re,
        signal.im,
        carrier_frequency,
        sampling_frequency,
        carrier_phase,
        num_samples,
        NumAnts(num_ants),
        correlator_sample_shifts
    )
    blocks_per_grid[2], threads_per_block[2] = launch_configuration(downconvert_and_accumulate_kernel.fun)
    Num_Ants=NumAnts(num_ants)
    @benchmark CUDA.@sync $kernel_algorithm(
        $threads_per_block,
        $blocks_per_grid,
        $shmem_size,
        $code_replica,
        $codes,
        $code_frequency,
        $sampling_frequency,
        $start_code_phase,
        $prn,
        $num_samples,
        $num_of_shifts,
        $code_length,
        $accum.re,
        $accum.im,
        $phi.re,
        $phi.im,
        $carrier_replica.re,
        $carrier_replica.im,
        $downconverted_signal.re,
        $downconverted_signal.im,
        $signal.re,
        $signal.im,
        $correlator_sample_shifts,
        $carrier_frequency,
        $carrier_phase,
        $Num_Ants,
        nothing,
        $algorithm
    )
end

# GPU Kernel Benchmark for KernelAlgorithm 2_4_cplx_multi
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{true},
    num_samples,
    num_ants,
    num_correlators,
    algorithm::KernelAlgorithm{2430}
)
    system = gnss(use_gpu = enable_gpu)
    codes = system.codes
    code_frequency = get_code_frequency(system)
    code_length = get_code_length(system)
    start_code_phase = 0.0f0
    carrier_phase = 0.0f0
    carrier_frequency = 1500Hz
    prn = 1
    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
    num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]
    code_replica = CUDA.zeros(Float32, num_samples + num_of_shifts)
    carrier_replica = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples), CUDA.zeros(Float32, num_samples)))
    downconverted_signal = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples, num_ants), CUDA.zeros(Float32, num_samples, num_ants)))
    threads_per_block = zeros(Int, 3)
    blocks_per_grid = zeros(Int, 3)
    accum = StructArray{ComplexF32}((CUDA.zeros(Float32, (num_samples, num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (num_samples, num_ants, length(correlator_sample_shifts)))))
    threads_per_block[3] = 1024
    blocks_per_grid[3] = cld(num_samples, threads_per_block[3])
    phi = StructArray{ComplexF32}((CUDA.zeros(Float32, (blocks_per_grid[3], num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (blocks_per_grid[3], num_ants, length(correlator_sample_shifts)))))
    shmem_size = sizeof(ComplexF32) * 1024 * num_correlators * num_ants
    code_replica_kernel = @cuda launch=false gen_code_replica_strided_kernel!(
            code_replica,
            codes,
            code_frequency,
            sampling_frequency,
            start_code_phase,
            prn,
            num_samples,
            num_of_shifts,
            code_length
        )
    blocks_per_grid[1], threads_per_block[1] = launch_configuration(code_replica_kernel.fun)
    downconvert_and_accumulate_kernel = @cuda launch=false downconvert_and_accumulate_strided_kernel!(
        accum.re,
        accum.im,
        code_replica,
        carrier_replica.re,
        carrier_replica.im,
        downconverted_signal.re,
        downconverted_signal.im,
        signal.re,
        signal.im,
        carrier_frequency,
        sampling_frequency,
        carrier_phase,
        num_samples,
        NumAnts(num_ants),
        correlator_sample_shifts
    )
    blocks_per_grid[2], threads_per_block[2] = launch_configuration(downconvert_and_accumulate_kernel.fun)
    Num_Ants=NumAnts(num_ants)
    @benchmark CUDA.@sync $kernel_algorithm(
        $threads_per_block,
        $blocks_per_grid,
        $shmem_size,
        $code_replica,
        $codes,
        $code_frequency,
        $sampling_frequency,
        $start_code_phase,
        $prn,
        $num_samples,
        $num_of_shifts,
        $code_length,
        $accum.re,
        $accum.im,
        $phi.re,
        $phi.im,
        $carrier_replica.re,
        $carrier_replica.im,
        $downconverted_signal.re,
        $downconverted_signal.im,
        $signal.re,
        $signal.im,
        $correlator_sample_shifts,
        $carrier_frequency,
        $carrier_phase,
        $Num_Ants,
        nothing,
        $algorithm
    )
end

# GPU Kernel Benchmark for KernelAlgorithm 2_4_cplx_multi_textmem
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{true},
    num_samples,
    num_ants,
    num_correlators,
    algorithm::KernelAlgorithm{2431}
)
    system = gnss(use_gpu = enable_gpu)
    codes = system.codes
    #convert to text_mem
    codes = CuTexture(
        CuTextureArray(codes),
        address_mode = CUDA.ADDRESS_MODE_WRAP,
        interpolation = CUDA.NearestNeighbour(),
        normalized_coordinates = true
    )
    code_frequency = get_code_frequency(system)
    code_length = get_code_length(system)
    start_code_phase = 0.0f0
    carrier_phase = 0.0f0
    carrier_frequency = 1500Hz
    prn = 1
    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
    num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]
    code_replica = CUDA.zeros(Float32, num_samples + num_of_shifts)
    carrier_replica = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples), CUDA.zeros(Float32, num_samples)))
    downconverted_signal = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples, num_ants), CUDA.zeros(Float32, num_samples, num_ants)))
    threads_per_block = zeros(Int, 3)
    blocks_per_grid = zeros(Int, 3)
    accum = StructArray{ComplexF32}((CUDA.zeros(Float32, (num_samples, num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (num_samples, num_ants, length(correlator_sample_shifts)))))
    threads_per_block[3] = 1024
    blocks_per_grid[3] = cld(num_samples, threads_per_block[3])
    phi = StructArray{ComplexF32}((CUDA.zeros(Float32, (blocks_per_grid[3], num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (blocks_per_grid[3], num_ants, length(correlator_sample_shifts)))))
    shmem_size = sizeof(ComplexF32) * 1024 * num_correlators * num_ants
    code_replica_kernel = @cuda launch=false gen_code_replica_strided_kernel!(
            code_replica,
            codes,
            code_frequency,
            sampling_frequency,
            start_code_phase,
            prn,
            num_samples,
            num_of_shifts,
            code_length
        )
    blocks_per_grid[1], threads_per_block[1] = launch_configuration(code_replica_kernel.fun)
    downconvert_and_accumulate_kernel = @cuda launch=false downconvert_and_accumulate_strided_kernel!(
        accum.re,
        accum.im,
        code_replica,
        carrier_replica.re,
        carrier_replica.im,
        downconverted_signal.re,
        downconverted_signal.im,
        signal.re,
        signal.im,
        carrier_frequency,
        sampling_frequency,
        carrier_phase,
        num_samples,
        NumAnts(num_ants),
        correlator_sample_shifts
    )
    blocks_per_grid[2], threads_per_block[2] = launch_configuration(downconvert_and_accumulate_kernel.fun)
    Num_Ants=NumAnts(num_ants)
    max_shmem = CUDA.attribute(device(), CUDA.DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK)
    if max_shmem < shmem_size
        threads_per_block[3] = max_shmem ÷ (sizeof(ComplexF32) * num_correlators * num_ants)
        shmem_size = sizeof(ComplexF32) * threads_per_block[3] * num_correlators * num_ants
    end
    @benchmark CUDA.@sync $kernel_algorithm(
        $threads_per_block,
        $blocks_per_grid,
        $shmem_size,
        $code_replica,
        $codes,
        $code_frequency,
        $sampling_frequency,
        $start_code_phase,
        $prn,
        $num_samples,
        $num_of_shifts,
        $code_length,
        $accum.re,
        $accum.im,
        $phi.re,
        $phi.im,
        $carrier_replica.re,
        $carrier_replica.im,
        $downconverted_signal.re,
        $downconverted_signal.im,
        $signal.re,
        $signal.im,
        $correlator_sample_shifts,
        $carrier_frequency,
        $carrier_phase,
        $Num_Ants,
        nothing,
        $algorithm
    )
end

# GPU Kernel Benchmark for KernelAlgorithm 3_4_cplx_multi_textmem
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{true},
    num_samples,
    num_ants,
    num_correlators,
    algorithm::KernelAlgorithm{3431}
)
    system = gnss(use_gpu = enable_gpu)
    codes = system.codes
    #convert to text_mem
    codes = CuTexture(
        CuTextureArray(codes),
        address_mode = CUDA.ADDRESS_MODE_WRAP,
        interpolation = CUDA.NearestNeighbour(),
        normalized_coordinates = true
    )
    code_frequency = get_code_frequency(system)
    code_length = get_code_length(system)
    start_code_phase = 0.0f0
    carrier_phase = 0.0f0
    carrier_frequency = 1500Hz
    prn = 1
    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
    num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]
    code_replica = CUDA.zeros(Float32, num_samples + num_of_shifts)
    carrier_replica = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples), CUDA.zeros(Float32, num_samples)))
    downconverted_signal = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples, num_ants), CUDA.zeros(Float32, num_samples, num_ants)))
    threads_per_block = zeros(Int64, 3)
    blocks_per_grid = zeros(Int64, 3)
    shmem_size = zeros(Int64, 2)
    threads_per_block[3] = 1024
    blocks_per_grid[3] = cld(num_samples, threads_per_block[3])
    threads_per_block[2] = 256
    blocks_per_grid[2] = cld(num_samples, threads_per_block[2]) ÷ 2
    accum = StructArray{ComplexF32}((CUDA.zeros(Float32, (blocks_per_grid[2], num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (blocks_per_grid[2], num_ants, length(correlator_sample_shifts)))))
    shmem_size[1] = sizeof(ComplexF32) * threads_per_block[2] * num_correlators * num_ants
    shmem_size[2] = sizeof(ComplexF32) *  threads_per_block[3] * num_correlators * num_ants
    code_replica_kernel = @cuda launch=false gen_code_replica_texture_mem_strided_kernel!(
            code_replica,
            codes,
            code_frequency,
            sampling_frequency,
            start_code_phase,
            prn,
            num_samples,
            num_of_shifts,
            code_length
        )
    blocks_per_grid[1], threads_per_block[1] = launch_configuration(code_replica_kernel.fun)
    Num_Ants=NumAnts(num_ants)
    max_shmem = CUDA.attribute(device(), CUDA.DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK)
    if max_shmem < shmem_size[1]
        threads_per_block[2] = max_shmem ÷ (sizeof(ComplexF32) * num_correlators * num_ants)
        shmem_size[2] = sizeof(ComplexF32) * threads_per_block[2] * num_correlators * num_ants
    end
    if max_shmem < shmem_size[2]
        threads_per_block[3] = max_shmem ÷ (sizeof(ComplexF32) * num_correlators * num_ants)
        shmem_size[2] = sizeof(ComplexF32) * threads_per_block[3] * num_correlators * num_ants
    end
    @benchmark CUDA.@sync $kernel_algorithm(
        $threads_per_block,
        $blocks_per_grid,
        $shmem_size,
        $code_replica,
        $codes,
        $code_frequency,
        $sampling_frequency,
        $start_code_phase,
        $prn,
        $num_samples,
        $num_of_shifts,
        $code_length,
        $accum.re,
        $accum.im,
        $carrier_replica.re,
        $carrier_replica.im,
        $downconverted_signal.re,
        $downconverted_signal.im,
        $signal.re,
        $signal.im,
        $correlator_sample_shifts,
        $carrier_frequency,
        $carrier_phase,
        $Num_Ants,
        $nothing,
        $algorithm
    )
end

# GPU Kernel Benchmark for KernelAlgorithm 4_4_cplx_multi_textmem
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{true},
    num_samples,
    num_ants,
    num_correlators,
    algorithm::KernelAlgorithm{4431}
)
    system = gnss(use_gpu = enable_gpu)
    codes = system.codes
    #convert to text_mem
    codes = CuTexture(
        CuTextureArray(codes),
        address_mode = CUDA.ADDRESS_MODE_WRAP,
        interpolation = CUDA.NearestNeighbour(),
        normalized_coordinates = true
    )
    code_frequency = get_code_frequency(system)
    code_length = get_code_length(system)
    start_code_phase = 0.0f0
    carrier_phase = 0.0f0
    carrier_frequency = 1500Hz
    prn = 1
    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
    num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]
    code_replica = CUDA.zeros(Float32, num_samples + num_of_shifts)
    carrier_replica = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples), CUDA.zeros(Float32, num_samples)))
    downconverted_signal = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples, num_ants), CUDA.zeros(Float32, num_samples, num_ants)))
    threads_per_block = zeros(Int64, 3)
    blocks_per_grid = zeros(Int64, 3)
    threads_per_block[3] = 1024
    blocks_per_grid[3] = cld(num_samples, threads_per_block[3])
    threads_per_block[2] = 256
    blocks_per_grid[2] = cld(num_samples, threads_per_block[2]) ÷ 2
    accum = StructArray{ComplexF32}((CUDA.zeros(Float32, (num_ants, length(correlator_sample_shifts))),CUDA.zeros(Float32, (num_ants, length(correlator_sample_shifts)))))
    shmem_size = sizeof(ComplexF32) * threads_per_block[2] * num_correlators * num_ants
    code_replica_kernel = @cuda launch=false gen_code_replica_texture_mem_strided_kernel!(
            code_replica,
            codes,
            code_frequency,
            sampling_frequency,
            start_code_phase,
            prn,
            num_samples,
            num_of_shifts,
            code_length
        )
    blocks_per_grid[1], threads_per_block[1] = launch_configuration(code_replica_kernel.fun)
    Num_Ants=NumAnts(num_ants)
    @benchmark CUDA.@sync $kernel_algorithm(
        $threads_per_block,
        $blocks_per_grid,
        $shmem_size,
        $code_replica,
        $codes,
        $code_frequency,
        $sampling_frequency,
        $start_code_phase,
        $prn,
        $num_samples,
        $num_of_shifts,
        $code_length,
        $accum.re,
        $accum.im,
        $carrier_replica.re,
        $carrier_replica.im,
        $downconverted_signal.re,
        $downconverted_signal.im,
        $signal.re,
        $signal.im,
        $correlator_sample_shifts,
        $carrier_frequency,
        $carrier_phase,
        $Num_Ants,
        $nothing,
        $algorithm
    )
end

# GPU Kernel Benchmark for KernelAlgorithm 5_4_cplx_multi_textmem
function _run_kernel_benchmark(
    gnss,
    enable_gpu::Val{true},
    num_samples,
    num_ants,
    num_correlators,
    algorithm::KernelAlgorithm{5431}
)
    system = gnss(use_gpu = enable_gpu)
    codes = system.codes
    #convert to text_mem
    codes = CuTexture(
        CuTextureArray(codes),
        address_mode = CUDA.ADDRESS_MODE_WRAP,
        interpolation = CUDA.NearestNeighbour(),
        normalized_coordinates = true
    )
    code_frequency = get_code_frequency(system)
    code_length = get_code_length(system)
    start_code_phase = 0.0f0
    carrier_phase = 0.0f0
    carrier_frequency = 1500Hz
    prn = 1
    signal, sampling_frequency = gen_signal(system, prn, carrier_frequency, num_samples, num_ants = NumAnts(num_ants), start_code_phase = start_code_phase, start_carrier_phase = carrier_phase)
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(system, correlator, sampling_frequency, 0.5)
    num_of_shifts = correlator_sample_shifts[end] - correlator_sample_shifts[1]
    code_replica = CUDA.zeros(Float32, num_samples + num_of_shifts)
    carrier_replica = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples), CUDA.zeros(Float32, num_samples)))
    downconverted_signal = StructArray{ComplexF32}((CUDA.zeros(Float32, num_samples, num_ants), CUDA.zeros(Float32, num_samples, num_ants)))
    threads_per_block = 512
    blocks_per_grid = cld(num_samples, threads_per_block)
    accum = StructArray{ComplexF32}((CUDA.zeros(Float32, (num_ants, length(correlator_sample_shifts))), CUDA.zeros(Float32, (num_ants, length(correlator_sample_shifts)))))
    shmem_size = sizeof(ComplexF32) * threads_per_block * num_correlators * num_ants
    Num_Ants = NumAnts(num_ants)
    @benchmark CUDA.@sync @cuda threads=$threads_per_block blocks=$cld($blocks_per_grid,2) shmem=$shmem_size $downconvert_and_correlate_kernel_5431!(
        $accum.re,
        $accum.im,
        $carrier_replica.re,
        $carrier_replica.im,
        $downconverted_signal.re,
        $downconverted_signal.im,
        $signal.re,
        $signal.im,
        $codes,
        $code_length,
        $code_replica,
        $prn,
        $correlator_sample_shifts,
        $num_of_shifts,
        $code_frequency,
        $carrier_frequency,
        $sampling_frequency,
        $start_code_phase,
        $carrier_phase,
        $num_samples,
        $Num_Ants,
    )
end

function run_kernel_benchmark(benchmark_params::Dict)
    @unpack GNSS, num_samples, num_ants, num_correlators, processor, algorithm = benchmark_params
    enable_gpu = (processor == "GPU" ? Val(true) : Val(false))
    algorithm = KernelAlgorithm(ALGODICT[algorithm])
    benchmark_results = _run_kernel_benchmark(
        GNSSDICT[GNSS],
        enable_gpu,
        num_samples,
        num_ants,
        num_correlators,
        algorithm
    )
    benchmark_results_w_params = copy(benchmark_params)
    add_results!(benchmark_results_w_params, benchmark_results)
    add_metadata!(benchmark_results_w_params, processor, algorithm)
    return benchmark_results_w_params
end

function _bench_reduction(num_samples, num_ants, num_correlators, algorithm::ReductionAlgorithm{1})
    input = StructArray{ComplexF32}(
        (
            CUDA.ones(Float32, (num_samples, num_ants, num_correlators)),
            CUDA.zeros(Float32, (num_samples, num_ants, num_correlators))
        )
    )
    threads_per_block = 256
    blocks_per_grid = cld(num_samples, threads_per_block)
    accum = StructArray{ComplexF32}(
        (
            CUDA.zeros(Float32, (blocks_per_grid, num_ants, num_correlators)),
            CUDA.zeros(Float32, (blocks_per_grid, num_ants, num_correlators))
        )
    )
    shmem_size = sizeof(Float32) * threads_per_block
    max_shmem = CUDA.attribute(device(), CUDA.DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK)
    if max_shmem < shmem_size
        threads_per_block =prevpow(2, max_shmem ÷ (sizeof(ComplexF32) * num_ants * num_correlators))
        shmem_size = sizeof(ComplexF32) * threads_per_block * num_ants * num_correlators
    end
    return @benchmark CUDA.@sync begin
        @inbounds for antenna_idx = 1:$num_ants
            @inbounds for corr_idx = 1:$num_correlators
                # re samples
                @cuda threads=$threads_per_block blocks=$blocks_per_grid shmem=$shmem_size $reduce_3(
                    $view($accum.re, :, antenna_idx, corr_idx),
                    $view($input.re, :, antenna_idx, corr_idx),
                    $num_samples
                )
                # im samples
                @cuda threads=$threads_per_block blocks=$blocks_per_grid shmem=$shmem_size $reduce_3(
                    $view($accum.im, :, antenna_idx, corr_idx),
                    $view($input.im, :, antenna_idx, corr_idx),
                    $num_samples
                )
            end
        end
        @inbounds for antenna_idx = 1:$num_ants
            @inbounds for corr_idx = 1:$num_correlators
                # re samples
                @cuda threads=$threads_per_block blocks=1 shmem=$shmem_size $reduce_3(
                    $view($accum.re, :, antenna_idx, corr_idx),
                    $view($accum.re, :, antenna_idx, corr_idx),
                    $size($accum, 1)
                )
                # im samples
                @cuda threads=$threads_per_block blocks=1 shmem=$shmem_size $reduce_3(
                    $view($accum.im, :, antenna_idx, corr_idx),
                    $view($accum.im, :, antenna_idx, corr_idx),
                    $size($accum, 1)
                )
            end
        end
    end
end

function _bench_reduction(num_samples, num_ants, num_correlators, algorithm::ReductionAlgorithm{2})
    signal_duration = 0.001s
    sampling_frequency = (num_samples/ustrip(signal_duration))Hz
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(GPSL1(), correlator, sampling_frequency, 0.5)
    input = StructArray{ComplexF32}(
        (
            CUDA.ones(Float32, (num_samples, num_ants, num_correlators)),
            CUDA.zeros(Float32, (num_samples, num_ants, num_correlators))
        )
    )
    threads_per_block = 256
    blocks_per_grid = cld(num_samples, threads_per_block)
    accum = StructArray{ComplexF32}(
        (
            CUDA.zeros(Float32, (blocks_per_grid, num_ants, num_correlators)),
            CUDA.zeros(Float32, (blocks_per_grid, num_ants, num_correlators))
        )
    )
    shmem_size = sizeof(ComplexF32) * threads_per_block
    max_shmem = CUDA.attribute(device(), CUDA.DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK)
    if max_shmem < shmem_size
        threads_per_block =prevpow(2, max_shmem ÷ (sizeof(ComplexF32) * num_ants * num_correlators))
        shmem_size = sizeof(ComplexF32) * threads_per_block * num_ants * num_correlators
    end
    return @benchmark CUDA.@sync begin
        @inbounds for antenna_idx = 1:$num_ants
            @inbounds for corr_idx = 1:$num_correlators
                @cuda threads=$threads_per_block blocks=$blocks_per_grid shmem=$shmem_size $reduce_cplx_3(
                    $view($accum.re, :, antenna_idx, corr_idx),
                    $view($accum.im, :, antenna_idx, corr_idx),
                    $view($input.re, :, antenna_idx, corr_idx),
                    $view($input.im, :, antenna_idx, corr_idx),
                    $num_samples
                )
            end
        end
        @inbounds for antenna_idx = 1:$num_ants
            @inbounds for corr_idx = 1:$num_correlators
                @cuda threads=$threads_per_block blocks=$1 shmem=$shmem_size $reduce_cplx_3(
                    $view($accum.re, :, antenna_idx, corr_idx),
                    $view($accum.im, :, antenna_idx, corr_idx),
                    $view($accum.re, :, antenna_idx, corr_idx),
                    $view($accum.im, :, antenna_idx, corr_idx),
                    $size($accum, 1)
                )
            end
        end
    end
end

function _bench_reduction(num_samples, num_ants, num_correlators, algorithm::ReductionAlgorithm{3})
    signal_duration = 0.001s
    sampling_frequency = (num_samples/ustrip(signal_duration))Hz
    correlator = EarlyPromptLateCorrelator(NumAnts(num_ants), NumAccumulators(num_correlators))
    correlator_sample_shifts = get_correlator_sample_shifts(GPSL1(), correlator, sampling_frequency, 0.5)
    input = StructArray{ComplexF32}(
        (
            CUDA.ones(Float32, (num_samples, num_ants, num_correlators)),
            CUDA.zeros(Float32, (num_samples, num_ants, num_correlators))
        )
    )
    threads_per_block = 256
    blocks_per_grid = cld(num_samples, threads_per_block)
    accum = StructArray{ComplexF32}(
        (
            CUDA.zeros(Float32, (blocks_per_grid, num_ants, num_correlators)),
            CUDA.zeros(Float32, (blocks_per_grid, num_ants, num_correlators))
        )
    )
    shmem_size = sizeof(ComplexF32) * threads_per_block * num_ants * num_correlators
    max_shmem = CUDA.attribute(device(), CUDA.DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK)
    if max_shmem < shmem_size
        threads_per_block =prevpow(2, max_shmem ÷ (sizeof(ComplexF32) * num_ants * num_correlators))
        shmem_size = sizeof(ComplexF32) * threads_per_block * num_ants * num_correlators
    end
    Num_Ants = NumAnts(num_ants)
    return @benchmark CUDA.@sync begin
        @cuda threads=$threads_per_block blocks=$blocks_per_grid shmem=$shmem_size reduce_cplx_multi_3(
            $accum.re,
            $accum.im,
            $input.re,
            $input.im,
            $num_samples,
            $Num_Ants,
            $correlator_sample_shifts
        )
        @cuda threads=$threads_per_block blocks=1 shmem=$shmem_size reduce_cplx_multi_3(
            $accum.re,
            $accum.im,
            $accum.re,
            $accum.im,
            $blocks_per_grid,
            $Num_Ants,
            $correlator_sample_shifts
        )
    end
end

function run_reduction_benchmark(benchmark_params::Dict)
    @unpack num_samples, num_ants, num_correlators, algorithm = benchmark_params
    algorithm_obj = REDDICT[algorithm]
    benchmark_results = _bench_reduction(num_samples, num_ants, num_correlators, algorithm_obj)
    benchmark_results_w_params = copy(benchmark_params)
    benchmark_results_w_params["Minimum"] = minimum(benchmark_results).time
    benchmark_results_w_params["Mean"] = mean(benchmark_results).time
    benchmark_results_w_params["Median"] = median(benchmark_results).time
    benchmark_results_w_params["Std"] = std(benchmark_results).time
    return benchmark_results_w_params
end

max_stars_repo_name: mfkiwl/GPUAcceleratedTracking

max_issues_repo_name: mfkiwl/GPUAcceleratedTracking

max_forks_repo_name: mfkiwl/GPUAcceleratedTracking


text:
export estimate_parameters!,
    uniform_parameters!,
    estimate_parameters_em!,
    update_pc_params_from_pbc!,
    estimate_parameters_em_multi_epochs!

using DirectedAcyclicGraphs: parent_stats
using LogicCircuits: num_nodes
using StatsFuns: logsumexp, logaddexp, logsubexp
using CUDA
using LoopVectorization

"""
    estimate_parameters!(pc::ProbCircuit, data; pseudocount::Float64, entropy_reg::Float64 = 0.0)

Maximum likelihood estimation of a `ProbCircuit`'s parameters given data
"""
function estimate_parameters!(pc::ProbCircuit, data; pseudocount::Float64, entropy_reg::Float64 = 0.0)
    estimate_single_circuit_parameters!(pc, data; pseudocount, entropy_reg)
end

"""
    estimate_parameters!(spc::SharedProbCircuit, data; pseudocount::Float64, entropy_reg::Float64 = 0.0)

Maximum likelihood estimation of a SharedProbCircuit's parameters given data.

bagging support: If `spc` is a SharedProbCircuit and data is an array of DataFrames
  with the same number of "components", learn each circuit with its corresponding
  dataset.
"""
function estimate_parameters!(spc::SharedProbCircuit, data; pseudocount::Float64, entropy_reg::Float64 = 0.0)
    @assert num_components(spc) == length(data) "SharedProbCircuit and data have different number of components: $(num_components(spc)) and $(length(data)), resp."

    map(1 : num_components(spc)) do component_idx
        estimate_single_circuit_parameters!(spc, data[component_idx]; pseudocount, component_idx, entropy_reg)
    end
end

"""
    estimate_single_circuit_parameters!(pc::ProbCircuit, data; pseudocount::Float64, component_idx::Integer = 0, entropy_reg::Float64 = 0.0)

Maximum likelihood estimation of a single circuit (for example, `ProbCircuit` or one component of `SharedProbCircuit`) parameters given data
"""
function estimate_single_circuit_parameters!(pc::ProbCircuit, data; pseudocount::Float64,
                                            component_idx::Integer = 0, entropy_reg::Float64 = 0.0)
    weights = nothing
    if isweighted(data)
        # `data' is weighted according to its `weight' column
        data, weights = split_sample_weights(data)
    end

    @assert isbinarydata(data) || isfpdata(data) "Probabilistic circuit parameter estimation for binary/floating point data only"
    bc = BitCircuit(pc, data)
    params = if isgpu(data)
        estimate_parameters_gpu(to_gpu(bc), data, pseudocount; weights, entropy_reg)
    else
        estimate_parameters_cpu(bc, data, pseudocount; weights, entropy_reg)
    end

    if pc isa SharedProbCircuit
        update_pc_params_from_pbc!(pc, bc, params, component_idx)
    else
        update_pc_params_from_pbc!(pc, bc, params)
    end
    params
end

"""
    update_pc_params_from_pbc!(pc::SharedProbCircuit, bc, params, component_idx; exp_update_factor = 0.0)

During parameter learning `pc`'s paramters are not updated automatically, and only the corresponding BitCircuit's paramters update.
This method updates the paramters of `pc` using the corresponding BitCircuit.

Note: This is mostly for internal use.
"""
function update_pc_params_from_pbc!(pc::SharedProbCircuit, bc, params, component_idx; exp_update_factor = 0.0)
    log_exp_factor = log(exp_update_factor)
    log_1_exp_factor = log(1.0 - exp_update_factor)
    foreach(pc) do pn
        if is⋁gate(pn)
            if num_children(pn) == 1
                pn.log_probs[:, component_idx] .= zero(Float64)
            else
                id = (bc.node2id[pn]::⋁NodeIds).node_id
                @inbounds els_start = bc.nodes[1,id]
                @inbounds els_end = bc.nodes[2,id]
                @inbounds @views pn.log_probs[:, component_idx] .= logaddexp.(log_exp_factor .+ pn.log_probs[:, component_idx], log_1_exp_factor .+ params[els_start:els_end])
                @assert isapprox(sum(exp.(pn.log_probs[:, component_idx])), 1.0, atol=1e-3) "Parameters do not sum to one locally: $(sum(exp.(pn.log_probs))); $(pn.log_probs)"
                pn.log_probs[:, component_idx] .-= logsumexp(pn.log_probs[:, component_idx]) # normalize away any leftover error
            end
        end
    end
    nothing
end

"""
    update_pc_params_from_pbc!(pc::ProbCircuit, pbc; exp_update_factor = 0.0)

During parameter learning `pc`'s paramters are not updated automatically, and only the corresponding ParamBitCircuit's paramters update.
This method updates the paramters of `pc` using the corresponding ParamBitCircuit.

Note: This is mostly for internal use.
"""
function update_pc_params_from_pbc!(pc::ProbCircuit, pbc; exp_update_factor = 0.0)
    if isgpu(pbc)
        pbc = to_cpu(pbc)
    end
    bc = pbc.bitcircuit
    params = pbc.params
    update_pc_params_from_pbc!(pc, bc, params; exp_update_factor)
end

"""
    function update_pc_params_from_pbc!(pc::ProbCircuit, bc, params; exp_update_factor = 0.0)

During parameter learning `pc`'s paramters are not updated automatically, and only the corresponding ParamBitCircuit's paramters update.
This method updates the paramters of `pc` using the corresponding ParamBitCircuit.

Note: This is mostly for internal use.
"""
function update_pc_params_from_pbc!(pc::ProbCircuit, bc, params; exp_update_factor = 0.0)
    log_exp_factor = log(exp_update_factor)
    log_1_exp_factor = log(1.0 - exp_update_factor)

    foreach(pc) do pn
        if is⋁gate(pn)
            if num_children(pn) == 1
                pn.log_probs .= zero(Float64)
            else
                id = (bc.node2id[pn]::⋁NodeIds).node_id
                @inbounds els_start = bc.nodes[1,id]
                @inbounds els_end = bc.nodes[2,id]
                @inbounds @views pn.log_probs .= logaddexp.(log_exp_factor .+ pn.log_probs, log_1_exp_factor .+ params[els_start:els_end])
                @assert isapprox(sum(exp.(pn.log_probs)), 1.0, atol=1e-3) "Parameters do not sum to one locally: $(sum(exp.(pn.log_probs))); $(pn.log_probs)"
                pn.log_probs .-= logsumexp(pn.log_probs) # normalize away any leftover error
            end
        end
    end

    nothing
end

"""
    estimate_parameters_cpu(bc::BitCircuit, data, pseudocount; weights = nothing, entropy_reg::Float64 = 0.0)

"""
function estimate_parameters_cpu(bc::BitCircuit, data, pseudocount; weights = nothing, entropy_reg::Float64 = 0.0)
    # no need to synchronize, since each computation is unique to a decision node
    if weights === nothing && isbinarydata(data)
        node_counts = Vector{UInt}(undef, num_nodes(bc))
    else
        node_counts = Vector{Float64}(undef, num_nodes(bc))
    end

    # Batch the data for a cleaner parameter learning pipeline
    if !isbatched(data)
        data = [data]
        if weights !== nothing
            weights = [weights]
        end
    end

    # rescale pseudocount using the average weight of samples
    if weights !== nothing
        pseudocount = pseudocount * mapreduce(sum, +, weights) / num_examples(data)
    end

    edge_counts::Vector{Float64} = zeros(Float64, num_elements(bc))
    parent_node_counts::Vector{Float64} = zeros(Float64, num_elements(bc))

    @inline function weighted_count_ones(bits::UInt64, start_idx::Number, end_idx::Number, weights)
        count::Float64 = 0.0
        op_bits::UInt64 = bits
        @inbounds for idx = 0 : (end_idx - start_idx)
            count += weights[start_idx + idx] * (op_bits & UInt64(0x1))
            op_bits = (op_bits >> 1)
        end
        count
    end

    @inline function on_node(flows, values, dec_id, weights::Nothing)
        node_counts[dec_id] = sum(1:size(flows,1)) do i
            count_ones(flows[i, dec_id])
        end
    end
    @inline function on_node(flows, values, dec_id, weights)
        node_counts[dec_id] = sum(1:size(flows,1)) do i
            weighted_count_ones(flows[i, dec_id], i * 64 - 63, min(i * 64, length(weights)), weights)
        end
    end
    @inline function on_node(flows, values::Matrix{<:AbstractFloat}, dec_id, weights::Nothing)
        node_counts[dec_id] = sum(1:size(flows,1)) do i
            flows[i, dec_id]
        end
    end
    @inline function on_node(flows, values::Matrix{<:AbstractFloat}, dec_id, weights)
        node_counts[dec_id] = sum(1:size(flows,1)) do i
            flows[i, dec_id] * weights[i]
        end
    end

    @inline function estimate(element, decision, edge_count)
        edge_counts[element] += edge_count
        parent_node_counts[element] += node_counts[decision]
    end

    @inline function on_edge(flows, values, prime, sub, element, grandpa, single_child, weights::Nothing)
        if !single_child
            edge_count = sum(1:size(flows,1)) do i
                count_ones(values[i, prime] & values[i, sub] & flows[i, grandpa])
            end
            estimate(element, grandpa, edge_count)
        end # no need to estimate single child params, they are always prob 1
    end
    @inline function on_edge(flows, values, prime, sub, element, grandpa, single_child, weights)
        if !single_child
            edge_count = sum(1:size(flows,1)) do i
                weighted_count_ones(values[i, prime] & values[i, sub] & flows[i, grandpa],
                                    i * 64 - 63, min(i * 64, length(weights)), weights)
            end
            estimate(element, grandpa, edge_count)
        end # no need to estimate single child params, they are always prob 1
    end
    @inline function on_edge(flows, values::Matrix{<:AbstractFloat}, prime, sub, element, grandpa, single_child, weights::Nothing)
        if !single_child
            edge_count = sum(1:size(flows,1)) do i
                if values[i, grandpa] == 0.0
                    0.0
                else
                    values[i, prime] * values[i, sub] / values[i, grandpa] * flows[i, grandpa]
                end
            end
            estimate(element, grandpa, edge_count)
        end # no need to estimate single child params, they are always prob 1
    end
    @inline function on_edge(flows, values::Matrix{<:AbstractFloat}, prime, sub, element, grandpa, single_child, weights)
        if !single_child
            edge_count = sum(1:size(flows,1)) do i
                if values[i, grandpa] == 0.0
                    0.0
                else
                    values[i, prime] * values[i, sub] / values[i, grandpa] * flows[i, grandpa] * weights[i]
                end
            end
            estimate(element, grandpa, edge_count)
        end # no need to estimate single child params, they are always prob 1
    end

    v, f = nothing, nothing
    if weights !== nothing
        map(zip(data, weights)) do (d, w)
            v, f = satisfies_flows(bc, d, v, f; on_node = on_node, on_edge = on_edge, weights = w)
        end
    else
        map(data) do d
            v, f = satisfies_flows(bc, d, v, f; on_node = on_node, on_edge = on_edge, weights = nothing)
        end
    end

    # Entropy regularization
    if entropy_reg > 1e-8
        total_data_counts = (weights === nothing) ? Float64(num_examples(data)) : mapreduce(sum, +, weights)
        apply_entropy_reg_cpu(bc; log_params = edge_counts, edge_counts, total_data_counts, pseudocount, entropy_reg)
    else
        # Reuse `edge_counts` to store log_params to save space and time.
        for i = 1 : num_elements(bc)
            @inbounds edge_counts[i] = log((edge_counts[i] + pseudocount / num_elements(bc.nodes, bc.elements[1, i])) / (parent_node_counts[i] + pseudocount))
        end

        edge_counts # a.k.a. log_params
    end
end

"""
    estimate_parameters_gpu(bc::BitCircuit, data, pseudocount; weights = nothing, entropy_reg::Float64 = 0.0)

"""
function estimate_parameters_gpu(bc::BitCircuit, data, pseudocount; weights = nothing, entropy_reg::Float64 = 0.0)
    # rescale pseudocount using the average weight of samples
    if weights !== nothing
        if isbatched(data)
            pseudocount = pseudocount * mapreduce(sum, +, weights) / num_examples(data)
        else
            pseudocount = pseudocount * sum(weights) / size(weights, 1)
        end
    end

    node_counts::CuVector{Float32} = CUDA.zeros(Float32, num_nodes(bc))
    edge_counts::CuVector{Float32} = CUDA.zeros(Float32, num_elements(bc))
    # need to manually cudaconvert closure variables
    node_counts_device = CUDA.cudaconvert(node_counts)
    edge_counts_device = CUDA.cudaconvert(edge_counts)

    @inline function on_node(flows, values, dec_id, chunk_id, flow, weight::Nothing)
        c::Float32 = CUDA.count_ones(flow) # cast for @atomic to be happy
        CUDA.@atomic node_counts_device[dec_id] += c
    end
    @inline function on_node(flows, values, dec_id, bit_idx, flow, weight::Float32)
        c::Float32 = ((flow >> bit_idx) & UInt64(0x1)) * weight # cast for @atomic to be happy
        CUDA.@atomic node_counts_device[dec_id] += c
    end
    @inline function on_node(flows, values, dec_id, bit_idx, flow::AbstractFloat, weight::Nothing)
        c::Float32 = flow # cast for @atomic to be happy
        CUDA.@atomic node_counts_device[dec_id] += c
    end
    @inline function on_node(flows, values, dec_id, bit_idx, flow::AbstractFloat, weight::Float32)
        c::Float32 = flow * weight # cast for @atomic to be happy
        CUDA.@atomic node_counts_device[dec_id] += c
    end

    @inline function on_edge(flows, values, prime, sub, element, grandpa, chunk_id, edge_flow, single_child, weight::Nothing)
        if !single_child
            c::Float32 = CUDA.count_ones(edge_flow) # cast for @atomic to be happy
            CUDA.@atomic edge_counts_device[element] += c
        end
    end
    @inline function on_edge(flows, values, prime, sub, element, grandpa, bit_idx, edge_flow, single_child, weight::Float32)
        if !single_child
            c::Float32 = ((edge_flow >> bit_idx) & UInt64(0x1)) * weight # cast for @atomic to be happy
            CUDA.@atomic edge_counts_device[element] += c
        end
    end
    @inline function on_edge(flows, values, prime, sub, element, grandpa, sample_idx,
                             edge_flow::AbstractFloat, single_child, weight::Nothing)
        if !single_child
            c::Float32 = edge_flow # cast for @atomic to be happy
            CUDA.@atomic edge_counts_device[element] += c
        end
    end
    @inline function on_edge(flows, values, prime, sub, element, grandpa, sample_idx,
                             edge_flow::AbstractFloat, single_child, weight::Float32)
        if !single_child
            c::Float32 = edge_flow * weight # cast for @atomic to be happy
            CUDA.@atomic edge_counts_device[element] += c
        end
    end

    if isbatched(data)
        v, f = nothing, nothing
        if weights !== nothing
            map(zip(data, weights)) do (d, w)
                if w !== nothing
                    w = to_gpu(w)
                end
                v, f = satisfies_flows(to_gpu(bc), to_gpu(d), v, f; on_node = on_node, on_edge = on_edge, weights = w)
            end
        else
            map(data) do d
                v, f = satisfies_flows(to_gpu(bc), to_gpu(d), v, f; on_node = on_node, on_edge = on_edge, weights = nothing)
            end
        end
    else
        if weights !== nothing
            weights = to_gpu(weights)
        end
        v, f = satisfies_flows(to_gpu(bc), to_gpu(data); on_node = on_node, on_edge = on_edge, weights = weights)
    end

    CUDA.unsafe_free!(v) # save the GC some effort
    CUDA.unsafe_free!(f) # save the GC some effort

    # Entropy regularization
    if entropy_reg > 1e-8
        total_data_counts = (weights === nothing) ? Float64(num_examples(data)) : mapreduce(sum, +, weights)
        apply_entropy_reg_gpu(bc; log_params = edge_counts, edge_counts, total_data_counts, pseudocount, entropy_reg)
    else
        # TODO: reinstate simpler implementation once https://github.com/JuliaGPU/GPUArrays.jl/issues/313 is fixed and released
        @inbounds parents = bc.elements[1,:]
        @inbounds parent_counts = node_counts[parents]
        @inbounds parent_elcount = bc.nodes[2,parents] .- bc.nodes[1,parents] .+ 1
        params = log.((edge_counts .+ (pseudocount ./ parent_elcount))
                        ./ (parent_counts .+ pseudocount))

        to_cpu(params)
    end
end

"""
    uniform_parameters!(pc::ProbCircuit; perturbation::Float64 = 0.0)

Assign parameters of ProbCircuit so that we have a Uniform distribution.
"""
function uniform_parameters!(pc::ProbCircuit; perturbation::Float64 = 0.0)
    foreach(pc) do pn
        if is⋁gate(pn)
            if num_children(pn) == 1
                pn.log_probs .= 0.0
            else
                if perturbation < 1e-8
                    pn.log_probs .= log.(ones(Float64, num_children(pn)) ./ num_children(pn))
                else
                    unnormalized_probs = map(x -> 1.0 - perturbation + x * 2 * perturbation, rand(num_children(pn)))
                    pn.log_probs .= log.(unnormalized_probs ./ sum(unnormalized_probs))
                end
            end
        end
    end
end

"""
    estimate_parameters_em!(pc::ProbCircuit, data; pseudocount::Float64, entropy_reg::Float64 = 0.0, exp_update_factor::Float64 = 0.0, update_per_batch::Bool = false)

One epoch of Expectation maximization (EM) parameter learning for circuits. Useful when having missing data or non-deterministic circuits.

- `entropy_reg`: Entropy Regularization
- `update_per_batch`: Whether to update circuit paramters per batch (using the ParamBitCircuit's paramters)
"""
function estimate_parameters_em!(pc::ProbCircuit, data; pseudocount::Float64, entropy_reg::Float64 = 0.0,
                                exp_update_factor::Float64 = 0.0, update_per_batch::Bool = false)
    if update_per_batch && isbatched(data)
        estimate_parameters_em_per_batch!(pc, data; pseudocount, entropy_reg, exp_update_factor)
    else
        weights = nothing
        if isweighted(data)
            # `data' is weighted according to its `weight' column
            data, weights = split_sample_weights(data)
        end

        pbc = ParamBitCircuit(pc, data)
        params = if isgpu(data)
            estimate_parameters_gpu(to_gpu(pbc), data, pseudocount; weights, entropy_reg)
        else
            estimate_parameters_cpu(pbc, data, pseudocount; weights, entropy_reg)
        end

        update_pc_params_from_pbc!(pc, pbc.bitcircuit, params; exp_update_factor)
        params
    end
end

"""
    estimate_parameters_em_multi_epochs!(circuit::ProbCircuit, train_data)

Runs multiple epochs of EM for circuit. It will run on CPU/GPU based on where train_data is.

Arguments:
- `circuit`:
- `train_data`: training data. If want gpu, move to gpu before calling this `to_gpu(train_data)`.

Keyword arguments:
- `valid_data=nothing`: validation data, should be same device as `train_data`
- `test_data=nothing`: test data, data, should be same device as `train_data`
- `entropy_reg::Float64 = 0.0`: Entropy regularization
- `exp_update_factor_start::Float64 = 0.1`:
- `exp_update_factor_end::Float64 = 0.9`:
- `em_warmup_iters=100` : number of EM Warm up iterations
- `em_finetune_iters=100`: number of EM fine tune iterations
- `verbose=false`: verbose or not
- `verbose_log_rate=1`: how often to print info
- `save_path=nothing`: path of the file to save the partially trained circuit
- `save_rate=20`: how often to save the circuit to file
"""
function estimate_parameters_em_multi_epochs!(circuit::ProbCircuit, train_data; pseudocount::Float64, entropy_reg::Float64 = 0.0,
                                                exp_update_factor_start::Float64 = 0.1, exp_update_factor_end::Float64 = 0.9,
                                                em_warmup_iters=100, em_finetune_iters=100,
                                                verbose=false, verbose_log_rate=1,
                                                valid_data=nothing, test_data=nothing,
                                                save_path=nothing, save_rate=20)

    if verbose
        println("Learn circuit paramters using estimate_parameters_em_multi_epochs!")
        println("Train Size [$(num_examples(train_data)), $(num_features(train_data))]")
        println("nodes: $(num_nodes(circuit)); edges: $(num_edges(circuit)); parameters: $(num_parameters(circuit))")
    end

    reuse_counts = isgpu(train_data) ? (nothing, nothing) : (nothing, nothing, nothing, nothing, nothing)
    reuse_v = nothing
    reuse_f = nothing

    pbc = ParamBitCircuit(circuit, train_data)
    if isgpu(train_data)
        pbc = to_gpu(ParamBitCircuit(circuit, train_data))
    end

    if verbose
        println("PBC Layers: ", size(pbc.bitcircuit.layers))
        println([length(layer) for layer in pbc.bitcircuit.layers])
        println("----Node Stats-----")
        for stat in node_stats(circuit)
            println(stat)
        end
        println("----Parent Stats-----")
        for stat in parent_stats(circuit)
            println(stat)
        end
    end

    epoch_callback(iter_name, iter, total_iter, t) = begin
        if verbose & (iter % verbose_log_rate == 0)
            t = round(t, digits=2)
            print("[$iter_name Iter $(iter)/$(total_iter); last iter took $(t)s]")
            t = @elapsed begin
                train_ll = marginal_log_likelihood_avg(pbc, train_data)
                valid_ll = isnothing(valid_data) ? nothing : marginal_log_likelihood_avg(pbc, valid_data)
                test_ll  = isnothing(test_data)  ? nothing : marginal_log_likelihood_avg(pbc, test_data)
            end
            t = round(t, digits=2)
            println("[Marginal log-likelihoods Avg last iter took $(t)s; train $(train_ll), valid $(valid_ll), test $(test_ll)]")
        end
        if !isnothing(save_path) && (iter % save_rate == 0)
            update_pc_params_from_pbc!(circuit, pbc)
            if verbose
                println("- Saving circuit at $(save_path)")
                write(save_path, circuit)
            end
        end
    end

    for iter = 1 : em_warmup_iters
        t = @elapsed begin
            exp_update_factor = exp_update_factor_start + (iter - 1) * (exp_update_factor_end - exp_update_factor_start) / (em_warmup_iters - 1)

            _, reuse_v, reuse_f, reuse_counts = estimate_parameters_em_per_batch!(circuit, pbc, train_data; reuse_v, reuse_f, reuse_counts, pseudocount, entropy_reg, exp_update_factor)
        end

        epoch_callback("EM Warmup", iter, em_warmup_iters, t)
    end

    for iter = 1 : em_finetune_iters
        t = @elapsed begin
            _, reuse_v, reuse_f, reuse_counts = estimate_parameters_em!(pbc, train_data; reuse_v, reuse_f, reuse_counts, pseudocount) # Full-batch EM
        end
        epoch_callback("EM Finetune", iter, em_finetune_iters, t)
    end

    update_pc_params_from_pbc!(circuit, pbc)
    nothing # return nothing
end

"""
    estimate_parameters_em_per_batch!(pc::ProbCircuit, data; pseudocount::Float64, entropy_reg::Float64 = 0.0, exp_update_factor = 0.0)

One epoch of Expectation maximization (EM) parameter learning for circuits. Useful when having missing data or non-deterministic circuits.

- `entropy_reg`: Entropy Regularization
- `update_per_batch`: Whether to update circuit paramters per batch (using the ParamBitCircuit's paramters)
"""
function estimate_parameters_em_per_batch!(pc::ProbCircuit, data; pseudocount::Float64, entropy_reg::Float64 = 0.0, exp_update_factor = 0.0)
    pbc = ParamBitCircuit(pc, data)
    if isgpu(data)
        pbc = to_gpu(pbc)
    end
    reuse_v, reuse_f = nothing, nothing
    reuse_counts = isgpu(data) ? (nothing, nothing) : (nothing, nothing, nothing, nothing, nothing)

    params, _, _ , _ = estimate_parameters_em_per_batch!(pc, pbc, data; pseudocount, entropy_reg, exp_update_factor, reuse_v, reuse_f, reuse_counts)
    params
end

function estimate_parameters_em_per_batch!(pc::ProbCircuit, pbc::ParamBitCircuit, data; pseudocount::Float64,
                                                    entropy_reg::Float64 = 0.0, exp_update_factor = 0.0,
                                                    reuse_v=nothing, reuse_f=nothing, reuse_counts=nothing)

    for idx = 1 : length(data)
        pbc, reuse_v, reuse_f, reuse_counts = estimate_parameters_em!(pbc, data[idx]; pseudocount, entropy_reg,
                                                                     reuse_v, reuse_f, reuse_counts,
                                                                     exp_update_factor)
    end

    update_pc_params_from_pbc!(pc, pbc)

    pbc.params, reuse_v, reuse_f, reuse_counts
end

function estimate_parameters_em!(pbc::ParamBitCircuit, data; pseudocount::Float64, entropy_reg::Float64 = 0.0,
                                reuse_v = nothing, reuse_f = nothing, reuse_counts = nothing,
                                exp_update_factor = 0.0
                               )

    @assert (isgpu(data) && isgpu(pbc)) || (!isgpu(data) && !isgpu(pbc)) "Both data and pbc must be on same device. CPU vs GPU"

    weights = nothing
    if isweighted(data)
        # `data' is weighted according to its `weight' column
        data, weights = split_sample_weights(data)
    end

    if reuse_counts === nothing
        reuse_counts = isgpu(data) ? (nothing, nothing) : (nothing, nothing, nothing, nothing, nothing)
    end

    params, v, f, reuse_counts = if isgpu(data)
        estimate_parameters_gpu(pbc, data, pseudocount; weights, reuse = true, reuse_v, reuse_f, reuse_counts, entropy_reg)
    else
        estimate_parameters_cpu(pbc, data, pseudocount; weights, reuse = true, reuse_v, reuse_f, reuse_counts, entropy_reg)
    end

    # Update the parameters to `pbc`
    if isgpu(data) # GPU
        tempparam = Vector{Float64}(undef, length(params))
        tempparam .= to_cpu(params)
        @inbounds @views pbc.params .+= log(exp_update_factor)
        @inbounds @views params .+= log(1.0 - exp_update_factor)
        delta = @inbounds @views @. CUDA.ifelse(pbc.params == params, CUDA.zero(params), CUDA.abs(pbc.params - params))
        @inbounds @views @. pbc.params = CUDA.max(pbc.params, params) + CUDA.log1p(CUDA.exp(-delta))

        CUDA.unsafe_free!(params)
        CUDA.unsafe_free!(delta)
    else # CPU
        @inbounds @views pbc.params = logaddexp.(pbc.params .+ log(exp_update_factor), params .+ log(1.0 - exp_update_factor))
    end

    pbc, v, f, reuse_counts
end

function estimate_parameters_cpu(pbc::ParamBitCircuit, data, pseudocount; weights = nothing, reuse::Bool = false,
                                 reuse_v = nothing, reuse_f = nothing,
                                 reuse_counts = (nothing, nothing, nothing, nothing, nothing), entropy_reg::Float64 = 0.0)
    # `data` is batched?
    data_batched = isbatched(data)

    # no need to synchronize, since each computation is unique to a decision node
    node_counts::Vector{Float64} = similar!(reuse_counts[1], Vector{Float64}, num_nodes(pbc.bitcircuit))
    @inbounds @views @. @avx node_counts .= typemin(eltype(node_counts)) # log(0)

    # rescale pseudocount using the average weight of samples
    if weights !== nothing
        if data_batched
            pseudocount = pseudocount * mapreduce(sum, +, weights) / num_examples(data)
        else
            pseudocount = pseudocount * sum(weights) / size(weights, 1)
        end
    end

    bc = pbc.bitcircuit
    params = pbc.params

    edge_counts::Vector{Float64} = similar!(reuse_counts[2], Vector{Float64}, num_elements(bc))
    parent_node_counts::Vector{Float64} = similar!(reuse_counts[3], Vector{Float64}, num_elements(bc))
    @inbounds @views edge_counts[:] .= zero(Float64)
    @inbounds @views parent_node_counts[:] .= zero(Float64)

    # Buffer to save some allocations
    buffer::Vector{Float64} = similar!(reuse_counts[4], Vector{Float64},
        data_batched ? num_examples(data[1]) : num_examples(data))

    # Pre-compute log(weights) to save computation time
    if weights !== nothing
        log_weights = similar!(reuse_counts[5], Vector{Float64},
            data_batched ? num_examples(data[1]) : num_examples(data))

        if !data_batched
            @inbounds @views @avx log_weights .= log.(weights)
        end
    else
        log_weights = nothing
    end

    # For batched dataset, we want to enable 'estimate' for parent_node_counts only in the last minibatch
    estimate_flag = true

    @inline function on_node(flows, values, dec_id, weights::Nothing)
        @inbounds @views @. @avx buffer = flows[:, dec_id]
        node_counts[dec_id] = logaddexp(node_counts[dec_id], logsumexp(buffer))
    end
    @inline function on_node(flows, values, dec_id, weights)
        @inbounds @views @. @avx buffer = flows[:, dec_id] + log_weights[:]
        node_counts[dec_id] = logaddexp(node_counts[dec_id], logsumexp(buffer))
    end

    @inline function estimate(element, decision, edge_count)
        edge_counts[element] += exp(edge_count)
        if estimate_flag # For batched dataset, we only accumulate parent_node_counts after all node_counts have been cumulated
            parent_node_counts[element] += exp(node_counts[decision])
        end
    end

    @inline function on_edge(flows, values, prime, sub, element, grandpa, single_child, weights::Nothing)
        θ = eltype(flows)(params[element])
        if !single_child
            @inbounds @views @. @avx buffer = values[:, prime] + values[:, sub] - values[:, grandpa] + flows[:, grandpa] + θ
            buffer .= ifelse.(isnan.(buffer[:]), typemin(eltype(flows)), buffer)

            edge_count = logsumexp(buffer)

            estimate(element, grandpa, edge_count)
        end # no need to estimate single child params, they are always prob 1
    end
    @inline function on_edge(flows, values, prime, sub, element, grandpa, single_child, weights)
        θ = eltype(flows)(params[element])
        if !single_child
            @inbounds @views @. @avx buffer = values[:, prime] + values[:, sub] - values[:, grandpa] + flows[:, grandpa] + θ + log_weights
            buffer .= ifelse.(isnan.(buffer), typemin(eltype(flows)), buffer)

            edge_count = logsumexp(buffer)

            estimate(element, grandpa, edge_count)
        end # no need to estimate single child params, they are always prob 1
    end

    if data_batched
        if weights !== nothing
            v, f = reuse_v, reuse_f
            for idx = 1 : length(data)
                d = data[idx]
                w = weights[idx]

                # Resize buffer if the current minibatch has a different size
                if size(buffer, 1) != num_examples(d)
                    resize!(buffer, num_examples(d))
                end

                # Resize log_weights if the current minibatch has a different size
                if data_batched
                    if size(log_weights, 1) != num_examples(d)
                        resize!(log_weights, num_examples(d))
                    end
                    @inbounds @views @avx log_weights .= log.(w)
                end

                # Only accumulate parent_node_counts in the final call to marginal_flows
                if idx == length(data)
                    estimate_flag = true
                else
                    estimate_flag = false
                end

                v, f = marginal_flows(pbc, d, v, f; on_node = on_node, on_edge = on_edge, weights = w)
            end
        else
            v, f = reuse_v, reuse_f
            for idx = 1 : length(data)
                d = data[idx]

                # Resize buffer if the current minibatch has a different size
                if size(buffer, 1) != num_examples(d)
                    resize!(buffer, num_examples(d))
                end

                if idx == length(data)
                    estimate_flag = true
                else
                    estimate_flag = false
                end

                v, f = marginal_flows(pbc, d, v, f; on_node = on_node, on_edge = on_edge, weights = nothing)
            end
        end
    else
        v, f = marginal_flows(pbc, data, reuse_v, reuse_f; on_node, on_edge, weights)
    end

    if entropy_reg > 1e-8
        total_data_counts = (weights === nothing) ? Float64(num_examples(data)) : mapreduce(sum, +, weights)
        edge_counts = apply_entropy_reg_cpu(bc; log_params = edge_counts, edge_counts, total_data_counts,
                                            pseudocount, entropy_reg)
    else
        # `edge_counts` now becomes "params"
        @simd for i = 1 : num_elements(bc)
            num_els = num_elements(bc.nodes, bc.elements[1, i])
            if num_els == 1
                @inbounds edge_counts[i] = zero(eltype(edge_counts)) # log(1)
            else
                @inbounds edge_counts[i] = log((edge_counts[i] + pseudocount / num_elements(bc.nodes, bc.elements[1, i])) / (parent_node_counts[i] + pseudocount))
            end
        end
    end

    if reuse
        # Also return the allocated vars v, f, counts for future reuse
        edge_counts, v, f, (node_counts, edge_counts, parent_node_counts, buffer, log_weights)
    else
        edge_counts # a.k.a. log_probs
    end
end

function estimate_parameters_gpu(pbc::ParamBitCircuit, data, pseudocount; weights = nothing, reuse::Bool = false,
                                 reuse_v = nothing, reuse_f = nothing, reuse_counts = (nothing, nothing),
                                 entropy_reg::Float64 = 0.0)
    # `data` is batched?
    data_batched = isbatched(data)

    # rescale pseudocount using the average weight of samples
    if weights !== nothing
        if data_batched
            if isgpu(weights)
                pseudocount = pseudocount * mapreduce(sum, +, to_cpu(weights)) / num_examples(data)
            else
                pseudocount = pseudocount * mapreduce(sum, +, weights) / num_examples(data)
            end
        else
            pseudocount = pseudocount * sum(weights) / size(weights, 1)
        end
    end

    bc = pbc.bitcircuit
    node_counts::CuVector{Float64} = similar!(reuse_counts[1], CuVector{Float64}, num_nodes(bc))
    edge_counts::CuVector{Float64} = similar!(reuse_counts[2], CuVector{Float64}, num_elements(bc))
    @inbounds @views node_counts[:] .= zero(Float64)
    @inbounds @views edge_counts[:] .= zero(Float64)
    # need to manually cudaconvert closure variables
    node_counts_device = CUDA.cudaconvert(node_counts)
    edge_counts_device = CUDA.cudaconvert(edge_counts)

    @inline function on_node(flows, values, dec_id, chunk_id, flow, weight::Nothing)
        c::Float64 = exp(flow) # cast for @atomic to be happy
        CUDA.@atomic node_counts_device[dec_id] += c
    end
    @inline function on_node(flows, values, dec_id, chunk_id, flow, weight::AbstractFloat)
        c::Float64 = exp(flow) * weight # cast for @atomic to be happy
        CUDA.@atomic node_counts_device[dec_id] += c
    end

    @inline function on_edge(flows, values, prime, sub, element, grandpa, chunk_id, edge_flow, single_child, weight::Nothing)
        if !single_child
            c::Float64 = exp(edge_flow) # cast for @atomic to be happy
            CUDA.@atomic edge_counts_device[element] += c
        end
    end
    @inline function on_edge(flows, values, prime, sub, element, grandpa, chunk_id, edge_flow, single_child, weight::AbstractFloat)
        if !single_child
            c::Float64 = exp(edge_flow) * weight # cast for @atomic to be happy
            CUDA.@atomic edge_counts_device[element] += c
        end
    end

    if data_batched
        v, f = reuse_v, reuse_f
        if !isnothing(weights)
            map(zip(data, weights)) do (d, w)
                if !isnothing(w)
                    w = to_gpu(w)
                end
                v, f = marginal_flows(pbc, d, v, f; on_node = on_node, on_edge = on_edge, weights = w)

                nothing # Return nothing to save some time
            end
        else
            map(data) do d
                v, f = marginal_flows(pbc, d, v, f; on_node = on_node, on_edge = on_edge, weights = nothing)

                nothing # Return nothing to save some time
            end

            nothing # Return nothing to save some time
        end
    else
        if !isnothing(weights)
            weights = to_gpu(weights)
        end

        v, f = marginal_flows(pbc, data, reuse_v, reuse_f; on_node, on_edge, weights)
    end

    if entropy_reg > 1e-8
        total_data_counts = (weights === nothing) ? Float64(num_examples(data)) : mapreduce(sum, +, weights)
        params = apply_entropy_reg_gpu(bc; log_params = edge_counts, edge_counts, total_data_counts, pseudocount, entropy_reg)
    else
        # TODO: reinstate simpler implementation once https://github.com/JuliaGPU/GPUArrays.jl/issues/313 is fixed and released
        @inbounds parents = bc.elements[1,:]
        @inbounds parent_counts = node_counts[parents]
        @inbounds parent_elcount = bc.nodes[2,parents] .- bc.nodes[1,parents] .+ 1
        params = log.((edge_counts .+ (pseudocount ./ parent_elcount))
                        ./ (parent_counts .+ pseudocount))
        params = @inbounds @views @. ifelse(parent_elcount == 1, zero(params), params)

        CUDA.unsafe_free!(parents)
        CUDA.unsafe_free!(parent_counts)
        CUDA.unsafe_free!(parent_elcount)
    end

    # Only free the memory if the reuse memory is not provided
    if !reuse
        CUDA.unsafe_free!(v) # save the GC some effort
        CUDA.unsafe_free!(f) # save the GC some effort
        CUDA.unsafe_free!(node_counts) # save the GC some effort
        CUDA.unsafe_free!(edge_counts) # save the GC some effort
    end

    if reuse
        # Also return the allocated vars v, f, counts for future reuse
        if !isgpu(params)
            to_gpu(params), v, f, (node_counts, edge_counts)
        else
            params, v, f, (node_counts, edge_counts)
        end
    else
        to_cpu(params) # a.k.a. log_probs
    end
end

"""
    apply_entropy_reg_cpu(bc::BitCircuit; log_params::Vector{Float64}, edge_counts::Vector{Float64}, total_data_counts::Float64, pseudocount::Float64 = 0.1, entropy_reg::Float64 = 0.0)::Vector{Float64}

Add entropy regularization to a deterministic (see LogicCircuits.isdeterministic) probabilistic
circuit. `alpha` is a hyperparameter that balances the weights between the likelihood and the
entropy: \argmax_{\theta} L(\theta) = ll_mean(\theta) + alpha * entropy(\theta).
"""
function apply_entropy_reg_cpu(bc::BitCircuit; log_params::Vector{Float64}, edge_counts::Vector{Float64},
                               total_data_counts::Float64, pseudocount::Float64 = 0.1,
                               entropy_reg::Float64 = 0.0)::Vector{Float64}
    a = Vector{Float64}(undef, 0) # For reuse
    log_probs = Vector{Float64}(undef, 0) # For reuse
    p_exp_logprob = Vector{Float64}(undef, 0) # For reuse

    node_log_probs = Vector{Float64}(undef, num_nodes(bc)) # Probability of each decision/sum node
    @inbounds @views node_log_probs .= -Inf
    @inbounds node_log_probs[end] = 0.0

    for i = num_nodes(bc) : -1 : 1 # Traverse the circuit top-down
        @inbounds child_ele_start = bc.nodes[1,i]
        @inbounds child_ele_end = bc.nodes[2,i]
        num_eles = child_ele_end - child_ele_start + 1

        if num_eles > 1 # No need to do anything if the decision node has 1 child.
            resize!(a, num_eles)
            resize!(log_probs, num_eles)
            resize!(p_exp_logprob, num_eles)

            @inbounds @views a .= edge_counts[child_ele_start: child_ele_end] .+ (pseudocount / num_eles)
            @inbounds b = entropy_reg * exp(node_log_probs[i] + log(total_data_counts + pseudocount))

            @inbounds @views log_probs .= log.(edge_counts[child_ele_start: child_ele_end] .+ (pseudocount / num_eles))
            @inbounds @views log_probs .-= logsumexp(log_probs)

            for _ = 1 : 3
                y = sum(b .* log_probs .- a .* exp.(-log_probs)) / num_eles
                for _ = 1 : 4
                    @inbounds @views p_exp_logprob .= a .* exp.(-log_probs)
                    @inbounds @views log_probs .+= (p_exp_logprob .- b .* log_probs .+ y) ./ (p_exp_logprob .+ b .+ 1e-4)
                    @inbounds @views log_probs .-= logsumexp(log_probs) # Project the parameters back to its valid space
                end
            end

            # Update parameters
            @inbounds @views log_params[child_ele_start: child_ele_end] .= log_probs
        end

        # Update child nodes' log_prob
        if child_ele_start >= 1
            for child_ele_id = child_ele_start : child_ele_end
                prime_idx = bc.elements[2, child_ele_id]
                sub_idx = bc.elements[3, child_ele_id]
                node_log_probs[prime_idx] = logaddexp(node_log_probs[prime_idx], node_log_probs[i] + log_params[child_ele_id])
                node_log_probs[sub_idx] = logaddexp(node_log_probs[sub_idx], node_log_probs[i] + log_params[child_ele_id])
            end
        end
    end

    log_params
end

"""
    apply_entropy_reg_gpu(bc::BitCircuit; log_params, edge_counts, total_data_counts::Float64, pseudocount::Float64 = 0.1, entropy_reg::Float64 = 0.0)

Add entropy regularization to a deterministic (see LogicCircuits.isdeterministic) probabilistic
circuit. `alpha` is a hyperparameter that balances the weights between the likelihood and the
entropy: \argmax_{\theta} L(\theta) = ll_mean(\theta) + alpha * entropy(\theta).
"""
function apply_entropy_reg_gpu(bc::BitCircuit; log_params, edge_counts,
                               total_data_counts::Float64, pseudocount::Float64 = 0.1, entropy_reg::Float64 = 0.0)
    # TODO: if this is performance-critical, consider change to native GPU code
    log_params = to_cpu(log_params)
    edge_counts = to_cpu(edge_counts)
    bc = to_cpu(bc)

    a = Vector{Float64}(undef, 0) # For reuse
    log_probs = Vector{Float64}(undef, 0) # For reuse
    p_exp_logprob = Vector{Float64}(undef, 0) # For reuse

    node_log_probs = Vector{Float64}(undef, num_nodes(bc)) # Probability of each decision/sum node
    @inbounds @views node_log_probs .= -Inf
    @inbounds node_log_probs[end] = 0.0

    for i = num_nodes(bc) : -1 : 1 # Traverse the circuit top-down
        @inbounds child_ele_start = bc.nodes[1,i]
        @inbounds child_ele_end = bc.nodes[2,i]
        num_eles = child_ele_end - child_ele_start + 1

        if num_eles > 1 # No need to do anything if the decision node has 1 child.
            resize!(a, num_eles)
            resize!(log_probs, num_eles)
            resize!(p_exp_logprob, num_eles)

            @inbounds @views a .= edge_counts[child_ele_start: child_ele_end] .+ (pseudocount / num_eles)
            @inbounds b = entropy_reg * exp(node_log_probs[i] + log(total_data_counts + pseudocount))

            @inbounds @views log_probs .= log.(edge_counts[child_ele_start: child_ele_end] .+ (pseudocount / num_eles))
            @inbounds @views log_probs .-= logsumexp(log_probs)

            for _ = 1 : 3
                y = sum(b .* log_probs .- a .* exp.(-log_probs)) / num_eles
                for _ = 1 : 4
                    @inbounds @views p_exp_logprob .= a .* exp.(-log_probs)
                    @inbounds @views log_probs .+= (p_exp_logprob .- b .* log_probs .+ y) ./ (p_exp_logprob .+ b .+ 1e-4)
                    @inbounds @views log_probs .-= logsumexp(log_probs) # Project the parameters back to its valid space
                end
            end

            # Update parameters
            @inbounds @views log_params[child_ele_start: child_ele_end] .= log_probs
        end

        # Update child nodes' log_prob
        if child_ele_start >= 1
            for child_ele_id = child_ele_start : child_ele_end
                prime_idx = bc.elements[2, child_ele_id]
                sub_idx = bc.elements[3, child_ele_id]
                node_log_probs[prime_idx] = logaddexp(node_log_probs[prime_idx], node_log_probs[i] + log_params[child_ele_id])
                node_log_probs[sub_idx] = logaddexp(node_log_probs[sub_idx], node_log_probs[i] + log_params[child_ele_id])
            end
        end
    end

    log_params
end

ext: jl


text:
@testset "DWT" begin
    # Single steps (1D)
    x = [2,3,-4,5.0]
    wt = wavelet(WT.db4)
    g, h = WT.makereverseqmfpair(wt, true)
    w₁, w₂ = DWT.dwt_step(x,h,g)
    w₁ = round.(w₁, digits=3); w₂ = round.(w₂, digits=3)
    @test [w₁;w₂] == [-0.524, 4.767, 1.803, 5.268]
    @test round.(DWT.idwt_step(w₁, w₂, h, g), digits=3) == x
    w₁, w₂ = DWT.dwt_step!(zeros(2), zeros(2), x,h,g)
    w₁ = round.(w₁, digits=3); w₂ = round.(w₂, digits=3)
    @test [w₁;w₂] == [-0.524, 4.767, 1.803, 5.268]
    @test round.(DWT.idwt_step!(zeros(4), w₁, w₂, h, g), digits=3) == x

    # Single steps (2D)
    x = [2 3;-4 5.0]
    w₁, w₂, w₃, w₄ = DWT.dwt_step(x, h, g)
    @test round.([w₁ w₂;w₃ w₄], digits=3) == [3 5;-2 4]
    @test round.(DWT.idwt_step(w₁, w₂, w₃, w₄, h, g), digits=3) == x
    w₁, w₂, w₃, w₄ = DWT.dwt_step!(zeros(1,1),zeros(1,1),zeros(1,1),zeros(1,1), x, h, g, zeros(2,2))
    @test round.([w₁ w₂;w₃ w₄], digits=3) == [3 5;-2 4]
    @test round.(DWT.idwt_step!(zeros(2,2), w₁, w₂, w₃, w₄, h, g, zeros(2,2)), digits=3) == x

    # 1D transforms
    x = randn(8)
    y1 = wpt(x, wt, 1)
    y2 = wpt(x, wt, 2)
    y3 = wpt(x, wt, 3)
    @test wpd(x, wt) ≈ [x y1 y2 y3]
    @test wpd!(zeros(8,4), x, wt) ≈ [x y1 y2 y3]
    @test iwpd(wpd(x,wt),wt) ≈ x
    @test iwpd(wpd(x,wt),wt,2) ≈ x
    @test iwpd(wpd(x,wt),wt,maketree(x,:dwt)) ≈ x

    # 2D transforms
    x = randn(8,8)
    y = dwt(x, wt, 1)
    z1 = wpt(x, wt, 1)
    z2 = wpt(x, wt, 2)
    z3 = wpt(x, wt, 3)
    @test y ≈ z1
    @test wpd(x, wt) ≈ cat(x, z1, z2, z3, dims=3)
    @test dwt(x, wt) ≈ wpt(x, wt, maketree(x,:dwt))
    @test iwpd(wpd(x,wt),wt) ≈ x
    @test iwpd(wpd(x,wt),wt,2) ≈ x
    @test iwpd(wpd(x,wt),wt,maketree(x,:dwt)) ≈ x
    @test iwpt(wpt(x,wt),wt) ≈ x
    @test iwpt(wpt(x,wt,2),wt,2) ≈ x
    @test iwpt(wpt(x,wt,maketree(x,:dwt)),wt,maketree(x,:dwt)) ≈ x
end

@testset "SWT" begin
    # Single step
    x = [2,3,-4,5.0]
    wt = wavelet(WT.db4)
    g, h = WT.makereverseqmfpair(wt, true)
    w₁, w₂ = SWT.sdwt_step(x, 0, h, g)
    w₁ = round.(w₁,digits=3); w₂ = round.(w₂,digits=3)
    @test [w₁ w₂] == [3.854 -6.181;-0.524 1.803;0.389 -0.89;4.767 5.268]
    @test round.(SWT.isdwt_step(w₁, w₂, 0, h, g), digits=3) == x
    @test round.(SWT.isdwt_step(w₁, w₂, 0, 0, 0, h, g), digits=3) == x
    @test round.(SWT.isdwt_step(w₁, w₂, 0, 0, 1, h, g), digits=3) == x
    @test_throws AssertionError SWT.isdwt_step(w₁, w₂, 0, -1, 0, h, g)
    @test_throws AssertionError SWT.isdwt_step(w₁, w₂, 0, 1, 0, h, g)
    @test_throws AssertionError SWT.isdwt_step(w₁, w₂, 0, 0, 2, h, g)

    # Single steps (2D)
    x = [2 3;-4 5.0]
    w₁, w₂, w₃, w₄ = SWT.sdwt_step(x, 0, h, g)
    @test round.(w₁, digits=3) == [3 3;3 3]
    @test round.(w₂, digits=3) == [-5 5;-5 5]
    @test round.(w₃, digits=3) == [2 2;-2 -2]
    @test round.(w₄, digits=3) == [4 -4;-4 4]
    @test round.(SWT.isdwt_step(w₁, w₂, w₃, w₄, 0, h, g), digits=3) == x
    @test round.(SWT.isdwt_step(w₁, w₂, w₃, w₄, 0, 0, 0, h, g), digits=3) == x
    @test round.(SWT.isdwt_step(w₁, w₂, w₃, w₄, 0, 0, 1, h, g), digits=3) == x
    w₁, w₂, w₃, w₄ = SWT.sdwt_step!(zeros(2,2),zeros(2,2),zeros(2,2),zeros(2,2), x, 0, h, g, zeros(2,2,2))
    @test round.(w₁, digits=3) == [3 3;3 3]
    @test round.(w₂, digits=3) == [-5 5;-5 5]
    @test round.(w₃, digits=3) == [2 2;-2 -2]
    @test round.(w₄, digits=3) == [4 -4;-4 4]
    @test round.(SWT.isdwt_step!(zeros(2,2), w₁, w₂, w₃, w₄, 0, h, g, zeros(2,2,2)), digits=3) == x
    @test round.(SWT.isdwt_step!(zeros(2,2), w₁, w₂, w₃, w₄, 0, 0, 0, h, g, zeros(2,2,2)), digits=3) == x
    @test round.(SWT.isdwt_step!(zeros(2,2), w₁, w₂, w₃, w₄, 0, 0, 1, h, g, zeros(2,2,2)), digits=3) == x
    @test_throws AssertionError SWT.isdwt_step!(zeros(2,2), w₁, w₂, w₃, w₄, 0, 0, 2, h, g, zeros(2,2,2))
    @test_throws AssertionError SWT.isdwt_step!(zeros(2,2), w₁, w₂, w₃, w₄, 0, -1, 1, h, g, zeros(2,2,2))
    @test_throws AssertionError SWT.isdwt_step!(zeros(2,2), w₁, w₂, w₃, w₄, 0, 1, 0, h, g, zeros(2,2,2))

    # 1D transforms
    x = randn(8)
    tree = maketree(x, :dwt)
    sm = 3
    @test isdwt(sdwt(x, wt, 3), wt) ≈ x
    @test isdwt(sdwt(x, wt), wt, sm) ≈ x
    @test swpt(x, wt) == swpd(x, wt)[:,8:15]
    @test swpt(x, wt, 3) == swpd(x, wt)[:,8:15]
    @test iswpt(swpt(x, wt), wt) ≈ x
    @test iswpt(swpt(x, wt), wt, sm) ≈ x
    @test iswpd(swpd(x, wt), wt) ≈ x
    @test iswpd(swpd(x, wt), wt, 2) ≈ x
    @test iswpd(swpd(x, wt), wt, tree) ≈ x
    @test iswpd(swpd(x, wt), wt, tree, sm) ≈ x
    @test iswpd(swpd(x, wt), wt, 2, sm) ≈ x

    # 2D transforms
    x = randn(8,8)
    tree = maketree(x, :dwt)
    sm = 3
    @test isdwt(sdwt(x, wt, 3), wt) ≈ x
    @test isdwt(sdwt(x, wt), wt, sm) ≈ x
    @test swpt(x, wt) == swpd(x, wt)[:,:,22:85]
    @test swpt(x, wt, 3) == swpd(x, wt)[:,:,22:85]
    @test iswpt(swpt(x, wt), wt) ≈ x
    @test iswpt(swpt(x, wt), wt, sm) ≈ x
    @test iswpd(swpd(x, wt), wt) ≈ x
    @test iswpd(swpd(x, wt), wt, 2) ≈ x
    @test iswpd(swpd(x, wt), wt, tree) ≈ x
    @test iswpd(swpd(x, wt), wt, tree, sm) ≈ x
    @test iswpd(swpd(x, wt), wt, 2, sm) ≈ x
end

@testset "ACWT" begin
    # Single step
    x = [2,3,-4,5.0]
    wt = wavelet(WT.db4)
    g, h = ACWT.make_acreverseqmfpair(wt)
    w₁, w₂ = ACWT.acdwt_step(x, 0, h, g)
    w₁ = round.(w₁,digits=3); w₂ = round.(w₂,digits=3)
    @test [w₁ w₂] == [4.243 -1.414;1.414 2.828;0 -5.657;2.828 4.243]
    @test round.(ACWT.iacdwt_step(w₁, w₂), digits=3) == x

    # Single steps (2D)
    x = [2 3;-4 5.0]
    w₁, w₂, w₃, w₄ = ACWT.acdwt_step(x, 0, h, g)
    @test round.(w₁, digits=3) == [3 3;3 3]
    @test round.(w₂, digits=3) == [-5 5;-5 5]
    @test round.(w₃, digits=3) == [2 2;-2 -2]
    @test round.(w₄, digits=3) == [4 -4;-4 4]
    @test round.(ACWT.iacdwt_step(w₁, w₂, w₃, w₄), digits=3) == x
    w₁, w₂, w₃, w₄ = ACWT.acdwt_step!(zeros(2,2),zeros(2,2),zeros(2,2),zeros(2,2), x, 0, h, g, zeros(2,2,2))
    @test round.(w₁, digits=3) == [3 3;3 3]
    @test round.(w₂, digits=3) == [-5 5;-5 5]
    @test round.(w₃, digits=3) == [2 2;-2 -2]
    @test round.(w₄, digits=3) == [4 -4;-4 4]
    @test round.(ACWT.iacdwt_step!(zeros(2,2), w₁, w₂, w₃, w₄, zeros(2,2,2)), digits=3) == x

    # acwt (1D)
    x = randn(8)
    wt = wavelet(WT.db4)
    tree = maketree(x, :dwt)
    @test iacdwt(acdwt(x, wt)) ≈ x
    @test iacdwt(acdwt(x, wt, 2), wt) ≈ x
    @test acwpt(x, wt) == acwpd(x, wt)[:,8:15]
    @test acwpt(x, wt, 2) == acwpd(x, wt)[:,4:7]
    @test iacwpt(acwpt(x, wt)) ≈ x
    @test iacwpd(acwpd(x, wt)) ≈ x
    @test iacwpd(acwpd(x, wt), 2) ≈ x
    @test iacwpd(acwpd(x, wt), wt, 2) ≈ x
    @test iacwpd(acwpd(x, wt), tree) ≈ x
    @test iacwpd(acwpd(x, wt), wt, tree) ≈ x
    @test iacwpd!(zeros(8), acwpd(x, wt), wt, tree) ≈ x
    @test_throws AssertionError iacwpd!(zeros(4), acwpd(x, wt), wt, tree)

    # acwt (2D)
    x = randn(8,8)
    tree = maketree(x, :dwt)
    @test iacdwt(acdwt(x, wt, 3), wt) ≈ x
    @test iacdwt(acdwt(x, wt, 3)) ≈ x
    @test acwpt(x, wt) == acwpd(x, wt)[:,:,22:85]
    @test acwpt(x, wt, 3) == acwpd(x, wt)[:,:,22:85]
    @test iacwpt(acwpt(x, wt), wt) ≈ x
    @test iacwpd(acwpd(x, wt)) ≈ x
    @test iacwpd(acwpd(x, wt), wt, tree) ≈ x
    @test iacwpd(acwpd(x, wt), tree) ≈ x
end

@testset "SIWPD" begin
    # siwpd
    x = randn(4)
    wt = wavelet(WT.haar)
    y = siwpd(x, wt, 2, 1)
    y0 = y[1,4:7]
    y1 = y[2,4:7]
    y2 = y[3,4:7]
    y3 = y[4,4:7]
    @test y0 == wpt(x, wt)
    @test !all([isdefined(y1, i) for i in 1:4])
    @test y2 == wpt(circshift(x,2), wt)
    @test !all([isdefined(y3, i) for i in 1:4])
    # make tree
    tree = [
        trues(1),
        repeat([trues(2)],2)...,
        repeat([BitVector([1,0,1,0])],4)...
    ]
    @test makesiwpdtree(4, 2, 1) == tree
end

@testset "Transform All" begin
    x = randn(8)    # 1D
    xₙ = [x x x]
    w = randn(8,8)  # 2D
    wₙ = cat(w,w,w,dims=3)
    wt = wavelet(WT.db4)

    # dwt
    y = dwt(x, wt); z = dwt(w, wt)
    yₙ = [y y y]; zₙ = cat(z,z,z,dims=3)
    @test dwtall(xₙ, wt) == yₙ
    @test dwtall(wₙ, wt) == zₙ
    @test idwtall(yₙ, wt) ≈ xₙ
    @test idwtall(zₙ, wt) ≈ wₙ

    # wpt
    y = wpt(x, wt); z = wpt(w, wt)
    yₙ = [y y y]; zₙ = cat(z,z,z,dims=3)
    @test wptall(xₙ, wt) == yₙ
    @test wptall(wₙ, wt) == zₙ
    @test iwptall(yₙ, wt) ≈ xₙ
    @test iwptall(zₙ, wt) ≈ wₙ

    # wpd
    y = wpd(x,wt); z = wpd(w,wt)
    yₙ = cat(y,y,y, dims=3); zₙ = cat(z,z,z,dims=4)
    @test wpdall(xₙ, wt, 3) ≈ yₙ
    @test wpdall(wₙ, wt, 3) ≈ zₙ
    @test iwpdall(yₙ, wt, 3) ≈ xₙ
    @test iwpdall(zₙ, wt, 3) ≈ wₙ

    # acdwt
    y = acdwt(x, wt); z = acdwt(w,wt)
    yₙ = cat(y,y,y, dims=3); zₙ = cat(z,z,z,dims=4)
    @test acdwtall(xₙ, wt) == yₙ
    @test acdwtall(wₙ, wt) == zₙ
    @test iacdwtall(yₙ, wt) ≈ xₙ
    @test iacdwtall(zₙ, wt) ≈ wₙ

    # acwpt
    y = acwpt(x, wt); z = acwpt(w,wt)
    yₙ = cat(y,y,y, dims=3); zₙ = cat(z,z,z,dims=4)
    @test acwptall(xₙ, wt) == yₙ
    @test acwptall(wₙ, wt) == zₙ
    @test iacwptall(yₙ, wt) ≈ xₙ
    @test iacwptall(zₙ, wt) ≈ wₙ

    # acwpd
    y = acwpd(x, wt); z = acwpd(w, wt)
    yₙ = cat(y,y,y, dims=3); zₙ = cat(z,z,z,dims=4)
    @test acwpdall(xₙ, wt, 3) ≈ yₙ
    @test acwpdall(wₙ, wt, 3) ≈ zₙ
    @test iacwpdall(yₙ, wt, 3) ≈ xₙ
    @test iacwpdall(yₙ, 3) ≈ xₙ
    @test iacwpdall(zₙ, wt, 3) ≈ wₙ
    @test iacwpdall(zₙ, 3) ≈ wₙ
    @test iacwpdall(yₙ, wt, maketree(8,3,:full)) ≈ xₙ
    @test iacwpdall(yₙ, maketree(8,3,:full)) ≈ xₙ
    @test iacwpdall(zₙ, wt, maketree(8,8,3,:full)) ≈ wₙ
    @test iacwpdall(zₙ, maketree(8,8,3,:full)) ≈ wₙ

    # sdwt
    y = sdwt(x, wt); z = sdwt(w,wt)
    yₙ = cat(y,y,y, dims=3); zₙ = cat(z,z,z,dims=4)
    @test sdwtall(xₙ, wt) == yₙ
    @test sdwtall(wₙ, wt) == zₙ
    @test isdwtall(yₙ, wt) ≈ xₙ
    @test isdwtall(yₙ, wt,2) ≈ xₙ
    @test isdwtall(zₙ, wt) ≈ wₙ
    @test isdwtall(zₙ, wt,2) ≈ wₙ
    @test_throws AssertionError isdwtall(yₙ, wt,12)
    @test_throws AssertionError isdwtall(zₙ, wt,12)

    # swpt
    y = swpt(x, wt); z = swpt(w,wt)
    yₙ = cat(y,y,y, dims=3); zₙ = cat(z,z,z,dims=4)
    @test swptall(xₙ, wt) == yₙ
    @test swptall(wₙ, wt) == zₙ
    @test iswptall(yₙ, wt) ≈ xₙ
    @test iswptall(yₙ, wt, 2) ≈ xₙ
    @test iswptall(zₙ, wt) ≈ wₙ
    @test iswptall(zₙ, wt, 2) ≈ wₙ
    @test_throws AssertionError iswptall(yₙ, wt, 12)
    @test_throws AssertionError iswptall(zₙ, wt, 12)

    # swpd
    y = swpd(x, wt); z = swpd(w,wt)
    yₙ = cat(y,y,y, dims=3); zₙ = cat(z,z,z,dims=4)
    @test swpdall(xₙ, wt, 3) ≈ yₙ
    @test swpdall(wₙ, wt, 3) ≈ zₙ
    @test iswpdall(yₙ, wt, 3) ≈ xₙ
    @test iswpdall(zₙ, wt, 3) ≈ wₙ
    @test iswpdall(yₙ, wt, 3, 4) ≈ xₙ
    @test iswpdall(zₙ, wt, 3, 4) ≈ wₙ
end