text:
export Guild

"""
A [`Guild`](@ref)'s verification level.
More details [here](https://discordapp.com/developers/docs/resources/guild#guild-object-verification-level).
"""
@enum VerificationLevel VL_NONE VL_LOW VL_MEDIUM VL_HIGH VL_VERY_HIGH
@boilerplate VerificationLevel :export :lower

"""
A [`Guild`](@ref)'s default message notification level.
More details [here](https://discordapp.com/developers/docs/resources/guild#guild-object-default-message-notification-level).
"""
@enum MessageNotificationLevel MNL_ALL_MESSAGES MNL_ONLY_MENTIONS
@boilerplate MessageNotificationLevel :export :lower

"""
A [`Guild`](@ref)'s explicit content filter level.
More details [here](https://discordapp.com/developers/docs/resources/guild#guild-object-explicit-content-filter-level).
"""
@enum ExplicitContentFilterLevel ECFL_DISABLED ECFL_MEMBERS_WITHOUT_ROLES ECFL_ALL_MEMBERS
@boilerplate ExplicitContentFilterLevel :export :lower

"""
A [`Guild`](@ref)'s MFA level.
More details [here](https://discordapp.com/developers/docs/resources/guild#guild-object-mfa-level).
"""
@enum MFALevel ML_NONE ML_ELEVATED
@boilerplate MFALevel :export :lower

"""
A Discord guild (server).
Can either be an [`UnavailableGuild`](@ref) or a [`Guild`](@ref).
"""
abstract type AbstractGuild end
function AbstractGuild(; kwargs...)
    return if get(kwargs, :unavailable, length(kwargs) <= 2) === true
        UnavailableGuild(; kwargs...)
    else
        Guild(; kwargs...)
    end
end
AbstractGuild(d::Dict{Symbol, Any}) = AbstractGuild(; d...)
mock(::Type{AbstractGuild}) = mock(rand(Bool) ? UnavailableGuild : Guild)

"""
An unavailable Discord guild (server).
More details [here](https://discordapp.com/developers/docs/resources/guild#unavailable-guild-object).
"""
struct UnavailableGuild <: AbstractGuild
    id::Snowflake
    unavailable::Optional{Bool}
end
@boilerplate UnavailableGuild :constructors :docs :lower :merge :mock

"""
A Discord guild (server).
More details [here](https://discordapp.com/developers/docs/resources/guild#guild-object).
"""
struct Guild <: AbstractGuild
    id::Snowflake
    name::String
    icon::Nullable{String}
    splash::OptionalNullable{String}
    owner::Optional{Bool}
    owner_id::Optional{Snowflake}  # Missing in Invite.
    permissions::Optional{Int64}
    region::Optional{String}  # Invite
    afk_channel_id::OptionalNullable{Snowflake}  # Invite
    afk_timeout::Optional{Int}  # Invite
    embed_enabled::Optional{Bool}
    embed_channel_id::OptionalNullable{Snowflake}  # Not supposed to be nullable.
    verification_level::Optional{VerificationLevel}
    default_message_notifications::Optional{MessageNotificationLevel}  # Invite
    explicit_content_filter::Optional{ExplicitContentFilterLevel}  # Invite
    roles::Optional{Vector{Role}}  # Invite
    emojis::Optional{Vector{Emoji}}  # Invite
    features::Optional{Vector{String}}
    mfa_level::Optional{MFALevel}  # Invite
    application_id::OptionalNullable{Snowflake}  # Invite
    widget_enabled::Optional{Bool}
    widget_channel_id::OptionalNullable{Snowflake}  # Not supposed to be nullable.
    system_channel_id::OptionalNullable{Snowflake}  # Invite
    joined_at::Optional{DateTime}
    large::Optional{Bool}
    unavailable::Optional{Bool}
    member_count::Optional{Int}
    max_members::Optional{Int}
    voice_states::Optional{Vector{VoiceState}}
    members::Optional{Vector{Member}}
    channels::Optional{Vector{DiscordChannel}}
    presences::Optional{Vector{Presence}}
    max_presences::OptionalNullable{Int}
    vanity_url_code::OptionalNullable{String}
    description::OptionalNullable{String}
    banner::OptionalNullable{String} # Hash
    djl_users::Optional{Set{Snowflake}}
    djl_channels::Optional{Set{Snowflake}}
end
@boilerplate Guild :constructors :docs :lower :merge :mock

Base.merge(x::UnavailableGuild, y::Guild) = y
Base.merge(x::Guild, y::UnavailableGuild) = x


text:
include("generic_functions.jl")
#example
instances = ["case14"]
output_instance_path = "..\\data_ROPF"
output_decomposition_path = "..\\data_sdp"
list_generation = ["generation1"]
generation_files_path = "..\\data_ROPF"
max_time = 600 #10 minutes

for instance_name in instances
    matpower_instance_path = "..\\matpower_data\\$(instance_name).m"
    for generation in list_generation
        construct_dat_file_ROPF(instance_name, matpower_instance_path, output_instance_path)
        generate_clique_decomposition(instance_name, matpower_instance_path, output_decomposition_path)
        ROPF = ROPF_infos(instance_name,
        matpower_instance_path,
        output_instance_path,
        "cholesky",
        output_decomposition_path,
        generation,
        generation_files_path)
        UB_plus, LB_plus, UB_minus, LB_minus = solve1(ROPF)
        println("Increase in generation : UB=$UB_plus ; LB = $LB_plus")
        println("Decrease in generation : UB=$UB_minus ; LB = $LB_minus \n")
        if UB_minus <= UB_plus
            cp("solutions\\solve1_solution_$(instance_name)_minus.csv", "solutions\\solve1_BEST_solution_$(instance_name).csv")
        else
            cp("solutions\\solve1_solution_$(instance_name)_plus.csv", "solutions\\solve1_BEST_solution_$(instance_name).csv")
        end
        UB_plus, LB_plus, UB_minus, LB_minus = solve2(ROPF, max_time)
        println("After B&B")
        println("Increase in generation : UB=$UB_plus ; LB = $LB_plus")
        println("Decrease in generation : UB=$UB_minus ; LB = $LB_minus \n")
        if UB_minus <= UB_plus
            cp("solutions\\BB_solution_$(instance_name)_minus.csv", "solutions\\solve2_BEST_solution_$(instance_name).csv")
        else
            cp("solutions\\BB_solution_$(instance_name)_plus.csv", "solutions\\solve2_BEST_solution_$(instance_name).csv")
        end
    end
end

ext: jl


text:
"""
Computes the Fourier representation of a Gaussian filter with width ``R_S``.
You would multiply this with the FFT'd field and then IFFT to apply the filter.

# Arguments:
- `R_S::T`: filter radius
- `kv`: wavevector, tuple of (kx, ky, kz)

# Returns:
- `Array{T,3}`: fourier representation of gaussian filter
"""
function kspace_gaussian_filter(R_S::T, kv) where T
    kx, ky, kz = kv
    nx, ny, nz = length(kx), length(ky), length(kz)
    filter_k = zeros(T, (nx, ny, nz))
    Threads.@threads for ix = 1:nx
        for iy = 1:ny, iz = 1:nz
            filter_k[ix, iy, iz] = exp(
                -(kx[ix]^2 + ky[iy]^2 + kz[iz]^2) * R_S^2 / 2)
        end
    end
    return filter_k
end

"""
Apply a gaussian filter to a density field with smoothing radius `R_S`
making use of wavevectors `kv`

# Arguments:
- `f::AbstractArray{T,3}`: density field
- `R_S::T`: filter radius
- `kv`: wavevector, tuple of (kx, ky, kz)

# Returns:
- Array{T,3}: smoothed density field
"""
function smooth_gauss(f::AbstractArray{T,3}, R_S::T, kv) where T

    GF = kspace_gaussian_filter(R_S, kv)  # get filter in Fourier space
    f_Rn = (real(ifft(GF .* fft(f))))
    f_Rn = f_Rn .* (sum(f)/sum(f_Rn))  # normalize
    return f_Rn
end

"""
Apply a gaussian filter to a density field with smoothing radius `R_S`
making use of wavevectors `kv` using log smoothing

# Arguments:
- `f::AbstractArray{T,3}`: density field
- `R_S::T`: filter radius
- `kv`: wavevector, tuple of (kx, ky, kz)

# Returns:
- Array{T,3}: log-smoothed density field
"""
function smooth_loggauss(f::AbstractArray{T,3}, R_S::T, kv) where {AA, T}
    # get filter in Fourier space
    GF = kspace_gaussian_filter(R_S, kv)

    # # simple version of what follows
    # f_Rn = 10 .^(real(ifft(GF .* fft(log10.(f)))))

    buffer = similar(f, Complex{T})   # make a complex version of the array
    @strided buffer .= log10.(f)      # threaded log10 the entire array
    fft!(buffer)                      # FFT
    @strided buffer .= GF .* buffer   # convolve
    ifft!(buffer)                     # IFFT
    f_result = real(buffer)           # get the real part
    @strided f_result .= (10 .^ f_result)  # threaded exp
    norm = sum(f)/sum(f_result)       # compute normalization
    @strided f_result .= f_result .* norm  # normalize
    return f_result
end

# THESE FUNCTIONS ONLY WORK FOR CUBES. From arxiv:1209.2043

# Computes Hessian of log-gaussian-smoothed density field
# f_Rn is the smoothed field
function slow_hessian_from_smoothed(f_Rn::AbstractArray{T,3}, R_S, kv) where T
    dims = size(f_Rn)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))
    f_Rn_hat = fft(f_Rn)

    kx = kv[1]
    ky = kv[2]
    kz = kv[3]

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (1,1)
                hessian[x,y,z,1] =
                    - kx[x] * kx[x] * R_S^2 * f_Rn_hat[x,y,z]
                # (1,2)
                hessian[x,y,z,2] =
                    - kx[x] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
                # (1,3)
                hessian[x,y,z,3] =
                    - kx[x] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
                # (2,2)
                hessian[x,y,z,4] =
                    - ky[y] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
                # (2,3)
                hessian[x,y,z,5] =
                    - ky[y] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
                # (3,3)
                hessian[x,y,z,6] =
                    - kz[z] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end
    hessian[:,:,:,1] .= ifft(hessian[:,:,:,1])
    hessian[:,:,:,2] .= ifft(hessian[:,:,:,2])
    hessian[:,:,:,3] .= ifft(hessian[:,:,:,3])
    hessian[:,:,:,4] .= ifft(hessian[:,:,:,4])
    hessian[:,:,:,5] .= ifft(hessian[:,:,:,5])
    hessian[:,:,:,6] .= ifft(hessian[:,:,:,6])
    real(R_S^2 .* hessian)
end

function fast_hessian_from_smoothed(f_Rn::AbstractArray{T,3}, R_S, kv) where T
    dims = size(f_Rn)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3], 6))
    f_Rn_hat = fft(f_Rn)

    kx = kv[1]
    ky = kv[2]
    kz = kv[3]

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (1,1)
                hessian[x,y,z,1] =
                    - kx[x] * kx[x] * R_S^2 * f_Rn_hat[x,y,z]
                # (1,2)
                hessian[x,y,z,2] =
                    - kx[x] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
                # (1,3)
                hessian[x,y,z,3] =
                    - kx[x] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
                # (2,2)
                hessian[x,y,z,4] =
                    - ky[y] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
                # (2,3)
                hessian[x,y,z,5] =
                    - ky[y] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
                # (3,3)
                hessian[x,y,z,6] =
                    - kz[z] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end

    for j in 1:6
        ifft!(@view (hessian[:,:,:,j]))
    end
    real(R_S^2 .* hessian)
end

function reduce_RAM_hessian_from_smoothed(f_Rn::AbstractArray{T,3}, R_S, kv, output_directory) where T
    dims = size(f_Rn)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))
    f_Rn_hat = fft(f_Rn)
    id_strings = ["_1", "_2", "_3", "_4", "_5", "_6"]

    kx = kv[1]
    ky = kv[2]
    kz = kv[3]

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (1,1)
                hessian[x,y,z] =
                    - kx[x] * kx[x] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[1], hessian)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                #(1,2)
                hessian[x,y,z] =
                    - kx[x] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[2], hessian)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (1,3)
                hessian[x,y,z] =
                    - kx[x] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[3], hessian)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (2,2)
                hessian[x,y,z] =
                    - ky[y] * ky[y] * R_S^2 * f_Rn_hat[x,y,z]
            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[4], hessian)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (2,3)
                hessian[x,y,z] =
                    - ky[y] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]

            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[5], hessian)
    hessian = zeros(Complex{T}, (dims[1], dims[2], dims[3]))

    Threads.@threads for x in 1:length(kv[1])
        for y in 1:length(kv[2])
            for z in 1:length(kv[3])
                # (3,3)
                hessian[x,y,z] =
                    - kz[z] * kz[z] * R_S^2 * f_Rn_hat[x,y,z]

            end
        end
    end

    ifft!(@view (hessian[:,:,:]))
    hessian = real(R_S^2 .* hessian)
    save_hessian_component(output_directory, R_S, id_strings[6], hessian)
    hessian=nothing

end

function signatures_from_hessian(hes::AbstractArray{T,4}) where T
    """
    Function to calculate the signatures from a given hessian.
    """

    hsize = size(hes)[1:3]
    sigs = zeros(T, (hsize[1], hsize[2], hsize[3], 3) )
    Threads.@threads for i in 1:hsize[1]
        for j in 1:hsize[2]
            for k in 1:hsize[3]

                # StaticArray (SA) here for speed, we know the dimensions are 3x3
                #the hessian matrix is symmetric
                hes_slice = SA[
                    hes[i,j,k,1]  hes[i,j,k,2]  hes[i,j,k,3];
                    hes[i,j,k,2]  hes[i,j,k,4]  hes[i,j,k,5];
                    hes[i,j,k,3]  hes[i,j,k,5]  hes[i,j,k,6]
                ]

                e1, e2, e3 = sort(real(eigvals(hes_slice)), rev=false)

                sigs[i,j,k,1] = (abs(e3 / e1) *
                    abs(e3)) * (θ(-e1) * θ(-e2) * θ(-e3))
                sigs[i,j,k,2] = (abs(e2 / e1) *
                    xθ(1-abs(e3 / e1))) * (abs(e2) * θ(-e1) * θ(-e2) )
                sigs[i,j,k,3] = (xθ(1-abs(e2 / e1)) *
                    xθ(1-abs(e3 / e1))) * (abs(e1) * θ(-e1))
            end
        end
    end

    # category shape strength 𝒮 (arxiv:1209.2043 eq 7)
    return sigs
end

"""
In this function we start at minimum length scale to calculate signatures
then iterate to the maximum scale of the problem, calculating the signatures
for each smoothing scale. We then chose the maximum signature over all scales
at each point to categorize whether that point is a wall, filament, or cluster.
"""
function maximum_signature(Rs, field::AbstractArray{T,3}; alg=:NEXUSPLUS) where T

    if alg ∉ (:NEXUS, :NEXUSPLUS)
        throw(ArgumentError("alg must be either :NEXUS or :NEXUSPLUS"))
    end

    nx,ny,nz = size(field)

    #we need to make sure the field has no 0 values
    field = field .+ 0.0001

    #calculate wave vectors for our field
    wave_vecs = wavevectors3D(T, (nx,ny,nz))

    #calculate signatures at each scale Rn, determine max
    sigmax = zeros(T,nx,ny,nz,3)

    for (R_index, R) in enumerate(Rs)

        if alg == :NEXUS
            f_Rn = smooth_gauss(field, R, wave_vecs)
        elseif alg == :NEXUSPLUS
            f_Rn = smooth_loggauss(field, R, wave_vecs)
        end

        H_Rn = fast_hessian_from_smoothed(f_Rn, R, wave_vecs)

        sigs_Rn = signatures_from_hessian(H_Rn)

        for ix = 1:nx, iy = 1:ny, iz = 1:nz, sigtype = 1:3
            sigmax[ix, iy, iz, sigtype] = max(
                sigmax[ix, iy, iz, sigtype], sigs_Rn[ix, iy, iz, sigtype])
        end
    end

    return sigmax
end

"""
In this function we start at minimum length scale to calculate signatures
then iterate to the maximum scale of the problem, calculating the signatures
for each smoothing scale. We then chose the maximum signature over all scales
at each point to categorize whether that point is a wall, filament, or cluster.
"""
function reduce_RAM_maximum_signature(Rs, output_directory, field::AbstractArray{T,3}; alg=:NEXUSPLUS) where T

    if alg ∉ (:NEXUS, :NEXUSPLUS)
        throw(ArgumentError("alg must be either :NEXUS or :NEXUSPLUS"))
    end

    nx,ny,nz = size(field)

    #we need to make sure the field has no 0 values
    field = field .+ 0.0001

    #calculate wave vectors for our field
    wave_vecs = wavevectors3D(T, (nx,ny,nz))

    #-------------------------------------------------------------------------

    print("\n")
    print("Computing All Hessian Values: ")
    print("\n")

    #create all hessians
    for (R_index, R) in enumerate(Rs)

        if alg == :NEXUS
            f_Rn = smooth_gauss(field, R, wave_vecs)
        elseif alg == :NEXUSPLUS
            f_Rn = smooth_loggauss(field, R, wave_vecs)
        end

        #calculate hessian subvolumes
        reduce_RAM_hessian_from_smoothed(f_Rn, R, wave_vecs, output_directory)

    end

    #-------------------------------------------------------------------------

    print("\n")
    print("Compiling Hessians & Calculating Max Signatures: ")
    print("\n")

    #subvolumes
    subvolumes = ["000", "001", "010", "100", "011", "110", "101", "111"]
    nx_sub,ny_sub,nz_sub = nx ÷ 2,ny ÷ 2,nz ÷ 2

    for (sub, subvolume) in enumerate(subvolumes)

        #calculate signatures at each scale Rn, determine max
        sigmax_sub = zeros(T,nx_sub,ny_sub,nz_sub,3)

        for (R_index, R) in enumerate(Rs)

            if alg == :NEXUS
                f_Rn = smooth_gauss(field, R, wave_vecs)
            elseif alg == :NEXUSPLUS
                f_Rn = smooth_loggauss(field, R, wave_vecs)
            end

            #combine hessian components
            compile_sub_hessians(f_Rn, R, subvolume, output_directory)

            #load back in necessary hessian subvolume
            @load output_directory * "hessian_component_"*subvolume*"_full.jld2" hessian_sub

            #calculate signatures on subvolume
            sigs_Rn = signatures_from_hessian(hessian_sub)

            #determine maximum signatures
            for ix = 1:nx_sub, iy = 1:ny_sub, iz = 1:nz_sub, sigtype = 1:3
                sigmax_sub[ix, iy, iz, sigtype] = max(
                    sigmax_sub[ix, iy, iz, sigtype], sigs_Rn[ix, iy, iz, sigtype])
            end

        end

        @save output_directory * "max_sigs_"*subvolume*".jld2" sigmax_sub

        #delete used hessian subvolume
        file = output_directory * "hessian_component_"*subvolume*"_full.jld2"
        run(`rm $file`)

    end
end


text:
module MovingWindow

export MovingWindowSolver
export MovingWindowSolution
export MoWA

using Reexport
using ArrowMacros
@reexport using TimeParallel
using RecipesBase

include("solver.jl")
include("solution.jl")
include("solve.jl")
include("plot.jl")

end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct VirtualMachineExtensionUpdate <: SwaggerModel
    tags::Any # spec type: Union{ Nothing, Dict{String, String} } # spec name: tags

    function VirtualMachineExtensionUpdate(;tags=nothing)
        o = new()
        validate_property(VirtualMachineExtensionUpdate, Symbol("tags"), tags)
        setfield!(o, Symbol("tags"), tags)
        o
    end
end # type VirtualMachineExtensionUpdate

const _property_map_VirtualMachineExtensionUpdate = Dict{Symbol,Symbol}(Symbol("tags")=>Symbol("tags"))
const _property_types_VirtualMachineExtensionUpdate = Dict{Symbol,String}(Symbol("tags")=>"Dict{String, String}")
Base.propertynames(::Type{ VirtualMachineExtensionUpdate }) = collect(keys(_property_map_VirtualMachineExtensionUpdate))
Swagger.property_type(::Type{ VirtualMachineExtensionUpdate }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_VirtualMachineExtensionUpdate[name]))}
Swagger.field_name(::Type{ VirtualMachineExtensionUpdate }, property_name::Symbol) =  _property_map_VirtualMachineExtensionUpdate[property_name]

function check_required(o::VirtualMachineExtensionUpdate)
    true
end

function validate_property(::Type{ VirtualMachineExtensionUpdate }, name::Symbol, val)
end


text:
export smooth
using FFTViews

function smooth(a::RFFTSpectrumArray{T}, octave::Real = 1//3) where {T <: Real}
    # algorithm: Hatziantoniou, Panagiotis D., and John N. Mourjopoulos. "Generalized fractional-octave smoothing of audio and acoustic responses." Journal of the Audio Engineering Society 48.4 (2000): 259-280.
    # but based on description in:
    # Tylka, Joseph G., Braxton B. Boren, and Edgar Y. Choueiri. "A generalized method for fractional-octave smoothing of transfer functions that preserves log-frequency symmetry." Journal of the Audio Engineering Society 65.3 (2017): 239-245.
    # Method 1: Symmetric weights
    # rectangular window
    octave ≥ 0 || error("octave must be ≥ 0")
    n = nframes(a)
    as = similar(a)
    X = data(a)
    Xs = data(as)

    X = [X; conj.(X[end-1:-1:2,:])] # we expect spectra created using rfft(), so extend by negative frequencies

    rQ = 2sinh(octave*log(2)/2) # Q reciprocial: 1/Q CHECK use log2?
    for c in 1:nchannels(a)
        XV = FFTView(view(X, :, c)) # make it circular and starting on zero index
        plb, pub = 0, -1 # previous lower and upper bounds
        msum = 0.0 # moving sum
        for k in 0:n-1
            mk = floor(Int, k*rQ/2)
            lb, ub = k-mk, k+mk
            msum = msum - sum(XV[plb:lb-1]) + sum(XV[pub+1:ub])
            mmean = msum / (2mk + 1)
            Xs[k+1, c] = mmean
            plb, pub = lb, ub
        end
    end
    as
end

function smooth_naive(a::RFFTSpectrumArray{T}, octave::Real = 1//3) where {T <: Real}
    # slower version of smooth - always recomputes window means
    octave ≥ 0 || error("octave must be ≥ 0")
    n = nframes(a)
    as = similar(a)
    X = data(a)
    Xs = data(as)

    X = [X; conj.(X[end-1:-1:2,:])] # we expect spectra created using rfft(), so extend by negative frequencies
    X = FFTView(X) # make it circular and starting on zero index

    rQ = 2sinh(octave*log(2)/2) # Q reciprocial: 1/Q CHECK use log2?
    for c in 1:nchannels(a)
        for k in 0:n-1
            mk = floor(Int, k*rQ/2)
            # Xs[k+1, c] = mean(X[k-mk:k+mk, c])
            Xs[k+1, c] = mean(X[max(k-mk, 0):min(k+mk, n-1), c])
        end
    end
    as
end

function smooth(X::RFFTSpectrumArray{T}, octave::Real = 1//3) where {E, T <: SignalElement{E}}
    Y = similar(X)
    # mags = smooth_naive(abs.(X), octave)
    mags = smooth(abs.(X), octave)
    phis = smooth(unwrap(angle.(X), dims=1), octave)
    Y .= (MagPhase(max(m, 0.0), p) for (m, p) in zip(mags, phis)) # max for numerical instability giving negative magnitude
    Y
end


text:
#trying out VCFTools.jl and GeneticVariation.jl
using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools

cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "test.08Jun17.d8b.vcf"
reader = VCF.Reader(openvcf(vcffile, "r"))
sampleID = VCF.header(reader).sampleID
record = read(reader)

record = VCF.Record("20\t14370\trs6054257\tG\tAAAA\t29\tPASS\tNS=3;DP=14;AF=0.5;DB;H2\tGT:GQ:DP:HQ\t0|0:48:1:51,51\t1|0:48:8:51,51")
record = VCF.Record("20\t14370\t.\tG\tA\t29\tPASS;fdsa\t.\tDS\t1.99\t1.01")
record = VCF.Record("20\t14370\t.\tG\tA\t29\tPASS\t.\tGT\t1/0\t0/0")
record.genotype

#test convert_ht
using Revise
using GeneticVariation
using VCFTools

cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "test.08Jun17.d8b.vcf"
H = convert_ht(Float32, vcffile, has_missing=true)

# byte representation mapping
String([0x09]) # '\t'
String([0x3a]) # ':'
String([0x30]) # '0'
String([0x31]) # '1'
String([0x2e]) # '.'
String([0x7c]) # '|'
String([0x2f]) # '/'

# test filter
using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools

#import data
cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "test.08Jun17.d8b.vcf.gz"
samples = nsamples(vcffile)
records = nrecords(vcffile)
reader = VCF.Reader(openvcf(vcffile, "r"))
record = read(reader)

# define masks
Random.seed!(123)
record_mask = trues(records)
# sample_mask = bitrand(samples)
sample_mask = trues(samples)
sample_mask[1] = false
sample_mask[end] = false

des = "filtered.test.08Jun17.d8b.vcf.gz"
VCFTools.filter(vcffile, record_mask, sample_mask, des=des)
reader = VCF.Reader(openvcf(des, "r"))
record = read(reader)
record.genotype

nsamples(des)

X = convert_gt(Float32, vcffile)
X_filter = convert_gt(Float32, "filtered." * vcffile)

@benchmark VCFTools.filter(vcffile, record_mask, sample_mask)
# old = 671.380 ms, 326.66 MiB, 3631762 alloc
# new = 573.432 ms, 257.75 MiB, 3352787 alloc

des = "filter." * vcffile
reader = VCF.Reader(openvcf(vcffile, "r"))
writer = VCF.Writer(openvcf(des, "w"), filter_header(reader, sample_mask))

reader = VCF.Reader(openvcf(vcffile, "r"))
record = read(reader)

# test masking

using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools

cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "test.08Jun17.d8b.vcf"
samples = nsamples(vcffile)
records = nrecords(vcffile)
masks = bitrand(records, samples)
mask_gt(vcffile, masks)
@benchmark mask_gt(vcffile, masks) #40.840 ms, 40.09 MiB, 521257 alloc

# test convert_ds
using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools

cd("/Users/biona001/Benjamin_Folder/UCLA/research/2nd_project/benchmarks/AFRPed/minimac4")
vcffile = "minimac4_result.dose.vcf"
dosage = convert_ds(Float64, vcffile)

# test overwrite_alt_ref_allele
using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools

function overwrite_alt_ref_allele(
    src::AbstractString,
    des::AbstractString,
    ref::String = "A",
    alt::String = "T"
    )

    # create input and output VCF files
    reader = VCF.Reader(openvcf(src, "r"))
    writer = VCF.Writer(openvcf(des, "w"), VCF.header(reader))

    # define ref/alt allele
    dict = Dict([("A", 0x41), ("T", 0x54), ("C", 0x43), ("G", 0x47), ("N", 0x4e)])
    haskey(dict, ref) || error("ref allele must be A, T, C, G, or N.")
    haskey(dict, alt) || error("ref allele must be A, T, C, G, or N.")
    ref, alt = dict[ref], dict[alt]
    for record in reader
        record.data[record.alt[1]] .= alt
        record.data[record.ref[1]] = ref
        write(writer, record)
    end
    flush(writer); close(reader); close(writer)
end

src = "full_test.vcf"
reader = VCF.Reader(openvcf(src, "r"))
record = read(reader)
# overwrite_alt_ref_allele("test.vcf")
overwrite_alt_ref_allele("test.vcf", "test_result.vcf")

# Simulate genotype matrix using real haplotypes
# 1000 Genomes Phase 1, chr 22 data downloaded from: https://genome.sph.umich.edu/wiki/Minimac4
using Revise
using MendelImpute
using VCFTools
using BenchmarkTools

cd("/Users/biona001/.julia/dev/MendelImpute/simulation/compare4/VCF_Files")

#2184 haplotypes each with 365644 SNPs -> 99MB of RAM
@time H = convert_ht(Bool, "ALL.chr22.phase1_v3.snps_indels_svs.genotypes.all.noSingleton.vcf.gz")

#simulate 100 samples each with 365644 SNPs
X = simulate_genotypes(H', 100)

# test convert_gt transpose

using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools
using Test

cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "test.08Jun17.d8b.vcf"

A  = convert_gt(Float64, vcffile)
At = convert_gt(Float64, vcffile, trans=true)
@test all(At .== A')
# size(A) = (191, 1356)

@benchmark convert_gt(Float64, vcffile) # 63.926 ms, 104.59 MiB, 1070395 alloc
@benchmark convert_gt(Float64, vcffile, trans=true)# 65.772 ms, 104.59 MiB, 1070396 alloc

# test if eof(reader) is working
out = Matrix{Union{Float64, Missing}}(undef, 191, 1400)
reader = VCF.Reader(openvcf(vcffile, "r"))
copy_gt!(out, reader)
@test all(ismissing.(out[:, 1357:end]))

out = Matrix{Union{Float64, Missing}}(undef, 1400, 191)
reader = VCF.Reader(openvcf(vcffile, "r"))
copy_gt_trans!(out, reader)
@test all(ismissing.(out[1357:end, :]))

# test impute, center, scale
A = Matrix{Union{Float64, Missing}}(undef, 191, 1400)
reader = VCF.Reader(openvcf(vcffile, "r"))
copy_gt!(A, reader, impute=true, center=true, scale=true)
@test isapprox(mean(A[:, 5]), 0, atol=10)
@test isapprox(var(A[:, 5]), 1, atol=10)

At = Matrix{Union{Float64, Missing}}(undef, 1400, 191)
reader = VCF.Reader(openvcf(vcffile, "r"))
copy_gt_trans!(At, reader, impute=true, center=true, scale=true)
@test isapprox(mean(At[5, :]), 0, atol=10)
@test isapprox(var(At[5, :]), 1, atol=10)

# test convert_ht transpose

using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools
using Test

# cd("/Users/biona001/.julia/dev/VCFTools/test")
# vcffile = "test.08Jun17.d8b.vcf"
vcffile = "target.vcf"

H  = convert_ht(Float64, vcffile)
Ht = convert_ht(Float64, vcffile, trans=true)
@test all(Ht .== H')
# size(H) = (382, 1356)

@benchmark convert_ht(Float64, vcffile) # 44.523 ms, 58.89 MiB, 552403 alloc
@benchmark convert_ht(Float64, vcffile, trans=true) # 44.900 ms, 58.89 MiB, 552404 alloc

# test if eof(reader) is working
out = Matrix{Float64}(undef, 382, 1400)
reader = VCF.Reader(openvcf(vcffile, "r"))
copy_ht!(out, reader)
@test size(out) == (382, 1357)
@test all(out[:, 1358:end] .== 0.0)

out = Matrix{Float64}(undef, 1400, 382)
reader = VCF.Reader(openvcf(vcffile, "r"))
copy_ht_trans!(out, reader)
@test size(out) == (1400, 382)
@test all(out[1358:end, :] .== 0.0)

# test convert_ds transpose

using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools
using Test

cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "test.08Jun17.d8b.vcf"

D  = convert_ds(Float64, vcffile)
Dt = convert_ds(Float64, vcffile, trans=true)
@test all(Dt .== D')
# size(D) = (191, 1356)

@benchmark convert_ds(Float64, vcffile) # 102.551 ms, 183.63 MiB, 2106379 alloc
@benchmark convert_ds(Float64, vcffile, trans=true) # 114.470 ms ms, 183.63 MiB, 2106380 alloc

# test if eof(reader) is working
out = Matrix{Union{Missing, Float64}}(undef, 191, 1400)
reader = VCF.Reader(openvcf(vcffile, "r"))
copy_ds!(out, reader)
@test size(out) == (191, 1400)
@test all(ismissing.(out[:, 1358:end]))

out = Matrix{Union{Missing, Float64}}(undef, 1400, 191)
reader = VCF.Reader(openvcf(vcffile, "r"))
copy_ds_trans!(out, reader)
@test size(out) == (1400, 191)
@test all(ismissing.(out[1358:end, :]))

using Revise
using VCFTools
using MendelImpute
using GeneticVariation
using Random
cd("/Users/biona001/.julia/dev/MendelImpute/simulation")

# impute
tgtfile = "./compare6/target_masked.vcf.gz"
reffile = "./compare6/haplo_ref.vcf"
outfile = "./compare6/imputed_target.vcf.gz"
width   = 800

H  = convert_ht(Float64, reffile, trans=true)
Hb = convert_ht(Bool, reffile, trans=true)

Hb_c = convert(Matrix{Float64}, Hb)
all(Hb_c .== Hb)

using Revise
using LoopVectorization
using Random
using LinearAlgebra
using BenchmarkTools

function gemv_naive!(c, A, b)
    @inbounds for j in 1:size(A, 2)
        @simd for i in 1:size(A, 1)
            c[i] += A[i, j] * b[j]
        end
    end
end

function gemv_avx!(c, A, b)
    @avx for j in 1:size(A, 2), i in 1:size(A, 1)
        c[i] += A[i, j] * b[j]
    end
end

function gemv_avx2!(c, A, b)
    @avx for i in 1:size(A, 1)
        ci = zero(eltype(c))
        for j in 1:size(A, 2)
            ci += A[i, j] * b[j]
        end
        c[i] = ci
    end
end

Random.seed!(2020)
n = 10000
A = bitrand(n, n)
b = rand(n)
c = zeros(n)
c_avx = zeros(n)
c_avx2 = zeros(n)
c_true = zeros(n)

# check correctness
mul!(c_true, A, b)
gemv_avx!(c_avx, A, b)
gemv_avx2!(c_avx2, A, b)
gemv_naive!(c, A, b)
[c c_avx c_avx2 c_true]

# efficiency (A = bitmatrix, 1000x1000)
@benchmark mul!(c_true, A, b)      # 3.411 ms     (Julia's default)
@benchmark gemv_naive!(c, A, b)    # 4.230 ms     (Looping using @simd and @inbounds)
@benchmark gemv_avx!(c_avx, A, b)  # 566.309 μs   (using LoopVectorization)
@benchmark gemv_avx2!(c_avx, A, b) # 572.952 μs   (using LoopVectorization with accumulator)

# efficiency (A = bitmatrix, 10000x10000)
@benchmark mul!(c_true, A, b)      # 341.411 ms  (Julia's default)
@benchmark gemv_naive!(c, A, b)    # 424.198 ms  (Looping using @simd and @inbounds)
@benchmark gemv_avx!(c_avx, A, b)  # 77.201 ms   (using LoopVectorization)
@benchmark gemv_avx2!(c_avx, A, b) # 73.227 ms   (using LoopVectorization with accumulator)

# efficiency (A = bitmatrix, 50000x50000)
@benchmark mul!(c_true, A, b)      # 8.999 s   (Julia's default)
@benchmark gemv_naive!(c, A, b)    # 10.207 s  (Looping using @simd and @inbounds)
@benchmark gemv_avx!(c_avx, A, b)  # 2.685 s   (using LoopVectorization)
@benchmark gemv_avx2!(c_avx, A, b) # 2.197 s   (using LoopVectorization with accumulator)

# efficiency (A = bitmatrix, 100000x100000)
@time mul!(c_true, A, b)      # 37.167032 s   (Julia's default)
@time gemv_naive!(c, A, b)    # 42.665357 s   (Looping using @simd and @inbounds)
@time gemv_avx!(c_avx, A, b)  # 17.452804 s   (using LoopVectorization)
@time gemv_avx2!(c_avx, A, b) # 17.881693 s   (using LoopVectorization with accumulator)

# efficiency (A = Matrix{Float64}, 1000x1000)
BLAS.set_num_threads(1)
@benchmark mul!(c_true, A, b)      # 137.602 μs   (Julia's default: calls BLAS)
@benchmark gemv_naive!(c, A, b)    # 155.327 μs   (Looping using @simd and @inbounds)
@benchmark gemv_avx!(c_avx, A, b)  # 174.087 μs   (using LoopVectorization)
@benchmark gemv_avx2!(c_avx, A, b) # 189.796 μs   (using LoopVectorization with accumulator)

# efficiency (A = Matrix{Float64}, 10000x10000)
BLAS.set_num_threads(1)
@benchmark mul!(c_true, A, b)      # 41.293 ms   (Julia's default: calls BLAS)
@benchmark gemv_naive!(c, A, b)    # 47.362 ms   (Looping using @simd and @inbounds)
@benchmark gemv_avx!(c_avx, A, b)  # 97.696 ms   (using LoopVectorization)
@benchmark gemv_avx2!(c_avx, A, b) # 99.377 ms   (using LoopVectorization with accumulator)

using VectorizationBase: gesp, stridedpointer
using Random
using LoopVectorization
using BenchmarkTools
using LinearAlgebra

function gemv_avx!(c, A, b)
    @avx for j in 1:size(A, 2), i in 1:size(A, 1)
        c[i] += A[i, j] * b[j]
    end
end
function gemv_avx_512by512!(c, A, b)
    @avx for j in 1:512, i in 1:512
        c[i] += A[i, j] * b[j]
    end
end
function gemv_tile!(c, A, b)
    M, N = size(A)
    Miter = M >>> 9 # fast div(M, 512)
    Mrem = M & 511 # fast rem(M, 512)
    Niter = N >>> 9
    Nrem = N & 511
    GC.@preserve c A b for n in 0:Niter-1
        for m in 0:Miter-1
            gemv_avx_512by512!(
                gesp(stridedpointer(c), (512m,)),
                gesp(stridedpointer(A), (8m, 8n)),
                gesp(stridedpointer(b), (512n,))
            )
        end
        if mrem > 0

        end
    end
    # TODO: handle nrem
end

n = 1 << 14 # 2^14, multiple of 512, because we're not handling remainders
A = bitrand(n, n);
b = rand(n); c1 = zero(b); c2 = zero(b);
gemv_tile!(c1, A, b);
gemv_avx!(c2, A, b);
c1 ≈ c2

@benchmark gemv_tile!($c1, $A, $b) # 200.051 ms  (cache aware avx)
@benchmark gemv_avx!($c2, $A, $b)  # 449.613 ms  (naive avx)
@benchmark mul!($c1, $A, $b)       # 1.133 s     (Julia built-in)

using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools

#import data
cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "test.08Jun17.d8b.vcf.gz"
snps = nrecords(vcffile)
chunks = ceil(Int, snps / 100)
snps_per_chunk = ceil(Int, snps / chunks)
(chunk - 1)*snps_per_chunk+1:(chunk * snps_per_chunk) #ranges

full_X = convert_gt(Float32, vcffile, trans=true)

Xreader = VCF.Reader(openvcf(vcffile, "r"))
X = Matrix{Union{Float32, Missing}}(undef, snps_per_chunk, nsamples(vcffile))
for chunk in 1:chunks
    copy_gt_trans!(X, Xreader)
    println(all(X .== full_X[(chunk - 1)*snps_per_chunk+1:(chunk * snps_per_chunk), :]))
end

# try data import strategy outlined in https://github.com/JuliaLang/julia/issues/34195#issuecomment-569343440

using Mmap
using VCFTools
using CodecZlib

using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools

#import data
cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "test.08Jun17.d8b.vcf.gz"

H, sampleID, record_chr, record_pos, record_ids, record_ref, record_alt = convert_ht(Float32, vcffile, save_snpinfo(), trans=true)
@code_warntype convert_ht(Float32, vcffile, save_snpinfo(), trans=true)
H2 = convert_ht(Float32, vcffile, trans=true)
all(H .== H2)
@code_warntype convert_ht(Float32, vcffile, trans=true)

X, sampleID, record_chr, record_pos, record_ids, record_ref, record_alt = convert_gt(Float32, vcffile, save_snpinfo(), trans=true)
@code_warntype convert_gt(Float32, vcffile, save_snpinfo(), trans=true)
X2 = convert_gt(Float32, vcffile, trans=true)
all(X .== X2)
@code_warntype convert_gt(Float32, vcffile, trans=true)

# test aim selection
using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools
using DataFrames
using CSV

cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "chr22.1kg.phase3.v5a.vcf.gz"
reader = VCF.Reader(openvcf(vcffile, "r"))
sampleID = VCF.header(reader).sampleID
close(reader)

# get each sample's population origin
cd("/Users/biona001/.julia/dev/MendelImpute/data/1000_genome_phase3_v5/country_origin")
sampleID_to_population = Dict{String, String}()
df = DataFrame(Person = String[], Sex = Int[], Ethnic = String[])
for population in readdir("data/")
    for sample in readdir("data/" * population)
        sample == ".DS_Store" && continue
        sampleID_to_population[sample] = population
        sample in sampleID && push!(df, (sample, 1, population))
    end
end
df
cd("/Users/biona001/.julia/dev/VCFTools/test")
CSV.write("chr22.1kg.phase3.ped", df)

# chr22 in 1000 genomes data
cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "chr22.1kg.phase3.v5a.vcf.gz"
pvals = aim_select(vcffile, sampleID_to_population)
snps = nrecords(vcffile)
bonferroni = 0.05 / snps
signifiant_snps = findall(p -> p < bonferroni, pvals)

using SnpArrays
vcf2plink("chr22.1kg.phase3.v5a.vcf.gz", "chr22.1kg.phase3.v5a")

using Revise
using MendelAimSelection
AimSelection("aim_control.txt")

# test aim selection
using Revise
using GeneticVariation
using Random
using VCFTools
using BenchmarkTools
using DataFrames
using CSV
using MendelAimSelection

# get each sample's population origin
cd("/Users/biona001/.julia/dev/MendelImpute/data/1000_genome_phase3_v5/country_origin")
sampleID_to_population = Dict{String, String}()
for population in readdir("data/")
    for sample in readdir("data/" * population)
        sample == ".DS_Store" && continue
        sampleID_to_population[sample] = population
    end
end

# chr22 in 1000 genomes data
cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "chr22.1kg.phase3.v5a.vcf.gz"
pvals = aim_select(vcffile, sampleID_to_population)

# test 16th record
cd("/Users/biona001/.julia/dev/VCFTools/test")
vcffile = "chr22.1kg.phase3.v5a.vcf.gz"
reader = VCF.Reader(openvcf(vcffile, "r"))
for i in 1:15
    read(reader)
end
record = read(reader)
close(reader)

sampleID = VCF.header(reader).sampleID
ethnics = VCFTools.ethnic(sampleID, sampleID_to_population)
sort!(ethnics) # why should I sort this?
populations = unique(ethnics)
alleles = Vector{Int}(undef, length(populations))
genes = Vector{Int}(undef, length(populations))
aim_select(record, populations, ethnics, alleles, genes)

AimSelection("aim_control.txt")


text:
module Layers21

include("init.jl")
include("block.jl")
include("add.jl")
include("mul.jl")
include("op.jl")
include("conv.jl")
include("linear.jl") # TODO: test, redesign
include("batchnorm.jl")
include("zeropad.jl")
include("show.jl")
include("jld2.jl")
include("embed.jl")

# TODO: rethink param/Param and generally parameter initialization, array_type etc.
#  xavier may not be optimal for embedding, should specify init options etc.
# TODO: rethink eliminating layers without parameters: dropout, activation

#
#
# include("dropout.jl")
# include("layernorm.jl")

end


text:
module Memoizeds

# Memoized data for the current iteration

using Comm, FunRefs, Par
using Cells, Defs, Domains, Grids, Norms

import Base.wait

export Memoized
export wait
export rk2
export info_output, file_output

immutable Memoized
    n::Int
    state::FunRef{Domain}
    rhs::FunRef{Domain}
    error::FunRef{Domain}
    errorNorm::FunRef{Norm}

    function Memoized(n::Int, state::FunRef{Domain})
        rhs = par(R=Domain) do; Domains.rhs(state[]) end
        error = par(R=Domain) do; Domains.error(state[]) end
        errorNorm = par(R=Norm) do; Domains.norm(error[]) end
        new(n, state, rhs, error, errorNorm)
    end
end

function wait(m::Memoized)
    wait(m.state[])
    wait(m.rhs[])
    wait(m.error[])
    wait(m.errorNorm)
end

# RK2
function rk2(m::Memoized)
    s0 = m.state
    r0 = m.rhs
    # Step 1
    s1 = par(R=Domain) do; Domains.axpy(0.5 * Defs.dt(), r0[], s0[]) end
    r1 = par(R=Domain) do; Domains.rhs(s1[]) end
    # Step 2
    par(R=Domain) do; Domains.axpy(Defs.dt(), r1[], s0[]) end
end

# Output
# TODO: futurize
function do_info_output(io::IO, m::Memoized)
    @assert Comm.myproc() == 1
    s = m.state[]
    en = m.errorNorm[]
    cell_size = norm(Cell()).count
    ncells = en.count / cell_size
    println(io, "n=", m.n, " t=", s.t, " ",
            "ncells: ", ncells, " L2-norm[error]: ", norm2(en))
end

function info_output(io::IO, m::Memoized;
                     is_first::Bool=false, was_last::Bool=false)
    if Defs.do_this_time(m.n, Defs.info_every,
                         is_first=is_first, was_last=was_last)
        do_info_output(io, m)
    end
end

# TODO: futurize
function do_file_output(io::IO, m::Memoized)
    @assert Comm.myproc() == 1
    s = m.state[]
    r = m.rhs[]
    en = m.errorNorm[]
    cell_size = norm(Cell()).count
    ncells = en.count / cell_size
    print(io, "State: ", s)
    print(io, "RHS: ", r)
    println(io, "ncells: ", ncells, " L2-norm[error]: ", norm2(en))
    println(io)
end

function file_output(io::IO, m::Memoized;
                     is_first::Bool=false, was_last::Bool=false)
    if Defs.do_this_time(m.n, Defs.file_every,
                         is_first=is_first, was_last=was_last)
        do_file_output(io, m)
    end
end

end


text:
using GeometricIntegrators.Solutions
using GeometricIntegrators.Solutions: createHDF5
using GeometricIntegrators.TestProblems.KuboOscillator
using HDF5: HDF5File
using Test

ntime  = 10
Δt     = .1
h5file = "test.hdf5"

@testset "$(rpad("SDE Solution",80))" begin
    sde  = kubo_oscillator_sde_1()
    ssol = Solution(sde, Δt, ntime)
    @test typeof(ssol) <: SolutionSDE

    # test hdf5 in- and output
    h5 = createHDF5(ssol, h5file)
    @test typeof(h5) == HDF5File
    close(h5)
    @test isfile(h5file)

    h5 = createHDF5(ssol, h5file, overwrite=false)
    @test typeof(h5) == HDF5File
    close(h5)
    @test isfile(h5file)
    rm(h5file)
end

@testset "$(rpad("PSDE Solution",80))" begin
    psde  = kubo_oscillator_psde_1()
    ssol = Solution(psde, Δt, ntime)
    @test typeof(ssol) <: SolutionPSDE

    # test hdf5 in- and output
    h5 = createHDF5(ssol, h5file)
    @test typeof(h5) == HDF5File
    close(h5)
    @test isfile(h5file)

    h5 = createHDF5(ssol, h5file, overwrite=false)
    @test typeof(h5) == HDF5File
    close(h5)
    @test isfile(h5file)
    rm(h5file)
end

@testset "$(rpad("SPSDE Solution",80))" begin
    spsde  = kubo_oscillator_spsde_1()
    ssol = Solution(spsde, Δt, ntime)
    @test typeof(ssol) <: SolutionPSDE
end


text:
"""
    struct ExtractVector{T} <: AbstractExtractor
	    n::Int
	    uniontypes::Bool
    end

represents an array of a fixed length, typically a feature vector of numbers of type T

```jloctest
julia> sc = ExtractVector(4)
julia> sc([2,3,1,4]).data
3×1 Array{Float32,2}:
 2.0
 3.0
 1.0
```
"""
struct ExtractVector{T} <: AbstractExtractor
	n::Int
	uniontypes::Bool
end
ExtractVector(n::Int, uniontypes = true) = ExtractVector{FloatType}(n, uniontypes)

make_missing_vector(s::ExtractVector, v, store_input) =
    s.uniontypes ?
    _make_array_node(stabilize_types_vector(s, fill(missing, s.n, 1)), [v], store_input) :
    error("This extractor does not support missing values")
stabilize_types_vector(s::ExtractVector{T}, x) where {T} = s.uniontypes ? Matrix{Union{Missing, T}}(x) : x

(s::ExtractVector)(v::MissingOrNothing; store_input=false) = make_missing_vector(s, v, store_input)
(s::ExtractVector{T})(::ExtractEmpty; store_input=false) where {T} =
    ArrayNode(stabilize_types_vector(s, Matrix{T}(undef, s.n, 0)))
(s::ExtractVector)(v; store_input=false) = make_missing_vector(s, v, store_input)

function (s::ExtractVector{T})(v::Vector; store_input=false) where {T}
	isempty(v) && make_missing_vector(s, v, store_input)
	if length(v) > s.n
		@warn "array too long, truncating"
		x = reshape(T.(v[1:s.n]), :, 1)
	elseif length(v) < s.n
		!s.uniontypes && error("This extractor does not support missing values")
		x = Matrix{Union{Missing, T}}(missing, s.n, 1)
		x[1:length(v)] .= v
	else
		x = reshape(T.(v), :, 1)
	end
	_make_array_node(stabilize_types_vector(s, x), [v], store_input)
end

Base.length(e::ExtractVector) = e.n
Base.hash(e::ExtractVector, h::UInt) = hash((e.n, e.uniontypes), h)
Base.:(==)(e1::ExtractVector, e2::ExtractVector) = e1.n === e2.n && e1.uniontypes === e2.uniontypes


text:
const GP = GaussianProcesses
function mean_var(model::GPBase, x::AbstractArray{<:Any, 1})
    μ, var = GP.predict_f(model, reshape(x, :, 1))
    μ[1], var[1]
end
myrand(model::GPBase, x::AbstractArray{<:Any, 1}) = rand(model, reshape(x, :, 1))[1]
myrand(model::GPBase, x) = rand(model, x)
mean_var(model::GPBase, x::AbstractArray{<:Any, 2}) = GP.predict_f(model, x)
dims(model::GPBase) = size(model.x)
maxy(model::GPBase) = length(model.y) == 0 ? -Inf : maximum(model.y)
update!(model::GPE{<:ElasticArray}, x, y) = append!(model, x, y)
function update!(model::GPE, x, y)
    if isempty(y)
        GP.fit!(model, model.x, model.y)
    else
        GP.fit!(model, hcat(model.x, x), [model.y; y])
    end
end

mutable struct MAPGPOptimizer{NT} <: ModelOptimizer
    i::Int
    every::Int
    options::NT
end
"""
    MAPGPOptimizer(; every = 10, kwargs...)

Set the GP hyperparameters to the maximum a posteriori (MAP) estimate `every`
number of steps. Run `BayesianOptimization.defaultoptions(MAPGPOptimizer)` to
see the default options. By default, all priors are flat. Uniform priors in an
interval can be specified by setting the bounds of the form
`[lowerbound, upperbound]`, e.g. for a kernel with 3 parameters one would set
`kernbounds = [[-3, -3, -4], [2, 3, 1]]`. Non-flat priors can be specified
directly on the GP parameters, e.g.
`using Distributions; set_priors!(mean, [Normal(0., 1.)])`
"""
MAPGPOptimizer(; every = 10, kwargs...) = MAPGPOptimizer(0, every,
                                                       merge(defaultoptions(MAPGPOptimizer),
                                                             kwargs.data))
function optimizemodel!(o::MAPGPOptimizer, model::GPBase)
    if o.i % o.every == 0
        optimizemodel!(model, o.options)
    end
    o.i += 1
end
defaultoptions(::Type{MAPGPOptimizer}) =
    (domean = true, kern = true, noise = true, lik = true, meanbounds = nothing,
     kernbounds = nothing, noisebounds = nothing, likbounds = nothing,
     method = :LD_LBFGS, maxeval = 500)

function optimizemodel!(gp::GPBase, options)
    params_kwargs = GP.get_params_kwargs(gp; domean=options.domean,
                                         kern=options.kern,
                                         noise=options.noise,
                                         lik=options.lik)
    f = (x, g) -> begin
            GP.set_params!(gp, x; params_kwargs...)
            GP.update_target_and_dtarget!(gp; params_kwargs...)
            @. g = gp.dtarget
            gp.target
        end
    lb, ub = GP.bounds(gp, options.noisebounds, options.meanbounds,
                       options.kernbounds, options.likbounds;
                       domean = options.domean, kern = options.kern,
                       noise = options.noise, lik = options.lik)
    opt = NLopt.Opt(options.method, length(lb))
    NLopt.lower_bounds!(opt, lb)
    NLopt.upper_bounds!(opt, ub)
    NLopt.maxeval!(opt, options.maxeval)
    NLopt.max_objective!(opt, f)
    fx, x, ret = NLopt.optimize(opt, GP.get_params(gp; params_kwargs...))
    ret == NLopt.FORCED_STOP && throw(InterruptException())
    fx, x, ret
end


text:
function reward(m::POMDP{S}, s::S, a::CovidAction, sp::S) where S <: CovidState
    inf_loss = m.inf_loss*infected(sp)/m.N
    test_loss = m.test_loss*a.testing_prop
    testrate_loss = m.testrate_loss*abs(a.testing_prop-s.prev_action.testing_prop)
    return -(inf_loss + test_loss + testrate_loss)
end

function continuous_gen(m::POMDP{S}, s::S, a::CovidAction, rng::AbstractRNG=Random.GLOBAL_RNG) where S <: CovidState
    sp, new_inf, o = sim_step(m, s, a)
    r = reward(m, s, a, sp)
    o = sum(o)

    return (sp=sp, o=o, r=r)
end

function POMDPs.gen(pomdp::POMDP{S}, s::S, a::CovidAction, rng::AbstractRNG=Random.GLOBAL_RNG) where S <: CovidState
    rsum = 0.0
    local o::Int
    for i in 1:pomdp.test_period
        s,o,r = continuous_gen(pomdp, s, a, rng)
        rsum += r
    end
    return (sp=s, o=o, r=rsum)
end

#=
TODO: Find type stable way to do this with Normal() ?
Currently, zero testing prop results in obersvation distribution of Normal(0,0),
resulting in NaNs for weights when normalizing.
=#
struct UninformedDist end

Distributions.pdf(::UninformedDist, ::Any) = 1.0

POMDPs.actions(pomdp::AbstractCovidPOMDP) = pomdp.actions

POMDPs.discount(pomdp::AbstractCovidPOMDP) = pomdp.discount


text:
const LY_JULIAFORMATTER = (
    overwrite = true,
    indent = 4,
    margin = 100,
    always_for_in = false,
    whitespace_typedefs = false,
    whitespace_ops_in_indices = false,
    remove_extra_newlines = false,
)

const LY_PKGTEMPLATE = (
    user = "Lyceum",
    host = "github.com",
    license = "MIT",
    authors = ["Colin Summers"],
    ssh = true,
    manifest = true,
    julia_version = VERSION,
    dev = false,
)

const LY_GITCONFIG = Dict("user.name" => "lyceum", "user.email" => "lyceumdevs@gmail.com")

const LY_REGISTRY = RegistrySpec(
    uuid = "96ca1a0e-015d-46b7-a12d-07d1320642ca",
    name = "LyceumRegistry",
    url = "https://github.com/Lyceum/LyceumRegistry.git",
)

const GENERAL_REGISTRY = RegistrySpec(
    uuid = "23338594-aafe-5451-b93e-139f81909106",
    name = "General",
    url = "https://github.com/JuliaRegistries/General.git",
)


text:
## matern.jl : implementation of Mat\'ern covariance function?

## Matern ##
struct Matern{T} <: IsotropicCovarianceStructure{T}
    λ::T
    ν::T
    σ::T
    p::T

    function Matern{T}(λ::T, ν::T, σ::T, p::T) where T
        λ > 0 || throw(DomainError(λ, "correlation length λ of Mat\u00E9rn covariance cannot be negative or zero"))
        ν > 0 || throw(DomainError(ν, "smoothness ν of Mat\u00E9rn covariance cannot be negative or zero"))
        σ > 0 || throw(DomainError(σ, "marginal standard deviation σ of Mat\u00E9rn covariance cannot be negative or zero"))
        p >= 1 || throw(DomainError(p, "in p-norm, p must be greater than or equal to 1"))
    	isinf(p) && throw(DomainError(p, "in p-norm, p cannot be infinity"))

        new{T}(λ, ν, σ, p)
    end
end

"""
    Matern(λ, ν, [σ = 1], [p = 2])

Mat\u00E9rn covariance structure with correlation length `λ`, smoothness `ν`, (optional) marginal standard deviation `σ` and (optional) `p`-norm, defined as

``C(x, y) = σ \\displaystyle\\frac{2^{1 - ν}}{Γ(ν)} \\left(\\frac{\sqrt{2\nu}ρ}{λ}\\right)^ν K_ν\\left(\\frac{\sqrt{2\nu}ρ}{λ}\\right)``

with ``ρ = ||x - y||_p``.

# Examples
```jldoctest
julia> Matern(0.1, 1.0)
Matérn (λ=0.1, ν=1.0, σ=1.0, p=2.0)

julia> Matern(1, 1, σ=2.0)
Matérn (λ=1.0, ν=1.0, σ=2.0, p=2.0)

```
See also: [`Exponential`](@ref), [`Linear`](@ref), [`Spherical`](@ref), [`Whittle`](@ref), [`Gaussian`](@ref), [`SquaredExponential`](@ref)
"""
Matern(λ::Real, ν::Real; σ::Real=1.0, p::Real=2) =
    Matern{promote_type(typeof(λ),typeof(ν),typeof(σ),typeof(p))}(promote(λ, ν, σ, p)...)

# TODO : Evaluation of the besselk function is extremely expensive; maybe try to catch special cases for ν = 1/2, 3/2 etc.
# evaluate Matern covariance
function apply(m::Matern, x::Real)
    if iszero(x)
        float(one(x))
    else
        2^(1 - m.ν) / gamma(m.ν) * (sqrt(2 * m.ν) * x / m.λ)^m.ν * besselk(m.ν, sqrt(2 * m.ν) * x / m.λ)
    end
end

# short name
shortname(::Matern) = "Mat\u00E9rn"


text:
using Flux, Zygote
using LinearAlgebra
using jInv.Mesh
using Printf
using Plots
using JLD
using Revise
using MFGnet
include("viewers.jl")
include("runOMThelpers.jl")

# set default experimental parameters
@isdefined(R)  ? R = R : R = Float64                                  # element type. Use Float32 or Float64
@isdefined(d)  ? d = d : d = 2                                        # dimension of problem
@isdefined(ar) ? ar = ar : ar = x-> R.(x)                             # function to change element type of arrays  and device (e.g., to cuArray)
@isdefined(m)           ? m=m               : m=16                   # width of network
@isdefined(nTh)         ? nTh=nTh           : nTh=2                   # number of nodes in ResNet discretization
@isdefined(nTrain)      ? nTrain=nTrain     : nTrain=16^2
@isdefined(nVal)        ? nVal=nVal         : nVal=minimum([64^2,nTrain]) # number of training samples
@isdefined(nt)          ? nt=nt             : nt=4                    # number of time steps for characteristics
@isdefined(stepper)     ? stepper=stepper   : stepper=RK4Step()       # time stepping for characteristics
@isdefined(T)           ? T=T               : T=1.0                   # final time for dynamical OT
@isdefined(alph)        ? α=alph            : α=[1.0,2.0,4.0,2.0,2.0] # weights for objective
@isdefined(sampleFreq)  ? sampleFreq = sampleFreq : sampleFreq = 25   # sample frequency
@isdefined(saveIter)    ? saveIter = saveIter     : saveIter   = 25   # iteration to save weights
@isdefined(maxIter)     ? maxIter  = maxIter      : maxIter = 200     # max number of iters
@isdefined(optim)       ? optim  = optim        : optim = :bfgs       # optimization algorithm
@isdefined(sigQ)        ? sigQ = sigQ           : sigQ = R.([1.0; 0.5]) # variance of obstacle Gaussian
@isdefined(Qheight)     ? Qheight = Qheight     : Qheight = 50.0 # height/mass of obstacle Gaussian
@isdefined(muFp)        ? muFp = muFp           : muFp = 1.0     # penalty for F preference
@isdefined(muFe)        ? muFe = muFe           : muFe = 1e-2    # penalty for F entropy
@isdefined(doPlots)     ? doPlots  = doPlots      : doPlots = true     # max number of iters

@isdefined(saveStr)     ? saveStr=saveStr   : saveStr = "Obstacle-BFGS-d-"*string(d)

mu0  = R.([0.0;3.0;zeros(R,d-2)])
sig0 = R.(0.3*ones(d))
rho0 = Gaussian(d,sig0,mu0)

mu1  = R.([0.0;-3.0;zeros(R,d-2)])
sig1 = R.(0.3*ones(d))
rho1 = Gaussian(d,sig1,mu1)

# sigQ = R.(0.5*ones(d))
# Q1   = Gaussian(d,sigQ,zeros(R,d),R(Qheight))
Q1   = Gaussian(2,sigQ,zeros(R,2),R(Qheight))
# Q(x) = vec((sqrt((2*pi)^d * prod(sigQ))).*Q1(x[1:d,:]))
Q(x) = Q1(x[1:2,:])
# println("sigQ !")

if doPlots
    # validation points on regular grid
    domain = 5.0*[-1 1 -1 1];
    M      = getRegularMesh(domain,[sqrt(nVal),sqrt(nVal)])
    X0val     = Matrix(getCellCenteredGrid(M)')
    X0val   = ar([X0val; zeros(d-2,size(X0val,2))]); nVal = size(X0val,2)
    rho0xVal = rho0(X0val)
    rho1xVal = rho1(X0val)
    vol     = 10^d
    wVal    = (vol/nVal) * rho0xVal
else
    # validation from samples
    X0val   = sample(rho0,nVal)
    wVal    = ar(1/nVal * ones(nVal))
    println("sum(wVal)=$(sum(wVal))")
end

if norm(sum(wVal) - R(1)) > 0.1
    error("Double check the weights in validation loss.")
end

α = ar.(α)
tspan = [0.0, T]
Φ = getPotentialResNet(nTh,1.0,nTh,R)
(w0,ΘN,A0,b0,z0) = initializeWeights(d,m,nTh,ar)

rho0xv = rho0(X0val)
rho1xv = rho1(X0val)

F1v  = Fp(Q,rho0,rho0xv,R.(muFp))
F2v  = Fe(rho0,rho0xv,R.(muFe))
Fv   = Fcomb([F1v;F2v])
Gv   = Gkl(rho0,rho1,rho0xv,rho1xv,R.(1.0))
Jv  = MeanFieldGame(Fv,Gv,X0val,rho0,wVal,Φ=Φ,stepper=stepper,nt=nt,α=α,tspan=tspan)

Θ = (w0,(ΘN),A0,b0,z0)

parms = MFGnet.myMap(x->x,Θ)
ps = Flux.params(parms)

println("\n\n ---------- Obstacle Driver -------------\n\n")
println("results stored in: $(pwd()*"/"*saveStr)-level-1.jld")
println("sampleFreq = $(sampleFreq), α = $(α), nTh = $(nTh), m = $(m), nt = $(nt), Qheight = $(Qheight), sigQ=$(sigQ)")
println("DIMENSION = $(d), nTrain = $(nTrain), nVal = $(nVal), saveIter = $(saveIter), maxIter = $(maxIter), muFe=$(muFe), muFp=$(muFp)\n\n")
println("optim: $optim")

global bestLoss = Inf
global Θtemp = initializeWeights(d,m,nTh,ar)
global Θbest = initializeWeights(d,m,nTh,ar)

cb = function(J,Jv,iter,His,parms,doPlots=doPlots)
    global bestLoss
    global Θtemp
    global Θbest

    θc = parms
    # Jvc = Jv(θc)
    Jvc = Jv(parms)

    if Jvc < bestLoss
        bestLoss = Jvc
        # println("loss:", bestLoss)
        Θtemp = copy(MFGnet.param2vec(parms))
        Θbest = MFGnet.vec2param!(Θtemp,Θbest)
        # println("Θopt:", Θopt[1][1:10])
    end

    Hisk = [J.cs[1] J.cs[2] J.cs[3] J.cs[4] J.cs[5] Jv.cs[1] Jv.cs[2] Jv.cs[3] Jv.cs[4] Jv.cs[5]]
    His[iter,:] = Hisk
    if iter==1
        str1 = @sprintf("iter\tobj\t\tcostL\t\tcostF\t\tcostG\t\tcostHJ\t\tcostHJfinal\t\t\tobjVal\t\tcostLVal\tcostFVal\tcostGVal\tcostHJVal\tcostHJfinalVal\n")
        str2 = @sprintf("%05d\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t\t\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t%1.3e",
                iter,sum(His[iter,1:5]),His[iter,1],His[iter,2],His[iter,3],His[iter,4],His[iter,5],
                sum(His[iter,6:10]),His[iter,6],His[iter,7],His[iter,8],His[iter,9],His[iter,10])
        str = str1*str2
    else
        str = @sprintf("%05d\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t\t\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t%1.3e\t%1.3e",
                iter,sum(His[iter,1:5]),His[iter,1],His[iter,2],His[iter,3],His[iter,4],His[iter,5],
                sum(His[iter,6:10]),His[iter,6],His[iter,7],His[iter,8],His[iter,9],His[iter,10])
    end
    println(str)

    # save every saveIter iterations
    if mod(iter,saveIter)==0
        println("saving...")
        currentSaveStr = saveStr *"-iter"*string(iter)*".jld"
        settings = (m,α,nTrain,R,d,mu0,sig0,mu1,sig1,domain,nTh,m,nt,T,Qheight,sigQ,muFe,muFp)
        # save(currentSaveStr, "Θ", Θopt,"settings",settings, "His",His[1:iter,:])
        save(currentSaveStr,"Θbest",Θbest,"His",His,"settings",settings)
    end

    if mod(iter,sampleFreq)==0;
        # println("resampling (symmetric)..")
        # X0      = sample(rho0,Int(nTrain/2))
        # X0      = X0 .- mu0
        # X0      = [X0 -X0]
        # X0      = X0 .+ mu0
        println("regular sampling...")
        X0    = sample(rho0,size(J.X0,2))
        J.X0    = X0
        # J.X0    = [J.X0 -J.X0]
        J.rho0x = rho0(J.X0)
        J.G.rho0x = rho0(J.X0)
        J.G.rho1x = rho1(J.X0)
    end

    if doPlots
        p1= viewImage2D(Jv.rho0x,M,aspect_ratio=:equal)
        title!("rho0(x)")
        p2= viewImage2D(Jv.G.rho1x,M,aspect_ratio=:equal)
        title!("rho1(x)")
        detDy = exp.(Jv.UN[d+1,:])
        rho1y = (Jv.G.rho1(Jv.UN[1:d,:]).*detDy)
        # p3= viewImage2D(rho1y ,M,aspect_ratio=:equal,clims=(minimum(Jv.rho0x),maximum(Jv.rho0x)))
        p3= viewImage2D(rho1y ,M,aspect_ratio=:equal)
        title!("rho1(y).*det")

        res = Jv.rho0x - rho1y
        p4 = viewImage2D(abs.(res),M,aspect_ratio=:equal)
        title!("rho0(x)-rho1(y).*det") # where's the determinant

        p5 = viewImage2D(detDy,M,aspect_ratio=:equal)
        title!("det")

        xc = [ Jv.X0; fill(0.0,1,size(Jv.X0,2))]
        # Φ0 = Jv.Φ(xc,θc)
        # p6 = viewImage2D(Φ0,M,aspect_ratio=:equal)
        # title!("Phi(x,0)")

        ################
        p6 = viewImage2D(Q(Jv.X0), M, aspect_ratio=:equal)
        title!("Potential Q(x)")

        xc = ar([ Jv.UN[1:d,:]; fill(1.0,1,size(Jv.X0,2))])
        Φ1 = vec(Jv.Φ(xc,θc))
        p7 = viewImage2D(Φ1,M,aspect_ratio=:equal)
        title!("Phi(z,1)")

        deltaG = Jv.α[3]*MFGnet.getDeltaG(Jv.G,Jv.UN)
        p8 = viewImage2D(deltaG,M,aspect_ratio=:equal)
        title!("deltaG(z)")
        p9 = viewImage2D(abs.(Φ1-deltaG),M,aspect_ratio=:equal)
        title!("Phi1 - deltaG(z)")

        Xt    = sample(rho0,10)
        wt     = 1/size(Xt,2) * ones(size(Xt,2))
        rho0xt = rho0(Xt)
        rho1xt = rho1(Xt)
        F1t  = Fp(Q,rho0,rho0xt,ar(1.0))
        F2t  = Fe(rho0,rho0xt,ar(1e-1))
        Ft   = Fcomb([F1t;F2t])
        Gt = Gkl(rho0,rho1,rho0xt,rho1xt,ar(1.0))
        Jt = MeanFieldGame(Ft,Gt,Xt,rho0,wt,Φ=Φ,stepper=stepper,nt=nt,α=α,tspan=tspan)

        # p20 = viewImage2D(J.rho0x ,M,aspect_ratio=:equal)
        ntv = Jt.nt
        Ut = MFGnet.integrate2(Jt.stepper,MFGnet.odefun,Jt,ar([Xt;zeros(4,size(Xt,2))]),θc,ar(tspan),ntv)
        Y0  = ar([Ut[1:d,:,end];zeros(4,size(Xt,2))])
        # U0 = MFGnet.integrate2(Jt.stepper,MFGnet.odefun,Jt,Y0,Θopt,ar([T 0.0]),ntv)

        for j=1:size(Xt,2)
            uj = Ut[1:d,j,:]
            # vj = U0[1:d,j,:]
            plot!(p6,uj[1,:],uj[2,:],legend=false,linecolor=:white,linewidth=1.5)
            # plot!(p1,vj[1,:],vj[2,:],legend=false,linecolor=:red,linewidth=1.5)
        end

        pt = plot(p1,p2,p3,p4,p5,p6,p7,p8,p9)
        display(pt)
    end
end

for level=1:length(nTrain)
    X0    = sample(rho0,nTrain[level])
    w     = ar(1/nTrain[level] * ones(nTrain[level]))
    println("sum(w)=$(sum(w))")
    if norm(sum(w) - R(1)) > 0.1
        error("Double check the weights in training loss.")
    end
    F1  = Fp(Q,rho0(X0),rho1(X0),ar(muFp))
    F2  = Fe(rho0,rho0(X0),ar(muFe))
    F   = Fcomb([F1;F2])
    G = Gkl(rho0,rho1,rho0(X0),rho1(X0),ar(1.0))
    J = MeanFieldGame(F,G,X0,rho0,w,Φ=Φ,stepper=stepper,nt=nt,α=α,tspan=tspan)

    His = zeros(maxIter[level],10)
    cbBFGS = (iter)-> cb(J,Jv,iter,His,parms,doPlots)
    # cbBFGS(1)

    # # Optimize

    if optim==:bfgs
        f   =(Θ)-> evalObj(J,Θ,parms,ps)
        fdf =(Θ)-> evalObjAndGrad(J,Θ,parms,ps)
        Θ0 = MFGnet.param2vec(parms)
        f(Θ0)
        fdf(Θ0)
        runtime = @elapsed Θopt,flag,His,X,H = MFGnet.bfgs(f,fdf,Θ0,maxIter=size(His,1),
                                                   out=0,atol=1e-10,cb=cbBFGS)

    else
        println("using ADAM")
        opt = ADAM()
        runtime = @elapsed for k=1:maxIter
            if mod(k,sampleFreq)==0
                Xt    = sample(J.rho0,size(J.X0,2))
                J.X0  = Xt;
                J.w   = (1/nTrain) * ones(R,nTrain);
                J.G.rho0x = J.G.rho0(Xt)
                J.G.rho1x = J.G.rho1(Xt)
                J.rho0x = J.rho0(Xt)
            end

            Jc,back = Zygote.pullback(() -> J(parms), ps)
            grads = back(1)
            ng = 0.0
            for p in ps
                ng += norm(grads[p]).^2
                Zygote.Tracker.update!(opt,p, grads[p])
            end
            cbBFGS(k)
        end
    end
    println("average time per iteration: $(runtime/maxIter)")
    Θopt = MFGnet.param2vec(parms)
    Θopt = MFGnet.vec2param!(Θopt,Θ)

    settings = (m,α,nTrain[level],R,d,mu0,sig0,mu1,sig1,domain,nTh,m,nt,T,Qheight,sigQ,muFe,muFp)

    save(saveStr*"-level-$level.jld","Θopt", Θopt,"Θbest",Θbest,"His",His,"settings",settings)
end

max_stars_repo_name: EmoryMLIP/MFGnet

max_issues_repo_name: EmoryMLIP/MFGnet

max_forks_repo_name: EmoryMLIP/MFGnet


text:
# Overview
#     Standalone script designed to produce plots to compare the calculation of the absorption
#     coefficient, using data from the Opacity Project, against absorption coefficients taken
#     from other sources. Currently, this just compares the absorption coefficients computed for
#     Hydrogenic atoms.

using Plots, Korg
include("../test/absorption_comparison_funcs.jl")

function _plot_x_intervals!(p, x_intervals)
    y_lim = ylims(p)
    for i in 1:size(x_intervals)[1]
        cur_x = [x_intervals[i,1], x_intervals[i,2]]
        cur_y1 = [y_lim[1], y_lim[1]]
        cur_y2 = [y_lim[2], y_lim[2]]
        plot!(p, cur_x, cur_y1, fillrange = cur_y2, fillalpha = 0.10,
              ylim = y_lim, fillcolor = :grey,
              label = "")
    end
end

function plot_hydrogenic_comparison(species_name, T = 7800.0, ndens_species = 3.0e16; p = nothing)
    λ_vals = OP_compare._dflt_λ_vals(species_name)
    hydrogenic_α_OP = OP_compare.calc_hydrogenic_bf_absorption_coef(λ_vals,  T, ndens_species,
                                                                    species_name;
                                                                    use_OP_data = true)
    hydrogenic_α_dflt = OP_compare.calc_hydrogenic_bf_absorption_coef(λ_vals,  T, ndens_species,
                                                                      species_name;
                                                                      use_OP_data = false)
    _pos_value_extrema(arr) = extrema(filter(x->x>0, arr))
    ylims = extrema((_pos_value_extrema(hydrogenic_α_OP)...,
                     _pos_value_extrema(hydrogenic_α_dflt)...))
    if isnothing(p)
        p = plot()
    end
    p = plot!(p, λ_vals, [hydrogenic_α_OP, hydrogenic_α_dflt],
              labels = ["Opacity Project" "Reference (Hydrogenic)"],
              yscale = :log10, xscale = :log10, ylims = ylims,
              ylabel = "absorption coef, α (cm⁻¹)", xlabel = "λ (Ångstroms)",
              title = string(join(split(species_name,'_'), " "), " Comparison"))
    if species_name == "H_I"
        vline!(p, [911.75, 3646, 8204, 14580, 22790, 32820], label = "known binding energies")
    end
    λ_comp_intervals = OP_compare._λ_comp_intervals(species_name)
    _plot_x_intervals!(p, λ_comp_intervals)

    cleaned_name = join(split(species_name,'_'), " ")
    title_str = string("Opacity Project Comparison: $cleaned_name\n",
                       "T = $T K, n($cleaned_name) = $ndens_species cm⁻³")
    title!(p,title_str)

    return p
end

function main()
    T = 7800.0 # kind of arbitrary
    ndens_species = 3.0e16 # kind of arbitrary
    for species_name in ["H_I", "He_II"]
        plot_fname = joinpath(@__DIR__, string(species_name, "_comp.png"))
        println("\nCreating Comparison plot at T = $T K, ndens_species = $ndens_species cm⁻³ for ",
                species_name)
        p = plot_hydrogenic_comparison(species_name, T, ndens_species);
        println("Writing file to $plot_fname")
        savefig(p, plot_fname)
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    main()
end


text:
include("../roblib.jl")

function mark_rectangle_perimeter(r)
    num_steps = through_rectangles_into_angle(r, (Down, Left))

    direction = Right

    while (isborder(r, Up) == false)
        if isborder(r, direction)
            direction = inverse(direction)
            move!(r, Up)
        else move!(r, direction)
        end
    end

    while (ismarker(r) == false)
        putmarker!(r)
        if isborder(r, right(direction))
            move!(r, direction)
        else
            direction = right(direction)
            move!(r, direction)
        end
    end

    through_rectangles_into_angle(r, (Down, Left))
    movements!(r, (Up, Right), reverse(num_steps))
end


text:
export keyPress, keyWait, spaceWait
export mousePress, mouseRelease, mouseX, mouseY

function keyPress(window, button)
	GLFW.PollEvents()
	if GLFW.GetKey(window, button) == GLFW.PRESS
		return true
	else
		return false
	end
end

function spaceWait(window)
	while true
		GLFW.WaitEvents()
		if GLFW.GetKey(window, GLFW.KEY_SPACE) == GLFW.PRESS
			break
		end
	end
end

function mousePress(window, button)
	GLFW.PollEvents()
	if GLFW.GetMouseButton(window, button) == GLFW.PRESS
		return true
	else
		return false
	end
end

function mouseRelease(window, button)
	GLFW.PollEvents()
	if GLFW.GetMouseButton(window, button) == GLFW.RELEASE
		return true
	else
		return false
	end
end

function mouseX(window)
	GLFW.PollEvents()
	x, y = GLFW.GetCursorPos(window)
	return x
end

function mouseY(window)
	GLFW.PollEvents()
	x, y = GLFW.GetCursorPos(window)
	return y
end


text:
repo = joinpath(homedir(), "ReferenceImages", "gallery")

recordings = joinpath(@__DIR__, "test_recordings")

cp(recordings, repo, force = true)

MakieGallery.generate_thumbnails(repo)


text:
using LatPhysPlottingPyPlot
using Test

@testset "LatPhysPlottingPyPlot.jl" begin
    # Write your own tests here.
end


text:
include("LargeField/basis.jl")
include("LargeField/FB.jl")
include("LargeField/misc2.jl")


text:
mutable struct DemographicParity <: Measure
    C::Int
    A::Array # Adp Matrix/array which is on the LHS of equation
    # This matrix A has been kept in the struct for efficiency in case of multiple calls on a single dataset
    # Add a constructor for zero matrix here
    DemographicParity() = new(0, [])
end

function (dp::DemographicParity)(ft::FairTensor{C}) where C
    if dp.C==0
        # Here matrix A of DemographicParity will be initialised as it wasn't previously initialised
        dp.C = C
        dp.A = zeros(Int, C, 8)
        Nc = sum(ft[C, :, :])
        for i in 1:C-1
            Ni = sum(ft[i, :, :])
            dp.A[i, 1] = Ni
            dp.A[i, 3] = Ni
            dp.A[i, 5] = -Nc
            dp.A[i, 7] = -Nc
        end
    end
    z = reshape(ft.mat, C, 4)
    z_c = reshape(repeat(z[C, :], C), 4, C)
    z = hcat(z, transpose(z_c))
    return all(dp.A*transpose(z) .== 0)
end


text:
module YeSQL

using DocStringExtensions: FIELDS, SIGNATURES, TYPEDEF, TYPEDFIELDS, TYPEDSIGNATURES

export load_path, load_path!, load_string, load_string!

# SQLOperation types for dispatch
"""Abstract parent type"""
abstract type SQLOperation end
"""An operation which inserts a row returns a row of values. Name suffix: `^!`"""
struct InsertReturning <: SQLOperation end
"""An operation which inserts a row returns a value. Name suffix: `\$!`"""
struct InsertReturningValue <: SQLOperation end
"""An operation which inserts/updates/deletes many rows. Name suffix: `*!`"""
struct InsertUpdateDeleteMany <: SQLOperation end
"""An operation which inserts/updates/deletes a row. Name suffix: `!`"""
struct InsertUpdateDelete <: SQLOperation end
"""A sql script run without arguments. Name suffix: `#`"""
struct Script <: SQLOperation end
"""A query returning a single row. Name suffix: `^`"""
struct SelectOne <: SQLOperation end
"""A query returning a single value. Name suffix: `\$`"""
struct SelectValue <: SQLOperation end
"""A select query returning potentially many rows. No suffix."""
struct Select <: SQLOperation end

"""
$(TYPEDEF)

SQL query with fields extracted from a *.sql* file or string

# Fields

$(TYPEDFIELDS)

# Query syntax

Each query should be preceded by comment lines defining the query name
and operation type, and an optional docstring:

```sql
-- name: <name><operation_suffix>
-- <docstring>  (optional)
<query text>
```
"""
struct SQLQuery{O<:SQLOperation}
    "Query name"
    name::Symbol
    "Processed query string"
    query::String
    "Docstring"
    docstring::String
    "Argument list"
    args::Vector{Symbol}
    "`SQLOperation` derived from the operation suffix"
    operation::O
end

function Base.show(io::IO, q::SQLQuery)
    print(io, """$(q.operation): $(q.docstring)\nargs: $(join(q.args, ", "))\n""")
end

"""
    $(TYPEDSIGNATURES)

Interpret `query_name` and return an `SQLOperation` and name as a symbol

```jldoctest
YeSQL._interpret("my_script_query#")

# output

(YeSQL.Script(), :my_script_query)
```
"""
function _interpret(query_name::AbstractString)::Tuple{<:SQLOperation, Symbol}
    (operation, name) = if endswith(query_name, "^!")
        (InsertReturning(), query_name[begin:end-2])
    elseif endswith(query_name, "\$!")
        (InsertReturningValue(), query_name[begin:end-2])
    elseif endswith(query_name, "*!")
        (InsertUpdateDeleteMany(), query_name[begin:end-2])
    elseif endswith(query_name, "!")
        (InsertUpdateDelete(), query_name[begin:end-1])
    elseif endswith(query_name, "#")
        (Script(), query_name[begin:end-1])
    elseif endswith(query_name, "^")
        (SelectOne(), query_name[begin:end-1])
    elseif endswith(query_name, "\$")
        (SelectValue(), query_name[begin:end-1])
    else
        (Select(), query_name)
    end
    return operation, Symbol(name)
end

include("reader.jl")
include("libpq_adapter.jl")

end


text:
module KendallTau
using Random
using BenchmarkTools
using LinearAlgebra # so that identity matrix I is available.

const RealVector{T <: Real} = AbstractArray{T,1}
const RealMatrix{T <: Real} = AbstractArray{T,2}

const RealOrMissingVector{T <: Real} = AbstractArray{<:Union{T, Missing},1}
const RealOrMissingMatrix{T <: Real} = AbstractArray{<:Union{T, Missing},2}

include("rankcorr.jl")
include("skipmissingpairs.jl")#old approach
include("skipmissingpairwise.jl")#new approach
if VERSION >= v"1.5" #@spawn not available on 1.0
    include("threads_v1.jl")
    include("threads_v2.jl")
    include("threads_v3.jl")
    include("threads_v4.jl")
    include("threads_v5.jl")
end
include("speedtests.jl")
include("examplecode.jl")

export corkendall

end # module


text:
module TestDeque

using TapirSchedulers: BoundedWorkStealingDeque, isfull, maybepop!, maybepopfirst!, trypush!
using Test

function test_full()
    q = BoundedWorkStealingDeque{Int}(4)
    @testset for i in 1:3
        @test !isfull(q)
        @test trypush!(q, i)
    end
    @test isfull(q)
    @test !trypush!(q, 4)
    @test maybepop!(q) === Some(3)
    @test trypush!(q, 4)
    @test maybepop!(q) === Some(4)
    @test maybepopfirst!(q) === Some(1)
    @test maybepopfirst!(q) === Some(2)
    @test maybepopfirst!(q) === nothing
    @test maybepop!(q) === nothing
end

end  # module


text:
# how should we apply a function recursively to the stFlux type? to each part of the chain, of course
function mapEvery3(to, ii, x)
    if ii % 3 == 1
        adapt(to, x)
    else
        x
    end
end
import FourierFilterFlux.cu
function cu(stf::stFlux{Dimension,Depth,ChainType,D,E,F}) where {Dimension,Depth,ChainType,D,E,F}
    newChain = Chain((map(iix -> (iix[1] % 3 == 1 ? cu(iix[2]) : iix[2]), enumerate(stf.mainChain)))...)
    return stFlux{Dimension,Depth,typeof(newChain),D,E,F}(newChain, stf.normalize, stf.outputSizes, stf.outputPool, stf.settings)
end
import Adapt.adapt
function adapt(to, stf::stFlux{Dimension,Depth,ChainType,D,E,F}) where {Dimension,Depth,ChainType,D,E,F}
    newChain = Chain((map(iix -> mapEvery3(to, iix...), enumerate(stf.mainChain)))...)
    return stFlux{Dimension,Depth,typeof(newChain),D,E,F}(newChain, stf.normalize, stf.outputSizes, stf.outputPool, stf.settings)
end

function adapt(to, stResult::ScatteredFull{T,N}) where {T,N}
    data = adapt(to, stResult.data)
    output = adapt(to, stResult.output)
    return ScatteredFull{eltype(data),N}(stResult.m, stResult.k, data, output)
end

function adapt(to, stResult::ScatteredOut{T,N}) where {T,N}
    output = adapt(to, stResult.output)
    return ScatteredOut{eltype(output),N}(stResult.m, stResult.k, output)
end
export adapt

"""
wave1, wave2, wave3, ... = getWavelets(sc::stFlux)

just a simple util to get the wavelets from each layer
"""
function getWavelets(sc::stFlux; spaceDomain=false)
    freqDomain = map(x -> x.weight, filter(x -> (typeof(x) <: ConvFFT), sc.mainChain.layers)) # filter to only have ConvFFTs, and then return the wavelets of those
    if spaceDomain
        return map(originalDomain, filter(x -> (typeof(x) <: ConvFFT), sc.mainChain.layers))
    else
        return map(x -> x.weight, filter(x -> (typeof(x) <: ConvFFT), sc.mainChain.layers))
    end
end

import ContinuousWavelets: getMeanFreq
"""
    getMeanFreq(sc::stFlux{1}, δt=1000)
Get a list of the mean frequencies for the filter bank in each layer. Note that δt gives the sampling rate for the input only, and that it decreases at each subsequent layer.
"""
function getMeanFreq(sc::stFlux{1}, δt=1000)
    waves = getWavelets(sc)
    shrinkage = [size(waves[i+1], 1) / size(waves[i], 1) for i = 1:length(waves)-1]
    δts = δt * [1, shrinkage...]
    map(getMeanFreq, waves, δts)
end

"""
    flatten(scatRes) -> output
given `scatRes`, a scattered output or full, it produces a single vector containing the entire transform in order, i.e. the same format as output by thinSt.
"""
function flatten(scatRes::S) where {S<:Scattered}
    return scatRes[:]
end
flatten(scatRes) = scatRes

"""
    rolled = roll(toRoll, st::stFlux)
Given the output of a scattering transform and something with the same number
of entries but in an array that is NCoeffs×extraDims, roll up the output
into an array of arrays like the scattered.
    rolled = roll(toRoll, st::stParallel; percentage=.9, outputSubsample=(-1,-1))
there is also a version for the parallel transform; since `percentage` and
`outputSubsample` are separate variables, they must also be input.
"""
function roll(toRoll, st::stFlux)
    Nd = ndims(st)
    oS = st.outputSizes
    roll(toRoll, oS, Nd)
end

function roll(toRoll, stOutput::S) where {S<:Scattered}
    Nd = ndims(stOutput)
    oS = map(size, stOutput.output)
    return roll(toRoll, oS, Nd)
end

function roll(toRoll, oS::Tuple, Nd)
    nExamples = size(toRoll)[2:end]
    rolled = ([adapt(typeof(toRoll), zeros(eltype(toRoll),
        sz[1:Nd+nPathDims(ii)]...,
        nExamples...)) for (ii, sz) in
               enumerate(oS)]...,)

    locSoFar = 0
    for (ii, x) in enumerate(rolled)
        szThisLayer = oS[ii][1:Nd+nPathDims(ii)]
        totalThisLayer = prod(szThisLayer)
        range = (locSoFar+1):(locSoFar+totalThisLayer)
        addresses = (szThisLayer..., nExamples...)
        rolled[ii][:] = reshape(toRoll[range, :], addresses)
        locSoFar += totalThisLayer
    end
    return ScatteredOut(rolled, Nd)
end

"""
    p = computeLoc(loc, toRoll, st::stFlux)
given a location `loc` in the flattened output `toRoll`, return a
pathLocs describing that location in the rolled version
"""
function computeLoc(loc, toRoll, st::stFlux)
    Nd = ndims(st)
    oS = st.outputSizes
    nExamples = size(toRoll)[2:end]
    locSoFar = 0
    for (ii, x) in enumerate(oS)
        szThisLayer = x[1:Nd+nPathDims(ii)]
        totalThisLayer = prod(szThisLayer)
        if locSoFar + totalThisLayer ≥ loc
            return pathLocs(ii - 1, Tuple(CartesianIndices(szThisLayer)[1+loc-locSoFar]))
        end
        locSoFar += totalThisLayer
    end
end

"""
given a scattered output, make a list that gives the largest value on each path
"""
function importantCoords(scatRes)
    return [dropdims(maximum(abs.(x), dims=(1, 2)), dims=(1, 2)) for x in scatRes]
end

"""
if the batch size is off, we don't want to suddenly drop performance. Split it up.
"""
function batchOff(stack, x, batchSize)
    nRounds = ceil(Int, size(x, 4) // batchSize)
    firstRes = stack(x[:, :, :, 1:batchSize])
    result = cu(zeros(size(firstRes)[1:end-1]..., size(x)[end]))
    result[:, 1:batchSize] = firstRes
    for i = 2:(nRounds-1)
        result[:, 1+(i-1)*batchSize:(i*batchSize)] = stack(x[:, :, :, 1+(i-1)*batchSize:(i*batchSize)])
    end
    result[:, (1+(nRounds-1)*batchSize):end] = stack(cat(x[:, :, :,
            (1+(nRounds-1)*batchSize):end],
        cu(zeros(size(x)[1:3]...,
            nRounds * batchSize
            -
            size(x, 4))),
        dims=4))[:, 1:(size(x, 4)-(nRounds-1)*batchSize)]
    return result
end

"""
given a st object and a symbol `s` representing a possible keyword, e.g. :Q, or :β, return the value for this transform. It may be in st.settings, or, if it is a default value, it is looked up.
"""
function getParameters(st, s)
    get(st.settings, s, default(s))
end

function default(s)
    if s == :dType
        Float32
    elseif s == :σ
        identity
    elseif s == :trainable
        false
    elseif s == :plan
        true
    elseif s == :init
        Flux.glorot_normal
    elseif s == :bias
        false
    elseif s == :convBoundary
        Sym()
    elseif s == :cw
        Morlet()
    elseif s == :averagingLayer
        false
    elseif s == :Q
        8
    elseif s == :boundary
        SymBoundary()
    elseif s == :averagingType
        Father()
    elseif s == :averagingLength
        4
    elseif s == :frameBound
        1
    elseif s == :p
        Inf
    elseif s == :β
        4
    end
end

import Base.size
function size(st::stFlux)
    l = st.mainChain[1]
    if typeof(l.fftPlan) <: Tuple
        sz = l.fftPlan[2].sz
    else
        sz = l.fftPlan.sz
    end
    es = originalSize(sz[1:ndims(l.weight)-1], l.bc)
    return es
end