text:
import DataFrames.skipmissing

##### To define functions for calculating the distances:

function calc_distance(ss1::ExoplanetsSysSim.CatalogSummaryStatistics, ss2::ExoplanetsSysSim.CatalogSummaryStatistics, return_KS_or_AD::String ; AD_mod::Bool=false, all_dist::Bool=false, save_dist::Bool=true)
    #This function calculates the total KS distance between two summary statistics (simulated observed catalogs).
    #If 'all_dist=true', the function outputs the individual distances in the distance function.
    #If 'save_dist=true', the function also saves the distances (individual and total) to a file (assuming file 'f' is open for writing).

    M_cat_obs1 = ones(Int64,0) #array to be filled with the number of transiting planets in each simulated system for ss1
    M_cat_obs2 = ones(Int64,0) #array to be filled with the number of transiting planets in each simulated system for ss2
    for k in 1:get_int(sim_param,"max_tranets_in_sys")
        append!(M_cat_obs1, k*ones(Int64, ss1.stat["num n-tranet systems"][k]))
        append!(M_cat_obs2, k*ones(Int64, ss2.stat["num n-tranet systems"][k]))
    end

    #To handle empty arrays:
    if min(length(ss1.stat["radius_ratio_above_list"]), length(ss1.stat["radius_ratio_below_list"]), length(ss1.stat["radius_ratio_across_list"])) < 2 || min(length(ss2.stat["radius_ratio_above_list"]), length(ss2.stat["radius_ratio_below_list"]), length(ss2.stat["radius_ratio_across_list"])) < 2 #need at least 2 elements in each of these summary statistics per catalog in order to be able to compute AD distances
        println("Not enough observed multi-planet systems in one of the catalogs to compute the AD distance.")
        d = ones(Int64,13)*1e6

        println("Distances: ", d, [sum(d)])
        if save_dist
            println(f, "Dist_KS: ", d, [sum(d)])
            println(f, "Dist_AD: ", d, [sum(d)])
        end

        if all_dist
            return d
        else
            return sum(d)
        end
    end

    max_incl_sys = get_real(sim_param,"max_incl_sys")
    cos_factor = cos(max_incl_sys*pi/180) #factor to divide the number of targets in simulation by to get the actual number of targets needed (with an isotropic distribution of system inclinations) to produce as many transiting systems for a single observer

    #To compute the KS distances:
    d_KS = Array{Float64}(undef, 13)
    d_KS[1] = abs(ss1.stat["num_tranets"]/(ss1.stat["num targets"]/cos_factor) - ss2.stat["num_tranets"]/(ss2.stat["num targets"]))
    d_KS[2] = ksstats_ints(M_cat_obs1, M_cat_obs2)[5]
    d_KS[3] = ksstats(ss1.stat["P list"], ss2.stat["P list"])[5]
    d_KS[4] = ksstats(ss1.stat["period_ratio_list"], ss2.stat["period_ratio_list"])[5]
    d_KS[5] = ksstats(ss1.stat["duration list"], ss2.stat["duration list"])[5]
    d_KS[6] = ksstats(ss1.stat["duration_ratio_list"], ss2.stat["duration_ratio_list"])[5]
    d_KS[7] = ksstats(ss1.stat["depth list"], ss2.stat["depth list"])[5]
    d_KS[8] = ksstats(ss1.stat["radius_ratio_list"], ss2.stat["radius_ratio_list"])[5]
    d_KS[9] = ksstats(ss1.stat["depth above list"], ss2.stat["depth above list"])[5]
    d_KS[10] = ksstats(ss1.stat["depth below list"], ss2.stat["depth below list"])[5]
    d_KS[11] = ksstats(ss1.stat["radius_ratio_above_list"], ss2.stat["radius_ratio_above_list"])[5]
    d_KS[12] = ksstats(ss1.stat["radius_ratio_below_list"], ss2.stat["radius_ratio_below_list"])[5]
    d_KS[13] = ksstats(ss1.stat["radius_ratio_across_list"], ss2.stat["radius_ratio_across_list"])[5]

    #To compute the AD distances:
    if AD_mod
        ADdist = ADstats_mod
    else
        ADdist = ADstats
    end

    d_AD = Array{Float64}(undef, 13)
    d_AD[1] = abs(ss1.stat["num_tranets"]/(ss1.stat["num targets"]/cos_factor) - ss2.stat["num_tranets"]/(ss2.stat["num targets"]))
    d_AD[2] = ksstats_ints(M_cat_obs1, M_cat_obs2)[5]
    d_AD[3] = ADdist(ss1.stat["P list"], ss2.stat["P list"])
    d_AD[4] = ADdist(ss1.stat["period_ratio_list"], ss2.stat["period_ratio_list"])
    d_AD[5] = ADdist(ss1.stat["duration list"], ss2.stat["duration list"])
    d_AD[6] = ADdist(ss1.stat["duration_ratio_list"], ss2.stat["duration_ratio_list"])
    d_AD[7] = ADdist(ss1.stat["depth list"], ss2.stat["depth list"])
    d_AD[8] = ADdist(ss1.stat["radius_ratio_list"], ss2.stat["radius_ratio_list"])
    d_AD[9] = ADdist(ss1.stat["depth above list"], ss2.stat["depth above list"])
    d_AD[10] = ADdist(ss1.stat["depth below list"], ss2.stat["depth below list"])
    d_AD[11] = ADdist(ss1.stat["radius_ratio_above_list"], ss2.stat["radius_ratio_above_list"])
    d_AD[12] = ADdist(ss1.stat["radius_ratio_below_list"], ss2.stat["radius_ratio_below_list"])
    d_AD[13] = ADdist(ss1.stat["radius_ratio_across_list"], ss2.stat["radius_ratio_across_list"])

    #To print and/or write the distances to file:
    println("KS Distances: ", d_KS, [sum(d_KS)])
    println("AD Distances: ", d_AD, [sum(d_AD)])
    if save_dist
        println(f, "Dist_KS: ", d_KS, [sum(d_KS)])
        println(f, "Dist_AD: ", d_AD, [sum(d_AD)])
    end

    #To return the distances or total distance:
    if all_dist
        if return_KS_or_AD == "KS"
            return d_KS
        elseif return_KS_or_AD == "AD"
            return d_AD
        end
    else
        if return_KS_or_AD == "KS"
            return sum(d_KS)
        elseif return_KS_or_AD == "AD"
            return sum(d_AD)
        end
    end
end

function calc_distance_Kepler(ss1::ExoplanetsSysSim.CatalogSummaryStatistics, return_KS_or_AD::String ; AD_mod::Bool=false, all_dist::Bool=false, save_dist::Bool=true)
    #This function calculates the total KS distance between a population generated by our model and the actual Kepler population (which must be loaded in).
    #If 'all_dist=true', the function outputs the individual distances in the distance function.
    #If 'save_dist=true', the function also saves the distances (individual and total) to a file (assuming file 'f' is open for writing).

    M_cat_obs = ones(Int64,0) #array to be filled with the number of transiting planets in each simulated system
    for k in 1:get_int(sim_param,"max_tranets_in_sys")
        append!(M_cat_obs, k*ones(Int64, ss1.stat["num n-tranet systems"][k]))
    end

    if min(length(ss1.stat["radius_ratio_above_list"]), length(ss1.stat["radius_ratio_below_list"]), length(ss1.stat["radius_ratio_across_list"])) < 2 #need at least 2 elements in each of these summary statistics in order to be able to compute AD distances
        println("Not enough observed multi-planet systems in the simulated catalog.")
        d = ones(Int64,13)*1e6

        println("Distances: ", d, [sum(d)])
        if save_dist
            println(f, "Dist_KS: ", d, [sum(d)])
            println(f, "Dist_AD: ", d, [sum(d)])
        end

        if all_dist
            return d
        else
            return sum(d)
        end
    end

    max_incl_sys = get_real(sim_param,"max_incl_sys")
    cos_factor = cos(max_incl_sys*pi/180) #factor to divide the number of targets in simulation by to get the actual number of targets needed (with an isotropic distribution of system inclinations) to produce as many transiting systems for a single observer

    #To compute the KS distances:
    d_KS = Array{Float64}(undef, 13)
    d_KS[1] = abs(ss1.stat["num_tranets"]/(ss1.stat["num targets"]/cos_factor) - length(P_confirmed)/N_Kepler_targets)
    d_KS[2] = ksstats_ints(M_cat_obs, M_confirmed)[5]
    d_KS[3] = ksstats(ss1.stat["P list"], P_confirmed)[5]
    d_KS[4] = ksstats(ss1.stat["period_ratio_list"], R_confirmed)[5]
    d_KS[5] = ksstats(ss1.stat["duration list"].*24, t_D_confirmed)[5] #transit durations in simulations are in days, while in the Kepler catalog are in hours
    d_KS[6] = ksstats(ss1.stat["duration_ratio_list"], xi_confirmed)[5]
    d_KS[7] = ksstats(ss1.stat["depth list"], D_confirmed)[5]
    d_KS[8] = ksstats(ss1.stat["radius_ratio_list"].^2, D_ratio_confirmed)[5] #simulations save radius ratios while we computed transit duration ratios from the Kepler catalog
    d_KS[9] = ksstats(ss1.stat["depth above list"], D_above_confirmed)[5]
    d_KS[10] = ksstats(ss1.stat["depth below list"], D_below_confirmed)[5]
    d_KS[11] = ksstats(ss1.stat["radius_ratio_above_list"].^2, D_ratio_above_confirmed)[5]
    d_KS[12] = ksstats(ss1.stat["radius_ratio_below_list"].^2, D_ratio_below_confirmed)[5]
    d_KS[13] = ksstats(ss1.stat["radius_ratio_across_list"].^2, D_ratio_across_confirmed)[5]

    #To compute the AD distances:
    if AD_mod
        ADdist = ADstats_mod
    else
        ADdist = ADstats
    end

    d_AD = Array{Float64}(undef, 13)
    d_AD[1] = abs(ss1.stat["num_tranets"]/(ss1.stat["num targets"]/cos_factor) - length(P_confirmed)/N_Kepler_targets)
    d_AD[2] = ksstats_ints(M_cat_obs, M_confirmed)[5]
    d_AD[3] = ADdist(ss1.stat["P list"], P_confirmed)
    d_AD[4] = ADdist(ss1.stat["period_ratio_list"], R_confirmed)
    d_AD[5] = ADdist(ss1.stat["duration list"].*24, t_D_confirmed) #transit durations in simulations are in days, while in the Kepler catalog are in hours
    d_AD[6] = ADdist(ss1.stat["duration_ratio_list"], xi_confirmed)
    d_AD[7] = ADdist(ss1.stat["depth list"], D_confirmed)
    d_AD[8] = ADdist(ss1.stat["radius_ratio_list"].^2, D_ratio_confirmed) #simulations save radius ratios while we computed transit duration ratios from the Kepler catalog
    d_AD[9] = ADdist(ss1.stat["depth above list"], D_above_confirmed)
    d_AD[10] = ADdist(ss1.stat["depth below list"], D_below_confirmed)
    d_AD[11] = ADdist(ss1.stat["radius_ratio_above_list"].^2, D_ratio_above_confirmed)
    d_AD[12] = ADdist(ss1.stat["radius_ratio_below_list"].^2, D_ratio_below_confirmed)
    d_AD[13] = ADdist(ss1.stat["radius_ratio_across_list"].^2, D_ratio_across_confirmed)

    #To print and/or write the distances to file:
    println("KS Distances: ", d_KS, [sum(d_KS)])
    println("AD Distances: ", d_AD, [sum(d_AD)])
    if save_dist
        println(f, "Dist_KS: ", d_KS, [sum(d_KS)])
        println(f, "Dist_AD: ", d_AD, [sum(d_AD)])
    end

    #To return the distances or total distance:
    if all_dist
        if return_KS_or_AD == "KS"
            return d_KS
        elseif return_KS_or_AD == "AD"
            return d_AD
        end
    else
        if return_KS_or_AD == "KS"
            return sum(d_KS)
        elseif return_KS_or_AD == "AD"
            return sum(d_AD)
        end
    end
end

function target_function(active_param::Vector{Float64}, use_KS_or_AD::String, Kep_or_Sim::String ; AD_mod::Bool=false, weights::Vector{Float64}=ones(13), all_dist::Bool=false, save_dist::Bool=true)
    #This function takes in the values of the active model parameters, generates a simulated observed catalog, and computes the distance function.
    #If 'all_dist=true', the function outputs the individual distances in the distance function.
    #If 'save_dist=true', the function also saves the distances (unweighted and weighted, individual and total) to a file (assuming file 'f' is open for writing).

    println("Active parameter values: ", active_param)
    if save_dist
        println(f, "Active_params: ", active_param) #to write the params to file
    end

    global sim_param
    sim_param_here = deepcopy(sim_param)
    ExoplanetsSysSim.update_sim_param_from_vector!(active_param,sim_param_here)
    cat_phys = generate_kepler_physical_catalog(sim_param_here)
    cat_phys_cut = ExoplanetsSysSim.generate_obs_targets(cat_phys,sim_param_here)
    cat_obs = observe_kepler_targets_single_obs(cat_phys_cut,sim_param_here)
    summary_stat = calc_summary_stats_model(cat_obs,sim_param_here)

    if Kep_or_Sim == "Kep"
        dist = calc_distance_Kepler(summary_stat, use_KS_or_AD; AD_mod=AD_mod, all_dist=true, save_dist=save_dist)
    elseif Kep_or_Sim == "Sim"
        global summary_stat_ref
        dist = calc_distance(summary_stat, summary_stat_ref, use_KS_or_AD; AD_mod=AD_mod, all_dist=true, save_dist=save_dist)
    end

    weighted_dist = dist .* weights

    if weights != ones(13)
        println("Weighted distances: ", weighted_dist, [sum(weighted_dist)])
        if save_dist
            println(f, "Dist_weighted: ", weighted_dist, [sum(weighted_dist)])
        end
    end

    if all_dist
        return weighted_dist
    else
        return sum(weighted_dist)
    end
end

function compute_weights_target_fitness_std_perfect_model(num_evals::Int64, use_KS_or_AD::String ; AD_mod::Bool=false, weight::Bool=true, dists_exclude::Vector{Int64}=Int64[], save_dist::Bool=true)
    t_elapsed = @elapsed begin
        active_param_true = make_vector_of_sim_param(sim_param)
        println("# True values: ", active_param_true)
        if save_dist
            println(f, "# Format: Dist: [distances][total distance]")
        end

        dists_true = zeros(num_evals,13)
        for i in 1:num_evals
            dists_true[i,:] = target_function(active_param_true, use_KS_or_AD, "Sim"; AD_mod=AD_mod, all_dist=true, save_dist=save_dist)
        end
        mean_dists = transpose(mean(dists_true, dims=1))[:,] #array of mean distances for each individual distance
        mean_dist = mean(sum(dists_true, dims=2)) #mean total distance
        #std_dists = transpose(std(dists_true, 1))[:,] #array of std distances for each individual distance
        rms_dists = transpose(sqrt.(mean(dists_true .^2, dims=1)))[:,] #array of rms (std around 0)  distances for each individual distance
        std_dist = std(sum(dists_true, dims=2)) #std of total distance
        rms_dist = sqrt(mean(sum(dists_true, dims=2) .^2)) #rms (std around 0) of total distance; should be similar to the mean total distance

        if weight
            weights = 1 ./ rms_dists #to use the array 'rms_dists' as the weights for the individual distances
        else
            weights = ones(13)
        end
        weights[dists_exclude] .= 0. #to exclude certain distances from being used in the total distance function (but still computing and saving them during the optimization) by setting their weights to zero

        weighted_dists_true = zeros(num_evals,13)
        for i in 1:num_evals
            weighted_dists_true[i,:] = dists_true[i,:]  .* weights
        end
        mean_weighted_dists = transpose(mean(weighted_dists_true, dims=1))[:,] #array of mean weighted distances for each individual distance
        mean_weighted_dist = mean(sum(weighted_dists_true, dims=2)) #mean weighted total distance
        std_weighted_dist = std(sum(weighted_dists_true, dims=2)) #std of weighted total distance
    end

    println("Mean dists: ", mean_dists)
    println("Rms dists: ", rms_dists)
    println("Weights (1/rms dists): ", weights)
    println("Mean weighted dists: ", mean_weighted_dists)
    println("Distance using true values: ", mean_dist, " +/- ", std_dist)
    println("Weighted distance using true values: ", mean_weighted_dist, " +/- ", std_weighted_dist)
    if save_dist
        println(f, "Mean: ", mean_dists, [mean_dist])
        println(f, "Rms: ", rms_dists, [rms_dist])
        println(f, "Weights (1/rms dists): ", weights)
        println(f, "Mean weighted dists: ", mean_weighted_dists, [mean_weighted_dist])
        println(f, "# Distance using true values (default parameter values): ", mean_dist, " +/- ", std_dist)
        println(f, "# Weighted distance using true values (default parameter values): ", mean_weighted_dist, " +/- ", std_weighted_dist)
        println(f, "# elapsed time: ", t_elapsed, " seconds")
        println(f, "#")
    end

    return (active_param_true, weights, mean_weighted_dist, std_weighted_dist)
end

function map_square_to_triangle(r1::Float64, r2::Float64, A::Vector{Float64}, B::Vector{Float64}, C::Vector{Float64})
    #This function takes in a point (r1,r2) in the unit square (i.e. r1,r2 in [0,1]) and maps it to a point P=(x,y) in the triangle defined by vertices A,B,C
    #If r1,r2 are uniformly drawn in [0,1], then the point P=(x,y) is also uniformly drawn in the triangle; see http://www.cs.princeton.edu/~funk/tog02.pdf (Section 4.2) for a reference

    @assert 0. <= r1 <= 1.
    @assert 0. <= r2 <= 1.
    P = (1. - sqrt(r1)) .* A + (sqrt(r1)*(1. - r2)) .* B + (sqrt(r1)*r2) .* C
    return P
end

function target_function_transformed_params(active_param_transformed::Vector{Float64}, transformed_indices::Vector{Int64}, A::Vector{Float64}, B::Vector{Float64}, C::Vector{Float64}, use_KS_or_AD::String, Kep_or_Sim::String ; AD_mod::Bool=false, weights::Vector{Float64}=ones(13), all_dist::Bool=false, save_dist::Bool=true)

    r1, r2 = active_param_transformed[transformed_indices]
    @assert 0. <= r1 <= 1.
    @assert 0. <= r2 <= 1.
    active_param = deepcopy(active_param_transformed)
    active_param[transformed_indices] = map_square_to_triangle(r1, r2, A, B, C)

    return target_function(active_param, use_KS_or_AD, Kep_or_Sim ; AD_mod=AD_mod, weights=weights, all_dist=all_dist, save_dist=save_dist)
end


text:
const dictpath = "/"*relpath((@__FILE__)*"/../..","/") * "/assets/dicts"

const Connectome2FS = deserialize(dictpath * "/Connectome2FS.jls")
const FS2Connectome = deserialize(dictpath * "/FS2Connectome.jls")

function get_node_attributes(graph)
    n_nodes = length(graph["node"])
    coords = Array{Float64}(undef, n_nodes, 3)
    nID = Vector{Int}(undef, n_nodes)
    region = Vector{String}(undef, n_nodes)
    labels = Vector{String}(undef, n_nodes)
    hemisphere = Vector{String}(undef, n_nodes)
    for i ∈ 1:n_nodes
        for j ∈ child_elements(graph["node"][i])
            if attribute(j, "key") == "d0"
                coords[i, 1] = parse(Float64, LightXML.content(j))
            elseif attribute(j, "key") == "d1"
                coords[i, 2] = parse(Float64, LightXML.content(j))
            elseif attribute(j, "key") == "d2"
                coords[i, 3] = parse(Float64, LightXML.content(j))
            elseif attribute(j, "key") == "d3"
                nID[i] = parse(Int, LightXML.content(j))
            elseif attribute(j, "key") == "d4"
                region[i] = LightXML.content(j)
            elseif attribute(j, "key") == "d5"
                labels[i] = LightXML.content(j)
            elseif attribute(j, "key") == "d7"
                hemisphere[i] = LightXML.content(j)
            end
        end
    end
    x, y, z = coords[:,1], coords[:,2], coords[:,3]
    return DataFrame(ID=nID, Label=labels, Region=region, Hemisphere=hemisphere, x=x, y=y, z=z)
end

function get_adjacency_matrix(graph)
    A = spzeros(83,83)
    local n
    local l
    for edge in graph["edge"]
        i = parse(Int, attribute(edge, "source"))
        j = parse(Int, attribute(edge, "target"))
        for child in child_elements(edge)
            if attribute(child, "key") == "d9"
                n = parse(Float64, LightXML.content(child))
            elseif attribute(child, "key") == "d12"
               l = parse(Float64, LightXML.content(child))
            end
        end
        A[i,j] = n / l^2
    end

    return SimpleWeightedGraph((A + transpose(A))*0.5)
end

function read_cmtk_parcellation(graph_path)
    xdoc = parse_file(graph_path)
    xroot = root(xdoc)
    ces = collect(child_elements(xroot))
    graph = ces[end]

    n_nodes = length(graph["node"])
    nID = Vector{Int}(undef, n_nodes)
    region = Vector{String}(undef, n_nodes)
    labels = Vector{String}(undef, n_nodes)
    hemisphere = Vector{String}(undef, n_nodes)

    for i ∈ 1:n_nodes
        for j ∈ child_elements(graph["node"][i])
            if attribute(j, "key") == "d0"
                region[i] = LightXML.content(j)
            elseif attribute(j, "key") == "d1"
                labels[i] = LightXML.content(j)
            elseif attribute(j, "key") == "d2"
                hemisphere[i] = LightXML.content(j)
            elseif attribute(j, "key") == "d3"
                nID[i] = parse(Int, LightXML.content(j))
            end
        end
    end
    return DataFrame(ID=nID, Label=labels, Region=region, Hemisphere=hemisphere)
end

function load_graphml(graph_path::String)
    xdoc = parse_file(graph_path)
    xroot = root(xdoc)
    ces = collect(child_elements(xroot))

    node_attributes = get_node_attributes(ces[end])
    A = get_adjacency_matrix(ces[end])
    return node_attributes, A
end

"""
    Connectome(path::String; norm=true)

Main type introduced by Connectomes.jl
"""
struct Connectome
    parc::DataFrame
    graph::SimpleWeightedGraph{Int64, Float64}
    A::SparseMatrixCSC{Float64, Int64}
    D::SparseMatrixCSC{Float64, Int64}
    L::SparseMatrixCSC{Float64, Int64}
    function Connectome(graph_path::String; norm=true)
        parc, Graph = load_graphml(graph_path)
        if norm
            Graph = SimpleWeightedGraphs.adjacency_matrix(Graph) |> max_norm |> SimpleWeightedGraph
        end
        A = SimpleWeightedGraphs.adjacency_matrix(Graph)
        D = SimpleWeightedGraphs.degree_matrix(Graph)
        L = SimpleWeightedGraphs.laplacian_matrix(Graph)
        new(parc, Graph, A, D, L)
    end

    function Connectome(parc, A)
        Graph = SimpleWeightedGraph(A)
        new(parc, Graph, SimpleWeightedGraphs.adjacency_matrix(Graph), SimpleWeightedGraphs.degree_matrix(Graph), SimpleWeightedGraphs.laplacian_matrix(Graph))
    end

    function Connectome(parc, coords, A)
        for (i, j) in enumerate([:x, :y, :z])
            parc[!, j] = coords[:,i]
        end
        Connectome(parc, A)
    end
end

# convenience functions for processing graphs
function graph_filter(connectome::Connectome, cutoff::Float64=1e-2)
    A = graph_filter(connectome.A, cutoff)
    Connectome(connectome.parc, A)
end

graph_filter(A, cutoff=1e-2) = A .* (A .> cutoff)

max_norm(M) = M ./ maximum(M)

function degree(C::Connectome)
    diag(C.D) |> Array
end

ext: jl


text:
module PARSE_TESTS
using Reproduce, Test, FileIO, JLD2

TEST_DIR = "TEST_DIR"

function arg_parse(args; use_git_info=false, as_symbols=false)

    s = ArgParseSettings()
    @add_arg_table! s begin
        "--a"
        help = "a file to save the results to"
        arg_type = Int64
        required = true
    end

    @add_arg_table! s begin
        "--b"
        help = "a file to save the results to"
        arg_type = String
        required = true
    end

    parsed = create_info(args, s, TEST_DIR;
                         use_git_info=use_git_info,
                         as_symbols=as_symbols)
    return parsed
end

function reset_tests()
    working_dir = pwd()
    rm(joinpath(working_dir, TEST_DIR); recursive=true)
end

function parse_test()

    test_dict = Dict("a"=>1, "b"=>"1.jld")
    args=["--a", "1", "--b", "1.jld"]
    parsed = arg_parse(args)

    reset_tests()

    return hash(test_dict) == parsed["_HASH"]
end

function track_test()

    used_keys = ["a", "b"]
    parsed_dicts = Dict{UInt64, Dict}()
    for i in 1:10
        args=["--a", "$(i)", "--b", "$(i).jld"]
        parsed = arg_parse(args)
        parsed_dicts[parsed["_HASH"]] = filter(k->(k[1] in used_keys), parsed)
    end
    dirs = (TEST_DIR*"/").*joinpath.(readdir(TEST_DIR), "settings.jld2")
    tests = fill(false, 10)
    for i in 1:10
        # @load dirs[i] settings_dict
        # dict = settings_dict
        dict = FileIO.load(dirs[i])
        parsed_args = dict["parsed_args"]
        used_keys = dict["used_keys"]
        tests[i] = filter(k->(k[1] in used_keys), parsed_args) == parsed_dicts[parsed_args["_HASH"]]
    end

    reset_tests()
    return all(tests)
end

function track_symbols_test()

    used_keys = [:a, :b]
    parsed_dicts = Dict{UInt64, Dict}()
    for i in 1:10
        args=["--a", "$(i)", "--b", "$(i).jld"]
        parsed = arg_parse(args; as_symbols=true)
        parsed_dicts[parsed[:_HASH]] = filter(k->(k[1] in used_keys), parsed)
    end
    dirs = (TEST_DIR*"/").*joinpath.(readdir(TEST_DIR), "settings.jld2")
    tests = fill(false, 10)
    for i in 1:10
        # @load dirs[i] settings_dict
        # dict = settings_dict
        dict = FileIO.load(dirs[i])
        parsed_args = dict["parsed_args"]
        used_keys = dict["used_keys"]
        tests[i] = filter(k->(k[1] in used_keys), parsed_args) == parsed_dicts[parsed_args[:_HASH]]
    end

    reset_tests()
    return all(tests)
end

function track_with_git_test()

    git_head = Reproduce.git_head()
    used_keys = ["a", "b"]
    parsed_dicts = Dict{UInt64, Dict}()
    for i in 1:10
        args=["--a", "$(i)", "--b", "$(i).jld"]
        parsed = arg_parse(args; use_git_info=true)
        parsed_dicts[parsed["_HASH"]] = parsed
    end
    dirs = (TEST_DIR*"/").*joinpath.(readdir(TEST_DIR), "settings.jld2")
    tests = fill(false, 10)
    for i in 1:10
        # @load dirs[i] parsed_args used_keys
        dict = FileIO.load(dirs[i])
        # @load dirs[i] settings_dict
        # dict = settings_dict
        parsed_args = dict["parsed_args"]
        used_keys = dict["used_keys"]
        tests[i] = (==(filter(k->(k[1] in used_keys), parsed_args),
                       filter(k->(k[1] in used_keys), parsed_dicts[parsed_args["_HASH"]]))
                    &
                    (parsed_args["_GIT_INFO"] == parsed_dicts[parsed_args["_HASH"]]["_GIT_INFO"] == git_head))
    end

    reset_tests()
    return all(tests)
end

macro testparse()
    @testset "Parse Tests" begin
        @test parse_test() == true
        @test track_test() == true
        @test track_symbols_test() == true
        @test track_with_git_test() == true
    end
end

end


text:
module Bayesutil

greet() = print("Hello World!")
include("btables.jl")
end # module


text:
export facets

"""
    facets(
        vertices :: Vector{Vector{Int64}};
        dir = "./" :: String,
        cleanup=true :: Bool
    ) :: Dict{String, Vector{Vector{Int64}}}

Computes the facet inequalities and equalities which bound the convex polyhedron
defined by the set of `vertices`.
If the optimal representation is used for the vertices, then no equalitities will
be present.
For example, if the "generalized" vertex representation is used the normalization
constraints will be included in the resulting equalities.
The output of this method is structured:

```
Dict(
    "facets" => inequalities, # :: Vector{Vector{Int64}}
    "equalitites" => equalities, # :: Vector{Vector{Int64}}
)
```

Facet inequalities and equalities are represented as single vector `f` where
`f[1:(end-1)]` contains the coefficients of the linear inquality and `f[end]`
is the bound.
Facet inequalities and equalities act upon a vertex `v` where inequalities are
arranged such that there is an implicit `f[1:(end-1)]' * v ≤ f[end]`
and equalities are arranged such that `f[1:(end-1)]' * v == f[end]`

!!! note "Supporting Software"
    The vertex -> facet transformation is performed using the
    [`traf`](https://juliapolyhedra.github.io/XPORTA.jl/dev/exports/#XPORTA.traf)
    method of [XPORTA.jl](https://github.com/JuliaPolyhedra/XPORTA.jl).
    Please refer to the source code for more details.

!!! warning "Performance Limitations"
    Vertex -> facet transformations are notoriously difficult computational problems.
    The performance limits of this method will be met far before the limits the
    vertex enumeration methods.
"""
function facets(
    vertices :: Vector{Vector{Int64}};
    dir = "./" :: String,
    cleanup=true :: Bool
) :: Dict{String, Vector{Vector{Int64}}}
    ieq = traf(POI(vertices = hcat(vertices...)'[:,:]), dir=dir, cleanup=cleanup)

    inequalities = convert.(Int64, ieq.inequalities)
    equalities = convert.(Int64, ieq.equalities)

    Dict(
        "facets" => map(row_id -> inequalities[row_id,:], 1:size(inequalities,1)),
        "equalities" => map(row_id -> equalities[row_id,:], 1:size(equalities,1))
    )
end

"""
    facets(poly::Polyhedron) :: Vector{Vector{Int64}}

Returns the facets of the local polytope `poly`. If the facets have not already
been computed, then they are transformed into the half-space representation from
the vertex representation.
"""
function facets(poly::Polyhedron) :: Vector{Vector{Int64}}
    inequalities = convert.(Int64, hrep(poly).inequalities)
    map(row_id -> inequalities[row_id,:] , 1:size(inequalities,1))
end


text:
"""
generate a stepfunction with the following properties:

     nmax      ┌──┐        ┌──┐
               │  │        │  │
               │  │        │  │
               │  │        │  │
     3      ┌──┘  │     ┌──┘  │
            │     │     │     │
            │     │     │     │
            │     │     │     │
     2   ┌──┘     │  ┌──┘     │
         │        │  │        │
         │        │  │        │
         │        │  │        │
     1 ──┘        └──┘        └
       0          1           2
"""
generate_stepfct(nmax) = x -> max(Int(floor(nmax*mod(x, 1))) + 1, 1)

"""
generate sawtooth function with the following properties:

     b          /│        /│        /│
               / │       / │       / │
              /  │      /  │      /  │
             /   │     /   │     /   │
            /    │    /    │    /    │
           /     │   /     │   /     │
          /      │  /      │  /      │
         /       │ /       │ /       │
     a  /        │/        │/        │
       0         dx        2dx      3dx
"""
generate_sawfct(dx, a, b) = x -> a + (b-a)*mod(x/dx, 1)

"""
generate a sawtooth function on [0,1] to [a,b] which has
jump points at the given points in jumps:

     b   //│    /│
       //  │   / │
           │   / │            //
           │  /  │          //
           │  /  │        //
           │ /   │      //
           │ /   │    //
           │/    │  //
     a     │/    │//
       0  j[1]  j[2]           1

"""
function generate_sawfctx(jumps; a=0.0, b=1.0)
    xk = sort(jumps)
    let xk = [Float64(xk[end]-1), Float64.(xk)..., Float64(xk[1]+1)],
        xk_diff = diff(xk)
      return function(x)
          y = 0.0
          for j = 1:length(xk)-1
              if xk[j] ≤ x < xk[j+1]
                  y += (x - xk[j])/xk_diff[j]
              end
          end
          return a + y*(b-a)
      end
    end
end

"""
generate rectangular impulses of length dx

     1      ┌─────┐     ┌─────┐
            │     │     │     │
            │     │     │     │
            │     │     │     │
            │     │     │     │
            │     │     │     │
            │     │     │     │
            │     │     │     │
     0 ─────┘     └─────┘     └
       0    dx   2dx   3dx

"""
generate_rectfct(dx) = x -> mod(floor(x/dx), 2)

"""
generate 1-periodic downwards "spike".

     1 ─────     ──────
             \\   /
             \\ /
               │
               │
     0         │

"""
generate_spike(; a=-9.0, b=0.5) = x -> 1 - (a == -Inf ? 0.0 : exp(a*abs(mod(x,1)-0.5)^b))

"""
Angle of -z scaled in the interval [0,1]
"""
angle_zero_one(z) = abs(z) == Inf ? 0.0 : (angle(-ComplexF64(z)) + pi)/(2pi)

"""
Generate color_number many colors with equispaced hue-values in [0,360).
"""
hsv_colors(color_number=600; saturation=1.0, value=1.0) = HSV{Float32}.(
    360*(0:color_number-1) ./ color_number, saturation, value)

"""
NIST's piecewise linear phase mapping:
https://dlmf.nist.gov/help/vrml/aboutcolor

q ∈ [0, 4)
Output: Angle deg [0, 360)
"""
nist_trafo(q) = (
    1 ≤ q < 2 ? 2*q - 1  :
    2 ≤ q < 3 ? q + 1    :
    3 ≤ q < 4 ? 2*(q-1)  :
    q) * 60

"""
Generate NIST phase mapping: https://dlmf.nist.gov/help/vrml/aboutcolor
"""
nist_colors(color_number=600; saturation=1.0, value=1.0) = HSV{Float32}.(
    nist_trafo.(4*(0:color_number-1) ./ color_number), saturation, value)

"""
brighten RGB-component.

bright ≥ 0
"""
br_rgb(e, bright) = (1 - bright)*e + bright

"""
darken RGB-component.

bright < 0
"""
dr_rgb(e, bright) = (1 + bright)*e

"""
brighten or darken RBG depending on birght sign.
"""
brightenRGB(rgb, bright) = (
  bright ≥ 0 ?
    RGB(br_rgb(rgb.r, bright), br_rgb(rgb.g, bright), br_rgb(rgb.b, bright))
    :
    RGB(dr_rgb(rgb.r, bright), dr_rgb(rgb.g, bright), dr_rgb(rgb.b, bright)))

"""
Use base_color_func and modify the brightness depending on abs(fz).
"""
function cs_mod_abs_brightness(base_color_func)
    return function(z, fz)
        # Version 1
        absfz = abs(fz)
        logf = log(absfz)
        bright = absfz == 0.0 ? 1.0 :
                 absfz == Inf ? 1.0 :
                 logf ≥ 0 ? (1 - (1/(1+logf)))^2  : -(1 - (1/(1-logf)))^2
        # Version 2
        # bright = (3.2f0/pi)*atan(Float32(abs(fz))) - 0.8f0
        return brightenRGB(RGB(base_color_func(z, fz)), bright)
    end
end

# vim:syn=julia:cc=79:fdm=indent:sw=2:ts=2:


text:
using HttpServer
import JSON
import Unmarshal

include("fx.jl")

srv = Server() do req::Request, res::Response
    if length(req.data) == 0
        res.status = 400
    else
        parsed_data = JSON.Parser.parse(join([Char(v) for v in req.data]))
        data = Unmarshal.unmarshal(Input, parsed_data)
        string(fx(data))
    end
end

run(srv, 3000)

max_stars_repo_name: jimbrig/fx

max_issues_repo_name: jimbrig/fx

max_forks_repo_name: jimbrig/fx


text:
""" XLNet Tokenizer """

"""
The XLNet model uses the sentence-piece model as the sub-word tokenizer.
Sentence-piece is a statistical tokenizer that supports two algorithms: unigram-LM and BPE

Note: These are temporary doc-strings, will be edited.

WIP: Add function to read protocol buffer files from spiece.model
"""

# Sentence Piece Structure :: SentencePieceModel
struct SentencePieceModel
    vocab_map::Dict{String, Tuple{Float64, Int}}
    unk_id::Int
end

# To initialize the tokenizer from the datadeps (vocab/json) :: load(ty::Type{T}, filenum::Int=1; unk_token="<unk>") where T<:PretrainedTokenizer
function load(ty::Type{T}, filenum::Int=1; unk_token="<unk>") where T<:PretrainedTokenizer
    filepath = @datadep_str tokenizer_files(ty)[filenum]
    name = tokenizer_files(ty)[filenum]
    filepath = "$filepath/$name"
    load_vocab(filepath, unk_token=unk_token)
end

function load_vocab_from_protobuf(ty::Type{T}, filenum::Int=1; unk_token="<unk>") where T<:PretrainedTokenizer
    filepath = @datadep_str tokenizer_files(ty)[filenum]
    name = tokenizer_files(ty)[filenum]
    """
    - Deserialize spiece.model
    - readproto
    - load into sentence piece model below
    """
    if haskey(vocab_map, unk_token)
        unk_id = vocab_map[unk_token][2]
    else
        throw(DomainError(unk_token, "Unknown token is not in the vocabulary"))
    end

    spm = SentencePieceModel(vocab_map, unk_id)
    return spm
end

# To initialize the SentencePieceModel by providing the `vocab filepath` :: load_vocab(path; unk_token="<unk>")
function load_vocab(path; unk_token="<unk>")
    vocab_path = readlines(path)
    vocabnlogp = split.(vocab_path, "\t")
    vocab_map = Dict(tok=>(parse(Float64, logp), index) for (index, (tok, logp)) in enumerate(vocabnlogp))

    if haskey(vocab_map, unk_token)
        unk_id = vocab_map[unk_token][2]
    else
        throw(DomainError(unk_token, "Unknown token is not in the vocabulary"))
    end

    spm = SentencePieceModel(vocab_map, unk_id)
    return spm
end

# Utility structure, To hold the results of the `forward pass` (the forward Viterbi lattice) :: Nodes (struct)
"""
holds
    text::String     : token string
    Score::Float32   : score
    index::Int64     : vocabulary index
    start::Int       : start character position
    en::Int          : end character position
"""
struct Nodes
    text::String
    score::Float32
    index::Int64
    start::Int
    en::Int
end

# Perform a forward pass to get the forward Viterbi lattice : decode_forward(sp::SentencePieceModel, text::String)
# Ref.[--- BLOG ---]
"""
    decode_forward(sp::SentencePieceModel, text::String) - Takes model structure and String
    returns results : Array{Nodes,1}
"""
function decode_forward(sp::SentencePieceModel, text::String)
    results = Array{Nodes, 1}(undef, lastindex(text))
    scores = fill(-Inf, lastindex(text))
    scores[1] = 0

    for char_end in eachindex(text)

        for char_start in eachindex(text)
            char_start > char_end && break
            subtoken = SubString(text, char_start:char_end)
            if haskey(sp.vocab_map, subtoken)
                subtokenid =  sp.vocab_map[subtoken][2]
                local_score = scores[char_start] + sp.vocab_map[subtoken][1]
                if local_score > scores[char_end]
                    results[char_end] = Nodes(SubString(text, char_start:char_end), local_score, subtokenid, char_start, char_end)
                    scores[char_end] = local_score
                end
            end
        end

        if scores[char_end] == -Inf
            results[char_end] = Nodes(SubString(text, prevind(text, char_end):char_end), -Inf, 1, char_end-1, char_end)
            scores[char_end] = 0
        end

        if scores[char_end] == 0
            results[char_end] = Nodes(SubString(text, char_end:char_end), -Inf, 1, char_end, char_end)
        end

    end

    return results
end

# Perform a backward pass : decode_backward(sp::SentencePieceModel, nodes::Array{Nodes,1}, text::AbstractString)
# Ref.[--- BLOG ---]
"""
    decode_backward(sp::SentencePieceModel, nodes::Array{Nodes,1}, text::AbstractString)
    - Takes model strcuture, nodes (i.e. output of `decode_forward`), and String
    returns best_seq : Array{Nodes,1}
"""
function decode_backward(sp::SentencePieceModel, nodes::Array{Nodes,1}, text::AbstractString)
    next_nodes = nodes[end]
    best_seq = Nodes[]

    while next_nodes.start > 1
        node_value = next_nodes
        next_nodes = nodes[prevind(text, node_value.start)]
        push!(best_seq, node_value)
    end
    push!(best_seq, next_nodes)
    return best_seq
end

# To perform the preprocessing followed by forward decoding and backward decoding
"""
    tokenizer(sp::SentencePieceModel,text::AbstractString)
    preprocessing --> decode_forward --> decode_backward
    output tokenized tokens as: Array{String,1}
"""
function tokenizer(sp::SentencePieceModel, text::AbstractString)
    text = replace(text, " " => "▁")
    if text[1] != '▁'
        text = "▁" * text
    end
    output = decode_forward(sp, text)
    tokens = decode_backward(sp, output, text)
    tokens = reverse(tokens)
    tks = [node.text for node in tokens]
    return tks

end

# To perform the preprocessing followed by forward decoding and backward decoding
"""
    (sp::SentencePieceModel)(text::AbstractString)
    preprocessing --> decode_forward --> decode_backward
"""
function (sp::SentencePieceModel)(text::AbstractString)
    tokenizer(sp, text)
end

# To get the ids from the tokens : ids_from_tokens(spm::SentencePieceModel, tk::Array{String,1})
"""
    ids_from_tokens(spm::SentencePieceModel, tk::Array{String,1})
    Given tokens it provide its indices
"""
function ids_from_tokens(spm::SentencePieceModel, tk::Array{String,1})
    map(tk) do x
        last(get(spm.vocab_map, x, spm.unk_id))
    end
end

# To get sentence from the tokens
"""
    sentence_from_tokens(tk::Array{String,1})
    Given tokens it provide its sentences
"""
function sentence_from_tokens(tk::Array{String,1})
    sen = join(tk)
    sen = replace(sen, "▁" => " ")
    sen = strip(sen)
    return sen
end


text:
## Functions for handling zoom:
# Zooms zr to the decimal % entered; view centered around center XY
function zoom_percent(z::Float64, zr::ZoomRegion, center::XY{Int})
    # Calculate size of new view
    fsize = XY(zr.fullview.x.right,zr.fullview.y.right) # full size
    csize = XY(Int(round(fsize.x/z)), Int(round(fsize.y/z))) # new current size
    # Calculate center point of new view
    offset = XY(center.x-Int(round(csize.x/2)),
                center.y-Int(round(csize.y/2))) # offset of cv
    # Limit offset
    if offset.x < 0
        y = offset.y
        offset = XY(0, y)
    elseif offset.x > (fsize.x-csize.x)
        y = offset.y
        offset = XY(fsize.x-csize.x, y)
    end
    if offset.y < 0
        x = offset.x
        offset = XY(x, 0)
    elseif offset.y > (fsize.y-csize.y)
        x = offset.x
        offset = XY(x, fsize.y-csize.y)
    end

    return (offset.y+1..offset.y+csize.y, offset.x+1..offset.x+csize.x)
end # return value can be pushed to a zr

# Sets default center to be the middle of the cv
function zoom_percent(z::Float64, zr::ZoomRegion)
    # Calculate cv
    return zoom_percent(z,zr,find_center(zr))
end

# Finds rounded center point of the current view of given ZoomRegion
function find_center(zr::ZoomRegion)
    range = zr.currentview
    csize = XY(range.x.right-range.x.left,range.y.right-range.y.left)
    center = XY(range.x.left+Int(floor(csize.x/2)),
                range.y.left+Int(floor(csize.y/2)))
    return center
end

# Zoom tracking
const zpercents = [1.0,1.2,1.5,2.0,2.5,3.0,4.0,8.0]

# For a given zoom region, finds the level that zoom_percent actually zooms to,
# for every item in zpercents. Used in zoom tracking.
function actual_zpercents_x(zr::ZoomRegion)
    levels = zoom_percent.(zpercents, zr) #returns an array of tuples
    zp_actual = []
    for i in 1:length(levels)
        push!(zp_actual, IntervalSets.width(zr.fullview.x)/
              IntervalSets.width(levels[i][2]))
    end
    return zp_actual
end

# Returns index of next zoom level after current zoom level in zpercents.
function next_zoom(ctx::ImageContext)
    xzoom=IntervalSets.width(value(ctx.zr).fullview.x)/IntervalSets.width(value(ctx.zr).currentview.x)
    zp_actual = actual_zpercents_x(value(ctx.zr))
    index = 1
    for n in zp_actual
        if n > xzoom
            break
        end
        index += 1
    end
    return index
end

# Returns index of zoom level before current in zpercents.
function prev_zoom(ctx::ImageContext)
    zr = value(ctx.zr)
    xzoom=IntervalSets.width(zr.fullview.x)/IntervalSets.width(zr.currentview.x)
    zp_actual = actual_zpercents_x(value(ctx.zr))
    index = length(zpercents)
    for n in zp_actual[end:-1:1] # loop backwards
        if n < xzoom
            break
        end
        index -= 1
    end
    return index
end

# Performs proportional zoom in and tracks zoom level using zpercents.
function zoom_in(ctx::ImageContext, center::XY{Int})
    i = ctx.zl
    zr = ctx.zr
    if 1 <= i <= length(zpercents)
        if i < length(zpercents)
            i += 1
            push!(zr, zoom_percent(zpercents[i],value(zr),center))
        end
    else
        i = next_zoom(value(zr))
        push!(zr, zoom_percent(zpercents[i],value(zr),center))
    end
    ctx.zl = i
    ctx.zr = zr
end

# Automatically centered zoom_in
function zoom_in(ctx::ImageContext)
    zoom_in(ctx, find_center(value(ctx.zr)))
end

# Performs proportional zoom out and tracks zoom level using zpercents.
function zoom_out(ctx::ImageContext, center::XY{Int})
    i = ctx.zl
    zr = ctx.zr
    if 1 <= i <= length(zpercents)
        if i > 1
            i -= 1
            push!(zr, zoom_percent(zpercents[i],value(zr),center))
        end
    else
        i = prev_zoom(value(zr))
        push!(zr, zoom_percent(zpercents[i],value(zr),center))
    end
    ctx.zl = i
    ctx.zr = zr
end

# Automatically centered zoom_out
function zoom_out(ctx::ImageContext)
    zoom_out(ctx, find_center(value(ctx.zr)))
end

# Performs proportional, centered zoom to level entered
function zoom_to(ctx::ImageContext, z::Float64)
    ctx.zl = -1
    push!(ctx.zr, zoom_percent(z,value(ctx.zr)))
    nothing
end

# Mouse actions for zoom
function init_zoom_click(ctx::ImageContext)
    c = ctx.canvas
    zr = ctx.zr
    enabled = Signal(true)
    # Left click calls zoom_in() centered on pixel clicked
    # Right click calls zoom_out() centered on pixel clicked
    dragging = Signal(false)
    moved = Signal(false)
    start = Signal(XY{UserUnit}(-1,-1))
    start_view = Signal(ZoomRegion((1:1,1:1)))

    dummybtn = MouseButton{UserUnit}()
    sigclick = map(filterwhen(enabled,dummybtn,c.mouse.buttonpress)) do btn
        push!(dragging,true)
        push!(moved,false)
        push!(start,btn.position)
        push!(start_view,value(zr))
    end

    sigdrag = map(filterwhen(dragging, dummybtn, c.mouse.motion)) do btn
        value(start_view) != value(zr) && push!(moved,true)
        nothing
    end

    sigend = map(filterwhen(dragging,dummybtn,c.mouse.buttonrelease)) do btn
        if !value(moved)
            #println("modifiers=",btn.modifiers)
            if btn.button == 1 && btn.modifiers == 256 #if left click & no modifiers
                center = XY(Int(round(Float64(btn.position.x))),
                            Int(round(Float64(btn.position.y))))
                zoom_in(ctx, center)
            elseif btn.button == 3 || btn.modifiers == 260 # right click/ctrl
                center = XY(Int(round(Float64(btn.position.x))),
                            Int(round(Float64(btn.position.y))))
                zoom_out(ctx, center)
            end
        end
        push!(dragging,false) # no longer dragging
        push!(moved,false) # reset moved
    end
    append!(c.preserved, [moved, sigclick, sigdrag, sigend])
    Dict("enabled"=>enabled)
end


text:
istetgen(g)=typeof(g)==Module && nameof(g)==:TetGen

istriangulate(g)=typeof(g)==Module && nameof(g)==:Triangulate

abstract type TetGenType  end

abstract type TriangulateType  end


text:
# This file is a part of Julia. License is MIT: https://julialang.org/license

@eval IterativeEigensolvers begin
    Base.A_mul_B!(y::StridedVector{T}, A::AtA_or_AAt{T}, x::StridedVector{T}) where {T} = Base.LinAlg.mul!(y, A, x)
    Base.A_mul_B!(y::StridedVector{T}, A::SVDAugmented{T}, x::StridedVector{T}) where {T} = Base.LinAlg.mul!(y, A, x)
end

max_stars_repo_name: YingboMa/julia

max_issues_repo_name: YingboMa/julia

max_forks_repo_name: YingboMa/julia


text:
using Compat

"""
$(SIGNATURES)

Given value(s) `x` apply the sigmoidal function with maximum
value `l`, a steepness `k`, and an inflection point `p`.
"""
function sigmoid(x::Float64, l, k, p)
    k = min(p, min(1-p, k))
    (x <= p-k) && return 0
    (x >= p+k) && return l
    x = (x-p)/k
    l/2*(sign(x).*21abs(x)./(20*abs(x)+1)+1)
end

"""
$(SIGNATURES)

Given value(s) `x` apply a simple linear function with maximum value `l`
"""
linear(x, l) = clamp(l.*x, min(0, l), max(0, l))

"""
A type representing a relationship between degree of knockdown and effect on
phenotype
"""
@compat abstract type KDPhenotypeRelationship end

type Linear <: KDPhenotypeRelationship end

type Sigmoidal <: KDPhenotypeRelationship
    "Slopes of the sigmoid are pulled from this distribution"
    width_dist::Distribution

    "Inflection points are pulled from this distribution"
    inflection_dist::Distribution

    function Sigmoidal()
        # width of sigmoidal region
        width = Normal(0.1, 0.05)
        # inflections are centered at high KD
        inflections = TruncatedNormal(0.8, 0.2, 0, 1)
        new(width, inflections)
    end
end

response(::Linear) = linear
function response(sig::Sigmoidal)
    width, inflection = rand(sig.width_dist), rand(sig.inflection_dist)
    (x, l) -> sigmoid(x, l, width, inflection)
end

"""
A type representing the behavior of different Cas9s
"""
@compat abstract type Cas9Behavior end

"""
$(TYPEDEF)

CRISPRi behavior is simply determined by the activity of the guide
"""
type CRISPRi <: Cas9Behavior end

"""
$(TYPEDEF)

CRISPR KO behavior is more complex since sgRNA-directed DNA damage repair is
stochastic. We assume that 2/3 of repair events at a given locus lead to a
frameshift, and that the screen is carried out in diploid cells. The assumption
that only bi-allelic frame-shift mutations lead to a phenotype in
CRISPRn screens for most sgRNAs is supported by the empirical finding that
in-frame deletions mostly do not show strong phenotypes, unless they occur in
regions encoding conserved residues or domains[^2]

[^2]: Horlbeck MA, Gilbert LA, Villalta JE, Adamson B, Pak RA, Chen Y, Fields AP,
    Park CY, Corn JE, Kampmann M, Weissman JS: Compact and highly active next-
    generation libraries for CRISPR-mediated gene repression and activation.
    *Elife* 2016, 5.
"""
type CRISPRn <: Cas9Behavior
    knockout_dist::Categorical

    CRISPRn(dist::Categorical) = new(dist)
end
CRISPRn() = CRISPRn(Categorical([1/9, 4/9, 4/9]))

"""
Wrapper containing all library construction parameters

$(FIELDS)
"""
type Library
    "Distribution of guide knockdown efficiencies"
    knockdown_dist::Dict{Int, Tuple{Symbol, Sampleable}}
    knockdown_probs::Categorical

    """
    Maximum phenotype categories mapped to their probability of being
    selected and the distribution to draw from if they are selected.

    ## Example

    The basic layout is a `Dict` mapping a class name to a tuple of the
    probability of selecting this class and then the
    [`Distributions.Sampleable`](https://juliastats.github.io/Distributions.jl/latest/types.html#Sampleable-1)
    from which to draw a random phenotype from this class. The probabilities
    across all the classes should add up to 1.

    ```julia
    max_phenotype_dists = Dict{Symbol, Tuple{Float64, Sampleable}}(
        :inactive => (0.60, Delta(0.0)),
        :negcontrol => (0.1, Delta(0.0)),
        :increasing => (0.3, TruncatedNormal(0.1, 0.1, 0.025, 1)),
    );
    Library(max_phenotype_dists, CRISPRi());
    ```

    For example, here we are making three different classes of "genes": the
    first group are :inactive, i.e. they have no phenotype, so we'll set their
    phenotypes to 0.0 using a [`Simulation.Delta`](@ref). We'll also make them
    60% of all the genes. The second group are the negative controls :negcontrol
    (the only required group) which make up 10% of the population of genes and
    also have no effect. The final group is :increasing which makes up 30% of
    all genes and which are represented by a Normal(μ=0.1, σ=0.1) distribution
    clamped between 0.025 and 1.
    """
    max_phenotype_dists::Dict{Int, Tuple{Symbol, Sampleable}}
    phenotype_probs::Categorical

    """
    Knockdown-phenotype relationships mapped to their probability of
    being selected and their respective [`Simulation.KDPhenotypeRelationship`](@ref)
    """
    kd_phenotype_relationships::Dict{Int, Tuple{Symbol, KDPhenotypeRelationship}}
    relationship_probs::Categorical

    """
    Whether this library is [`Simulation.CRISPRi`](@ref) or [`Simulation.CRISPRn`](@ref)
    """
    cas9_behavior::Cas9Behavior

    function Library(knockdown_dist::Dict{Symbol, Tuple{Float64, Sampleable}},
                     max_phenotype_dists::Dict{Symbol, Tuple{Float64, Sampleable}},
                     kd_phenotype_relationships::Dict{Symbol, Tuple{Float64, KDPhenotypeRelationship}},
                     cas9_behavior::Cas9Behavior)

        kd = unroll(knockdown_dist)
        max_p = unroll(max_phenotype_dists)
        rela = unroll(kd_phenotype_relationships)
        new(kd[1], kd[2], max_p[1], max_p[2], rela[1], rela[2], cas9_behavior)
    end
end

function Library(cas9_behavior::Cas9Behavior)
    max_phenotype_dists = Dict{Symbol, Tuple{Float64, Sampleable}}(
        :inactive => (0.75, Delta(0.0)),
        :negcontrol => (0.05, Delta(0.0)),
        :increasing => (0.1, TruncatedNormal(0.55, 0.2, 0.1, 1)),
        :decreasing => (0.1, TruncatedNormal(-0.55, 0.2, -1, -0.1))
    )
    Library(max_phenotype_dists, cas9_behavior)
end

function Library(max_phenotype_dists::Dict{Symbol, Tuple{Float64, Sampleable}},
                 cas9_behavior::CRISPRi)
    # Assuming a high quality library has mostly good guides with some bad ones
    knockdown_dist = Dict{Symbol, Tuple{Float64, Sampleable}}(
        :high => (0.9, TruncatedNormal(0.90, 0.1, 0, 1)),
        :low => (0.1, TruncatedNormal(0.05, 0.07, 0, 1))
    )
    Library(max_phenotype_dists, knockdown_dist, cas9_behavior)
end

function Library(max_phenotype_dists::Dict{Symbol, Tuple{Float64, Sampleable}},
                 cas9_behavior::CRISPRn)
    # For CRISPR KO assume that if guide is "high quality" than it will a
    # maximum knockdown of 100%
    knockdown_dist = Dict{Symbol, Tuple{Float64, Sampleable}}(
        :high => (0.9, Delta(1.0)),
        :low => (0.1, TruncatedNormal(0.05, 0.07, 0, 1))
    )
    Library(max_phenotype_dists, knockdown_dist, cas9_behavior)
end

function Library(max_phenotype_dists::Dict{Symbol, Tuple{Float64, Sampleable}},
                 knockdown_dist::Dict{Symbol, Tuple{Float64, Sampleable}},
                 cas9_behavior::Cas9Behavior)

    kd_phenotype_relationships = Dict{Symbol, Tuple{Float64, KDPhenotypeRelationship}}(
        :linear => (0.75, Linear()),
        :sigmoidal => (0.25, Sigmoidal())
    )
    Library(knockdown_dist, max_phenotype_dists, kd_phenotype_relationships, cas9_behavior)
end

function as_array(lib::Library)
    cas9_behavior = typeof(lib.cas9_behavior)
    frac_hq = 0.0
    mean_hq_kd = 0.0
    for (id, (quality, dist)) in lib.knockdown_dist
        if quality == :high
            mean_hq_kd = cas9_behavior == CRISPRi ? mean(dist.untruncated) : NaN
            frac_hq = lib.knockdown_probs.p[id]
        end
    end

    relationships = Symbol(sort([behavior for (mapping, (behavior, dist)) in
                                 lib.kd_phenotype_relationships])...)

    frac_inc_genes, frac_dec_genes = 0.0, 0.0
    cent_inc_phenotype = 0.0
    for (id, (quality, dist)) in lib.max_phenotype_dists
        if quality == :increasing
            frac_inc_genes = lib.phenotype_probs.p[id]
            cent_inc_phenotype = mean(dist.untruncated)
        end
        (quality == :decreasing) && (frac_dec_genes = lib.phenotype_probs.p[id])
    end
    [frac_hq, mean_hq_kd, frac_inc_genes, frac_dec_genes, cent_inc_phenotype, relationships, cas9_behavior]
end

array_names(::Type{Library}) = [:frac_hq, :mean_hq_kd, :frac_inc, :frac_dec, :cent_inc_phenotype, :relationships, :crisprtype]

function unroll{T}(data::Dict{Symbol, Tuple{Float64, T}})
    probs = Float64[]
    results = Dict{Int, Tuple{Symbol, T}}()
    count = 0
    for (key, value) in data
        push!(probs, value[1])
        results[count+=1] = (key, value[2])
    end
    results, Categorical(probs)
end

function rand_gene(lib::Library)
    class, dist = lib.max_phenotype_dists[rand(lib.phenotype_probs)]
    max_phenotype = rand(dist)
    behavior, relationship = lib.kd_phenotype_relationships[rand(lib.relationship_probs)]
    class, max_phenotype, behavior, relationship
end

"""
$(SIGNATURES)

Constructs the guide library for `N` genes with `coverage` number of guides per
gene. Returns a tuple of guides and their relative frequencies (assigned randomly).
"""
function construct_library(setup::ScreenSetup, lib::Library)
    barcodes = Barcode[]
    N = setup.num_genes
    coverage = setup.coverage

    for gene in 1:N

        gene_class, max_phenotype, gene_behavior, relationship = rand_gene(lib)
        # get this gene's KD-phenotype relationship
        kd_response = response(relationship)

        for i in 1:coverage
            barcode_quality, barcode_knockdown_dist = lib.knockdown_dist[rand(lib.knockdown_probs)]
            barcode_knockdown = rand(barcode_knockdown_dist)
            # phenotype of barcode given its knockdown efficiency
            barcode_phenotype = kd_response(barcode_knockdown, max_phenotype)
            push!(barcodes, Barcode(gene, barcode_knockdown, barcode_phenotype,
                  gene_behavior, gene_class))
        end
    end

    # the expected z-score for a 90% confidence interval is 2x1.645=3.29
    # construct a normal distribution with a σ = 1/3.29 so that there is
    # a 10-fold difference between the 95th/5th percentiles
    vals = 10.^rand(Normal(0, 1/3.29), N*coverage)
    guide_freqs = vals/sum(vals)
    barcodes, Categorical(guide_freqs)
end

"""
$(TYPEDEF)

Constructs a delta function at a given δ value. This distribution always emits
the same value.
"""
type Delta <: Distributions.Sampleable{Univariate, Discrete}
    δ::Float64
end

Base.rand(d::Delta) = d.δ


text:
module OpenSSH

include("utilities.jl")
include("documenter_keygen.jl")
include("ssh_keygen.jl")

end


text:
#!/usr/bin/env julia
using Base.Sys: iswindows
using Documenter: doctest
using PortAudio:
    combine_default_sample_rates,
    devices,
    get_default_input_index,
    get_default_output_index,
    get_device,
    get_input_type,
    get_output_type,
    handle_status,
    initialize,
    PortAudioException,
    PortAudio,
    PortAudioDevice,
    PortAudioStream,
    safe_load,
    seek_alsa_conf,
    terminate,
    name
using PortAudio.LibPortAudio:
    Pa_AbortStream,
    PaError,
    PaErrorCode,
    paFloat32,
    Pa_GetDefaultHostApi,
    Pa_GetDeviceInfo,
    Pa_GetHostApiCount,
    Pa_GetLastHostErrorInfo,
    Pa_GetSampleSize,
    Pa_GetStreamCpuLoad,
    Pa_GetStreamInfo,
    Pa_GetStreamReadAvailable,
    Pa_GetStreamTime,
    Pa_GetStreamWriteAvailable,
    Pa_GetVersionInfo,
    Pa_HostApiDeviceIndexToDeviceIndex,
    paHostApiNotFound,
    Pa_HostApiTypeIdToHostApiIndex,
    PaHostErrorInfo,
    paInDevelopment,
    paInvalidDevice,
    Pa_IsFormatSupported,
    Pa_IsStreamActive,
    paNoError,
    paNoFlag,
    paNotInitialized,
    Pa_OpenDefaultStream,
    paOutputUnderflowed,
    Pa_SetStreamFinishedCallback,
    Pa_Sleep,
    Pa_StopStream,
    PaStream,
    PaStreamInfo,
    PaStreamParameters,
    PaVersionInfo
using SampledSignals: nchannels, s, SampleBuf, samplerate, SinSource
using Test: @test, @test_logs, @test_nowarn, @testset, @test_throws

@testset "Tests without sound" begin
    @testset "Reports version" begin
        io = IOBuffer()
        PortAudio.versioninfo(io)
        result = split(String(take!((io))), "\n")
        # make sure this is the same version I tested with
        @test startswith(result[1], "PortAudio V19")
    end

    @testset "Can list devices without crashing" begin
        display(devices())
        println()
    end

    @testset "libortaudio without sound" begin
        @test handle_status(Pa_GetHostApiCount()) >= 0
        @test handle_status(Pa_GetDefaultHostApi()) >= 0
        # version info not available on windows?
        if !Sys.iswindows()
            @test safe_load(Pa_GetVersionInfo(), ErrorException("no info")) isa
                  PaVersionInfo
        end
        @test safe_load(Pa_GetLastHostErrorInfo(), ErrorException("no info")) isa
              PaHostErrorInfo
        @test PaErrorCode(Pa_IsFormatSupported(C_NULL, C_NULL, 0.0)) == paInvalidDevice
        @test PaErrorCode(
            Pa_OpenDefaultStream(Ref(C_NULL), 0, 0, paFloat32, 0.0, 0, C_NULL, C_NULL),
        ) == paInvalidDevice
    end

    @testset "Errors without sound" begin
        @test sprint(showerror, PortAudioException(paNotInitialized)) ==
              "PortAudioException: PortAudio not initialized"
        @test_throws KeyError("foobarbaz") get_device("foobarbaz")
        @test_throws KeyError(-1) get_device(-1)
        @test_throws ArgumentError("Could not find alsa.conf in ()") seek_alsa_conf(())
        @test_logs (:warn, "libportaudio: Output underflowed") handle_status(
            PaError(paOutputUnderflowed),
        )
        @test_throws PortAudioException(paNotInitialized) handle_status(
            PaError(paNotInitialized),
        )
        Pa_Sleep(1)
        @test Pa_GetSampleSize(paFloat32) == 4
    end

    # make sure we can terminate, then reinitialize
    terminate()
    initialize()
end

if !isempty(devices())
    @testset "Tests with sound" begin
        # these default values are specific to local machines
        input_name = get_device(get_default_input_index()).name
        output_name = get_device(get_default_output_index()).name

        @testset "Interactive tests" begin
            println("Recording...")
            stream = PortAudioStream(input_name, output_name, 2, 0; adjust_channels = true)
            buffer = read(stream, 5s)
            @test size(buffer) ==
                  (round(Int, 5 * samplerate(stream)), nchannels(stream.source))
            close(stream)
            sleep(1)
            println("Playing back recording...")
            PortAudioStream(input_name, output_name, 0, 2; adjust_channels = true) do stream
                write(stream, buffer)
            end
            sleep(1)
            println("Testing pass-through")
            stream = PortAudioStream(input_name, output_name, 2, 2; adjust_channels = true)
            sink = stream.sink
            source = stream.source
            @test sprint(show, stream) == """
                PortAudioStream{Float32}
                  Samplerate: 44100.0Hz
                  2 channel sink: $(repr(input_name))
                  2 channel source: $(repr(output_name))"""
            @test sprint(show, source) == "2 channel source: $(repr(output_name))"
            @test sprint(show, sink) == "2 channel sink: $(repr(input_name))"
            write(stream, stream, 5s)
            @test PaErrorCode(handle_status(Pa_StopStream(stream.pointer_to))) == paNoError
            @test isopen(stream)
            close(stream)
            sleep(1)
            @test !isopen(stream)
            @test !isopen(sink)
            @test !isopen(source)
            println("done")
        end
        @testset "Samplerate-converting writing" begin
            PortAudioStream(input_name, output_name, 0, 2; adjust_channels = true) do stream
                write(
                    stream,
                    SinSource(eltype(stream), samplerate(stream) * 0.8, [220, 330]),
                    3s,
                )
                println("expected blip")
                write(
                    stream,
                    SinSource(eltype(stream), samplerate(stream) * 1.2, [220, 330]),
                    3s,
                )
            end
        end
        sleep(1)
        # no way to check that the right data is actually getting read or written here,
        # but at least it's not crashing.
        @testset "Queued Writing" begin
            PortAudioStream(input_name, output_name, 0, 2; adjust_channels = true) do stream
                buffer = SampleBuf(
                    rand(eltype(stream), 48000, nchannels(stream.sink)) * 0.1,
                    samplerate(stream),
                )
                frame_count_1 = @async write(stream, buffer)
                frame_count_2 = @async write(stream, buffer)
                @test fetch(frame_count_1) == 48000
                println("expected blip")
                @test fetch(frame_count_2) == 48000
            end
            sleep(1)
        end
        @testset "Queued Reading" begin
            PortAudioStream(input_name, output_name, 2, 0; adjust_channels = true) do stream
                buffer = SampleBuf(
                    rand(eltype(stream), 48000, nchannels(stream.source)) * 0.1,
                    samplerate(stream),
                )
                frame_count_1 = @async read!(stream, buffer)
                frame_count_2 = @async read!(stream, buffer)
                @test fetch(frame_count_1) == 48000
                @test fetch(frame_count_2) == 48000
            end
            sleep(1)
        end
        @testset "Constructors" begin
            PortAudioStream(2, maximum; adjust_channels = true) do stream
                @test isopen(stream)
            end
            PortAudioStream(output_name; adjust_channels = true) do stream
                @test isopen(stream)
            end
            PortAudioStream(input_name, output_name; adjust_channels = true) do stream
                @test isopen(stream)
            end
        end
        @testset "Errors with sound" begin
            big = typemax(Int)
            @test_throws DomainError(
                typemax(Int),
                "$big exceeds maximum input channels for $output_name",
            ) PortAudioStream(input_name, output_name, big, 0)
            @test_throws ArgumentError("Input or output must have at least 1 channel") PortAudioStream(
                input_name,
                output_name,
                0,
                0;
                adjust_channels = true,
            )
            @test_throws ArgumentError("""
            Default sample rate 0 for input $output_name disagrees with
            default sample rate 1 for output $input_name.
            Please specify a sample rate.
            """) combine_default_sample_rates(
                get_device(input_name),
                0,
                get_device(output_name),
                1,
            )
        end
        @testset "libportaudio with sound" begin
            @test PaErrorCode(Pa_HostApiTypeIdToHostApiIndex(paInDevelopment)) ==
                  paHostApiNotFound
            @test Pa_HostApiDeviceIndexToDeviceIndex(paInDevelopment, 0) == 0
            stream = PortAudioStream(input_name, output_name, 2, 2; adjust_channels = true)
            pointer_to = stream.pointer_to
            @test handle_status(Pa_GetStreamReadAvailable(pointer_to)) >= 0
            @test handle_status(Pa_GetStreamWriteAvailable(pointer_to)) >= 0
            @test Bool(handle_status(Pa_IsStreamActive(pointer_to)))
            @test safe_load(Pa_GetStreamInfo(pointer_to), ErrorException("no info")) isa
                  PaStreamInfo
            @test Pa_GetStreamTime(pointer_to) >= 0
            @test Pa_GetStreamCpuLoad(pointer_to) >= 0
            @test PaErrorCode(handle_status(Pa_AbortStream(pointer_to))) == paNoError
            @test PaErrorCode(
                handle_status(Pa_SetStreamFinishedCallback(pointer_to, C_NULL)),
            ) == paNoError
        end
    end
    doctest(PortAudio)
end


text:
using Plots

#export plot_callback
"""
Plot trajectories of the first nodes of sys and spec, and display the loss.
Parameters:
- pbt: PBT problem
- p: array of sys and spec parameters, usually res.minimizer
- loss: value of the output distance
Optional parameters:
- loss_array: array to append value of loss to. Can be used to plot progress of the optimization over iterations.
- scenario_nums: number(s) of scenarios to be plotted. If not provided, a scenario is chosen randomly.
- fig_name: path and file name to save the plot. If not provided, plot is not saved.
- offset: defines by how much each scenario is offset from the previous one for better plot readability. If set to -1, all samples are plotted on subplots instead of a single figure.
- plot_options: various plot options
"""
function plot_callback(pbt, p, loss; loss_array=nothing, scenario_nums=nothing, fig_name=nothing, offset=2, title = nothing, plot_options...)
    display(loss)
    plt = plot()
    isnothing(scenario_nums) ? scenario_nums = rand(1:pbt.N_samples) : nothing
    if length(scenario_nums) == 1
        dd_sys, dd_spec = solve_sys_spec_n(pbt, scenario_nums[1], p)
        plt = plot(dd_sys, vars=1; label="System output", plot_options...)
        plot!(plt, dd_spec, vars=1; label="Specification output", plot_options...)
        plot!(plt, dd_spec.t, pbt.input_sample[scenario_nums[1]]; c=:gray, alpha=0.75, label="Input", plot_options...)
        title!("Input scenario $(scenario_nums[1])")
        display(plt)
    else
        j = 1
        if offset == -1
            plt = plot(layout = (length(scenario_nums),1))
            for i in scenario_nums
                dd_sys, dd_spec = solve_sys_spec_n(pbt, i, p)
                plot!(plt, dd_sys, vars=1; label = "System output (scenario $i)", c = palette(:tab20)[2*i-1], subplot = j, plot_options...)
                plot!(plt, dd_spec, vars=1; label = "Specification output (scenario $i))", c = palette(:tab20)[2*i], linestyle = :dash, subplot = j, plot_options...)
                plot!(plt, dd_spec.t, pbt.input_sample[i]; c=:gray, alpha=0.5, label = "Input $i", subplot = j, title = ("Input scenario $i"), plot_options...)
                #title!("Input scenarios $i")
                j+=1
          end
          xlabel!("t")
          ylabel!("output")
          display(plt)
        else
            for i in scenario_nums
                dd_sys, dd_spec = solve_sys_spec_n(pbt, i, p)
                plot!(dd_sys.t, dd_sys[1,:] .+ offset * (j - 1), vars=1; label=false, color_palette=:tab20, plot_options...) # yaxis = nothing "System output (scnenario $i)"
                plot!(plt, dd_spec.t, dd_spec[1,:] .+ offset * (j - 1), vars=1; label=false, linestyle=:dash, plot_options...) # "Specification output (scnenario $i))"
                j += 1
            end
            plot!([0.,0.01], [0.,0.];label="System output", c=:gray)
            plot!([0.,0.01], [0.,0.];label="Specification output", c=:gray, linestyle=:dash)
            scenarios_line = ""
            for s in @view scenario_nums[1:end - 1]
                scenarios_line *= "$s, "
            end
            scenarios_line *= "$(scenario_nums[end])"
            if isnothing(title)
                title!("Scenarios " * scenarios_line)
            else
                title!(title)
            end
            xlabel!("t")
            ylabel!("output")
            display(plt)
        end

    end
    if !isnothing(fig_name)
          savefig(fig_name)
    end
    if !isnothing(loss_array)
          append!(loss_array, loss)
    end
      # Tell sciml_train to not halt the optimization.
      # If return true, then optimization stops.
    return false
end

#=
export plot_sys_graph
"""
Plot graph representations of optimized system and spec system.
"""
function plot_sys_graph(pbt::PBTProblem)
    nodecolor = [colorant"orange", colorant"lightseagreen"]
    nodefillc_sys = nodecolor[[1 for i in 1:pbt.size_p_sys]];
    nodefillc_spec = nodecolor[[1 for i in 1:pbt.size_p_spec]];
    nodefillc_sys[1] = nodecolor[2]
    nodefillc_spec[1] = nodecolor[2]
    L_sys = pbt.f_sys.L
    L_spec = pbt.f_spec.L
    g1 = gplot(SimpleGraph(L_sys), nodefillc=nodefillc_sys)
    g2 = gplot(SimpleGraph(L_spec), nodefillc=nodefillc_spec)
    display(g1)
    display(g2)
end
=#


text:
#using Revise
#using CUDA, LinearAlgebra, BenchmarkTools, Zygote, SparseArrays, UnPack
#import Base.*
#abstract type AbstractAnnhilator{T} end

#creates a pre-calculated projection object for the dummy matrix in a fixed effect regression
#this approach should balance speed and performance
struct FixedAnnihilator{TF,T} <: AbstractAnnhilator{T}
  F::TF
  FCFt::TF

  FixedAnnihilator(F::TF,FCFt::TF) where {TF} = new{TF, eltype(TF)}(F,FCFt)
end

FixedAnnihilator(F::TF) where {TF<:SparseMatrixCSC} = FixedAnnihilator(F,cholesky!(Matrix(F'*F))\I*F' |> TF)
function FixedAnnihilator(F::TF, ::Type{<:Any}, ::Type{T} = eltype(TF)) where {TF,TA,T}
  (maximum(F) - T(1) == minimum(F) == zero(T)) || error("Fixed annihilation matrix only
    makes sense with a dummy matrix")
  return FixedAnnihilator(SparseMatrixCSC{T}(F))
end

#version for cuarrays
function FixedAnnihilator(F::TF, ::Type{<:CuArray}, ::Type{T} = eltype(TF)) where {TF,T}
  (maximum(F) - T(1) == minimum(F) == zero(T)) || error("Fixed annihilation matrix only
    makes sense with a dummy matrix")
    M = FixedAnnihilator(F |> Matrix{T}, Matrix{T})
    Fsp = M.F |> CUSPARSE.CuSparseMatrixCSR
    FCFtsp = M.FCFt |> CUSPARSE.CuSparseMatrixCSR

  return FixedAnnihilator(Fsp,FCFtsp)
end

#function *(M::FixedAnnihilator,X::TX, ::Type{T}=eltype(TX))::TX where {T, TX<:AbstractArray}
#  return X .- M.F*(M.FCFt*X)
  #CUSPARSE.mm!('N',T(-1.0),M.F, M.FCFt* X, T(1.0),deepcopy(X),'O')
#end

function *(M::FixedAnnihilator,X::TX, ::Type{T}=eltype(TX))::TX where {T, TX<:AbstractArray}
  #return deepcopy(X)
  return X .- M.F*(M.FCFt*X)
end

#the below is really fast but fragile, hopefully it will get less fragile over time
function *(M::FixedAnnihilator, X::TX, ::Type{T}=eltype(TX))::TX where {T, TX<:CuMatrix}
  #return deepcopy(X)
  return CUSPARSE.mm!('N','N',T(-1.0),M.F, M.FCFt* X, T(1.0),deepcopy(X),'O')
end

#special inplace version for a slight performance boost
function multiply!(M::FixedAnnihilator, X::TX, ::Type{T}=eltype(TX))::TX where {T, TX<:CuMatrix}
  #return X
  return CUSPARSE.mm!('N','N',T(-1.0),M.F, M.FCFt* X, T(1.0),X,'O')
end

#the above high speed operation works much better on a matrix than a vector
function *(M::FixedAnnihilator, y::Ty, ::Type{T}=eltype(Ty))::Ty where {T, Ty<:CuVector}
  #return deepcopy(y)
  return y .- M.F*(M.FCFt*y)#CUSPARSE.mv!('N',T(-1.0),M.F, M.FCFt * y, T(1.0),deepcopy(y),'O')
end

function multiply!(M::FixedAnnihilator, y::Ty, ::Type{T}=eltype(Ty))::Ty where {T, Ty<:CuVector}
  #return y
  return CUSPARSE.mv!('N',T(-1.0),M.F, M.FCFt * y, T(1.0),y,'O')
end

#fallback for the in-place multiply
function multiply!(M::FixedAnnihilator, X::TX, ::Type{T}=eltype(TX)) where {TX,T}
  #return X
  X .= M*X
  return X
end

function testFixedAnnihilator(F::TF, ::Type{T}=eltype(TF)) where {TF,T}
  M = FixedAnnihilator(F)
  v = rand(T, size(F,1))
  X= rand(T, size(F,1), 5)
  @assert Matrix(M.F) ≈ F
  @assert M*v ≈ v .- F*((cholesky!(Matrix(F'*F))\I) * F' * v)
  @assert M*X ≈ X .- F*((cholesky!(Matrix(F'*F))\I) * F' * X)

  return nothing
end

###############################
####Dense Annihilator
###############################

#creates a pre-calculated dense projection object
struct DenseAnnihilator{TW,T} <: AbstractAnnhilator{T}
  W::TW
  WtWinvWt::TW

  DenseAnnihilator(W::TW,WtWinvWt::TW) where {TW} = new{TW, eltype(TW)}(W,WtWinvWt)
#  DenseAnnihilator(::Nothing) where {TW} = new{Nothing, Nothing}(nothing, nothing)
end

function DenseAnnihilator(W::Matrix)

  try
    #return DenseAnnihilator(W, cholesky!(Hermitian(W'*W))\(W'))
    return DenseAnnihilator(W, svd(Symmetric(W'*W))\(W'))
  catch err
    @warn "ERROR! printing printmln(W'*W):"
    printmln(W'*W)
    error("Error: $err")
  end
end

function DenseAnnihilator(W::CuMatrix{T}) where T
  Wcpu = W |> Matrix{T}
  try
    #return DenseAnnihilator(W, cholesky!(Hermitian(W'*W))\(W'))
    return DenseAnnihilator(W, svd(Symmetric(Wcpu'*Wcpu))\(Wcpu') |> CuMatrix{T})
  catch err
    @warn "ERROR! printing printmln(W'*W):"
    printmln(W'*W)
    println("Printing diag of SVD")
    printmln(svd(W'*W).S |> Diagonal)
    println("Printing inverse:")
    printmln(inv(svd(W'*W |> Matrix)))
    if size(W,1) < 10^6
      W |> Matrix |> DataFrame |> CSV.write("$(PARAM[:testpath])\\error_W_$(size(W)).csv")
    end
    error("Error: $err")
  end
end

function *(M::DenseAnnihilator{TW,T}, X::TX)::TX where {TW, TX<:AbstractArray, T}

  return X - M.W*(M.WtWinvWt*X)
end

#special case for if we have no
#*(M::DenseAnnihilator{Nothing,Nothing}, X::TX)::TX where {TX<:AbstractArray} = X

#expand the annihilator matrix
Base.Matrix{T}(M::DenseAnnihilator{TW,T}) where {TW<:AbstractArray,T} = I - M.W * M.WtWinvWt
function CUDA.CuMatrix{T}(M::DenseAnnihilator{TW,T}) where {TW<:CuArray, T}
  @unpack W, WtWinvWt = M

  (N,K) = size(W)
  id = CuMatrix{T}(I, N, N)
  return id - W * WtWinvWt
end

function testdenseannihilator(::Type{TW}, ::Type{T}=eltype(TW); N=2500, K=10, bench=true) where {TW,T}
  W = rand(N, 2K) |> TW
  X = rand(N, K) |> TW
  M = DenseAnnihilator(W)
  A = M |> TW

  ver = X - W*(cholesky(W'*W)\(W'*X)) |> Matrix
  @assert ver ≈ ( (M * X) |> Matrix)
  @assert ver ≈ (A * X) |> Matrix
  verf(X) = A * X |> sum
  testf(X) = M * X |> sum
  @assert (gradient(verf, X)[1] |> Array) ≈ (gradient(testf, X)[1] |> Array)

  if bench
    print("Testing M*X for type $TW:")
    @btime CUDA.@sync $M*$X
    print("Testing A*X for type $TW:")
    @btime CUDA.@sync $A*$X
    print("Testing ∇M*X for type $TW:")
    @btime CUDA.@sync gradient($testf, $X)[1]
    print("Testing ∇A*X for type $TW:")
    @btime CUDA.@sync gradient($verf, $X)[1]
  end

  @info "Passed testdenseannihilator for type $TW"
end

  #test vectorized application of projections
function testdenseannihilator(
    blocks::Int, ::Type{TW}, ::Type{T}=eltype(TW); N=2500, K=10, bench=true)  where {TW,T}

  vecW = [rand(T, N + rand((-K^2):K^2), 2K) |> TW for b ∈ 1:blocks]
  vecX = [rand(T, size(vecW[b],1), K) |> TW  for b ∈ 1:blocks]
  vecM = (DenseAnnihilator).(vecW)

  for b in 1:blocks
    W = vecW[b]
    X = vecX[b]

    @assert (vecM[b] * vecX[b]) |> Matrix ≈ (X - W*(cholesky(W'*W)\(W'*X)) |> Matrix)
  end

  #  (W,X,zb)->Threads.@spawn((W*X) .* zb)).() .|> fetch .|> sum |> sum
  vecmult(z::Tz, ::Type{TW}) where {Tz, TW <: CuMatrix} = (
    (W,X,zb)->(W*X) .* zb |> sum).(vecM, vecX,z) |> sum

  vecmult(z::Tz, ::Type{TW}) where {Tz, TW <: Matrix} = ((W,X,zb)->(W*X) .* zb).(vecM, vecX,z) .|> sum |> sum
  ∇vecmult(z) = gradient(vecmult,z)[1]

  z = [rand(T, 1, K) |> TW  for b ∈ 1:blocks] #testing placeholder

  #warmup- make sure everything works
  CUDA.reclaim()
  @assert !(∇vecmult(z) === nothing)
  @assert reduce(vcat, ∇vecmult(z) .|> Array) .|> abs |> sum > 0.0

  if bench
    print("Testing vecmult for type $TW:")
    @btime CUDA.@sync $vecmult($z)
    print("Testing ∇vecmult for type $TW:")
    @btime CUDA.@sync $∇vecmult($z)
  end
end

function testdenseannihilator(blocks::Int = 1100; N=2500, K=10, bench=false, benchblocks=false)
  testdenseannihilator(Matrix{Float64}; N, K, bench)
  testdenseannihilator(CuMatrix{Float64}; N, K, bench)
  testdenseannihilator(blocks, Matrix{Float64}; N, K, bench=benchblocks)
  testdenseannihilator(blocks, CuMatrix{Float64}; N, K,bench= benchblocks)
end

#sleep(0.2)
#CUDA.allowscalar(false)
#testdenseannihilator(1100, bench=false, benchblocks=true, N=2500, K=10)

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
using RecurrentNN
reload("RecurrentNN.jl")
using Base.Test

# graph output test
m1 = RecurrentNN.NNMatrix(3,2)
m1.w[1,1] = 1.; m1.w[1,2] = 2.
m1.w[2,1] = 3.; m1.w[2,2] = 4.
m1.w[3,1] = 5.; m1.w[3,2] = 6.

m2 = RecurrentNN.NNMatrix(2,3)
m2.w[1,1] = 2.; m2.w[1,2] = 3.; m2.w[1,3] = 4.
m2.w[2,1] = 5.; m2.w[2,2] = 6.; m2.w[2,3] = 7.

# add test
g =  RecurrentNN.Graph()
m3 = RecurrentNN.add!(g,m1,m1)
m3.dw[1,1] = .1; m3.dw[1,2] = .2
m3.dw[2,1] = .3; m3.dw[2,2] = .4
m3.dw[3,1] = .5; m3.dw[3,2] = .6

g.backprop[1]()
# set previous chained gradient
@test m1.dw[1,1] == 0.2
@test m1.dw[1,2] == 0.4
@test m1.dw[2,1] == 0.6
@test m1.dw[2,2] == 0.8
@test m1.dw[3,1] == 1.0
@test m1.dw[3,2] == 1.2

# mul test
m1.dw[:] = 0. # reset gradient matrices
m2.dw[:] = 0. # reset  gradient matrices
g =  RecurrentNN.Graph()
m3 = RecurrentNN.mul!(g,m1,m2)
@test m3.w[1,1] == 12.
@test m3.w[1,2] == 15.
@test m3.w[1,3] == 18.
@test m3.w[2,1] == 26.
@test m3.w[2,2] == 33.
@test m3.w[2,3] == 40.
@test m3.w[3,1] == 40.
@test m3.w[3,2] == 51.
@test m3.w[3,3] == 62.

m3.dw[1,1] = .1; m3.dw[1,2] = .2; m3.dw[1,3] = .3
m3.dw[2,1] = .4; m3.dw[2,2] = .5; m3.dw[2,3] = .6
m3.dw[3,1] = .7; m3.dw[3,2] = .8; m3.dw[3,3] = .9
g.backprop[1]()

# m1 gradient tests
@test m1.dw[1,1] == 2.
@test m1.dw[1,2] == 3.8000000000000003
@test m1.dw[2,1] == 4.699999999999999
@test m1.dw[2,2] == 9.2
@test m1.dw[3,1] == 7.4
@test m1.dw[3,2] == 14.600000000000001

# m2 gradient tests
@test m2.dw[1,1] == 4.800000000000001
@test m2.dw[1,2] == 5.7
@test m2.dw[1,3] == 6.6
@test m2.dw[2,1] == 5.999999999999999
@test m2.dw[2,2] == 7.200000000000001
@test m2.dw[2,3] == 8.4

# reul() tests
m4 = RecurrentNN.NNMatrix(3,2)
m4.w[1,1] = 1.; m4.w[1,2] =-2.
m4.w[2,1] =-3.; m4.w[2,2] = 4.
m4.w[3,1] = 5.; m4.w[3,2] =-6.
g =  RecurrentNN.Graph()
m5 = RecurrentNN.relu!(g,m4)
@test m5.w[1,1] == 1.
@test m5.w[1,2] == 0.
@test m5.w[2,1] == 0.
@test m5.w[2,2] == 4.
@test m5.w[3,1] == 5.
@test m5.w[3,2] == 0.

m5.dw[1,1] =-.1; m5.dw[1,2] = .2
m5.dw[2,1] = .3; m5.dw[2,2] = .4
m5.dw[3,1] = .5; m5.dw[3,2] = .6

g.backprop[1]()
@test m4.dw[1,1] == -0.1
@test m4.dw[1,2] == 0.
@test m4.dw[2,1] == 0.
@test m4.dw[2,2] == 0.4
@test m4.dw[3,1] == 0.5
@test m4.dw[3,2] == 0.

# rowpluck!() tests
m4 = RecurrentNN.NNMatrix(3,2)
m4.w[1,1] = 1.; m4.w[1,2] =-2.
m4.w[2,1] =-3.; m4.w[2,2] = 4.
m4.w[3,1] = 5.; m4.w[3,2] =-6.

g =  RecurrentNN.Graph()
m5 = RecurrentNN.rowpluck!(g,m4,2)
@test m5.w[1,1] == -3.
@test m5.w[2,1] == 4.

m5.dw[1,1] =-.1; m5.dw[2,1] = .2
g.backprop[1]()
m4.dw
@test m4.dw[1,1] == 0
@test m4.dw[1,2] == 0
@test m4.dw[2,1] ==-0.1
@test m4.dw[2,2] == 0.2
@test m4.dw[3,1] == 0.
@test m4.dw[3,2] == 0.

# softmax tests
m6 = RecurrentNN.NNMatrix(5,1)
m6.w[1,1] = 0.3; m6.w[2,1] =0.1; m6.w[3,1] =0.6; m6.w[4,1] = 0.002; m6.w[5,1] = 0.00001
sm = RecurrentNN.softmax(m6)
# @test_approx_eq
@test_approx_eq sm.w[1,1] 0.21494050089813527
@test_approx_eq sm.w[2,1] 0.17597839816728894
@test_approx_eq sm.w[3,1] 0.2901393282421457
@test_approx_eq sm.w[4,1] 0.1595506217827435
@test_approx_eq sm.w[5,1] 0.15939115090968653

max_forks_repo_name: haroldsoh/RecurrentNN


text:
"""
    NegBinomialLikelihood(r::Real)

## Arguments
- `r::Real` number of failures until the experiment is stopped

---
[Negative Binomial likelihood](https://en.wikipedia.org/wiki/Negative_binomial_distribution) with number of failures `r`
```math
    p(y|r, f) = binomial(y + r - 1, y) (1 - σ(f))ʳ σ(f)ʸ
    p(y|r, f) = Γ(y + r)/Γ(y + 1)Γ(r) (1 - σ(f))ʳ σ(f)ʸ
```
Where `σ` is the logistic function
"""
struct NegBinomialLikelihood{T<:Real,Tr<:Real,A<:AbstractVector{T}} <: EventLikelihood{T}
    r::Tr
    c::A
    θ::A
    function NegBinomialLikelihood{T}(r::Real) where {T<:Real}
        return new{T,typeof(r),Vector{T}}(r)
    end
    function NegBinomialLikelihood{T}(
        r::Real, c::A, θ::A
    ) where {T<:Real,A<:AbstractVector{T}}
        return new{T,typeof(r),A}(r, c, θ)
    end
end

function NegBinomialLikelihood(r::Real)
    return NegBinomialLikelihood{Float64}(r)
end

implemented(::NegBinomialLikelihood, ::Union{<:AnalyticVI,<:GibbsSampling}) = true

function init_likelihood(
    likelihood::NegBinomialLikelihood{T}, ::AbstractInference{T}, ::Int, nSamplesUsed::Int
) where {T}
    return NegBinomialLikelihood{T}(
        likelihood.r, rand(T, nSamplesUsed), zeros(T, nSamplesUsed)
    )
end

function (l::NegBinomialLikelihood)(y::Real, f::Real)
    return pdf(NegativeBinomial(lr, get_p(l, f)), y)
end

function Distributions.loglikelihood(l::NegBinomialLikelihood, y::Real, f::Real)
    return logpdf(NegativeBinomial(lr, get_p(l, f)), y)
end

function expec_count(l::NegBinomialLikelihood, f)
    return broadcast((p, r) -> p * r ./ (1 .- p), get_p.(l, f), l.r)
end

function get_p(::NegBinomialLikelihood, f)
    return logistic.(f)
end

function Base.show(io::IO, l::NegBinomialLikelihood{T}) where {T}
    return print(io, "Negative Binomial Likelihood (r = $(l.r))")
end

function compute_proba(
    l::NegBinomialLikelihood{T}, μ::AbstractVector{<:Real}, σ²::AbstractVector{<:Real}
) where {T<:Real}
    N = length(μ)
    pred = zeros(T, N)
    sig_pred = zeros(T, N)
    for i in 1:N
        x = pred_nodes .* sqrt(max(σ²[i], zero(T))) .+ μ[i]
        pred[i] = dot(pred_weights, get_p.(l, x))
        sig_pred[i] = dot(pred_weights, get_p.(l, x) .^ 2) - pred[i]^2
    end
    return pred, sig_pred
end

## Local Updates ##

function local_updates!(
    l::NegBinomialLikelihood{T}, y::AbstractVector, μ::AbstractVector, Σ::AbstractVector
) where {T}
    @. l.c = sqrt(abs2(μ) + Σ)
    @. l.θ = (l.r + y) / l.c * tanh(0.5 * l.c)
end

function sample_local!(l::NegBinomialLikelihood, y::AbstractVector, f::AbstractVector)
    return set_ω!(l, rand.(PolyaGamma.(y .+ Int(l.r), abs.(f))))
end

## Global Updates ##

@inline function ∇E_μ(
    l::NegBinomialLikelihood{T}, ::AOptimizer, y::AbstractVector
) where {T}
    return (0.5 * (y .- l.r),)
end
@inline function ∇E_Σ(
    l::NegBinomialLikelihood{T}, ::AOptimizer, y::AbstractVector
) where {T}
    return (0.5 .* l.θ,)
end

## ELBO Section ##

AugmentedKL(l::NegBinomialLikelihood, y::AbstractVector) = PolyaGammaKL(l, y)

function logabsbinomial(n, k)
    return log(binomial(n, k))
end

function negbin_logconst(y, r::Real)
    return loggamma.(y .+ r) - loggamma.(y .+ 1) .- loggamma(r)
end

function negbin_logconst(y, r::Int)
    return logabsbinomial.(y .+ (r - 1), y)
end

function expec_loglikelihood(
    l::NegBinomialLikelihood{T},
    ::AnalyticVI,
    y::AbstractVector,
    μ::AbstractVector,
    diag_cov::AbstractVector,
) where {T}
    tot = Zygote.@ignore(sum(negbin_logconst(y, l.r))) - log(2.0) * sum(y .+ l.r)
    tot += 0.5 * dot(μ, (y .- l.r)) - 0.5 * dot(l.θ, μ) - 0.5 * dot(l.θ, diag_cov)
    return tot
end

function PolyaGammaKL(l::NegBinomialLikelihood, y::AbstractVector)
    return PolyaGammaKL(y .+ l.r, l.c, l.θ)
end


text:
"""
Functions to read atomic data.
"""

"""
    get_atomic_stage(element::String, stage; source="NIST")

Returns the level structure for a given atomic stage for `element`
and ionisation stage `stage`. Uses data from the given `source`,
returns an `AtomicStage` struct. Currently, the only supported source
is `"NIST"`, using locally-saved data obtained from the
[NIST Atomic Spectra Database Levels Form](https://physics.nist.gov/PhysRefData/ASD/levels_form.html).

# Examples
```julia-repl
julia> MgII = get_atomic_stage("Mg", "II")
AtomicStage("Mg", "Mg_II", 2, II, 137, (...))

julia> MgII = get_atomic_stage("Mg", 2)
AtomicStage("Mg", "Mg_II", 2, II, 137, (...))
```
"""
function get_atomic_stage(element::String, stage; source="NIST")
    if element ∉ element_symbols
        error("Invalid element $element")
    end
    if source == "NIST"
        return read_NIST(element, stage)
    else
        error("NotImplemented: atomic data source $source not available.")
    end
end

"""
Reads NIST atomic level data saved locally. The data were extracted from the
[NIST Atomic Spectra Database Levels Form](https://physics.nist.gov/PhysRefData/ASD/levels_form.html).
"""
function read_NIST(element::String, stage)
    stage = RomanNumeral(stage)
    file = string(element, "_", repr(stage), ".txt")
    filepath = joinpath(@__DIR__, "..", "data", "NIST", file)
    if isfile(filepath)
        data = readdlm(filepath)
        # Index entries that have statistical weights
        index = typeof.(data[:, 4]) .== Int
        index .*= data[:, 3] .!= "---"  # some cases with mismatched wnum as Int
        index .*= data[:, 3] .!= ""     # cases with no J or g
        g = convert.(Int, data[index, 4])
        χ = NIST_wavenumber_to_energy.(data[index, 5])
        # Find first ionisation edge
        if sum(data[:, 2] .== "Limit") == 0
            χ_ion = 0.0u"J"
        else
            wavenum_ion = data[data[:, 2] .== "Limit", 4][1]
            χ_ion = NIST_wavenumber_to_energy(wavenum_ion)
        end
        return AtomicStage(element, stage, g, χ, χ_ion)
    else
        error("NIST data for $element $(repr(stage)) not found.")
    end
end

"""
Parses level energy field from NIST tables. Brackets (round or square)
indicate interpolated or theoretical values. Converts from wavenumber
to energy.
"""
function NIST_wavenumber_to_energy(wavenum)
    to_remove = ["[", "]", "(", ")", "?", "a", "l", "x", "y", "z",
                 "u", "+", "†", "&dgger;"]
    if typeof(wavenum) in [String, SubString{String}]
        for suffix in to_remove
            wavenum = replace(wavenum, suffix => "")
        end
        wn = parse(Float64, wavenum)u"cm^-1"
    elseif typeof(wavenum) <: Real
        wn = convert(Float64, wavenum)u"cm^-1"
    else
        error("Invalid type $(typeof(wavenum)) for wave number")
    end
    return (h * c_0 * wn) |> u"J"
end


text:
using PackageCompiler

mode = get(ENV, "MAKIE_COMPILE", "")
if isempty(mode) # don't do anything
elseif mode == "force"
    PackageCompiler.compile_package(joinpath(@__DIR__, ".."), true)
elseif mode == "build"
    PackageCompiler.compile_package(joinpath(@__DIR__, ".."), false)
end


text:
# import libraries
using GeometricIntegrators
using ProgressMeter

# import Particles package
using Particles

# import HDF5 package for storing solution
using HDF5

# parameters
Δt = 0.1    # time step size
nt = 100    # number of time steps
ns = 5      # number of samples

# output file
h5file = "charged_particles.hdf5"

# random initial conditions
z₀ = collect(hcat(rand(ns) .* 2 .- 1, rand(ns) .* 2 .- 1, zeros(ns), rand(ns) .- 0.5, rand(ns) .- 0.5, ones(ns))')

# electric field
E(x::Vector{DT}) where {DT} = DT[0, 0, cos(2π*x[3])]

# magnetic field
B(x::Vector{DT}) where {DT} = DT[0, 0, 1]

# vector field
function lorentz_force!(t, z, ż)
    for i in axes(z,2)
        x = z[1:3,i]
        v = z[4:6,i]

        e = E(x)
        b = B(x)

        ż[1,i] = v[1]
        ż[2,i] = v[2]
        ż[3,i] = v[3]

        ż[4,i] = e[1] + v[2] * b[3] - v[3] * b[2]
        ż[5,i] = e[2] + v[3] * b[1] - v[1] * b[3]
        ż[6,i] = e[3] + v[1] * b[2] - v[2] * b[1]
    end
end

# solution storage
function copy_to_hdf5(h5z, z, n)
    h5z[:,:,n+1] = z
end

# create an Equation instance
equ = ODE(lorentz_force!, z₀)

# create HDF5 file and copy initial conditions
h5  = h5open(h5file, "w")
h5z = create_dataset(h5, "z", eltype(z₀), ((6, ns, nt+1), (6, ns, -1)), chunk=(6,ns,1))
copy_to_hdf5(h5z, z₀, 0)

# create integrator
int = IntegratorExplicitEuler(equ, Δt)

# create atomic solution
asol = AtomicSolution(equ)

# copy initial conditons to atomic solution
set_initial_conditions!(asol, equ)

# initilize integrator
initialize!(int, asol)

# loop over time steps showing progress bar
try
    @showprogress 5 for n in 1:nt
        integrate_step!(int, asol)
        copy_to_hdf5(h5z, asol.q, n)
    end
finally
    # close HDF5 file
    close(h5)
end

# load Plots package
using Plots

# read array from HDF5 file
z = h5read(h5file, "z")

# set plot ranges
xlim = (-2, +2)
ylim = (-2, +2)
zlim = ( 0, 20)

# create empty 3d plot with ns empty series
plt = plot3d(
    ns,
    xlim = xlim,
    ylim = ylim,
    zlim = zlim,
    title = "Charged Particles",
    legend = false,
    marker = 5,
    size = (800, 600)
)

# plot solutions
for j in axes(z,3)
    for i in axes(z,2)
        push!(plt, i, z[1,i,j], z[2,i,j], z[3,i,j])
    end
end

# save figure to file
savefig("charged_particles.png")


text:
struct PolyharmonicSpline <: UQModel
    w::Array{Float64}
    v::Array{Float64}
    c::Array{Float64,2}
    k::Int64
    n::Array{Symbol}
    output::Symbol

    function PolyharmonicSpline(data::DataFrame, k::Int64, output::Symbol)
        f = data[:, output]

        centers = select(data, Not(output))
        n = propertynames(centers)
        centers = Matrix{Float64}(centers)

        dim = size(centers, 1)

        A = zeros(dim, dim)
        for i in 1:dim, j in 1:dim
            if i == j
                continue
            end
            r = sqrt(sum((centers[i, :] - centers[j, :]) .^ 2))
            A[i, j] = ϕ(r, k)
        end

        B = [ones(dim, 1) centers]

        M = [A B; transpose(B) zeros(size(B, 2), size(B, 2))]
        F = [f; zeros(size(B, 2), 1)]

        wv = M \ F

        w = wv[1:dim, :]
        v = wv[(dim + 1):size(wv, 1), :]

        return new(w, v, centers, k, n, output)
    end
end

function ϕ(r::Float64, k::Int64)
    if k % 2 != 0
        return r^k
    elseif r < 1
        return r^(k - 1) * log(r^r)
    else
        return r^k * log(r)
    end
end

function calc(ps::PolyharmonicSpline, x::Array{Float64,1})
    r = sqrt.(sum((ps.c .- transpose(x)) .^ 2; dims=2))
    f = sum(ϕ.(r, ps.k) .* ps.w)
    return f += (transpose(ps.v) * [1; x])[1]
end

function evaluate!(ps::PolyharmonicSpline, df::DataFrame)
    x = Matrix{Float64}(df[:, ps.n]) # convert to matrix and order variables by ps.n

    out = map(row -> calc(ps, convert(Array, row)), eachrow(x))
    return df[!, ps.output] = out
end

max_stars_repo_name: FriesischScott/julia-uq

max_issues_repo_name: FriesischScott/julia-uq

max_forks_repo_name: FriesischScott/julia-uq


text:
mutable struct TensorParser
    preprocessors::Vector{Any} # any preprocessing steps
    contractiontreebuilder::Any # determine a contraction tree for a contraction involving multiple tensors
    contractiontreesorter::Any # transforms the contraction expression into an expression of nested binary contractions using the tree output from the contractiontreebuilder

    postprocessors::Vector{Any}
    function TensorParser()
        preprocessors = [normalizeindices,
                            expandconj,
                            nconindexcompletion,
                            extracttensorobjects]
        contractiontreebuilder = defaulttreebuilder
        contractiontreesorter = defaulttreesorter
        postprocessors = [_flatten, removelinenumbernode, addtensoroperations]
        return new(preprocessors,
                    contractiontreebuilder,
                    contractiontreesorter,
                    postprocessors)
    end
end

function (parser::TensorParser)(ex::Expr)
    if ex.head == :function
        return Expr(:function, ex.args[1], parser(ex.args[2]))
    end
    for p in parser.preprocessors
        ex = p(ex)::Expr
    end
    treebuilder = parser.contractiontreebuilder
    treesorter = parser.contractiontreesorter
    ex = processcontractions(ex, treebuilder, treesorter)::Expr
    ex = tensorify(ex)::Expr
    for p in parser.postprocessors
        ex = p(ex)::Expr
    end
    return ex
end

function processcontractions(ex::Expr, treebuilder, treesorter)
    if ex.head == :macrocall && ex.args[1] == Symbol("@notensor")
        return ex
    end
    ex = Expr(ex.head, map(e->processcontractions(e, treebuilder, treesorter), ex.args)...)
    if istensorcontraction(ex) && length(ex.args) > 3
        args = ex.args[2:end]
        network = map(getindices, args)
        for a in getallindices(ex)
            count(a in n for n in network) <= 2 ||
                throw(ArgumentError("invalid tensor contraction: $ex"))
        end
        tree = treebuilder(network)
        ex = treesorter(args, tree)
    end
    return ex
end
processcontractions(ex, treebuilder, treesorter) = ex

function defaulttreesorter(args, tree)
    if isa(tree, Int)
        return args[tree]
    else
        return Expr(:call, :*,
                        defaulttreesorter(args, tree[1]),
                        defaulttreesorter(args, tree[2]))
    end
end

function defaulttreebuilder(network)
    if isnconstyle(network)
        tree = ncontree(network)
    else
        tree = Any[1,2]
        for k = 3:length(network)
            tree = Any[tree, k]
        end
    end
    return tree
end

# functions for parsing and processing tensor expressions
function tensorify(ex::Expr)
    if ex.head == :macrocall && ex.args[1] == Symbol("@notensor")
        return ex.args[3]
    end
    # assignment case
    if isassignment(ex) || isdefinition(ex)
        lhs, rhs = getlhs(ex), getrhs(ex)
        if isa(rhs, Expr) && rhs.head == :call && rhs.args[1] == :throw
            return rhs
        end

        # process left hand side
        if istensor(lhs) && istensorexpr(rhs)
            indices = getindices(rhs)
            if hastraceindices(lhs)
                err = "left hand side of an assignment should have unique indices: $lhs"
                return :(throw(IndexError($err)))
            end
            dst, leftind, rightind = decomposetensor(lhs)
            if Set(vcat(leftind,rightind)) != Set(indices)
                err = "non-matching indices between left and right hand side: $ex"
                return :(throw(IndexError($err)))
            end
            if isassignment(ex)
                if ex.head == :(=)
                    return Expr(:(=), dst,instantiate(dst, false, rhs, true, leftind, rightind))
                elseif ex.head == :(+=)
                    return Expr(:(=), dst,instantiate(dst, true, rhs, 1, leftind, rightind))
                else
                    return Expr(:(=), dst,instantiate(dst, true, rhs, -1, leftind, rightind))
                end
            else
                return Expr(:(=), dst, instantiate(nothing, false, rhs, true, leftind, rightind, false))
            end
        elseif isassignment(ex) && isscalarexpr(lhs)
            if istensorexpr(rhs) && isempty(getindices(rhs))
                return Expr(ex.head, instantiate_scalar(lhs), Expr(:call, :scalar, instantiate(nothing, false, rhs, true, [], [], true)))
            elseif isscalarexpr(rhs)
                return Expr(ex.head, instantiate_scalar(lhs), instantiate_scalar(rhs))
            end
        else
            return ex # likely an error
        end
    end
    if ex.head == :block
        return Expr(ex.head, map(tensorify, ex.args)...)
    end
    if ex.head == :for
        return Expr(ex.head, ex.args[1], tensorify(ex.args[2]))
    end
    if ex.head == :function
        return Expr(ex.head, ex.args[1], tensorify(ex.args[2]))
    end
    # constructions of the form: a = @tensor ...
    if isscalarexpr(ex)
        return instantiate_scalar(ex)
    end
    if istensorexpr(ex)
        if !isempty(getindices(ex))
            err = "cannot evaluate $ex to a scalar: uncontracted indices"
            return :(throw(IndexError($err)))
        end
        return Expr(:call, :scalar, instantiate(nothing, false, ex, true, [], [], true))
    end
    error("invalid syntax in @tensor macro: $ex")
end
tensorify(ex) = ex


text:
function accept_rparen(ps)
    if kindof(ps.nt) === Tokens.RPAREN
        return mPUNCTUATION(next(ps))
    else
        return mErrorToken(ps, mPUNCTUATION(Tokens.RPAREN, 0, 0), UnexpectedToken)
    end
end
accept_rparen(ps::ParseState, args) = push!(args, accept_rparen(ps))

function accept_rsquare(ps)
    if kindof(ps.nt) === Tokens.RSQUARE
        return mPUNCTUATION(next(ps))
    else
        return mErrorToken(ps, mPUNCTUATION(Tokens.RSQUARE, 0, 0), UnexpectedToken)
    end
end
accept_rsquare(ps::ParseState, args) = push!(args, accept_rsquare(ps))

function accept_rbrace(ps)
    if kindof(ps.nt) === Tokens.RBRACE
        return mPUNCTUATION(next(ps))
    else
        return mErrorToken(ps, mPUNCTUATION(Tokens.RBRACE, 0, 0), UnexpectedToken)
    end
end
accept_rbrace(ps::ParseState, args) = push!(args, accept_rbrace(ps))

function accept_end(ps::ParseState)
    if kindof(ps.nt) === Tokens.END
        return mKEYWORD(next(ps))
    else
        return mErrorToken(ps, mKEYWORD(Tokens.END, 0, 0), UnexpectedToken)
    end
end
accept_end(ps::ParseState, args) = push!(args, accept_end(ps))

function accept_comma(ps)
    if iscomma(ps.nt)
        return mPUNCTUATION(next(ps))
    else
        return mErrorToken(mPUNCTUATION(Tokens.COMMA, 0, 0), UnexpectedToken)
    end
end
accept_comma(ps::ParseState, args) = push!(args, accept_comma(ps))

function requires_ws(x, ps)
    if x.span == x.fullspan
        return mErrorToken(ps, x, Unknown)
    else
        return x
    end
end

function requires_no_ws(x, ps)
    if !(ps.nt.kind === Tokens.RPAREN ||
        ps.nt.kind === Tokens.RBRACE ||
        ps.nt.kind === Tokens.RSQUARE) && x.span != x.fullspan
        return mErrorToken(ps, x, UnexpectedWhiteSpace)
    else
        return x
    end
end


text:
using Dash
using DashHtmlComponents
using DashCoreComponents

app = dash()

app.layout = html_div() do
    html_label("x"),
    dcc_input(id="x", value = 3),
    html_label("y"),
    dcc_input(id="y", value = 6),
    html_label("x + y (clientside)"),
    dcc_input(id="x-plus-y"),
    html_label("x+y / 2 (serverside)"),
    dcc_input(id="x-plus-y-div-2"),
    html_div() do
        html_label("Display x, y, x+y/2 (serverside)"),
        dcc_textarea(id = "display-all-of-the-values")
    end,
    html_label("Mean(x, y, x+y, x+y/2) (clientside)"),
    dcc_input(id="mean-of-all-values")
end

callback!(
    ClientsideFunction("clientside", "add"),
    app,
    Output("x-plus-y", "value"),
    [Input("x","value"), Input("y","value")]
    )

callback!(app, Output("x-plus-y-div-2", "value"), Input("x-plus-y", "value")) do value
    return Float64(value) / 2.
end

callback!(app,
        Output("display-all-of-the-values", "value"),
        [Input("x","value"), Input("y","value"), Input("x-plus-y","value"), Input("x-plus-y-div-2","value")]
        ) do args...
    return join(string.(args), "\n")
end

callback!(
    ClientsideFunction("clientside", "mean"),
    app,
    Output("mean-of-all-values","value"),
    [Input("x","value"), Input("y","value"), Input("x-plus-y","value"), Input("x-plus-y-div-2","value")]
)

run_server(app)