__precompile__()

module Pardiso

if !isfile(joinpath(@__DIR__, "..", "deps", "deps.jl"))
    error("""please run Pkg.build("Pardiso") before loading the package""")
end

include("../deps/deps.jl")

function show_build_log()
    logfile = joinpath(@__DIR__, "..", "deps", "build.log")
    if !isfile(logfile)
        error("no log file found")
    else
        println(read(logfile, String))
    end
end

using Libdl
using SparseArrays
using LinearAlgebra
import Base.show

if !LOCAL_MKL_FOUND
    import MKL_jll
end

if LinearAlgebra.BLAS.vendor() === :mkl && LinearAlgebra.BlasInt == Int64
    const MklInt = Int64
    const PARDISO_FUNC = :pardiso_64
else
    const MklInt = Int32
    const PARDISO_FUNC = :pardiso
end

if Sys.iswindows()
    const libmkl_rt = "mkl_rt"
elseif Sys.isapple()
    const libmkl_rt = "@rpath/libmkl_rt.dylib"
else
    const libmkl_rt = "libmkl_rt"
end

export PardisoSolver, MKLPardisoSolver
export set_iparm!, set_dparm!, set_matrixtype!, set_solver!, set_phase!, set_msglvl!, set_nprocs!
export get_iparm, get_iparms, get_dparm, get_dparms
export get_matrixtype, get_solver, get_phase, get_msglvl, get_nprocs
export set_maxfct!, set_perm!, set_mnum!
export get_maxfct, get_perm, get_mnum
export checkmatrix, checkvec, printstats, pardisoinit, pardiso
export solve, solve!
export get_matrix
export schur_complement, pardisogetschur
export fix_iparm!

struct PardisoException <: Exception
    info::String
end

struct PardisoPosDefException <: Exception
    info::String
end

Base.showerror(io::IO, e::Union{PardisoException,PardisoPosDefException}) = print(io, e.info);

const PardisoNumTypes = Union{Float64,ComplexF64}

abstract type AbstractPardisoSolver end

function collect_gfortran_lib_candidates(main_ver)
    candidates = String[]
    if Sys.isapple()
        homebrew_gcc = "/usr/local/Cellar/gcc/"
        isdir(homebrew_gcc) || return String[]
        vers = readdir(homebrew_gcc)
        filter!(x -> startswith(x, "$main_ver."), vers)
        for v in vers
            push!(candidates, joinpath(homebrew_gcc, v, "lib/gcc/$main_ver/"))
        end
    elseif Sys.islinux()
        gcc_path = "/usr/lib/gcc/x86_64-linux-gnu/"
        isdir(gcc_path) || return String[]
        vers = readdir(gcc_path)
        filter!(x -> startswith(x, "$main_ver.") || isequal(x, "$main_ver"), vers)
        for v in vers
            push!(candidates, joinpath(gcc_path, v))
        end
    end
    return candidates
end

load_lib_fortran(lib::String, v::Int) = load_lib_fortran(lib, [v])
function load_lib_fortran(lib::String, vs::Vector{Int})
    candidates = String[]
    for v in vs
        append!(candidates, collect_gfortran_lib_candidates(v))
    end
    path = Libdl.find_library(lib, candidates)
    isempty(path) && (path = lib)
    Libdl.dlopen(path * "." * Libdl.dlext, Libdl.RTLD_GLOBAL)
end

# Pardiso
const init = Ref{Ptr}()
const pardiso_f = Ref{Ptr}()
const pardiso_chkmatrix = Ref{Ptr}()
const pardiso_chkmatrix_z = Ref{Ptr}()
const pardiso_printstats = Ref{Ptr}()
const pardiso_printstats_z = Ref{Ptr}()
const pardiso_chkvec = Ref{Ptr}()
const pardiso_chkvec_z = Ref{Ptr}()
const pardiso_get_schur_f = Ref{Ptr}()
const PARDISO_LOADED = Ref(false)

function __init__()
    if !haskey(ENV, "PARDISOLICMESSAGE")
        ENV["PARDISOLICMESSAGE"] = 1
    end

    if LOCAL_MKL_FOUND && !haskey(ENV, "MKLROOT")
        @warn "MKLROOT not set, MKL Pardiso solver will not be functional"
    end

    # This is apparently needed for MKL to not get stuck on 1 thread when
    # libpardiso is loaded in the block below...
    get_nprocs_mkl()

    if PARDISO_LIB_FOUND
        try
            libpardiso = Libdl.dlopen(PARDISO_PATH)
            init[] = Libdl.dlsym(libpardiso, "pardisoinit")
            pardiso_f[] = Libdl.dlsym(libpardiso, "pardiso")
            pardiso_chkmatrix[] = Libdl.dlsym(libpardiso, "pardiso_chkmatrix")
            pardiso_chkmatrix_z[] = Libdl.dlsym(libpardiso, "pardiso_chkmatrix_z")
            pardiso_printstats[] = Libdl.dlsym(libpardiso, "pardiso_printstats")
            pardiso_printstats_z[] = Libdl.dlsym(libpardiso, "pardiso_printstats_z")
            pardiso_chkvec[] = Libdl.dlsym(libpardiso, "pardiso_chkvec")
            pardiso_chkvec_z[] = Libdl.dlsym(libpardiso, "pardiso_chkvec_z")
            pardiso_get_schur_f[] = Libdl.dlsym(libpardiso, "pardiso_get_schur")

            if Sys.isunix()
                gfortran_v = [8, 9]
                for lib in ("libgfortran", "libgomp")
                    load_lib_fortran(lib, gfortran_v)
                end
            end

            # Windows Pardiso lib comes with BLAS + LAPACK prebaked but not on UNIX so we open them here
            # if not MKL is loaded
            if Sys.isunix()
                ptr = C_NULL
                for l in ("libblas", "libblas.so.3")
                    ptr = Libdl.dlopen_e(l, Libdl.RTLD_GLOBAL)
                    if ptr !== C_NULL
                        break
                    end
                end
                if ptr == C_NULL
                    error("could not load blas library")
                end
            end
            PARDISO_LOADED[] = true
        catch e
            @error("Pardiso did not manage to load, error thrown was: $(sprint(showerror, e))")
        end
    end
end
include("enums.jl")
include("project_pardiso.jl")
include("mkl_pardiso.jl")

# Getters and setters
set_matrixtype!(ps::AbstractPardisoSolver, v::Int) = set_matrixtype!(ps, MatrixType(v))
function set_matrixtype!(ps::AbstractPardisoSolver, v::MatrixType)
    ps.mtype = v
end

get_matrixtype(ps::AbstractPardisoSolver) = ps.mtype
get_iparm(ps::AbstractPardisoSolver, i::Integer) = ps.iparm[i]
get_iparms(ps::AbstractPardisoSolver) = ps.iparm
set_iparm!(ps::AbstractPardisoSolver, i::Integer, v::Integer) = ps.iparm[i] = v

get_mnum(ps::AbstractPardisoSolver) = ps.mnum
set_mnum!(ps::AbstractPardisoSolver, mnum::Integer) = ps.mnum = mnum

get_maxfct(ps::AbstractPardisoSolver) = ps.maxfct
set_maxfct!(ps::AbstractPardisoSolver, maxfct::Integer) = ps.maxfct = maxfct

get_perm(ps::AbstractPardisoSolver) = ps.perm
set_perm!(ps::AbstractPardisoSolver, perm::Vector{T}) where {T <: Integer} = ps.perm = convert(Vector{Int32}, perm)

get_phase(ps::AbstractPardisoSolver) = ps.phase

set_phase!(ps::AbstractPardisoSolver, v::Int) = set_phase!(ps, Phase(v))
function set_phase!(ps::AbstractPardisoSolver, v::Phase)
    ps.phase = v
end

get_msglvl(ps::AbstractPardisoSolver) = ps.msglvl

set_msglvl!(ps::AbstractPardisoSolver, v::Integer) = set_msglvl!(ps, MessageLevel(v))
function set_msglvl!(ps::AbstractPardisoSolver, v::MessageLevel)
    ps.msglvl = v
end

function pardisoinit(ps::AbstractPardisoSolver)
    ccall_pardisoinit(ps)
    return
end

function solve(ps::AbstractPardisoSolver, A::SparseMatrixCSC{Tv,Ti},
               B::StridedVecOrMat{Tv}, T::Symbol=:N) where {Ti, Tv <: PardisoNumTypes}
  X = copy(B)
  solve!(ps, X, A, B, T)
  return X
end

function fix_iparm!(ps::AbstractPardisoSolver, T::Symbol)
    # We need to set the transpose flag in PARDISO when we DON'T want
    # a transpose in Julia because we are passing a CSC formatted
    # matrix to PARDISO which expects a CSR matrix.
    if T === :N
        if isa(ps, PardisoSolver)
            set_iparm!(ps, 12, 1)
        else
            # iparm[12] = 1 is complex conjugate in MKL
            set_iparm!(ps, 12, 2)
        end
    elseif T === :C || T === :T
        set_iparm!(ps, 12, 0)
    else
        throw(ArgumentError("only :T, :N and :C, are valid transpose symbols"))
    end
end

function solve!(ps::AbstractPardisoSolver, X::StridedVecOrMat{Tv},
                A::SparseMatrixCSC{Tv,Ti}, B::StridedVecOrMat{Tv},
                T::Symbol=:N) where {Ti, Tv <: PardisoNumTypes}
    set_phase!(ps, ANALYSIS_NUM_FACT_SOLVE_REFINE)

    # This is the heuristics for choosing what matrix type to use
    ##################################################################
    # - If hermitian try to solve with symmetric positive definite.
    #   - On pos def exception, solve instead with symmetric indefinite.
    # - If complex and symmetric, solve with symmetric complex solver
    # - Else solve as unsymmetric.
    if ishermitian(A)
        eltype(A) == Float64 ? set_matrixtype!(ps, REAL_SYM_POSDEF) : set_matrixtype!(ps, COMPLEX_HERM_POSDEF)
        pardisoinit(ps)
        fix_iparm!(ps, T)
        try
            pardiso(ps, X, get_matrix(ps, A, T), B)
        catch e
            set_phase!(ps, RELEASE_ALL)
            pardiso(ps, X, A, B)
            set_phase!(ps, ANALYSIS_NUM_FACT_SOLVE_REFINE)
            if !isa(e, PardisoPosDefException)
                rethrow()
            end
            eltype(A) == Float64 ? set_matrixtype!(ps, REAL_SYM_INDEF) : set_matrixtype!(ps, COMPLEX_HERM_INDEF)
            pardisoinit(ps)
            fix_iparm!(ps, T)
            pardiso(ps, X, get_matrix(ps, A, T), B)
        end
    elseif issymmetric(A)
        set_matrixtype!(ps, COMPLEX_SYM)
        pardisoinit(ps)
        fix_iparm!(ps, T)
        pardiso(ps, X, get_matrix(ps, A, T), B)
    else
        eltype(A) == Float64 ? set_matrixtype!(ps, REAL_NONSYM) : set_matrixtype!(ps, COMPLEX_NONSYM)
        pardisoinit(ps)
        fix_iparm!(ps, T)
        pardiso(ps, X, get_matrix(ps, A, T), B)
    end

    # Release memory, TODO: We are running the convert on IA and JA here
    # again which is unnecessary.
    set_phase!(ps, RELEASE_ALL)
    pardiso(ps, X, A, B)

    set_phase!(ps, ANALYSIS_NUM_FACT_SOLVE_REFINE)
    return X
end

function get_matrix(ps::AbstractPardisoSolver, A, T)
    mtype = get_matrixtype(ps)

    if isposornegdef(mtype)
        if ps isa MKLPardisoSolver
            T == :C && return conj(tril(A))
            return tril(A)
        elseif ps isa PardisoSolver
            T == :T && return tril(A)
            return conj(tril(A))
        end
    end

    if !issymmetric(mtype)
        T == :C && return conj(A)
        return A
    end

    if mtype == COMPLEX_SYM
        T == :C && return conj(tril(A))
        return tril(A)
    end

    error("Unhandled matrix type")
end

function pardiso(ps::AbstractPardisoSolver, X::StridedVecOrMat{Tv}, A::SparseMatrixCSC{Tv,Ti},
                 B::StridedVecOrMat{Tv}) where {Ti, Tv <: PardisoNumTypes}
    if length(X) != 0
        dim_check(X, A, B)
    end

    if Tv <: Complex && isreal(get_matrixtype(ps))
        throw(ErrorException(string("input matrix is complex while PardisoSolver ",
                                    "has a real matrix type set: $(get_matrixtype(ps))")))
    end

    if Tv <: Real && !isreal(get_matrixtype(ps))
        throw(ErrorException(string("input matrix is real while PardisoSolver ",
                                    "has a complex matrix type set: $(get_matrixtype(ps))")))
    end

    N = size(A, 2)

    AA = A.nzval
    IA = convert(Vector{Int32}, A.colptr)
    JA = convert(Vector{Int32}, A.rowval)

    resize!(ps.perm, size(B, 1))

    NRHS = size(B, 2)

    ccall_pardiso(ps, N, A.nzval, A.colptr, A.rowval, NRHS, B, X)
end

pardiso(ps::AbstractPardisoSolver) = ccall_pardiso(ps, Int32(0), Float64[], Int32[], Int32[], Int32(0), Float64[], Float64[])
function pardiso(ps::AbstractPardisoSolver, A::SparseMatrixCSC{Tv,Ti}, B::StridedVecOrMat{Tv}) where {Ti, Tv <: PardisoNumTypes}
    pardiso(ps, Tv[], A, B)
end

# populated rows of S determine schur complment block
"""
    schur_complement(ps,A,x) -> S

Schur complement `S` of the submatrix defined by the nonzero entries of `x` in matrix `A`.
If `n=nnz(x)`, then `S` is `n`-by-`n`.

WARNING: for complex `M`, seems to be unstable, made worse as number of nonzero elements in `M` decreases
"""
schur_complement(ps::AbstractPardisoSolver,A,x::SparseVector,T::Symbol=:N) = _schur_complement_permuted(ps,A,x.nzind,T)
schur_complement(ps::AbstractPardisoSolver,A,x::SparseMatrixCSC,T::Symbol=:N) = _schur_complement_permuted(ps,A,unique!(sort!(x.rowval)),T)

# permute A and then compute complement of lower right-hand `n`-by-`n` block
function _schur_complement_permuted(ps,A,rows,T::Symbol)
    P = sparse(vcat(setdiff(1:A.n,rows),rows),1:A.n,1,size(A)...)
    schur_complement(ps,P'*A*P,length(rows),T)
end
# or integer gives last n rows and columns as schur complement block
"""
    schur_complement(ps,M,n) -> S

Schur complement `S` of upper-left block in `M`, where `n` is the size of lower-right block (and therefore also of Schur complement)

WARNING: for complex `M`, seems to be unstable, made worse as number of nonzero elements in `M` decreases
"""
function schur_complement(ps::AbstractPardisoSolver,A::SparseMatrixCSC{Tv},n::Integer,T::Symbol=:N) where Tv <: PardisoNumTypes

    n ≥ size(A,1) ? throw(ErrorException("complement block size n=$n≥A.m=$(A.m)")) : nothing
    # Tv<:Complex ? (@warn "unstable for complex types, unknown why") : nothing

    pardisoinit(ps)
    original_phase = get_phase(ps)
    original_iparms = get_iparms(ps)
    set_iparm!(ps,1,1) # use custom IPARM
    set_iparm!(ps,38,n) # set Schur complement block size to n
    set_phase!(ps,12) # analyze and factorize
    B = Matrix{Tv}(undef,size(A,1),0) # dummy array to feed to pardiso

    # transpose via IPARM(12) doesn't work at factorize step (only on entry to solve step)
    if T==:N
        M = permutedims(A)
        set_iparm!(ps, 12, 1)
    elseif T == :C
        M = conj(permutedims(A))
        set_iparm!(ps, 12, 0)
    elseif T == :T
        M = A
        set_iparm!(ps, 12, 0)
    else
        throw(ArgumentError("only :T, :N and :C, are valid transpose symbols"))
    end

    pardiso(ps,B,M,B)
    S = pardisogetschur(ps) # get schur complement matrix

    set_phase!(ps, RELEASE_ALL)
    pardiso(ps, B, M, B)
    set_phase!(ps, original_phase) # reset phase to user setting
    for i ∈ eachindex(original_iparms)
        set_iparm!(ps,i,original_iparms[i])
    end

    return S
end

"""
    pardisogetschur(ps) -> S

retrieve schur complement from PardisoSolver `ps`.
"""
function pardisogetschur(ps::AbstractPardisoSolver)
    nnzschur = get_iparm(ps, 39)
    nschur = get_iparm(ps,38)
    T = isreal(get_matrixtype(ps)) ? Float64 : ComplexF64
    if nnzschur==0
        return spzeros(T,nschur,nschur)
    else
        S = Vector{T}(undef,nnzschur)
        IS = Vector{Int32}(undef,nschur)
        JS = Vector{Int32}(undef,nnzschur)
        ccall_pardiso_get_schur(ps,S,IS,JS)
        IS = pushfirst!(IS,Int32(1)) # some issue with IS (nschur+1 doesn't seem to work)
        S = permutedims(SparseMatrixCSC(nschur,nschur,IS,JS,S)) # really constructing CSR and then transposing
        return S
    end
end

function dim_check(X, A, B)
    size(X) == size(B) || throw(DimensionMismatch(string("solution has $(size(X)), ",
                                                         "RHS has size as $(size(B)).")))
    size(A, 1) == size(B, 1) || throw(DimensionMismatch(string("matrix has $(size(A,1)) ",
                                                               "rows, RHS has $(size(B,1)) rows.")))
    size(B, 1) == stride(B, 2) || throw(DimensionMismatch(
                                            string("Only memory-contiguous RHS supported")))
end

end # module


using Plots
using OceanParameterizations
using WindMixing

reconstruct_fluxes = false
println("Reconstruct fluxes? $(reconstruct_fluxes)")

enforce_surface_fluxes = true
println("Enforce surface fluxes? $(enforce_surface_fluxes)")

output_directory=pwd()*"/ProfileVisuals/reconstruct_$(reconstruct_fluxes)/enforce_surface_fluxes_$(enforce_surface_fluxes)"
mkpath(output_directory)

files =  ["free_convection", "strong_wind", "strong_wind_no_coriolis", "weak_wind_strong_cooling",
          "strong_wind_weak_cooling", "strong_wind_weak_heating"]

file_labels = Dict(
    "free_convection" => "Free convection",
    "strong_wind" => "Strong wind",
    "strong_wind_no_coriolis" => "Strong wind, no rotation",
    "weak_wind_strong_cooling" => "Weak wind, strong cooling",
    "strong_wind_weak_cooling" => "Strong wind, weak cooling",
    "strong_wind_weak_heating" => "Strong wind, weak heating"
)

Ts = Dict()

for file in files
    Ts[file] = data(file, reconstruct_fluxes=reconstruct_fluxes,
                    enforce_surface_fluxes=enforce_surface_fluxes) # <: OceananigansData
end

x_lims = Dict(
    "uw" => (-10,0),
    "vw" => (-4,4.5),
    "wT" => (-1.5,0.7),
    "T" => (19.6,20)
)

f = Dict(
    "u" => 𝒟 -> 𝒟.u.coarse,
    "v" => 𝒟 -> 𝒟.v.coarse,
    "uw" => 𝒟 -> 𝒟.uw.coarse,
    "vw" => 𝒟 -> 𝒟.vw.coarse,
    "wT" => 𝒟 -> 𝒟.wT.coarse,
    "T"  => 𝒟 -> 𝒟.T.coarse
)

zs = Dict(
    "u" => 𝒟 -> 𝒟.u.z,
    "v" => 𝒟 -> 𝒟.v.z,
    "uw" => 𝒟 -> 𝒟.uw.z,
    "vw" => 𝒟 -> 𝒟.vw.z,
    "wT" => 𝒟 -> 𝒟.wT.z,
    "T"  => 𝒟 -> 𝒟.T.z
)

# legend_placement = Dict(
#     "uw" => :bottomleft,
#     "vw" => :bottomright,
#     "wT" => :right
# )

legend_placement = Dict(
    "u" => false,
    "v" => false,
    "T" => false,
    "uw" => false,
    "vw" => false,
    "wT" => false,
    "T"  => false,
)

scaling_factor = Dict(
    "uw" => 1e4,
    "vw" => 1e4,
    "wT" => 1e4,
    "u" => 1,
    "v" => 1,
    "T" => 1
)

x_labels = Dict(
    "u" => "U (m/s)",
    "v" => "V (m/s)",
    "uw" => "U'W' x 10⁴ (m²/s²)",
    "vw" => "V'W' x 10⁴ (m²/s²)",
    "wT" => "W'T' x 10⁴ (C⋅m/s)",
    "T" => "T (C)"
)

titles = Dict(
    "uw" => "Zonal momentum flux, U'W'",
    "vw" => "Meridional momentum flux, V'W'",
    "wT" => "Temperature flux, W'T'",
    "u" => "Zonal momentum (m/s)",
    "v" => "Meridional momentum (m/s)",
    "T" => "Temperature, T",
)

function plot_frame_i(name, i)
    p = plot(xlabel=x_labels[name], ylabel="Depth (m)", palette=:Paired_6, legend=legend_placement[name], foreground_color_grid=:white, plot_titlefontsize=20)
    # p = plot(xlabel=x_labels[name], xlims = x_lims[name], ylabel="Depth (m)", palette=:Paired_6, legend=legend_placement[name], foreground_color_grid=:white, plot_titlefontsize=20)
    for (file, T) in Ts
        plot!(f[name](T)[:,i].*scaling_factor[name], zs[name](T), title = titles[name], label="$(file)", linewidth=3)
    end
    plot!(size=(400,500))
    p
end

p1 = plot_frame_i("uw", 288)
savefig(p1, output_directory*"/uw_last_frame.pdf")

p2 = plot_frame_i("vw", 288)
savefig(p2, output_directory*"/vw_last_frame.pdf")

p3 = plot_frame_i("wT", 288)
savefig(p3, output_directory*"/wT_last_frame.pdf")

pu = plot_frame_i("u", 288)
savefig(pu, output_directory*"/u_last_frame.pdf")

pv = plot_frame_i("v", 288)
savefig(pv, output_directory*"/v_last_frame.pdf")

pT = plot_frame_i("T", 288)
savefig(pT, output_directory*"/T_last_frame.pdf")

p4 = plot(grid=false, showaxis=false, palette=:Paired_6, ticks=nothing)
for (file, T) in Ts
    plot!(1, label=file_labels[file], legend=:left, size=(200,600))
end
p4
savefig(p4, output_directory*"/legend_last_frame.pdf")

layout = @layout [a b c d]
p = plot(p1,p2,p3,p4,layout=layout, size=(1600,400), tickfontsize=12)
savefig(p, output_directory*"/all_last_frame_new_suite.pdf")

layout = @layout [a b c d e]
p = plot(p1,p2,p3,pT,p4,layout=layout, size=(1600,400), tickfontsize=12)
savefig(p, output_directory*"/all_last_frame_new_suite_with_T.pdf")

## ANIMATION

# add x_lims for animations
function plot_frame_i(name, i)
    p = plot(xlabel=x_labels[name], xlims = x_lims[name], ylabel="Depth (m)", palette=:Paired_6, legend=legend_placement[name], foreground_color_grid=:white, plot_titlefontsize=20)
    for (file, T) in Ts
        plot!(f[name](T)[:,i].*scaling_factor[name], zs[name](T), title = titles[name], label="$(file)", linewidth=3)
    end
    plot!(size=(400,500))
    p
end

function animate_all(name, Ts)
    anim = @animate for i in 1:288
        plot_frame_i(name, i)
    end
    return anim
end

save_animation(anim, filename) = gif(anim, filename, fps=20)
save_animation(animate_all("uw", Ts), output_directory*"/uw.gif")
save_animation(animate_all("vw", Ts), output_directory*"/vw.gif")
save_animation(animate_all("wT", Ts), output_directory*"/wT.gif")
save_animation(animate_all("T", Ts), output_directory*"/T.gif")


module TestExamplesP4estMesh3D

using Test
using Trixi

include("test_trixi.jl")

EXAMPLES_DIR = joinpath(examples_dir(), "p4est_3d_dgsem")

# Start with a clean environment: remove Trixi output directory if it exists
outdir = "out"
isdir(outdir) && rm(outdir, recursive=true)

@testset "P4estMesh3D" begin
  @trixi_testset "elixir_advection_basic.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_basic.jl"),
      # Expected errors are exactly the same as with TreeMesh!
      l2   = [0.00016263963870641478],
      linf = [0.0014537194925779984])
  end

  @trixi_testset "elixir_advection_unstructured_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_unstructured_curved.jl"),
      l2   = [0.0004750004258546538],
      linf = [0.026527551737137167])
  end

  @trixi_testset "elixir_advection_nonconforming.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_nonconforming.jl"),
      l2   = [0.00253595715323843],
      linf = [0.016486952252155795])
  end

  @trixi_testset "elixir_advection_amr.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_amr.jl"),
      # Expected errors are exactly the same as with TreeMesh!
      l2   = [9.773852895157622e-6],
      linf = [0.0005853874124926162],
      coverage_override = (maxiters=6, initial_refinement_level=1, base_level=1, med_level=2, max_level=3))
  end

  @trixi_testset "elixir_advection_amr_unstructured_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_amr_unstructured_curved.jl"),
      l2   = [1.6236411810065552e-5],
      linf = [0.0010554006923731395],
      tspan = (0.0, 1.0),
      coverage_override = (maxiters=6, initial_refinement_level=0, base_level=0, med_level=1, max_level=2))
  end

  @trixi_testset "elixir_advection_cubed_sphere.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_cubed_sphere.jl"),
      l2   = [0.002006918015656413],
      linf = [0.027655117058380085])
  end

  @trixi_testset "elixir_advection_restart.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_advection_restart.jl"),
      l2   = [0.002590388934758452],
      linf = [0.01840757696885409])
  end

  @trixi_testset "elixir_euler_source_terms_nonconforming_unstructured_curved.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_source_terms_nonconforming_unstructured_curved.jl"),
      l2   = [4.070355207909268e-5, 4.4993257426833716e-5, 5.10588457841744e-5, 5.102840924036687e-5, 0.00019986264001630542],
      linf = [0.0016987332417202072, 0.003622956808262634, 0.002029576258317789, 0.0024206977281964193, 0.008526972236273522],
      tspan = (0.0, 0.01))
  end

  @trixi_testset "elixir_euler_source_terms_nonperiodic.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_source_terms_nonperiodic.jl"),
      l2   = [0.0015106060984283647, 0.0014733349038567685, 0.00147333490385685, 0.001473334903856929, 0.0028149479453087093],
      linf = [0.008070806335238156, 0.009007245083113125, 0.009007245083121784, 0.009007245083102688, 0.01562861968368434],
      tspan = (0.0, 1.0))
  end

  @trixi_testset "elixir_euler_free_stream.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_free_stream.jl"),
      l2   = [5.162664597942288e-15, 1.941857343642486e-14, 2.0232366394187278e-14, 2.3381518645408552e-14, 7.083114561232324e-14],
      linf = [7.269740365245525e-13, 3.289868377720495e-12, 4.440087186807773e-12, 3.8686831516088205e-12, 9.412914891981927e-12],
      tspan = (0.0, 0.03))
  end

  @trixi_testset "elixir_euler_free_stream_extruded.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_free_stream_extruded.jl"),
      l2   = [8.444868392439035e-16, 4.889826056731442e-15, 2.2921260987087585e-15, 4.268460455702414e-15, 1.1356712092620279e-14],
      linf = [7.749356711883593e-14, 2.8792246364872653e-13, 1.1121659149182506e-13, 3.3228975127030935e-13, 9.592326932761353e-13],
      tspan=(0.0, 0.1))
  end

  @trixi_testset "elixir_euler_ec.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_ec.jl"),
      l2   = [0.010380390326164493, 0.006192950051354618, 0.005970674274073704, 0.005965831290564327, 0.02628875593094754],
      linf = [0.3326911600075694, 0.2824952141320467, 0.41401037398065543, 0.45574161423218573, 0.8099577682187109],
      tspan = (0.0, 0.2),
      coverage_override = (polydeg=3,)) # Prevent long compile time in CI
  end

  @trixi_testset "elixir_euler_sedov.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_sedov.jl"),
      l2   = [7.82070951e-02, 4.33260474e-02, 4.33260474e-02, 4.33260474e-02, 3.75260911e-01],
      linf = [7.45329845e-01, 3.21754792e-01, 3.21754792e-01, 3.21754792e-01, 4.76151527e+00],
      tspan = (0.0, 0.3),
      coverage_override = (polydeg=3,)) # Prevent long compile time in CI
  end

  @trixi_testset "elixir_euler_source_terms_nonconforming_earth.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_source_terms_nonconforming_earth.jl"),
      l2 = [6.040180337738628e-6, 5.4254175153621895e-6, 5.677698851333843e-6, 5.8017136892469794e-6, 1.3637854615117974e-5],
      linf = [0.00013996924184311865, 0.00013681539559939893, 0.00013681539539733834, 0.00013681539541021692, 0.00016833038543762058],
      # Decrease tolerance of adaptive time stepping to get similar results across different systems
      abstol=1.0e-11, reltol=1.0e-11,
      coverage_override = (trees_per_cube_face=(1, 1), polydeg=3)) # Prevent long compile time in CI
  end

  @trixi_testset "elixir_euler_circular_wind_nonconforming.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_circular_wind_nonconforming.jl"),
      l2   = [1.573832094977477e-7, 3.863090659429634e-5, 3.867293305754584e-5, 3.686550296950078e-5, 0.05508968493733932],
      linf = [2.2695202613887133e-6, 0.0005314968179916946, 0.0005314969614147458, 0.0005130280733059617, 0.7944959432352334],
      tspan = (0.0, 2e2),
      coverage_override = (trees_per_cube_face=(1, 1), polydeg=3)) # Prevent long compile time in CI
  end

  @trixi_testset "elixir_euler_baroclinic_instability.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_baroclinic_instability.jl"),
      l2   = [6.725065410642336e-7, 0.00021710117340245454, 0.000438679759422352, 0.00020836356588024185, 0.07602006689579247],
      linf = [1.9101671995258585e-5, 0.029803626911022396, 0.04847630924006063, 0.022001371349740104, 4.847761006938526],
      tspan = (0.0, 1e2),
      # Decrease tolerance of adaptive time stepping to get similar results across different systems
      abstol=1.0e-9, reltol=1.0e-9,
      coverage_override = (trees_per_cube_face=(1, 1), polydeg=3)) # Prevent long compile time in CI
  end

  @trixi_testset "elixir_euler_source_terms_nonperiodic_hohqmesh.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_euler_source_terms_nonperiodic_hohqmesh.jl"),
      l2   = [0.0042023406458005464, 0.004122532789279737, 0.0042448149597303616, 0.0036361316700401765, 0.007389845952982495],
      linf = [0.04530610539892499, 0.02765695110527666, 0.05670295599308606, 0.048396544302230504, 0.1154589758186293])
  end

  @trixi_testset "elixir_mhd_alfven_wave_nonconforming.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_mhd_alfven_wave_nonconforming.jl"),
      l2   = [0.00019018725889431733, 0.0006523517707148006, 0.0002401595437705759, 0.0007796920661427565,
              0.0007095787460334334, 0.0006558819731628876, 0.0003565026134076906, 0.0007904654548841712,
              9.437300326448332e-7],
      linf = [0.0012482306861187897, 0.006408776208178299, 0.0016845452099629663, 0.0068711236542984555,
              0.004626581522263695, 0.006614624811393632, 0.0030068344747734566, 0.008277825749754025,
              1.3475027166309006e-5],
      tspan = (0.0, 0.25),
      coverage_override = (trees_per_dimension=(1, 1, 1),))
  end

  @trixi_testset "elixir_mhd_shockcapturing_amr.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_mhd_shockcapturing_amr.jl"),
      l2   = [0.006298541670176575, 0.0064368506652601265, 0.007108729762852636, 0.006530420607206385,
              0.02061185869237284, 0.005562033787605515, 0.007571716276627825, 0.005571862660453231,
              3.909755063709152e-6],
      linf = [0.20904054009050665, 0.18622917151105936, 0.2347957890323218, 0.19432508025509926,
              0.6858860133405615, 0.15172116633332622, 0.22432820727833747, 0.16805989780225183,
              0.000535219040687628],
      tspan = (0.0, 0.04),
      coverage_override = (maxiters=6, initial_refinement_level=1, base_level=1, max_level=2))
  end
end

# Clean up afterwards: delete Trixi output directory
@test_nowarn rm(outdir, recursive=true)

end # module


#!/usr/bin/env julia
using Covid19Inc, PGFPlotsX
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\usepackage{fontspec}")
push!(PGFPlotsX.CUSTOM_PREAMBLE, raw"\setmainfont{Asana-Math}")
p = Covid19Inc.covid_pop(724, :i10o90, "2020-04-13") |> Covid19Inc.covidpl
pgfsave("../../images/724_covid19_2020-04-13.svg", p)


module CovPlots

using ..PFB: targets

using Base.Threads
using Printf, Statistics

using CSV
using DataFrames
using HypothesisTests
using StatsPlots

function __init__()
    theme(:ggplot2)
    gr(size=(600, 400) .* 1.2)
end

const UNITKEYS = [:time, :fuzzer, :run]
const VALUECOLS = [:l_per, :l_abs, :b_per, :b_abs]

PALETTE = palette(:seaborn_colorblind)
COLORMAP = Dict("aflnet" => PALETTE.colors.colors[1],
                "aflnet-no-state" => PALETTE.colors.colors[2],
                "aflnwe" => PALETTE.colors.colors[3],
                "aflpp" => PALETTE.colors.colors[8],
                "nyx" => PALETTE.colors.colors[5],
                "nyx-newspec" => PALETTE.colors.colors[6],
                "nyx-balanced" => PALETTE.colors.colors[7],
                "nyx-aggressive" => PALETTE.colors.colors[4])

struct TargetData
    df::DataFrame
    agg::DataFrame
    runslastcov::DataFrame
    target::String
end

"Parse CSV with runs for different fuzzers on a single target"
function getruns(csvpath::String, cutoff::UInt)::Union{TargetData, Nothing}
    df = DataFrame(CSV.File(csvpath))
    println("$csvpath: retrieved $(nrow(df)) rows")
    nrow(df) > 0 || return

    target = unique(df.subject)[1]

    # transform time to relative to the first per run
    udf = unstack(df, UNITKEYS, :cov_type, :cov, allowduplicates=true)
    disallowmissing!(udf)
    fuzzerruns = groupby(udf, [:fuzzer, :run])
    # time from seconds to minutes since first in fuzzer run
    transform!(fuzzerruns, :time => (t -> (t .- minimum(t)) .÷ 60) => :time)

    # fill missing values at all times (i.e. 0 to cutoff)
    filled = DataFrame()
    for gd in fuzzerruns
        fullrun = DataFrame(time=0:cutoff, fuzzer=gd[1, :fuzzer], run=gd[1, :run])
        # pick highest values for the same fuzzer / run / time
        gd = combine(groupby(gd, UNITKEYS), VALUECOLS .=> maximum .=> VALUECOLS)
        # join with full times
        append!(filled, leftjoin(fullrun, gd, on=UNITKEYS))
    end
    sort!(filled, [:fuzzer, :run, :time])
    # forward-fill values
    for col in VALUECOLS
        filled[!, col] = accumulate((n0,n1) -> ismissing(n1) ? n0 : n1,
                                    filled[:, col],
                                    init=zero(eltype(filled[:, col])))
    end
    disallowmissing!(filled)
    TargetData(filled, aggruns(filled), getlatestcov(filled), target)
end

"Aggregate across runs"
aggruns(data::TargetData)::DataFrame = aggruns(data.df)
aggruns(df::DataFrame)::DataFrame =
    combine(groupby(df, [:fuzzer, :time]),
            VALUECOLS .=> repeat([mean, median], 2) .=> VALUECOLS)

"Get latest coverage (`b_abs`) for each run of each fuzzer"
getlatestcov(data::TargetData) = getlatestcov(data.df)
getlatestcov(df::DataFrame) =
    combine(groupby(df, [:fuzzer, :run]), :b_abs => last => :cov)

"Create coverage plot"
function makeplot(data::TargetData, layout::Symbol=:pfb)::Plots.Plot
    fuzzers = unique(data.agg.fuzzer)
    cs = [COLORMAP[f] for f in fuzzers]
    timeformatter(x) = @sprintf("%2.0fh", x ÷ 60)
    if layout == :branches
        p1 = @df data.agg plot(:time, :b_abs, group=:fuzzer, title=data.target, color_palette=cs,
                               legend=:bottomright, linewidth=2, yguide="branches",
                               xformatter=timeformatter)
        get(f) = @view data.runslastcov[data.runslastcov.fuzzer .== f, :]
        tickrotation = 20
        p2 = @df get(fuzzers[1]) boxplot(:fuzzer, :cov, xrotation=tickrotation, xtickfontsize=6,
                                         label=nothing, color=COLORMAP[fuzzers[1]])
        for f in fuzzers[2:end]
            @df get(f) boxplot!(:fuzzer, :cov, xrotation=tickrotation,
                                label=nothing, color=COLORMAP[f])
        end
        @df data.runslastcov dotplot!(:fuzzer, :cov, marker=(:black, stroke(0)), label=nothing)
        return plot(p1, p2)
    elseif layout == :pfb
        yticks = [0:20:100;]
        percformatter = y -> @sprintf("%2.1f%%", y)
        p1 = @df data.agg plot(:time, :b_abs, group=:fuzzer, title=data.target, color_palette=cs,
                               legend=:bottomright, yguide="branches", xformatter=timeformatter)
        p2 = @df data.agg plot(:time, :b_per, group=:fuzzer, title=data.target, color_palette=cs,
                               legend=:best, yformatter=percformatter, xformatter=timeformatter)
        p3 = @df data.agg plot(:time, :l_abs, group=:fuzzer, title=data.target, color_palette=cs,
                               legend=:bottomright, yguide="lines", xformatter=timeformatter)
        p4 = @df data.agg plot(:time, :l_per, group=:fuzzer, title=data.target, color_palette=cs,
                               legend=:best, yformatter=percformatter, xformatter=timeformatter)
        return plot(p1, p2, p3, p4, layout=(2, 2))
    else
        error("Unexpected layout $layout")
    end
end

"Process multiple targets selectable from a menu, given tha parent directory"
function multiple(dir::String, cutoff::UInt;
        layout=:pfb, seldirs=false, latex=false, nyxsumm=true, signconf=0.05)
    options, choices = targets(dir, seldirs)
    length(choices) > 0 || return

    # iterators are not thread safe
    selected = options[collect(choices)]

    outpath_suffix = layout == :branches ? "br" : "all"

    # collect aggregated runs for selected targets
    dfs = Dict()
    sizehint!(dfs, length(selected))
    @threads for targetdir in selected
        if seldirs
            csvpath = joinpath(targetdir, "results.csv")
            outpath = joinpath(dirname(csvpath), "coverage_$outpath_suffix.pdf")
        else
            csvpath = targetdir
            base, _ = splitext(targetdir)
            outpath = base*"_$outpath_suffix.pdf"
        end
        runs = getruns(csvpath, cutoff)
        if runs != nothing
            dfs[csvpath] = (runs, outpath, significance(runs, signconf))
        end
    end

    # the plotting library is not thread safe
    summ = DataFrame()
    for (csvpath, (runs, outpath, sign)) in dfs
        if layout != nothing
            p = makeplot(runs, layout)
            savefig(p, outpath)
        end
        if nyxsumm
            append!(summ, sign[startswith.(sign.vs, "nyx"), :])
        else
            append!(summ, sign)
        end
    end
    sort!(summ, :target)
    io = IOBuffer()
    show(io, MIME("text/"*(latex ? "latex" : "plain")), summ;
         eltypes=false, show_row_number=false, summary=false)
    println(String(take!(io)))
    return summ
end

"Compute and return pairwise Mann-Whitney-U test on final median branch coverage"
function significance(data::TargetData, conf::Real)::DataFrame
    fuzzers = ["aflnet", "aflnet-no-state", "aflnwe", "aflpp",
               "nyx", "nyx-newspec", "nyx-balanced", "nyx-aggressive"]
    tab = DataFrame("target" => String[],
                    "vs" => String[],
                    "aflnet" => Union{String, Missing}[],
                    "aflnet-no-state" => Union{String, Missing}[],
                    "aflnwe" => Union{String, Missing}[],
                    "aflpp" => Union{String, Missing}[],
                    "nyx" => Union{String, Missing}[],
                    "nyx-newspec" => Union{String, Missing}[],
                    "nyx-balanced" => Union{String, Missing}[],
                    "nyx-aggressive" => Union{String, Missing}[])
    for fuzz_i in fuzzers
        d = Dict{String, Union{String, Missing}}("vs" => fuzz_i, "target" => data.target)
        for f in fuzzers
            d[f] = missing
        end
        for fuzz_j in fuzzers
            (fuzz_i == fuzz_j) && continue
            runs_i = @view data.runslastcov[data.runslastcov.fuzzer .== fuzz_i, :cov]
            runs_j = @view data.runslastcov[data.runslastcov.fuzzer .== fuzz_j, :cov]
            (length(runs_i) == 0 || length(runs_j) == 0) && continue
            test = MannWhitneyUTest(runs_i, runs_j)
            ρ = pvalue(test)
            d[fuzz_j] = @sprintf("%8.1f %c", test.median, ρ < conf ? "✓" : "x")
        end
        push!(tab, d)
    end
    return tab
end

function loadall(dir::String; fixaflnwe=false)
    options, choices = targets(dir, false)
    length(choices) > 0 || return
    selected = options[collect(choices)]
    df = DataFrame()
    for targetfile in selected
        runs = getruns(targetfile, UInt(1440))
        if runs != nothing
            insertcols!(runs.df, 1, :target => runs.target)
            append!(df, runs.df)
        end
    end
    select!(df, :target, :fuzzer, :run, :time, :b_abs => :cov)
    !fixaflnwe && return df
    df_cp = DataFrame()
    old_val = 0
    for r in eachrow(df)
        if r.fuzzer == "aflnwe" && r.target == "forked-daapd"
            if r.cov < old_val
                r.cov = old_val
            else
                old_val = r.cov
            end
        end
        push!(df_cp, r)
    end
    return df_cp
end

function mktable(df::DataFrame; noseeds=false)
    if noseeds
        baseline_fuzz = "aflnet-no-state"
        other_fuzzs = ["nyx-aggressive"]
        table = DataFrame("target"=>[],
                          baseline_fuzz=>[],
                          # "nyx-aggressive Δ"=>[],
                          "nyx-aggressive %"=>[])
    else
        baseline_fuzz = "aflnet"
        other_fuzzs = ["aflnet-no-state", "aflnwe", "aflpp", "nyx", "nyx-balanced", "nyx-aggressive"]
        table = DataFrame("target"=>[],
                          baseline_fuzz=>[],
                          # "aflnet-no-state Δ"=>[],
                          "aflnet-no-state %"=>[],
                          # "aflnwe Δ"=>[],
                          "aflnwe %"=>[],
                          # "aflpp Δ"=>[],
                          "aflpp %"=>[],
                          # "nyx Δ"=>[],
                          "nyx %"=>[],
                          # "nyx-balanced Δ"=>[],
                          "nyx-balanced %"=>[],
                          # "nyx-aggressive Δ"=>[],
                          "nyx-aggressive %"=>[])
    end
    for gd in groupby(df, :target)
        baseline = @view gd[gd.fuzzer .== baseline_fuzz, :]
        baseline_last = combine(groupby(baseline, :run), :cov => last => :last_cov)
        row = Dict()
        row["target"] = gd.target[1]
        baseline_med = median(baseline_last.last_cov)
        @show row[baseline_fuzz] = "\$$baseline_med\$"
        for fuzz in other_fuzzs
            fuzz == baseline_fuzz && continue
            if !any(gd.fuzzer .== fuzz)
                # row["$fuzz Δ"] = "n/a"
                row["$fuzz %"] = "n/a"
                continue
            end
            fuzz_data = @view gd[gd.fuzzer .== fuzz, :]
            fuzz_last = combine(groupby(fuzz_data, :run), :cov => last => :last_cov)
            test = MannWhitneyUTest(fuzz_last.last_cov, baseline_last.last_cov)
            perc = (median(fuzz_last.last_cov) - baseline_med) / (baseline_med / 100)
            if pvalue(test) < 0.05
                # row["$fuzz Δ"] = @sprintf("\$\\mathbf{%+.1f}\$", test.median)
                row["$fuzz %"] = @sprintf("\$\\mathbf{%+.1f\\%%}\$", perc)
            else
                # row["$fuzz Δ"] = @sprintf("\$%+.1f\$", test.median)
                row["$fuzz %"] = @sprintf("\$%+.1f\\%%\$", perc)
            end
        end
        push!(table, row, cols=:subset)
    end
    sort!(table, :target)
    for row in eachrow(table)
        println(join(Tuple(row), " & "), "\\\\")
    end
    return table
end

function timetocov(df::DataFrame)
    baseline_fuzz = "aflnet"
    # other_fuzzs = ["aflnet-no-state", "aflnwe", "aflpp", "nyx", "nyx-balanced", "nyx-aggressive"]
    other_fuzzs = ["nyx", "nyx-balanced", "nyx-aggressive"]
    table = DataFrame("target"=>[], "nyx"=>[], "nyx-balanced"=>[], "nyx-aggressive"=>[])
    for gd in groupby(df, :target)
        baseline = @view gd[gd.fuzzer .== baseline_fuzz, :]
        baseline_meds = combine(groupby(baseline, :time), :cov => median => :med_cov)
        baseline_med = last(baseline_meds.med_cov)
        baseline_med_time = first(baseline_meds[baseline_meds.med_cov .>= baseline_med, :time])
        row = Dict{String, Any}("target" => gd.target[1])
        for fuzz in other_fuzzs
            # for each run, find time to baseline median last cov
            fuzz_data = @view gd[gd.fuzzer .== fuzz, :]
            fuzz_meds = combine(groupby(fuzz_data, :time), :cov => median => :med_cov)
            ttcov = fuzz_meds[fuzz_meds.med_cov .>= baseline_med, :time]
            if length(ttcov) > 0
                ttcov = ttcov[1]
                # perc = (ttcov - baseline_med_time) / (baseline_med_time / 100)
                # perc = (baseline_med_time - ttcov) / (baseline_med_time / 100)
                # row[fuzz] = @sprintf("%.2f%%", perc)
                row[fuzz] = "$ttcov $baseline_med_time"
            else
                row[fuzz] = "-"
            end
        end
        push!(table, row, cols=:union)
    end
    sort!(table, "target")
    return table
end

getfuzzname(f) =
    if f == "nyx"
        "nyxnet-none"
    elseif f == "nyx-aggressive"
        "nyxnet-aggressive"
    elseif f == "nyx-balanced"
        "nyxnet-balanced"
    else
        f
    end

LINE_TYPES = Dict(
                  "aflnet" => :dot,
                  "nyx" => :dash,
                  "nyx-balanced" => :dashdot,
                  "nyx-aggressive" => :solid,
                  "aflnet-no-state" => :dash,
                  "aflnwe" => :dot,
                  "aflpp" => :dashdot,
                 )

function mkplotpaper(df::DataFrame, extra_target::String, first_target::String="kamailio";
        alltargets=false, allfuzz=false, noseeds=false, fixaflnwe=false)
    @assert extra_target != first_target
    selected_fuzzers = allfuzz ?
        ["aflnet", "aflnet-no-state", "aflnwe", "aflpp", "nyx", "nyx-balanced", "nyx-aggressive"] :
        ["aflnet", "nyx", "nyx-balanced", "nyx-aggressive"]
    ps = []
    timeformatter(x) = @sprintf("%2.0fh", x ÷ 60)
    sort!(df, :target)
    df[!, :time] = df.time .+ 1
    agg = combine(groupby(df, [:target, :fuzzer, :time]), :cov => median => :med_cov)
    if fixaflnwe
        agg_cp = DataFrame()
        old_val = 0
        for r in eachrow(agg)
            if r.fuzzer == "aflnwe" && r.target == "forked-daapd"
                if r.med_cov < old_val
                    r.med_cov = old_val
                else
                    old_val = r.med_cov
                end
            end
            push!(agg_cp, r)
        end
        agg = agg_cp
    end
    extra_plot = nothing
    extra_plot_2 = nothing
    first_p = nothing
    pic_size = (400, 400) .* 0.8
    idx_multiplot = 1
    for gd in groupby(agg, :target)
        target = gd.target[1]
        gd = filter(:fuzzer => f -> f in selected_fuzzers, gd; view=false)
        fuzzers = unique(gd.fuzzer)
        cs = [COLORMAP[f] for f in fuzzers]
        label_fuzzers = permutedims([getfuzzname(f) for f in fuzzers])
        line_types = permutedims([LINE_TYPES[f] for f in fuzzers])
        if target == extra_target
            mkextra(leg) = @df gd plot(:time, :med_cov, group=:fuzzer, title=target, color_palette=cs,
                                    line=line_types,
                                    legend=leg ? :bottomright : false,
                                    label=label_fuzzers,
                                    yguide="Branches",
                                    xformatter=timeformatter,
                                    xticks=0:240:1440,
                                    linewidth=alltargets ? 2 : 1,
                                    size=pic_size,
                                   )
            extra_plot = mkextra(true)
            extra_plot_2 = mkextra(false)
        else
            yguide = target == first_target ? "Branches" : ""
            legend = (!alltargets && idx_multiplot == 12) ? :bottomright : false
            p = @df gd plot(:time, :med_cov, group=:fuzzer, title=target, color_palette=cs,
                            line=line_types,
                            legend=legend,
                            label=label_fuzzers,
                            yguide=yguide,
                            xticks=0:240:1440,
                            xformatter=timeformatter,
                            linewidth=2,
                           )
            if target == first_target
                first_p = p
            else
                push!(ps, p)
            end
            idx_multiplot += 1
        end
    end
    pic_size_scaled = pic_size .* 0.8
    if noseeds
        grid_l = [2, 6]
        ps = [first_p ps... extra_plot]
    else
        grid_l = alltargets ? [2, 7] : [2, 6]
        ps = alltargets ? [extra_plot first_p ps... extra_plot_2] : [first_p ps...]
    end
    ps = plot(
              ps...,
              layout=Tuple(grid_l),
              size=pic_size_scaled.*reverse(grid_l),
             )
    return ps, extra_plot
end

end # module


using Test
using MetaArrays
using StaticArrays
using AxisArrays
using Unitful

struct TestMerge
  val::Int
end
MetaArrays.metamerge(x::TestMerge,y::TestMerge) = TestMerge(x.val + y.val)

struct TestIndex <: AbstractArray{Float64,1}
  x::Vector{Float64}
end
Base.similar(val::TestIndex,::Type{S},dims::NTuple{<:Any,Int}) where S =
  TestIndex{typeof(val.x),eltype(val),ndims(val)}(similar(val.x))
Base.size(val::TestIndex) = size(val.x)
Base.getindex(val::TestIndex,i::Int...) = getindex(val.x,i...)
# not realy a useful function, just a test:
Base.getindex(val::TestIndex,i::Float64) = val[floor(Int,i)] + val[ceil(Int,i)]

struct TestMergeFail
  val::Int
end

struct TestArray{A,T,N} <: AbstractArray{T,N}
  val::A
end
Base.similar(val::TestArray,::Type{S},dims::NTuple{<:Any,Int}) where S =
  TestArray{typeof(val.val),eltype(val),ndims(val)}(similar(val.val))
Base.size(val::TestArray) = size(val.val)
Base.getindex(val::TestArray,i::Int...) = getindex(val.val,i...)
Base.setindex!(val::TestArray,v,i::Int...) = setindex!(val.val,v,i...)

testunion(x::MetaUnion{AbstractRange}) = :range
testunion(x) = :notrange

# TODO: add some more robust tests of broadcast machinery
# (greater variety of arguments, with more call variants
# and then check for test coverage

@testset "MetaArrays" begin

  @testset "MetaArray handles standard array operations" begin
    data = collect(1:10)
    x = meta(data,val=1)
    y = meta(collect(1:10),val=1)

    @test (x.^2) == (1:10).^2
    @test x[1] == data[1]
    @test size(x) == size(data)
    @test similar(x) isa MetaArray
    @test x[1:5] == data[1:5]
    @test (y[1:5] .= 1; sum(y[1:5]) == 5)
    @test (y[1] = 2; y[1] == 2)
    @test x .+ (1:10) == data .+ (1:10)
    @test (.-x) isa MetaArray
    @test collect(1:10) .+ x .+ collect(11:20) == (13:3:40)
    @test broadcast(+,collect(1:10),x,collect(11:20)) == (13:3:40)
    @test (1:10) .+ meta(1:10,val=1) .+ (11:20) isa MetaArray
    @test view(x,1:2) isa MetaArray

    x = TestIndex(collect(1:5))
    @test x[1] == 1
    @test x[1.5] == 3
  end

  @testset "MetaArray preserves metadata over array operations" begin
    data = collect(1:10)
    x = meta(data,val=1)

    @test x.val == 1
    @test x[1:5].val == x.val
    @test x[:].val == x.val
    @test (x .+ (1:10)).val == x.val
    @test (x .+= (1:10); x.val == 1)
    @test (broadcast(+,x,1:10).val == x.val)
    @test similar(x).val == x.val
    @test (.-x).val == 1

    x = meta(collect(1:10),val=1)
    y = meta(collect(1:10),val=1)
    @test vcat(x,y) isa MetaArray
    @test hcat(x,y) isa MetaArray
    @test x*y' isa MetaArray

    x = collect(1:10)
    y = meta(collect(1:10),val=1)
    @test hcat(x,y) isa Array
    @test hcat(y,x) isa MetaArray
    @test vcat(x,y) isa Array
    @test vcat(y,x) isa MetaArray
    @test x*y' isa Array
    @test y*x' isa MetaArray
  end

  @testset "MetaArray properly handles strides" begin
    @test_throws MethodError strides(meta(1:10,val=1))
    x = meta(collect(1:10),val=1)
    @test strides(x) == (1,)
    @test strides(x)[1] == stride(x,1)
  end

  @testset "MetaUnion dispatches properly" begin
    @test testunion(1:5) == :range
    @test testunion(meta(1:10,val=1)) == :range
    @test testunion(meta(collect(1:10),val=1)) == :notrange
  end

  @testset "MetaArray properly merges metadata" begin
    x = meta(collect(1:10),val=1)
    y = meta(collect(2:11),string="string")
    @test (x.+y).val == 1
    @test (x.+y).string == "string"
    @test (x .+ (1:10)).val == 1
    @test ((1:10) .+ x).val == 1

    x = MetaArray(Dict(:test => "value"),collect(1:10))
    y = MetaArray(Dict(:test => "value"),collect(2:11))
    @test getmeta((x .+ y))[:test] == "value"

    x = meta(collect(1:10),val=1)
    y = meta(collect(1:10),val=2)
    @test_throws ErrorException x.+y

    x = meta(collect(1:10),val=TestMerge(1))
    y = meta(collect(1:10),val=TestMerge(2))
    @test (x.+y).val == TestMerge(3)
    @test broadcast(+,x,y).val == TestMerge(3)

    x = meta(collect(1:10),val=(joe=2,bob=3))
    y = meta(collect(1:10),val=(bill=4,))
    @test (x.+y).val == (joe=2,bob=3,bill=4)
  end

  f(x) = (x[0u"s" .. 0.5u"s"] .= 0)

  @testset "MetaArray is AxisArray friendly" begin
    x = meta(AxisArray(rand(10,10),Axis{:time}(range(0u"s",1u"s",length=10)),
      Axis{:freq}(1:10)),val=1)
    x[0u"s" .. 0.5u"s"] = fill(0.0,size(x[0u"s" .. 0.5u"s"]))
    @test all(x[1:5] .== 0)
    x[0u"s" .. 0.5u"s"] .= 1
    @test all(x[1:5] .== 1)

    @test axisdim(x,Axis{:time}) == 1
    @test AxisArray(x) isa AxisArray
    @test AxisArrays.axes(x)[1] isa Axis{:time}
    @test AxisArrays.axes(x,1) isa Axis{:time}
    @test AxisArrays.axes(x,Axis{:time}) isa Axis{:time}
    @test axisnames(x) == (:time,:freq)
    @test axisvalues(x)[2] == 1:10
    @test size(similar(x)) == size(x)
  end

  @testset "MetaArray preserves broadcast specialization" begin
    # Range objects define a specialized lazy broadcasting style,
    # we use them to test the presevation of this lazy style
    x = meta(1:10,val=1)
    @test (x .+ 4) isa MetaArray{<:AbstractRange}
    @test (x .+ 4) == [xi+4 for xi in x]
    @test (.-x) isa MetaArray{<:AbstractRange}
    @test (1:10) .+ meta(1:10,val=1) .+ (11:20) isa MetaArray{<:AbstractRange}
    @test .-x == [-xi for xi in x]

    # SVector defines a specialized style for broadcasting, we use it to test
    # the handling of specialized broadcasting styles
    x = meta(SVector(1,2,3),val=1)
    y = meta([1,2,3],val=1)
    @test x.+y == [2,4,6]
    @test (x.+y).val == 1

    x = meta(collect(1:10),val=1)
    y = SVector((1:10)...)
    @test (x .+ y).val == 1
    @test (y .+ x).val == 1

    y = meta(y,val=1)
    x = collect(1:10)
    @test (x .+ y).val == 1
    @test (y .+ x).val == 1
  end

  @testset "MetaArray allows custom metadata type" begin
    x = MetaArray(TestMerge(2),1:10)
    y = MetaArray(TestMerge(3),1:10)
    k = MetaArray(TestMergeFail(1),1:10)
    h = MetaArray(TestMergeFail(1),1:10)
    m = MetaArray(TestMergeFail(2),1:10)

    z = x.+y
    @test x == 1:10
    @test (x.+y) == ((1:10) .+ (1:10))
    @test (x.+y).val == 5
    @test (h.+k).val == 1
    @test_throws ErrorException h.+m
  end

  @testset "Can extract metadata and underlying array" begin
    x = meta(1:10,val=1)

    @test convert(Array,x) isa Array
    @test getcontents(x) isa AbstractRange
    @test getmeta(x).val == 1
    @test getmeta(x) isa NamedTuple
  end

  @testset "Appropriate conversion" begin
    x = meta(1:10,val=1)

    @test convert(AbstractArray{Int},x) === x
    @test convert(AbstractArray{Int,1},x) === x
    @test_throws MethodError convert(AbstractArray{String},x)
    @test_throws MethodError convert(AbstractArray{Int,2},x)
    @test Array(x) == Array(1:10)

    xplus = MetaArray((test=2,),x)
    @test getcontents(xplus) == getcontents(x)
    @test xplus.val == 1
    @test xplus.test == 2

    xplus = meta(x,test=2)
    @test getcontents(xplus) == getcontents(x)
    @test xplus.val == 1
    @test xplus.test == 2
  end

  @testset "Preoper MetaArray display" begin
    expected = "MetaArray of 1:10"
    x = meta(1:10,val=1)
    iobuf = IOBuffer()
    display(TextDisplay(iobuf), x)
    @test String(take!(iobuf)) == expected
  end

  @testset "Conversion to AbstractArray passes through" begin
    x = meta(collect(1:10),val=1)
    @test convert(AbstractArray{Int},x) === x
    @test convert(AbstractArray{Int,1},x) === x
  end

  @testset "Can compute identity of underlying array" begin
    x = meta(TestArray{Array,Int,1}(collect(1:5)),val=1)
    @test similar(x) isa MetaArray{<:TestArray}
    @test zero(x) isa MetaArray{<:TestArray}
    @test one(x) isa MetaArray{<:TestArray}
    @test all(zero(x) .== 0)
    @test all(one(x) .== 1)
  end
end


"""
    Support for indicator constraints with a set constraint as the right hand side
"""
function JuMP._build_indicator_constraint(
    _error::Function,
    variable::JuMP.AbstractVariableRef,
    constraint::JuMP.VectorConstraint,
    ::Type{MOI.IndicatorSet{A}},
) where {A}
    S = typeof(constraint.set)
    F = typeof(JuMP.moi_function(constraint))
    set = CS.IndicatorSet{A,F,S}(constraint.set, 1 + length(constraint.func))
    if constraint.func isa Vector{VariableRef}
        vov = JuMP.VariableRef[variable]
    else
        vov = JuMP.AffExpr[variable]
    end
    append!(vov, constraint.func)
    return JuMP.VectorConstraint(vov, set)
end


import IndexedTables: pkeynames, excludecols, pkeys, reindex
import Dagger: dsort_chunks

export reindex, rechunk

using StatsBase

function reindex(t::DDataset, by=pkeynames(t), select=excludecols(t, by); kwargs...)
    @noinline function _rechunk(c)
        reindex(c, by, select; kwargs...)
    end
    fromchunks(delayedmap(_rechunk, t.chunks))
end

"""
`rechunk(t::Union{DNDSparse, DNDSparse}[, by[, select]]; <options>)`

Reindex and sort a distributed dataset by keys selected by `by`.

Optionally `select` specifies which non-indexed fields are kept. By default this is all fields not mentioned in `by` for Table and the value columns for NDSparse.

# Options:
- `chunks` -- how to distribute the data. This can be:
    1. An integer -- number of chunks to create
    2. An vector of `k` integers -- number of elements in each of the `k` chunks. `sum(k)` must be same as `length(t)`
    3. The distribution of another array. i.e. `vec.subdomains` where `vec` is a distributed array.
- `merge::Function` -- a function which merges two sub-table or sub-ndsparse into one NDSparse. They may have overlaps in their indices.
- `splitters::AbstractVector` -- specify keys to split by. To create `n` chunks you would need to pass `n-1` splitters and also the `chunks=n` option.
- `chunks_sorted::Bool` -- are the chunks sorted locally? If true, this skips sorting or re-indexing them.
- `affinities::Vector{<:Integer}` -- which processes (Int pid) should each output chunk be created on. If unspecified all workers are used.
- `closed::Bool` -- if true, the same key will not be present in multiple chunks (although sorted). `true` by default.
- `nsamples::Integer` -- number of keys to randomly sample from each chunk to estimate splitters in the sorting process. (See [samplesort](https://en.wikipedia.org/wiki/Samplesort)). Defaults to 2000.
- `batchsize::Integer` -- how many chunks at a time from the input should be loaded into memory at any given time. This will essentially sort in batches of `batchsize` chunks.
"""
function rechunk(dt::DDataset,
                 by=pkeynames(dt),
                 select=dt isa DNextTable ? excludecols(dt, by) : valuenames(dt);
                 merge=_merge,
                 splitters=nothing,
                 chunks_presorted=false,
                 affinities=workers(),
                 chunks=nworkers(),
                 closed=true,
                 sortchunks=true,
                 nsamples=2000,
                 batchsize=nworkers())

    if sortchunks
        perm = sortperm(dt.domains, by=first)
        cs = dt.chunks[perm]
    else
        cs = dt.chunks
    end

    function sortandsample(data, nsamples, presorted)
        r = sample(1:length(data), min(length(data), nsamples),
                   replace=false, ordered=true)

        sorted = !presorted ? reindex(data, by, select) : data
        chunk = !presorted ? tochunk(sorted) : nothing

        (chunk, pkeys(sorted)[r])
    end

    dsort_chunks(cs, chunks, nsamples,
                 batchsize=batchsize,
                 sortandsample=sortandsample,
                 affinities=affinities,
                 splitters=splitters,
                 chunks_presorted=chunks_presorted,
                 merge=merge,
                 by=pkeys,
                 sub=subtable) |> fromchunks

end

### Permutedims

function Base.permutedims(t::DNDSparse{K,V}, p::AbstractVector) where {K,V}
    if !(length(p) == ndims(t) && isperm(p))
        throw(ArgumentError("argument to permutedims must be a valid permutation"))
    end

    permuteintv(intv,d) = Interval(first(intv)[d],
                                   last(intv)[d])
    idxs = map(t.domains) do dmn
        IndexSpace(permuteintv(dmn.interval, p),
                   permuteintv(dmn.boundingrect, p),
                   dmn.nrows,
                  )
    end

    chunks = map(delayed(c -> permutedims(c, p)), t.chunks)
    t1 = DNDSparse{eltype(idxs[1]), V}(idxs, chunks)

    cache_thunks(rechunk(t1))
end


module HM
export HMT, TVar, mk_tcstate, UN, VerboseUN
export Var, Nom, Bound, Var, Arrow, App, Tup, Forall, Implicit
export gTrans, gTransCtx, gCheck
export tvar_of_int, int_of_tvar, IllFormedType
export ⪯, less_than_under_evidences

using MLFS: @typedIO
import MLFS
using MLStyle
import Base

const VerboseUN = Ref(false)

mutable struct UN # unique name
    name :: Symbol
    ts :: UInt64
end
const pool = Dict{Tuple{UInt64, UInt64}, UN}()

function UN(n::Symbol)
    un = UN(n, time_ns())
    pool[(un.ts, objectid(un))] = un
    un
end

MLFS.TypedIO.toVec(::Type{UN}, un::UN) = [
    MLFS.TypedIO.toVec(Symbol, un.name), string(objectid(un)), string(un.ts)]
function MLFS.TypedIO.fromVec(::Type{UN}, v::Vector)
    n = Symbol(v[1])
    objectid = parse(UInt64, v[2])
    ts = parse(UInt64, v[3])
    key = (ts, objectid)

    Some(get!(pool, key) do
        un = UN(n, ts)
    end)

end

const noImplicits = Val(false)

Base.show(io::IO, unique_name::UN) =
    if VerboseUN[]
        Base.print(io, unique_name.name, "@", objectid(unique_name))
    else
        Base.print(io, unique_name.name)
    end

@data HMT begin
    Var(id::UInt)
    Nom(n::Symbol)
    Bound(n::UN)
    App(f::HMT, arg::HMT)
    Arrow(from::HMT, to::HMT)
    Tup(Vector{HMT})
    Forall(Vector{UN}, HMT)

    Implicit(HMT)
end

⊤ = let a = UN(:a)
    Forall([a], Bound(a))
end

@typedIO Nom(Symbol)

@typedIO HMT = begin
    Nom(Symbol)
    Bound(UN)
    App(HMT, HMT)
    Arrow(HMT, HMT)
    Tup(Vector{HMT})
    Forall(Vector{UN}, HMT)

    Implicit(HMT)
end

need_parens(hmt::HMT) =
    @match hmt begin
        ::App || ::Arrow || ::Forall => true
        _ => false
    end

with_parens(f::Function, io::IO, need_parens::Bool) =
    if need_parens
        print(io, '(')
        f()
        print(io, ')')
    else
        f()
    end

Base.show(io::IO, hmt::HMT) =
    @switch hmt begin
    @case Var(v)
        print(io, '\'')
        print(io, v)
        return
    @case Nom(n)
        print(io, n)
        return

    @case Bound(n)
        print(io, n)
        return

    @case App(f, a)
        Base.show(io, f)
        print(io, " ")
        with_parens(io, need_parens(a)) do
            Base.show(io, a)
        end
        return
    @case Arrow(a, r)
        with_parens(io, need_parens(a)) do
            Base.show(io, a)
        end
        print(io, "->")
        Base.show(io, r)
        return
    @case Tup(xs)
        print(io, "{")
        print(io, join(repr.(xs), ", "))
        print(io, "}")
        return
    @case Forall(ns, t)
        print(io, "forall ")
        print(io, join(string.(ns), " "), ".")
        Base.show(io, t)
        return
    @case Implicit(t)
        print(io, "implicit[$t]")
    end

function gTrans(self::Function, root::HMT)
    @match root begin
        Var(_) || Nom(_) || Bound(_) => root
        App(f, arg) => App(self(f), self(arg))
        Arrow(arg, ret) => Arrow(self(arg), self(ret))
        Tup(xs) => Tup(HMT[self(x) for x in xs])
        Forall(ns, p) => Forall(ns, self(p))
        Implicit(x) => Implicit(self(x))
    end
end

function gTransCtx(self′::Function, ctx::Ctx, root::HMT) where Ctx
    @inline self(root::HMT) = self′(ctx, root)
    @match root begin
        Var(_) || Nom(_) || Bound(_) => root
        App(f, arg) => App(self(f), self(arg))
        Arrow(arg, ret) => Arrow(self(arg), self(ret))
        Tup(xs) => Tup(HMT[self(x) for x in xs])
        Forall(ns, p) => Forall(ns, self(p))
        Implicit(x) => Implicit(self(x))
    end
end

function gCheck(self::Function, root::HMT)
    @match root begin
        Var(_) || Nom(_) || Bound(_) => true
        App(f, arg) => self(f) && self(arg)
        Arrow(arg, ret) => self(arg) && self(ret)
        Tup(xs) => all(self(x) for x in xs)
        Forall(_, p) => self(p)
        Implicit(x) => self(x)
    end
end

function fresh(substitutions::Dict{UN, V}, root::HMT) where V <: HMT
    substitutions = Dict{Bound, V}(Bound(k) => v for (k, v) in substitutions)
    fresh(substitutions, root)
end

function fresh(substitutions::Dict{Bound, V}, root::HMT) where V <: HMT
    function freshRec(root::HMT)
        @match root, substitutions begin
        (Bound(_), Dict(root => var)) => var
        # (Forall(ns, p), _) => let restore = Pair{Bound, HMT}[]
        #         for n in ns
        #             k = Bound(n)
        #             v = pop!(substitutions, k, nothing)
        #             if v !== nothing
        #                 push!(restore, k => v)
        #             end
        #         end
        #         r = isempty(substitutions) ? p : gTrans(freshRec, p)
        #         for (k, v) in restore
        #             substitutions[k] = v
        #         end
        #         Forall(ns, r)
        #     end
        (t, _) => gTrans(freshRec, t)
        end
    end
    freshRec(root)
end

function subst(substitutions::Dict{A, B}, root::HMT) where {A <: HMT, B <: HMT}
    @match substitutions begin
        Dict(root => replaced) => replaced
        _ => gTransCtx(subst, substitutions, root)
    end
end
struct IllFormedType <: Exception
    msg::String
end

function tvar_of_int(i::Integer)
    Var(i)
end

function int_of_tvar(x::Var)
    x.id
end

function un_of_bound(x::Bound)
    x.n
end

function bound_of_un(x::UN)
    Bound(x)
end

function mk_tcstate(tctx::Vector{HMT})
    bound_links = Dict{UN, Set{UInt}}()

    function unlink(un::Union{UN, Bound})
        unlink(un) do _::UInt
            nothing
        end
    end

    function unlink(f::Function, b::Bound)
        unlink(f, un_of_bound(b))
    end

    function unlink(f::Function, un::UN)
        rels = pop!(bound_links, un, nothing)
        rels === nothing && return
        for rel in rels
            tctx[rel] = Var(rel)
            f(rel)
        end
        empty!(rels)
        nothing
    end

    function new_tvar()::HMT
        vid = UInt(length(tctx) + 1)
        tvar = tvar_of_int(vid)
        push!(tctx, tvar)
        tvar
    end

    function new_tvar_id()::UInt
        vid = UInt(length(tctx) + 1)
        tvar = tvar_of_int(vid)
        push!(tctx, tvar)
        vid
    end

    function occur_in(i::UInt, ty::HMT)
        @switch ty begin
            @case Var(i′) && if i′ === i end
                return false
            @case _
                visit_func(x) =
                    @match x begin
                        Var(i′) && if i′ === i end => false
                        _ => gCheck(visit_func, x)
                    end
                return !visit_func(ty)
        end
    end

    function link!(typeref::UInt, b::Bound)
        un = un_of_bound(b)
        links =
            get!(bound_links, un) do
                Set{UInt}(UInt[typeref])
            end
        typeref in links || begin
            push!(links, typeref)
            tctx[typeref] = b
        end
        b
    end

    function link!(typeref::UInt, t::HMT)
        tctx[typeref] = t
        t
    end

    function prune(a::HMT)
        @match a begin
            Var(i) =>
                @match tctx[i] begin
                    Var(i′) && if i′ === i end => a
                    t => link!(i, prune(t))
                end
            a => gTrans(prune, a)
        end
    end

    function pruneWithVarCheck(f::Function, a::HMT)
        function prune(a::HMT)
            v = @match a begin
                Var(i) =>
                    @match tctx[i] begin
                        Var(i′) && if i′ === i end => a

                        t => link!(i, prune(t))
                    end
                a => gTrans(prune, a)
            end
            if v isa Var
                f(v)
            end
            v
        end
        prune(a)
    end

    function instantiate(x::HMT)
        @match x begin
            Forall(ns, ty) =>
                let freshmap = Dict{Bound, Var}(Bound(a) => new_tvar() for a in ns)
                    subst(freshmap, ty)
                end
            _ => x
        end
    end

    function unifyImplicits!(lhs::HMT, rhs::HMT, implicits::Vector{HMT})
        lhs = prune(lhs)
        rhs = prune(rhs)
        lhs === rhs && return true

        @match lhs, rhs begin
            (Forall(_, lhs), _) => unifyImplicits!(lhs, rhs, implicits)
            (Var(_), _) ||
            (_, Var(_)) =>
                unifyINST(lhs, rhs)

            (_, Forall()) =>
                begin
                    unifyImplicits!(lhs, instantiate(rhs), implicits)
                end

            (_, Arrow(Implicit(im), rhs)) =>
                begin # for debugger
                    push!(implicits, im)
                    unifyImplicits!(lhs, rhs, implicits)
                end

            _ => unifyINST(lhs, rhs)
        end
    end

    function unifyINST(lhs::HMT, rhs::HMT)
        lhs = prune(lhs)
        rhs = prune(rhs)
        lhs === rhs && return true
        # @info :INST lhs rhs

        @match lhs, rhs begin
            (Forall(_, lhs), _) => unifyINST(lhs, rhs)

            (Var(i && ai), b) ||
            (b, Var(i && ai)) =>
                if occur_in(ai, b)
                    throw(IllFormedType("a = a -> b"))
                else
                    link!(i, b)
                    true
                end

            (_, Forall()) =>
                begin # for debugger
                    unifyINST(lhs, instantiate(rhs))
                end

            (_, Implicit(rhs)) => unifyINST(lhs, rhs)
            (Implicit(lhs), _) => unifyINST(lhs, rhs)
            (Arrow(a1, r1), Arrow(a2, r2)) =>
                begin # for debugger
                    unifyINST(a2, a1) && unifyINST(r1, r2)
                end

            (App(f1, a1), App(f2, a2)) =>
                begin # for debugger
                    unifyINST(f1, f2) && unifyINST(a1, a2)
                end

            (Tup(xs1), Tup(xs2)) =>
                all(zip(xs1, xs2)) do (lhs, rhs)
                    unifyINST(lhs, rhs)
                end
            _ => false
        end
    end

    function unify(lhs::HMT, rhs::HMT)
        lhs = prune(lhs)
        rhs = prune(rhs)
        lhs === rhs && return true

        @match lhs, rhs begin
            (Forall(ns1, p1), Forall(ns2, p2)) =>
                (begin
                    N1, N2 = length(ns1), length(ns2)
                    N1 !== N2 && return false

                    subst1 = Dict{UN, Var}(a => new_tvar() for a in ns1)
                    generic = p2
                    unify(fresh(subst1, p1), generic) || return false

                    generic = prune(generic)
                    remap2 = Bound[Bound(n) for n in ns2]
                    remap1 = Dict{Bound, Bound}()

                    for (k, v) in subst1
                        v = prune(v)
                        v in remap2 || return false
                        haskey(remap1, v) && return false
                        remap1[v] = Bound(k)
                    end

                    backmap1 = Dict{Bound, Bound}()
                    backmap2 = Dict{Bound, Bound}()

                    for un in ns2
                        unlink(un) do i::UInt
                            backmap1[Var(i)] = remap1[Bound(un)]
                            backmap2[Var(i)] = Bound(un)
                        end
                    end

                    unify(subst(backmap2, p2), p2) &&
                    unify(subst(backmap1, p1), p1)
                end)

            (Var(i && ai), b) ||
            (b, Var(i && ai)) =>

                if occur_in(ai, b)
                    throw(IllFormedType("a = a -> b"))
                else
                    link!(i, b)
                    true
                end

            (Arrow(a1, r1), Arrow(a2, r2)) =>
                unify(a1, a2) && unify(r1, r2)

            (App(f1, a1), App(f2, a2)) =>
                unify(f1, f2) && unify(a1, a2)

            (Tup(xs1), Tup(xs2)) =>
                all(zip(xs1, xs2)) do (lhs, rhs)
                    unify(lhs, rhs)
                end
            _ => false
        end
    end

    (unifyINST = unifyINST,
     unifyImplicits! = unifyImplicits!,
     unify = unify,
     tctx = tctx,
     new_tvar = new_tvar,
     new_tvar_id = new_tvar_id,
     unlink = unlink,
     occur_in = occur_in,
     instantiate = instantiate,
     prune = prune,
     pruneWithVarCheck = pruneWithVarCheck)

end

(lhs::HMT) ⪯ (rhs::HMT)  = begin
    small_tc = mk_tcstate(HMT[])
    subst_table = Dict{UInt, Var}()
    function subst(root::HMT)
        @match root begin
            Var(i) =>
                get!(subst_table, i) do
                    small_tc.new_tvar()
                end

            _ => gTrans(subst, root)
        end
    end
    small_tc.unifyImplicits(subst(lhs), subst(rhs), HMT[])
end

function less_than_under_evidences(lhs::HMT, rhs ::HMT)
    evidences = HMT[]
    small_tc = mk_tcstate(HMT[])
    subst_table = Dict{UInt, Var}()
    function subst(root::HMT)
        @match root begin
            Var(i) =>
                get!(subst_table, i) do
                    small_tc.new_tvar()
                end

            _ => gTrans(subst, root)
        end
    end
    small_tc.unifyImplicits!(subst(lhs), subst(rhs), evidences), evidences
end

end


# Used in GECCO paper
using Statistics, Printf
using DataFrames, CSV, Dates
# This file does not include the "export" declarations needed to be included in CGP.jl.
# Thus, it should be directly included.  Sample runs in ../data/9_3_21/.

# Calculates robustness and evolvability for all phenotypes by sampling circuits as in Hu et al. (2020).
# The node-adjacency matrix for the network of phenotypes is computed as goal_edge_matrix.
# Diagonal entries are the number of self edges which is the robustness count.
# goal_edge_matrix[g,h] is the number of mutations discovered from goal g to goal h.
# The methodology is to do nprocesses*nwalks random walks each of length steps, and record all of the transitions made.
# If csvfile is specified, writes the dataframe to this file.
function run_random_walks_parallel( nprocesses::Int64, nwalks::Int64, gl::Vector{MyInt}, p::Parameters, steps::Int64;
      csvfile::String="", output_dict::Bool=true, save_complex::Bool=false)
  ngoals = 2^(2^p.numinputs)
  addvalues(x::Tuple{Int64,Float64},y::Tuple{Int64,Float64}) = (x[1]+y[1],(x[2]+y[2])/2.0)
  if output_dict
    goal_pair_dict = save_complex ? Dict{Tuple{MyInt,MyInt},Tuple{Int64,Float64}}() : Dict{Tuple{MyInt,MyInt},Int64}()
    #goal_pair_dict = Dict{Tuple{MyInt,MyInt},Int64}()
    dict_list = pmap( x->run_random_walks( nwalks, p, steps, output_dict=output_dict, save_complex=save_complex ),
        collect(1:nprocesses) )
    #println("length(dict_list): ",length(dict_list))
    #println("dict_list[1]: ",dict_list[1])
    #println("typeof(goal_pair_dict): ",typeof(goal_pair_dict))
    #for gp in dict_list
      goal_pair_dict = save_complex ? merge(addvalues,goal_pair_dict,dict_list...) : merge(+,goal_pair_dict,dict_list...)
      #goal_pair_dict = merge(+, goal_pair_dict, gp )
   # end
    #return goal_pair_dict
    df = robust_evolvability( goal_pair_dict, gl, p )
  else
    goal_edge_matrix = zeros(Int64,ngoals,ngoals)
    goal_edge_matrix_list = pmap( x->run_random_walks( nwalks, p, steps, output_dict=output_dict ), collect(1:nprocesses) )
    println("len: ",length(goal_edge_matrix_list),"  size: ",size(goal_edge_matrix_list[1]),"  type: ",typeof(goal_edge_matrix_list[1]))
    for gem in goal_edge_matrix_list
      goal_edge_matrix .+= gem
    end
    df = robust_evolvability( goal_edge_matrix, gl )
  end
  if length(csvfile) > 0
    println("csvfile: ",csvfile)
    open( csvfile, "w" ) do f
      hostname = chomp(open("/etc/hostname") do f read(f,String) end)
      println(f,"# date and time: ",Dates.now())
      println(f,"# host: ",hostname," with ",nprocs()-1,"  processes: " )
      println(f,"# funcs: ", Main.CGP.default_funcs(p.numinputs))
      print_parameters(f,p,comment=true)
      println(f,"# nwalks: ",nwalks)
      println(f,"# steps: ",steps)
      println(f,"# nprocesses: ",nprocesses)
      CSV.write( f, df, append=true, writeheader=true )
      println("csvfile written")
    end
  end
  df
end

function run_random_walks( nwalks::Int64, p::Parameters, steps::Int64; output_dict::Bool=true, save_complex::Bool=false )
  @assert p.numoutputs == 1
  funcs = default_funcs(p.numinputs)
  addvalues(x::Tuple{Int64,Float64},y::Tuple{Int64,Float64}) = (x[1]+y[1],(x[2]+y[2])/2.0)
  ngoals = 2^(2^p.numinputs)
  c_list = [ random_chromosome(p,funcs) for _ = 1:nwalks ]
  if output_dict
    goal_pair_dict = save_complex ? Dict{Tuple{MyInt,MyInt},Tuple{Int64,Float64}}() : Dict{Tuple{MyInt,MyInt},Int64}()
    rw_list = [random_walk( c_list[i], steps, output_dict=output_dict, save_complex=save_complex ) for i = 1:nwalks ]
    goal_pair_dict = save_complex ? merge(addvalues,goal_pair_dict,rw_list...) :  merge(+,goal_pair_dict,rw_list...)
    return goal_pair_dict
  else
    goal_edge_matrix = zeros(Int64,ngoals,ngoals)
    for i = 1:nwalks
      goal_edge_matrix .+= random_walk( c_list[i], steps, output_dict=output_dict )
    end
    return goal_edge_matrix
  end
end

# Outputs a node-edge adjacency matrix unless output_dict==true,
#      in which case a dictionary indexed on goal pairs whose values the the count of the goal pair
function random_walk( c::Chromosome, steps::Int64; output_dict::Bool=true, save_complex::Bool=false  )
  funcs = default_funcs(c.params.numinputs)
  ngoals = 2^(2^c.params.numinputs)
  #addvalues(x,y) = (x[1]+y[1],(x[2]+y[2])/2.0)
  addvalues(x::Tuple{Int64,Float64},y::Tuple{Int64,Float64}) = (x[1]+y[1],(x[2]+y[2])/2.0)
  if output_dict && !save_complex
    goal_pair_dict = Dict{Tuple{MyInt,MyInt},Int64}()
  elseif output_dict && save_complex
    goal_pair_dict = Dict{Tuple{MyInt,MyInt},Tuple{Int64,Float64}}()
  else
    goal_edge_matrix = zeros(Int64,ngoals,ngoals)
  end
  cmplx = 0.0
  goal = output_values(c)[1]
  #println("start goal: ",goal)
  for i = 1:steps
    cmplx = save_complex ? complexity5(c) : 0  # complexity of circuit mapping to prev_goal
    prev_goal = goal
    mutate_chromosome!( c, funcs )
    #print_circuit(c)
    goal = output_values(c)[1]
    #println("i: ",i,"  prev_goal: ",prev_goal,"  goal: ",goal)
    if output_dict
      value = save_complex ? (1,cmplx) : 1
      if haskey(goal_pair_dict,(prev_goal,goal))
        prev_value = goal_pair_dict[(prev_goal,goal)]
        goal_pair_dict[(prev_goal,goal)] = save_complex ?  addvalues(prev_value,value) : prev_value+value
      else
        goal_pair_dict[(prev_goal,goal)] = value
      end
    else
      goal_edge_matrix[Int64(prev_goal)+1,Int64(goal)+1] += 1
    end
  end
  if output_dict
    #println("length(goal_pair_dict): ",length(goal_pair_dict))
    goal_pair_dict
  else
    goal_edge_matrix
  end
end

function robust_evolvability( goal_pair_list::Array{Pair{Tuple{UInt16,UInt16},Int64},1}, p::Parameters)
  robust_evolvability( pairs_to_dict( goal_pair_list ), p )
end

#
# Note that gl is a list of MyInts, not a list of goals
function robust_evolvability( goal_pair_dict::Dict{Tuple{MyInt,MyInt},Tuple{Int64,Float64}}, gl::Vector{MyInt}, p::Parameters )
  ngoals = 2^(2^p.numinputs)
  allgoals = MyInt(0):MyInt(ngoals-1)
  println("allgoals: ",allgoals,"  ngoals: ",ngoals)
  re_list = pmap( g->robust_evo(g, goal_pair_dict, ngoals ), gl )
  #re_list = map( g->robust_evo(g, goal_pair_dict, ngoals ), gl )
  df = DataFrame()
  df.goal = gl
  df.frequency= [ re[1] for re in re_list ]
  df.robustness = [ re[2] for re in re_list ]
  df.s_evolvability = [ re[3] for re in re_list ]
  df.d_evolvability = [ re[4] for re in re_list ]
  df.complexity= [ re[5] for re in re_list ]
  df
end

# Helper function used in pmap() in robust_evolvability()
# Corrected 1/2/21 and 1/3/21
function robust_evo( g::MyInt, goal_pair_dict::Dict{Tuple{MyInt,MyInt},Tuple{Int64,Float64}}, ngoals::Int64 )
  allgoals = MyInt(0):MyInt(ngoals-1)
  dget(x) = get(goal_pair_dict,x,(0,0.0))
  dget10(x) = get(goal_pair_dict,x,(0,0.0))[1]>0 ? 1 : 0
  #println("dget((g,g)): ",dget((g,g)),"  dget((g,g))[2]: ",dget((g,g))[2])
  sum_gh = sum( dget((g,h))[1] for h in allgoals )
  #sum_hg = sum( dget((h,g))[1] for h in allgoals )
  #for h in allgoals println("  dget((g,h)): ",dget((g,h))) end
  sum_gh_c = sum( dget((g,h))[1]*dget((g,h))[2] for h in allgoals )
  gg = dget((g,g))[1]
  frequency = sum_gh
  robustness = gg/frequency
  s_evolvability = sum_gh - gg  # count of mutations that take g to a different phenotype
  d_evolvability = sum(dget10((g,h)) for h in allgoals) - ((gg>0) ? 1 : 0)  #count of number of phenotypes by mutation of g
  complexity = sum_gh_c/frequency
  (frequency,robustness,s_evolvability,d_evolvability,complexity)
end
#
# Note that gl is a list of MyInts, not a list of goals
function robust_evolvability( goal_pair_dict::Dict{Tuple{MyInt,MyInt},Int64}, gl::Vector{MyInt}, p::Parameters )
  ngoals = 2^(2^p.numinputs)
  allgoals = MyInt(0):MyInt(ngoals-1)
  println("allgoals: ",allgoals)
  re_list = pmap( g->robust_evo(g, goal_pair_dict, ngoals ), gl )
  #re_list = map( g->robust_evo(g, goal_pair_dict, ngoals ), gl )
  df = DataFrame()
  df.goal = gl
  df.frequency= [ re[1] for re in re_list ]
  df.robustness = [ re[2] for re in re_list ]
  df.s_evolvability = [ re[3] for re in re_list ]
  df.d_evolvability = [ re[4] for re in re_list ]
  df
end

# Helper function used in pmap() in robust_evolvability()
# Corrected 1/2/21 and 1/3/21
function robust_evo( g::MyInt, goal_pair_dict::Dict{Tuple{MyInt,MyInt},Int64}, ngoals::Int64 )
  allgoals = MyInt(0):MyInt(ngoals-1)
  dget(x) = get(goal_pair_dict,x,0)
  dget10(x) = get(goal_pair_dict,x,0)>0 ? 1 : 0
  sum_gh = sum( dget((g,h)) for h in allgoals )
  #sum_hg = sum( dget((h,g)) for h in allgoals )
  gg = dget((g,g))
  frequency = sum_gh
  robustness = gg/frequency
  s_evolvability = sum_gh - gg  # count of mutations that take g to a different phenotype
  d_evolvability = sum(dget10((g,h)) for h in allgoals) - ((gg>0) ? 1 : 0)  #count of number of phenotypes by mutation of g
  (frequency,robustness,s_evolvability,d_evolvability)
end

# Calculates robustness and degree evolvability for each goal and saves these in a dataframe.
# Robustness
function robust_evolvability( goal_edge_matrix::Array{Int64,2}, gl::Vector{MyInt} )
  ngoals = size(goal_edge_matrix)[1]
  #triangularize!(goal_edge_matrix)
  frequency = zeros(Int64,ngoals)
  robustness = zeros(Float64,ngoals)
  s_evolvability = zeros(Int64,ngoals)
  d_evolvability = zeros(Int64,ngoals)
  for g in map(x->x+1,gl)   # convert to 1-based indexing
    frequency[g] = sum( goal_edge_matrix[g,h]  for h = 1:ngoals )
    robustness[g] = goal_edge_matrix[g,g]/frequency[g]
    s_evolvability[g] = frequency[g] - goal_edge_matrix[g,g]
    d_evolvability[g] = sum( ((goal_edge_matrix[g,h] != 0) ? 1 : 0) for h = 1:ngoals ) - ((goal_edge_matrix[g,g]>0) ? 1 : 0)
  end
  df = DataFrame()
  df.goal = collect(MyInt(0):MyInt(ngoals-1))
  df.frequency = frequency
  df.robustness = robustness
  df.s_evolvability = s_evolvability
  df.d_evolvability = d_evolvability
  df
end

# Convert into a upper triangular matrix by adding the below-diagonal entries to the corresponding above-diagonal entry
# Never used.
function triangularize!( M::Array{Int64,2} )
  ngoals = size(M)[1]
  @assert ngoals == size(M)[2]
  for g = 1:(ngoals-1)
    for h = (g+1):ngoals
      M[g,h] += M[h,g]
      M[h,g] = 0
    end
  end
end

function matrix_to_dict( M::Array{Int64,2} )
  goal_pair_dict = Dict{Tuple{MyInt,MyInt},Int64}()
  for g = 0:(size(M)[1]-1)
    for h = 0:(size(M)[2]-1)
      if M[g+1,h+1] != 0
        goal_pair_dict[(g,h)] = M[g+1,h+1]
      end
    end
  end
  goal_pair_dict
end

function dict_to_matrix( goal_pair_dict::Dict{Tuple{MyInt,MyInt},Int64}, p::Parameters )
  M = zeros(Int64,2^2^p.numinputs,2^2^p.numinputs)
  for (g,h) in keys(goal_pair_dict)
    M[g+1,h+1] = goal_pair_dict[(g,h)]
  end
  M
end

function pairs_to_dict( pairs::Array{Pair{Tuple{MyInt,MyInt},Int64},1})
  println("pairs to dict")
  pair_dict = Dict{Tuple{MyInt,MyInt},Int64}()
  for p in pairs
    pair_dict[p[1]] = p[2]
  end
  pair_dict
end

# Run robust_evolvability and then write to csvfile if it is given
# Ran for over 24 hours on a large goal_pair_list.
function robust_evolvability_to_df( goal_pair_list::Array{Pair{Tuple{UInt16,UInt16},Int64},1}, p::Parameters;
      csvfile::String="" )
  df = robust_evolvability( pairs_to_dict( goal_pair_list ), p )
  if length(csvfile) > 0
    open( csvfile, "w" ) do f
      hostname = chomp(open("/etc/hostname") do f read(f,String) end)
      println(f,"# date and time: ",Dates.now())
      println(f,"# host: ",hostname," with ",nprocs()-1,"  processes: " )
      println(f,"# funcs: ", Main.CGP.default_funcs(p.numinputs))
      print_parameters(f,p,comment=true)
      #println(f,"# nwalks: ",nwalks)
      #println(f,"# steps: ",steps)
      #println(f,"# nprocesses: ",nprocesses)
      CSV.write( f, df, append=true, writeheader=true )
    end
  end
  df
end

# Computes frequency, robustess, strength_evolvability, degree_evolvabilty of a list of MyInts (interpreted as goals)
# Note that gl is a list of MyInts, not a list of goals
function dict_to_evolvability( goal_pair_dict::Dict{Tuple{MyInt,MyInt},Int64}, gl::Vector{MyInt}, p::Parameters )
  @assert p.numoutputs == 1
  dget(x) = get(goal_pair_dict,x,0)
  dget10(x) = get(goal_pair_dict,x,0)>0 ? 1 : 0
  frequency = zeros(Int64,length(gl))
  robustness = zeros(Float64,length(gl))
  s_evolvability = zeros(Int64,length(gl))
  d_evolvability = zeros(Int64,length(gl))
  allgoals = MyInt(0):MyInt(2^2^p.numinputs-1)
  i = 1
  for g in gl
    #frequency[i] = sum( get( goal_pair_dict, (g,h), 0 ) + get( goal_pair_dict, (h,g), 0 ) for h in allgoals )
    frequency[i] = sum( dget((g,h)) for h in allgoals )
    #robustness[i] = get( goal_pair_dict, (g,g), 0 )
    robustness[i] = dget((g,g))/frequency[i]
    s_evolvability[i] = sum( dget((g,h)) for h in allgoals ) - dget((g,g))
    d_evolvability[i] = sum( dget10((g,h)) for h in allgoals ) - (dget((g,g))>0 ? 1 : 0)
    i += 1
  end
  df = DataFrame()
  df.goal = gl
  df.frequency = frequency
  df.robustness = robustness
  df.s_evolvability = s_evolvability
  df.d_evolvability = d_evolvability
  df
end


"""
Kinetic.jl: A Portable Framework for Computational Fluid Dynamics and Scientific Machine Learning

Copyright (c) 2021 Tianbai Xiao (tianbaixiao@gmail.com)
"""
module Kinetic

export 転

"""
Portable Framework for Computational Fluid Dynamics and Scientific Machine Learning

轻量化的计算流体力学建模和计算框架

"転" means "rolling" in Japannese
"""
const 転 = Kinetic

using Reexport
@reexport using KitBase
@reexport using KitML

end


"""
Failure classifier adversary.

Inputs low dimensional representation, classifies failures based on this.
"""
module Adversary

export accuracy, train, trainandsave, load, getdata

using Base.Iterators: repeated
using BSON: @save, @load
using Colors
using CUDA
using Images
using MLDatasets
using NNlib
using Parameters
using Flux
using Flux.Data: DataLoader
using Flux: @epochs, onehotbatch, onecold, crossentropy, throttle
using Statistics

include("Δsoftmax.jl")

CUDA.allowscalar(false)

@with_kw mutable struct Args
    α::Float64 = 3e-5        # learning rate (default: 3e-4)
    batchsize::Int = 50      # batch size
    epochs::Int = 20         # number of epochs
    device::Function = gpu   # set as gpu if available
    throttle::Int = 1        # throttle print every X seconds
    lowdim::Int = 64         # size of low-dimensional representation
    trainfrac::Float64 = 0.5 # [train | test] fraction
end

global USE_LOW_DIM_REPRESENTATION = true
global USE_BATCHES = false
global USE_ΔSOFTMAX = false

function splitdata(X, Y, sut, encoder, args=Args())
    k_train = floor(Int, size(X,2) * args.trainfrac)
    k_test = size(X,2) - k_train

    X_train = X[:, 1:k_train]
    X_test = X[:, k_train+1:end]

    Y_train_classes = Y[:, 1:k_train]
    Y_test_classes = Y[:, k_train+1:end]

    # evaluate the training data on the SUT to get true target labels
    Y_train = onecold(sut(X_train)) .!= onecold(Y_train_classes)
    # Y_test = fill(NaN, k_test) # no true testing data, we'll have to actually call the SUT after candidate failure selection
    Y_test = onecold(sut(X_test)) .!= onecold(Y_test_classes) # FOR TESTING/DEBUGGING/ANALYSIS PURPOSES

    # Low-dimensional representation of X
    if USE_LOW_DIM_REPRESENTATION
        X̃_train = encoder(X_train)
        X̃_test = encoder(X_test)
    else
        X̃_train = X_train
        X̃_test = X_test
    end

    Δtrain = Δsoftmax(sut, X_train, Y_train_classes)
    Δtest = Δsoftmax(sut, X_test, Y_test_classes)

    X̃_train, Y_train, Δtrain = duplicate_failures(X̃_train, Y_train, Δtrain)

    # split into batches
    if USE_BATCHES
        encoded_traindata = DataLoader(X̃_train, Y_train, batchsize=args.batchsize) #, shuffle=true)
        encoded_testdata = DataLoader(X̃_test, Y_test, batchsize=args.batchsize)
        true_traindata = DataLoader(X_train, Y_train_classes, batchsize=args.batchsize)
        true_testdata = DataLoader(X_test, Y_test_classes, batchsize=args.batchsize)
    else
        if USE_ΔSOFTMAX
            encoded_traindata = DataLoader(X̃_train, Y_train, Δtrain) #, shuffle=true)
            encoded_testdata = DataLoader(X̃_test, Y_test, Δtest)
        else
            encoded_traindata = DataLoader(X̃_train, Y_train) #, shuffle=true)
            encoded_testdata = DataLoader(X̃_test, Y_test)
        end
        true_traindata = DataLoader(X_train, Y_train_classes)
        true_testdata = DataLoader(X_test, Y_test_classes)
    end

    return encoded_traindata, encoded_testdata, true_traindata, true_testdata
end

function convert_data(testdata, sut, encoder)
    x̃ = reduce(hcat, [encoder(x) for (x,_) in testdata])
    ỹ = reduce(vcat, [onecold(sut(x)) .!= onecold(y) for (x,y) in testdata])
    return DataLoader(x̃, ỹ)
end

function num_failures(testdata, sut)
    num = 0
    for (x,y) in testdata
        num += sum(onecold(sut(x)) .!= onecold(y))
    end
    return num
end

function duplicate_failures(X̃_train, Y_train, Δtrain, d=10)
    newX = Matrix{eltype(X̃_train)}(undef, size(X̃_train,1), 0)
    newY = Vector{eltype(Y_train)}(undef, 0)
    newΔ = Vector{eltype(Δtrain)}(undef, 0)
    for (i,y) in enumerate(Y_train)
        if y
            for _ in 1:d
                newX = hcat(newX, X̃_train[:,i])
                newY = vcat(newY, Y_train[i])
                newΔ = vcat(newΔ, Δtrain[i])
            end
        end
        newX = hcat(newX, X̃_train[:,i])
        newY = vcat(newY, Y_train[i])
        newΔ = vcat(newΔ, Δtrain[i])
    end
    return newX, newY, newΔ
end

function buildmodel(args)
    if USE_LOW_DIM_REPRESENTATION
        lowdim = args.lowdim
    else
        lowdim = prod([28,28,1])
    end
    return Chain(
            Dense(lowdim, 2lowdim, relu),
            Dense(2lowdim, lowdim, relu),
            Dense(lowdim, 1, sigmoid))
end

# Flux binarycrossentropy loss fails with use of @.
function bce(ŷ, y; agg=mean, ϵ=Flux.epseltype(ŷ))
    agg(.- Flux.Losses.xlogy.(y, ŷ .+ ϵ) .- Flux.Losses.xlogy.(1 .- y, 1 .- ŷ .+ ϵ))
end

function bce_adversary(ŷ, y; Ω=float32(0.95), agg=mean, ϵ=Flux.epseltype(ŷ))
    agg(.- Flux.Losses.xlogy.(y .* Ω, ŷ .+ ϵ) .- Flux.Losses.xlogy.((1 .- y) .* (1-Ω), 1 .- ŷ .+ ϵ))
end

function bce_adversary_multi(ŷ, y; Ω=float32(2), agg=mean, ϵ=Flux.epseltype(ŷ))
    agg(.- Flux.Losses.xlogy.(y .* Ω, ŷ .+ ϵ) .- Flux.Losses.xlogy.((1 .- y), 1 .- ŷ .+ ϵ))
end

adversarial_loss = bce

"""
Cost function:

\$\\mathcal{J}(\\mathbf{\\hat{y}}, \\mathbf{y}) = \\frac{1}{m}\\sum_{i=1}^m \\mathcal{L}(\\hat{y}_i, y_i)\$

with loss:

\$\\mathcal{L}(\\hat{y}, y) = - y \\log(\\hat y) - (1 - y)\\log(1-\\hat y)\$

\$\\mathcal{L}(\\hat{y}, y) = - y \\log(\\hat y) \\Omega - (1 - y)\\log(1-\\hat y)\$
"""
cost(dataloader, model) = mean(adversarial_loss(model(x), y) for (x,y) in dataloader)

function accuracy(dataloader, model)
    acc = 0
    len = 0
    for (x,y) in dataloader
        acc += sum(predict(model, x)  .== cpu(y)) / size(x,2)
        len += size(x,2)
    end
    return acc/len
end

function failure_accuracy(dataloader, model)
    acc = 0
    len = 0
    for (x,y) in dataloader
        ŷ = predict(model, cpu(x))
        for (i, yᵢ) in enumerate(cpu(y))
            if yᵢ # true failure
                acc += ŷ[i] == cpu(yᵢ) # / size(x,2)
                len += 1
            end
        end
    end
    return acc/len
end

function train(m, X, Y, sut, encoder; kwargs...)
    # initialize model parameters
    args = Args(; kwargs...)

    # load data
    traindata, testdata, true_traindata, true_testdata = splitdata(X, Y, sut, encoder, args)

    # construct model
    if isnothing(m)
        @info "Building new adversarial model."
        m = buildmodel(args)
    else
        @info "Using input adversarial model."
    end
    traindata = args.device.(traindata)
    testdata = args.device.(testdata)
    m = args.device(m)
    # L₂(x, λ=float32(5e-2)) = λ/2 * sum(abs2, x) # L₂ regularization
    # loss(x,y,Δ) = adversarial_loss(m(x), y) + sum(Δ) # Δsoftmax augmentation
    loss(x,y) = adversarial_loss(m(x), y) # + sum(L₂, params(m))

    # callback, optimizer, and training
    callback = throttle(()->@show(cost(traindata, m)), args.throttle)
    opt = ADAM(args.α)

    @info "Training adversary..."
    @epochs args.epochs Flux.train!(loss, params(m), traindata, opt, cb=callback)

    @show accuracy(traindata, m)
    @show accuracy(testdata, m)

    @show failure_accuracy(traindata, m)
    @show failure_accuracy(testdata, m)

    return cpu(m), cpu(testdata), cpu(true_testdata)
end

function predict(m, x, thresh=0.5)
    return cpu(m)(cpu(x)) .>= thresh
end

function select_candidates(model, testdata, mapping; debug=false)
    # Run each X_test datapoint over failure classifier model
    X_candidates = []
    Y_candidates = []
    debug && @warn("select_candidates using true target Y values `debug=true`")
    for (i, (x̃,ỹ)) in enumerate(testdata)
        ŷ = predict(model, x̃) # only select the predictions above X confidence
        for (j, ŷⱼ) in enumerate(ŷ)
            if (!debug && ŷⱼ) || (debug && ỹ[j])
                x, y = mapping[i] # vec(x)
                push!(X_candidates, x[:,j])
                push!(Y_candidates, y[:,j])
            end
        end
    end
    return X_candidates, Y_candidates
end

function rand_select_candidates(testdata, mapping, k)
    X_candidates, Y_candidates = [], []
    if k == 0
        return X_candidates, Y_candidates
    else
        testdata = Flux.Data.DataLoader([d[1] for d in testdata], [d[2][1] for d in testdata])
        rand_idx = rand(1:length(testdata), k)

        for (i, (x̃,ỹ)) in enumerate(testdata)
            for (j, ỹⱼ) in enumerate(ỹ)
                for _ in 1:sum(i .== rand_idx) # counts duplicates from rand
                    x, y = mapping[i] # vec(x)
                    push!(X_candidates, x[:,j])
                    push!(Y_candidates, y[:,j])
                end
            end
        end
        return X_candidates, Y_candidates
    end
end

function trainandsave(X, Y, sut, encoder; kwargs...)
    adversary = train(nothing, X, Y, sut, encoder; kwargs...)
    adversary = cpu(adversary)
    @save "models/adversary.bson" adversary
end

"""
Load pre-trained neural network model.

```julia
using Flux
using CUDA
m = load(\"models/adversary.bson\")
```
"""
function load(model_filename_bson)
    @load model_filename_bson m
    return m
end

end # module


#
# Mapping with integer indexes.
#

MappingQuery(idom, oty, col::Column) =
    Query(
        MappingSig(col),
        Input(convert(Domain, idom)),
        convert(Output, oty))

MappingQuery(idom, oty, offs::AbstractVector{Int}, vals::AbstractVector) =
    MappingQuery(idom, oty, Column(offs, vals))

MappingQuery{T}(idom, oty, data::AbstractVector{Vector{T}}) =
    MappingQuery(idom, convert(Output, oty) |> setoptional() |> setplural(), Column(data))

MappingQuery{T}(idom, oty, data::AbstractVector{Nullable{T}}) =
    MappingQuery(idom, convert(Output, oty) |> setoptional, Column(data))

MappingQuery(idom, oty, data::AbstractVector) =
    MappingQuery(idom, oty, Column(data))

MappingQuery{T}(idom, data::AbstractVector{Vector{T}}) =
    MappingQuery(idom, T, data)

MappingQuery{T}(idom, data::AbstractVector{Nullable{T}}) =
    MappingQuery(idom, T, data)

MappingQuery{T}(idom, data::AbstractVector{T}) =
    MappingQuery(idom, T, data)

MappingQuery(data::AbstractVector) =
    MappingQuery(Int, data)

immutable MappingSig <: AbstractPrimitive
    col::Column
end

ev(sig::MappingSig, idxs::AbstractVector) =
    sig.col[idxs]