text:
push!(LOAD_PATH, "/Users/saydjari/Dropbox/GradSchool_AKS/Doug/Projects/DHC/main")
using DHC_2DUtils
using Profile
using BenchmarkTools
using FFTW
using HDF5
using Test
using SparseArrays

filter_hash = fink_filter_hash(1, 8, nx=256, pc=1, wd=2)
out = DHC_compute(rand(256,256),filter_hash,doS2=true, doS12=false, doS20=false)
out = DHC_compute(rand(256,256),filter_hash,doS2=true, doS12=true, doS20=false)
out = DHC_compute(rand(256,256),filter_hash,doS2=true, doS12=true, doS20=true)
out = DHC_compute(rand(256,256),filter_hash,doS2=true, doS12=true, doS20=true)
out = DHC_compute(rand(256,256),filter_hash,doS2=true, doS12=true, doS20=true)
out = DHC_compute(rand(256,256),filter_hash,doS2=true, doS12=true, doS20=true)

function fink_filter_bank(c, L; nx=256, wd=1, pc=1, shift=false, Omega=false, safety_on=true)
    #c     - sets the scale sampling rate (1 is dyadic, 2 is half dyadic)
    #L     - number of angular bins (usually 8*pc or 16*pc)
    #wd    - width of the wavelets (default 1, wd=2 for a double covering)
    #pc    - plane coverage (default 1, full 2pi 2)
    #shift - shift in θ by 1/2 of the θ spacing
    #Omega - true= append Omega filter (all power beyond Nyquist) so the sum of filters is 1.0

    # -------- assertion errors to make sure arguments are reasonable
    #@test wd <= L/2

    # -------- set parameters
    dθ   = pc*π/L
    θ_sh = shift ? dθ/2 : 0.0
    dx   = nx/2-1

    im_scale = convert(Int8,log2(nx))
    # -------- number of bins in radial direction (size scales)
    J = (im_scale-2)*c

    # -------- allocate output array of zeros
    filt      = zeros(nx, nx, J*L+(Omega ? 2 : 1))
    psi_index = zeros(Int32, J, L)
    theta     = zeros(Float64, L)
    j_value   = zeros(Float64, J)

    # -------- compute the required wd
    j_rad_exp = zeros(J)
    for j_ind = 1:J
        j = j_ind/c
        jrad  = im_scale-j-1
        j_rad_exp[j_ind] = 2^(jrad)
    end

    wd_j = max.(ceil.(L./(pc.*π.*j_rad_exp)),wd)

    if !safety_on
        wd_j.=wd
    end

    # loop over wd from small to large
    ## there is some uneeded redundancy added by doing this esp in l loop
    for wd in sort(unique(wd_j))
        # -------- allocate theta and logr arrays
        θ    = zeros(nx, nx)
        logr = zeros(nx, nx)

        wdθ  = wd*dθ
        norm = 1.0/sqrt(wd)
        # -------- loop over l
        for l = 0:L-1
            θ_l        = dθ*l+θ_sh
            theta[l+1] = θ_l

        # -------- allocate anggood BitArray
            anggood = falses(nx, nx)

        # -------- loop over pixels
            for x = 1:nx
                sx = mod(x+dx,nx)-dx -1    # define sx,sy so that no fftshift() needed
                for y = 1:nx
                    sy = mod(y+dx,nx)-dx -1
                    θ_pix  = mod(atan(sy, sx)+π -θ_l, 2*π)
                    θ_good = abs(θ_pix-π) <= wdθ

                    # If this is a pixel we might use, calculate log2(r)
                    if θ_good
                        anggood[y, x] = θ_good
                        θ[y, x]       = θ_pix
                        r2            = sx^2 + sy^2
                        logr[y, x]    = 0.5*log2(max(1,r2))
                    end
                end
            end
            angmask = findall(anggood)
        # -------- compute the wavelet in the Fourier domain
        # -------- the angular factor is the same for all j
            F_angular = norm .* cos.((θ[angmask].-π).*(L/(2*wd*pc)))

        # -------- loop over j for the radial part
        #    for (j_ind, j) in enumerate(1/c:1/c:im_scale-2)
            j_ind_w_wd = findall(wd_j.==wd)
            for j_ind = j_ind_w_wd
                j = j_ind/c
                j_value[j_ind] = j  # store for later
                jrad  = im_scale-j-1
                Δj    = abs.(logr[angmask].-jrad)
                rmask = (Δj .<= 1/c)

        # -------- radial part
                F_radial = cos.(Δj[rmask] .* (c*π/2))
                ind      = angmask[rmask]
        #      Let's have these be (J,L) if you reshape...
        #        f_ind    = (j_ind-1)*L+l+1
                f_ind    = j_ind + l*J
                filt[ind, f_ind] = F_radial .* F_angular[rmask]
                psi_index[j_ind,l+1] = f_ind
            end
        end
    end

    # -------- phi contains power near k=0 not yet accounted for
    filter_power = (sum(filt.*filt, dims=3))[:,:,1]

    # -------- for plane half-covered (pc=1), add other half-plane
    if pc == 1
        filter_power .+= circshift(filter_power[end:-1:1,end:-1:1],(1,1))
    end

    # -------- compute power required to sum to 1.0
    i0 = round(Int16,nx/2-2)
    i1 = round(Int16,nx/2+4)
    center_power = 1.0 .- fftshift(filter_power)[i0:i1,i0:i1]
    zind = findall(center_power .< 1E-15)
    center_power[zind] .= 0.0  # set small numbers to zero
    phi_cen = zeros(nx, nx)
    phi_cen[i0:i1,i0:i1] = sqrt.(center_power)

    # -------- before adding ϕ to filter bank, renormalize ψ if pc=1
    if pc==1 filt .*= sqrt(2.0) end  # double power for half coverage

    # -------- add result to filter array
    phi_index  = J*L+1
    filt[:,:,phi_index] .= fftshift(phi_cen)

    # -------- metadata dictionary
    info=Dict()
    info["npix"]         = nx
    info["j_value"]      = j_value
    info["theta_value"]  = theta
    info["psi_index"]    = psi_index
    info["phi_index"]    = phi_index
    info["pc"]           = pc
    info["wd"]           = wd_j
    info["fs_center_r"]  = j_rad_exp

    if Omega     # append a filter containing the rest (outside Nyquist)
        filter_power += filt[:,:,phi_index].^2
        edge_power    = 1.0 .- filter_power
        zind          = findall(edge_power .< 1E-15)
        edge_power[zind]     .= 0.0  # set small numbers to zero
        Omega_index           = J*L+2
        info["Omega_index"]   = Omega_index
        filt[:,:,Omega_index] = sqrt.(edge_power)
    end

    return filt, info
end

function fink_filter_bank_3dizer(hash, cz; nz=256)
    # -------- set parameters
    nx = convert(Int64,hash["npix"])
    n2d = size(hash["filt_value"])[1]
    dz = nz/2-1

    im_scale = convert(Int8,log2(nz))
    # -------- number of bins in radial direction (size scales)
    K = (im_scale-2)*cz

    # -------- allocate output array of zeros
    #filt      = zeros(Float64, nx, nx, nz, K*n2d)
    #psi_index = zeros(Int32, J, L, K)
    #theta     = zeros(Float64, L)
    k_value   = zeros(Float64, K)
    logr      = zeros(Float64,nz)
    filt_temp = zeros(Float64,nx,nx,nz)
    F_z       = zeros(Float64,nx,nx,nz)

    filtind = fill(CartesianIndex{2}[], K*n2d)
    filtval = fill(Float64[], K*n2d)

    for z = 1:convert(Int64,dz+1)
        sz = mod(z+dz,nz)-dz-1
        z2 = sz^2
        logr[z] = 0.5*log2(max(1,z2))
    end

    for k_ind = 1:K
        k = k_ind/cz
        k_value[k_ind] = k  # store for later
        krad  = im_scale-k-1
        Δk    = abs.(logr.-krad)
        kmask = (Δk .<= 1/cz)
        k_count = count(kmask)

        # -------- radial part
        F_z[:,:,kmask].= reshape(cos.(Δk[kmask] .* (cz*π/2)),1,1,k_count)

        for index = 1:n2d
            p_ind = hash["filt_index"][index]
            p_filt = hash["filt_value"][index]
            f_ind    = k_ind + (index-1)*K
            filt_tmp = p_filt.*F_z[p_ind,kmask] #
            ind = findall(filt_tmp .> 1E-13)
            val = filt_tmp[ind]
            filtind[f_ind] = ind
            filtval[f_ind] = val
            #psi_index[j_ind,l+1] = f_ind
        end
    end

    # -------- metadata dictionary
    info=Dict()
    info["2d_npix"]         = hash["npix"]
    info["2d_j_value"]      = hash["npix"]
    info["2d_theta_value"]  = hash["theta_value"]
    info["2d_psi_index"]    = hash["psi_index"]
    info["2d_phi_index"]    = hash["phi_index"]
    info["2d_pc"]           = hash["pc"]
    info["2d_wd"]           = hash["wd"]
    info["2d_fs_center_r"]  = hash["fs_center_r"]
    info["2d_filt_index"]   = hash["filt_index"]
    info["2d_filt_value"]   = hash["filt_value"]

    info["nz"]              = nz
    info["cz"]              = cz
    info["k_value"]         = k_value
    info["filt_index"]      = filtind
    info["filt_value"]      = filtval

    return info
end

hash = fink_filter_hash(1, 4, nx=128, pc=1, wd=1)

@time filt_3d = fink_filter_bank_3dizer(hash, 1, nz=128)

filter_temp = zeros(256,256,256)
filter_temp[filt_3d["filt_index"][12]] = filt_3d["filt_index"][12]

h5write("../DHC/scratch_AKS/data/filt_3d_rs.h5", "main/data", fft2())

Profile.clear()
@profile filt_3d = fink_filter_bank_3dizer(hash, 1, nz=128)

Juno.profiler()

@benchmark fink_filter_bank_3dizer(hash, 1, nz=128)

hash = fink_filter_hash(1, 8, nx=128, pc=1, wd=1)
@benchmark filt_3d = fink_filter_bank_3dizer(hash, 1, nz=128)

hash = fink_filter_hash(1, 8, nx=256, pc=1, wd=1)
@benchmark filt_3d = fink_filter_bank_3dizer(hash, 1, nz=256)

hash = fink_filter_hash(1, 8, nx=256, pc=1, wd=2)
@benchmark filt_3d = fink_filter_bank_3dizer(hash, 1, nz=256)

## Here we go with the plotting

function plot_filter_bank_QA(hash; fname="filter_bank_QA.png")

    # -------- define plot1 to append plot to a list
    function plot1(ps, image; clim=nothing, bin=1.0, fsz=12, label=nothing)
        # ps    - list of plots to append to
        # image - image to heatmap
        # clim  - tuple of color limits (min, max)
        # bin   - bin down by this factor, centered on nx/2+1
        # fsz   - font size
        marg   = -1mm
        nx, ny = size(image)
        nxb    = nx/round(Integer, 2*bin)

        # -------- center on nx/2+1
        i0 = max(1,round(Integer, (nx/2+2)-nxb-1))
        i1 = min(nx,round(Integer, (nx/2)+nxb+1))
        lims = [i0,i1]
        subim = image[i0:i1,i0:i1]
        push!(ps, heatmap(image,aspect_ratio=:equal,clim=clim,
            xlims=lims, ylims=lims, size=(400,400),
            legend=false,xtickfontsize=fsz,ytickfontsize=fsz,#tick_direction=:out,
            rightmargin=marg,leftmargin=marg,topmargin=marg,bottommargin=marg))
        if label != nothing
            annotate!(lims'*[.96,.04],lims'*[.09,.91],text(label,:left,:white,16))
        end
        return
    end

    # -------- initialize array of plots
    ps   = []
    jval = hash["2d_j_value"]
    J    = length(jval)
    L    = length(hash["2d_theta_value"])
    nx   = hash["2d_npix"]
    nz   = hash["nz"]
    # plot the first 3 and last 3 j values
    jplt = [1:3;J-2:J]

    ind = zeros(Int64,J,L)
    for j_ind=1:J
        for l=1:L
            ind[j_ind,l] = Int(j_ind-1 + (l-1)*J)
        end
    end
    ind .+= 1

    # -------- loop over j
    for j_ind in jplt
        # first 3 filters in Fourier space
        j = jval[j_ind]
        bfac = max(1,2.0^(j-1))
        for l = 1:3
            label=string("j=",j_ind," ℓ=",l-1)
            #label=latexstring("\\tilde{\\psi}_{$j_ind,$(l-1)}")
            filt = zeros(nx,nx,nz)
            println(nx)
            println(nz)
            filt[hash["2d_filt_index"][2],hash["filt_index"][2]] = hash["filt_value"][2]
            temp = dropdims(sum(filt,dims=3))
            plot1(ps, fftshift(filt[:,:]), clim=(0,1),
                  bin=bfac, label=label)
        end
        fac = max(1,2.0^(5-j))
        # real part of config space of 3rd filter
        plot1(ps, fftshift(real(ifft(filt[:,:,ind[j_ind,3]]))),
              bin=fac, label=string("j=",j_ind," ℓ=",3-1))

        # sum all filters to form (partial) ring
        ring = dropdims(sum(filt[:,:,ind[j_ind,:]].^2, dims=3), dims=3)
        plot1(ps, fftshift(ring[:,:,1]), clim=(0,1),
              bin=bfac, label=string("j=",j_ind," ℓ=",0,":",L-1))
    end

    # wavepow = fftshift(dropdims(sum(filt[:,:,ind].^2, dims=(3,4)), dims=(3,4)))
    # plot1(ps, wavepow, clim=(-0.1,1), label=string("j=1:",J," ℓ=0:",L-1))
    #
    # if info["pc"]==1
    #     wavepow += circshift(wavepow[end:-1:1,end:-1:1],(1,1))
    # end
    # plot1(ps, wavepow, clim=(-0.1,1), bin=32, label="missing power")
    # phi = filt[:,:,info["phi_index"]]
    # phi_shift = fftshift(phi)
    # plot1(ps, phi_shift.^2, clim=(-0.1,1), bin=32, label="ϕ")
    # disc = wavepow+(phi_shift.^2)
    # plot1(ps, fftshift(real(ifft(phi))), label="ϕ")
    # plot1(ps, disc, clim=(-0.1,1), bin=32, label="all")

    myplot = plot(ps..., layout=(7,5), size=(1400,2000))
    savefig(myplot, fname)
end

plot_filter_bank_QA(filt_3d, fname="/Users/saydjari/Dropbox/GradSchool_AKS/Doug/Projects/DHC/scratch_AKS/filt-3d.png")

## Plotting is not working because keeping track of indices is a mess. Lets
## try sparse matrices

function fink_filter_bank(c, L; nx=256, wd=1, pc=1, shift=false, Omega=false, safety_on=true)
    #c     - sets the scale sampling rate (1 is dyadic, 2 is half dyadic)
    #L     - number of angular bins (usually 8*pc or 16*pc)
    #wd    - width of the wavelets (default 1, wd=2 for a double covering)
    #pc    - plane coverage (default 1, full 2pi 2)
    #shift - shift in θ by 1/2 of the θ spacing
    #Omega - true= append Omega filter (all power beyond Nyquist) so the sum of filters is 1.0

    # -------- assertion errors to make sure arguments are reasonable
    @test wd <= L/2

    # -------- set parameters
    dθ   = pc*π/L
    θ_sh = shift ? dθ/2 : 0.0
    dx   = nx/2-1

    im_scale = convert(Int8,log2(nx))
    # -------- number of bins in radial direction (size scales)
    J = (im_scale-2)*c

    # -------- allocate output array of zeros
    filt      = zeros(nx, nx, J*L+(Omega ? 2 : 1))
    psi_index = zeros(Int32, J, L)
    theta     = zeros(Float64, L)
    j_value   = zeros(Float64, J)

    # -------- compute the required wd
    j_rad_exp = zeros(J)
    for j_ind = 1:J
        j = j_ind/c
        jrad  = im_scale-j-1
        j_rad_exp[j_ind] = 2^(jrad)
    end

    wd_j = max.(ceil.(L./(pc.*π.*j_rad_exp)),wd)

    if !safety_on
        wd_j.=wd
    end

    # loop over wd from small to large
    ## there is some uneeded redundancy added by doing this esp in l loop
    for wd in sort(unique(wd_j))
        # -------- allocate theta and logr arrays
        θ    = zeros(nx, nx)
        logr = zeros(nx, nx)

        wdθ  = wd*dθ
        norm = 1.0/sqrt(wd)
        # -------- loop over l
        for l = 0:L-1
            θ_l        = dθ*l+θ_sh
            theta[l+1] = θ_l

        # -------- allocate anggood BitArray
            anggood = falses(nx, nx)

        # -------- loop over pixels
            for x = 1:nx
                sx = mod(x+dx,nx)-dx -1    # define sx,sy so that no fftshift() needed
                for y = 1:nx
                    sy = mod(y+dx,nx)-dx -1
                    θ_pix  = mod(atan(sy, sx)+π -θ_l, 2*π)
                    θ_good = abs(θ_pix-π) <= wdθ

                    # If this is a pixel we might use, calculate log2(r)
                    if θ_good
                        anggood[y, x] = θ_good
                        θ[y, x]       = θ_pix
                        r2            = sx^2 + sy^2
                        logr[y, x]    = 0.5*log2(max(1,r2))
                    end
                end
            end
            angmask = findall(anggood)
        # -------- compute the wavelet in the Fourier domain
        # -------- the angular factor is the same for all j
            F_angular = norm .* cos.((θ[angmask].-π).*(L/(2*wd*pc)))

        # -------- loop over j for the radial part
        #    for (j_ind, j) in enumerate(1/c:1/c:im_scale-2)
            j_ind_w_wd = findall(wd_j.==wd)
            for j_ind = j_ind_w_wd
                j = j_ind/c
                j_value[j_ind] = j  # store for later
                jrad  = im_scale-j-1
                Δj    = abs.(logr[angmask].-jrad)
                rmask = (Δj .<= 1/c)

        # -------- radial part
                F_radial = cos.(Δj[rmask] .* (c*π/2))
                ind      = angmask[rmask]
        #      Let's have these be (J,L) if you reshape...
        #        f_ind    = (j_ind-1)*L+l+1
                f_ind    = j_ind + l*J
                filt[ind, f_ind] = F_radial .* F_angular[rmask]
                psi_index[j_ind,l+1] = f_ind
            end
        end
    end

    # -------- phi contains power near k=0 not yet accounted for
    filter_power = (sum(filt.*filt, dims=3))[:,:,1]

    # -------- for plane half-covered (pc=1), add other half-plane
    if pc == 1
        filter_power .+= circshift(filter_power[end:-1:1,end:-1:1],(1,1))
    end

    # -------- compute power required to sum to 1.0
    i0 = round(Int16,nx/2-2)
    i1 = round(Int16,nx/2+4)
    center_power = 1.0 .- fftshift(filter_power)[i0:i1,i0:i1]
    zind = findall(center_power .< 1E-15)
    center_power[zind] .= 0.0  # set small numbers to zero
    phi_cen = zeros(nx, nx)
    phi_cen[i0:i1,i0:i1] = sqrt.(center_power)

    # -------- before adding ϕ to filter bank, renormalize ψ if pc=1
    if pc==1 filt .*= sqrt(2.0) end  # double power for half coverage

    # -------- add result to filter array
    phi_index  = J*L+1
    filt[:,:,phi_index] .= fftshift(phi_cen)

    # -------- metadata dictionary
    info=Dict()
    info["npix"]         = nx
    info["j_value"]      = j_value
    info["theta_value"]  = theta
    info["psi_index"]    = psi_index
    info["phi_index"]    = phi_index
    info["pc"]           = pc
    info["wd"]           = wd_j
    info["fs_center_r"]  = j_rad_exp

    if Omega     # append a filter containing the rest (outside Nyquist)
        filter_power += filt[:,:,phi_index].^2
        edge_power    = 1.0 .- filter_power
        zind          = findall(edge_power .< 1E-15)
        edge_power[zind]     .= 0.0  # set small numbers to zero
        Omega_index           = J*L+2
        info["Omega_index"]   = Omega_index
        filt[:,:,Omega_index] = sqrt.(edge_power)
    end

    return filt, info
end

@benchmark hash = fink_filter_bank(1, 8, nx=256, pc=1, wd=2)

function fink_filter_bank_AKS(c, L; nx=256, wd=1, pc=1, shift=false, Omega=false, safety_on=true)
    #c     - sets the scale sampling rate (1 is dyadic, 2 is half dyadic)
    #L     - number of angular bins (usually 8*pc or 16*pc)
    #wd    - width of the wavelets (default 1, wd=2 for a double covering)
    #pc    - plane coverage (default 1, full 2pi 2)
    #shift - shift in θ by 1/2 of the θ spacing
    #Omega - true= append Omega filter (all power beyond Nyquist) so the sum of filters is 1.0

    # -------- assertion errors to make sure arguments are reasonable
    @test wd <= L/2

    # -------- set parameters
    dθ   = pc*π/L
    θ_sh = shift ? dθ/2 : 0.0
    dx   = nx/2-1

    im_scale = convert(Int8,log2(nx))
    # -------- number of bins in radial direction (size scales)
    J = (im_scale-2)*c

    # -------- allocate output array of zeros
    filt      = spzeros(nx,nx,J*L+(Omega ? 2 : 1))
    filt_pow  = spzeros(nx,nx)
    psi_index = zeros(Int32, J, L)
    theta     = zeros(Float64, L)
    j_value   = zeros(Float64, J)

    # -------- compute the required wd
    j_rad_exp = zeros(J)
    for j_ind = 1:J
        j = j_ind/c
        jrad  = im_scale-j-1
        j_rad_exp[j_ind] = 2^(jrad)
    end

    wd_j = max.(ceil.(L./(pc.*π.*j_rad_exp)),wd)

    if !safety_on
        wd_j.=wd
    end

    # loop over wd from small to large
    ## there is some uneeded redundancy added by doing this esp in l loop
    for wd in sort(unique(wd_j))
        # -------- allocate theta and logr arrays
        θ    = spzeros(nx,nx)
        logr = spzeros(nx,nx)
        F_angular = spzeros(nx,nx)

        wdθ  = wd*dθ
        norm = 1.0/sqrt(wd)
        # -------- loop over l
        for l = 0:L-1
            θ_l        = dθ*l+θ_sh
            theta[l+1] = θ_l

        # -------- allocate anggood BitArray
            anggood = falses(nx, nx)

        # -------- loop over pixels
            for x = 1:nx
                sx = mod(x+dx,nx)-dx -1    # define sx,sy so that no fftshift() needed
                for y = 1:nx
                    sy = mod(y+dx,nx)-dx -1
                    θ_pix  = mod(atan(sy, sx)+π -θ_l, 2*π)
                    θ_good = abs(θ_pix-π) <= wdθ

                    # If this is a pixel we might use, calculate log2(r)
                    if θ_good
                        anggood[y, x] = θ_good
                        θ[y, x]       = θ_pix
                        r2            = sx^2 + sy^2
                        logr[y, x]    = 0.5*log2(max(1,r2))
                    end
                end
            end
        # -------- compute the wavelet in the Fourier domain
        # -------- the angular factor is the same for all j
            f = x -> norm*cos((x-π)*(L/(2*wd*pc)))
            map!(f,θ.nzval,θ.nzval) #inplace so theta becomes F_ang

        # -------- loop over j for the radial part
        #    for (j_ind, j) in enumerate(1/c:1/c:im_scale-2)
            j_ind_w_wd = findall(wd_j.==wd)
            for j_ind = j_ind_w_wd
                Δj    = deepcopy(logr)
                j = j_ind/c
                j_value[j_ind] = j  # store for later
                jrad  = im_scale-j-1
                g = x -> abs(x-jrad)
                map!(g,Δj.nzval,logr.nzval)
        # -------- radial part
                h = x -> norm*cos(x.*(c*π/2))
                map!(h,Δj.nzval,Δj.nzval) #inplace so Δj becomes F_rad
        #      Let's have these be (J,L) if you reshape...
                f_ind    = j_ind + l*J
                filt[:,:,f_ind] = Δj .* θ
            end
        end
    end
    return filt
    # -------- phi contains power near k=0 not yet accounted for
    filt_pow = (sum(filt.^2, dims=3))
    println(filt_pow)
    println(size(filt_pow))

    # -------- for plane half-covered (pc=1), add other half-plane
    if pc == 1
        filt_pow .+= circshift(filt_pow[end:-1:1,end:-1:1],(1,1))
    end

    # -------- compute power required to sum to 1.0
    i0 = round(Int16,nx/2-2)
    i1 = round(Int16,nx/2+4)
    center_power = 1.0 .- fftshift(filt_pow)[i0:i1,i0:i1]
    zind = findall(center_power .< 1E-15)
    center_power[zind] .= 0.0  # set small numbers to zero
    phi_cen = zeros(nx, nx)
    phi_cen[i0:i1,i0:i1] = sqrt.(center_power)

    # -------- before adding ϕ to filter bank, renormalize ψ if pc=1
    if pc==1 filt .*= sqrt(2.0) end  # double power for half coverage

    # -------- add result to filter array
    phi_index  = J*L+1
    filt[:,:,phi_index] .= fftshift(phi_cen)

    # -------- metadata dictionary
    info=Dict()
    info["npix"]         = nx
    info["j_value"]      = j_value
    info["theta_value"]  = theta
    info["psi_index"]    = psi_index
    info["phi_index"]    = phi_index
    info["pc"]           = pc
    info["wd"]           = wd_j
    info["fs_center_r"]  = j_rad_exp

    if Omega     # append a filter containing the rest (outside Nyquist)
        filter_power += filt[:,:,phi_index].^2
        edge_power    = 1.0 .- filter_power
        zind          = findall(edge_power .< 1E-15)
        edge_power[zind]     .= 0.0  # set small numbers to zero
        Omega_index           = J*L+2
        info["Omega_index"]   = Omega_index
        filt[:,:,Omega_index] = sqrt.(edge_power)
    end

    return filt, info
end

@time hash = fink_filter_bank_AKS(1, 8, nx=256, pc=1, wd=2)

hash

A = SparseArrays.sprand(256,256,0.1)

B = sprand(256,256,0.1)

@time A*B

A1 = Matrix(A)
B1 = Matrix(B)

@time A1*B1

A.nzval.-=1

A

f = x -> cos(x)

map!(f,A1)

f = x -> 2*x
foo = [1   2   3];
foo1 = [7   8   9];
map!(f,foo,foo1)

foo
foo1

map!(f,A.nzval,A.nzval)
C = copy(A)

C.nzval .= 0

C

A

@benchmark A.^2

@benchmark A.*A

sum(A)

##

hash = fink_filter_hash(1, 8, nx=256, pc=1, wd=2)

function fink_filter_bank_3dizer(hash, cz; nz=256)
    # -------- set parameters
    nx = convert(Int64,hash["npix"])
    n2d = size(hash["filt_value"])[1]
    dz = nz/2-1

    im_scale = convert(Int8,log2(nz))
    # -------- number of bins in radial direction (size scales)
    K = (im_scale-2)*cz

    # -------- allocate output array of zeros
    #filt      = zeros(Float64, nx, nx, nz, K*n2d)
    #psi_index = zeros(Int32, J, L, K)
    #theta     = zeros(Float64, L)
    k_value   = zeros(Float64, K)
    logr      = zeros(Float64,nz)
    filt_temp = zeros(Float64,nx,nx,nz)
    F_z       = zeros(Float64,nx,nx,nz)

    filtind = fill(CartesianIndex{3}[], K*n2d)
    filtval = fill(Float64[], K*n2d)

    for z = 1:convert(Int64,dz+1)
        sz = mod(z+dz,nz)-dz-1
        z2 = sz^2
        logr[z] = 0.5*log2(max(1,z2))
    end

    @inbounds for k_ind = 1:K
        k = k_ind/cz
        k_value[k_ind] = k  # store for later
        krad  = im_scale-k-1
        Δk    = abs.(logr.-krad)
        kmask = (Δk .<= 1/cz)
        k_count = count(kmask)

        # -------- radial part
        @views F_z[:,:,kmask].= reshape(cos.(Δk[kmask] .* (cz*π/2)),1,1,k_count)

        @inbounds for index = 1:n2d
            p_ind = hash["filt_index"][index]
            p_filt = hash["filt_value"][index]
            f_ind    = k_ind + (index-1)*K
            @views filt_tmp = p_filt.*F_z[p_ind,kmask]
            ind = findall(filt_tmp .> 1E-13)
            val = filt_tmp[ind]
            map!(x->CartesianIndex(p_ind[x[1]][1],p_ind[x[1]][2],x[2]),ind,filtind[f_ind])
            filtval[f_ind] = val
            #psi_index[j_ind,l+1] = f_ind
        end
    end

    # -------- metadata dictionary
    info=Dict()
    info["2d_npix"]         = hash["npix"]
    info["2d_j_value"]      = hash["npix"]
    info["2d_theta_value"]  = hash["theta_value"]
    info["2d_psi_index"]    = hash["psi_index"]
    info["2d_phi_index"]    = hash["phi_index"]
    info["2d_pc"]           = hash["pc"]
    info["2d_wd"]           = hash["wd"]
    info["2d_fs_center_r"]  = hash["fs_center_r"]
    info["2d_filt_index"]   = hash["filt_index"]
    info["2d_filt_value"]   = hash["filt_value"]

    info["nz"]              = nz
    info["cz"]              = cz
    info["k_value"]         = k_value
    info["filt_index"]      = filtind
    info["filt_value"]      = filtval

    return info
end

filt_3d = fink_filter_bank_3dizer(hash, 1, nz=256)

@time hash = fink_filter_hash(1, 8, nx=256, pc=1, wd=2)
@benchmark filt_3d = fink_filter_bank_3dizer(hash, 1, nz=256)

filt_3d[1].*filt_3d[2]

filt_3d

temp = map(x->[x[1],x[2]],filt_3d[1])

temp2 = collect(Base.product(temp, filt_3d[2]))
temp3 = [CartesianIndex(x[1][1],x[1][2],x[2]) for x in temp2]

temp4 = vec(temp3)

filt_3d["2d_filt_index"][2]

filt_3d["filt_index"]

function plot_filter_bank_QA(hash; fname="filter_bank_QA.png")

    # -------- define plot1 to append plot to a list
    function plot1(ps, image; clim=nothing, bin=1.0, fsz=12, label=nothing)
        # ps    - list of plots to append to
        # image - image to heatmap
        # clim  - tuple of color limits (min, max)
        # bin   - bin down by this factor, centered on nx/2+1
        # fsz   - font size
        marg   = -1mm
        nx, ny = size(image)
        nxb    = nx/round(Integer, 2*bin)

        # -------- center on nx/2+1
        i0 = max(1,round(Integer, (nx/2+2)-nxb-1))
        i1 = min(nx,round(Integer, (nx/2)+nxb+1))
        lims = [i0,i1]
        subim = image[i0:i1,i0:i1]
        push!(ps, heatmap(image,aspect_ratio=:equal,clim=clim,
            xlims=lims, ylims=lims, size=(400,400),
            legend=false,xtickfontsize=fsz,ytickfontsize=fsz,#tick_direction=:out,
            rightmargin=marg,leftmargin=marg,topmargin=marg,bottommargin=marg))
        if label != nothing
            annotate!(lims'*[.96,.04],lims'*[.09,.91],text(label,:left,:white,16))
        end
        return
    end

    # -------- initialize array of plots
    ps   = []
    jval = hash["2d_j_value"]
    J    = length(jval)
    L    = length(hash["2d_theta_value"])
    nx   = hash["2d_npix"]
    nz   = hash["nz"]
    # plot the first 3 and last 3 j values
    jplt = [1:3;J-2:J]

    ind = zeros(Int64,J,L)
    for j_ind=1:J
        for l=1:L
            ind[j_ind,l] = Int(j_ind-1 + (l-1)*J)
        end
    end
    ind .+= 1

    # -------- loop over j
    for j_ind in jplt
        # first 3 filters in Fourier space
        j = jval[j_ind]
        bfac = max(1,2.0^(j-1))
        for l = 1:3
            label=string("j=",j_ind," ℓ=",l-1)
            #label=latexstring("\\tilde{\\psi}_{$j_ind,$(l-1)}")
            filt = zeros(nx,nx,nz)
            println(nx)
            println(nz)
            filt[hash["2d_filt_index"][2],hash["filt_index"][2]] = hash["filt_value"][2]
            temp = dropdims(sum(filt,dims=3))
            plot1(ps, fftshift(filt[:,:]), clim=(0,1),
                  bin=bfac, label=label)
        end
        fac = max(1,2.0^(5-j))
        # real part of config space of 3rd filter
        plot1(ps, fftshift(real(ifft(filt[:,:,ind[j_ind,3]]))),
              bin=fac, label=string("j=",j_ind," ℓ=",3-1))

        # sum all filters to form (partial) ring
        ring = dropdims(sum(filt[:,:,ind[j_ind,:]].^2, dims=3), dims=3)
        plot1(ps, fftshift(ring[:,:,1]), clim=(0,1),
              bin=bfac, label=string("j=",j_ind," ℓ=",0,":",L-1))
    end

    # wavepow = fftshift(dropdims(sum(filt[:,:,ind].^2, dims=(3,4)), dims=(3,4)))
    # plot1(ps, wavepow, clim=(-0.1,1), label=string("j=1:",J," ℓ=0:",L-1))
    #
    # if info["pc"]==1
    #     wavepow += circshift(wavepow[end:-1:1,end:-1:1],(1,1))
    # end
    # plot1(ps, wavepow, clim=(-0.1,1), bin=32, label="missing power")
    # phi = filt[:,:,info["phi_index"]]
    # phi_shift = fftshift(phi)
    # plot1(ps, phi_shift.^2, clim=(-0.1,1), bin=32, label="ϕ")
    # disc = wavepow+(phi_shift.^2)
    # plot1(ps, fftshift(real(ifft(phi))), label="ϕ")
    # plot1(ps, disc, clim=(-0.1,1), bin=32, label="all")

    myplot = plot(ps..., layout=(7,5), size=(1400,2000))
    savefig(myplot, fname)
end

plot_filter_bank_QA(filt_3d, fname="/Users/saydjari/Dropbox/GradSchool_AKS/Doug/Projects/DHC/scratch_AKS/filt-3d.png")

Profile.clear()
@profile filt_3d = fink_filter_bank_3dizer(filter_hash, 1, nz=256)

Juno.profiler()

@benchmark zeros(256,256,256)
@benchmark zeros(128,128,128)

## Ok, so lets try to have a good workflow

filter_hash = fink_filter_hash(1, 8, nx=256, pc=1, wd=2)

function fink_filter_bank_3dizer(hash, cz; nz=256)
    # -------- set parameters
    nx = convert(Int64,hash["npix"])
    n2d = size(hash["filt_value"])[1]
    dz = nz/2-1
    J = size(hash["j_value"])[1]
    L = size(hash["theta_value"])[1]

    im_scale = convert(Int8,log2(nz))
    # -------- number of bins in radial direction (size scales)
    K = (im_scale-2)*cz

    # -------- allocate output array of zeros
    J_L_K     = zeros(Int32, n2d*K,3)
    psi_index = zeros(Int32, J+1, L, K)
    k_value   = zeros(Float64, K)
    logr      = zeros(Float64,nz)
    filt_temp = zeros(Float64,nx,nx,nz)
    F_z       = zeros(Float64,nx,nx,nz)
    #F_p       = zeros(Float64,nx,nx,nz)

    filtind = fill(CartesianIndex{3}[], K*n2d)
    filtval = fill(Float64[], K*n2d)

    for z = 1:convert(Int64,dz+1)
        sz = mod(z+dz,nz)-dz-1
        z2 = sz^2
        logr[z] = 0.5*log2(max(1,z2))
    end

    @inbounds for k_ind = 1:K
        k = k_ind/cz
        k_value[k_ind] = k  # store for later
        krad  = im_scale-k-1
        Δk    = abs.(logr.-krad)
        kmask = (Δk .<= 1/cz)
        k_count = count(kmask)
        k_vals = findall(kmask)

        # -------- radial part
        #I have somehow dropped a factor of sqrt(2) which has been reinserted as a 0.5... fix my brain
        @views F_z[:,:,kmask].= reshape(1/sqrt(2).*cos.(Δk[kmask] .* (cz*π/2)),1,1,k_count)
        @inbounds for index = 1:n2d
            p_ind = hash["filt_index"][index]
            p_filt = hash["filt_value"][index]
            #F_p[p_ind,:] .= p_filt
            f_ind    = k_ind + (index-1)*K
            @views filt_tmp = p_filt.*F_z[p_ind,kmask]
            #temp_ind = findall((F_p .> 1E-13) .& (F_z .> 1E-13))
            #@views filt_tmp = F_p[temp_ind].*F_z[temp_ind]
            ind = findall(filt_tmp .> 1E-13)
            @views filtind[f_ind] = map(x->CartesianIndex(p_ind[x[1]][1],p_ind[x[1]][2],k_vals[x[2]]),ind)
            #filtind[f_ind] = temp_ind[ind]
            @views filtval[f_ind] = filt_tmp[ind]
            @views J_L_K[f_ind,:] = [hash["J_L"][index,1],hash["J_L"][index,2],k_ind-1]
            psi_index[hash["J_L"][index,1]+1,hash["J_L"][index,2]+1,k_ind] = f_ind
        end
    end

    # -------- metadata dictionary
    info=Dict()
    info["2d_npix"]         = hash["npix"]
    info["2d_j_value"]      = hash["j_value"]
    info["2d_theta_value"]  = hash["theta_value"]
    info["2d_psi_index"]    = hash["psi_index"]
    info["2d_phi_index"]    = hash["phi_index"]
    info["2d_J_L"]          = hash["J_L"]
    info["2d_pc"]           = hash["pc"]
    info["2d_wd"]           = hash["wd"]
    info["2d_fs_center_r"]  = hash["fs_center_r"]
    info["2d_filt_index"]   = hash["filt_index"]
    info["2d_filt_value"]   = hash["filt_value"]

    info["nz"]              = nz
    info["cz"]              = cz
    info["J_L_K"]           = psi_index
    info["psi_index"]       = psi_index
    info["k_value"]         = k_value
    info["filt_index"]      = filtind
    info["filt_value"]      = filtval

    return info
end

filt_3d = fink_filter_bank_3dizer(filter_hash, 1, nz=256)

@time filt_3d = fink_filter_bank_3dizer(filter_hash, 1, nz=256)

@benchmark filt_3d = fink_filter_bank_3dizer(filter_hash, 1, nz=256)

using Plots; gr()
using Measures
theme(:dark)

function plot_filter_bank_QAxy(hash; fname="filter_bank_QA.png")

    # -------- define plot1 to append plot to a list
    function plot1(ps, image; clim=nothing, bin=1.0, fsz=12, label=nothing,color=:white)
        # ps    - list of plots to append to
        # image - image to heatmap
        # clim  - tuple of color limits (min, max)
        # bin   - bin down by this factor, centered on nx/2+1
        # fsz   - font size
        marg   = -1mm
        nx, ny = size(image)
        nxb    = nx/round(Integer, 2*bin)

        # -------- center on nx/2+1
        i0 = max(1,round(Integer, (nx/2+2)-nxb-1))
        i1 = min(nx,round(Integer, (nx/2)+nxb+1))
        lims = [i0,i1]
        subim = image[i0:i1,i0:i1]
        push!(ps, heatmap(image,aspect_ratio=:equal,clim=clim,
            xlims=lims, ylims=lims, size=(400,400),
            legend=false,xtickfontsize=fsz,ytickfontsize=fsz,#tick_direction=:out,
            rightmargin=marg,leftmargin=marg,topmargin=marg,bottommargin=marg))
        if label != nothing
            annotate!(lims'*[.96,.04],lims'*[.09,.91],text(label,:left,color,16))
        end
        return
    end

    # -------- initialize array of plots
    ps   = []
    jval = hash["2d_j_value"]
    kval = hash["k_value"]
    J    = length(hash["2d_j_value"])
    L    = length(hash["2d_theta_value"])
    K    = length(hash["k_value"])
    nx   = hash["2d_npix"]
    nz   = hash["nz"]
    lup  = hash["psi_index"]
    index= hash["filt_index"]
    value= hash["filt_value"]
    # plot the first 3 and last 3 j values
    jplt = [1:3;J-2:J]
    lplt = 1:3
    kplt = [1,2,K]

    # -------- loop over k
    for k_ind in kplt
        # -------- loop over j
        k = kval[k_ind]
        krs = convert(Int64,(max(1,2.0^(7-k))+nz/2))
        kfs = convert(Int64,(max(1,2.0^(7-k))+1))
        for j_ind in jplt
            # first 3 filters in Fourier space
            j = jval[j_ind]
            bfac = max(1,2.0^(j-1))
            # -------- loop over l
            for l_ind in lplt
                label=string("j=",j_ind," ℓ=",l_ind-1," k=",k_ind)
                filt = zeros(nx,nx,nz)
                ind = lup[j_ind,l_ind,k_ind]
                filt[index[ind]] = value[ind]
                #temp = dropdims(sum(filt,dims=3),dims=3)
                temp = filt[:,:,kfs]
                plot1(ps, fftshift(temp[:,:]), clim=(0,1),
                      bin=bfac, label=label)
            end
            fac = max(1,2.0^(5-j))
            # real part of config space of 3rd filter
            l_ind = 3
            filt = zeros(nx,nx,nz)
            ind = lup[j_ind,l_ind,k_ind]
            filt[index[ind]] = value[ind]
            rs = fftshift(real(ifft(filt)))
            rs_slice = rs[:,:,krs]
            #rs_sum = dropdims(sum(rs,dims=3),dims=3)
            plot1(ps, rs_slice,
                  bin=fac, label=string("j=",j_ind," ℓ=",l_ind-1," k=",k_ind))

            # sum all filters to form (partial) ring
            ring = zeros(nx,nx)
            filt = zeros(nx,nx,nz)
            for indx in lup[j_ind,:,k_ind]
                filt[index[indx]] .+= value[indx].^2
            end
            ring = filt[:,:,kfs]
            plot1(ps, fftshift(ring[:,:]), clim=(0,1),
                  bin=bfac, label=string("j=",j_ind," ℓ=",0,":",L-1," k=",k_ind))
        end
    end

    filt = zeros(nx,nx,nz)
    for indx in lup[1:J,:,2]
        filt[index[indx]] .+= value[indx].^2
    end

    wavepow = fftshift(dropdims(sum(filt, dims=3), dims=3))
    plot1(ps, wavepow, clim=(-0.1,1), label=string("j=1:",J," ℓ=0:",L-1," K=1:",K),color=:black)

    if hash["2d_pc"]==1
        wavepow += circshift(wavepow[end:-1:1,end:-1:1],(1,1))
    end
    plot1(ps, wavepow, #clim=(0,K),
        bin=32, label="missing power",color=:black)

    phi_z = zeros(nx,nx,nz)
    for indx in lup[J+1,1,2]
        phi_z[index[indx]] .+= value[indx].^2
    end

    phi_rs = zeros(nx,nx,nz)
    for indx in lup[J+1,1,2]
        phi_rs[index[indx]] .+= value[indx]
    end
    phi_rst = fftshift(real(ifft(phi_rs)))[:,:,krs]

    k = kval[K]
    krs = convert(Int64,(max(1,2.0^(7-k))+nz/2))
    kfs = convert(Int64,(max(1,2.0^(7-k))+1))

    phi_shift = fftshift(dropdims(sum(phi_z, dims=3), dims=3))
    plot1(ps, phi_shift, #clim=(0,K),
        bin=32, label="ϕ")
    disc = wavepow.+phi_shift
    plot1(ps, phi_rst, label=string("ϕ k=",convert(Int64,k)))
    plot1(ps, disc, #clim=(0,K),
        bin=32, label="all", color=:black)

    myplot = plot(ps..., layout=(19,5), size=(3000,8000))
    savefig(myplot, fname)
    return phi_z
end

temp1 = plot_filter_bank_QAxy(filt_3d, fname="/Users/saydjari/Dropbox/GradSchool_AKS/Doug/Projects/DHC/scratch_AKS/images/filt-3d.png")

function plot_filter_bank_QAxz(hash; fname="filter_bank_QA.png")

    # -------- define plot1 to append plot to a list
    function plot1(ps, image; clim=nothing, binx=1.0, biny=1.0, fsz=12, label=nothing,color=:white)
        # ps    - list of plots to append to
        # image - image to heatmap
        # clim  - tuple of color limits (min, max)
        # bin   - bin down by this factor, centered on nx/2+1
        # fsz   - font size
        marg   = -1mm
        nx, ny = size(image)
        nxb    = nx/round(Integer, 2*binx)
        nyb    = ny/round(Integer, 2*biny)

        # -------- center on nx/2+1
        xi0 = max(1,round(Integer, (nx/2+2)-nxb-1))
        xi1 = min(nx,round(Integer, (nx/2)+nxb+1))
        xlims = [xi0,xi1]

        yi0 = max(1,round(Integer, (ny/2+2)-nyb-1))
        yi1 = min(ny,round(Integer, (ny/2)+nyb+1))
        ylims = [yi0,yi1]

        subim = image[xi0:xi1,yi0:yi1]
        push!(ps, heatmap(image,aspect_ratio=biny/binx,clim=clim,
            xlims=xlims, ylims=ylims, size=(400,400),
            legend=false,xtickfontsize=fsz,ytickfontsize=fsz,#tick_direction=:out,
            rightmargin=marg,leftmargin=marg,topmargin=marg,bottommargin=marg))
        if label != nothing
            annotate!(xlims'*[.96,.04],ylims'*[.09,.91],text(label,:left,color,16))
        end
        return
    end

    # -------- initialize array of plots
    ps   = []
    jval = hash["2d_j_value"]
    kval = hash["k_value"]
    J    = length(hash["2d_j_value"])
    L    = length(hash["2d_theta_value"])
    K    = length(hash["k_value"])
    nx   = hash["2d_npix"]
    nz   = hash["nz"]
    lup  = hash["psi_index"]
    index= hash["filt_index"]
    value= hash["filt_value"]
    # plot the first 3 and last 3 j values
    jplt = [1:3;J-2:J]
    lplt = [1,2,L]
    kplt = [1,2,K]

    # -------- loop over k
    for k_ind in kplt
        # -------- loop over j
        k = kval[k_ind]
        krs = convert(Int64,(max(1,2.0^(7-k))+nz/2))
        kfs = convert(Int64,(max(1,2.0^(7-k))+1))
        kfac = max(1,2.0^(k-1))
        for j_ind in jplt
            # first 3 filters in Fourier space
            j = jval[j_ind]
            jrs = convert(Int64,(max(1,2.0^(7-j))+nx/2))
            jfs = convert(Int64,(max(1,2.0^(7-j))+1))

            bfac = max(1,2.0^(j-1))
            # -------- loop over l
            for l_ind in lplt
                label=string("j=",j_ind," ℓ=",l_ind-1," k=",k_ind)
                filt = zeros(nx,nx,nz)
                ind = lup[j_ind,l_ind,k_ind]
                filt[index[ind]] = value[ind]
                #temp = dropdims(sum(filt,dims=2),dims=2)
                #temp = filt[:,nx÷2,:]
                plot1(ps, fftshift(filt)[nx÷2,:,:], #clim=(0,1),
                      binx=kfac, biny=bfac, label=label)
            end
            j_fac = max(1,2.0^(5-j))
            k_fac = max(1,2.0^(5-k))
            # real part of config space of 1st filter
            l_ind = 1

            filt = zeros(nx,nx,nz)
            ind = lup[j_ind,l_ind,k_ind]
            filt[index[ind]] = value[ind]
            rs = fftshift(real(ifft(filt)))
            rs_slice = rs[nx÷2,:,:]
            plot1(ps, rs_slice,
                  binx=k_fac, biny=j_fac, label=string("j=",j_ind," ℓ=",l_ind-1," k=",k_ind))

            filt = zeros(nx,nx,nz)
            ind = lup[j_ind,l_ind,k_ind]
            filt[index[ind]] = value[ind]
            rs = fftshift(real(ifft(filt)))
            rs_slice = rs[:,nx÷2,:]
            plot1(ps, rs_slice,
                binx=k_fac, biny=j_fac, label=string("j=",j_ind," ℓ=",l_ind-1," k=",k_ind))

            # sum all filters to form (partial) ring
            ring = zeros(nx,nx)
            filt = zeros(nx,nx,nz)
            for indx in lup[j_ind,:,k_ind]
                filt[index[indx]] .+= value[indx].^2
            end
            ring = filt
            plot1(ps, fftshift(ring)[nx÷2,:,:], #clim=(0,1),
                  binx=kfac, biny=bfac, label=string("j=",j_ind," ℓ=",0,":",L-1," k=",k_ind))
        end
    end

    filt = zeros(nx,nx,nz)
    for indx in lup[1:J,:,:]
        filt[index[indx]] .+= value[indx].^2
    end

    wavepow = fftshift(filt)
    plot1(ps, wavepow[:,nx÷2,:], #clim=(-0.1,1),
    label=string("j=1:",J," ℓ=0:",L-1," K=1:",K))

    wavepow2 = wavepow[:,nx÷2,:]
    if hash["2d_pc"]==1
        wavepow2 += reverse(wavepow[:,nx÷2,:], dims = 2)
        wavepow2 += circshift(wavepow2[end:-1:1,end:-1:1],(1,1))
    end
    plot1(ps, wavepow2, #clim=(0,K),
        binx=16, label="missing power",color=:black)

    phi_z = zeros(nx,nx,nz)
    for indx in lup[J+1,1,:]
        phi_z[index[indx]] .+= value[indx].^2
    end

    phi_rs = zeros(nx,nx,nz)
    for indx in lup[J+1,1,:]
        phi_rs[index[indx]] .+= value[indx]
    end
    phi_rst = fftshift(real(ifft(phi_rs)))[:,:,krs]

    k = kval[K]
    krs = convert(Int64,(max(1,2.0^(7-k))+nz/2))
    kfs = convert(Int64,(max(1,2.0^(7-k))+1))

    phi_shift = fftshift(dropdims(sum(phi_z, dims=1), dims=1))
    plot1(ps, phi_shift, #clim=(0,K),
        binx=32, label="ϕ")
    disc = wavepow[:,nx÷2,:].+phi_shift
    plot1(ps, phi_rst, label=string("ϕ k=",convert(Int64,k)))
    plot1(ps, disc, #clim=(0,K),
        binx=32, label="all", color=:black)

    plot1(ps, disc, #clim=(0,K),
        binx=32, label="all", color=:black)

    myplot = plot(ps..., layout=(19,6), size=(3000,8000))
    savefig(myplot, fname)

    filt = zeros(nx,nx,nz)
    ind = lup[1,1,1]
    filt[index[ind]] = value[ind]

    return wavepow
end

temp1 = plot_filter_bank_QAxz(filt_3d, fname="/Users/saydjari/Dropbox/GradSchool_AKS/Doug/Projects/DHC/scratch_AKS/images/filt-3dxz.png")

nx = convert(Int64,filt_3d["2d_npix"])
nz = 256
lup  = filt_3d["psi_index"]
index= filt_3d["filt_index"]
value= filt_3d["filt_value"]

filt = zeros(nx,nx,nz)
ind = lup[1,1,1]
filt[index[ind]] = value[ind]
plot(dropdims(sum(filt[1,].^2,dims=(1,2)),dims=(1,2)),xlims=(0,128))

filt = zeros(nx,nx,nz)
filt[filt_3d[1]] = filt_3d[2]
plot(dropdims(sum(filt.^2,dims=(1,2)),dims=(1,2)),xlims=(0,128))

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
using Test
using ChipSort
using SIMD

include("testutils.jl")

@testset "Array sorting" begin

data = randa(Int32, 64)
@test chipsort_small!(data,Val(8),Val(8)) == sort(data)

data = randa(Int32, 2^13)
ref = sort(data)
@test chipsort_medium!(data, Val(8), Val(8), Val(128)) == ref

data = randa(Int32, 2^20)
ref = sort(data)
@test chipsort_large(data,Val(8),Val(32)) == ref

data = randa(Int32, 2^7)
ref = sort(data)
@test combsort!(data, 2^6) == ref

data = randa(Int32, 2^7)
ref = sort(data)
@test insertion_sort!(data) == ref

end

ext: jl


text:
export relatedness_posthoc

function sig_within(data::PopData, results::DataFrame, population::String, iterations::Int = 20000)
    # add extra columns of population names
    pop_names = select(data.meta, :name => :sample_1 ,:population => :pop_1)
    tmp_res = innerjoin(results, pop_names, on = :sample_1)
    select!(pop_names, :sample_1 => :sample_2,:pop_1 => :pop_2)
    tmp_res = innerjoin(tmp_res, pop_names, on = :sample_2)
    tmp_res.groups = collect(zip(tmp_res.pop_1, tmp_res.pop_2))
    tmp_res.group = tmp_res.pop_1 .* tmp_res.pop_2
    select!(tmp_res, Not([:sample_1, :sample_2, :n_loci, :pop_1, :pop_2]))
    size(tmp_res,1) == 0 && error("Samples in the relatedness results not found in the provided PopData.")
    estimators = Symbol.(names(tmp_res)[begin:end-2])

    # extract the values for within-population estimates
    within_coeff = [tmp_res[tmp_res.group .== population^2, i] for i in estimators]
    # get the mean of those values
    mean_within = mean.(within_coeff)
    # extract the values of the between-population estimates
    among_coeff = [tmp_res[(population .∈ tmp_res.groups) .& (tmp_res.group .!= population^2), i] for i in estimators]
    # get the mean of those values
    mean_among = mean.(among_coeff)
    # find the difference
    differ = mean_within - mean_among

    # setup bootstrap
    n_within = length(within_coeff[1])
    n_among = length(among_coeff[1])
    n_tot = n_within + n_among
    n_range = collect(1:n_tot)
    n_iter = iterations - 1

    bootstrapped = [Vector{Float64}(undef, n_iter) for i in 1:length(estimators)]
    p = Progress(iterations * length(estimators), dt = 1, color = :blue)

    for i in 1:length(estimators)
        est = estimators[i]
        iter_count = 0
        @sync @inbounds for j in 1:n_iter
            Base.Threads.@spawn begin
                iter_count += 1
                idx_within = sample(Xoroshiro128Star(), 1:n_tot, n_within, replace = false)
                idx_among = n_range[Not(idx_within)]

                mu_within = mean(results[:, est][idx_within])
                mu_among = mean(results[:, est][idx_among])

                bootstrapped[i][j] = mu_within - mu_among

                pair_text =  population * "  $(iter_count)" * "/" * "$(iterations)"
                ProgressMeter.next!(p; showvalues = [(:Population, pair_text), (:Method, string(est))])
            end
        end
    end
    [(sum(differ[i] .<= bootstrapped[i]) + 1) / iterations for i in 1:length(estimators)]
    #(sum(differ .<= bootstrapped) + 1) / (iterations)
end

"""
    relatedness_posthoc(::PopData, results::DataFrame; iterations::Int)

Performs a posthoc analysis using the resulting DataFrame or NamedTuple
from relatedness(). This analysis uses permutations to test if a population has
significantly higher within-population relatedess than between-population relatedness.
The `results` object must have been generated from the provided `PopData`. Use `iterations = `
to specify the number of iterations for the permutation tests (default = `20000`). **Recommended**
that you use `MultipleTesting.jl` to correct resulting P-values.

**Example**
```
julia> cats = @nancycats ;

julia> rel_out = relatedness(cats, method = [Ritland, Moran], iterations = 100);

julia> relatedness_posthoc(cats, rel_out)
17x3 DataFrame
 Row │ population  Ritland_P  Moran_P
     │ String      Float64    Float64
─────┼────────────────────────────────
   1 │ 1              5.0e-5   5.0e-5
   2 │ 2              5.0e-5   5.0e-5
   3 │ 3              5.0e-5   5.0e-5
   4 │ 4              5.0e-5   5.0e-5
   5 │ 5              5.0e-5   5.0e-5
   6 │ 6              5.0e-5   5.0e-5
   7 │ 7              5.0e-5   5.0e-5
   8 │ 8              5.0e-5   5.0e-5
   9 │ 9              5.0e-5   5.0e-5
  10 │ 10             5.0e-5   5.0e-5
  11 │ 11             5.0e-5   5.0e-5
  12 │ 12             5.0e-5   5.0e-5
  13 │ 13             5.0e-5   5.0e-5
  14 │ 14             5.0e-5   5.0e-5
  15 │ 15             5.0e-5   5.0e-5
  16 │ 16             5.0e-5   5.0e-5
  17 │ 17             5.0e-5   5.0e-5
```
"""
function relatedness_posthoc(data::PopData, results::DataFrame; iterations::Int = 20000)
    all_pops = unique(data.meta.population)
    estimators = Symbol.(names(results)[names(results) .∉ Ref(["sample_1", "sample_2", "n_loci"])] .* "_P")
    sigs = map(pop -> sig_within(data, results, pop, iterations), all_pops)

    DataFrame(
        :population => all_pops,
        [estimators[i] => getindex.(sigs,i) for i in 1:length(estimators)]...
    )
end

function relatedness_posthoc(data::PopData, results::NamedTuple; iterations::Int = 20000)
    estimators = keys(results)
    coeffs = [results[i][:, estimators[i]] for i in 1:length(estimators)]
    df = select(results[1], "sample_1", "sample_2", "n_loci")
    [df[:, estimators[i]] = coeffs[i] for i in 1:length(estimators)]
    relatedness_posthoc(data, df, iterations = iterations)
end


text:
function getallergens(lines::Vector{<:AbstractString})
    ingredients = Vector{Vector{String}}()
    allergens = Dict{String,Vector{String}}()
    for line in lines
        ingr, aller = split(line[1:end - 1], " (contains ")
        ingr_list = split(ingr, " ")
        push!(ingredients, ingr_list)
        aller_list = split(aller, ", ")
        for al in aller_list
            allergens[al] = intersect(get(allergens, al, ingr_list), ingr_list)
        end
    end
    return allergens, ingredients
end

function countsafe(
    allergens::Dict{String,Vector{String}},
    ingredients::Vector{Vector{String}})
    unsafe = union(values(allergens)...)
    all_ingredients = union(ingredients...)
    safe = setdiff(all_ingredients, unsafe)
    total = 0
    for ingr_list in ingredients
        total += count(x -> x ∈ safe, ingr_list)
    end
    return total
end

const data = readlines("day21.txt")
const allergens, ingredients = getallergens(data)

println("Part 1: $(countsafe(allergens, ingredients))")

function next_identified(allergens::Dict{String,Vector{String}})
    sure = Tuple{String,String}[]
    for (aller, ing) in filter(x -> length(x[2]) == 1, collect(allergens))
        push!(sure, (aller, ing[1]))
        for (k, v) in allergens
            allergens[k] = setdiff(v, ing)
        end
    end
    return sure
end

function findmap(allergens::Dict{String,Vector{String}})
    known = []
    m = next_identified(allergens)
    while !(isempty(m))
        push!(known, m)
        m = next_identified(allergens)
    end
    return collect(Iterators.flatten(known))
end

function part2(allergens)
    known = findmap(allergens)
    return join((x[2] for x in sort(known)), ",")
end

println("Part 2: $(part2(allergens))")


text:
"""
    DiscreteSampler{T}

Efficient ``O(\\log_2 n)`` sampler from discrete probability distributions.
It is based on a custom binary heap with the following properties:
* each value is proportional to the probability of corresponding element,
* each root has the sum of its children's values,
  ``w_i = w_{2i} + w_{2i + 1}``,
* real indices of elements (leaves) can be found using `findindex`.
"""
struct DiscreteSampler{T}
    values :: Vector{T}
end

function push!(d::DiscreteSampler{T}, value) where T
    x = T(value)
    if length(d.values) == 0
        push!(d.values, x)
    else
        root = length(d.values) >> 1 + 1
        push!(d.values, d.values[root])
        push!(d.values, x)
        while root != 0
            d.values[root] += x
            root >>= 1
        end
    end
end

function pop!(d::DiscreteSampler)
    x = pop!(d.values)
    pop!(d.values)
    root = length(d.values) >> 1
    while root != 0
        d.values[root] -= x
        root >>= 1
    end
    return x
end

function findindex(d::DiscreteSampler, i::Integer)
    j = i * 2 - 1
    while true
        j_ = j << 1
        if j_ <= length(d.values)
            j = j_
        else
            break
        end
    end
    return j
end

function inc_index!(d::DiscreteSampler{T}, i::Integer) where T
    root = findindex(d, i)
    while root != 0
        d.values[root] += one(T)
        root >>= 1
    end
end

function dec_index!(d::DiscreteSampler{T}, i::Integer) where T
    root = findindex(d, i)
    while root != 0
        d.values[root] -= one(T)
        root >>= 1
    end
end

# Could not be `rand` to avoid clash with `Random` module, which had
# too many functions unneeded by `Junet` to build on.
function randd(d::DiscreteSampler; return_residual=false)
    x = rand(1:d.values[1])
    i = 1       # current node index
    j = 1       # index in original order
    while true
        i_ = i << 1
        i_ > length(d.values) && break
        if x <= d.values[i_]
            i = i_
        else
            j = i + 1
            i = i_ + 1
            x -= d.values[i_]
        end
    end
    return return_residual ? (j, x) : j
end


text:
# collection of functions to create movies/gifs/figures of quantities
# interpolated from 2D onto a line to make interpretation easier
using FileIO, Plots

function make_gif_of_files_in(file_folder::String,firstInt::Int,lastInt::Int)
    anim = @animate for j=firstInt:lastInt
        filename = string(file_folder,"/interpVals_$(j).jld2")
        time, interpVals = load(filename, "time", "interpVals")
        N,M = size(interpVals)

        nameArr = ["x","y","psi_A","psi_B","theta_s","|v|"]
        pL = plot()
        pR = plot()
        printTime = round(time,digits=2)
        for i=3:4
            plot!(pL,interpVals[:,2],interpVals[:,i],label=nameArr[i],title="t = $(printTime)",ylims=(0,0.005))
            plot!(pR,interpVals[:,2],interpVals[:,i+2],label=nameArr[i+2],title="t = $(printTime)",ylims=(0,1.4))
        end
        l = @layout [a{0.5w} b{0.5w}]
        plot(pL,pR,layout=l)
    end every 1

    gif(anim, "fpstest.gif", fps = 15)

    nothing
end


text:
####
#### Apply Boxplot
####

function apply_drawing!(g::GLE, bp::Boxplot, el_cntr::Int, figid::String, axidx::Int)
    # 1. add boxplot subroutines (vertical and horizontal) if not there already
    subname = ifelse(bp.horiz, "bp_horiz", "bp_vert")
    f = Figure(figid; _noreset=true)
    if subname ∉ keys(f.subroutines)
        f.subroutines[subname] = GLE_DRAW_SUB[subname]
    end

    # draw the boxplots one by one
    for k ∈ 1:bp.nobj
        # 1. retrieve the statistics
        wlow, q25, q50, q75, whigh, mean = bp.stats[k, :]

        # 2. call the subroutine
        s = bp.boxstyles[k]
        "\n\tdraw $subname" |> g
        "$k $wlow $q25 $q50 $q75 $whigh $mean"                          |> g
        # box styling
        sbl = s.blstyle
        "$(s.bwidth) $(s.wwidth) $(sbl.lstyle) $(sbl.lwidth) \"$(col2str(sbl.color))\"" |> g
        # median line
        sml = s.mlstyle
        "$(sml.lstyle) $(sml.lwidth) \"$(col2str(sml.color))\"" |> g
        # mean marker NOTE: the marker size drawn by itself and in a plot are not scaled the same
        # way, to make them scale the same way there is the division by textstyle.hei.
        sm = s.mmstyle
        "$(Int(s.mshow)) $(sm.marker) $(sm.msize/f.textstyle.hei) \"$(col2str(sm.color))\"" |> g

        el_cntr += 1
    end
    return el_cntr
end

####
#### Apply Heatmap
####

function apply_drawing!(g::GLE, hm::Heatmap, el_cntr::Int, figid::String, axidx::Int)
    # 1. apply the subroutine with a hash depending on figid and axidx
    hashid = hash((figid, axidx))
    add_sub_heatmap!(Figure(figid; _noreset=true), hm, hashid)

    # 2. write the zfile
    faux = auxpath(hash(hm.data), figid, axidx)
    isfile(faux) || csv_writer(faux, hm.transpose ? hm.data' : hm.data, false)

    nrows, ncols = size(hm.data)
    bw, bh = inv.((ncols, nrows))
    hm.transpose && (tmp = bw; bw = bh; bh = bw)

    nct = ifelse(hm.transpose, nrows, ncols)

    # 3. load data
    vs = prod("d$j=c0,c$j " for j ∈ 1:nct)
    "\n\tdata \"$faux\" $vs" |> g

    # 4. go over all the columns and draw the boxes (scales linearly with # cols)
    for j ∈ 1:nct
        "\n\tdraw hm_$hashid $j d$j $bw $bh" |> g
    end
    return el_cntr + 1
end


text:
abstract AbstractResult

abstract SpectralResult <: AbstractResult

typealias Projection{T <: Real} AbstractMatrix{T}

type Diffmap{T <: Real}
    d::Int
    t::Int
    K::AbstractMatrix{T}
    Y::AbstractMatrix{T}
end

function Base.show(io::IO, res::Diffmap)
    println(io, "Dimensionality:")
    show(io, res.d)
    print(io, "\n\n")
    println(io, "Timesteps:")
    show(io, res.t)
    print(io, "\n\n")
    println(io, "Kernel:")
    show(io, res.K)
    print(io, "\n\n")
    println(io, "Embedding:")
    show(io, res.Y)
end


text:
using ModelingToolkit, OrdinaryDiffEq, Test
using ModelingToolkit: topsort_equations

@variables t x(t) y(t) z(t) k(t)
eqs = [
       x ~ y + z
       z ~ 2
       y ~ 2z + k
      ]

sorted_eq = topsort_equations(eqs, [x, y, z, k])

ref_eq = [
          z ~ 2
          y ~ 2z + k
          x ~ y + z
         ]
@test ref_eq == sorted_eq

@test_throws ArgumentError topsort_equations([
                                             x ~ y + z
                                             z ~ 2
                                             y ~ 2z + x
                                            ], [x, y, z, k])

@parameters t σ ρ β
@variables x(t) y(t) z(t) a(t) u(t) F(t)
D = Differential(t)

test_equal(a, b) = @test isequal(simplify(a, polynorm=true), simplify(b, polynorm=true))

eqs = [D(x) ~ σ*(y-x),
       D(y) ~ x*(ρ-z)-y + β,
       0 ~ sin(z) - x + y,
       sin(u) ~ x + y,
       2β ~ 2,
       x ~ a,
      ]

lorenz1 = ODESystem(eqs,t,[u,x,y,z,a],[σ,ρ,β],name=:lorenz1)

lorenz1_aliased = alias_elimination(lorenz1)
reduced_eqs = [
               D(x) ~ σ * (y - x),
               D(y) ~ x*(ρ-z)-y + 1,
               0 ~ sin(z) - x + y,
               0 ~ x + y - sin(u),
              ]
test_equal.(equations(lorenz1_aliased), reduced_eqs)
test_equal.(states(lorenz1_aliased), [u, x, y, z])
test_equal.(observed(lorenz1_aliased), [
    β ~ 1,
    a ~ x,
])

# Multi-System Reduction

eqs1 = [
        D(x) ~ σ*(y-x) + F,
        D(y) ~ x*(ρ-z)-u,
        D(z) ~ x*y - β*z,
        u ~ x + y - z,
       ]

lorenz1 = ODESystem(eqs1,pins=[F],name=:lorenz1)

eqs2 = [
        D(x) ~ F,
        D(y) ~ x*(ρ-z)-x,
        D(z) ~ x*y - β*z,
        u ~ x - y - z
       ]

lorenz2 = ODESystem(eqs2,pins=[F],name=:lorenz2)

connected = ODESystem([lorenz2.y ~ a + lorenz1.x,
                       lorenz1.F ~ lorenz2.u,
                       lorenz2.F ~ lorenz1.u],t,[a],[],systems=[lorenz1,lorenz2])

# Reduced Flattened System

flattened_system = ModelingToolkit.flatten(connected)

aliased_flattened_system = alias_elimination(flattened_system)

@test isequal(states(aliased_flattened_system), [
        a
        lorenz1.x
        lorenz1.y
        lorenz1.z
        lorenz2.x
        lorenz2.y
        lorenz2.z
       ])

@test setdiff(parameters(aliased_flattened_system), [
        lorenz1.σ
        lorenz1.ρ
        lorenz1.β
        lorenz1.F
        lorenz2.F
        lorenz2.ρ
        lorenz2.β
       ]) |> isempty

reduced_eqs = [
               0 ~ a + lorenz1.x - lorenz2.y, # irreducible by alias elimination
               D(lorenz1.x) ~ lorenz1.σ*(lorenz1.y-lorenz1.x) + lorenz2.x - lorenz2.y - lorenz2.z,
               D(lorenz1.y) ~ lorenz1.x*(lorenz1.ρ-lorenz1.z)-(lorenz1.x + lorenz1.y - lorenz1.z),
               D(lorenz1.z) ~ lorenz1.x*lorenz1.y - lorenz1.β*lorenz1.z,
               D(lorenz2.x) ~ lorenz1.x + lorenz1.y - lorenz1.z,
               D(lorenz2.y) ~ lorenz2.x*(lorenz2.ρ-lorenz2.z)-lorenz2.x,
               D(lorenz2.z) ~ lorenz2.x*lorenz2.y - lorenz2.β*lorenz2.z
              ]
test_equal.(equations(aliased_flattened_system), reduced_eqs)

observed_eqs = [
                lorenz1.F ~ lorenz2.u,
                lorenz2.F ~ lorenz1.u,
                lorenz1.u ~ lorenz1.x + lorenz1.y - lorenz1.z,
                lorenz2.u ~ lorenz2.x - lorenz2.y - lorenz2.z,
               ]
test_equal.(observed(aliased_flattened_system), observed_eqs)

# issue #578

let
    @variables t x(t) y(t) z(t)
    D = Differential(t)
    eqs = [
        D(x) ~ x + y
        x ~ y
    ]
    sys = ODESystem(eqs, t, [x], [])
    asys = alias_elimination(flatten(sys))

    test_equal.(asys.eqs, [D(x) ~ 2x])
    test_equal.(asys.observed, [y ~ x])
end

# issue #724 and #716
let
    @parameters t
    D = Differential(t)
    @variables x(t), u(t), y(t)
    @parameters a, b, c, d
    ol = ODESystem([D(x) ~ a * x + b * u; y ~ c * x + d * u], t, pins=[u], name=:ol)
    @variables u_c(t), y_c(t)
    @parameters k_P
    pc = ODESystem(Equation[u_c ~ k_P * y_c], t, pins=[y_c], name=:pc)
    connections = [
        ol.u ~ pc.u_c,
        pc.y_c ~ ol.y
    ]
    connected = ODESystem(connections, t, systems=[ol, pc])
    @test equations(connected) isa Vector{Equation}
    sys = flatten(connected)
    reduced_sys = alias_elimination(sys)
    ref_eqs = [
               D(ol.x) ~ ol.a*ol.x + ol.b*pc.u_c
               0 ~ ol.c*ol.x + ol.d*pc.u_c - ol.y
               0 ~ pc.k_P*ol.y - pc.u_c
              ]
    @test ref_eqs == equations(reduced_sys)
end


text:
using sampleproject
using Test

@testset "test foo" begin
    @test foo() == "bar"
end

max_stars_repo_name: nluetts/sampleproject

max_issues_repo_name: nluetts/sampleproject

max_forks_repo_name: nluetts/sampleproject


text:
# <editor-fold desc= prepare and create exchange related variables"

# XXX prepare expansion and capacity variables for exchange
function prepareExcExpansion!(partExc::OthPart,partLim::OthPart,prepExc_dic::Dict{Symbol,NamedTuple},tsYear_dic::Dict{Int,Int},anyM::anyModel)

	# XXX determine dimensions of expansion variables (expansion for exchange capacities is NOT directed!)
	# get all possible dimensions of exchange
	potDim_df = DataFrame(map(x -> (lvlTs = x[2].tsExp, lvlR = x[2].rDis, C = x[1]), collect(anyM.cInfo)))
	tsLvl_dic = Dict(x => getfield.(getNodesLvl(anyM.sets[:Ts],x),:idx) for x in unique(potDim_df[!,:lvlTs]))
	rLvl_dic = Dict(x => getfield.(getNodesLvl(anyM.sets[:R],x),:idx) for x in unique(potDim_df[!,:lvlR]))

	potExc_df = flatten(flatten(flatten(combine(x -> (Ts_exp = map(y -> tsLvl_dic[y],x.lvlTs), R_a = map(y -> rLvl_dic[y],x.lvlR), R_b = map(y -> rLvl_dic[y],x.lvlR)), groupby(potDim_df,:C)),:Ts_exp),:R_a),:R_b)

	# get dimensions where exchange should actually be defined
	exExp_df = DataFrame(R_a = Int[], R_b = Int[], C = Int[])

	for excPar in intersect((:capaExcResi,:capaExcResiDir),keys(partExc.par))
		append!(exExp_df,matchSetParameter(potExc_df,partExc.par[excPar],anyM.sets)[!,[:R_a,:R_b,:C]])
	end

	for excPar in intersect((:capaExcUp, :capaExcLow, :capaExcFix, :expExcUp, :expExcLow, :expExcFix),keys(partLim.par))
		append!(exExp_df,matchSetParameter(potExc_df,partLim.par[excPar],anyM.sets)[!,[:R_a,:R_b,:C]])
	end

	# ensure expansion entries are not directed
	exExp_df = unique(exExp_df) |> (x -> filter(y -> y.R_a < y.R_b,vcat(x,rename(x,replace(namesSym(x),:R_a => :R_b, :R_b => :R_a))))) |> (z -> unique(z))

	# add supordiante timesteps of expansion
	allExExp_df = innerjoin(potExc_df,exExp_df, on = namesSym(exExp_df))
	allExExp_df[!,:Ts_expSup] = map(x -> getDescendants(x,anyM.sets[:Ts],false,anyM.supTs.lvl) |> (y -> typeof(y) == Array{Int,1} ? y : [y] ), allExExp_df[!,:Ts_exp])

	# filter cases where expansion is fixed to zero
	if :expExcFix in keys(anyM.parts.lim.par)
		allExExp_df = removeEntries([filterZero(allExExp_df,getLimPar(anyM.parts.lim,:expExcFix,anyM.sets[:Te]),anyM)],allExExp_df)
	end

	# filter cases where in and out region are the same
	filter!(x -> x.R_a != x.R_b, allExExp_df)

	# save result to dictionary for variable creation
	exp_df = addSupTsToExp(allExExp_df,partExc.par,:Exc,tsYear_dic,anyM)
	prepExc_dic[:expExc] = (var = convertExcCol(exp_df),ratio = DataFrame(), resi = DataFrame())

	return potExc_df
end

# XXX create exchange variables
function createExcVar!(partExc::OthPart,ts_dic::Dict{Tuple{Int,Int},Array{Int,1}},anyM::anyModel)
	# XXX extend capacity variables to dispatch variables
	capa_df = partExc.var[:capaExc][!,Not([:var,:dir])] |> (x -> unique(vcat(x,rename(x,replace(namesSym(x),:R_from => :R_to, :R_to => :R_from)))))
	# replace orginal carrier with leafs
	capa_df = replCarLeafs(capa_df,anyM.sets[:C])

	cToLvl_dic = Dict(x => (anyM.cInfo[x].tsDis, anyM.cInfo[x].rDis) for x in unique(capa_df[!,:C]))
	capa_df[!,:lvlTs] = map(x -> cToLvl_dic[x][1],capa_df[!,:C])
	capa_df[!,:lvlR] = map(x -> cToLvl_dic[x][2],capa_df[!,:C])
	rExc_dic = Dict(x => anyM.sets[:R].nodes[x[1]].lvl != x[2] ? getDescendants(x[1],anyM.sets[:R],false,x[2]) : [x[1]]
																							for x in union([map(x -> (x[y],x.lvlR), eachrow(unique(capa_df[!,[y,:lvlR]]))) for y in (:R_from,:R_to)]...))

	capa_df[!,:R_from] = map(x -> rExc_dic[x.R_from,x.lvlR],eachrow(capa_df[!,[:R_from,:lvlR]]))
	capa_df[!,:R_to] = map(x -> rExc_dic[x.R_to,x.lvlR],eachrow(capa_df[!,[:R_to,:lvlR]]))
	capa_df = flatten(select(capa_df,Not(:lvlR)),:R_from); capa_df = unique(flatten(capa_df,:R_to))

	disp_df = combine(x -> (Ts_dis = ts_dic[(x.Ts_disSup[1],x.lvlTs[1])],),groupby(capa_df,namesSym(capa_df)))[!,Not(:lvlTs)]

	# filter entries where availability is zero
	if !isempty(partExc.par[:avaExc].data) && 0.0 in partExc.par[:avaExc].data[!,:val]
		disp_df = filter(x -> x.val != 0.0, matchSetParameter(disp_df,partExc.par[:avaExc],anyM.sets))[!,Not(:val)]
	end

	# computes value to scale up the global limit on dispatch variable that is provied per hour and create variables
	partExc.var[:exc] = orderDf(createVar(disp_df,"exc",getUpBound(disp_df,anyM.options.bound.disp / anyM.options.scaFac.dispExc,anyM.supTs,anyM.sets[:Ts]),anyM.optModel,anyM.lock,anyM.sets, scaFac = anyM.options.scaFac.dispExc))
end

# XXX add residual capacties for exchange (both symmetric and directed)
function addResidualCapaExc!(partExc::OthPart,prepExc_dic::Dict{Symbol,NamedTuple},potExc_df::DataFrame,anyM::anyModel)

	expSup_dic = Dict(x => getDescendants(x,anyM.sets[:Ts],true,anyM.supTs.lvl) for x in unique(potExc_df[!,:Ts_exp]))
	potExc_df[!,:Ts_disSup] = map(x -> expSup_dic[x],potExc_df[!,:Ts_exp])
	potExc_df = flatten(potExc_df[!,Not(:Ts_exp)],:Ts_disSup)

	# obtain symmetric residual capacites
	capaResi_df = filter(x -> x.R_a != x.R_b, checkResiCapa(:capaExc,potExc_df, partExc, anyM))
	sortR_mat = sort(hcat([capaResi_df[!,x] for x in (:R_a,:R_b)]...);dims = 2)
	for (index,col) in enumerate((:R_a,:R_b)) capaResi_df[!,col] = sortR_mat[:,index] end

	# manipulate entries in case directed residual capacities are defined
	if :capaExcResiDir in keys(partExc.par)

		directExc_df = matchSetParameter(potExc_df,partExc.par[:capaExcResiDir],anyM.sets)
		directExc_df[!,:var] = map(x -> AffExpr(x), directExc_df[!,:val]); select!(directExc_df,Not(:val))

		excDim_arr = [:C, :R_a, :R_b, :Ts_disSup]
		excDimP_arr = replace(excDim_arr,:R_a => :R_b, :R_b => :R_a)

		#  entries, where a directed capacity is provided and a symmetric one already exists
		bothExc_df = vcat(innerjoin(directExc_df, capaResi_df; on = excDim_arr, makeunique = true), innerjoin(directExc_df, capaResi_df; on = Pair.(excDim_arr,excDimP_arr), makeunique = true))
		bothExc_df = combine(x -> (var = x.var + x.var_1,), groupby(bothExc_df,excDim_arr))
		if !(:var in namesSym(bothExc_df)) bothExc_df[!,:var] = AffExpr[] end
		 # entries, where only a directed capacity was provided
		onlyDirExc_df = antijoin(directExc_df, bothExc_df; on = excDim_arr )

		# entries originally symmetric that now become directed, because a directed counterpart was introduced
		flipSym_df = antijoin(innerjoin(capaResi_df, bothExc_df[!,Not(:var)]; on = excDim_arr),bothExc_df[!,Not(:var)]; on = excDim_arr .=> excDimP_arr)
		swtExc_df = vcat(bothExc_df,flipSym_df)

		# solely directed entries
		dirExc_df = vcat(onlyDirExc_df,swtExc_df)
		dirExc_df[!,:dir] .= true

		# entries who become directed because their counterpart became directed
		becomDirExc_df = innerjoin(rename(dirExc_df[!,Not([:var,:dir])],:R_a => :R_b, :R_b => :R_a),vcat(capaResi_df,rename(capaResi_df,:R_a => :R_b, :R_b => :R_a)); on = excDim_arr)
		becomDirExc_df[!,:dir] .= true

		# entries entries originally symmetric that remain symmetric
		unDirExc_df = antijoin(capaResi_df, vcat(dirExc_df, rename(dirExc_df,:R_a => :R_b, :R_b => :R_a)); on = excDim_arr )
		unDirExc_df[!,:dir] .= false

		# adjust dataframe of residual capacities according to directed values
		capaResi_df = vcat(dirExc_df,vcat(unDirExc_df,becomDirExc_df))

		# adjust dataframe of capacities determining where variables will be created to reflect which of these correspond to directed cases now
		allVar_df = prepExc_dic[:capaExc].var
		if !isempty(prepExc_dic[:capaExc].var)
			undirBoth_df = vcat(dirExc_df,rename(dirExc_df,replace(namesSym(dirExc_df),:R_a => :R_b, :R_b => :R_a)))[!,Not(:dir)]
			dirVar_df = convertExcCol(innerjoin(convertExcCol(allVar_df[!,Not(:dir)]), vcat(undirBoth_df,swtExc_df)[!,Not(:var)],on = excDim_arr))
			dirVar_df[!,:dir] .= true
			adjVar_df = vcat(dirVar_df,antijoin(allVar_df,dirVar_df,on = [:C, :R_from, :R_to, :Ts_disSup] ))
		else
			adjVar_df = allVar_df
		end
	else
		capaResi_df[!,:dir] .= false
		adjVar_df = prepExc_dic[:capaExc].var
	end

	# filter cases where in and out region are the same
	filter!(x -> x.R_a != x.R_b, capaResi_df)

	# adjust dictionary accordingly
	capaResi_df[!,:Ts_expSup] .= 0
	prepExc_dic[:capaExc] = (var = adjVar_df, ratio = DataFrame(), resi = convertExcCol(capaResi_df))
end

# XXX converts table where exchange regins are given as "R_a" and "R_b" to "R_to" and "R_from" and the other way around
convertExcCol(in_df::DataFrame) = rename(in_df, namesSym(in_df) |> (x  -> :R_a in x ? replace(x,:R_a => :R_from, :R_b => :R_to) : replace(x,:R_from => :R_a, :R_to => :R_b)))

# </editor-fold>

# <editor-fold desc= create exchange related constraints"

# XXX connect capacity and expansion constraints for exchange
function createCapaExcCns!(partExc::OthPart,anyM::anyModel)
	capaVar_df = rename(partExc.var[:capaExc],:var => :capaVar)
	if :expExc in keys(partExc.var)
		expVar_df = flatten(partExc.var[:expExc],:Ts_disSup)[!,Not(:Ts_exp)]

		cns_df = innerjoin(capaVar_df, combine(groupby(expVar_df,[:Ts_disSup, :R_from, :R_to, :C]), :var => (x -> sum(x)) => :expVar); on = [:Ts_disSup, :R_from, :R_to, :C])

		# prepare, scale and create constraints
		cns_df[!,:cnsExpr] = map(x -> x.capaVar - x.capaVar.constant - x.expVar, eachrow(cns_df))
		cns_df = intCol(cns_df,:dir) |> (x -> orderDf(cns_df[!,[x...,:cnsExpr]]))
		scaleCnsExpr!(cns_df,anyM.options.coefRng,anyM.options.checkRng)
		partExc.cns[:excCapa] = createCns(cnsCont(cns_df,:equal),anyM.optModel)
	end

	# create and control operated capacity variables
	if anyM.options.decomm != :none && :capaExc in keys(partExc.var)
		# constraints for operated capacities are saved into a dictionary of containers and then actually created
		cns_dic = Dict{Symbol,cnsCont}()
		createOprVarCns!(partExc,cns_dic,anyM)
		for cnsSym in keys(cns_dic)
			scaleCnsExpr!(cns_dic[cnsSym].data,anyM.options.coefRng,anyM.options.checkRng)
			partExc.cns[cnsSym] = createCns(cns_dic[cnsSym],anyM.optModel)
		end
	end
end

# XXX create capacity restriction for exchange
function createRestrExc!(ts_dic::Dict{Tuple{Int64,Int64},Array{Int64,1}},partExc::OthPart,anyM::anyModel)

	# group exchange capacities by carrier
	grpCapa_df = groupby(rename(partExc.var[anyM.options.decomm != :none ? :oprCapaExc : :capaExc],:var => :capa),:C)

	# pre-allocate array of dataframes for restrictions
	restr_arr = Array{DataFrame}(undef,length(grpCapa_df))
	itrRestr = collect(enumerate(grpCapa_df))

	# create restrictions
	@threads for x in itrRestr
		restr_arr[x[1]] = prepareRestrExc(copy(x[2]),ts_dic,partExc,anyM)
	end

	anyM.parts.exc.cns[:excRestr] = createCns(cnsCont(vcat(restr_arr...),:smaller),anyM.optModel)
end

# XXX prepare capacity restriction for specific carrier
function prepareRestrExc(cns_df::DataFrame,ts_dic::Dict{Tuple{Int64,Int64},Array{Int64,1}},partExc::OthPart,anyM::anyModel)

	c_int = cns_df[1,:C]
	leafes_arr = filter(y -> isempty(anyM.sets[:C].nodes[y].down), [c_int,getDescendants(c_int,anyM.sets[:C],true)...])
	cRes_tup = anyM.cInfo[c_int] |> (x -> (Ts_dis = x.tsDis, R_from = x.rDis, R_to = x.rDis))

	# extend constraint dataframe to dispatch levels
	cns_df[!,:Ts_dis] .= map(x -> ts_dic[x,cRes_tup.Ts_dis],cns_df[!,:Ts_disSup])
	cns_df = flatten(cns_df,:Ts_dis)

	# resize capacity variables
	cns_df[!,:capa]  = cns_df[!,:capa] .* map(x -> anyM.supTs.sca[(x,cRes_tup.Ts_dis)], cns_df[!,:Ts_disSup])

	# filter relevant dispatch variables
	relDisp_df = filter(x -> x.C in leafes_arr, partExc.var[:exc])

	# first aggregate symmetric and directed entries in one direction, then directed entries in the other direction
	cns_df[!,:disp] = aggUniVar(relDisp_df,cns_df,[:Ts_dis,:R_from,:R_to],cRes_tup,anyM.sets)
	dir_arr = findall(.!cns_df[!,:dir])
	cns_df[dir_arr,:disp] = cns_df[dir_arr,:disp] .+ aggUniVar(rename(relDisp_df,:R_from => :R_to, :R_to => :R_from),cns_df[dir_arr,:],[:Ts_dis,:R_from,:R_to],cRes_tup,anyM.sets)

	# add availablities to dataframe
	cns_df = matchSetParameter(convertExcCol(cns_df),partExc.par[:avaExc],anyM.sets, newCol = :avaSym)

	if :avaExcDir in keys(partExc.par)
		dirAva_df = matchSetParameter(cns_df[!,intCol(cns_df,:dir)],partExc.par[:avaExcDir],anyM.sets, newCol = :avaDir)
		cns_df = joinMissing(cns_df,dirAva_df,[:Ts_disSup,:Ts_dis,:R_a,:R_b,:C,:dir],:left, Dict(:avaDir => nothing))
	else
		cns_df[!,:avaDir] .= nothing
	end

	# prepare, scale and create constraints
	cns_df[!,:cnsExpr] = map(x -> x.disp  - x.capa * (isnothing(x.avaDir) ? x.avaSym : x.avaDir), eachrow(cns_df))
	scaleCnsExpr!(cns_df,anyM.options.coefRng,anyM.options.checkRng)

	return convertExcCol(cns_df) |> (x -> select(x,intCol(x,:cnsExpr)))
end

# </editor-fold>

# <editor-fold desc= utility functions for exchange"

# XXX obtain values for exchange losses
function getExcLosses(exc_df::DataFrame,excPar_dic::Dict{Symbol,ParElement},sets_dic::Dict{Symbol,Tree})
	lossPar_obj = copy(excPar_dic[:lossExc])
	if :R_a in namesSym(lossPar_obj.data) && :R_b in namesSym(lossPar_obj.data)
		lossPar_obj.data = lossPar_obj.data |> (x -> vcat(x,rename(x,:R_a => :R_b, :R_b => :R_a)))
	end
	excLoss_df = matchSetParameter(exc_df,lossPar_obj,sets_dic,newCol = :loss)

	# overwrite symmetric losses with any directed losses provided
	if :lossExcDir in keys(excPar_dic)
		oprCol_arr = intCol(excLoss_df)
		dirLoss_df = matchSetParameter(excLoss_df[!,oprCol_arr],excPar_dic[:lossExcDir],sets_dic,newCol = :lossDir)
		excLoss_df = joinMissing(excLoss_df,dirLoss_df,oprCol_arr,:left,Dict(:lossDir => nothing))
		excLoss_df[!,:loss] = map(x -> isnothing(x.lossDir) ? x.loss : x.lossDir,eachrow(excLoss_df[!,[:loss,:lossDir]]))
		select!(excLoss_df,Not(:lossDir))
	end

	return excLoss_df
end

# </editor-fold>


text:
include("learning_rate/alrsvi.jl")
include("learning_rate/inversedecay.jl")
include("analytic.jl")
include("analyticVI.jl")
include("gibbssampling.jl")
include("numericalVI.jl")

function post_process!(model::AbstractGP{T,<:Likelihood,<:Inference}) where {T}
    nothing
end

const GibbsorVI = Union{<:GibbsSampling,<:AnalyticVI}


text:
#
println("Ck) Test of the MultipolePolarizibility module with ASF from an internally generated multiplet.")
#
setDefaults("print summary: open", "zzz-MultipolePolarizibility.sum")
wa = Atomic.Computation(Atomic.Computation(), name="xx", grid=JAC.Radial.Grid(true), nuclearModel=Nuclear.Model(26.),
                        properties=[Polarity()],
                        configs=[Configuration("[Ne] 3s^2 3p^5"), Configuration("[Ne] 3s 3p^6")],
                        polaritySettings=MultipolePolarizibility.Settings(EmMultipole[], 0, 0, Float64[], false, false, Int64[]) )

wb = perform(wa)
setDefaults("print summary: close", "")


text:
"""
   randommask(n, Bernoulli(p)[ symmetric=true])

Generate a random Bernoulli mask of n × n.
"""
function randommask(n::Int, distribution::Bernoulli{T}; symmetric=true) where T
    n ≥ 2 || throw(DomainError(n, "n < 2"))

    if symmetric
        # Generates a random lower triangular matrix and then symmetrizes it
        return Symmetric([i[1] > i[2] ? convert(Int, rand(distribution)) : 1 for i in CartesianIndices((1:n,1:n))], :L)
    else
        # Generates random entries for all entries
        return [i[1] != i[2] ? convert(Int, rand(distribution)) : 1 for i in CartesianIndices((1:n,1:n))]
    end
end

"""
   randommask(D, Bernoulli(p)[ symmetric=true])

Generate a random Bernoulli mask of size(D, 1).
"""
function randommask(D::EDM, distribution::Bernoulli{T}; symmetric=true) where T
    return randommask(size(D, 1), distribution, symmetric=symmetric)
end

"""
    randommask(n::Int, number_of_deletions::Int)

Randomly and uniformly create a mask with 2 × "number_of_deletions" number of deletions. The
2 × multiplier is due to the deletions done only on the entries below the diagonal,
and then symmetrized.
"""
function randommask(n::Int, number_of_deletions::Int)
    n ≥ 2 || throw(DomainError(n, "n < 2"))
    number_of_deletions > 0 || throw(DomainError(number_of_deletions, "deletions < 1"))

    # We create a lower diagonal matrix from which we can sample CartesianIndices
    # uniformly at random without replacement using sample from StatsBase
    mask = UnitLowerTriangular(ones(Int, n, n)) - Matrix(I, n, n)
    indices = sample(findall(x -> x == 1, mask), number_of_deletions, replace=false)

    # We set the values of indices to zero, and symmetrize
    mask[indices] .= 0
    mask = mask + Matrix(I, n, n)

    return Symmetric(mask, :L)
end

"""
    randommask(D::EDM, number_of_deletions::Int)

Randomly and uniformly create a mask with 2 × "number_of_deletions" number of deletions for D.
The 2 × multiplier is due to the deletions done only on the entries below the diagonal,
and then symmetrized.
"""
function randommask(D::EDM, number_of_deletions::Int)
    return randommask(size(D, 1), number_of_deletions)
end

"""
    randommask(n::Int, fraction_of_deletions::Real)

Randomly and uniformly create a mask with n * (n - 1) * fraction_of_deletions number of deleted
elements.
"""
function randommask(n::Int, fraction_of_deletions::Float64)
    n ≥ 2 || throw(DomainError(n, "n < 2"))
    1 ≥ fraction_of_deletions > 0 || throw(DomainError(fraction_of_deletions, "We need 0 < fraction ≤ 1 of entries."))

    deletions = round(Int, n * (n-1) / 2 * fraction_of_deletions)
    return randommask(n, deletions)
end

"""
    randommask(D::EDM, fraction_of_deletions::Real)

Randomly and uniformly create a mask with n * (n - 1) * fraction_of_deletions number of deleted
elements.
"""
function randommask(D::EDM, fraction_of_deletions::Real)
    return randommask(size(D, 1), fraction_of_deletions)
end

"""
    idxmask(D::EDM, indices::CartesianIndices)

Generate a mask from CartesianIndices.

The indices along the diagonal are ignored.
"""
function idxmask(D::EDM, indices::CartesianIndices; symmetric::Bool=true)
    mask = ones(Int, size(D))

    # We set values to zero for all non-diagonal indices
    mask[[idx for idx in indices if idx[1] != idx[2]]] .= 0

    if symmetric && issymmetric(mask)
        return mask
    elseif symmetric && !issymmetric(mask)
        # If there are more zeros in the lower triangle
        if count(x -> x == 0, LowerTriangular(mask)) > count(x -> x == 0, UpperTriangular(mask))
            return Symmetric(mask, :L)
        else # It's not symmetric, so the zeros must be above the diagonal
            return Symmetric(mask, :U)
        end
    else
        return mask
    end

end


text:
export mdependentlift

"""
    mdependentlift(s::AbstractDiscreteSwitchedSystem, m, forward=true)

Returns the `m`-path-dependent lift of the hybrid system `s` [LD06]. The
corresponding automaton is the debruijn graph of dimension `m`. If `forward` is
`false` then the debruijn graph works backward [ELD14].

[LD06] Lee, J.-W. & Dullerud, G. E.
Uniform stabilization of discrete-time switched and Markovian jump linear systems
Automatica, Elsevier, 2006, 42, 205-21

[ELD14] Essick, R. and Lee, J.-W. and Dullerud, G. E.
Control of linear switched systems with receding horizon modal information
IEEE Transactions on Automatic Control, 2014, 59, 2340-2352

[PEDJ16] Philippe, M. and Essick, R. and Dullerud, G. E. and Jungers, R. M.
Stability of discrete-time switching systems with constrained switching sequences
Automatica, 2016, 72, 242-250
"""
function mdependentlift(s::AbstractDiscreteSwitchedSystem, n::Integer, forward=true)
    curstid = 0
    idmap = Dict{Int, Int}()
    statelabels = String[]
    for st in states(s)
        for sw in switchings(s, n, st, forward)
            id = stateid(s, sw.seq, 1:n)
            if !haskey(idmap, id)
                curstid += 1
                idmap[id] = curstid
                push!(statelabels, join(string.(sw.seq)))
            end
        end
    end
    G = LightAutomaton(curstid)
    for st in states(s)
        for sw in switchings(s, n+1, st, forward)
            seq = sw.seq
            src = stateid(s, seq, 1:n)
            dst = stateid(s, seq, 1 .+ (1:n))
            src_id = idmap[src]
            dst_id = idmap[dst]
            # The label is automatically the same so we just have to check
            # whether there is any transition between `src_id` and `dst_id`
            if !has_transition(G, src_id, dst_id)
                σ = forward ? symbol(s, seq[n+1]) : symbol(s, seq[1])
                add_transition!(G, src_id, dst_id, σ)
            end
        end
    end
    s = discreteswitchedsystem(map(rm -> rm.A, s.resetmaps), G)
    s.ext[:statelabels] = statelabels
    s
end

function stateid(s, seq, I)
    m = length(s.resetmaps)
    id = 0
    for i in I
        id = id * m + symbol(s, seq[i])-1
    end
    id + 1
end


text:
#=##############################################################################################
Filename: system.jl
Author: Ryan Anderson
Contact: rymanderson@gmail.com
README: this file defines the `Aircraft` system struct.
=###############################################################################################

"""
    Aircraft{w,r,n,s,m,b}

`Aircraft` system object.

Members are subsystems of the aircraft. Subsystems are not typed to provide flexibility as to what models are used. This
allows simulations of a wide range of interchangeable models. Functionality is defined in
instances of the `Action` types.

# Fields:

* `wing_system` : comprehensively defines the aircraft's wings
* `rotor_system` : comprehensively defines the aircraft's rotors
* `nonlifting_body_system` : comprehensively defines non-lifting body(s) composing the aircraft
* `structure_system` : comprehensively defines the aircraft's structure(s)
* `motor_system` : comprehensively defines the aircraft's motor(s)
* `battery_system` : comprehensively defines the aircraft's batterie(s)

"""
struct Aircraft{TF,TF1,TF2,TF3,TAF}
        # T1 <: Union{Nothing, CC.Mach_correction}, T2 <: Union{Nothing, CC.ReCorrection},
        # T3 <: Union{Nothing, CC.RotationCorrection}, T4 <: Union{Nothing, CC.TipCorrection}}
    wing_system::Union{VortexLatticeSystem{TF}, Nothing}
    rotor_system::CCBladeSystem{TF,TF1,TF2,TF3,TAF}#,T1,T2,T3,T4}, Nothing}
    inertia_system::Union{Inertia{TF}, Nothing}
    nonlifting_body_system::Nothing
    structure_system::Nothing
    motor_system::Nothing
    battery_system::Nothing
end

max_stars_repo_name: byuflowlab/AircraftSystems

max_issues_repo_name: byuflowlab/AircraftSystems

max_forks_repo_name: byuflowlab/AircraftSystems


text:
using DiskJockeyTests
using Test

##########################
# Project.toml management
##########################

# The goal of this package is to test the scripts environment of DiskJockey.
# Because there are a lot of environments floating around, it makes sense
# to describe them here.

# 1) DiskJockeyTests.jl/Project.toml : the basic description for this package.
# This package does nothing on its own, so this is essentially empty.

# 2) DiskJockeyTests/test/Project.toml : this is the testing environment activated
# when running `pkg> test`. The only packages we really need are Tests
# for reasons explained in the next item.

# 3) outputdir/Project.toml : The design of the DiskJockey package is such that the user
# creates a Julia environment to which they install DiskJockey.
# Then, the add the `scripts/` directory to their system PATH so that they can run these
# files easily.
# They should also set the environment variable `JULIA_PROJECT` to their environment
# directory. We could also set it to "@."
# The scripts that go along with DiskJockey use other packages for reading configuration
# files and for plotting. The user can either install these packages individually, or
# use the DJ_initialize_environment.jl scirpt to install all of them directly.
# This project directory will contain the data, config, and output pertaining to that run.

# The testing is designed to simulate this workflow.
# Check to see if an environment variable is defined. If so, save output there.
outputdir = get(ENV, "JL_TEST_TEMPDIR", nothing)
if isnothing(outputdir)
    # grab a temporary directory to save all output products there by default
    outputdir = mktempdir()
else
    # if it doesn't exist, make the directory
    if !ispath(outputdir)
        mkdir(outputdir)
    end
end

# println("Created $outputdir")

# copy the data file to it
cp("data.hdf5", joinpath(outputdir, "data.hdf5"))

# change to this directory
cd(outputdir)

# Create and activate a project environment
using Pkg
Pkg.activate(".")

# Make sure any script started (like those prefixed by `DJ_`) uses this Project.toml first
ENV["JULIA_PROJECT"] = outputdir

# Use the DJ_initialize_environment.jl script (which should be on the PATH)
# to initialize a new environment here.
run(`DJ_initialize_environment.jl`)

# test loading of a dataset
# test creation of an init structure

# @testset "DiskJockeyTests.jl" begin
#     # First, create the necessary input files
#     run(`DJ_initialize.jl --new-project=standard`)

#     # make plots of the structure
#     run(`DJ_plot_structure.jl`)

#     # run the synthesis
#     run(`DJ_synthesize_model.jl`)

#     # plot the channel maps
#     run(`DJ_plot_chmaps.jl`)

#     # Try downsampling the model to the baseline locations
#     run(`DJ_write_model.jl`)

# end

###################################
### Functional, integration tests
###################################

######################################
### Standard model, fixed distance

# Just make sure the directory is clean
# run(`bash clean.sh`)

# Make a fake model so we can test off of it.
# println("Initializing directory with a new standard project.")

run(`DJ_initialize.jl --new-project=standard`)

# Here test the fix parameters reading routines
include("parameters_fixed.jl")

# Set up the config file and pos0.npy
# println("Editing config.yaml to use the standard model with fixed distance.")
# run(`python edit_config.py --fix-distance --model=standard`)

# Run the initialization and plotting scripts via the make interface
run(`make all`)

# Run venus.jl and see if we sample properly. The idea is to get a simple test with a few walkers without running forever.
# run(`venus.jl --test`)

# Clean up before moving on to the next test
# run(`make clean`)


text:
using TrajectoryOptimization
using StaticArrays
using LinearAlgebra
using TrajOptPlots
using GeometryTypes
using MeshCat
const TO = TrajectoryOptimization
import TrajectoryOptimization.Controllers: RBState

# Model
Rot = UnitQuaternion{Float64,CayleyMap}
model = Dynamics.FreeBody{Rot,Float64}()
n,m = size(model)

if !isdefined(Main,:vis)
    vis = Visualizer()
    open(vis)
    set_mesh!(vis, model)
end

# Discretization
N = 101
tf = 5.0

# Initial and final conditions
X0 = zero(RBState)
Xf = RBState([1,1,1.], expm(deg2rad(70) * @SVector rand(3)), zeros(3), zeros(3))
x0 = Dynamics.build_state(model, X0)
xf = Dynamics.build_state(model, Xf)

# Objective
Qd = Dynamics.fill_state(model, 0e-0, 0e-1, 1e-2, 1e-2)
Rd = @SVector fill(1e-2, 6)
costfun = QuatLQRCost(Diagonal(Qd), Diagonal(Rd), xf, w=0.)
obj = Objective(costfun, N)

# Constraint
struct PointingConstraint <: TO.AbstractConstraint{Equality,State,1}
    p::SVector{3,Float64}
    b::SVector{3,Float64}
    PointingConstraint(p::SVector{3}, b=@SVector [0,0,1.]) = new(p,b)
end
TO.state_dim(::PointingConstraint) = 13

function TO.evaluate(con::PointingConstraint, x::SVector)
    q = UnitQuaternion(x[4],x[5],x[6],x[7])
    x3 = q*con.b
    r = @SVector [x[1],x[2],x[3]]
    dx = con.p - r
    return @SVector [x3'normalize(dx) - 1]
end

point = @SVector [1,-.5,1]
add_point!(vis, point)
con = PointingConstraint(point)
point2 = @SVector [-1,1,0.]
add_point!(vis, point2, name="point2")
con2 = PointingConstraint(point2, @SVector [1,0,0])

conSet = ConstraintSet(n,m,N)
add_constraint!(conSet, con, N:N)
add_constraint!(conSet, con2, N:N)

# Initialization
U0 = [@SVector zeros(6) for k = 1:N-1]

# Problem
prob = Problem(model, obj, xf, tf, x0=x0, constraints=conSet)

# Solve
solver = iLQRSolver(prob)
rollout!(solver)
cost(solver)
solver.opts.verbose = true
solve!(solver)

solver = AugmentedLagrangianSolver(prob)
initial_controls!(solver, U0)
rollout!(solver)
solver.opts.verbose = true
TO.update_constraints!(solver)
TO.constraint_jacobian!(solver)
max_violation(solver)
get_constraints(solver)
cost(solver)
evaluate(con, states(solver)[end])

solve!(solver)
err = Xf ⊖ RBState(model, states(solver)[end])
norm(err)

visualize!(vis, solver)

qf = Xf.q
e3 = @SVector [0,0,1.]
qf*e3

rotmat(qf)


text:
### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ dcab9a18-5d58-4632-a325-d16af9c79344
using Pkg, DrWatson

# ╔═╡ c95f3f02-f179-4f81-9c53-6e6b87f25613
begin
	@quickactivate "NumericalMethodsforEngineers"
	using NumericalMethodsforEngineers, Symbolics, SymbolicUtils
	using StaticArrays, DataFrames, Plots
end

# ╔═╡ 75361cb8-3c45-4475-94dc-bf8093114613
md" ##### Building a Lagrangian polynomial in expression form."

# ╔═╡ c1acd2f9-a762-4beb-b44a-feb74fba31fb
@variables x[1:4] y[1:4];

# ╔═╡ 15eaabdd-ed7b-42ec-bd19-96b927cbbdd8
@variables g(x::Vector{Symbolics.Num}, y::Vector{Symbolics.Num});

# ╔═╡ a240a3be-f978-4f27-985c-b997b8cb75e4
begin
	z = Vector{AbstractString}(undef, length(x))
	for i in 1:length(x)
		for j in 1:length(x)
			if j == 1 && i == 1
				z[i] = "$(y[1])"
			elseif j == 1
				z[i] = "((x - $(x[j])) / ($(x[i]) - $(x[j])))"
			elseif i == j
				z[i] = "$(z[i]) * $(y[i])"
			elseif i !== j
				z[i] = "$(z[i]) * ((x - $(x[j])) / ($(x[i]) - $(x[j])))"
			end
		end
	end
	z
end

# ╔═╡ 9a626dae-2a2e-4d17-9813-dcb1564848de
typeof(z)

# ╔═╡ db65cf7d-d33e-49cd-bd03-464b4096c8d6
ze = Meta.parse.(z)

# ╔═╡ fe586e50-f0e4-41e2-b7f0-f408c9038464
begin
	dct = Dict(x[1]=>1, x[2]=>2, x[3]=>3, x[4]=>4,
		y[1]=>1, y[2]=>4, y[3]=>9, y[4]=>16)
	L(dct)
end

# ╔═╡ 53d5c163-2d9a-4de8-9f37-1698aa142b69
substitute(ze[1], Dict(x[1]=>1, x[2]=>2, x[3]=>3, x[4]=>4, y[1]=>1, y[2]=>4, y[3]=>9, y[4]=>16))

# ╔═╡ 7147c00f-c9ea-4493-929b-8bf1f0354f28
substitute(ze[2], dct)

# ╔═╡ b3ee083c-9002-4273-aab1-b898348c2d0a
begin
	@variables u v
	to_compute = [u^2 + v, v^2 + u]
end;

# ╔═╡ 386da460-ee76-4506-8aaf-3d02658720f9
f_expr = build_function(to_compute, [u, v]);

# ╔═╡ 77c31df8-cab7-4bb8-986a-8d67b56b2b2b
begin
	myf = eval(f_expr[1])
	myf(SA[2.0, 3.0])
end

# ╔═╡ 162d30d7-d520-457e-8f1e-91ee85eb5573
ze_expr = build_function(ze, [x, y]);

# ╔═╡ a29d1a22-bdb9-4253-88ac-ed8e4ccb34e7
zef = eval(ze_expr[1])

# ╔═╡ 2fc2edef-203d-4822-b7a8-80a9c410ab01
zef(SA[1,2,3,4], SA[1,4, 9, 16])

# ╔═╡ Cell order:
# ╠═dcab9a18-5d58-4632-a325-d16af9c79344
# ╠═c95f3f02-f179-4f81-9c53-6e6b87f25613
# ╠═75361cb8-3c45-4475-94dc-bf8093114613
# ╠═c1acd2f9-a762-4beb-b44a-feb74fba31fb
# ╠═15eaabdd-ed7b-42ec-bd19-96b927cbbdd8
# ╠═a240a3be-f978-4f27-985c-b997b8cb75e4
# ╠═9a626dae-2a2e-4d17-9813-dcb1564848de
# ╠═db65cf7d-d33e-49cd-bd03-464b4096c8d6
# ╠═fe586e50-f0e4-41e2-b7f0-f408c9038464
# ╠═53d5c163-2d9a-4de8-9f37-1698aa142b69
# ╠═7147c00f-c9ea-4493-929b-8bf1f0354f28
# ╠═b3ee083c-9002-4273-aab1-b898348c2d0a
# ╠═386da460-ee76-4506-8aaf-3d02658720f9
# ╠═77c31df8-cab7-4bb8-986a-8d67b56b2b2b
# ╠═162d30d7-d520-457e-8f1e-91ee85eb5573
# ╠═a29d1a22-bdb9-4253-88ac-ed8e4ccb34e7
# ╠═2fc2edef-203d-4822-b7a8-80a9c410ab01


text:
"""
    reynolds(fl::AbstractFluid, L, U, Tc=20.0, Pa=101.325)

Calcula o número de Reynolds para um fluido `fl`, uma escala de comprimento `L` em m,
temperatura em °C e pressão em kPa.

```jldoctest
julia> reynolds(Air, 3e-3, 10)
1966.4892978721982

julia> reynolds(Air, 3e-3, 1)
196.64892978721983

julia> reynolds(Air, 2e-3, 1, 85.0, 93)
84.4083527295806
```
"""
reynolds(fl::AbstractFluid, L, U, Tc=20.0, Pa=101.325) = density(fl, Tc, Pa) * U*L/viscosity(fl, Tc)

"""
    nusphere(Re,Pr,μrat)

Número de Nusselt de uma esfera para número de Reynolds `Re`, número de Prandtl `Pr` e
razão de viscosidade na superfíce e ao longe.

O número de Nusselt permite calcular o coeficiente de convecção a partir da condutividade
térmica e um comprimento característico:

```math
Nu = h L/k
```

```jldoctest
julia> nusphere(100, 0.7)
6.589121562852225

julia> nusphere(100, 0.7, 1.3)
6.90021798681175
```

"""
nusphere(Re, Pr, μrat=1.0) = 2 + (0.4Re^0.5 + 0.06Re^0.6667)*Pr^0.4*μrat^0.25

"""
    nucylinder(Re,Pr)

Número de Nusselt de um cilindro para número de Reynolds `Re`, número de Prandtl `Pr`.

O número de Nusselt permite calcular o coeficiente de convecção a partir da condutividade
térmica e um comprimento característico:

```math
Nu = h L/k
```

```jldoctest
julia> nucylinder(3.0, 0.7)
1.0199902962373502
```

"""
nucylinder(Re, Pr) = 0.75 * Re^0.4 * Pr^0.37

"""
    grashof(Air, Ts, Ta, L, Pa, g)

Número de Grashof para ar seco, onde `Ts` é a temperatura do fluido em °C, `Ta` é a
temperatura ao longe, L é um comprimento em m, Pa a pressão em kPa e g é a aceleração
da gravidade em m/s²

O número de Grashof é usado em convecção natural.

```jldoctest
julia> grashof(Air, 80, 20, 2e-3, 101.325, 9.8)
68.9474295156832

julia> grashof(Air, 80, 20, 2e-1, 101.325, 9.8)
6.89474295156832e7
```
"""
function grashof(Air::AirType, Ts, Ta, L, Pa=101.325, g=9.80665)

    ν = kinviscosity(Air, Ta, Pa)
    β = volthermalexpansion(Air, Ta)

    return g * β * (Ts-Ta) * L^3 / (ν^2)
end

"""
    rayleigh(Air, Ts, Ta, L, Pa, g)

Número de Rayleigh para ar seco, onde `Ts` é a temperatura do fluido em °C, `Ta` é a
temperatura ao longe, L é um comprimento em m, Pa a pressão em kPa e g é a aceleração
da gravidade em m/s²

O número de Rayleigh é usado em convecção natural e vale:

```math
Ra = Gr Pr
```

```jldoctest
julia> rayleigh(Air, 80, 20, 2e-3, 101.325, 9.8)
49.491941991306895

julia> rayleigh(Air, 80, 20, 1e-3, 101.325, 9.8)
6.186492748913362
```
"""
function rayleigh(Air::AirType, Ts, Ta, L, Pa=101.325, g=9.80665)
    return grashof(Air, Ts, Ta, L, Pa) * prandtl(Air, Ta)
end

"""
    nuspherefree(Ra, Pr)

Número de Nusselt para uma esfera em convecção natural.
Utiliza o número de Rayleigh `Ra` e o número de Prandtl `Pr`

O número de Nusselt permite calcular o coeficiente de convecção a partir da condutividade
térmica e um comprimento característico:

```math
Nu = h L/k
```

"""
nuspherefree(Ra, Pr) = 2.0 + 0.589*Ra^0.25 / (1 + (0.469/Pr)^(9/16))^(4/9)

"""
    hconvect(U=5.0, D=3e-3, Ts=80.0, Ta=20.0, Pa=101.325 )

Calcula o coeficiente de convecção forçada em uma esfera aquecida dentro de uma corrente de ar.

  * `U` Velocidade em m/s
  * `D` Diâmetro em m
  * `Ts` Temperatura da esfera em °C
  * `Ta` Temperatura do ar ao longe e °C
  * `Pa` Pressão atmosférica em kPa

Se a velocidade for baixa, admite convecção natural.

```math
h = Nu(Re ou Ra, Pr) k(T) / L
```
```jldoctest
julia> hconvect(1.0, 3e-3, 80, 20, 101)
72.37579698816417

julia> hconvect(5.0, 3e-3, 80, 20, 101)
150.73181000789378

julia> hconvect(10.0, 3e-3, 80, 20, 101)
213.4926385463913

julia> hconvect(10.0, 2e-3, 80, 20, 101)
260.7448248784586

julia> hconvect(10.0, 2e-3, 160, 20, 101)
252.09076590966555
```
"""
function hconvect(U=5.0, D=3e-3, Ts=80.0, Ta=20.0, Pa=101.325 )

    Pr = prandtl(Air, Ta)
    relim = 20.0
    if U==0.0
        Ra = rayleigh(Air, Ts, Ta, D, Pa)
        Nu = nuspherefree(Ra, Pr)
    else
        Re = reynolds(Air, D, U, Ta, Pa)
        μ = viscosity(Air, Ta)
        μₛ = viscosity(Air, Ts)
        if Re < relim
            Ra = rayleigh(Air, Ts, Ta, D, Pa)
            Nu1 = nuspherefree(Ra, Pr)
            Nu2 = nusphere(relim, Pr, μ/μₛ)
            Nu = Nu1 * (relim-Re)/relim + Nu2 * Re/relim
        else
            Nu = nusphere(Re, Pr, μ/μₛ)
        end
    end

    h = Nu*heatcond(Air, Ta)/D

    return h

end

function hconvectcyl(U=5.0, D=5e-6, Ts=80.0, Ta=20.0, Pa=101.325 )

    Pr = prandtl(Air, Ta)
    Re = reynolds(Air, D, U, Ta, Pa)
    Nu = nucylinder(Re, Pr)

    h = Nu*heatcond(Air, Ta)/D
    return h

end


text:
pd = ProbabilityDistribution(["head","tails"],[0.5,0.5])

@test_throws DBayesianNode(1, :hulu, pd)
@test_throws DBayesianNode("Hej", pd)
@test_throws DBayesianNode(:hulu, x -> x^5)

pd1 = ProbabilityDistribution(["head","tails"],[0.5,0.5])
pd2 = ProbabilityDistribution(["head","tails"],[0.7,0.3])
a1 = DBayesianNode(:hulu, pd1)
a2 = DBayesianNode(:bulu, pd2)
a3 = DBayesianNode(:tjaeh, pd)
a3 = DBayesianNode(:hulu, pd)
a4 = DBayesianNode(:diags, ["d1", "d2", "d_water"])

@test a1 == a1
@test a1 == a3

a1.index = 2
a3.index = 4

@test a1 == a1
@test a1 != a3

@test_throws CBayesianNode(1, :hulu, x -> x^5)
@test_throws CBayesianNode(:hulu, pd1)
b1 = CBayesianNode(:hulu, x -> x^5)

@test b1.label == :hulu
@test_throws CBayesianNode(:hulu, x -> "$x^5")
@test_throws CBayesianNode(:hulu, x -> :hej)

b2 = CBayesianNode(:bjarke, x -> x)
b3 = CBayesianNode(:bjarke, x -> x)

@test b2.label == :bjarke

@test b2 == b2
@test b2 == b3

b2.index = 2
b3.index = 3

@test b2 == b2
@test b2 != b3

@test probability(b3, 1) == b3.pdf(1)
@test probability(b3, 2) != probability(b3, 4)


text:
import StatsFuns.binompdf

seed!(12345)

function testBinomial(n, p, low, high, N)
  function pmf(x::Int64)
    return binompdf(n, p, x)
  end
  function sampler()
    return sampleBinomial(n, p)
  end
  return testGOFDiscrete(pmf, sampler, low, high, N)
end

@test testBinomial(100, 0.5, 30, 60, 2^21) > 0.01
@test testBinomial(1000, 0.0001, 0, 3, 2^22) > 0.01
@test testBinomial(5, .25, 0, 4, 2^21) > 0.01
@test testBinomial(5, .001, 0, 1, 2^21) > 0.01
@test testBinomial(1000, 0.75, 700, 800, 2^22) > 0.01
@test testBinomial(10, 0.75, 2, 10, 2^22) > 0.01