text:
# # Custom ReLU layer

#md # [![](https://img.shields.io/badge/show-github-579ACA.svg)](@__REPO_ROOT_URL__/docs/src/examples/custom-relu.jl)

# We demonstrate how DiffOpt can be used to generate a simple neural network
# unit - the ReLU layer. A neural network is created using Flux.jl and
# trained on the MNIST dataset.

# This tutorial uses the following packages

using JuMP
import DiffOpt
import Ipopt
import ChainRulesCore
import Flux
import MLDatasets
import Statistics
import Base.Iterators: repeated
using LinearAlgebra

# ## The ReLU and its derivative

# Define a relu through an optimization problem solved by a quadratic solver.
# Return the solution of the problem.
function matrix_relu(
    y::Matrix;
    model = Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
)
    layer_size, batch_size = size(y)
    empty!(model)
    set_silent(model)
    @variable(model, x[1:layer_size, 1:batch_size] >= 0)
    @objective(model, Min, x[:]'x[:] -2y[:]'x[:])
    optimize!(model)
    return value.(x)
end

# Define the reverse differentiation rule, for the function we defined above.
function ChainRulesCore.rrule(::typeof(matrix_relu), y::Matrix{T}) where T
    model = Model(() -> DiffOpt.diff_optimizer(Ipopt.Optimizer))
    pv = matrix_relu(y, model = model)
    function pullback_matrix_relu(dl_dx)
        ## some value from the backpropagation (e.g., loss) is denoted by `l`
        ## so `dl_dy` is the derivative of `l` wrt `y`
        x = model[:x] # load decision variable `x` into scope
        dl_dy = zeros(T, size(dl_dx))
        dl_dq = zeros(T, size(dl_dx))
        ## set sensitivities
        MOI.set.(model, DiffOpt.ReverseVariablePrimal(), x[:], dl_dx[:])
        ## compute grad
        DiffOpt.reverse_differentiate!(model)
        ## return gradient wrt objective function parameters
        obj_exp = MOI.get(model, DiffOpt.ReverseObjectiveFunction())
        ## coeff of `x` in q'x = -2y'x
        dl_dq[:] .= JuMP.coefficient.(obj_exp, x[:])
        dq_dy = -2 # dq/dy = -2
        dl_dy[:] .= dl_dq[:] * dq_dy
        return (ChainRulesCore.NoTangent(), dl_dy,)
    end
    return pv, pullback_matrix_relu
end

# For more details about backpropagation, visit [Introduction, ChainRulesCore.jl](https://juliadiff.org/ChainRulesCore.jl/dev/).

# ## Define the network

layer_size = 10
m = Flux.Chain(
    Flux.Dense(784, layer_size), # 784 being image linear dimension (28 x 28)
    matrix_relu,
    Flux.Dense(layer_size, 10), # 10 being the number of outcomes (0 to 9)
    Flux.softmax,
)

# ## Prepare data

N = 1000 # batch size
## Preprocessing train data
imgs = MLDatasets.MNIST.traintensor(1:N)
labels = MLDatasets.MNIST.trainlabels(1:N)
train_X = float.(reshape(imgs, size(imgs, 1) * size(imgs, 2), N)) # stack images
train_Y = Flux.onehotbatch(labels, 0:9);
## Preprocessing test data
test_imgs = MLDatasets.MNIST.testtensor(1:N)
test_labels = MLDatasets.MNIST.testlabels(1:N)
test_X = float.(reshape(test_imgs, size(test_imgs, 1) * size(test_imgs, 2), N))
test_Y = Flux.onehotbatch(test_labels, 0:9);

# Define input data
# The original data is repeated `epochs` times because `Flux.train!` only
# loops through the data set once

epochs = 50 # ~1 minute (i7 8th gen with 16gb RAM)
## epochs = 100 # leads to 77.8% in about 2 minutes
dataset = repeated((train_X, train_Y), epochs);

# ## Network training

# training loss function, Flux optimizer
custom_loss(x, y) = Flux.crossentropy(m(x), y)
opt = Flux.ADAM()
evalcb = () -> @show(custom_loss(train_X, train_Y))

# Train to optimize network parameters

@time Flux.train!(custom_loss, Flux.params(m), dataset, opt, cb = Flux.throttle(evalcb, 5));

# Although our custom implementation takes time, it is able to reach similar
# accuracy as the usual ReLU function implementation.

# ## Accuracy results

# Average of correct guesses

accuracy(x, y) = Statistics.mean(Flux.onecold(m(x)) .== Flux.onecold(y));

# Training accuracy

accuracy(train_X, train_Y)

# Test accuracy

accuracy(test_X, test_Y)

# Note that the accuracy is low due to simplified training.
# It is possible to increase the number of samples `N`,
# the number of epochs `epoch` and the connectivity `inner`.


text:
function cpu_reduce_partial_sum(
    partial_sum::StructArray
)
    return sum(replace_storage(Array, partial_sum), dims=1)
end

function cuda_reduce_partial_sum(
    partial_sum::StructArray
)
    return CUDA.sum(partial_sum.re, dims=1), CUDA.sum(partial_sum.im, dims=1)
end

function gen_code_replica_kernel!(
    code_replica,
    codes,
    code_frequency,
    sampling_frequency,
    start_code_phase,
    prn,
    num_samples,
    latest_shift,
    code_length
)
    # thread_idx goes from 1:2502
    # sample_idx converted to -1:2500 -> [-1 0 +1]
    thread_idx = 1 + ((blockIdx().x - 1) * blockDim().x + (threadIdx().x - 1))
    if thread_idx <= num_samples
        @inbounds code_replica[thread_idx] = codes[1+mod(floor(Int32, code_frequency/sampling_frequency * (thread_idx + latest_shift) + start_code_phase), code_length), prn]
    end

    return nothing
end

function gen_code_replica_strided_kernel!(
    code_replica,
    codes,
    code_frequency,
    sampling_frequency,
    start_code_phase,
    prn,
    num_samples,
    latest_shift,
    code_length
)
    # thread_idx goes from 1:2502
    # sample_idx converted to -1:2500 -> [-1 0 +1]
    stride = blockDim().x * gridDim().x
    thread_idx = 1 + ((blockIdx().x - 1) * blockDim().x + (threadIdx().x - 1))
    @inbounds for i = thread_idx:stride:num_samples
        code_replica[i] = codes[1+mod(floor(Int32, code_frequency/sampling_frequency * (i + latest_shift) + start_code_phase), code_length), prn]
    end

    return nothing
end

function gen_code_replica_texture_mem_strided_kernel!(
    code_replica,
    codes, # texture memory codes
    code_frequency,
    sampling_frequency,
    start_code_phase,
    prn::Int,
    num_samples,
    latest_shift,
    code_length
)
    # thread_idx goes from 1:2502
    # sample_idx converted to -1:2500 -> [-1 0 +1]

    thread_idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    if thread_idx <= num_samples
        code_replica[thread_idx] = codes[(code_frequency/sampling_frequency * (thread_idx + latest_shift) + start_code_phase) / code_length, prn]
    end

    return nothing
end

function gen_code_replica_texture_mem_strided_nsat_kernel!(
    code_replica,
    codes, # texture memory codes
    code_frequency,
    sampling_frequency,
    start_code_phase,
    prn,
    num_samples,
    latest_shift,
    code_length
)
    # thread_idx goes from 1:2502
    # sample_idx converted to -1:2500 -> [-1 0 +1]
    sat_idx = blockIdx().y
    thread_idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    if thread_idx <= num_samples
        code_replica[thread_idx, sat_idx] = codes[(code_frequency/sampling_frequency * (thread_idx + latest_shift) + start_code_phase) / code_length, prn[sat_idx]]
    end

    return nothing
end

function gen_code_replica_texture_mem_strided_kernel!(
    code_replica,
    codes, # texture memory codes
    code_frequency,
    sampling_frequency,
    start_code_phase,
    prn,
    num_samples,
    latest_shift,
    code_length
)
    # thread_idx goes from 1:2502
    # sample_idx converted to -1:2500 -> [-1 0 +1]
    thread_idx = 1 + ((blockIdx().x - 1) * blockDim().x + (threadIdx().x - 1))
    if thread_idx <= num_samples
        code_replica[thread_idx] = codes[(code_frequency/sampling_frequency * (thread_idx + latest_shift) + start_code_phase) / code_length, prn]
    end

    return nothing
end

function gen_code_replica_texture_mem_kernel!(
    code_replica,
    codes, # texture memory codes
    code_frequency,
    sampling_frequency,
    start_code_phase,
    prn,
    num_samples,
    latest_shift,
    code_length
)
    # thread_idx goes from 1:2502
    # sample_idx converted to -1:2500 -> [-1 0 +1]
    thread_idx = 1 + ((blockIdx().x - 1) * blockDim().x + (threadIdx().x - 1))
    if thread_idx <= num_samples
        code_replica[thread_idx] = codes[(code_frequency/sampling_frequency * (thread_idx + latest_shift) + start_code_phase) / code_length, prn]
    end

    return nothing
end

function downconvert_and_correlate_kernel_1330!(
    res_re,
    res_im,
    signal_re,
    signal_im,
    codes,
    code_frequency,
    correlator_sample_shifts,
    carrier_frequency,
    sampling_frequency,
    start_code_phase,
    carrier_phase,
    code_length,
    prn,
    num_samples,
    num_ants,
    num_corrs
)
    cache = @cuDynamicSharedMem(Float32, (2 * blockDim().x, num_ants, num_corrs))
    sample_idx   = 1 + ((blockIdx().x - 1) * blockDim().x + (threadIdx().x - 1))
    antenna_idx  = 1 + ((blockIdx().y - 1) * blockDim().y + (threadIdx().y - 1))
    corr_idx     = 1 + ((blockIdx().z - 1) * blockDim().z + (threadIdx().z - 1))
    iq_offset = blockDim().x
    cache_index = threadIdx().x - 1

    code_phase = accum_re = accum_im = dw_re = dw_im = carrier_re = carrier_im = 0.0f0
    mod_floor_code_phase = Int(0)

    if sample_idx <= num_samples && antenna_idx <= num_ants && corr_idx <= num_corrs
        # generate carrier
        carrier_im, carrier_re = CUDA.sincos(2π * ((sample_idx - 1) * carrier_frequency / sampling_frequency + carrier_phase))

        # downconvert with the conjugate of the carrier
        dw_re = signal_re[sample_idx, antenna_idx] * carrier_re + signal_im[sample_idx, antenna_idx] * carrier_im
        dw_im = signal_im[sample_idx, antenna_idx] * carrier_re - signal_re[sample_idx, antenna_idx] * carrier_im

        # calculate the code phase
        code_phase = code_frequency / sampling_frequency * ((sample_idx - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase

        # wrap the code phase around the code length e.g. phase = 1024 -> modfloorphase = 1
        mod_floor_code_phase = 1 + mod(floor(Int32, code_phase), code_length)

        # multiply elementwise with the code
        accum_re += codes[mod_floor_code_phase, prn] * dw_re
        accum_im += codes[mod_floor_code_phase, prn] * dw_im
    end

    cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] = accum_re
    cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] = accum_im

    ## Reduction
    # wait until all the accumulators have done writing the results to the cache
    sync_threads()

    i::Int = blockDim().x ÷ 2
    @inbounds while i != 0
        if cache_index < i
            cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 0 * iq_offset + i, antenna_idx, corr_idx]
            cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 1 * iq_offset + i, antenna_idx, corr_idx]
        end
        sync_threads()
        i ÷= 2
    end

    if (threadIdx().x - 1) == 0
        res_re[blockIdx().x, antenna_idx, corr_idx] += cache[1 + 0 * iq_offset, antenna_idx, corr_idx]
        res_im[blockIdx().x, antenna_idx, corr_idx] += cache[1 + 1 * iq_offset, antenna_idx, corr_idx]
    end
    return nothing
end

function downconvert_and_correlate_kernel_1331!(
    res_re,
    res_im,
    signal_re,
    signal_im,
    codes,
    code_frequency,
    correlator_sample_shifts,
    carrier_frequency,
    sampling_frequency,
    start_code_phase,
    carrier_phase,
    code_length,
    prn,
    num_samples,
    num_ants,
    num_corrs
)
    cache = @cuDynamicSharedMem(Float32, (2 * blockDim().x, num_ants, num_corrs))
    sample_idx   = 1 + ((blockIdx().x - 1) * blockDim().x + (threadIdx().x - 1))
    antenna_idx  = 1 + ((blockIdx().y - 1) * blockDim().y + (threadIdx().y - 1))
    corr_idx     = 1 + ((blockIdx().z - 1) * blockDim().z + (threadIdx().z - 1))
    iq_offset = blockDim().x
    cache_index = threadIdx().x - 1

    accum_re = accum_im = dw_re = dw_im = carrier_re = carrier_im = 0.0f0

    if sample_idx <= num_samples && antenna_idx <= num_ants && corr_idx <= num_corrs
        # generate carrier
        carrier_im, carrier_re = CUDA.sincos(2π * ((sample_idx - 1) * carrier_frequency / sampling_frequency + carrier_phase))

        # downconvert with the conjugate of the carrier
        dw_re = signal_re[sample_idx, antenna_idx] * carrier_re + signal_im[sample_idx, antenna_idx] * carrier_im
        dw_im = signal_im[sample_idx, antenna_idx] * carrier_re - signal_re[sample_idx, antenna_idx] * carrier_im

        # multiply elementwise with the code
        accum_re += codes[(code_frequency / sampling_frequency * ((sample_idx - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase) / code_length, prn] * dw_re
        accum_im += codes[(code_frequency / sampling_frequency * ((sample_idx - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase) / code_length, prn] * dw_im
    end

    cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] = accum_re
    cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] = accum_im

    ## Reduction
    # wait until all the accumulators have done writing the results to the cache
    sync_threads()

    i::Int = blockDim().x ÷ 2
    @inbounds while i != 0
        if cache_index < i
            cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 0 * iq_offset + i, antenna_idx, corr_idx]
            cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 1 * iq_offset + i, antenna_idx, corr_idx]
        end
        sync_threads()
        i ÷= 2
    end

    if (threadIdx().x - 1) == 0
        res_re[blockIdx().x, antenna_idx, corr_idx] += cache[1 + 0 * iq_offset, antenna_idx, corr_idx]
        res_im[blockIdx().x, antenna_idx, corr_idx] += cache[1 + 1 * iq_offset, antenna_idx, corr_idx]
    end
    return nothing
end

function downconvert_and_correlate_kernel_1431!(
    res_re,
    res_im,
    signal_re,
    signal_im,
    codes,
    code_frequency,
    correlator_sample_shifts,
    carrier_frequency,
    sampling_frequency,
    start_code_phase,
    carrier_phase,
    code_length,
    prn,
    num_samples,
    num_ants,
    num_corrs
)
    block_dim_x = 2 * blockDim().x # launched with half the grid
    sample_idx   = 1 + ((blockIdx().x - 1) * block_dim_x  + (threadIdx().x - 1))
    antenna_idx  = 1 + ((blockIdx().y - 1) * blockDim().y + (threadIdx().y - 1))
    corr_idx     = 1 + ((blockIdx().z - 1) * blockDim().z + (threadIdx().z - 1))
    iq_offset = blockDim().x
    cache_index = threadIdx().x - 1

    # allocate shared memory
    cache = @cuDynamicSharedMem(Float32, (2 * blockDim().x, num_ants, num_corrs))
    # wipe values
    cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] = 0.0f0
    cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] = 0.0f0

    # define local variables
    accum_re_1 = accum_im_1 = dw_re_1 = dw_im_1 = carrier_re_1 = carrier_im_1 = 0.0f0
    accum_re_2 = accum_im_2 = dw_re_2 = dw_im_2 = carrier_re_2 = carrier_im_2 = 0.0f0

    if sample_idx <= num_samples && antenna_idx <= num_ants && corr_idx <= num_corrs
        # generate carrier
        carrier_im_1, carrier_re_1 = CUDA.sincos(2π * ((sample_idx - 1) * carrier_frequency / sampling_frequency + carrier_phase))

        # downconvert with the conjugate of the carrier
        dw_re_1 = signal_re[sample_idx, antenna_idx] * carrier_re_1 + signal_im[sample_idx, antenna_idx] * carrier_im_1
        dw_im_1 = signal_im[sample_idx, antenna_idx] * carrier_re_1 - signal_re[sample_idx, antenna_idx] * carrier_im_1

        # multiply elementwise with the code
        accum_re_1 += codes[(code_frequency / sampling_frequency * ((sample_idx - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase) / code_length, prn] * dw_re_1
        accum_im_1 += codes[(code_frequency / sampling_frequency * ((sample_idx - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase) / code_length, prn] * dw_im_1

        # write results to shared memory
        cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] = accum_re_1
        cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] = accum_im_1

        if sample_idx + blockDim().x <= num_samples
            # generate carrier for the remaining samples
            carrier_im_2, carrier_re_2 = CUDA.sincos(2π * ((sample_idx + blockDim().x - 1) * carrier_frequency / sampling_frequency + carrier_phase))

            # downconvert with the conjugate of the carrier for the remaining samples
            dw_re_2 = signal_re[sample_idx + blockDim().x, antenna_idx] * carrier_re_2 + signal_im[sample_idx + blockDim().x, antenna_idx] * carrier_im_2
            dw_im_2 = signal_im[sample_idx + blockDim().x, antenna_idx] * carrier_re_2 - signal_re[sample_idx + blockDim().x, antenna_idx] * carrier_im_2

            # multiply elementwise with the code for the remaining samples
            accum_re_2 += codes[(code_frequency / sampling_frequency * ((sample_idx + blockDim().x - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase) / code_length, prn] * dw_re_2
            accum_im_2 += codes[(code_frequency / sampling_frequency * ((sample_idx + blockDim().x - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase) / code_length, prn] * dw_im_2

            # append results to shared memory
            cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += accum_re_2
            cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += accum_im_2
        end
    end

    ## Reduction in shared memory
    # wait until all the accumulators have done writing the results to the cache
    sync_threads()

    i::Int = blockDim().x ÷ 2
    @inbounds while i != 0
        if cache_index < i
            cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 0 * iq_offset + i, antenna_idx, corr_idx]
            cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 1 * iq_offset + i, antenna_idx, corr_idx]
        end
        sync_threads()
        i ÷= 2
    end

    if (threadIdx().x - 1) == 0
        res_re[blockIdx().x, antenna_idx, corr_idx] += cache[1 + 0 * iq_offset, antenna_idx, corr_idx]
        res_im[blockIdx().x, antenna_idx, corr_idx] += cache[1 + 1 * iq_offset, antenna_idx, corr_idx]
    end
    return nothing
end

# function downconvert_and_correlate_kernel_3431!(
#     res_re,
#     res_im,
#     signal_re,
#     signal_im,
#     codes,
#     code_frequency,
#     correlator_sample_shifts,
#     carrier_frequency,
#     sampling_frequency,
#     start_code_phase,
#     carrier_phase,
#     code_length,
#     prn,
#     num_samples,
#     num_ants,
#     num_corrs
# )
#     block_dim_x = 2 * blockDim().x # launched with half the grid
#     sample_idx   = 1 + ((blockIdx().x - 1) * block_dim_x  + (threadIdx().x - 1))
#     antenna_idx  = 1 + ((blockIdx().y - 1) * blockDim().y + (threadIdx().y - 1))
#     corr_idx     = 1 + ((blockIdx().z - 1) * blockDim().z + (threadIdx().z - 1))
#     iq_offset = blockDim().x
#     cache_index = threadIdx().x - 1

#     # allocate shared memory
#     cache = @cuDynamicSharedMem(Float32, (2 * blockDim().x, num_ants, num_corrs))
#     # wipe values
#     cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] = 0.0f0
#     cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] = 0.0f0

#     # define local variables
#     accum_re_1 = accum_im_1 = dw_re_1 = dw_im_1 = carrier_re_1 = carrier_im_1 = 0.0f0
#     accum_re_2 = accum_im_2 = dw_re_2 = dw_im_2 = carrier_re_2 = carrier_im_2 = 0.0f0

#     if sample_idx <= num_samples && antenna_idx <= num_ants && corr_idx <= num_corrs
#         # generate carrier
#         carrier_im_1, carrier_re_1 = CUDA.sincos(2π * ((sample_idx - 1) * carrier_frequency / sampling_frequency + carrier_phase))

#         # downconvert with the conjugate of the carrier
#         dw_re_1 = signal_re[sample_idx, antenna_idx] * carrier_re_1 + signal_im[sample_idx, antenna_idx] * carrier_im_1
#         dw_im_1 = signal_im[sample_idx, antenna_idx] * carrier_re_1 - signal_re[sample_idx, antenna_idx] * carrier_im_1

#         # multiply elementwise with the code
#         accum_re_1 += codes[(code_frequency / sampling_frequency * ((sample_idx - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase) / code_length, prn] * dw_re_1
#         accum_im_1 += codes[(code_frequency / sampling_frequency * ((sample_idx - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase) / code_length, prn] * dw_im_1

#         # write results to shared memory
#         cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] = accum_re_1
#         cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] = accum_im_1

#         if sample_idx + blockDim().x <= num_samples
#             # generate carrier for the remaining samples
#             carrier_im_2, carrier_re_2 = CUDA.sincos(2π * ((sample_idx + blockDim().x - 1) * carrier_frequency / sampling_frequency + carrier_phase))

#             # downconvert with the conjugate of the carrier for the remaining samples
#             dw_re_2 = signal_re[sample_idx + blockDim().x, antenna_idx] * carrier_re_2 + signal_im[sample_idx + blockDim().x, antenna_idx] * carrier_im_2
#             dw_im_2 = signal_im[sample_idx + blockDim().x, antenna_idx] * carrier_re_2 - signal_re[sample_idx + blockDim().x, antenna_idx] * carrier_im_2

#             # multiply elementwise with the code for the remaining samples
#             accum_re_2 += codes[(code_frequency / sampling_frequency * ((sample_idx + blockDim().x - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase) / code_length, prn] * dw_re_2
#             accum_im_2 += codes[(code_frequency / sampling_frequency * ((sample_idx + blockDim().x - 1) + correlator_sample_shifts[corr_idx]) + start_code_phase) / code_length, prn] * dw_im_2

#             # append results to shared memory
#             cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += accum_re_2
#             cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += accum_im_2
#         end
#     end

#     ## Reduction in shared memory
#     # wait until all the accumulators have done writing the results to the cache
#     sync_threads()

#     i::Int = blockDim().x ÷ 2
#     @inbounds while i != 0
#         if cache_index < i
#             cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 0 * iq_offset + i, antenna_idx, corr_idx]
#             cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 1 * iq_offset + i, antenna_idx, corr_idx]
#         end
#         sync_threads()
#         i ÷= 2
#     end

#     if (threadIdx().x - 1) == 0
#         res_re[blockIdx().x, antenna_idx, corr_idx] += cache[1 + 0 * iq_offset, antenna_idx, corr_idx]
#         res_im[blockIdx().x, antenna_idx, corr_idx] += cache[1 + 1 * iq_offset, antenna_idx, corr_idx]
#     end
#     return nothing
# end

function downconvert_and_correlate_kernel_3431!(
    partial_sum_re,
    partial_sum_im,
    carrier_replica_re,
    carrier_replica_im,
    downconverted_signal_re,
    downconverted_signal_im,
    signal_re,
    signal_im,
    code_replica,
    correlator_sample_shifts::SVector{NCOR, Int64},
    carrier_frequency,
    sampling_frequency,
    carrier_phase,
    num_samples::Int,
    num_ants::NumAnts{NANT}
)  where {NCOR, NANT}
    cache = @cuDynamicSharedMem(Float32, (2 * blockDim().x, NANT, NCOR))
    sample_idx   = 1 + ((blockIdx().x - 1) * (2 * blockDim().x) + (threadIdx().x - 1)) # double the grid
    iq_offset = blockDim().x # indexing offset for complex values I/Q samples
    cache_index = threadIdx().x - 1

    dw_re_1 = dw_im_1 = carrier_re_1 = carrier_im_1 = 0.0f0
    dw_re_2 = dw_im_2 = carrier_re_2 = carrier_im_2 = 0.0f0

    @inbounds if sample_idx <= num_samples
        # carrier replica generation, sin->im , cos->re
        carrier_im_1, carrier_re_1 = CUDA.sincos(2π * ((sample_idx - 1) * carrier_frequency / sampling_frequency + carrier_phase))

        # downconversion / carrier wipe off
        for antenna_idx = 1:NANT
            dw_re_1 = signal_re[sample_idx, antenna_idx] * carrier_re_1 + signal_im[sample_idx, antenna_idx] * carrier_im_1
            dw_im_1 = signal_im[sample_idx, antenna_idx] * carrier_re_1 - signal_re[sample_idx, antenna_idx] * carrier_im_1
            for corr_idx = 1:NCOR
                sample_shift = correlator_sample_shifts[corr_idx] - correlator_sample_shifts[1]

                # write to shared memory cache
                cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] = code_replica[sample_idx + sample_shift] * dw_re_1
                cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] = code_replica[sample_idx + sample_shift] * dw_im_1
            end
        end

        if sample_idx + blockDim().x <= num_samples
            # carrier replica generation, sin->im , cos->re
            carrier_im_2, carrier_re_2 = CUDA.sincos(2π * ((sample_idx + blockDim().x - 1) * carrier_frequency / sampling_frequency + carrier_phase))

            # downconversion / carrier wipe off
            for antenna_idx = 1:NANT
                dw_re_2 = signal_re[sample_idx + blockDim().x, antenna_idx] * carrier_re_2 + signal_im[sample_idx + blockDim().x, antenna_idx] * carrier_im_2
                dw_im_2 = signal_im[sample_idx + blockDim().x, antenna_idx] * carrier_re_2 - signal_re[sample_idx + blockDim().x, antenna_idx] * carrier_im_2
                for corr_idx = 1:NCOR
                    sample_shift = correlator_sample_shifts[corr_idx] - correlator_sample_shifts[1]

                    # write to shared memory cache
                    cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += code_replica[sample_idx + blockDim().x + sample_shift] * dw_re_2
                    cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += code_replica[sample_idx + blockDim().x + sample_shift] * dw_im_2
                end
            end
        end
    end
    ## Partial Reduction
    # wait until all the accumulators have done writing the results to the cache
    sync_threads()

    i::Int = blockDim().x ÷ 2
    @inbounds while i != 0
        if cache_index < i
            for antenna_idx = 1:NANT
                for corr_idx = 1:NCOR
                    cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 0 * iq_offset + i, antenna_idx, corr_idx]
                    cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 1 * iq_offset + i, antenna_idx, corr_idx]
                end
            end
        end
        sync_threads()
        i ÷= 2
    end

    @inbounds if (threadIdx().x - 1) == 0
         for antenna_idx = 1:NANT
            for corr_idx = 1:NCOR
                partial_sum_re[blockIdx().x, antenna_idx, corr_idx] = cache[1 + 0 * iq_offset, antenna_idx, corr_idx]
                partial_sum_im[blockIdx().x, antenna_idx, corr_idx] = cache[1 + 1 * iq_offset, antenna_idx, corr_idx]
            end
        end
    end
    return nothing
end

function downconvert_and_correlate_kernel_4431!(
    accum_re,
    accum_im,
    carrier_replica_re,
    carrier_replica_im,
    downconverted_signal_re,
    downconverted_signal_im,
    signal_re,
    signal_im,
    code_replica,
    correlator_sample_shifts::SVector{NCOR, Int64},
    carrier_frequency,
    sampling_frequency,
    carrier_phase,
    num_samples::Int,
    num_ants::NumAnts{NANT}
)  where {NCOR, NANT}
    cache = @cuDynamicSharedMem(Float32, (2 * blockDim().x, NANT, NCOR))
    sample_idx   = 1 + ((blockIdx().x - 1) * (2 * blockDim().x) + (threadIdx().x - 1)) # double the grid
    iq_offset = blockDim().x # indexing offset for complex values I/Q samples
    cache_index = threadIdx().x - 1

    dw_re_1 = dw_im_1 = carrier_re_1 = carrier_im_1 = 0.0f0
    dw_re_2 = dw_im_2 = carrier_re_2 = carrier_im_2 = 0.0f0

    @inbounds if sample_idx <= num_samples
        # carrier replica generation, sin->im , cos->re
        carrier_im_1, carrier_re_1 = CUDA.sincos(2π * ((sample_idx - 1) * carrier_frequency / sampling_frequency + carrier_phase))

        # downconversion / carrier wipe off
        for antenna_idx = 1:NANT
            dw_re_1 = signal_re[sample_idx, antenna_idx] * carrier_re_1 + signal_im[sample_idx, antenna_idx] * carrier_im_1
            dw_im_1 = signal_im[sample_idx, antenna_idx] * carrier_re_1 - signal_re[sample_idx, antenna_idx] * carrier_im_1
            for corr_idx = 1:NCOR
                sample_shift = correlator_sample_shifts[corr_idx] - correlator_sample_shifts[1]

                # write to shared memory cache
                cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] = code_replica[sample_idx + sample_shift] * dw_re_1
                cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] = code_replica[sample_idx + sample_shift] * dw_im_1
            end
        end

        if sample_idx + blockDim().x <= num_samples
            # carrier replica generation, sin->im , cos->re
            carrier_im_2, carrier_re_2 = CUDA.sincos(2π * ((sample_idx + blockDim().x - 1) * carrier_frequency / sampling_frequency + carrier_phase))

            # downconversion / carrier wipe off
            for antenna_idx = 1:NANT
                dw_re_2 = signal_re[sample_idx + blockDim().x, antenna_idx] * carrier_re_2 + signal_im[sample_idx + blockDim().x, antenna_idx] * carrier_im_2
                dw_im_2 = signal_im[sample_idx + blockDim().x, antenna_idx] * carrier_re_2 - signal_re[sample_idx + blockDim().x, antenna_idx] * carrier_im_2
                for corr_idx = 1:NCOR
                    sample_shift = correlator_sample_shifts[corr_idx] - correlator_sample_shifts[1]

                    # write to shared memory cache
                    cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += code_replica[sample_idx + blockDim().x + sample_shift] * dw_re_2
                    cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += code_replica[sample_idx + blockDim().x + sample_shift] * dw_im_2
                end
            end
        end
    end

    ## Partial Reduction
    # wait until all the accumulators have done writing the results to the cache
    sync_threads()

    i::Int = blockDim().x ÷ 2
    @inbounds while i != 0
        if cache_index < i
            for antenna_idx = 1:NANT
                for corr_idx = 1:NCOR
                    cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 0 * iq_offset + i, antenna_idx, corr_idx]
                    cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 1 * iq_offset + i, antenna_idx, corr_idx]
                end
            end
        end
        sync_threads()
        i ÷= 2
    end

    @inbounds if threadIdx().x == 1
        for antenna_idx = 1:NANT
            for corr_idx = 1:NCOR
                CUDA.@atomic accum_re[antenna_idx, corr_idx] += cache[1 + 0 * iq_offset, antenna_idx, corr_idx]
                CUDA.@atomic accum_im[antenna_idx, corr_idx] += cache[1 + 1 * iq_offset, antenna_idx, corr_idx]
            end
        end
    end

    return nothing
end

function downconvert_and_correlate_kernel_3d_4431!(
    accum_re,
    accum_im,
    signal_re,
    signal_im,
    code_replica,
    correlator_sample_shifts,
    carrier_frequency,
    sampling_frequency,
    carrier_phase,
    num_samples::Int,
    num_ants::NumAnts{NANT}
)  where {NANT}
    cache = CuDynamicSharedArray(Float32, (2 * blockDim().x, NANT))
    sample_idx = 1 + ((blockIdx().x - 1) * (2 * blockDim().x) + (threadIdx().x - 1)) # double the grid
    iq_offset = blockDim().x # indexing offset for complex values I/Q samples
    cache_index = threadIdx().x - 1
    antenna_idx = threadIdx().y
    corr_idx = blockIdx().y
    sat_idx = blockIdx().z

    dw_re_1 = dw_im_1 = carrier_re_1 = carrier_im_1 = 0.0f0
    dw_re_2 = dw_im_2 = carrier_re_2 = carrier_im_2 = 0.0f0
    sample_shift = correlator_sample_shifts[corr_idx]

    @inbounds if sample_idx <= num_samples
        # carrier replica generation, sin->im , cos->re
        carrier_im_1, carrier_re_1 = CUDA.sincos(2π * ((sample_idx - 1) * carrier_frequency / sampling_frequency + carrier_phase))

        # downconversion / carrier wipe off
        for antenna_idx = 1:NANT
            dw_re_1 = signal_re[sample_idx, antenna_idx, sat_idx] * carrier_re_1 + signal_im[sample_idx, antenna_idx, sat_idx] * carrier_im_1
            dw_im_1 = signal_im[sample_idx, antenna_idx, sat_idx] * carrier_re_1 - signal_re[sample_idx, antenna_idx, sat_idx] * carrier_im_1

            # write to shared memory cache
            cache[1 + cache_index + 0 * iq_offset, antenna_idx] = code_replica[sample_idx + sample_shift, sat_idx] * dw_re_1
            cache[1 + cache_index + 1 * iq_offset, antenna_idx] = code_replica[sample_idx + sample_shift, sat_idx] * dw_im_1
        end

        if sample_idx + blockDim().x <= num_samples
            # carrier replica generation, sin->im , cos->re
            carrier_im_2, carrier_re_2 = CUDA.sincos(2π * ((sample_idx + blockDim().x - 1) * carrier_frequency / sampling_frequency + carrier_phase))

            # downconversion / carrier wipe off
            for antenna_idx = 1:NANT
                dw_re_2 = signal_re[sample_idx + blockDim().x, antenna_idx, sat_idx] * carrier_re_2 + signal_im[sample_idx + blockDim().x, antenna_idx, sat_idx] * carrier_im_2
                dw_im_2 = signal_im[sample_idx + blockDim().x, antenna_idx, sat_idx] * carrier_re_2 - signal_re[sample_idx + blockDim().x, antenna_idx, sat_idx] * carrier_im_2

                # write to shared memory cache
                cache[1 + cache_index + 0 * iq_offset, antenna_idx] += code_replica[sample_idx + blockDim().x + sample_shift, sat_idx] * dw_re_2
                cache[1 + cache_index + 1 * iq_offset, antenna_idx] += code_replica[sample_idx + blockDim().x + sample_shift, sat_idx] * dw_im_2
            end
        end
    end

    ## Partial Reduction
    # wait until all the accumulators have done writing the results to the cache
    sync_threads()

    i::Int = blockDim().x ÷ 2
    @inbounds while i != 0
        if cache_index < i
            for antenna_idx = 1:NANT
                cache[1 + cache_index + 0 * iq_offset, antenna_idx] += cache[1 + cache_index + 0 * iq_offset + i, antenna_idx]
                cache[1 + cache_index + 1 * iq_offset, antenna_idx] += cache[1 + cache_index + 1 * iq_offset + i, antenna_idx]
            end
        end
        sync_threads()
        i ÷= 2
    end

    @inbounds if threadIdx().x == 1
        for antenna_idx = 1:NANT
            CUDA.@atomic accum_re[antenna_idx, corr_idx, sat_idx] += cache[1 + 0 * iq_offset, antenna_idx]
            CUDA.@atomic accum_im[antenna_idx, corr_idx, sat_idx] += cache[1 + 1 * iq_offset, antenna_idx]
        end
    end

    return nothing
end

function 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::SVector{NCOR, Int64},
    num_of_shifts,
    code_frequency,
    carrier_frequency,
    sampling_frequency,
    start_code_phase,
    carrier_phase,
    num_samples::Int,
    num_ants::NumAnts{NANT},
)  where {NCOR, NANT}
    cache = @cuDynamicSharedMem(Float32, (2 * blockDim().x, NANT, NCOR))
    sample_idx   = 1 + ((blockIdx().x - 1) * (2 * blockDim().x) + (threadIdx().x - 1))
    iq_offset = blockDim().x # indexing offset for complex values I/Q samples
    cache_index = threadIdx().x - 1

    # local
    dw_re_1 = dw_im_1 = carrier_re_1 = carrier_im_1 = 0.0f0
    dw_re_2 = dw_im_2 = carrier_re_2 = carrier_im_2 = 0.0f0

    # Code replica generation
    if sample_idx <= num_samples + num_of_shifts
        @inbounds code_replica[sample_idx] = codes[(code_frequency/sampling_frequency * (sample_idx - num_of_shifts) + start_code_phase)/code_length, prn]
        if sample_idx + blockDim().x <= num_samples + num_of_shifts
            @inbounds code_replica[sample_idx + blockDim().x] = codes[(code_frequency/sampling_frequency * (sample_idx + blockDim().x - num_of_shifts) + start_code_phase)/code_length, prn]
        end
    end

    sync_threads()

    @inbounds if sample_idx <= num_samples
        # carrier replica generation, sin->im , cos->re
        carrier_im_1, carrier_re_1 = CUDA.sincos(2π * ((sample_idx - 1) * carrier_frequency / sampling_frequency + carrier_phase))

        # downconversion / carrier wipe off
        for antenna_idx = 1:NANT
            dw_re_1 = signal_re[sample_idx, antenna_idx] * carrier_re_1 + signal_im[sample_idx, antenna_idx] * carrier_im_1
            dw_im_1 = signal_im[sample_idx, antenna_idx] * carrier_re_1 - signal_re[sample_idx, antenna_idx] * carrier_im_1
            for corr_idx = 1:NCOR
                sample_shift = correlator_sample_shifts[corr_idx] - correlator_sample_shifts[1]
                # write to shared memory cache
                cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] = code_replica[sample_idx + sample_shift] * dw_re_1
                cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] = code_replica[sample_idx + sample_shift] * dw_im_1
            end
        end

        if sample_idx + blockDim().x <= num_samples
            # carrier replica generation, sin->im , cos->re
            carrier_im_2, carrier_re_2 = CUDA.sincos(2π * ((sample_idx + blockDim().x - 1) * carrier_frequency / sampling_frequency + carrier_phase))

            # downconversion / carrier wipe off
            for antenna_idx = 1:NANT
                dw_re_2 = signal_re[sample_idx + blockDim().x, antenna_idx] * carrier_re_2 + signal_im[sample_idx + blockDim().x, antenna_idx] * carrier_im_2
                dw_im_2 = signal_im[sample_idx + blockDim().x, antenna_idx] * carrier_re_2 - signal_re[sample_idx + blockDim().x, antenna_idx] * carrier_im_2
                for corr_idx = 1:NCOR
                    sample_shift = correlator_sample_shifts[corr_idx] - correlator_sample_shifts[1]
                    # write to shared memory cache
                    cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += code_replica[sample_idx + blockDim().x + sample_shift] * dw_re_2
                    cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += code_replica[sample_idx + blockDim().x + sample_shift] * dw_im_2
                end
            end
        end
    end

    ## Partial Reduction
    # wait until all the accumulators have done writing the results to the cache
    sync_threads()

    i::Int = blockDim().x ÷ 2
    @inbounds while i != 0
        if cache_index < i
            for antenna_idx = 1:NANT
                for corr_idx = 1:NCOR
                    cache[1 + cache_index + 0 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 0 * iq_offset + i, antenna_idx, corr_idx]
                    cache[1 + cache_index + 1 * iq_offset, antenna_idx, corr_idx] += cache[1 + cache_index + 1 * iq_offset + i, antenna_idx, corr_idx]
                end
            end
        end
        sync_threads()
        i ÷= 2
    end

    # Last block reduction via atomic add
    @inbounds if (threadIdx().x - 1) == 0
        for antenna_idx = 1:NANT
            for corr_idx = 1:NCOR
                CUDA.@atomic accum_re[antenna_idx, corr_idx] += cache[1 + 0 * iq_offset, antenna_idx, corr_idx]
                CUDA.@atomic accum_im[antenna_idx, corr_idx] += cache[1 + 1 * iq_offset, antenna_idx, corr_idx]
            end
        end
    end

    return nothing
end

function 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::Int,
    num_ants::NumAnts{NANT},
    correlator_sample_shifts::SVector{NCOR, Int64}
) where {NANT, NCOR}
    stride = blockDim().x * gridDim().x
    thread_idx = 1 + ((blockIdx().x - 1) * blockDim().x + (threadIdx().x - 1))

    @inbounds for sample_idx = thread_idx:stride:num_samples
        # gen carrier replica
        carrier_replica_im[sample_idx], carrier_replica_re[sample_idx] = CUDA.sincos(2π * ((sample_idx - 1) * carrier_frequency / sampling_frequency + carrier_phase))

        for antenna_idx = 1:NANT
            # downconvert
            downconverted_signal_re[sample_idx, antenna_idx] = signal_re[sample_idx, antenna_idx] * carrier_replica_re[sample_idx] + signal_im[sample_idx, antenna_idx] * carrier_replica_im[sample_idx]
            downconverted_signal_im[sample_idx, antenna_idx] = signal_im[sample_idx, antenna_idx] * carrier_replica_re[sample_idx] - signal_re[sample_idx, antenna_idx] * carrier_replica_im[sample_idx]
            for corr_idx = 1:NCOR
                # accumulate
                shift = correlator_sample_shifts[corr_idx] - correlator_sample_shifts[1]
                accum_re[sample_idx, antenna_idx, corr_idx] = downconverted_signal_re[sample_idx, antenna_idx] * code_replica[sample_idx + shift]
                accum_im[sample_idx, antenna_idx, corr_idx] = downconverted_signal_im[sample_idx, antenna_idx] * code_replica[sample_idx + shift]
            end
        end
    end

    return nothing
end

# KERNEL 1_3_cplx_multi
function 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,
    partial_sum,
    carrier_replica_re,
    carrier_replica_im,
    downconverted_signal_re,
    downconverted_signal_im,
    signal_re,
    signal_im,
    correlator_sample_shifts::SVector{NCOR, Int64},
    carrier_frequency,
    carrier_phase,
    num_ants::NumAnts{NANT},
    num_corrs,
    algorithm::KernelAlgorithm{1330};
) where {NANT, NCOR}
    @cuda threads=threads_per_block[1] blocks=blocks_per_grid shmem=shmem_size[1] downconvert_and_correlate_kernel_1330!(
        partial_sum.re,
        partial_sum.im,
        signal_re,
        signal_im,
        codes,
        code_frequency,
        correlator_sample_shifts,
        carrier_frequency,
        sampling_frequency,
        start_code_phase,
        carrier_phase,
        code_length,
        prn,
        num_samples,
        NANT,
        NCOR
    )
    @cuda threads=512 blocks=1 shmem=shmem_size[2] reduce_cplx_multi_3(
        partial_sum.re,
        partial_sum.im,
        partial_sum.re,
        partial_sum.im,
        blocks_per_grid,
        num_ants,
        correlator_sample_shifts
    )
end

# KERNEL 1_3_cplx_multi_textmem
function 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,
    partial_sum,
    carrier_replica_re,
    carrier_replica_im,
    downconverted_signal_re,
    downconverted_signal_im,
    signal_re,
    signal_im,
    correlator_sample_shifts::SVector{NCOR, Int64},
    carrier_frequency,
    carrier_phase,
    num_ants::NumAnts{NANT},
    num_corrs,
    algorithm::KernelAlgorithm{1331};
) where {NANT, NCOR}
    NVTX.@range "downconvert_and_correlate_kernel_1331!" begin
        @cuda threads=threads_per_block[1] blocks=blocks_per_grid shmem=shmem_size[1] downconvert_and_correlate_kernel_1331!(
            partial_sum.re,
            partial_sum.im,
            signal_re,
            signal_im,
            codes,
            code_frequency,
            correlator_sample_shifts,
            carrier_frequency,
            sampling_frequency,
            start_code_phase,
            carrier_phase,
            code_length,
            prn,
            num_samples,
            NANT,
            NCOR
        )
    end
    NVTX.@range "reduce_cplx_multi_3" begin
        @cuda threads=threads_per_block[2] blocks=1 shmem=shmem_size[2] reduce_cplx_multi_3(
            partial_sum.re,
            partial_sum.im,
            partial_sum.re,
            partial_sum.im,
            blocks_per_grid,
            num_ants,
            correlator_sample_shifts
        )
    end
end

# KERNEL 1_4_cplx_multi_textmem
# 1431
function 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,
    partial_sum,
    carrier_replica_re,
    carrier_replica_im,
    downconverted_signal_re,
    downconverted_signal_im,
    signal_re,
    signal_im,
    correlator_sample_shifts::SVector{NCOR, Int64},
    carrier_frequency,
    carrier_phase,
    num_ants::NumAnts{NANT},
    num_corrs,
    algorithm::KernelAlgorithm{1431};
) where {NANT, NCOR}
    NVTX.@range "downconvert_and_correlate_kernel_1431!" begin
        @cuda threads=threads_per_block[1] blocks=blocks_per_grid shmem=shmem_size[1] downconvert_and_correlate_kernel_1431!(
            partial_sum.re,
            partial_sum.im,
            signal_re,
            signal_im,
            codes,
            code_frequency,
            correlator_sample_shifts,
            carrier_frequency,
            sampling_frequency,
            start_code_phase,
            carrier_phase,
            code_length,
            prn,
            num_samples,
            NANT,
            NCOR
        )
    end
    NVTX.@range "reduce_cplx_multi_4" begin
        @cuda threads=threads_per_block[2] blocks=1 shmem=shmem_size[2] reduce_cplx_multi_4(
            partial_sum.re,
            partial_sum.im,
            partial_sum.re,
            partial_sum.im,
            blocks_per_grid,
            num_ants,
            correlator_sample_shifts
        )
    end
end

# KERNEL 2_3_cplx_multi
# 2330
function 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::SVector{NCOR, Int64},
    carrier_frequency,
    carrier_phase,
    num_ants::NumAnts{NANT},
    num_corrs,
    algorithm::KernelAlgorithm{2330}
) where {NANT, NCOR}
    NVTX.@range "gen_code_replica_kernel!" begin
        @cuda threads=threads_per_block[1] blocks=blocks_per_grid[1] gen_code_replica_strided_kernel!(
            code_replica,
            codes,
            code_frequency,
            sampling_frequency,
            start_code_phase,
            prn,
            num_samples,
            num_of_shifts,
            code_length
        )
    end
    NVTX.@range "downconvert_and_accumulate!" begin
        @cuda threads=threads_per_block[2] blocks=blocks_per_grid[2] 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,
            num_ants,
            correlator_sample_shifts
        )
    end
    NVTX.@range "reduce_cplx_multi_3" begin
        @cuda threads=threads_per_block[3] blocks=blocks_per_grid[3] shmem=shmem_size reduce_cplx_multi_3(
            phi_re,
            phi_im,
            accum_re,
            accum_im,
            num_samples,
            num_ants,
            correlator_sample_shifts
        )
    end
    # print(Array(phi_re))
    NVTX.@range "reduce_cplx_multi_3" begin
        @cuda threads=threads_per_block[3] blocks=1 shmem=shmem_size reduce_cplx_multi_3(
            phi_re,
            phi_im,
            phi_re,
            phi_im,
            blocks_per_grid[3],
            num_ants,
            correlator_sample_shifts
        )
    end
    # print(Array(phi_re))
end

# KERNEL 2_3_cplx_multi_textmem
# 2331
function 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::SVector{NCOR, Int64},
    carrier_frequency,
    carrier_phase,
    num_ants::NumAnts{NANT},
    num_corrs,
    algorithm::KernelAlgorithm{2331}
) where {NANT, NCOR}
    NVTX.@range "gen_code_replica_kernel!" begin
        @cuda threads=threads_per_block[1] blocks=blocks_per_grid[1] 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
        )
    end
    NVTX.@range "downconvert_and_accumulate!" begin
        @cuda threads=threads_per_block[2] blocks=blocks_per_grid[2] 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,
            num_ants,
            correlator_sample_shifts
        )
    end
    NVTX.@range "reduce_cplx_multi_3" begin
        @cuda threads=threads_per_block[3] blocks=blocks_per_grid[3] shmem=shmem_size reduce_cplx_multi_3(
            phi_re,
            phi_im,
            accum_re,
            accum_im,
            num_samples,
            num_ants,
            correlator_sample_shifts
        )
    end
    # print(Array(phi_re))
    NVTX.@range "reduce_cplx_multi_3" begin
        @cuda threads=threads_per_block[3] blocks=1 shmem=shmem_size reduce_cplx_multi_3(
            phi_re,
            phi_im,
            phi_re,
            phi_im,
            blocks_per_grid[3],
            num_ants,
            correlator_sample_shifts
        )
    end
    # print(Array(phi_re))
end

# KERNEL 2_4_cplx_multi
# 2430
function 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::SVector{NCOR, Int64},
    carrier_frequency,
    carrier_phase,
    num_ants::NumAnts{NANT},
    num_corrs,
    algorithm::KernelAlgorithm{2430}
) where {NANT, NCOR}
    NVTX.@range "gen_code_replica_kernel!" begin
        @cuda threads=threads_per_block[1] blocks=blocks_per_grid[1] gen_code_replica_strided_kernel!(
            code_replica,
            codes,
            code_frequency,
            sampling_frequency,
            start_code_phase,
            prn,
            num_samples,
            num_of_shifts,
            code_length
        )
    end
    NVTX.@range "downconvert_and_accumulate!" begin
        @cuda threads=threads_per_block[2] blocks=blocks_per_grid[2] 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,
            num_ants,
            correlator_sample_shifts
        )
    end
    NVTX.@range "reduce_cplx_multi_4" begin
        @cuda threads=threads_per_block[3] blocks=cld(blocks_per_grid[3], 2) shmem=shmem_size reduce_cplx_multi_4(
            phi_re,
            phi_im,
            accum_re,
            accum_im,
            num_samples,
            num_ants,
            correlator_sample_shifts
        )
    end
    # print(Array(phi_re))
    NVTX.@range "reduce_cplx_multi_4" begin
        @cuda threads=threads_per_block[3] blocks=1 shmem=shmem_size reduce_cplx_multi_4(
            phi_re,
            phi_im,
            phi_re,
            phi_im,
            blocks_per_grid[3],
            num_ants,
            correlator_sample_shifts
        )
    end
    # print(Array(phi_re))
end

# KERNEL 2_4_cplx_multi_textmem
# 2431
function 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::SVector{NCOR, Int64},
    carrier_frequency,
    carrier_phase,
    num_ants::NumAnts{NANT},
    num_corrs,
    algorithm::KernelAlgorithm{2431}
) where {NANT, NCOR}
    NVTX.@range "gen_code_replica_kernel!" begin
        @cuda threads=threads_per_block[1] blocks=blocks_per_grid[1] 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
        )
    end
    NVTX.@range "downconvert_and_accumulate!" begin
        @cuda threads=threads_per_block[2] blocks=blocks_per_grid[2] 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,
            num_ants,
            correlator_sample_shifts
        )
    end
    NVTX.@range "reduce_cplx_multi_4" begin
        @cuda threads=threads_per_block[3] blocks=cld(blocks_per_grid[3], 2) shmem=shmem_size reduce_cplx_multi_4(
            phi_re,
            phi_im,
            accum_re,
            accum_im,
            num_samples,
            num_ants,
            correlator_sample_shifts
        )
    end
    # print(Array(phi_re))
    NVTX.@range "reduce_cplx_multi_4" begin
        @cuda threads=threads_per_block[3] blocks=1 shmem=shmem_size reduce_cplx_multi_4(
            phi_re,
            phi_im,
            phi_re,
            phi_im,
            blocks_per_grid[3],
            num_ants,
            correlator_sample_shifts
        )
    end
    # print(Array(phi_re))
end

# KERNEL 3_4_cplx_multi_textmem
# 3431
function 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::SVector{NCOR, Int64},
    carrier_frequency,
    carrier_phase,
    num_ants::NumAnts{NANT},
    num_corrs,
    algorithm::KernelAlgorithm{3431}
) where {NANT, NCOR}
    NVTX.@range "gen_code_replica_texture_mem_strided_kernel!" begin
        @cuda threads=threads_per_block[1] blocks=blocks_per_grid[1] 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
        )
    end
    NVTX.@range "downconvert_and_correlate_kernel_3431!" begin
        @cuda threads=threads_per_block[2] blocks=blocks_per_grid[2] shmem=shmem_size[1] downconvert_and_correlate_kernel_3431!(
            accum_re,
            accum_im,
            carrier_replica_re,
            carrier_replica_im,
            downconverted_signal_re,
            downconverted_signal_im,
            signal_re,
            signal_im,
            code_replica,
            correlator_sample_shifts,
            carrier_frequency,
            sampling_frequency,
            carrier_phase,
            num_samples,
            num_ants
        )
    end
    NVTX.@range "reduce_cplx_multi_4" begin
        @cuda threads=threads_per_block[3] blocks=1 shmem=shmem_size[2] reduce_cplx_multi_4(
            accum_re,
            accum_im,
            accum_re,
            accum_im,
            blocks_per_grid[2],
            num_ants,
            correlator_sample_shifts
        )
    end
end

# KERNEL 4_4_cplx_multi_textmem
# 4431
function 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::SVector{NCOR, Int64},
    carrier_frequency,
    carrier_phase,
    num_ants::NumAnts{NANT},
    num_corrs,
    algorithm::KernelAlgorithm{4431}
) where {NANT, NCOR}
    NVTX.@range "gen_code_replica_texture_mem_strided_kernel!" begin
        @cuda threads=threads_per_block[1] blocks=blocks_per_grid[1] 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
        )
    end
    NVTX.@range "downconvert_and_correlate_kernel_3431!" begin
        @cuda threads=threads_per_block[2] blocks=blocks_per_grid[2] shmem=shmem_size[1] downconvert_and_correlate_kernel_4431!(
            accum_re,
            accum_im,
            carrier_replica_re,
            carrier_replica_im,
            downconverted_signal_re,
            downconverted_signal_im,
            signal_re,
            signal_im,
            code_replica,
            correlator_sample_shifts,
            carrier_frequency,
            sampling_frequency,
            carrier_phase,
            num_samples,
            num_ants
        )
    end
end

ext: jl

max_stars_repo_name: mfkiwl/GPUAcceleratedTracking

max_issues_repo_name: mfkiwl/GPUAcceleratedTracking

max_forks_repo_name: mfkiwl/GPUAcceleratedTracking


text:
using SemiclassicalOrthogonalPolynomials
using ClassicalOrthogonalPolynomials, ContinuumArrays, BandedMatrices, QuasiArrays, Test, LazyArrays, LinearAlgebra, InfiniteArrays
import LazyArrays: AbstractCachedVector
import SemiclassicalOrthogonalPolynomials: initialα, αdirect, αdirect!, backαcoeff!, αcoefficients!, evalϕn, neg1c_tolegendre, evalQn, getindex, initialα_gen, symlowered_jacobimatrix, αgenfillerbackwards!, symlowered_jacobimatrix

@testset "Special case: SemiclassicalJacobi(t,0,0,-1) " begin
    @testset "inital α" begin
        t1 = 1.1
        t2 = 1.841
        t3 = 3.91899
        t4 = BigFloat("1.0000000000000000000001")
        # Mathematica
        @test initialα(t1) ≈ 0.4430825224938773
        @test initialα(t2) ≈ 0.1980462516542294
        @test initialα(t3) ≈ 0.0865853392346796
        @test initialα(t4) ≈ 0.9610516212042500
    end

    @testset "αdirect! consistency" begin
        N = 5
        t1 = 1.1
        t2 = 1.841
        t3 = 3.91899
        t4 = BigFloat("1.0000000000000000000001")
        α1 = zeros(eltype(float(t1)), N)
        α2 = zeros(eltype(float(t2)), N)
        α3 = zeros(eltype(float(t3)), N)
        α4 = zeros(eltype(float(t3)), N)
        αdirect!(α1,t1,1:N)
        αdirect!(α2,t2,1:N)
        αdirect!(α3,t3,1:N)
        αdirect!(α4,t4,1:N)
        @test αdirect.((1:N),t1) ≈ α1
        @test αdirect.((1:N),t2) ≈ α2
        @test αdirect.((1:N),t3) ≈ α3
        @test αdirect.((1:N),t4) ≈ α4
    end

    @testset "basic forward and back recurrence" begin
        # set parameters
        N = 30
        t1 = BigFloat("1.841")
        t2 = BigFloat("1.0000000000000000000001")
        # initialize α
        α1f = zeros(eltype(float(t1)), N)
        α1f[1] = initialα(t1)
        α2f = zeros(eltype(float(t2)), N)
        α2f[1] = initialα(t2)
        α1b = zeros(eltype(float(t1)), N)
        α2b = zeros(eltype(float(t2)), N)
        # compute coefficients via forward recurrence
        αcoefficients!(α1f,t1,BigInt.(2:N))
        αcoefficients!(α2f,t2,BigInt.(2:N))
        # compute coefficients via back recurrence
        backαcoeff!(α1b,t1,BigInt.(1:N))
        backαcoeff!(α2b,t2,BigInt.(1:N))
        # Mathematica α1
        @test α1b[4]  ≈ 0.262708732908399743 ≈ α1f[4]
        @test α1b[6]  ≈ 0.272687692260606064 ≈ α1f[6]
        @test α1b[10] ≈ 0.281264391787711543 ≈ α1f[10]
        @test α1b[20] ≈ 0.288083843194443346 ≈ α1f[20]
        # Mathematica α2
        @test α2b[3]  ≈ 0.98621165663772362 ≈ α2f[3]
        @test α2b[5]  ≈ 0.99152243369113373 ≈ α2f[5]
        @test α2b[10] ≈ 0.99562287407137481 ≈ α2f[10]
        @test α2b[20] ≈ 0.99774034482793111 ≈ α2f[20]
    end

    @testset "evaluation normalized" begin
        t = BigFloat("1.1")
        # Mathematica values
        @test evalQn(0,0.99,t) ≈ 1
        @test evalQn(1,0.5,t) ≈ -0.6623723753894052
        @test evalQn(6,0.12,t) ≈ -1.965171674178137
    end

    @testset "evaluation non-normalized" begin
        t = 1.1
        x = 0.1
        n = 5
        α = zeros(n+1)'
        α[1] = initialα(2*t-1)
        αcoefficients!(α,2*t-1,2:n)
        # compare versions with and without recomputing α with Mathematica results
        @test evalϕn(0,x,t) == 1
        @test evalϕn(1,x,t) ≈ -1.165935217151491
        @test evalϕn(2,x,t) ≈ 0.806910345733665
    end

    @testset "Expansion" begin
        # basis
        t = 1.00001
        Q = SemiclassicalJacobi(t,0,0,-1)
        x = axes(Q,1)
        # test functions
        f1(x) = x^2
        f2(x) = (t-x)^2
        f3(x) = exp(t-x)
        f4(x) = sinh(t*x)
        # test expansion
        y = rand(1)[1]
        @test (Q[:,1:30]*(Q[:,1:30]\f1.(x)))[y] ≈ f1(y)
        @test (Q[:,1:30]*(Q[:,1:30]\f2.(x)))[y] ≈ f2(y)
        @test (Q[:,1:30]*(Q[:,1:30]\f3.(x)))[y] ≈ f3(y)
        @test (Q[:,1:30]*(Q[:,1:30]\f4.(x)))[y] ≈ f4(y)
    end

    @testset "Expansion (adaptive)" begin
        # basis
        t = 1.23
        Q = SemiclassicalJacobi(t,0,0,-1)
        x = axes(Q,1)
        # test functions
        f1(x) = x^2
        f2(x) = (t-x)^2
        f3(x) = exp(t-x)
        f4(x) = sinh(t*x)
        # test expansion
        y = rand(1)[1]
        @test (Q*(Q\f1.(x)))[y] ≈ f1(y)
        @test (Q*(Q\f2.(x)))[y] ≈ f2(y)
        @test (Q*(Q\f3.(x)))[y] ≈ f3(y)
        @test (Q*(Q\f4.(x)))[y] ≈ f4(y)
    end

    @testset "Multiplication by x" begin
        # basis
        t = 1.00001
        Q = SemiclassicalJacobi(t,0,0,-1)
        x = axes(Q,1)
        X = jacobimatrix(Q)
        # test functions
        f1(x) = x^2
        f2(x) = (t-x)^2
        f3(x) = exp(t-x)
        f4(x) = sinh(t*x)
        # test expansion
        y = 0.4781
        @test (Q*(X*(Q\x.^2)))[y] ≈ y*f1(y)
        @test (Q*(X*X*(Q\(t.-x).^2)))[y] ≈ y^2*f2(y)
        @test (Q*(X*(Q\exp.(t.-x))))[y] ≈ y*f3(y)
        @test (Q*(X*(Q\sinh.(t.*x))))[y] ≈ y*f4(y)
    end
end

@testset "Generic lowering operators" begin
    @testset "Iterative lowering" begin
        # lowering a iteratively
        t = 1.812
        PLegendre = SemiclassicalJacobi(t,0,0,0)
        RaisetoLower = SemiclassicalJacobi(t,15,3,6,PLegendre)
        RaisetoCompare = SemiclassicalJacobi(t,12,3,6,PLegendre)
        LoweredPoly = SemiclassicalJacobi(t,12,3,6,RaisetoLower)
        @test LoweredPoly.X[1:100,1:100] ≈ RaisetoCompare.X[1:100,1:100]
        # lowering b iteratively
        t = 1.1
        PLegendre = SemiclassicalJacobi(t,0,0,0)
        RaisetoLower = SemiclassicalJacobi(t,7,6,8,PLegendre)
        RaisetoCompare = SemiclassicalJacobi(t,7,4,8,PLegendre)
        LoweredPoly = SemiclassicalJacobi(t,7,4,8,RaisetoLower)
        @test LoweredPoly.X[1:100,1:100] ≈ RaisetoCompare.X[1:100,1:100]
        # lowering c iteratively
        t = 1.001
        PLegendre = SemiclassicalJacobi(t,0,0,0)
        RaisetoLower = SemiclassicalJacobi(t,5,8,8,PLegendre)
        RaisetoCompare = SemiclassicalJacobi(t,5,8,5,PLegendre)
        LoweredPoly = SemiclassicalJacobi(t,5,8,5,RaisetoLower)
        @test LoweredPoly.X[1:100,1:100] ≈ RaisetoCompare.X[1:100,1:100]
    end

    @testset "Lower all by 1" begin
        t = 1.13
        PLegendre = SemiclassicalJacobi(t,0,0,0)
        RaisetoLower = SemiclassicalJacobi(t,15,13,16,PLegendre)
        RaisetoCompare = SemiclassicalJacobi(t,14,12,15,PLegendre)
        LoweredPoly = SemiclassicalJacobi(t,14,12,15,RaisetoLower)
        @test LoweredPoly.X[1:200,1:200] ≈ RaisetoCompare.X[1:200,1:200]
    end

    @testset "Lower all by 2" begin
        t = 1.102
        PLegendre = SemiclassicalJacobi(t,0,0,0)
        RaisetoLower = SemiclassicalJacobi(t,15,13,16,PLegendre)
        RaisetoCompare = SemiclassicalJacobi(t,13,11,14,PLegendre)
        LoweredPoly = SemiclassicalJacobi(t,13,11,14,RaisetoLower)
        @test LoweredPoly.X[1:200,1:200] ≈ RaisetoCompare.X[1:200,1:200]
    end

    @testset "Mixed lowering" begin
        # mixed a and c lowering
        t = 1.212
        PLegendre = SemiclassicalJacobi(t,0,0,0)
        RaisetoLower = SemiclassicalJacobi(t,15,3,6,PLegendre)
        RaisetoCompare = SemiclassicalJacobi(t,12,3,4,PLegendre)
        LoweredPoly = SemiclassicalJacobi(t,12,3,4,RaisetoLower)
        @test LoweredPoly.X[1:100,1:100] ≈ RaisetoCompare.X[1:100,1:100]
        # mixed b and c lowering
        t = 1.11
        PLegendre = SemiclassicalJacobi(t,0,0,0)
        RaisetoLower = SemiclassicalJacobi(t,7,16,8,PLegendre)
        RaisetoCompare = SemiclassicalJacobi(t,7,15,5,PLegendre)
        LoweredPoly = SemiclassicalJacobi(t,7,15,5,RaisetoLower)
        @test LoweredPoly.X[1:100,1:100] ≈ RaisetoCompare.X[1:100,1:100]
        # mixed a and b lowering
        t = 1.1
        PLegendre = SemiclassicalJacobi(t,0,0,0)
        RaisetoLower = SemiclassicalJacobi(t,5,18,8,PLegendre)
        RaisetoCompare = SemiclassicalJacobi(t,4,17,8,PLegendre)
        LoweredPoly = SemiclassicalJacobi(t,4,17,8,RaisetoLower)
        @test LoweredPoly.X[1:100,1:100] ≈ RaisetoCompare.X[1:100,1:100]
    end

    @testset "Higher order stress tests" begin
        @testset "lowering a" begin
            t = 1.0001
            @test symlowered_jacobimatrix(SemiclassicalJacobi(t,10,10,5,SemiclassicalJacobi(t,0,0,0)),:a)[1:500,1:500] ≈ jacobimatrix(SemiclassicalJacobi(t,9,10,5,SemiclassicalJacobi(t,0,0,0)))[1:500,1:500]
        end
        @testset "lowering b" begin
            t = 3.1
            @test symlowered_jacobimatrix(SemiclassicalJacobi(t,6,6,6,SemiclassicalJacobi(t,0,0,0)),:b)[1:500,1:500]≈ jacobimatrix(SemiclassicalJacobi(t,6,5,6,SemiclassicalJacobi(t,0,0,0)))[1:500,1:500]
            t = 1.11
            @test symlowered_jacobimatrix(SemiclassicalJacobi(t,5,7,4,SemiclassicalJacobi(t,0,0,0)),:b)[1200,1200]≈ jacobimatrix(SemiclassicalJacobi(t,5,6,4,SemiclassicalJacobi(t,0,0,0)))[1200,1200]
        end
        @testset "lowering c" begin
            t = 1.1
            @test symlowered_jacobimatrix(SemiclassicalJacobi(t,2,3,4,SemiclassicalJacobi(t,0,0,0)),:c)[1:500,1:500]≈ jacobimatrix(SemiclassicalJacobi(t,2,3,3,SemiclassicalJacobi(t,0,0,0)))[1:500,1:500]
            t = 1.00001
            @test symlowered_jacobimatrix(SemiclassicalJacobi(t,12,4,7,SemiclassicalJacobi(t,0,0,0)),:c)[1:500,1:500] ≈ jacobimatrix(SemiclassicalJacobi(t,12,4,6,SemiclassicalJacobi(t,0,0,0)))[1:500,1:500]
        end
    end
end


text:
# Model
mutable struct Z3Model
    ptr::PyObject
    numvars::Int
    varvalue::Vector{Float64}
    varlower::Vector{Float64}
    varupper::Vector{Float64}
    vartype::Vector{Symbol}
    varptr::Vector{PyObject}
    varname::Vector{String}
    numconstr::Int
    constr::Vector{PyObject}
    objvar::PyObject
    objvalue::Float64
    objbound::Float64
    sense::Symbol
    objlower::Float64
    objupper::Float64
    status::Symbol
    Z3Model() = new(z3[:Solver](), 0, Float64[],Float64[],Float64[],Symbol[],PyObject[],String[],0,PyObject[],nothing,NaN,NaN,:Min,1e-12,1e12,:Unchecked)
end

push!(model::Z3Model) = model.ptr[:push]()
pop!(model::Z3Model) = model.ptr[:pop]()


text:
using KernelAbstractions
using KernelGradients
using Enzyme
using ROCKernels
using AMDGPU
using Test

include(joinpath(dirname(pathof(KernelAbstractions)), "..", "test", "testsuite.jl"))
include(joinpath(dirname(pathof(KernelGradients)), "..", "test", "testsuite.jl"))

CI = parse(Bool, get(ENV, "CI", "false"))
if CI
    default = "CPU"
else
    default = "ROCM"
end

backend = get(ENV, "KERNELABSTRACTIONS_TEST_BACKEND", default)

if backend != "ROCM"
    @info "ROCM backend not selected"
    exit()
end

if AMDGPU.has_rocm_gpu()
    AMDGPU.allowscalar(false)
    Testsuite.testsuite(ROCDevice, backend, AMDGPU, ROCArray, AMDGPU.ROCDeviceArray)
    # GradientsTestsuite.testsuite(ROCDevice, backend, AMDGPU, ROCArray, AMDGPU.ROCDeviceArray)
elseif !CI
    error("No AMD GPUs available!")
end


text:
abstract type AbstractEllipsoid{T} <: AbstractSet{T} end
dimension(ell::AbstractEllipsoid) = LinearAlgebra.checksquare(ell.Q)

struct HyperSphere <: AbstractEllipsoid{Bool}
    dim::Int
end
dimension(sphere::HyperSphere) = sphere.dim
Polyhedra.polar(sphere::HyperSphere) = sphere

"""
    struct Ellipsoid{T} <: AbstractEllipsoid{T}
        Q::Symmetric{T, Matrix{T}}
    end
"""
struct Ellipsoid{T} <: AbstractEllipsoid{T}
    Q::Symmetric{T, Matrix{T}}
end
ellipsoid(ell::Ellipsoid) = ell
function Polyhedra.project(ell::Ellipsoid, I)
    return project(polar_representation(ell), I)
end
convexity_proof(ell::Ellipsoid) = ell.Q
function scaling_function(ell::Ellipsoid)
    @assert dimension(ell) == 2
    Q = ell.Q
    return (x, y) -> begin
        val = x^2 * Q[1, 1] + 2x*y * Q[1, 2] + y^2 * Q[2, 2]
        if -1e-8 < val < 0
            # `sqrt` would error
            return zero(float(val))
        end
        return sqrt(val)
    end
end

function ellipsoid(ell::PolarOf{<:Ellipsoid})
    ellipsoid(polar_representation(ell))
end
function polar_representation(ell::PolarOf{<:Ellipsoid})
    Ellipsoid(inv(ell.set.Q))
end
function polar_representation(ell::Ellipsoid)
    polar(Ellipsoid(inv(ell.Q)))
end

function zero_eliminate(ell::Ellipsoid, I)
    J = setdiff(1:dimension(ell), I)
    return Ellipsoid(Symmetric(ell.Q[J, J]))
end

function _print_gauge_function(ell::Ellipsoid; digits=6)
    DynamicPolynomials.@polyvar x[1:2]
    print(" ")
    Q = ell.Q
    if digits !== nothing
        Q = round.(Q, digits=digits)
    end
    println(x' * Q * x)
end

struct LiftedEllipsoid{T}
    P::Matrix{T}
end
dimension(ell::LiftedEllipsoid) = LinearAlgebra.checksquare(ell.P) - 1

function perspective_variables(ell::Union{Ellipsoid, LiftedEllipsoid})
    return nothing
end
function space_variables(ell::Union{Ellipsoid, LiftedEllipsoid})
    return nothing
end

function LiftedEllipsoid(t::Translation{<:Ellipsoid})
    ell = t.set
    md = ell.Q * t.c
    δ = t.c' * md-1
    d = -md
    D = ell.Q
    P = [δ d'
         d D]
    LiftedEllipsoid(P)
end

function Bbβλ(P)
    n = LinearAlgebra.checksquare(P) - 1
    ix = 1 .+ (1:n)
    β = P[1, 1]
    b = P[1, ix]
    B = P[ix, ix]
    λ = dot(b, B \ b) - β
    @assert λ >= 0
    B, b, β, λ
end
function ellipsoid(ell::LiftedEllipsoid)
    # P is
    # λ * [c'Qc-1  -c'Q
    #         -Qc   Q]
    # Let P be [β b'; b B]
    # We have
    # β = λ c'Qc - λ
    # b = -λ Qc <=> Q^{-1/2}b = -λ Q^{1/2} c
    # hence
    # λ c'Qc = β + λ
    # λ^2 c'Qc = b'Q^{-1}b = λ b'B^{-1}b <=> λ c'Qc = b'B^{-1}b
    # Hence λ = b'B^{-1}b - β
    B, b, β, λ = Bbβλ(ell.P)
    c = -(B \ b)
    Q = B / λ
    Translation(Ellipsoid(Symmetric(Q)), c)
end

function _HPH(D, d, δ)
    P = [δ d'
         d D]
end

_HPH(ell::Ellipsoid) = _HPH(ell.Q, zeros(size(q.Q, 1)), -1.0)

"""
    struct ShiftedEllipsoid{T}
        Q::Symmetric{T, Matrix{T}}
        b::Vector{S}
        β::S
    end

Set ``\\{\\, x \\mid x^\\top Q x + 2 b^\\top x + \\beta \\le 0 \\,\\}``.
"""
struct ShiftedEllipsoid{T} <: AbstractEllipsoid{T}
    Q::Symmetric{T, Matrix{T}}
    b::Vector{T}
    β::T
end
_HPH(q::ShiftedEllipsoid) = _HPH(q.Q, q.b, q.β)
convexity_proof(ell::ShiftedEllipsoid) = ell.Q

function LiftedEllipsoid(qc::HouseDualOf{<:AbstractEllipsoid})
    return LiftedEllipsoid(inv(_HPH(perspective_dual(qc))))
end
function ellipsoid(qc::HouseDualOf{<:AbstractEllipsoid})
    return ellipsoid(LiftedEllipsoid(qc))
end
function Polyhedra.project(ell::HouseDualOf{<:AbstractEllipsoid},
                           I::AbstractVector)
    return project(ellipsoid(ell), I)
end

function PerspectiveInteriorEllipsoid(ell::LiftedEllipsoid)
    Pd = inv(le.P)
    H = _householder(h[state])
    HPdH = H * Pd * H
    # HPdH is not like a solution what would be obtained by solving the program
    # since the λ computed for unlifting it is maybe not one.
    # Therefore, the S-procedure's λ for the constraints will be different.
    B, b, β, λ = Bbβλ(HPdH)
    ps[state] = y' * H * _HPH(B/λ, b/λ, β/λ) * H * y
    error("TODO: LiftedEllipsoid -> PerspectiveInteriorEllipsoid")
end


text:
function eftAddHiLo{T<:StdFloat}(ahi::T, alo::T, bhi::T, blo::T)
    s1, s2 = eftAdd(ahi,bhi)
    t1, t2 = eftAdd(alo,blo)
    s2 += t1
    s1, s2 = eftAddGTE(s1,s2)
    s2 += t2
    eftAddGTE(s1,s2)
end

function eftAddHiLo{T<:StdFloat}(ahi::T, alo::T, b::T)
    s1, s2 = eftAdd(ahi,b)
    s2 += alo
    eftAddGTE(s1,s2)
end

function eftSubHiLo{T<:StdFloat}(ahi::T, alo::T, bhi::T, blo::T)
    s1, s2 = eftSub(ahi,bhi)
    t1, t2 = eftSub(alo,blo)
    s2 += t1
    s1, s2 = eftAddGTE(s1,s2)
    s2 += t2
    eftAddGTE(s1,s2)
end

function eftSubHiLo{T<:StdFloat}(ahi::T, alo::T, b::T)
    s1, s2 = eftSub(ahi,b)
    s2 += a.lo
    eftAddGTE(s1,s2)
end

@inline function inlEftAddGTE{T<:StdFloat}(a::T, b::T)
  hi = a + b
  lo = b - (hi - a)
  hi, lo
end

@inline function inlEftMul{T<:StdFloat}(a::T, b::T)
    hi = a * b
    lo = fma(a, b, -hi)
    hi,lo
end

function eftMulHiLo{T<:StdFloat}(ahi::T, alo::T, bhi::T, blo::T)
  #t1,t2 = eftMul(ahi,bhi)
  t1, t2 = inlEftMul(ahi, bhi)
  t3 = ahi * blo
  t4 = alo * bhi
  t3 = t3 + t4
  t4 = t2 + t3
  inlEftAddGTE(t1,t4)
end

function eftMulHiLo{T<:StdFloat}(ahi::T, alo::T, b::T)
  t1,t2 = eftMul(ahi,b)
  t4 = alo * b
  t6 = t2 + t4
  eftAddGTE(t1,t6)
end

#=
function inv{T<:StdFloat}(bhi::T, blo::T)
  q1 = one(T) / bhi
  r  = one(T) - (q1 * b)

  q2 = r.hi / bhi
  r = r - (q2 * b)

  q3 = r.hi / b.hi

  q1,q2 = eftAddGTE(q1, q2)
  q1,q2 = eftAddAs2(q1,q2,q3)
  FloatFloat(q1,q2)
end

function (/){T<:StdFloat}(a::FloatFloat{T}, b::FloatFloat{T})
  q1 = a.hi / b.hi
  r  = a - (q1 * b)

  q2 = r.hi / b.hi
  r = r - (q2 * b)

  q3 = r.hi / b.hi

  q1,q2 = eftAddGTE(q1, q2)
  q1,q2 = eftAddAs2(q1,q2,q3)

  FloatFloat(q1,q2)
end

function (/){T<:StdFloat}(a::FloatFloat{T}, b::T)
  q1 = a.hi / b
  r  = a - (q1 * b)

  q2 = r.hi / b
  r = r - (q2 * b)

  q3 = r.hi / b

  q1,q2 = eftAddGTE(q1, q2)
  q1,q2 = eftAddAs2(q1,q2,q3)

  FloatFloat(q1,q2)
end
=#


text:
export Parameters, default_parameters, print_parameters

mutable struct Parameters
    mu::Integer
    lambda::Integer
    mutrate::Real
    targetfitness::Real

    numinputs::Integer
    numoutputs::Integer
    nodearity::Integer

    numinteriors::Integer   # Number of gates. For LinCircuits, number of instructions
    numlevelsback::Integer  # For LinCircuits, number of registers
end

function Parameters(numinputs, numoutputs, numinteriors, numlevelsback)
    nodearity = 2
    mu = 1
    lambda = 4
    mutrate = 0.05
    targetfitness = 0.0
    MyInt_bits = MyIntBits( MyInt )
    #println("MyInt_bits: ",MyInt_bits)
    if numinteriors > MyInt_bits && maxints_for_degen > MyInt_bits
      println("maxints_for_degen: ",maxints_for_degen,"  MyInt_bits: ",MyInt_bits,"  numinteriors: ",numinteriors)
      error("maxints_for_degen > MyInt_bits in function Parameters.  Run with a larger width MyInt" )
    end
    return Parameters(mu, lambda, mutrate, targetfitness, numinputs, numoutputs, nodearity, numinteriors, numlevelsback)
end

function Parameters( ; numinputs=2, numoutputs=1, nodearity=2, numinteriors=4, numlevelsback=3 )
    mu = 1
    lambda = 4
    mutrate = 0.05
    targetfitness = 0.0
    MyInt_bits = MyIntBits( MyInt )
    #println("MyInt_bits: ",MyInt_bits)
    if numinteriors > MyInt_bits && maxints_for_degen > MyInt_bits
      println("maxints_for_degen: ",maxints_for_degen,"  MyInt_bits: ",MyInt_bits,"  numinteriors: ",numinteriors)
      error("maxints_for_degen > MyInt_bits in function Parameters.  Run with a larger width MyInt" )
    end
    return Parameters(mu, lambda, mutrate, targetfitness, numinputs, numoutputs, nodearity, numinteriors, numlevelsback)
end

function print_parameters(f:: IO, p::Parameters; comment::Bool=false )
  if !comment
    println(f,"MyInt: ",MyInt)
    println(f,"numinputs: ",p.numinputs)
    println(f,"numoutputs: ",p.numoutputs)
    println(f,"numinteriors: ",p.numinteriors)
    println(f,"numlevelsback: ",p.numlevelsback)
    println(f,"nodearity: ",p.nodearity)
  else
    println(f,"# MyInt: ",MyInt)
    println(f,"# numinputs: ",p.numinputs)
    println(f,"# numoutputs: ",p.numoutputs)
    println(f,"# numinteriors: ",p.numinteriors)
    println(f,"# numlevelsback: ",p.numlevelsback)
    println(f,"# nodearity: ",p.nodearity)
  end
end

function print_parameters(p::Parameters; comment::Bool=false  )
  f = Base.stdout
  print_parameters( f, p, comment=comment )
end

function MyIntBits( my_int::Type )
  if my_int == UInt8
    8
  elseif my_int == UInt16
    16
  elseif my_int == UInt32
    32
  elseif my_int == UInt64
    64
  elseif my_int == UInt128
    128
  else
    error("error in MyIntBits")
  end
end


text:
module WrapperArrays

import LinearAlgebra

export iswrapper, adopt, storage, childstyle

using Base.Broadcast: BroadcastStyle

#=
    This file defines
        1. What it means to be a "simple wrapper array". Any simple wrapper
        must define the `adopt` function, which describes how to construct
        analogous wrapper arrays with new parents, and specialize Base.parent.
        2. Corresponding methods for "simple wrapper arrays" in Base.
        3. A conveience type `ShallowArray` which defines several pass-through
        methods for easy construction of wrapper arrays.
    Ideally, numbers 1 and 2 should live in Adapt.jl, (note that `adapt` can be
    implemented with adopt).

    Wrappers need to define
        parent
        adopt
        iswrapper?
        dataids?
        unaliascopy?
        alias?

=#

"""
    adopt(child, arg)

Wrap `arg` in an analogous wrapper array to `child`. This function should
create an array with the same semantics as `child`. Generally, if `p` is
mostly the same array as `parent(x)`, it should hold that `adopt(x, p)` is
mostly the same array as `x`.

See also: [`parent`](@ref), [`iswrapper`](@ref)
"""
adopt(arr, arg) = throw(MethodError(adopt, (arr, arg)))

"""
    storage(arr)

Return the deepest ancestor of the array `arr`. If `arr` is a wrapper array,
returns `storage(parent(arr))`. Otherwise, returns `arr`.

See also: [`parent`](@ref), [`iswrapper`](@ref)
"""
function storage(f::F, arr) where {F}
    if iswrapper(arr)
        return parent(arr)
    else
        return arr
    end
end

"""
    iswrapper(arr)

A trait function which returns `true` if and only if `arr !== parent(arr)`. We
reccommend specializing this function with hard-coded implementations for custom
array types to improve type inference of wrapper array functions like
[`storage`](@ref).

See also: [`parent`](@ref)
"""
iswrapper(arr) = arr !== parent(arr)

#Base
iswrapper(::Array) = false

iswrapper(::LinearAlgebra.Transpose) = true
adopt(arr::LinearAlgebra.Transpose, arg::AbstractVecOrMat) = LinearAlgebra.transpose(arg)

iswrapper(::LinearAlgebra.Adjoint) = true
adopt(arr::LinearAlgebra.Adjoint, arg::AbstractVecOrMat) = LinearAlgebra.adjoint(arg)

iswrapper(::SubArray) = true
adopt(arr::SubArray, arg) = SubArray(arg, parentindices(arr))

iswrapper(::LinearAlgebra.LowerTriangular) = true
adopt(arr::LinearAlgebra.LowerTriangular, arg) = LinearAlgebra.LowerTriangular(arg)

iswrapper(::LinearAlgebra.UnitLowerTriangular) = true
adopt(arr::LinearAlgebra.UnitLowerTriangular, arg) = LinearAlgebra.UnitLowerTriangular(arg)

iswrapper(::LinearAlgebra.UpperTriangular) = true
adopt(arr::LinearAlgebra.UpperTriangular, arg) = LinearAlgebra.UpperTriangular(arg)

iswrapper(::LinearAlgebra.UnitUpperTriangular) = true
adopt(arr::LinearAlgebra.UnitUpperTriangular, arg) = LinearAlgebra.UnitUpperTriangular(arg)

iswrapper(::LinearAlgebra.Diagonal) = true
adopt(arr::LinearAlgebra.Diagonal, arg) = LinearAlgebra.Diagonal(arg)

iswrapper(::Base.ReshapedArray) = true
adopt(arr::Base.ReshapedArray, arg) = reshape(arg, arr.dims)

iswrapper(::PermutedDimsArray) = true
function adopt(arg::Arg, arr::PermutedDimsArray{<:Any,N,perm,iperm}) where {T,N,perm,iperm,Arg<:AbstractArray{T, N}}
    return PermutedDimsArray{T,N,perm,iperm,Arg}(arg)
end

"""
    `childstyle(::Type{<:AbstractArray}, ::BroadcastStyle)`

Broadcast styles are used to determine behavior of objects under broadcasting.
To customize the broadcasting behavior of a wrapper array, one can first declare
how the broadcast style should behave under broadcasting after the wrapper array
is applied by overriding the `childstyle` method.
"""
@inline childstyle(Arr::Type{<:AbstractArray}, ::BroadcastStyle) = BroadcastStyle(Arr)

end


text:
using Distributions
using Reinforce
using Test

struct TestEpisodePolicy <: AbstractPolicy end
action(::TestEpisodePolicy, r, s, A) = 2

@testset "MultiArmedBanditEnv" begin

@info "Reinforce.MultiArmedBandit"

@testset "constructor" begin
  let
    σ = 5
    env = MultiArmedBandit(10, 42; σ = σ)

    @test iszero(reward(env))
    @test maxsteps(env) == 42
    @test actions(env, nothing) == 1:10
    @test length(env.arms) == 10
    for arm ∈ env.arms
      @test arm.σ == σ
    end
  end

  let
    d₁ = Uniform(1, 42)
    d₂ = Gamma()
    env = MultiArmedBandit(d₁, d₂)

    @test iszero(reward(env))
    @test actions(env, nothing) == 1:2
    @test length(env.arms) == 2
    @test env.arms[1] == d₁
    @test env.arms[2] == d₂
  end
end  # @testset "constructor"

@testset "episode iterator" begin

  let
    env = MultiArmedBandit(10, 5)
    π = TestEpisodePolicy()

    @test iszero(reward(env))
    for (s, a, r, s′) ∈ Episode(env, π)
      @test a == 2
    end
    @test !iszero(reward(env))
  end
end

end  # @testset "MultiArmedBanditEnv"


text:
@define_scenario WeekAheadScenario = begin
    ρ::PriceCurve{168,Float64}
    Q̃::InflowSequence{7,typeof(Skellefteälven),Float64}

    @expectation begin
        ρ = mean([s.ρ for s in scenarios])
        Q̃ = mean([s.Q̃ for s in scenarios])
        return WeekAheadScenario(ρ, Q̃, probability = 1.0)
    end
end

@sampler RecurrentWeekAheadSampler = begin
    date::Date
    plants::PlantCollection
    price_forecaster::Forecaster
    flow_forecaster::Forecaster

    @sample WeekAheadScenario begin
        price_curve = Vector{Float64}()
        for d = 1:7
            daily_curve = forecast(sampler.price_forecaster, month(sampler.date))
            while !all(daily_curve .>= 0)
                daily_curve = forecast(sampler.price_forecaster, month(sampler.date))
            end
            append!(price_curve, daily_curve[:])
        end
        flows = forecast(sampler.flow_forecaster, week(sampler.date))
        while !all(flows[:, 2] .>= 0)
            flows = forecast(sampler.flow_forecaster, week(sampler.date))
        end
        return WeekAheadScenario(PriceCurve(price_curve), InflowSequence(sampler.plants, flows))
    end
end


text:
using heuristicopt

function main()
    # obj supports only maximazation
    initsate = rand(0:1, 100)
    tmin = 10
    tmax = 10000
    steps = 10000
    param = SA.SAParameter(tmin, tmax, steps)
    model = SA.SAModel(param, move, objfunc)
    besttrial = SA.predict(model, initsate)
end

function move(array)
    N = length(array)
    point = rand(1:N)
    array[point] = 1 - array[point]
    return array
end

function objfunc(array)
    sum(array) * length(array)
end

main()

max_stars_repo_name: ariririri/heuristicopt

max_issues_repo_name: ariririri/heuristicopt


text:
#=
real vector domain ℝᵈ₊

extended formulation
∃ λ, u ≥ Σᵢ λᵢ, (λᵢ, v, wᵢ) ∈ EpiPerSepSpectralCone
=#

mutable struct VectorEpiPerSepSpectral{D <: PrimDual, E <: NatExt} <: Cache
    oa_s::Vector{AE}
    h::SepSpectralFun
    d::Int
    λ::Vector{VR}
    function VectorEpiPerSepSpectral{D, E}() where {D <: PrimDual, E <: NatExt}
        return new{D, E}()
    end
end

function create_sepspectral_cache(
    ::Type{VectorCSqr{RealF}},
    use_dual::Bool,
    d::Int,
    opt::Optimizer,
)
    D = primal_or_dual(use_dual)
    E = nat_or_ext(opt, d)
    return VectorEpiPerSepSpectral{D, E}()
end

function MOIPajarito.Cones.add_init_cuts(
    cache::VectorEpiPerSepSpectral{D},
    opt::Optimizer,
) where {D}
    (_, v) = swap_epiper(D, cache.oa_s[1:2]...)
    @views w = cache.oa_s[3:end]

    # add variable bounds
    JuMP.@constraint(opt.oa_model, v >= 0)
    if dom_pos(D, cache.h)
        JuMP.@constraint(opt.oa_model, w .>= 0)
    end
    opt.use_init_fixed_oa || return

    # add cuts at p = 1 and R = r₀ e
    r_vals = init_r_vals(D, cache.h)
    for r0 in r_vals
        r = fill(r0, cache.d)
        cuts = _get_cuts(1.0, r, cache, opt)
        JuMP.@constraint(opt.oa_model, cuts .>= 0)
    end
    return
end

function MOIPajarito.Cones.get_subp_cuts(
    z::Vector{RealF},
    cache::VectorEpiPerSepSpectral{D},
    opt::Optimizer,
) where {D}
    (p, _) = swap_epiper(D, z[1:2]...)
    return _get_cuts(p, z[3:end], cache, opt)
end

function MOIPajarito.Cones.get_sep_cuts(
    s::Vector{RealF},
    cache::VectorEpiPerSepSpectral{D},
    opt::Optimizer,
) where {D}
    h = cache.h
    ws = s[3:end]
    if dom_pos(D, h)
        ws = max.(ws, 1e-10)
    end

    (us, vs) = swap_epiper(D, s[1:2]...)
    v_pos = max(vs, 1e-9)
    if us - per_sepspec(val_or_conj(D), h, v_pos, ws) > -opt.tol_feas
        return AE[]
    end

    # gradient cut is (1, h⋆(r), r) at r = -h'(ws / vs)
    r = -h_grad(D, h, ws / v_pos)
    return _get_cuts(1.0, r, cache, opt)
end

# unextended formulation

function _get_cuts(
    p::RealF,
    r::Vector{RealF},
    cache::VectorEpiPerSepSpectral{D, Nat},
    opt::Optimizer,
) where {D}
    # strengthened cut is (p, p * h⋆(r / p), r)
    q = per_sepspec(conj_or_val(D), cache.h, p, r)
    (p, q) = swap_epiper(D, p, q)

    z = vcat(p, q, r)
    cut = dot_expr(z, cache.oa_s, opt)
    return [cut]
end

# extended formulation

function MOIPajarito.Cones.num_ext_variables(
    cache::VectorEpiPerSepSpectral{<:PrimDual, Ext},
)
    return cache.d
end

function MOIPajarito.Cones.extend_start(
    cache::VectorEpiPerSepSpectral{D, Ext},
    s_start::Vector{RealF},
    opt::Optimizer,
) where {D}
    (_, v_start) = swap_epiper(D, s_start[1:2]...)
    if v_start < 1e-8
        return zeros(cache.d)
    end
    w_start = [max(s_start[i], 1e-9) for i in 3:length(s_start)]
    return [per_sepspec(val_or_conj(D), cache.h, v_start, [w_i]) for w_i in w_start]
end

function MOIPajarito.Cones.setup_auxiliary(
    cache::VectorEpiPerSepSpectral{D, Ext},
    opt::Optimizer,
) where {D}
    @assert cache.d >= 2
    (u, _) = swap_epiper(D, cache.oa_s[1:2]...)
    λ = cache.λ = JuMP.@variable(opt.oa_model, [1:(cache.d)])
    JuMP.@constraint(opt.oa_model, u >= sum(λ))
    return λ
end

function _get_cuts(
    p::RealF,
    r::Vector{RealF},
    cache::VectorEpiPerSepSpectral{D, Ext},
    opt::Optimizer,
) where {D}
    (_, v) = swap_epiper(D, cache.oa_s[1:2]...)
    @views w = cache.oa_s[3:end]
    λ = cache.λ

    cuts = AE[]
    for i in 1:(cache.d)
        r_i = r[i]
        # r_i = min(r[i], 1e-7) # iszero(r_i) && continue # TODO?
        # strengthened disaggregated cut on (λᵢ, v, wᵢ) is (p, p * h⋆(rᵢ / p), rᵢ)
        # TODO check math
        q = per_sepspec(conj_or_val(D), cache.h, p, [r_i])
        cut = JuMP.@expression(opt.oa_model, p * λ[i] + q * v + r_i * w[i])
        push!(cuts, cut)
    end
    return cuts
end


text:
using FewShotAnomalyDetection
using Flux
using MLDataPattern
using FluxExtensions
using Adapt
using DataFrames
using CSV
using Serialization
using IPMeasures
using IPMeasures: crit_mmd2_var
using Plots
plotlyjs()

include("experimentalutils.jl")
include("vae.jl")

outputFolder = mainfolder * "experiments/twostagesvae_scoretest_mmdpval/"
mkpath(outputFolder)

s = SVAEtwocaps(8, 8, 8, 2, "relu", "Dense", :scalarsigma, Float32)
FewShotAnomalyDetection.set_normal_μ_nonparam(s, vcat(Float32(1), zeros(7)))
samplepz(n::Int) = FewShotAnomalyDetection.samplez(s, ones(8, n) .* FewShotAnomalyDetection.normalizecolumns(s.priorμ), ones(1, n) .* s.priorκ)

n = 1000
z1 = samplepz(n)
z2 = samplepz(n)
γs = -10:0.05:2
cs = [crit_mmd2_var(IPMeasures.IMQKernel(10.0 ^ γ), z1, z2, IPMeasures.pairwisecos) for γ in γs]
plot(cs)
global gγ = 10 ^ γs[argmax(cs)]
println("We chose kernal size $gγ for p(z) null hyp")
null_dst = zeros(Float64, n)
for i in 1:n
    null_dst[i] = FewShotAnomalyDetection.mmd_imq(samplepz(n), samplepz(n), gγ)
end
sort!(null_dst)
mmdpval(x) = searchsortedfirst(null_dst, x) / n

function runExperiment(datasetName, train, test, inputDim, hiddenDim, latentDim, numLayers, nonlinearity, layerType, β, batchSize = 100, numBatches = 10000, i = 0)

    T = Float32

    encoder = Adapt.adapt(T, FluxExtensions.layerbuilder(inputDim, hiddenDim, (latentDim - 1) * 2, numLayers + 1, nonlinearity, "linear", layerType))
    decoder = Adapt.adapt(T, FluxExtensions.layerbuilder((latentDim - 1), hiddenDim, inputDim + 1, numLayers + 1, nonlinearity, "linear", layerType))
    outerVAE = VAE(encoder, decoder, β, :scalarsigma)
    opt = Flux.Optimise.ADAM(1e-4)
    cb = Flux.throttle(() -> println("$datasetName outer VAE: $(printingloss(outerVAE, train[1]))"), 5)
    Flux.train!((x, y) -> loss(outerVAE, x), Flux.params(outerVAE), RandomBatches((train[1], zero(train[2])), batchSize, numBatches), opt, cb = cb)

    svae = SVAEtwocaps(latentDim - 1, hiddenDim, latentDim, numLayers, nonlinearity, layerType, :scalarsigma, T)
    FewShotAnomalyDetection.set_normal_μ_nonparam(svae, vcat(T(1), zeros(latentDim - 1)))
    opt = Flux.Optimise.RMSProp(1e-4)

    for i in 1:(numBatches ÷ 1000)
        z = FewShotAnomalyDetection.samplez(svae, FewShotAnomalyDetection.zparams(svae, zparams(outerVAE, train[1])[1])...).data
        zp = FewShotAnomalyDetection.samplez(svae, ones(latentDim, size(train[1], 2)) .* FewShotAnomalyDetection.normalizecolumns(svae.priorμ), ones(1, size(train[1], 2)) .* svae.priorκ)
        γs = -10:0.05:2
        cs = [crit_mmd2_var(IPMeasures.IMQKernel(10.0 ^ γ), z, zp, IPMeasures.pairwisecos) for γ in γs]
        gamma = 10 ^ γs[argmax(cs)]
        println("We chose kernal size $gamma for svae learning")
        learnRepresentation!(data, labels) = wloss(svae, data, (x, y) -> FewShotAnomalyDetection.mmd_imq(x, y, gamma))
        printing_learnRepresentation!(data, labels) = printing_wloss(svae, data, (x, y) -> FewShotAnomalyDetection.mmd_imq(x, y, gamma))
        cb = Flux.throttle(() -> println("$datasetName inner SVAE: $(printing_learnRepresentation!(samplez(outerVAE, train[1]), zero(train[2])))"), 5)
        Flux.train!((x, y) -> learnRepresentation!(samplez(outerVAE, x), y), Flux.params(svae), RandomBatches((train[1], zero(train[2])), batchSize, 1000), opt, cb = cb)
    end

    ztrain = Flux.Tracker.data(FewShotAnomalyDetection.zparams(svae, zparams(outerVAE, train[1])[1])[1])
    ztest = Flux.Tracker.data(FewShotAnomalyDetection.zparams(svae, zparams(outerVAE, test[1])[1])[1])
    ztrain_dst = FewShotAnomalyDetection.mmd_imq(samplepz(n), ztrain, gγ)
    ztest_dst = FewShotAnomalyDetection.mmd_imq(samplepz(n), ztest, gγ)
    ztrain_mmd_pval = mmdpval(ztrain_dst)
    ztest_mmd_pval = mmdpval(ztest_dst)

    data = Flux.Tracker.data(zparams(outerVAE, test[1])[1])
    zs = Flux.Tracker.data(FewShotAnomalyDetection.zparams(svae, data)[1])
    us = Flux.Tracker.data(svae.g(zs)[1:end-1, :])
    log_pxv = vec(collect(log_pxexpectedz(outerVAE, test[1], us)'))
    log_pzs = vec(collect(log_pz(svae, data)'))
    log_pzs_jacobian_enc = vec(collect(log_pz_jacobian_encoder(svae, data)'))
    log_pzs_jacobian_dec = vec(collect(log_pz_jacobian_decoder(svae, zs)'))
    log_det_jac_enc_vae = vec(collect(log_det_jacobian_encoder(outerVAE, test[1])'))
    log_det_jac_dec_vae = vec(collect(log_det_jacobian_decoder(outerVAE, us)'))

    auc_pxv = computeauc(.-log_pxv, test[2] .- 1)
    auc_pz = computeauc(.-log_pzs, test[2] .- 1)
    auc_pz_jacobian_enc = computeauc(.-(log_pzs_jacobian_enc .+ log_det_jac_enc_vae), test[2] .- 1)
    auc_pz_jacobian_dec = computeauc(.-(log_pzs_jacobian_dec .- log_det_jac_dec_vae), test[2] .- 1)
    auc_pxv_pz = computeauc(.-(log_pxv .+ log_pzs), test[2] .- 1)
    auc_pxv_pz_jacobian_enc = computeauc(.-(log_pxv .+ log_pzs_jacobian_enc .+ log_det_jac_enc_vae), test[2] .- 1)
    auc_pxv_pz_jacobian_dec = computeauc(.-(log_pxv .+ log_pzs_jacobian_dec .- log_det_jac_dec_vae), test[2] .- 1)

    serialize(outputFolder * "twostagesvae-$datasetName-$i-svae.jls", svae)
    serialize(outputFolder * "twostagesvae-$datasetName-$i-vae.jls", outerVAE)
    serialize(outputFolder * "twostagesvae-$datasetName-$i-train.jls", train)
    serialize(outputFolder * "twostagesvae-$datasetName-$i-test.jls", test)

    return DataFrame(dataset = datasetName, idim = inputDim, hdim = hiddenDim, ldim = latentDim, layers = numLayers, i = i, auc_pxv = auc_pxv, auc_pz = auc_pz, auc_pz_jacobian_enc = auc_pz_jacobian_enc,
                    auc_pz_jacobian_dec = auc_pz_jacobian_dec, auc_pxv_pz = auc_pxv_pz, auc_pxv_pz_jacobian_enc = auc_pxv_pz_jacobian_enc, auc_pxv_pz_jacobian_dec = auc_pxv_pz_jacobian_dec,
                    ztrain_mmd_pval = ztrain_mmd_pval, ztest_mmd_pval = ztest_mmd_pval)
end

datasets = ["breast-cancer-wisconsin"]
difficulties = ["easy"]
batchSize = 100
iterations = 10000

if length(ARGS) != 0
    datasets = [ARGS[1]]
    difficulties = ["easy"]
end

for i in 1:10
    for (dn, df) in zip(datasets, difficulties)
        train, test, clusterdness = loaddata(dn, df)

        evaluateOneConfig = p -> runExperiment(dn, train, test, size(train[1], 1), p..., batchSize, iterations, i)
        results = gridsearch(evaluateOneConfig, [64], [8], [3], ["swish"], ["Dense"], [1.])

        CSV.write(outputFolder * "twostagesvae-$dn-$i.csv", vcat(results...))
    end
end


text:
struct CyclicVector{T} <: AbstractVector{T}
    s::Int64
    V::Vector{T}
end
CyclicVector(V::Vector{T}) where T = CyclicVector{T}(length(V),V)
Base.getindex(cV::CyclicVector, idx) = cV.V[((idx .- 1) .% cV.s .+ 1)]
Base.size(LV::CyclicVector) = (LV.s,)
Base.size(LV::CyclicVector, ::Integer) = LV.s

abstract type subArray{mT} end
struct SubMX{mT} <: subArray{mT}
    ranges::Vector{Tuple{UnitRange{Int64},UnitRange{Int64}}}
    MXs::Vector{mT}
end
SubMX(range::Tuple{<:AbstractVector,<:AbstractVector}, MX::SMatrix{<:Real}) = SubMX([range], [MX])
SubMX(range::Tuple{<:AbstractVector,<:AbstractVector}, MX::AbstractMatrix{<:Real}) = SubMX([range], [MX])
struct SubV{vT} <: subArray{vT}
    V::vT
end

abstract type AbstractResult{d} end
struct Result{d,lddepT<:AbstractLDDEProblem{d},mT,submxT,subvT,tT,AavgT} <: AbstractResult{d}
    LDDEP::lddepT
    method::mT #::DiscretizationMethod
    subMXs::Vector{Vector{submxT}} # [[A(t1),A(t2),...],[B1(t1),B1(t2),...],[B2(t1),B2(t2),...],...]
    subVs::Vector{subvT} # [c(t1),c(t2),...]
    #
    A_avgs::AavgT
    ts::Vector{tT} # [0,t1,t2,...]
    n_steps::Int64 # number of time steps
    n::Int64 # Large discretisation matrix size
    calculate_additive::Bool
    # d::Integer # state space dimension
end

calculate_Aavgs(A::ProportionalMX{d,<:Function}, ts::AbstractVector{<:Real},Δt::Real) where d = [(quadgk(A, ts[i], ts[i + 1])[1] ./ Δt) for i in 1:length(ts)-1]
calculate_Aavgs(A::ProportionalMX{d,<:mT}, ts::AbstractVector{<:Real},Δt::Real)  where d where mT <: AbstractMatrix{T} where T = CyclicVector([A.MX])

function Result(LDDEP::LDDEProblem{d,AT,BT, cT}, method::DiscretizationMethod{fT}, DiscretizationLength::Real; n_steps::Int64=nStepOfLength(DiscretizationLength, method.Δt), calculate_additive::Bool=false) where {d,AT,BT,cT,N,fT}
    # DiscretizationLength discretisated time interval length
    # n_steps: how many mapping matrix to calculate
    ts =  collect(fT,take(n_steps + 1, iterated(x -> method.Δt + x, zero(method.Δt))))
    n = (rOfDelay(DiscretizationLength, method) + 1) * d
    A_avgs = calculate_Aavgs(LDDEP.A, ts, method.Δt)
    subMXs = [Vector{SubMX{eltype(A_avgs)}}(undef, n_steps) for i in 1:(length(LDDEP.Bs) + 1)] # []
    if calculate_additive
        # subVs = Vector{SubV{eltype(LDDEP.cT.V)}}(undef, n_steps) # []
        subVs = Vector{SubV{SVector{d,eltype(eltype(A_avgs))}}}(undef, n_steps) # []
    else
        # subVs = Vector{SubV{eltype(LDDEP.cT.V)}}(undef, 0)
        subVs = Vector{SubV{SVector{d,eltype(eltype(A_avgs))}}}(undef, 0)
    end

    Result(LDDEP, method, subMXs, subVs, A_avgs, ts, n_steps, n, calculate_additive)
end

struct DiscreteMapping{tT,mxT,vT}
    ts::Vector{tT}
    mappingMXs::Vector{mxT}
    mappingVs::Vector{vT}
end

function initializeOneStepMapping(dm::DiscreteMapping{tT,mxT,vT}) where {tT,mxT,vT}
    DiscreteMapping(Vector{tT}(undef,2),Vector{mxT}(undef,1),Vector{vT}(undef,1))
end