using SparseArrays
using CUDA
using Random
function D2_test(idata,odata,Nx,Ny,h)
    for i = 1:Nx
        for j = 1:Ny
            global_index_x = (i-1)*Ny+j

            # idata_index_x = (i-1)*Ny + j + test_mod(i,Nx) * Nx
            # idata_index_y = (i-1)*Ny + j + test_mod(j,Ny)

            idata_index_x = (i-1)*Ny + j + (div(2*Nx-i-3,Nx-2) - 1) * Nx
            idata_index_y = (i-1)*Ny + j + div(2*Ny-j-3,Ny-2) - 1

            odata[global_index_x] = ( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])
                                    + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) ) / h^2
        end
    end
end

function D2_test_v2(idata,odata,Nx,Ny,h)
    for i = 1:Nx
        for j = 1:Ny
            global_index_x = (i-1)*Ny+j

            # idata_index_x = (i-1)*Ny + j + test_mod(i,Nx) * Nx
            # idata_index_y = (i-1)*Ny + j + test_mod(j,Ny)

            offset_x = div(2*Nx-i-3,Nx-2) - 1
            offset_y = div(2*Ny-j-3,Ny-2) - 1

            idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
            idata_index_y = (i-1)*Ny + j + offset_y

            odata[global_index_x] = ( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])
                                    + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) ) / 2^(abs(offset_x) + abs(offset_y))
        end
    end
end

function matrix_free_A_v2(idata,odata,Nx,Ny,h)
    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    odata .= 0
    for i = 1:Nx
        for j = 1:Ny
            global_index = (i-1)*Ny+j

            # idata_index_x = (i-1)*Ny + j + test_mod(i,Nx) * Nx
            # idata_index_y = (i-1)*Ny + j + test_mod(j,Ny)

            offset_x = div(2*Nx-i-3,Nx-2) - 1
            offset_y = div(2*Ny-j-3,Ny-2) - 1

            idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
            idata_index_y = (i-1)*Ny + j + offset_y

            # odata[global_index] = 0

            odata[global_index] += ( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) ) / (h^2)  # D2

            odata[global_index] += abs(offset_y) * 2 * alpha1 * ( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y]) / h^2    # BS_y

            odata[global_index] += abs(offset_x) * ( beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h) # BS_x'
            odata[global_index+Ny*offset_x] += abs(offset_x) * beta * 2 * (-1*idata[global_index]) / h^2 # BS_x'
            odata[global_index+2*Ny*offset_x] += abs(offset_x) * beta * (0.5*idata[global_index]) / h^2 # BS_x'
        end
    end

end

function matrix_free_A_v4(idata,odata,Nx,Ny,h)
    # one function that does D2 + boundary conditions + H_tilde
    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    odata .= 0
    offset_j = spzeros(Int,Ny)
    offset_j[2] = -1
    offset_j[3] = -2
    offset_j[Ny-1] = 1
    offset_j[Ny-2] = 2
    coef_j = spzeros(Ny)
    coef_j[1] = 2* 1.5
    coef_j[2] = - 2.0
    coef_j[3] = 0.5
    coef_j[Ny] = 2*1.5
    coef_j[Ny-1] = -2.0
    coef_j[Ny-2] = 0.5
    for i = 1:Nx
        for j = 1:Ny
            global_index = (i-1)*Ny+j

            # idata_index_x = (i-1)*Ny + j + test_mod(i,Nx) * Nx
            # idata_index_y = (i-1)*Ny + j + test_mod(j,Ny)

            offset_x = div(2*Nx-i-3,Nx-2) - 1
            offset_y = div(2*Ny-j-3,Ny-2) - 1

            idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
            idata_index_y = (i-1)*Ny + j + offset_y

            # odata[global_index] = 0

            # odata[global_index] = (( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) ) / (h^2)
            # + abs(offset_y) * 2 * (alpha1 * (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y]) / h^2
            # + abs(offset_x) * alpha4 * 2 * (idata[global_index] / h)  + coef_j[i] * beta * (idata[global_index+offset_j[i] * Ny]) / h^2)

            # Incorporating H_tilde
            odata[global_index] = (( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) )
            + abs(offset_y) * 2 * alpha1 *( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y])
            + abs(offset_x) * alpha4 * 2 * (idata[global_index] * h)  + coef_j[i] * beta * (idata[global_index+offset_j[i] * Ny]) ) / 2^(abs(offset_x) + abs(offset_y))

            #####
            # form boundary conditions separately

            # odata[global_index] += abs(offset_y) * 2 * (alpha1 * (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y]) / h^2 # BS_y
            # odata[global_index] += abs(offset_x) * alpha4 * 2 * (idata[global_index] / h) #
            # odata[global_index] += coef_j[i] * beta * (idata[global_index+offset_j[i] * Ny]) / h^2

            # # odata[global_index] += abs(offset_x) * ( beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h) # BS_x'
            # # odata[global_index+Ny*offset_x] += abs(offset_x) * beta * 2 * (-1*idata[global_index]) / h^2
            # # odata[global_index+2*Ny*offset_x] += abs(offset_x) * beta * (0.5*idata[global_index]) / h^2
        end
    end

end

function matrix_free_A_v5(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    # one function that does D2 + boundary conditions + H_tilde
    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    odata .= 0
    offset_j = spzeros(Int,Ny)
    offset_j[2] = -1
    offset_j[3] = -2
    offset_j[Ny-1] = 1
    offset_j[Ny-2] = 2
    coef_j = spzeros(Ny)
    coef_j[1] = 2* 1.5
    coef_j[2] = - 2.0
    coef_j[3] = 0.5
    coef_j[Ny] = 2*1.5
    coef_j[Ny-1] = -2.0
    coef_j[Ny-2] = 0.5
    for i = 1:Nx
        for j = 1:Ny
            global_index = (i-1)*Ny+j

            offset_x = div(2*Nx-i-3,Nx-2) - 1
            offset_y = div(2*Ny-j-3,Ny-2) - 1

            idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
            idata_index_y = (i-1)*Ny + j + offset_y

            odata[global_index] = (( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) )
            + abs(offset_y) * 2 * alpha1 * ( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y])
            + abs(offset_x) * alpha4 * 2 * (idata[global_index] * h)
             + coef_j[i] * beta * (idata[global_index+offset_j[i] * Ny])
             ) / 2^(abs(offset_x) + abs(offset_y))

        end
    end
end

function matrix_free_GPU(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,::Val{TILE_DIM1}, ::Val{TILE_DIM2}) where {TILE_DIM1, TILE_DIM2}
    tidx = threadIdx().x
    tidy = threadIdx().y

    # for global memory indexing
    i = (blockIdx().x - 1) * TILE_DIM1 + tidx
    j = (blockIdx().y - 1) * TILE_DIM2 + tidy

    global_index = (i-1)*Ny+j

    offset_x = div(2*Nx-i-3,Nx-2) - 1
    offset_y = div(2*Ny-j-3,Ny-2) - 1
    idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
    idata_index_y = (i-1)*Ny + j + offset_y

    # tidx = (blockIdx().x - 1) * TILE_DIM + threadIdx().x
    # tidy = (blockIdx().y - 1) * TILE_DIM + threadIdx().y

    # global_index = (tidx - 1) * Ny + tidy

    alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    # odata .= 0
    # offset_j = spzeros(Int,Ny)
    # offset_j = zeros(Int,Ny)
    # offset_j[2] = -1
    # offset_j[3] = -2
    # offset_j[Ny-1] = 1
    # offset_j[Ny-2] = 2
    # # coef_j = spzeros(Ny)
    # coef_j = zeros(Ny)
    # coef_j[1] = 2* 1.5
    # coef_j[2] = - 2.0
    # coef_j[3] = 0.5
    # coef_j[Ny] = 2*1.5
    # coef_j[Ny-1] = -2.0
    # coef_j[Ny-2] = 0.5

    if 1 <= i <= Nx && 1 <= j <= Ny

        odata[global_index] = (( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) )
        + abs(offset_y) * 2 * alpha1 * ( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y])
        + abs(offset_x) * alpha4 * 2 * (idata[global_index] * h)  + coef_j[i] * beta * (idata[global_index+offset_j[i] * Ny]) ) / 2^(abs(offset_x) + abs(offset_y))
    end
    nothing
end

function matrix_free_GPU_debug(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,::Val{TILE_DIM1}, ::Val{TILE_DIM2}) where {TILE_DIM1, TILE_DIM2}
    tidx = threadIdx().x
    tidy = threadIdx().y

    # for global memory indexing
    i = (blockIdx().x - 1) * TILE_DIM1 + tidx
    j = (blockIdx().y - 1) * TILE_DIM2 + tidy

    global_index = (i-1)*Ny+j

    offset_x = div(2*Nx-i-3,Nx-2) - 1
    offset_y = div(2*Ny-j-3,Ny-2) - 1
    idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
    idata_index_y = (i-1)*Ny + j + offset_y

    # tidx = (blockIdx().x - 1) * TILE_DIM + threadIdx().x
    # tidy = (blockIdx().y - 1) * TILE_DIM + threadIdx().y

    # global_index = (tidx - 1) * Ny + tidy

    alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    # odata .= 0
    # offset_j = spzeros(Int,Ny)
    # offset_j = zeros(Int,Ny)
    # offset_j[2] = -1
    # offset_j[3] = -2
    # offset_j[Ny-1] = 1
    # offset_j[Ny-2] = 2
    # # coef_j = spzeros(Ny)
    # coef_j = zeros(Ny)
    # coef_j[1] = 2* 1.5
    # coef_j[2] = - 2.0
    # coef_j[3] = 0.5
    # coef_j[Ny] = 2*1.5
    # coef_j[Ny-1] = -2.0
    # coef_j[Ny-2] = 0.5

    if 1 <= i <= Nx && 1 <= j <= Ny
    # if global_index <= Nx*Ny

        odata[global_index] = 0
        odata[global_index] = (( (idata[idata_index_x-Ny] - 2*idata[idata_index_x] + idata[idata_index_x + Ny])  + (idata[idata_index_y-1] - 2*idata[idata_index_y] + idata[idata_index_y + 1]) )
        + abs(offset_y) * 2 * alpha1 * ( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y])
        + abs(offset_x) * alpha4 * 2 * (idata[global_index] * h)
         + coef_j[i] * beta * idata[global_index+offset_j[i]* Ny]  # * (idata[global_index+offset_j[i] * Ny])   # this line is causing issue
         ) / 2^(abs(offset_x) + abs(offset_y))
    end
    nothing
end

function matrix_free_GPU_debug_v2(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,::Val{TILE_DIM1}, ::Val{TILE_DIM2}) where {TILE_DIM1, TILE_DIM2}
    tidx = threadIdx().x
    tidy = threadIdx().y

    # for global memory indexing
    i = (blockIdx().x - 1) * TILE_DIM1 + tidx
    j = (blockIdx().y - 1) * TILE_DIM2 + tidy

    global_index = (i-1)*Ny+j

    # offset_x = div(2*Nx-i-3,Nx-2) - 1
    # offset_y = div(2*Ny-j-3,Ny-2) - 1
    # idata_index_x = (i-1)*Ny + j + (offset_x) * Nx
    # idata_index_y = (i-1)*Ny + j + offset_y

    # alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    # odata .= 0
    # offset_j = spzeros(Int,Ny)
    # offset_j = zeros(Int,Ny)
    # offset_j[2] = -1
    # offset_j[3] = -2
    # offset_j[Ny-1] = 1
    # offset_j[Ny-2] = 2
    # # coef_j = spzeros(Ny)
    # coef_j = zeros(Ny)
    # coef_j[1] = 2* 1.5
    # coef_j[2] = - 2.0
    # coef_j[3] = 0.5
    # coef_j[Ny] = 2*1.5
    # coef_j[Ny-1] = -2.0
    # coef_j[Ny-2] = 0.5

    if 2 <= i <= Nx-1 && 2 <= j <= Ny-1
    # if global_index <= Nx*Ny

        odata[global_index] = 0
        odata[global_index] = (( (idata[global_index-Ny] - 2*idata[global_index] + idata[global_index + Ny])  + (idata[global_index-1] - 2*idata[global_index] + idata[global_index + 1]) )
        # + abs(offset_y) * 2 * (alpha1 * (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y])
        # + abs(offset_x) * alpha4 * 2 * (idata[global_index] * h)
        #  + coef_j[i] * beta * idata[global_index+offset_j[i]* Ny]  # * (idata[global_index+offset_j[i] * Ny])   # this line is causing issue
        ) #  ) / 2^(abs(offset_x) + abs(offset_y))
    end
    nothing
end

function Boundary_Conditions(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    odata .= 0
    for i = 1:Nx
        for j = 1:Ny
            global_index = (i-1)*Ny + j
            if j == Ny # N
                odata[global_index] += alpha1 * 2 * (1.5*idata[global_index] - 2*idata[global_index-1] + 0.5*idata[global_index-2]) / h^2
                # odata[global_index] = alpha1 * 2 * (1.5*idata[global_index] - 2*idata[global_index-1] + 0.5*idata[global_index-2]) / h^2
            end

            if j == 1 # S
                odata[global_index] += alpha2 * 2 *(1.5*idata[global_index] - 2*idata[global_index+1] + 0.5*idata[global_index+2]) / h^2
            end

            if i == Nx # E
                odata[global_index] += beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h
                odata[global_index-Ny] += beta * 2 * (-1*idata[global_index]) / h^2
                odata[global_index-2*Ny] += beta * (0.5*idata[global_index]) / h^2
            end

            if i == 1 # W
                odata[global_index] += beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha3 * 2* (idata[global_index]) / h
                odata[global_index+Ny] += beta * 2 * (-1*idata[global_index]) / h^2
                odata[global_index+2*Ny] += beta * (0.5*idata[global_index]) / h^2
            end
        end
    end
end

function Boundary_Conditions_v2(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    odata .= 0
    for i = 1:Nx
        for j = 1:Ny
            global_index = (i-1)*Ny + j
            offset_x = div(2*Nx-i-3,Nx-2) - 1
            offset_y = div(2*Ny-j-3,Ny-2) - 1
            # if j == Ny # N
            #     odata[global_index] += alpha1 * 2 * (1.5*idata[global_index] - 2*idata[global_index-1] + 0.5*idata[global_index-2]) / (2^(abs(offset_x)+abs(offset_y)))
            #     # odata[global_index] = alpha1 * 2 * (1.5*idata[global_index] - 2*idata[global_index-1] + 0.5*idata[global_index-2]) / h^2
            # end

            # if j == 1 # S
            #     odata[global_index] += alpha2 * 2 *(1.5*idata[global_index] - 2*idata[global_index+1] + 0.5*idata[global_index+2]) / h^2
            # end
            odata[global_index] += abs(offset_y) * 2 * alpha1 * ( (1.5*idata[global_index]) - 2*idata[global_index+offset_y] + 0.5*idata[global_index+2*offset_y]) / h^2# + abs(offset_x) * ( beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h)
                                # +  abs(offset_x) * ( beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h)

            # if i == Nx # E
            #     odata[global_index] += beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h
            #     odata[global_index-Ny] += beta * 2 * (-1*idata[global_index]) / h^2
            #     odata[global_index-2*Ny] += beta * (0.5*idata[global_index]) / h^2
            # end

            # if i == 1 # W
            #     odata[global_index] += beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha3 * 2* (idata[global_index]) / h
            #     odata[global_index+Ny] += beta * 2 * (-1*idata[global_index]) / h^2
            #     odata[global_index+2*Ny] += beta * (0.5*idata[global_index]) / h^2
            # end
            odata[global_index] += abs(offset_x) * ( beta * 2 * (1.5*idata[global_index] ) / h^2 + alpha4 * 2* (idata[global_index]) / h)
            odata[global_index+Ny*offset_x] += abs(offset_x) * beta * 2 * (-1*idata[global_index]) / h^2
            odata[global_index+2*Ny*offset_x] += abs(offset_x) * beta * (0.5*idata[global_index]) / h^2
        end
    end
end

function H_tilde_matrix_free(idata,odata,Nx,Ny,h)
    for i in 1:Nx
        for j in 1:Ny
            global_index = (i-1)*Ny+j
            odata[global_index] = idata[global_index] * h^2
            if i == 1
                odata[global_index] /= 2
            end
            if j == 1
                odata[global_index] /= 2
            end
            if i == Nx
                odata[global_index] /= 2
            end
            if j == Ny
                odata[global_index] /= 2
            end
        end
    end
end

function H_tilde_matrix_free_v2(idata,odata,Nx,Ny,h)
    for i in 1:Nx
        for j in 1:Ny
            global_index = (i-1)*Ny + j
            var = abs(div(2*Nx-i-3,Nx-2) - 1) + abs(div(2*Ny-j-3,Ny-2)-1)
            odata[global_index] = idata[global_index] * h^2 / (2^var)
            # odata[global_index] = var
        end
    end
end

function test_H_tilde(Nx,Ny)
    h = 1/(Nx-1)
    # idata = randn(Nx*Ny)
    idata = ones(Nx*Ny)
    odata1 = spzeros(Nx*Ny)
    odata2 = spzeros(Nx*Ny)
    H_tilde_matrix_free(idata,odata1,Nx,Ny,h)
    H_tilde_matrix_free_v2(idata,odata2,Nx,Ny,h)
    # @show odata1
    # @show odata2
    @show odata1 - odata2
end

function test_D2(Nx,Ny)
    h = 1/(Nx-1)
    # idata = randn(Nx*Ny)
    idata = ones(Nx*Ny)
    odata1 = spzeros(Nx*Ny)
    odata2 = spzeros(Nx*Ny)
    odata3 = spzeros(Nx*Ny)
    D2_test(idata,odata1,Nx,Ny,h)
    H_tilde_matrix_free(odata1,odata2,Nx,Ny,h)
    D2_test_v2(idata,odata3,Nx,Ny,h)
    # @show odata1
    # @show odata2
    @show odata2 - odata3
end

function test_Boundary_Conditions(Nx,Ny)
    h = 1/(Nx-1)
    alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    # idata = randn(Nx*Ny)
    idata = randn(Nx*Ny)
    odata1 = spzeros(Nx*Ny)
    odata2 = spzeros(Nx*Ny)
    Boundary_Conditions(idata,odata1,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    Boundary_Conditions_v2(idata,odata2,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    # @show odata1
    # @show odata2
    @show odata1 - odata2
end

function matrix_free_A(idata,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    odata1 = spzeros(Nx*Ny)
    odata2 = spzeros(Nx*Ny)
    D2_test(idata,odata1,Nx,Ny,h)
    Boundary_Conditions(idata,odata2,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    # odata3 = odata1 .+ odata2
    # return odata
    odata .= odata1 .+ odata2
    H_tilde_matrix_free(odata,odata,Nx,Ny,h)
    # return odata3
end

function test_matrix_free(Nx,Ny)
    h = 1/(Nx-1)
    alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    alpha1 = alpha2 = -13/h
    idata = randn(Nx*Ny)
    odata1 = spzeros(Nx*Ny)
    odata2 = spzeros(Nx*Ny)
    odata3 = spzeros(Nx*Ny)
    odata4 = spzeros(Nx*Ny)
    odata5 = spzeros(Nx*Ny)
    odata6 = spzeros(Nx*Ny)
    D2_test(idata,odata1,Nx,Ny,h)
    Boundary_Conditions(idata,odata2,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    # D2_test_v2(idata,odata3,Nx,Ny,h)
    Boundary_Conditions_v2(idata,odata4,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    matrix_free_A(idata,odata3,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    matrix_free_A_v2(idata,odata5,Nx,Ny,h) # without incorporating H_tilde
    matrix_free_A_v4(idata,odata6,Nx,Ny,h) # including H_tilde
    @show reshape(odata1 + odata2 - odata5,Nx,Ny)
    @show reshape(odata3 - odata6,Nx,Ny)
    @assert odata1 + odata2 ≈ odata5
    @assert odata3 ≈ odata6
end

function test_GPU_kernel(level)
    alpha1 = alpha2 = alpha3 = alpha4 = beta = 1
    Nx = Ny = 2^level + 1
    h = 1/(Nx-1)
    Random.seed!(0)
    u = randn(Nx*Ny)
    odata = randn(Nx*Ny)
    d_u = CuArray(u)
    d_odata = CuArray(odata)
    TILE_DIM_1 = 16
    TILE_DIM_2 = 16
    griddim = (div(Nx,TILE_DIM_1) + 1, div(Ny,TILE_DIM_2) + 1)
	blockdim = (TILE_DIM_1,TILE_DIM_2)

    offset_j = CuArray(zeros(Int,Ny))
    # offset_j = zeros(Int,Ny)
    offset_j[2] = -1
    offset_j[3] = -2
    offset_j[Ny-1] = 1
    offset_j[Ny-2] = 2
    # coef_j = spzeros(Ny)
    coef_j = CuArray(zeros(Float64,Ny))
    coef_j[1] = 2* 1.5
    coef_j[2] = - 2.0
    coef_j[3] = 0.5
    coef_j[Ny] = 2*1.5
    coef_j[Ny-1] = -2.0
    coef_j[Ny-2] = 0.5
    matrix_free_A_v5(u,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    @cuda threads=blockdim blocks=griddim matrix_free_GPU_debug(d_u,d_odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,Val(TILE_DIM_1), Val(TILE_DIM_2))
    # @show reshape(odata,Nx,Ny)
    # @show reshape(d_odata,Nx,Ny)

    @assert odata ≈ Array(d_odata)

    iter_times = 100

    t0 = time()
    for i in 1:iter_times
        matrix_free_A_v5(u,odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta)
    end
    t_cpu = time() - t0

    t1 = time()
    for i in 1:iter_times
        @cuda threads=blockdim blocks=griddim matrix_free_GPU_debug(d_u,d_odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,Val(TILE_DIM_1), Val(TILE_DIM_2))
    end
    synchronize()
    t_gpu = time() - t1

    t_D2 = time()
    for i in 1:iter_times
        @cuda threads=blockdim blocks=griddim matrix_free_GPU_debug_v2(d_u,d_odata,Nx,Ny,h,alpha1,alpha2,alpha3,alpha4,beta,offset_j,coef_j,Val(TILE_DIM_1), Val(TILE_DIM_2))
    end
    synchronize()
    t_D2 = time() - t_D2

    @show t_cpu
    @show t_gpu
    @show t_D2
end

function test_sparse_matrix(A,b)
    iter_times = 1000
    t3 = time()
    for i in 1:iter_times
        out = A*b
    end
    cpu_time = time() - t3
    @show cpu_time
end


function divergence(t,colat,long,vph,vth,aph,ath)

# computes the divergence of the velocity and acceleration
# on a lat,long grid at a series of snapshots.
#
# input
#
# t -  array of times
# colat  - array of colatitudes
# long  - array of longitudes
# vph,vth  - phi and theta components of velocity (km/yr)
# aph,ath  - phi and theta components of acceleration (km/yr^2)

#
# output
#
# vdiv   -  divergence of velocity
# adiv   -  divergence of acceleration

# allocate memory for divergence
vdiv = Array{Float64,3}(undef,length(t),length(colat),length(long));
adiv = Array{Float64,3}(undef,length(t),length(colat),length(long));

# radius of core-mantle boundary
rcmb = 3480.0;  #  in km

# flag grid points near the pole
nlat = length(colat);
nlong = length(long);
k = Int(floor(0.08*nlat));

# loop over times
for i = 1 : length(t)

#  loop over latitude and longitude

   for j = 1 : nlat

      if (j < k) || (j > nlat - k)
         rsinth = rcmb * sin(colat[k]);
      else
         rsinth = rcmb * sin(colat[j]);
      end

      for k = 1 : nlong

         # derivative of vth wrt colatitude
         if (j == 1)
           dlat = colat[j+1];
           dvth = (sin(colat[j+1])*vth[i,j+1,k] -
                     0.0) / (rsinth * dlat);
           dath = (sin(colat[j+1])*ath[i,j+1,k] -
                     0.0) / (rsinth * dlat);

         elseif (j == nlat )
            dlat = pi - colat[j-1];
            dvth = (0.0 - sin(colat[j-1])*vth[i,j-1,k] ) /
                    (rsinth * dlat);
            dath = (0.0 - sin(colat[j-1])*ath[i,j-1,k] ) /
                    (rsinth * dlat);

         else
            dlat = colat[j+1] - colat[j-1];
            dvth = (sin(colat[j+1])*vth[i,j+1,k] -
                  sin(colat[j-1])*vth[i,j-1,k])/(rsinth*dlat);
            dath = (sin(colat[j+1])*ath[i,j+1,k] -
                  sin(colat[j-1])*ath[i,j-1,k])/(rsinth*dlat);
         end

         # derivative of vphhi wrt longitude
        if (k == 1)
            dlong = long[k+1] - (long[nlong]-2*pi);
            dvph = (vph[i,j,k+1] - vph[i,j,nlong])/(rsinth * dlong);
            daph = (aph[i,j,k+1] - aph[i,j,nlong])/(rsinth * dlong);

         elseif (k == nlong)
            dlong = long[1] + 2*pi - long[k-1];
            dvph = (vph[i,j,1] - vph[i,j,k-1])/(rsinth * dlong);
            daph = (aph[i,j,1] - aph[i,j,k-1])/(rsinth * dlong);

         else
            dlong = long[k+1] - long[k-1];
            dvph = (vph[i,j,k+1] - vph[i,j,k-1])/(rsinth * dlong);
            daph = (aph[i,j,k+1] - aph[i,j,k-1])/(rsinth * dlong);
         end

         # divergence
         vdiv[i,j,k] = dvth + dvph;
         adiv[i,j,k] = dath + daph;

      end

   end

end

return vdiv, adiv

end


# --- Call FFTW3 directly --- #

using FFTW

function c_G_to_R( Ns::Array{Int,1}, fG::Array{ComplexF64,1} )
    out = zeros(ComplexF64,Ns[1]*Ns[2]*Ns[3])
    ccall( (:fftw_inv_fft3d, "../libs/fft3d.so"), Nothing,
           (Ptr{ComplexF64}, Ptr{ComplexF64}, Int64, Int64, Int64),
            fG, out, Ns[3],Ns[2],Ns[1] )
    return out
end

"""
Multicolumn version
"""
function c_G_to_R( Ns::Array{Int,1}, fG::Array{ComplexF64,2} )
    Npoints = size(fG)[1]
    Ncol = size(fG)[2]
    out = zeros(ComplexF64,Npoints,Ncol)
    for ic = 1:Ncol
       ccall( (:fftw_inv_fft3d, "../libs/fft3d.so"), Nothing,
              (Ptr{ComplexF64}, Ptr{ComplexF64}, Int64, Int64, Int64),
               fG[:,ic], out[:,ic], Ns[3],Ns[2],Ns[1] )
    end
    return out
end

# --- Call FFTW3 directly
function c_R_to_G( Ns::Array{Int,1}, fR::Array{ComplexF64,1} )
    out = zeros(ComplexF64,Ns[1]*Ns[2]*Ns[3])
    ccall( (:fftw_fw_fft3d, "../libs/fft3d.so"), Nothing,
           (Ptr{ComplexF64}, Ptr{ComplexF64}, Int64, Int64, Int64),
            fR, out, Ns[3],Ns[2],Ns[1] )
    return out
end

function c_R_to_G( Ns::Array{Int,1}, fR_::Array{Float64,1} )
    out = zeros(ComplexF64,Ns[1]*Ns[2]*Ns[3])
    fR = convert(Array{ComplexF64,1},fR_)
    ccall( (:fftw_fw_fft3d, "../libs/fft3d.so"), Nothing,
           (Ptr{ComplexF64}, Ptr{ComplexF64}, Int64, Int64, Int64),
            fR, out, Ns[3],Ns[2],Ns[1] )
    return out
end

"""
Multicolumn version
"""
function c_R_to_G( Ns::Array{Int,1}, fR::Array{ComplexF64,2} )
    Npoints = size(fR)[1]
    Ncol = size(fR)[2]
    out = zeros(ComplexF64,Npoints,Ncol)
    for ic = 1:Ncol
        ccall( (:fftw_fw_fft3d, "../libs/fft3d.so"), Nothing,
               (Ptr{ComplexF64}, Ptr{ComplexF64}, Int64, Int64, Int64),
                fR[:,ic], out[:,ic], Ns[3],Ns[2],Ns[1] )
    end
    return out
end

function G_to_R( Ns::Array{Int,1}, fG::Array{ComplexF64,1} )
    out = reshape( ifft( reshape(fG,Ns[1],Ns[2],Ns[3]) ),size(fG) )
end

function G_to_R( Ns::Array{Int,1}, fG::Array{ComplexF64,2} )
    Npoints = prod(Ns)
    out = zeros( ComplexF64, size(fG) ) # Is this safe?
    for ic = 1:size(fG,2)
        out[:,ic] = reshape( ifft( reshape(fG[:,ic],Ns[1],Ns[2],Ns[3]) ), Npoints )
    end
    return out
end

function R_to_G( Ns::Array{Int,1}, fR::Array{ComplexF64,1} )
    out = reshape( fft( reshape(fR,Ns[1],Ns[2],Ns[3]) ), size(fR) )
end

# In case we forget to convert the input, we convert it in this version
function R_to_G( Ns::Array{Int,1}, fR_::Array{Float64,1} )
    fR = convert(Array{ComplexF64,1},fR_)
    out = reshape( fft( reshape(fR,Ns[1],Ns[2],Ns[3]) ), size(fR) )
end

function R_to_G( Ns::Array{Int,1}, fR::Array{ComplexF64,2} )
    Npoints = prod(Ns)
    Ncol = size(fR,2)
    out = zeros( ComplexF64, size(fR) )
    for ic = 1:Ncol
        out[:,ic] = reshape( fft( reshape(fR[:,ic],Ns[1],Ns[2],Ns[3]) ), Npoints )
    end
    return out
end


#############################################################################
# gradient.jl
# calculates the gradients of functions
#############################################################################

export grad
export sign, curvature, monotonicity, evaluate

mutable struct GradientAtom <: AbstractExpr
  head::Symbol
  id_hash::UInt64
  children::Tuple{AbstractExpr}
  size::Tuple{Int, Int}
  variable::Variable
  func::Function

  function GradientAtom(x::AbstractExpr)
    var = get_cv(x)
    if x.size != (1, 1) || sign(x) == ComplexSign()
      error("A function has to be real-valued to calculate its gradient.")
    elseif length(var) > 1
      error("Only able to calculate gradients of functions with single variable.")
    elseif var[1].size[2] != 1
      error("Only able to calculate gradients of functions with variable's size being (n, 1).")
    end
    children = (x,)
    f(z) = evaluate(x, z)[1]
    return new(:grad, hash(children), children, var[1].size, var[1], f)
  end
end

grad(x::AbstractExpr) = GradientAtom(x)

function sign(x::GradientAtom)
  if monotonicity(x.children[1]) == (Nonincreasing(), )
    return Positive()
  elseif monotonicity(x.children[1]) == (Nondecreasing(), )
    return Negative()
  elseif monotonicity(x.children[1]) == (ConstMonotonicity(), )
    return Positive()
  else
    return NoSign()
  end
end

function curvature(x::GradientAtom)
  return NotDcp()
end

function monotonicity(x::GradientAtom)
  if curvature(x.children[1]) == ConvexVexity()
    return (Nondecreasing(), )
  elseif curvature(x.children[1]) == ConcaveVexity()
    return (Nonincreasing(), )
  elseif curvature(x.children[1]) == AffineVexity
    return (ConstMonotonicity(), )
  else
    return (NoMonotonicity(), )
  end
end

function evaluate(x::GradientAtom)
  var = copy(x.variable.value)
  val = gradient(x.func, var)
  x.variable.value = var
  return val
end


@safetestset "MINERVA" begin
    using GlobalMatchingModels, Statistics, Test, Random
    M = [1 1 0; 0 -1 0]
    p = [0, 1, 0]

    s = compute_similarity(M, p)
    @test s == [.5,-1]

    a = compute_activation(M, p)
    @test a == [.125,-1]

    x = compute_intensity(M, p)
    @test x == -0.875

    include("../Examples/MINERVA_Functions.jl")
    Random.seed!(1125435)
    # initialize model
    model = MINERVA(n_features=20, n_traces=60, L=.50)
    # simulate experiment
    sim_data = mapreduce(x->simulate(model, 4, 5), hcat, 1:20000)'
    means = mean(sim_data, dims=1)
    @test means ≈ [.0 0.153 0.31  0.46  0.61  0.77] atol = .01
end


# Plots an random example of a prior and the resulting information distribution

using SubmodularMaximization
using PyPlot

close("all")

steps = 100
grid_size = 10
num_observations = 3

grid = Grid(grid_size, grid_size)
sensor = RangingSensor()

target_state = random_state(grid)

# precomputation
histogram_filter = Filter(grid, target_state)

# Create some fake observations and update the filter
for ii = 1:num_observations
  robot_state = random_state(grid)
  range_observation = generate_observation(grid, sensor, robot_state,
                                           target_state)
  measurement_update!(histogram_filter, robot_state, get_states(grid), sensor,
                      grid, range_observation)
end

@show entropy(histogram_filter)

# Visualize the new prior
visualize_filter(histogram_filter, show_colorbar = false)
title("Prior")

# Information gain on the prior
information = Array(get_data(histogram_filter))
for ii = 1:length(information)
  state = SubmodularMaximization.index_to_state(grid, ii)

  information[ii] = finite_horizon_information(grid, histogram_filter, sensor,
                                               [state]).reward
end

# Visualize the information gain
figure()
limits = [0.5, grid_size + 0.5, 0.5, grid_size + 0.5]
image = imshow(information', cmap="viridis", vmin=0.0,
               extent=limits,
               interpolation="nearest", origin="lower")
#make_tight_colorbar(image)
title("Information gain")


@testset "KMM -- $optlib" for optlib in [JuliaLib, JuMPLib]
  for (i, (pair, rtol)) in enumerate([(pair₁, 2e-1), (pair₂, 4e-1)])
    d_nu, d_de = pair
    rng = MersenneTwister(42)
    x_nu, x_de = rand(rng, d_nu, 2_000), rand(rng, d_de, 1_000)

    # estimated density ratio
    D = [sqrt(DensityRatioEstimation.euclidsq(x, y)) for x in x_nu, y in x_de]
    σ, B, ϵ, λ = median(D), Inf, 0.001, 0.01
    kmm = KMM(σ=σ, B=B, ϵ=ϵ,  λ=λ)
    r̂ = densratio(x_nu, x_de, kmm; optlib=optlib)

    # simplex constraints
    @test abs(mean(r̂) - 1) ≤ 1e-2
    @test all(r̂ .≤ B)

    # compare against true ratio
    r = pdf.(d_nu, x_de) ./ pdf.(d_de, x_de)
    @test r ≈ r̂ rtol=rtol

    # type consistency
    @test eltype(r) == typeof(σ)

    # iterator and matrix version consistency for JuliaLib
    if optlib == JuliaLib
      r̂_mat = densratio(reshape(x_nu, 1, :), reshape(x_de, 1, :), kmm; optlib=optlib)
      @test r̂ ≈ r̂_mat
    end

    if visualtests
      gr(size=(800, 800))
      @test_reference "data/KMM-$optlib-$i.png" plot_d_nu(pair, x_de, r̂)
    end
  end
end


function step!(hs :: HardSpheres, pc :: PeriodicCollisions)

    n = hs.n
    ϵ = hs.ϵ

    dtmin, fant, i1, i2 = dt_min_position(pc)

    if fant > 0

       qa = [hs.q[i2] + offset[k] + dtmin*hs.v[i2] for k in 1:9]
       qb = hs.q[i1] + dtmin .* hs.v[i1]
       vr = argmin([norm(qx-qb) for qx in qa])

       hs.q[i2] = qa[vr]
       hs.q[i1] = qb

       J = (dot(hs.v[i2]-hs.v[i1],hs.q[i2]-hs.q[i1]))/ 2ϵ
       hs.v[i1] = hs.v[i1] + J * (hs.q[i2]-hs.q[i1]) / 2ϵ
       hs.v[i2] = hs.v[i2] - J * (hs.q[i2]-hs.q[i1]) / 2ϵ
       hs.q[i1] = mod.(hs.q[i1],1)
       hs.q[i2] = mod.(hs.q[i2],1)

    end

    for i in [1:min(i1,i2)-1;min(i1,i2)+1:max(i1,i2)-1;max(i1,i2)+1:n]
        qnew = hs.q[i] .+ dtmin .* hs.v[i]
        hs.q[i] = mod.(qnew, 1.)
    end

    pc.dt .-= dtmin
    reset!( pc, i1)
    reset!( pc, i2)

    compute_dt!(pc, i1, hs)
    compute_dt!(pc, i2, hs)

    return dtmin

end


__precompile__()
module ResponseModels
using FilmDataStructures
export eVtoOmega, eps_InAs_struct, eps_InAsntype, eps_InAsptype, hStep, cstRsp, sicRsp, siDsc, dptDsc, siRsp, prmMSi
# Collection of material permittivity models for use in heatLayers code.
# Conversion factors.
# Microns to electron volts.
const muEv = 1.239
# Radial frequency to electron volts.
const hBEv = 6.582e-16
# Temperature to electron volts.
const blzK = 8.6173e-5
# Electron volts to Joules.
const evJ = 1.602176565e-19
const eps_0 = 8.8542*^(10,-12)
const hbar = 1.05457*^(10,-34) #m2kg/s

@inline function cstRsp(cst::ComplexF64, enr::Float64)::ComplexF64

	return cst
end
precompile(cstRsp, (ComplexF64, Float64))
## Heaviside function for testing
@inline function hStep(arg::Float64)::Float64

	if arg >= 0.0

		return 1.0
	else
		return 0.0
	end
end
precompile(hStep, (Float64, ))
### Silicon carbide.
"""

	sicRsp(enr::Float64)::ComplexF64

Simplified permittivity response model for silicon carbide.
Approximately valid for energies between 0.1 and 10 electron volts.
"""
# Reproduction of the model reported in J. Appl. Phys. 106, 044306 (2009).
@inline function sicRsp(enr::Float64)::ComplexF64

	# Energy values (in electron volts) for the longitudinal and transverse optical phonons,
	# and decay rate.
	enrL = 0.12014342321784009
	enrT = 0.09831711568536007
	enrD = 0.0005901528146448003

	return 6.7 * (1.0 + /(enrL^2.0 - enrT^2.0, enrT^2.0 - enr^2.0 - im * enrD * enr))
end

precompile(sicRsp, (Float64, ))
### Silicon
# Sample usage
# Boron p-type doping
# brnDpt = dptDsc(0.044, 1.0e18)
# Empty doping
# nllDpt = dptDsc(0.0, 0.0)
# Calculate silicon model parameters
# siModE = prmMSi(tmpLst[1], nllDpt, nllDpt)
# siRspE(enr) = siRsp(enr, siModE)
"""
Data structure holding computed model parameters for the permittivity response of silicon for
given dopant properties at a fixed temperature. Call prmMSi(tmp::Float64, dnr::dptDsc,
acp::dptDsc)::siDsc for construction.

# Arguments
.tmp: temperature at which parameters where calculated.
.dnrFac: strength of donor impurity Lorentzian.
.acpFac: strength of acceptor impurity Lorentzian.
.dnrDcy: decay rate (gamma) for donor impurity Lorentzian.
.acpDcy: decay rate (gamma) for acceptor impurity Lorentzian.
"""
struct siDsc

	tmp::Float64
	dnrFac::Float64
	acpFac::Float64
	dnrDcy::Float64
	acpDcy::Float64
end
# Include model code.
include("siMod.jl")
"""

	siRsp(enr::Float64, tmp::Float64, siVal::siDsc)::ComplexF64

Permittivity response model for doped silicon. See source file for additional details.
"""
@inline function siRsp(enr::Float64, siVal::siDsc)::ComplexF64

	return siRspBck(enr, siVal.tmp) - /(siVal.dnrFac, enr^2 + im * enr * siVal.dnrDcy) - /(siVal.acpFac, enr^2 + im * enr * siVal.acpDcy)
end
precompile(siRsp, (Float64, siDsc))

using eps_InAs_ntype
using eps_InAs_ptype
struct InAsDsc
    N0::Float64
    T::Float64
	mstar::Float64
	gamma_ntype::Float64
	gamma_ptype::Float64
	E0_T_InAs_value::Float64
end

@inline function eVtoOmega(energy) #energy in eV
    return(energy*evJ/hbar)
end
precompile(eVtoOmega, (Float64,))

@inline function eps_InAs_struct(N0,T)
	E0_T_InAs_value = E0_T_InAs(N0,T)
	gamma_ntype = Gamma_ntype(N0,T)
	gamma_ptype = Gamma_ptype(N0,T)
    mstar= m_star(N0)
    return InAsDsc(N0,T,mstar,gamma_ntype,gamma_ptype,E0_T_InAs_value)
end
precompile(eps_InAs_struct,(Float64,Float64))

@inline function eps_InAsntype(E_photon::Float64,InAsDsc_struct::InAsDsc)::ComplexF64 #N_base in m^(-3)
     omega = eVtoOmega(E_photon)
     #o_angular = 5.38*10^14   #hbar*5.38*10^14/eV = 0.35eV
    if(E_photon>0.35)
        return epsIB(omega,InAsDsc_struct.N0,InAsDsc_struct.T,InAsDsc_struct.E0_T_InAs_value) + epsFCL(omega,InAsDsc_struct.mstar,InAsDsc_struct.gamma_ntype,InAsDsc_struct.N0)
    else
        return epsFCL(omega,InAsDsc_struct.mstar,InAsDsc_struct.gamma_ntype,InAsDsc_struct.N0)
    end
end
precompile(eps_InAs_ntype,(Float64,InAsDsc))

@inline function eps_InAsptype(E_photon::Float64,InAsDsc_struct::InAsDsc)::ComplexF64 #N_base in m^(-3)
	omega = eVtoOmega(E_photon)
	#o_angular = 5.38*10^14   #hbar*5.38*10^14/eV = 0.35eV
   if(E_photon>0.35)
	   return epsIB(omega,InAsDsc_struct.N0,InAsDsc_struct.T,InAsDsc_struct.E0_T_InAs_value) + epsFCL(omega,InAsDsc_struct.mstar,InAsDsc_struct.gamma_ptype,InAsDsc_struct.N0)
   else
	   return epsFCL_ptype(omega,InAsDsc_struct.mstar,InAsDsc_struct.gamma_ptype,InAsDsc_struct.N0)
   end
end
precompile(eps_InAs_ptype,(Float64,InAsDsc))
end


using DiffEqPhysics

function generate_bodies_in_cell_nodes(n::Int, m::Real, v_dev::Real, L::Real)

    rng = MersenneTwister(n);
    velocities = v_dev * randn(rng, Float64, (3, n))
    bodies = MassBody[]

    count = 1
    dL = L / (ceil(n^(1 / 3)))
    for x = dL/2:dL:L, y = dL/2:dL:L, z = dL/2:dL:L
        if count > n
            break
        end
        r = SVector(x, y, z)
        v = SVector{3}(velocities[:,count])
        body = MassBody(r, v, m)
        push!(bodies, body)
        count += 1
    end
    return bodies
end

function generate_bodies_in_line(n::Int, m::Real, v_dev::Real, L::Real)
    dL = L / (ceil(n^(1 / 3)))
    n_line = floor(Int, L / dL)
    rng = MersenneTwister(n);
    velocities = v_dev * randn(rng, Float64, (3, n_line))
    bodies = MassBody[]
    x = y = L / 2
    for i = 1:n_line
        r = SVector(x, y, i * dL)
        v = SVector{3}(velocities[:,i])
        body = MassBody(r, v, m)
        push!(bodies, body)
    end
    return bodies
end

function generate_random_directions(n::Int)
    theta = acos.(1 - 2 * rand(n));
    phi = 2 * pi * rand(n);
    directions = [@SVector [sin(theta[i]) .* cos(phi[i]), sin(theta[i]) .* sin(phi[i]), cos(theta[i])] for i = 1:n]
end

units = :real
units = :reduced

const T = 120.0 # °K
const kb = 1.38e-23 # J/K
const ϵ = T * kb
const σ = 3.4e-10 # m
const ρ = 1374 # kg/m^3
const m = 39.95 * 1.6747 * 1e-27 # kg
const N = 125#floor(Int, ρ * L^3 / m)
const L = (m*N/ρ)^(1/3)#10.229σ
const R = 2.25σ
const v_dev = sqrt(kb * T / m)
const τ = 1e-14 # σ/v
const t1 = 0τ
const t2 = 2000τ
#bodies = generate_bodies_randomly(N, m, v_dev, L)
bodies = generate_bodies_in_cell_nodes(N, m, v_dev, L)
#bodies = generate_bodies_in_line(N, m, v_dev, L)
jl_parameters = LennardJonesParameters(ϵ, σ, R)
pbc = CubicPeriodicBoundaryConditions(L)
thermostat = AndersenThermostat(0.02, T, kb)
lj_system = PotentialNBodySystem(bodies, Dict(:lennard_jones => jl_parameters));
simulation = NBodySimulation(lj_system, (t1, t2), pbc, thermostat);
#result = run_simulation(simulation, Tsit5())
result = @time run_simulation(simulation, VelocityVerlet(), dt=τ)

#=
using Plots
import GR
(rs, grf) = rdf(result)
(ts, dr2) = msd(result)
plot(rs/σ, grf, xlim=[0, 0.4999L/σ], label=["Radial distribution function"],ylabel="g(r)", xlabel="r/sigma")

using JLD
time_now = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
Nactual = length(bodies)
timesteps = round(length(result.solution.t))
#save("D:/water $Nactual molecules $timesteps steps.jld", "rs", rs, "grf", grf, "ts", ts, "dr2", dr2)
save("D:/liquid argon $Nactual molecules $timesteps steps $time_now.jld", "rs", rs, "grf", grf, "ts", ts, "dr2", dr2, "e_tot", e_tot, "e_kin", e_kin, "e_pot", e_pot)
=#
#=
using Plots
import GR
time_now = Dates.format(now(), "yyyy_mm_dd_HH_MM_SS")
Nactual = length(bodies)
timesteps = round(length(result.solution.t))
@time animate(result, "D:/$Nactual liquid argon particles with $timesteps timesteps $time_now.gif")

#plot(simulation)
=#


let
    epc   = Epoch(2018, 3, 20, 16, 15, 0) # Test on Vernal equinox
    oe    = [R_EARTH + 500e3, 0, 90.0, 0, 0, 0]
    x     = sOSCtoCART(oe, use_degrees = true)
    R     = rECItoECEF(epc)

    a_grav_point = accel_point_mass(x[1:3])

    a_grav_00 = accel_gravity(x, R, 0, 0)

    tol=1e-9
    @test isapprox(a_grav_00[1], a_grav_point[1], atol=tol)
    @test isapprox(a_grav_00[2], a_grav_point[2], atol=tol)
    @test isapprox(a_grav_00[3], a_grav_point[3], atol=tol)

    a_grav_20 = accel_gravity(x, R, 20, 20)
    @test norm(a_grav_20) > norm(a_grav_00)
end

let
    # Define Initial State
    epc   = Epoch(2018, 3, 20, 16, 15, 0) # Test on Vernal equinox
    oe    = [R_EARTH + 500e3, 0, 90.0, 0, 0, 0]
    x     = sOSCtoCART(oe, use_degrees = true)
    r_sun = sun_position(epc)

    # Call function
    a_sun = accel_thirdbody_sun(epc, x)

    @test a_sun[1] < 1e-5
    @test a_sun[2] < 1e-5
    @test a_sun[3] < 1e-5

    a_sun2 = accel_thirdbody_sun(x, r_sun)

    @test a_sun[1] == a_sun2[1]
    @test a_sun[2] == a_sun2[2]
    @test a_sun[3] == a_sun2[3]
end

let
    # Define Initial State
    epc    = Epoch(2018, 3, 20, 16, 15, 0) # Test on Vernal equinox
    oe     = [R_EARTH + 500e3, 0, 90.0, 0, 0, 0]
    x      = sOSCtoCART(oe, use_degrees = true)
    r_moon = moon_position(epc)

    # Call function
    a_moon = accel_thirdbody_moon(epc, x)

    @test a_moon[1] < 1e-5
    @test a_moon[2] < 1e-5
    @test a_moon[3] < 1e-5

    a_moon2 = accel_thirdbody_moon(x, r_moon)

    @test a_moon[1] == a_moon2[1]
    @test a_moon[2] == a_moon2[2]
    @test a_moon[3] == a_moon2[3]
end

let
    # Define Initial State
    epc   = Epoch(2018, 3, 20, 16, 15, 0) # Test on Vernal equinox
    oe    = [R_EARTH + 500e3, 0, 90.0, 0, 0, 0]
    x     = sOSCtoCART(oe, use_degrees = true)
    T     = rECItoECEF(epc)

    # Call function
    rho    = density_harris_priester(epc, x)
    a_drag = accel_drag(x, rho, 1, 1, 2.3, T)

    @test a_drag[1] < 1e-5
    @test a_drag[2] < 1e-5
    @test a_drag[3] < 1e-5
end

let
    # Define Initial State
    epc   = Epoch(2018, 3, 20, 16, 15, 0) # Test on Vernal equinox
    oe    = [R_EARTH + 500e3, 0, 0, 0, 0, 0]
    x     = sOSCtoCART(oe, use_degrees = true)
    r_sun = sun_position(epc)

    # Call function
    nu = eclipse_cylindrical(x, r_sun)
    @test nu == 1.0
    @test nu == eclipse_cylindrical(epc, x)

    oe    = [R_EARTH + 500e3, 0, 0, 180.0, 0, 0]
    x     = sOSCtoCART(oe, use_degrees = true)
    nu = eclipse_cylindrical(x, r_sun)
    @test nu == 0.0
    @test nu == eclipse_cylindrical(epc, x)
end

let
    # Define Initial State
    epc   = Epoch(2018, 3, 20, 16, 15, 0) # Test on Vernal equinox
    oe    = [R_EARTH + 500e3, 0, 0, 0, 0, 0]
    x     = sOSCtoCART(oe, use_degrees = true)
    r_sun = sun_position(epc)

    # Call function
    nu = eclipse_conical(x, r_sun)
    @test nu == 0.0
    @test nu == eclipse_conical(epc, x)

    oe    = [R_EARTH + 500e3, 0, 0, 180.0, 0, 0]
    x     = sOSCtoCART(oe, use_degrees = true)
    nu = eclipse_conical(x, r_sun)
    @test nu == 1.0
    @test nu == eclipse_conical(epc, x)
end

let
    # Define Initial State
    epc   = Epoch(2018, 1, 1, 0, 0, 0)
    oe    = [R_EARTH + 500e3, 0, 90.0, 0, 0, 0]
    x     = sOSCtoCART(oe, use_degrees = true)
    r_sun = sun_position(epc)

    # Call function
    a_srp = accel_srp(x, r_sun, 1, 1)

    @test a_srp[1] < 1e-5
    @test a_srp[2] < 1e-5
    @test a_srp[3] < 1e-5
end

let
    # Define Initial State
    oe = [R_EARTH + 500e3, 0, 90.0, 0, 0, 0]
    x  = sOSCtoCART(oe, use_degrees=true)

    # Call function
    a_rel = accel_relativity(x)

    @test a_rel[1] < 1e-7
    @test a_rel[2] < 1e-7
    @test a_rel[3] < 1e-7
end


using JuMP, Gurobi, Distributions, StatsFuns, StatsBase

include("portfolio_var_initialization_template.jl")

const gurobi_env = Gurobi.Env()

const numVariables = 1000
const numAssets = numVariables

const returnLBD = 1.2 #lower bound on the return

#*============ MODEL PARAMETERS ================
const mean = Array{Float64}(numAssets)
const stdev = Array{Float64}(numAssets)
for i = 1:numAssets
	mean[i] = 1.05 + 0.3*(numAssets-i)/(numAssets-1)
	stdev[i] = (0.05 + 0.6*(numAssets-i)/(numAssets-1))/3
end

const sigma = zeros(Float64, numAssets, numAssets)
for i = 1:numAssets
	sigma[i,i] = stdev[i]
end
#*===============================================

# compute probability that constraints aren't satisfied
function computeRiskLevel(y_opt::Array{Float64})

	riskLevel::Float64 = 1.0 - normcdf((sum(mean[i]*y_opt[i] for i = 1:numAssets) - returnLBD)/norm(sigma*y_opt))

	return riskLevel
end

# evaluate the constraint values for given decision vector x
# and realization of the random variables xi
function evaluateConstraint(x::Array{Float64},xi::Array{Float64})

	return evaluateConstraint(x,xi,1.0)
end

# evaluate the constraint values for given decision vector x,
# realization of the random variables xi, and constraint scalings
function evaluateConstraint(x::Array{Float64},xi::Array{Float64},scalings::Float64)

	con::Float64 = (returnLBD - xi'*x)/scalings

	return con
end

# evaluate constraint gradients
function evaluateConstraintGradients(x::Array{Float64},xi::Array{Float64})

	return evaluateConstraintGradients(x,xi,1.0)
end

# evaluate constraint gradients
function evaluateConstraintGradients(x::Array{Float64},xi::Array{Float64},scalings::Float64)

	grad_con = zeros(Float64,numVariables)
	for j = 1:numVariables
		grad_con[j] = -xi[j]
	end
	grad_con /= scalings

	return grad_con
end

# generate random samples of the stock returns
function generateRandomSamples(numSamples::Int64)

	xi = Array{Float64}(numAssets,numSamples)
	for i = 1:numAssets
		xi[i,:] = rand(Normal(mean[i],stdev[i]), numSamples)
	end

	return xi
end

# project on to the set {x[1:n] >= 0, sum(x[1:n]) = 1}
function projectOntoSimplex(y_ref::Array{Float64})

	y = deepcopy(y_ref)

	v = Array{Float64}(1)
	v[1] = y[1]
	w = Array{Float64}(0)
	rho::Float64 = y[1] - 1.0

	for j = 2:numVariables
		if(y[j] > rho)
			rho = rho + (y[j] - rho)/(size(v,1) + 1)
			if(rho > y[j] - 1.0)
				push!(v,y[j])
			else
				append!(w,v)
				v = [y[j]]
				rho = y[j] - 1.0
			end
		end
	end

	if(size(w,1) > 0)
		for j = 1:size(w,1)
			if(w[j] > rho)
				push!(v,w[j])
				rho = rho + (w[j] - rho)/size(v,1)
			end
		end
	end

	numDeleted::Int64 = 1

	while(numDeleted > 0)
		initialCard::Int64 = size(v,1)
		numDeleted = 0
		for j = 1:initialCard
			if(v[j-numDeleted] <= rho)
				rho = rho + (rho - v[j-numDeleted])/(size(v,1)-1)
				deleteat!(v,j-numDeleted)
				numDeleted += 1
			end
		end
	end

	y_proj = zeros(Float64,numVariables)
	for j = 1:numVariables
		y_proj[j] = max(0.0,y[j] - rho)
	end

	return y_proj
end

function solveInnerProblem(y_ref::Array{Float64,1},lambda::Float64,objectiveBound::Float64)

	y = deepcopy(y_ref)

	numerators = zeros(Float64,numVariables)
	denominators = zeros(Float64,numVariables)
	for j = 1:numVariables
		numerators[j] = y[j]/(1.0+2*abs2(stdev[j])*lambda)
		denominators[j] = 1.0/(1.0+2*abs2(stdev[j])*lambda)
	end

	v = Array{Float64}(1)
	den_v = Array{Int64}(1)
	v[1] = y[1]
	den_v[1] = 1
	w = Array{Float64}(0)
	den_w = Array{Float64}(0)
	nr::Float64 = numerators[1]
	dr::Float64 = denominators[1]
	rho::Float64 = (nr-1.0)/dr

	for j = 2:numVariables
		if(y[j] > rho)
			nr += numerators[j]
			dr += denominators[j]
			rho = (nr-1.0)/dr
			if(rho > (numerators[j]-1.0)/denominators[j])
				push!(v,y[j])
				push!(den_v,j)
			else
				append!(w,v)
				append!(den_w,den_v)
				v = [y[j]]
				den_v = [j]
				nr = numerators[j]
				dr = denominators[j]
				rho = (nr-1.0)/dr
			end
		end
	end

	if(size(w,1) > 0)
		for j = 1:size(w,1)
			if(w[j] > rho)
				push!(v,w[j])
				push!(den_v,den_w[j])
				nr += numerators[den_w[j]]
				dr += denominators[den_w[j]]
				rho = (nr-1.0)/dr
			end
		end
	end

	numDeleted::Int64 = 1

	while(numDeleted > 0)
		initialCard::Int64 = size(v,1)
		numDeleted = 0
		for j = 1:initialCard
			if(v[j-numDeleted] <= rho)
				nr -= numerators[den_v[j-numDeleted]]
				dr -= denominators[den_v[j-numDeleted]]
				rho = (nr-1.0)/dr
				deleteat!(v,j-numDeleted)
				deleteat!(den_v,j-numDeleted)
				numDeleted += 1
			end
		end
	end

	y_proj = zeros(numVariables)
	for j = 1:numVariables
		y_proj[j] = max(0.0, (y_ref[j] - rho)/(1.0 + 2*abs2(stdev[j])*lambda))
	end

	residual::Float64 = norm(sigma*y_proj) - sqrt(objectiveBound)

	return rho, residual
end

function projectOntoFeasibleSet(y_ref::Array{Float64},objectiveBound::Float64,proj_params::Array{Float64})

	y_test = projectOntoSimplex(y_ref)
	lambda::Float64 = 0.0
	residual_test::Float64 = norm(sigma*y_test) - sqrt(objectiveBound)
	if(residual_test <= 0.0)
		return y_test, proj_params
	end

	lambda_tol::Float64 = 1E-08
	lambda_low::Float64 = proj_params[1]
	lambda_up::Float64 = proj_params[2]
	residual_tol::Float64 = 1E-09
	~, residual_up::Float64 = solveInnerProblem(y_ref,lambda_up,objectiveBound)
	~, residual_low::Float64 = solveInnerProblem(y_ref,lambda_low,objectiveBound)
	while(residual_up > 0.0 || residual_low < 0.0)
		if(residual_up > 0.0)
			lambda_low = lambda_up
			lambda_up *= 2.0
			~, residual_up = solveInnerProblem(y_ref,lambda_up,objectiveBound)
		else
			lambda_up = lambda_low
			lambda_low /= 2.0
			~, residual_low = solveInnerProblem(y_ref,lambda_low,objectiveBound)
		end
	end

	while(lambda_up - lambda_low > lambda_tol*lambda_up)
		lambda = (lambda_low + lambda_up)/2.0
		~, residual::Float64 = solveInnerProblem(y_ref,lambda,objectiveBound)
		if(abs(residual) <= residual_tol)
			break
		elseif(residual > 0.0)
			lambda_low = lambda
		else
			lambda_up = lambda
		end
	end

	mu::Float64, ~ = solveInnerProblem(y_ref,lambda,objectiveBound)
	y_proj = zeros(Float64,numVariables)
	for j = 1:numVariables
		y_proj[j] = max(0.0, (y_ref[j]-mu)/(1.0+2*abs2(stdev[j])*lambda))
	end

	err_net::Float64 = max(maximum(-y_proj),abs(sum(y_proj)-1),norm(sigma*y_proj)-sqrt(objectiveBound))
	if(err_net > 1E-06)
		println("WARNING! PROJECTION STEP MAY NOT BE ACCURATE!!!")
	end

	if(proj_params[3] > 0.5)
		if(lambda > lambda_tol)
			lambda_factor::Float64 = sqrt(2.0)
			proj_params[1] = lambda/lambda_factor
			proj_params[2] = lambda*lambda_factor
		else
			proj_params[1] = 0.0
			proj_params[2] = 1.0
		end
	end

	return y_proj, proj_params
end

# get the stochastic gradient of the approximation
function getStochasticGradient(x::Array{Float64},numSamples::Int64,mu::Float64,tau::Float64,scalings::Float64)

	xi = generateRandomSamples(numSamples)

	stochgrad = zeros(numVariables)

	for samp = 1:numSamples
		con::Float64 = evaluateConstraint(x,xi[:,samp],scalings)

		grad_con = evaluateConstraintGradients(x,xi[:,samp],scalings)

		baseFactor1::Float64 = tau/((sqrt(mu)*exp(0.5*tau*con) + sqrt(1.0/mu)*exp(-0.5*tau*con))^2)

		stochgrad += baseFactor1*grad_con/numSamples
	end

	return stochgrad
end

# get the stochastic gradient of the approximation
function getStochasticGradientForFixedSample(x::Array{Float64},xi::SubArray{Float64,2,Array{Float64,2},Tuple{Base.Slice{Base.OneTo{Int64}},UnitRange{Int64}},true},mu::Float64,tau::Float64,scalings::Float64)

	numSamples::Int64 = size(xi,2)

	stochgrad = zeros(numVariables)

	for samp = 1:numSamples
		con::Float64 = evaluateConstraint(x,xi[:,samp],scalings)

		grad_con = evaluateConstraintGradients(x,xi[:,samp],scalings)

		baseFactor1::Float64 = tau/((sqrt(mu)*exp(0.5*tau*con) + sqrt(1.0/mu)*exp(-0.5*tau*con))^2)

		stochgrad += baseFactor1*grad_con/numSamples
	end

	return stochgrad

end

# estimate constraint scalings
function estimateScalings(x::Array{Float64},numSamples::Int64)

	scaling_tol::Float64 = 1E-06

	xi = generateRandomSamples(numSamples)

	constraint_values = zeros(numSamples)
	for samp = 1:numSamples
		constraint_values[samp] = evaluateConstraint(x,xi[:,samp])
	end

	scalings::Float64 = median(abs.(constraint_values))

	scalings = max(scalings,scaling_tol)

	return scalings
end

# estimate Lipschitz constant of gradient of approximation
function estimateCompositeLipschitzConstant(x_ref::Array{Float64},numSamplesForLipVarEst::Int64,numGradientSamples::Int64,batchSize::Int64,mu::Float64,tau::Float64,scalings::Float64,objectiveBound::Float64,proj_params::Array{Float64})

	sample_radius::Float64 = norm(x_ref)/10.0

	xi = generateRandomSamples(batchSize)

	L_avg::Float64 = 0.0
	for batch = 1:batchSize
		xi_start::Int64 = batch
		xi_end::Int64 = batch

		L_max::Float64 = 0.0
		for samp = 1:numSamplesForLipVarEst
			x_1 = pickPointOnSphere(x_ref,sample_radius)
			x_1,~ = projectOntoFeasibleSet(x_1,objectiveBound,proj_params)

			x_2 = pickPointOnSphere(x_ref,sample_radius)
			x_2,~ = projectOntoFeasibleSet(x_2,objectiveBound,proj_params)

			@views grad_1 = getStochasticGradientForFixedSample(x_1,xi[:,xi_start:xi_end],mu,tau,scalings)
			@views grad_2 = getStochasticGradientForFixedSample(x_2,xi[:,xi_start:xi_end],mu,tau,scalings)

			L_est::Float64 = norm(grad_2 - grad_1)/norm(x_2 - x_1)

			if(L_est > L_max)
				L_max = L_est
			end
		end

		L_avg += L_max/batchSize
	end

	return L_avg
end

# estimate step length factor
function estimateStepLength(x_ref::Array{Float64},numGradientSamples::Int64,
							numSamplesForLipVarEst::Int64,batchSize::Int64,maxNumIterations::Int64,
							minNumReplicates::Int64,mu::Float64,tau::Float64,
							scalings::Float64,objectiveBound::Float64)

	proj_params = zeros(Float64,3)
	proj_params[2] = 1.0

	sample_radius::Float64 = norm(x_ref)/10.0

	rho_est::Float64 = estimateCompositeLipschitzConstant(x_ref,numSamplesForLipVarEst,numGradientSamples,batchSize,mu,tau,scalings,objectiveBound,proj_params)

	variance_est::Float64 = 0.0
	for samp = 1:numSamplesForLipVarEst
		x_1 = pickPointOnSphere(x_ref,sample_radius)
		x_1,~ = projectOntoFeasibleSet(x_1,objectiveBound,proj_params)

		xi = generateRandomSamples(numGradientSamples*batchSize)

		variance_1::Float64 = 0.0
		for batch = 1:batchSize
			xi_start::Int64 = (batch-1)*numGradientSamples + 1
			xi_end::Int64 = batch*numGradientSamples

			@views grad_1 = getStochasticGradientForFixedSample(x_1,xi[:,xi_start:xi_end],mu,tau,scalings)

			variance_1 += norm(grad_1)^2/batchSize
		end

		if(variance_1 > variance_est)
			variance_est = variance_1
		end
	end

	gamma::Float64 = 1/sqrt(rho_est*variance_est)

	stepLength::Float64 = gamma/sqrt((maxNumIterations+1.0)*minNumReplicates)

	return stepLength
end

# given a reference point, pick a random point within a given radius
function pickPointOnSphere(x_orig::Array{Float64},radius::Float64)

	if(radius < 1E-09)
		error("Sampling radius in pickPointOnSphere is almost zero! Set manually.")
	end

	n::Int64 = size(x_orig,1)
	unifRand = rand(1)
	factor1::Float64 = radius*(unifRand[1])^(1.0/n)
	normRand = rand(Normal(0,1), n)
	factor2::Float64 = norm(normRand)

	x_samp = x_orig + factor1*normRand/factor2

	return x_samp
end


function metaplex_gillespie(mp::Metaplex{T1,T2}, Xi0::BitVector, Xμ0::Vector{Int}, β, D::Vector; nmax=length(Xi0), tmax=100.0, sampling_method=:tree) where {T1,T2}
    N = nv(mp.g) # number of individuals
    g = mp.g
    M = nv(mp.h) # number of populations
    h = mp.h
    #od = map(k-> k>0 ? 1/float(k) : 0, outdegree(h))
    #od = outdegree(h)
    od = outdegree(h) .> 0
    Xi = copy(Xi0) # epidemic state of each individual
    Xμ = copy(Xμ0) # location of each individual
    popcounts = zeros(Int,2,M) # number of of individuals susceptible and infected for each component
    pops_s = [Set{Int}() for i in 1:M] # Population to individuals without traversing the whole list of individuals
    pops_i = [Set{Int}() for i in 1:M]
    for k in 1:N
        if !Xi[k]
            push!(pops_s[Xμ[k]],k)
            #popcounts[1,Xμ[k]] += 1
        else
            push!(pops_i[Xμ[k]],k)
            #popcounts[2,Xμ[k]] += 1
        end
    end
    popcounts[1,:] = map(length, pops_s)
    popcounts[2,:] = map(length, pops_i)
    a = zeros(N+2*M) # aggregate reactions for each individual becoming infected and individuals migrating from each population
    for k in 1:N
        if !Xi[k]
            l = length(filter(i->Xi[i] && Xμ[i]==Xμ[k], inneighbors(g,k)))
            a[k] = β*l # probability of k recieving infection is β time its number of infected incoming neighbors in the same population
            #nreactions += l
        end
    end
    a[N+1:N+M] .= D[1].*popcounts[1,:].*od
    a[N+M+1:N+2*M] .= D[2].*popcounts[2,:].*od
    if sampling_method == :tree
        a = CategoricalTree(a)
    end
    a0 = sum(a)
    ts = zeros(nmax)
    pcs = zeros(Int,nmax,2,M)
    pcs[1,:,:] .= popcounts
    t = 0.0
    n = 1
    while n < nmax && t < tmax && a0 > 0
        τ = log(1/rand())/a0
        k = rand_categorical(a, a0)
        if k in 1:N # node k becomes infected
            Xi[k] = true
            delete!(pops_s[Xμ[k]], k)
            push!(pops_i[Xμ[k]], k)
            popcounts[1,Xμ[k]] = length(pops_s[Xμ[k]])
            popcounts[2,Xμ[k]] = length(pops_i[Xμ[k]])
            a[N+Xμ[k]] = D[1]*popcounts[1,Xμ[k]]*od[Xμ[k]]
            a[N+M+Xμ[k]] += D[2]*popcounts[2,Xμ[k]]*od[Xμ[k]]
            a[k] = 0.0
            for i in Iterators.filter(j->!Xi[j] && Xμ[j]==Xμ[k], outneighbors(g,k))
                a[i] += β
            end
        elseif k in N+1:N+M # a random susceptible node from population k-N moves to another population
            μ = k-N
            ν = rand(outneighbors(h,μ))
            i = rand(pops_s[μ])
            Xμ[i] = ν
            delete!(pops_s[μ], i)
            push!(pops_s[ν], i)
            popcounts[1,μ] = length(pops_s[μ])
            popcounts[1,ν] = length(pops_s[ν])
            a[N+μ] = D[1]*popcounts[1,μ]*od[μ]
            a[N+ν] = D[1]*popcounts[1,ν]*od[ν]
            l = length(filter(j->Xi[j] && Xμ[j]==ν, inneighbors(g,i)))
            a[i] = β*l
        else # a random infected node from k-N-M moves to another population
            μ = k-N-M
            ν = rand(outneighbors(h,μ))
            i = rand(pops_i[μ])
            Xμ[i] = ν
            delete!(pops_i[μ], i)
            push!(pops_i[ν], i)
            for j in Iterators.filter(l->!Xi[l], outneighbors(g,i))
                if Xμ[j] == μ
                    a[j] -= β
                elseif Xμ[j] == ν
                    a[j] += β
                end
            end
            popcounts[2,μ] = length(pops_i[μ])
            popcounts[2,ν] = length(pops_i[ν])
            a[N+M+μ] = D[2]*popcounts[2,μ]*od[μ]
            a[N+M+ν] = D[2]*popcounts[2,ν]*od[ν]
        end
        a0 = sum(a)
        t += τ
        ts[n+1] = t
        pcs[n+1,:,:] .= popcounts
        n += 1
    end
    return ts[1:n], pcs[1:n,:,:]
end


## as anonymous function
@model StateSpaceModel(y) = begin
    N = length(y)
    x = tzeros(Real, N)
    σy = 1.0
    σx = 1.0

    x[1] ~ Normal(0, 1)
    y[1] ~ Normal(x[1], σy)

    # Iterators.drop(1:N, 1)
    for i ∈ 2:N
        x[i] ~ Normal(x[i-1], σx)
        y[i] ~ Normal(x[i], σy)
    end
end

## as function
    # @model function StateSpaceModel(y)
    #     N = length(y)
    #     x = tzeros(Real, N)
    #     σy = 1.0
    #     σx = 1.0

    #     x[1] ~ Normal(0, 1)
    #     y[1] ~ Normal(x[1], σy)
    #     for i ∈ 2:N
    #         x[i] ~ Normal(x[i-1], σx)
    #         y[i] ~ Normal(x[i-1], σy)
    #     end
    # end

@model RandomWalkAngularAccelerationModel(y, s; σω = 1.0) = begin
    N = length(y)
    α = tzeros(Real, N)
    ω = tzeros(Real, N)
    c = zeros(Real, N)
    θ = zeros(Real, N)
    Δt = s.Δt
    κ1, κ2 = cupula_dynamics(Δt)

    σα = σω
    # σω = 0.1
    σy = s.sensor.σ

    α[1] ~ Normal(0, σα)
    ω[1] ~ Normal(0, σω)
    c[1] ~ Determin(κ1 * 0 + κ2 * ω[1])
    y[1] ~ Normal(ω[1] - c[1], σy)

    for i ∈ Iterators.drop(1:N, 1)
        α[i] ~ Normal(α[i-1], σα)
        ω[i] ~ Normal(ω[i-1] + Δt * α[i], σω)
        c[i] ~ Determin(κ1 * c[i-1] + κ2 * ω[i])
        y[i] ~ Normal(ω[i] - c[i], σy)
    end
end

@model RandomWalkAngularVelocityModel(y, s; σω = 1.0) = begin
    N = length(y)
    ω = tzeros(Real, N)
    p = tzeros(Real, N)
    c = zeros(Real, N)
    Δt = s.Δt
    κ1, κ2 = cupula_dynamics(Δt)

    # σω = 1
    σy = s.sensor.σ
    ω[1] ~ Normal(0, σω)
    p[1] ~ Determin(ω[1] > 0 ? 1 : 0)
    c[1] ~ Determin(κ1 * 0 + κ2 * ω[1])
    y[1] ~ Normal(ω[1] - c[1], σy)

    for i ∈ Iterators.drop(1:N, 1)
        ω[i] ~ Normal(ω[i-1], σω)
        p[i] ~ Determin(ω[i] > 0 ? 1 : 0)
        c[i] ~ Determin(κ1 * c[i-1] + κ2 * ω[i])
        y[i] ~ Normal(ω[i] - c[i], σy)
    end
end

@model HeadTurnModel(y, s, m, ::Type{T} = Float64; σωᵤ = 0.01, σωₑ = 0.1) where {T} = begin
    N = length(y)
    stochastic = s.stochastic
    # intialize vectors
    α1 = TArray{T}(undef, N)
    ω1 = TArray{T}(undef, N)
    p1 = TArray{T}(undef, N)
    ω2 = TArray{T}(undef, N)
    p2 = TArray{T}(undef, N)
    ω = TArray{T}(undef, N)
    p = TArray{T}(undef, N)
    c = TArray{T}(undef, N)

    # parameters
    D1 = m.D == "left" ? -1 : 1
    A1 = stochastic == 0 ? m.A : rand(Normal(m.A, 0.5))
    # A1 = rand(Normal(m.A, 0.5))
    # A1 ~ Normal(m.A, 0.1)

    Δt = s.Δt
    onsetᵤ = hasproperty(m, :onset) ? m.onset : s.onsetᵤ
    duration1 = m.duration
    f1 = 1/(duration1) # frequency: single sinusiodal head turn
    timesteps = range(0, stop = N, step = Δt)
    κ1, κ2 = cupula_dynamics(Δt)

    # priors
    σy = 1.1 * s.sensor.σ # jitter the measurement noise
    # σωᵤ ~ Truncated(Normal(0.5, 0.1))
    # σωₑ ~ Truncated(Normal(0.5, 0.1))

    # initial state at t_0
    α1[1] = 0.0
    ω1[1] ~ Normal(0, σωᵤ)
    # p1[1] ~ Determin(ω1[1] > 0 ? 1 : 0)
    ω2[1] ~ Normal(0, σωₑ)
    # p2[1] ~ Determin(ω2[1] > 0 ? 1 : 0)
    ω[1] ~ Determin(ω1[1] + ω2[1])
    # p[1] ~ Determin(ω[1] > 0 ? 1 : 0)
    c[1] ~ Determin(κ1 * 0 + κ2 * ω[1])
    y[1] ~ Normal(ω[1] - c[1], σy)

    # for i ∈ Iterators.drop(1:N, 1)
    for i ∈ 2:N
        t = round(timesteps[i]; digits = 3)
        α1[i] = (t > onsetᵤ) & (t < onsetᵤ + duration1) ? acceleration(D1, A1, f1, t, onsetᵤ, noisy = stochastic) : 0.0

        ω1[i] ~ Normal(ω1[i-1] + Δt * α1[i], σωᵤ)
        # p1[i] ~ Determin(ω1[i] > 0 ? 1 : 0)
        ω2[i] ~ Normal(ω2[i-1], σωₑ)
        # p2[i] ~ Determin(ω2[i] > 0 ? 1 : 0)
        ω[i] ~ Determin(ω1[i] + ω2[i])
        # p[i] ~ Determin(ω[i] > 0 ? 1 : 0)

        c[i] ~ Determin(κ1 * c[i-1] + κ2 * ω[i])
        y[i] ~ Normal(ω[i] - c[i], σy)
    end
end

@model HeadTurnOnlyModel(y, s, m; σωᵤ = 0.01) = begin
    N = length(y)
    stochastic = s.stochastic
    # intialize vectors
    α1 = tzeros(Real, N)
    ω = tzeros(Real, N)
    p = tzeros(Real, N)
    c = zeros(Real, N)

    # parameters
    D1 = m.D == "left" ? -1 : 1
    # A1 = stochastic == 0 ? m.A : rand(Normal(m.A, 0.5))
    A1 = rand(Normal(m.A, 0.5))

    Δt = s.Δt
    onsetᵤ = hasproperty(m, :onset) ? m.onset : s.onsetᵤ
    duration1 = m.duration
    f1 = 1/(duration1) # frequency: single sinusiodal head turn
    timesteps = range(0, stop = N, step = Δt)
    κ1, κ2 = cupula_dynamics(Δt)

    # priors
    σy = s.sensor.σ

    # initial state at t_0
    ω[1] ~ Normal(0, σωᵤ)
    p[1] ~ Determin(ω[1] > 0 ? 1 : 0)
    c[1] ~ Determin(κ1 * 0 + κ2 * ω[1])
    y[1] ~ Normal(ω[1] - c[1], σy)

    for i ∈ Iterators.drop(1:N, 1)
        t = round(timesteps[i]; digits = 3)
        α1[i] = (t > onsetᵤ) & (t < onsetᵤ + duration1) ? acceleration(D1, A1, f1, t, onsetᵤ, noisy = stochastic) : 0.0

        ω[i] ~ Normal(ω[i-1] + Δt * α1[i], σωᵤ)
        p[i] ~ Determin(ω[i] > 0 ? 1 : 0)
        c[i] ~ Determin(κ1 * c[i-1] + κ2 * ω[i])
        y[i] ~ Normal(ω[i] - c[i], σy)
    end
end

@model HeadTurnOnlyParamModel(y, s, m; σωᵤ = 0.01, σₐ = 1.0) = begin
    N = length(y)

    # intialize vectors
    α1 = tzeros(Real, N)
    ω = tzeros(Real, N)
    p = tzeros(Real, N)
    c = zeros(Real, N)

    # parameters
    D1 = m.D == "left" ? -1 : 1
    A1 ~ Normal(m.A, σₐ)

    Δt = s.Δt
    onsetᵤ = hasproperty(m, :onset) ? m.onset : s.onsetᵤ
    duration1 = m.duration
    f1 = 1/(duration1) # frequency: single sinusiodal head turn
    timesteps = range(0, stop = N, step = Δt)
    κ1, κ2 = cupula_dynamics(Δt)

    # priors
    σy = s.sensor.σ

    # initial state at t_0
    ω[1] ~ Normal(0, σωᵤ)
    # p[1] ~ Determin(ω[1] > 0 ? 1 : 0)
    c[1] ~ Determin(κ1 * 0 + κ2 * ω[1])
    y[1] ~ Normal(ω[1] - c[1], σy)

    for i ∈ Iterators.drop(1:N, 1)
        t = round(timesteps[i]; digits = 3)
        α1[i] = (t > onsetᵤ) & (t < onsetᵤ + duration1) ? acceleration(D1, A1, f1, t, onsetᵤ, noisy = false) : 0.0

        ω[i] ~ Normal(ω[i-1] + Δt * α1[i], σωᵤ)
        # p[i] ~ Determin(ω[i] > 0 ? 1 : 0)
        c[i] ~ Determin(κ1 * c[i-1] + κ2 * ω[i])
        y[i] ~ Normal(ω[i] - c[i], σy)
    end
end

@model HeadTurnOnlySimulationModel(y, s, m; σωᵤ = 0.01, σₐ = 1.0) = begin
    N = length(y)

    # intialize vectors
    α1 = tzeros(Real, N)
    ω = tzeros(Real, N)
    p = tzeros(Real, N)
    c = zeros(Real, N)

    # parameters
    D = m.D == "left" ? -1 : 1
    A ~ Normal(m.A, σₐ)

    Δt = s.Δt
    on = hasproperty(m, :onset) ? m.onset : s.onsetᵤ
    # on = shape * scale
    onsetᵤ ~ Gamma(20 * on, 0.05)
    duration1 = m.duration
    f = 1/(duration1) # frequency: single sinusiodal head turn
    timesteps = range(0, stop = N, step = Δt)
    κ1, κ2 = cupula_dynamics(Δt)

    # priors
    σy = s.sensor.σ

    # initial state at t_0
    ω[1] ~ Normal(0, σωᵤ)
    # p[1] ~ Determin(ω[1] > 0 ? 1 : 0)
    c[1] ~ Determin(κ1 * 0 + κ2 * ω[1])
    y[1] ~ Normal(ω[1] - c[1], σy)

    for i ∈ Iterators.drop(1:N, 1)
        t = round(timesteps[i]; digits = 3)
        α1[i] = (t > onsetᵤ) & (t < onsetᵤ + duration1) ? acceleration(D, A, f, t, onsetᵤ, noisy = false) : 0.0

        ω[i] ~ Normal(ω[i-1] + Δt * α1[i], σωᵤ)
        # p[i] ~ Determin(ω[i] > 0 ? 1 : 0)
        c[i] ~ Determin(κ1 * c[i-1] + κ2 * ω[i])
        y[i] ~ Normal(ω[i] - c[i], σy)
    end
end

# TODO: when moving, set σωᵤ much lower, when not moving higher
@model HeadTurnComplexModel(y, s, m; σωᵤ = 0.01, σωₑ = 0.1) = begin
    N = length(y)
    stochastic = s.stochastic
    # intialize vectors
    α1 = tzeros(Real, N)
    ω1 = tzeros(Real, N)
    p1 = tzeros(Real, N)
    ω2 = tzeros(Real, N)
    p2 = tzeros(Real, N)
    ω = tzeros(Real, N)
    move = tzeros(Real, N)
    p = tzeros(Real, N)
    c = zeros(Real, N)

    # parameters
    D1 = m.D == "left" ? -1 : 1
    A1 = stochastic == 0 ? m.A : rand(Normal(m1.A, 0.5))

    Δt = s.Δt
    onsetᵤ = hasproperty(m, :onset) ? m.onset : s.onsetᵤ
    duration1 = m.duration
    f1 = 1/(duration1) # frequency: single sinusiodal head turn
    timesteps = range(0, stop = N, step = Δt)
    κ1, κ2 = cupula_dynamics(Δt)

    # priors
    σy = s.sensor.σ

    # initial state at t_0
    ω1[1] ~ Normal(0, σωᵤ)
    p1[1] ~ Determin(ω1[1] > 0 ? 1 : 0)
    ω2[1] ~ Normal(0, σωₑ)
    p2[1] ~ Determin(ω2[1] > 0 ? 1 : 0)
    ω[1] ~ Determin(ω1[1] + ω2[1])
    move[1] ~ Determin(0)
    p[1] ~ Determin(ω[1] > 0 ? 1 : 0)
    c[1] ~ Determin(κ1 * 0 + κ2 * ω[1])
    y[1] ~ Normal(ω[1] - c[1], σy)

    for i ∈ Iterators.drop(1:N, 1)
        t = round(timesteps[i]; digits = 3)
        α1[i] = (t > onsetᵤ) & (t < onsetᵤ + duration1) ? acceleration(D1, A1, f1, t, onsetᵤ, noisy = stochastic) : 0.0

        ω1[i] ~ Normal(ω1[i-1] + Δt * α1[i], σωᵤ)
        move[i] ~ Determin((t > onsetᵤ) & (t < onsetᵤ + duration1) ? 1 : 0)
        p1[i] ~ Determin(ω1[i] > 0 ? 1 : 0)
        ω2[i] ~ Normal(ω2[i-1], σωₑ)
        p2[i] ~ Determin(ω2[i] > 0 ? 1 : 0)
        ω[i] ~ Determin(ω1[i] + ω2[i])
        p[i] ~ Determin(ω[i] > 0 ? 1 : 0)

        c[i] ~ Determin(κ1 * c[i-1] + κ2 * ω[i])
        y[i] ~ Normal(ω[i] - c[i], σy)
    end
end

@model CoupledModel(y, s, m; σωᵤ = 0.01, σωₑ = 0.1) = begin
    N = length(y)
    stochastic = s.stochastic
    # intialize vectors
    α1 = tzeros(Real, N)
    ω1 = tzeros(Real, N)
    # p1 = tzeros(Real, N)
    ω2 = tzeros(Real, N)
    p2 = tzeros(Real, N)
    ω = tzeros(Real, N)
    # p = tzeros(Real, N)
    c = zeros(Real, N)

    # parameters
    D1 = m.D == "left" ? -1 : 1
    A1 = stochastic == 0 ? m.A : rand(Normal(m1.A, 0.5))

    Δt = s.Δt
    onsetᵤ = hasproperty(m, :onset) ? m.onset : s.onsetᵤ
    duration1 = m.duration
    f1 = 1/(duration1) # frequency: single sinusiodal head turn
    timesteps = range(0, stop = N, step = Δt)
    κ1, κ2 = cupula_dynamics(Δt)

    # priors
    σy = s.sensor.σ

    # initial state at t_0
    ω1[1] ~ Normal(0, σωᵤ)
    # p1[1] ~ Determin(ω1[1] > 0 ? 1 : 0)
    ω2[1] ~ Normal(0, σωₑ)
    # p2[1] ~ Determin(ω2[1] > 0 ? 1 : 0)
    ω[1] ~ Determin(ω1[1] - ω2[1])
    # p[1] ~ Determin(ω[1] > 0 ? 1 : 0)
    c[1] ~ Determin(κ1 * 0 + κ2 * ω[1])
    y[1] ~ Normal(ω1[1] - c[1], σy)
    y[1] ~ Normal(ω2[1] - c[1], σy)

    for i ∈ Iterators.drop(1:N, 1)
        t = round(timesteps[i]; digits = 3)
        α1[i] = (t > onsetᵤ) & (t < onsetᵤ + duration1) ? acceleration(D1, A1, f1, t, onsetᵤ, noisy = stochastic) : 0.0

        ω1[i] ~ Normal(ω1[i-1] + Δt * α1[i], σωᵤ)
        # p1[i] ~ Determin(ω1[i] > 0 ? 1 : 0)
        ω2[i] ~ Normal(ω2[i-1], σωₑ)
        # p2[i] ~ Determin(ω2[i] > 0 ? 1 : 0)
        ω[i] ~ Determin(ω1[i] - ω2[i])
        # p[i] ~ Determin(ω[i] > 0 ? 1 : 0)

        c[i] ~ Determin(κ1 * c[i-1] + κ2 * ω[i])
        y[i] ~ Normal(ω1[i] - c[i], σy)
        y[i] ~ Normal(ω2[i] - c[i], σy)
    end
end


# Auxilary functions
# Arc functions ────────────────────────────────────────────────────────────────────────────────
"""
    cₐ(a::Arc, x=a.x)

Returns arc cost for arc `a` for arc flow `x`
"""
function cₐ(a::Arc, x=a.x)
    tₒ= a.tₒ
    α = a.α
    β = a.β
    V = a.V
    ϕ = a.ϕ

    t = tₒ * (1 + α * (x/V) ^ β)
    t′= ϕ == 0 || β == 0 ? 0.0 : tₒ * α * β * (x ^ (β - 1))/(V ^ β)

    c = t + ϕ * x * t′
    return c
end

"""
    cₐ′(a::Arc, x=a.x)

Returns first derivative of arc cost wrt arc flow for arc `a` at arc flow `x`
"""
function cₐ′(a::Arc, x=a.x)
    tₒ= a.tₒ
    α = a.α
    β = a.β
    V = a.V
    ϕ = a.ϕ

    t′= β == 0 ? 0.0 : tₒ * α * β * (x ^ (β - 1))/(V ^ β)
    t″ = ϕ == 0 || β == 0 || β == 1 ?  0.0 : tₒ * α * β * (β - 1) * (x ^ (β - 2))/(V ^ β)

    c′ = t′ + ϕ * x * t″
    return c′
end

# Djikstra functions ────────────────────────────────────────────────────────────────────────────────
# NOTE: Node values must match node location (index) in the set of nodes N

"""
    djk(G::Graph, o::Origin)

Djikstra's label setting algorithm

Returns predecessor arc label (index in vector `A`) `Lᵖ` for least cost paths from origin `o` on graph `G`
"""
function djk(G::Graph, o::Origin)
    N = G.N
    A = G.A
    K = G.K
    r = o.n

    Lᵖ = Vector{Int64}(undef, length(N))     # Predecessor label
    Lᶜ = [Inf for _ in N]                    # Cost label
    Lᵖ[r.v] = -1
    Lᶜ[r.v] = 0.0

    X = copy(N)                              # Set of open nodes

    index = findfirst(x -> (x == r), X)::Int64
    t = X[index]
    deleteat!(X, index)
    while !isempty(X)
        i = t.v
        for j in t.H
            k = K[i,j]
            a = A[k]
            l = Lᶜ[i] + a.c
            if l ≤ Lᶜ[j] Lᵖ[j], Lᶜ[j] = k, l
            end
        end
        index = argmin([Lᶜ[n.v] for n in X])
        t = X[index]
        deleteat!(X, index)
    end

    return Lᵖ
end

"""
    tree(G::Graph, o::Origin)

Returns tree for graph `G` rooted at origin `o`
"""
function tree(G::Graph, Lᵖ)
    N = G.N
    A = G.A

    T = Vector{Node}[[] for _ in N]

    for k in Lᵖ
        if k == -1 continue end
        a = A[k]
        t, h = a.t, a.h
        i = t.v
        push!(T[i], h)
    end

    return T
end

"""
    path(G::Graph, Lᵖ, r::Node, s::Node)

Returns path between origin node `r` and destination node `s` on graph `G` using predecessor label `Lᵖ` (for origin node `r`)
"""
function path(G::Graph, Lᵖ, r::Node, s::Node)
    A = G.A

    p = Arc[]

    h = s
    j = h.v
    while j ≠ r.v
        k = Lᵖ[j]
        a = A[k]
        push!(p, a)
        j = a.t.v
    end

    reverse!(p)
    return p
end
# ────────────────────────────────────────────────────────────────────────────────


# ------------------------------------------------------------
# 1D change of basis matrix

function Freq2Wave(samples::StridedVector, wavename::AbstractString, J::Int, B::Float64=NaN; uniform::Bool=true, args...)
	vm = van_moment(wavename)
	( Nint = 2^J ) >= 2*vm-1 || throw(AssertionError("Scale it not large enough for this wavelet"))

	M = length(samples)
	if Nint >= M
		warn("The scale is high compared to the number of samples")
	end

	# Weights for non-uniform samples
	uni_samples = isuniform(samples)
	if uni_samples
		sample_dist = samples[2] - samples[1]
		W = Nullable{ Vector{Complex{Float64}} }()
	else
		isnan(B) && error("Samples are not uniform; supply bandwidth")
		Nint <= 2*B || warn("The scale is high compared to the bandwidth")
		W = sqrt(weights(samples, B))
		W = Nullable(complex( W ))
	end

	# Fourier transform of the internal scaling function
	internal = FourScalingFunc( samples, wavename, J; args... )

	# The number of internal wavelets in the reconstruction
	if hasboundary(wavename)
		Nint -= 2*vm
		Nint <= 0 && error("Too few wavelets: Boundary functions overlap")
	end

    if uni_samples && isinteger(1/sample_dist) && uniform
	# For suitable uniform samples, the ordinary FFT can be used instead of NFFT
		p = FFTPlan(samples, J, Nint)
	else
		# NFFTPlans: Frequencies must be in the torus [-1/2, 1/2)
		xi = samples*2.0^(-J)
		frac!(xi)
		p = NFFTPlan(xi, Nint)
	end

	# Wavelets w/o boundary
	if !hasboundary(wavename)
		return Freq2NoBoundaryWave1D(internal, W, J, wavename, p)
	else
		left = FourScalingFunc( samples, wavename, 'L', J; args... )
		right = FourScalingFunc( samples, wavename, 'R', J; args... )

		return Freq2BoundaryWave1D(internal, W, J, wavename, p, left, right)
	end
end

@doc """
	isuniform(Freq2Wave)

Is the change of basis matrix based on uniform samples.
"""->
function isuniform(T::Freq2Wave)
	isnull(T.weights)
end

@doc """
	hasboundary(T::Freq2Wave) -> Bool

Does the wavelet in the change of basis matrix `T` have boundary correction.
"""->
function hasboundary(T::Freq2Wave)
	isdefined(T, :left)
end

# ------------------------------------------------------------
# Basic operations for 1D Freq2Wave

function Base.collect(T::Freq2NoBoundaryWave1D)
	M, N = size(T)

	F = Array{Complex{Float64}}(M, N)
	offset = div(N, 2) + 1
	for n in 1:N
		for m in 1:M
			@inbounds F[m,n] = T.internal[m]*cis( -twoπ*(n-offset)*T.NFFT.x[m] )
		end
	end

	if !isuniform(T)
		scale!(get(T.weights), F)
	end

	return F
end

function Base.collect(T::Freq2BoundaryWave1D)
	M, N = size(T)
	F = Array{Complex{Float64}}(M, N)

	# Left boundary
	p = van_moment(T)
	F[:,1:p] = T.left

	# Internal function
	offset = div(N, 2) + 1
	for n in (p+1):(N-p)
		for m in 1:M
			@inbounds F[m,n] = T.internal[m]*cis( -twoπ*(n-offset)*T.NFFT.x[m] )
		end
	end

	# Right boundary
	F[:,N-p+1:N] = T.right

	if !isuniform(T)
		scale!(get(T.weights), F)
	end

	return F
end

# ------------------------------------------------------------
# 2D change of basis matrix

@doc """
	dim(Freq2wave)

Return the dimension of the `T`.
"""->
dim(::Freq2Wave1D) = 1
dim(::Freq2Wave2D) = 2

@doc """
	wscale(Freq2Wave)

Return the scale of the wavelet coefficients.
"""->
function wscale(T::Freq2Wave)
	T.J
end

@doc """
	wsize(T::Freq2Wave)

The number of reconstructed wavelet coefficients (in each dimension).

- When `D` == 1, the output is (Int,)
- When `D` == 2, the output is (Int,Int)
"""->
function wsize(T::Freq2BoundaryWave2D)
	N = 2^wscale(T)
	return (N, N)
end

wsize(T::Freq2BoundaryWave1D) = (2^wscale(T),)
wsize(T::Freq2NoBoundaryWave2D) = T.NFFT.N
wsize(T::Freq2NoBoundaryWave1D) = T.NFFT.N

function UnifFourScalingFunc(samples::StridedMatrix{Float64}, wavename::AbstractString, J::Integer; args...)
	# Test if samples are on a grid and in the correct order
	usamplesx = unique(view(samples,:,1))
	usamplesy = unique(view(samples,:,2))
	Mx = length(usamplesx)
	My = length(usamplesy)
	grid_samples = grid( (Mx, My), usamplesx[2]-usamplesx[1] )

	samples == grid_samples ||
	throw(AssertionError("Samples are not on a grid or not ordered properly"))

	# Internal scaling function
	internalx = FourScalingFunc( usamplesx, wavename, J )
	internalx = kron(internalx, ones(My))
	internaly = FourScalingFunc( usamplesy, wavename, J )
	internaly = repmat(internaly, Mx)
	internal = Dict( :x => internalx, :y => internaly )

	# Boundary scaling functions
	if !hasboundary(wavename)
		return internal
	else
		vm = van_moment(wavename)

		leftx = FourScalingFunc( usamplesx, wavename, 'L', J; args... )
		lefty = FourScalingFunc( usamplesy, wavename, 'L', J; args... )
		left = Dict( :x => repmat(leftx, My), :y => kron(lefty, ones(Mx)) )

		rightx = FourScalingFunc( usamplesx, wavename, 'R', J; args... )
		righty = FourScalingFunc( usamplesy, wavename, 'R', J; args... )
		right = Dict( :x => repmat(rightx, My), :y => kron(righty, ones(Mx)) )

		return internal, left, right
	end
end

function NotUnifFourScalingFunc(samples::StridedMatrix{Float64}, wavename::AbstractString, J::Integer, B::Float64; args...)
	# Fourier transform of the internal scaling function
	int = FourScalingFunc( samples, wavename, J )
	internal = Dict( :x => int[:,1], :y => int[:,2] )

	# Boundary scaling functions
	if !hasboundary(wavename)
		return internal
	else
		samplesx = view(samples, :, 1)
		samplesy = view(samples, :, 2)

		left = Dict{ Symbol, Matrix{Complex{Float64}} }()
		left[:x] = FourScalingFunc( samplesx, wavename, 'L', J; args... )
		left[:y] = FourScalingFunc( samplesy, wavename, 'L', J; args... )

		right = Dict{ Symbol, Matrix{Complex{Float64}} }()
		right[:x] = FourScalingFunc( samplesx, wavename, 'R', J; args... )
		right[:y] = FourScalingFunc( samplesy, wavename, 'R', J; args... )

		return internal, left, right
	end
end

function Freq2Wave(samples::StridedMatrix{Float64}, wavename::AbstractString, J::Int, B::Float64=NaN; uniform::Bool=true, args...)
	vm = van_moment(wavename)
	( Nint = 2^J ) >= 2*vm-1 || throw(AssertionError("Scale it not large enough for this wavelet"))
	M = size(samples, 1)
	size(samples,2) == 2 || throw(DimensionMismatch("Samples must have two columns"))

	if Nint >= M
		warn("The scale is high compared to the number of samples")
	end

	uni_samples = isuniform(samples)
	if uni_samples
		W = Nullable{ Vector{Complex{Float64}} }()

		sample_dist = samples[2,2] - samples[1,2]
		scaling_funcs = UnifFourScalingFunc(samples, wavename, J; args...)
	else
		# Weights for non-uniform samples
		isnan(B) && error("Samples are not uniform; supply bandwidth")
		Nint <= 2*B || warn("The scale is high compared to the bandwidth")
		W = sqrt(weights(samples, B))
		W = Nullable(complex( W ))

		scaling_funcs = NotUnifFourScalingFunc(samples, wavename, J, B; args...)
	end

	if hasboundary(wavename)
		Nint -= 2*vm
		Nint <= 0 && error("Too few wavelets: Boundary functions overlap")

		internal = scaling_funcs[1]
		left = scaling_funcs[2]
		right = scaling_funcs[3]
	else
		internal = scaling_funcs
	end

    xi = samples'
	if uni_samples && isinteger(1/sample_dist) && uniform
        # For suitable uniform samples, the ordinary FFT can be used instead of NFFT
        p = FFTPlan(xi, J, (Nint,Nint))
	else
        # NFFTPlans: Frequencies must be in the torus [-1/2, 1/2)^2
        scale!(xi, 2.0^(-J))
        frac!(xi)
        p = NFFTPlan(xi, (Nint,Nint))
    end

	# Wavelets w/o boundary
	if hasboundary(wavename)
		return Freq2BoundaryWave2D(internal, W, J, wavename, p, left, right)
	else
		return Freq2NoBoundaryWave2D(internal, W, J, wavename, p)
	end
end

function Base.A_mul_B!(y::DenseVector{Complex{Float64}}, T::Freq2NoBoundaryWave1D, x::DenseVector{Complex{Float64}})
	size(T,1) == length(y) || throw(DimensionMismatch())
	size(T,2) == length(x) || throw(DimensionMismatch())

	myfft!(y, T.NFFT, x)
	had!(y, T.internal)

	isuniform(T) || had!(y, get(T.weights))

	return y
end

function Base.A_mul_B!(y::DenseVector{Complex{Float64}}, T::Freq2NoBoundaryWave2D, X::DenseMatrix{Complex{Float64}})
	(M = size(T,1)) == length(y) || throw(DimensionMismatch())
	wsize(T) == size(X) || throw(DimensionMismatch())

	myfft!(y, T.NFFT, X)
	for m in 1:M
		@inbounds y[m] *= T.internal[:x][m] * T.internal[:y][m]
	end

	isuniform(T) || had!(y, get(T.weights))

	return y
end

function Base.A_mul_B!(y::DenseVector{Complex{Float64}}, T::Freq2BoundaryWave1D, x::DenseVector{Complex{Float64}})
	(M = size(T,1)) == length(y) || throw(DimensionMismatch())
	size(T,2) == length(x) || throw(DimensionMismatch())

	xleft, xint, xright = split(x, van_moment(T))

	# Internal scaling function
	myfft!(y, T.NFFT, xint)
	had!(y, T.internal)

	# Contribution from the boundaries
	BLAS.gemv!('N', ComplexOne, T.left, xleft, ComplexOne, y)
	BLAS.gemv!('N', ComplexOne, T.right, xright, ComplexOne, y)

	isuniform(T) || had!(y, get(T.weights))

	return y
end

function Base.A_mul_B!(y::DenseVector{Complex{Float64}}, T::Freq2BoundaryWave2D, X::DenseMatrix{Complex{Float64}})
	(M = size(T,1)) == length(y) || throw(DimensionMismatch())
	(N = wsize(T)) == size(X) || throw(DimensionMismatch())

	# Internal scaling functions
	vm = van_moment(T)
	S = split(X, vm)
	myfft!(y, T.NFFT, S.II)
	for m in 1:M
		@inbounds y[m] *= T.internal[:x][m] * T.internal[:y][m]
	end

	onesp = ones(Complex{Float64}, vm)

	# ------------------------------------------------------------
	# Corners

	# LL
	A_mul_B!( T.tmpMulVec, T.left[:x], S.LL )
	had!( T.tmpMulVec, T.left[:y] )
	# y += sum(T.tmpMulVec,2) :
	BLAS.gemv!('N', ComplexOne, T.tmpMulVec, onesp, ComplexOne, y)

	# RL
	A_mul_B!( T.tmpMulVec, T.right[:x], S.RL )
	had!( T.tmpMulVec, T.left[:y] )
	BLAS.gemv!('N', ComplexOne, T.tmpMulVec, onesp, ComplexOne, y)

	# RR
	A_mul_B!( T.tmpMulVec, T.right[:x], S.RR )
	had!( T.tmpMulVec, T.right[:y] )
	BLAS.gemv!('N', ComplexOne, T.tmpMulVec, onesp, ComplexOne, y)

	# LR
	A_mul_B!( T.tmpMulVec, T.left[:x], S.LR )
	had!( T.tmpMulVec, T.right[:y] )
	BLAS.gemv!('N', ComplexOne, T.tmpMulVec, onesp, ComplexOne, y)

	# ------------------------------------------------------------
	# Sides

	# IL
	NFFT.nfft!(T.NFFTx, S.IL, T.tmpMulVec)
	scale!(T.internal[:x], T.tmpMulVec)
	had!(T.tmpMulVec, T.left[:y])
	BLAS.gemv!('N', ComplexOne, T.tmpMulVec, onesp, ComplexOne, y)

	# IR
	NFFT.nfft!(T.NFFTx, S.IR, T.tmpMulVec)
	scale!(T.internal[:x], T.tmpMulVec)
	had!(T.tmpMulVec, T.right[:y])
	BLAS.gemv!('N', ComplexOne, T.tmpMulVec, onesp, ComplexOne, y)

	# LI
	NFFT.nfft!(T.NFFTy, S.LI, T.tmpMulVecT)
	transpose!(T.tmpMulVec, T.tmpMulVecT)
	scale!(T.internal[:y], T.tmpMulVec)
	had!(T.tmpMulVec, T.left[:x])
	BLAS.gemv!('N', ComplexOne, T.tmpMulVec, onesp, ComplexOne, y)

	# RI
	NFFT.nfft!(T.NFFTy, S.RI, T.tmpMulVecT)
	transpose!(T.tmpMulVec, T.tmpMulVecT)
	scale!(T.internal[:y], T.tmpMulVec)
	had!(T.tmpMulVec, T.right[:x])
	BLAS.gemv!('N', ComplexOne, T.tmpMulVec, onesp, ComplexOne, y)

	isuniform(T) || had!(y, get(T.weights))

	return y
end

function Base.A_mul_B!(y::DenseVector{Complex{Float64}}, T::Freq2Wave2D, x::DenseVector{Complex{Float64}})
	size(T,2) == length(x) || throw(DimensionMismatch())

	X = reshape(x, wsize(T))
	A_mul_B!(y, T, X)

	return y
end

@compat function Base.:*(T::Freq2Wave, x::DenseArray)
	if !isa(x, Array{Complex{Float64}})
		x = map(Complex{Float64}, x)
	end

	y = Array{Complex{Float64}}( size(T,1) )
	A_mul_B!(y, T, x)

	return y
end

function Base.Ac_mul_B!(z::DenseVector{Complex{Float64}}, T::Freq2NoBoundaryWave1D, v::DenseVector{Complex{Float64}})
	(M = size(T,1)) == length(v) || throw(DimensionMismatch())
	size(T,2) == length(z) || throw(DimensionMismatch())

	hadc!(T.tmpMulVec, v, T.internal)
	isuniform(T) || had!(T.tmpMulVec, get(T.weights))

	myfft_adjoint!(z, T.NFFT, T.tmpMulVec)

	return z
end

function Base.Ac_mul_B!(Z::DenseMatrix{Complex{Float64}}, T::Freq2NoBoundaryWave2D, v::DenseVector{Complex{Float64}})
	(M = size(T,1)) == length(v) || throw(DimensionMismatch())
	wsize(T) == size(Z) || throw(DimensionMismatch())

	for m in 1:M
		@inbounds T.tmpMulVec[m] = v[m] * conj(T.internal[:x][m]) * conj(T.internal[:y][m])
	end
	isuniform(T) || had!(T.tmpMulVec, get(T.weights))

	myfft_adjoint!(Z, T.NFFT, T.tmpMulVec)

	return Z
end

function Base.Ac_mul_B!(z::DenseVector{Complex{Float64}}, T::Freq2BoundaryWave1D, v::DenseVector{Complex{Float64}})
	(M = size(T,1)) == length(v) || throw(DimensionMismatch())
	size(T,2) == length(z) || throw(DimensionMismatch())

	zleft, zint, zright = split(z, van_moment(T))

	# Boundary contributions are first as they don't use T.internal
	copy!(T.tmpMulVec, v)
	isuniform(T) || had!(T.tmpMulVec, get(T.weights))
	Ac_mul_B!( zleft, T.left, T.tmpMulVec )
	Ac_mul_B!( zright, T.right, T.tmpMulVec )

	# Internal scaling function
	hadc!(T.tmpMulVec, T.internal)
	myfft_adjoint!(zint, T.NFFT, T.tmpMulVec)

	return z
end

function Base.Ac_mul_B!(Z::DenseMatrix{Complex{Float64}}, T::Freq2BoundaryWave2D, v::DenseVector{Complex{Float64}})
	(M = size(T,1)) == length(v) || throw(DimensionMismatch())
	(N = wsize(T)) == size(Z) || throw(DimensionMismatch())

	# As in Ac_mul_B! for Freq2BoundaryWave1D
	if isuniform(T)
		copy!(T.weigthedVec, v)
	else
		had!(T.weigthedVec, v, get(T.weights))
	end

	vm = van_moment(T)
	S = split(Z, vm)

	# Internal coefficients
	for m in 1:M
		@inbounds T.tmpMulcVec[m] = T.weigthedVec[m] * conj(T.internal[:x][m]) * conj(T.internal[:y][m])
	end
	myfft_adjoint!(S.II, T.NFFT, T.tmpMulcVec)

	# ------------------------------------------------------------
	# Left blocks: All use 'L' for the y coordinate

	# LL
	conj!(T.tmpMulVec, T.left[:y])
	scale!(T.weigthedVec, T.tmpMulVec)
	Ac_mul_B!( S.LL, T.left[:x], T.tmpMulVec )

	# RL: Reuse T.tmpMulVec
	Ac_mul_B!( S.RL, T.right[:x], T.tmpMulVec )

	# IL
	conj!(T.tmpMulcVec, T.internal[:x])
	scale!(T.tmpMulcVec, T.tmpMulVec)
	NFFT.nfft_adjoint!( T.NFFTx, T.tmpMulVec, S.IL )

	# ------------------------------------------------------------
	# Middle blocks: All use 'I' for the y coordinate

	# LI
	conj!(T.tmpMulVec, T.left[:x])
	scale!(T.weigthedVec, T.tmpMulVec)

	conj!(T.tmpMulcVec, T.internal[:y])
	scale!(T.tmpMulcVec, T.tmpMulVec)
	transpose!(T.tmpMulVecT, T.tmpMulVec)
	NFFT.nfft_adjoint!( T.NFFTy, T.tmpMulVecT, S.LI )

	# RI: Reuse T.tmpMulcVec
	conj!(T.tmpMulVec, T.right[:x])
	scale!(T.weigthedVec, T.tmpMulVec)
	scale!(T.tmpMulcVec, T.tmpMulVec)
	transpose!(T.tmpMulVecT, T.tmpMulVec)
	NFFT.nfft_adjoint!( T.NFFTy, T.tmpMulVecT, S.RI )

	# ------------------------------------------------------------
	# Right blocks: All use 'R' for the y coordinate

	# LR
	conj!(T.tmpMulVec, T.right[:y])
	scale!(T.weigthedVec, T.tmpMulVec)
	Ac_mul_B!( S.LR, T.left[:x], T.tmpMulVec )

	# RR: Reuse T.tmpMulVec
	Ac_mul_B!( S.RR, T.right[:x], T.tmpMulVec )

	# IR
	conj!(T.tmpMulcVec, T.internal[:x])
	scale!(T.tmpMulcVec, T.tmpMulVec)
	NFFT.nfft_adjoint!( T.NFFTx, T.tmpMulVec, S.IR )

	return Z
end

function Base.Ac_mul_B!(z::DenseVector{Complex{Float64}}, T::Freq2Wave2D, v::DenseVector{Complex{Float64}})
	size(T,2) == length(z) || throw(DimensionMismatch())

	Z = reshape(z, wsize(T))
	Ac_mul_B!(Z, T, v)

	return z
end

function Base.Ac_mul_B(T::Freq2Wave, v::AbstractVector)
	if !isa(v, Array{Complex{Float64}})
		v = map(Complex{Float64}, v)
	end

	z = Array{Complex{Float64}}(size(T,2))
	Ac_mul_B!(z, T, v)

	return z
end

@compat function Base.:\(T::Freq2Wave, Y::AbstractMatrix)
	length(Y) == (M = size(T,1)) || throw(DimensionMismatch())

	y = flatten_view(Y)
	x = T \ y
end

@compat function Base.:\(T::Freq2Wave, y::AbstractVector)
	if !isa(y, Array{Complex{Float64}})
		y = map(Complex{Float64}, y)
	end

	# Non-uniform samples: Scale observations
	if !isuniform(T)
		y .*= get(T.weights)
	end

	x0 = zeros(Complex{Float64}, wsize(T))
	x = cgnr(T, y, x0)
	#x, h = lsqr(T, y)
end

@doc """
	collect(Freq2Wave) -> Matrix

Return the full change of basis matrix.

In 2D, the reconstruction grid is sorted by the `y` coordinate, i.e., the order is
(1,1),
(2,1),
(3,1),
(1,2),
(2,2),
(3,2)
etc
"""->
function Base.collect(T::Freq2NoBoundaryWave2D)
	M = size(T,1)
	Nx, Ny = wsize(T)
	F = Array{Complex{Float64}}(M, Nx*Ny)

	phi = T.internal[:x] .* T.internal[:y]

	offsetx = div(Nx, 2) + 1
	offsety = div(Ny, 2) + 1
	row_idx = 0
	for ny in 1:Ny, nx in 1:Nx
		row_idx += 1
		for m in 1:M
			@inbounds F[m,row_idx] = phi[m]*cis( -twoπ*((nx-offsetx)*T.NFFT.x[1,m] + (ny-offsety)*T.NFFT.x[2,m]) )
		end
	end

	if !isuniform(T)
		scale!(get(T.weights), F)
	end

	return F
end

function Base.collect(T::Freq2BoundaryWave2D)
	M = size(T,1)
	Nx, Ny = wsize(T)
	F = Array{Complex{Float64}}(M, Nx*Ny)

	phix = Array{Complex{Float64}}(M)
	phiy = similar(phix)

	row_idx = 0
	for ny in 1:Ny
		unsafe_FourScaling!(phiy, T, ny, :y)
		for nx in 1:Nx
			unsafe_FourScaling!(phix, T, nx, :x)
			had!(phix, phiy)
			F[:,row_idx+=1] = phix
		end
	end

	if !isuniform(T)
		scale!(get(T.weights), F)
	end

	return F
end

@doc """
	unsafe_FourScaling!(phi, T::Freq2BoundaryWave2D, n::Int, d::Symbol)

Replace `phi` with the `n`'th "column" from dimension `d` of `T`.
"""->
function unsafe_FourScaling!(phi::Vector{Complex{Float64}}, T::Freq2BoundaryWave2D, n::Integer, d::Symbol)
	M = length(phi)
	N = 2^wscale(T)
	p = van_moment(T)

	if p < n <= N-p
		offset = div(N, 2) + 1
        D = (d == :x ? 1 : 2)
		for m in 1:M
			@inbounds phi[m] = T.internal[d][m]*cis( -twoπ*(n-offset)*T.NFFT.x[D,m] )
		end
	elseif 1 <= n <= p
		unsafe_copy!( phi, 1, T.left[d], (n-1)*M+1, M )
	else
		unsafe_copy!( phi, 1, T.right[d], (n-N+p-1)*M+1, M )
	end
end

# ------------------------------------------------------------
# Common

function Base.size(T::Freq2Wave)
	( size(T,1), size(T,2) )
end

function Base.size(T::Freq2Wave, d::Integer)
	if d == 1
        prod(T.NFFT.M)
	elseif d == 2
		prod( wsize(T) )
	else
		throw(AssertionError())
	end
end

van_moment(T::Freq2Wave1D) = hasboundary(T) ? size(T.left,2) : van_moment(T.wavename)
van_moment(T::Freq2Wave2D) = hasboundary(T) ? size(T.left[:x],2)::Int64 : van_moment(T.wavename)

function Base.eltype(::Freq2Wave)
	return Complex{Float64}
end

function Base.show(io::IO, T::Freq2Wave)
	D = dim(T)
	println(io, D, "D change of basis matrix")

	isuniform(T) ?  U = " " : U = " non-"
	M = size(T,1)
	println(io, "From: ", M, U, "uniform frequency samples")

	D == 1 ? N = size(T,2) : N = wsize(T)
	print(io, "To: ", N, " ", T.wavename, " wavelets")
end


using ITensorGaussianMPS
using ITensors
using LinearAlgebra
using Test

@testset "Basic" begin
  # Test Givens rotations
  v = randn(6)
  g, r = ITensorGaussianMPS.givens_rotations(v)
  @test g * v ≈ r * [n == 1 ? 1 : 0 for n in 1:length(v)]
end

@testset "Fermion" begin
  N = 10
  Nf = N ÷ 2

  # Hopping
  t = 1.0

  # Hopping Hamiltonian
  h = Hermitian(diagm(1 => fill(-t, N - 1), -1 => fill(-t, N - 1)))
  e, u = eigen(h)

  @test h * u ≈ u * Diagonal(e)

  E = sum(e[1:Nf])

  # Get the Slater determinant
  Φ = u[:, 1:Nf]
  @test h * Φ ≈ Φ * Diagonal(e[1:Nf])

  # Diagonalize the correlation matrix as a
  # Gaussian MPS (GMPS)
  n, gmps = slater_determinant_to_gmps(Φ; maxblocksize=4)

  ns = round.(Int, n)
  @test sum(ns) == Nf

  Λ = conj(Φ) * transpose(Φ)
  @test gmps * Λ * gmps' ≈ Diagonal(ns) rtol = 1e-2
  @test gmps' * Diagonal(ns) * gmps ≈ Λ rtol = 1e-2

  # Form the MPS
  s = siteinds("Fermion", N; conserve_qns=true)
  ψ = slater_determinant_to_mps(s, Φ; blocksize=4)

  ampo = AutoMPO()
  for i in 1:N, j in 1:N
    if h[i, j] ≠ 0
      ampo .+= h[i, j], "Cdag", i, "C", j
    end
  end
  H = MPO(ampo, s)

  @test inner(ψ, H, ψ) ≈ E rtol = 1e-5

  # Compare to DMRG
  sweeps = Sweeps(10)
  maxdim!(sweeps, 10, 20, 40, 60)
  cutoff!(sweeps, 1E-12)
  energy, ψ̃ = dmrg(H, productMPS(s, n -> n ≤ Nf ? "1" : "0"), sweeps; outputlevel=0)

  # Create an mps
  @test abs(inner(ψ, ψ̃)) ≈ 1 rtol = 1e-5
  @test inner(ψ̃, H, ψ̃) ≈ inner(ψ, H, ψ) rtol = 1e-5
  @test E ≈ energy
end

@testset "Fermion (complex)" begin
  N = 10
  Nf = N ÷ 2

  # Hopping
  θ = π / 8
  t = exp(im * θ)

  # Hopping Hamiltonian
  h = Hermitian(diagm(1 => fill(-t, N - 1), -1 => fill(-conj(t), N - 1)))
  e, u = eigen(h)

  @test h * u ≈ u * Diagonal(e)

  E = sum(e[1:Nf])

  # Get the Slater determinant
  Φ = u[:, 1:Nf]
  @test h * Φ ≈ Φ * Diagonal(e[1:Nf])

  # Diagonalize the correlation matrix as a
  # Gaussian MPS (GMPS)
  n, gmps = slater_determinant_to_gmps(Φ; maxblocksize=4)

  ns = round.(Int, n)
  @test sum(ns) == Nf

  Λ = conj(Φ) * transpose(Φ)
  @test gmps * Λ * gmps' ≈ Diagonal(ns) rtol = 1e-2
  @test gmps' * Diagonal(ns) * gmps ≈ Λ rtol = 1e-2

  # Form the MPS
  s = siteinds("Fermion", N; conserve_qns=true)
  ψ = slater_determinant_to_mps(s, Φ; blocksize=4)

  ampo = AutoMPO()
  for i in 1:N, j in 1:N
    if h[i, j] ≠ 0
      ampo .+= h[i, j], "Cdag", i, "C", j
    end
  end
  H = MPO(ampo, s)

  @test inner(ψ, H, ψ) ≈ E rtol = 1e-5
  @test inner(ψ, H, ψ) / norm(ψ) ≈ E rtol = 1e-5

  # Compare to DMRG
  sweeps = Sweeps(10)
  maxdim!(sweeps, 10, 20, 40, 60)
  cutoff!(sweeps, 1E-12)
  energy, ψ̃ = dmrg(H, productMPS(s, n -> n ≤ Nf ? "1" : "0"), sweeps; outputlevel=0)

  # Create an mps
  @test abs(inner(ψ, ψ̃)) ≈ 1 rtol = 1e-5
  @test inner(ψ̃, H, ψ̃) ≈ inner(ψ, H, ψ) rtol = 1e-5
  @test E ≈ energy
end


using Mitosis, LinearAlgebra, Test, Statistics
m = [1.0, 0.5]
K = 10000
Q = Matrix(1.0I, 2, 2)
G = Gaussian(μ=m, Σ=Q)

@testset "Gauss" begin
    @test G === Gaussian{keys(G)}(m, Q)

    @test m === mean(G)
    @test Q === cov(G)

    @test 10/sqrt(K) > norm(m - mean(rand(G) for x in 1:K))
    @test 10/sqrt(K) > norm(Q - cov([rand(G) for x in 1:K]))
end

p = WGaussian(μ=m, Σ=Q, c=0.0)
p2 = convert(WGaussian{(:μ,:Σ,:c)}, convert(WGaussian{(:F,:Γ,:c)}, p))
@test p ≈ p2

p = Gaussian(μ=m, Σ=Q)
p2 = convert(Gaussian{(:μ,:Σ)}, convert(Gaussian{(:F,:Γ)}, p))
@test p ≈ p2