text:
"""
This script contains all functions used for TSP
"""

"Construct TSP pre-model, without cycle elimination mechanism"
function prebuild_tsp(dist::Matrix)

    n = size(dist, 1)
    # Definition of model
    model = Model(GLPK.Optimizer)
    set_optimizer_attribute(model, "tm_lim", 300 * 1_000)
    #model = Model(Gurobi.Optimizer)
    #set_optimizer_attributes(model, "OutputFlag" => 0)

    # Main variable: x_ij=1 if the tour visits i and j in that order, 0 otherwise
    @variable(model, x[1:n, 1:n], Bin)
    # Objective: minimizing the total cost (distance) of the tour
    @objective(model, Min, sum(dist[i, j] * x[i, j] for i = 1:n, j = 1:n))
    # SHARED CONSTRAINTS
    @constraint(
        model, no_self_edges[i = 1:n], x[i,i] == 0
    )
    @constraint(
        model, exactly_one_successor[i = 1:n], sum(x[i, j] for j = 1:n) == 1
    )
    @constraint(
        model, exactly_one_predecessor[j = 1:n], sum(x[i, j] for i = 1:n) == 1
    )
    return model, x
end

"""
    Given the induced graph as an adjacency list (i.e., next[i] is the next node to visit after node i),
        compute all subtours.
    Return them as a list of lists of nodes in the same component
"""
function find_subtours(next::Vector{Int})
    n = length(next)
    g = DiGraph(n)
    for i = 1:n
        add_edge!(g, i, next[i])
    end
    components = strongly_connected_components(g)
    return sort(components, by=length)
end

"Solve the TSP using an iterative approach"
function solve_iterative(dist_matrix::Dict, zone_list::Array; lower_bound=0, time_limit_seconds::Real = 600,
                         eliminate_length_2::Bool=false,
                         verbose::Bool = true)
    # We first solve the model without any subtour elimination consideration

    dist = zeros(length(zone_list), length(zone_list))
    for i in 1:length(zone_list), j in 1:length(zone_list)
        if i != j
            dist[i, j] = dist_matrix[zone_list[i]][zone_list[j]]
        end
    end

    model, x = prebuild_tsp(dist)
    n = size(dist,1)
    @constraint(model, sum(dist[i, j] * x[i, j] for i = 1:n, j = 1:n) >= lower_bound)
    if eliminate_length_2
        @constraint(model, no_length_2[i = 1:n, j = 1:n], x[i, j] + x[j, i] <= 1)
    end
    verbose || set_optimizer_attribute(model, "OutputFlag", 0)
    start=time()
    optimize!(model)

    while true
        # We store the incumbent solution
        next = [findfirst(x -> x > 0.5, value.(x[i, :])) for i = 1:n]
        # Note: checking for >0.5 is conservative (x is binary!) but it avoids numerical errors
        subtours = find_subtours(next)
        #println("Found $(length(subtours)) subtours after $(time() - start) seconds")
        if length(subtours) == 1 # only one cycle, the TSP solution
            solvetime = time() - start
            return solvetime, objective_value(model), value.(x)
        else
            # eliminate subtours
            for subtour in subtours
                @constraint(model, sum(x[i, j] for i=subtour, j=setdiff(1:n, subtour)) >= 1)
                @constraint(model, sum(x[i, j] for i=setdiff(1:n, subtour), j=subtour) >= 1)
            end
        end
        optimize!(model)
        time() - start > time_limit_seconds && return solvetime, objective_value(model), value.(x)
    end
end

function calculate_tsp_seq(route_df, zone_distance_matrix; first_zone = "INIT")

    zone_list = unique(skipmissing(route_df.zone_id))
    opt_zone_seq = zone_list
    compute_time, obj, opt_x = solve_iterative(zone_distance_matrix, zone_list)
    opt_ind = find_subtours([findfirst(x -> x > 0.5, opt_x[i, :]) for i = 1:size(opt_x,1)])[1]
    init_output_seq = [zone_list[i] for i in opt_ind]
    init_ind = findfirst(x -> x == first_zone,init_output_seq)
    zone_output_seq = vcat(init_output_seq[init_ind:end],init_output_seq[1:init_ind-1])

    return zone_output_seq
end

"Construct path TSP pre-model, without cycle elimination mechanism"
function prebuild_path_tsp(dist::Matrix, s, t)

    n = size(dist, 1)
    other_node_list = setdiff(collect(1:n), [s, t])
    # Definition of model
    model = Model(GLPK.Optimizer)
    set_optimizer_attribute(model, "tm_lim", 300 * 1_000)
    #model = Model(Gurobi.Optimizer)
    #set_optimizer_attributes(model, "OutputFlag" => 0)

    # Main variable: x_ij=1 if the tour visits i and j in that order, 0 otherwise
    @variable(model, x[1:n, 1:n], Bin)
    # Objective: minimizing the total cost (distance) of the tour
    @objective(model, Min, sum(dist[i, j] * x[i, j] for i = 1:n, j = 1:n))
    # SHARED CONSTRAINTS
    @constraint(
        model, no_self_edges[i = 1:n], x[i,i] == 0
    )
    @constraint(
        model, exactly_one_successor[i in other_node_list], sum(x[i, j] for j = 1:n) == 1
    )
    @constraint(model, sum(x[s, j] for j in 1:n) == 1)
    @constraint(model, sum(x[j, s] for j in 1:n) == 0)
    @constraint(model, sum(x[i, t] for i in 1:n) == 1)
    @constraint(model, sum(x[t, i] for i in 1:n) == 0)
    @constraint(
        model, exactly_one_predecessor[j in other_node_list], sum(x[i, j] for i = 1:n) == 1
    )
    return model, x
end

"""
    Given the induced graph as an adjacency list (i.e., next[i] is the next node to visit after node i),
        compute all subtours.
    Return them as a list of lists of nodes in the same component
"""
function find_path_subtours(next)
    n = length(next)
    g = DiGraph(n)
    for i = 1:n
        if next[i] != nothing
            add_edge!(g, i, next[i])
        end
    end
    cycles = simplecycles(g)
    #components = connected_components(g)
    return cycles #sort(components, by=length)
end

"Solve the TSP using an iterative approach"
function solve_path_iterative(dist_matrix::Dict, zone_list::Array, first_stop, last_stop; time_limit_seconds::Real = 600,
                         eliminate_length_2::Bool=false)
    # We first solve the model without any subtour elimination consideration

    dist = zeros(length(zone_list), length(zone_list))
    for i in 1:length(zone_list), j in 1:length(zone_list)
        if i != j
            dist[i, j] = dist_matrix[zone_list[i]][zone_list[j]]
        end
    end

    s = findfirst(x -> x == first_stop, zone_list)
    t = findfirst(x -> x == last_stop, zone_list)

    model, x = prebuild_path_tsp(dist, s, t)
    n = size(dist,1)
    if eliminate_length_2
        @constraint(model, no_length_2[i = 1:n, j = 1:n], x[i, j] + x[j, i] <= 1)
    end
    start=time()
    optimize!(model)

    while true
        # We store the incumbent solution
        next = [findfirst(x -> x > 0.5, value.(x[i, :])) for i = 1:n]
        # Note: checking for >0.5 is conservative (x is binary!) but it avoids numerical errors
        subtours = find_path_subtours(next)
        #println("Found $(length(subtours)) subtours after $(time() - start) seconds")
        if length(subtours) == 0 # no cycles
            solvetime = time() - start
            return solvetime, objective_value(model), value.(x), s
        else
            # eliminate subtours
            for subtour in subtours
                @constraint(model, sum(x[i, j] for i=subtour, j=setdiff(1:n, subtour)) >= 1)
                @constraint(model, sum(x[i, j] for i=setdiff(1:n, subtour), j=subtour) >= 1)
            end
        end
        optimize!(model)
        time() - start > time_limit_seconds && return solvetime, objective_value(model), value.(x), s
    end
end

function calculate_path_tsp_seq(route_df, first_zone, last_zone, zone_distance_matrix)

    zone_list = unique(skipmissing(route_df.zone_id))
    if first_zone != "INIT" && last_zone != "INIT"
        zone_list = zone_list[2:end]
    end

    adj_zone_distance_matrix = Dict()
    for i in zone_list
        adj_zone_distance_matrix[i] = Dict()
        for j in zone_list
            adj_zone_distance_matrix[i][j] = 0
        end
    end

    for i in zone_list
        for j in zone_list
            adj_zone_distance_matrix[i][j] = zone_distance_matrix[i][j]
        end
    end

    compute_time, obj, opt_x = solve_path_iterative(adj_zone_distance_matrix, zone_list, first_zone, last_zone)

    next = [findfirst(x -> x > 0.5, opt_x[i, :]) for i = 1:size(opt_x,1)]

    output_seq = [first_zone]
    while length(output_seq) < length(zone_list)
        next_zone_ind = next[findfirst(x -> x == output_seq[end], zone_list)]
        next_zone = zone_list[next_zone_ind]
        push!(output_seq, next_zone)
    end

    @assert output_seq[1] == first_zone
    @assert output_seq[end] == last_zone

    return output_seq
end

function find_first_and_last_stop_list(prev_node, zone, next_zone, zone_to_node_dict, stop_distance_matrix)

    zone_node = zone_to_node_dict[zone]
    next_zone_node = zone_to_node_dict[next_zone]

    prev_dist_list = []
    next_dist_list = []
    for node in zone_node

        prev_dist = stop_distance_matrix[prev_node][node]
        push!(prev_dist_list, (prev_dist, node))

        next_dist = 0
        for next_node in next_zone_node
            next_dist += stop_distance_matrix[node][next_node]
        end
        push!(next_dist_list, (next_dist, node))
    end

    first_node_list = [i[2] for i in sort(prev_dist_list)[1:min(length(zone_node), 3)]]
    last_node_list = [i[2] for i in sort(next_dist_list)[1:min(length(zone_node), 3)]]

    return first_node_list, last_node_list
end

function generate_complete_seq_multiple(zone_seq, route_df, stop_distance_matrix)
    zone_output_seq = zone_seq # opt_zone_seq

    # Generate zone-node correspondence dict
    zone_to_node_dict = Dict(i => [] for i in zone_seq)
    selected_df = route_df[:, ["stops", "zone_id"]]
    for i in 1:size(selected_df, 1)
        stop_id = selected_df[i,"stops"]
        zone_id = selected_df[i,"zone_id"]
        push!(zone_to_node_dict[zone_id], stop_id)
    end

    init_stop = zone_to_node_dict["INIT"][1]
    output_complete_sequence = [init_stop]
    prev_node = init_stop
    for i in 2:length(zone_output_seq)
        curr_zone = zone_output_seq[i]
        # Zone only contains one stop
        if length(zone_to_node_dict[curr_zone]) <= 1
            push!(output_complete_sequence, zone_to_node_dict[curr_zone][1])
            continue
        end
        if i == length(zone_output_seq)
            next_zone = zone_output_seq[1]
        else
            next_zone = zone_output_seq[i+1]
        end
        first_node_list, last_node_list = find_first_and_last_stop_list(prev_node, curr_zone, next_zone, zone_to_node_dict, stop_distance_matrix)
        stop_list = zone_to_node_dict[curr_zone]

        inner_option_list = []
        for first_node in first_node_list
            for last_node in last_node_list
                if first_node == last_node
                    compute_time, obj, opt_x = solve_iterative(stop_distance_matrix, stop_list)
                    opt_ind = find_subtours([findfirst(x -> x > 0.5, opt_x[i, :]) for i = 1:size(opt_x,1)])[1]
                    init_output_seq = [stop_list[i] for i in opt_ind]
                    init_ind = findfirst(x -> x == first_node,init_output_seq)
                    output_path = vcat(init_output_seq[init_ind:end],init_output_seq[1:init_ind-1])
                    reduced_obj = obj - stop_distance_matrix[output_path[end]][output_path[1]]

                    push!(inner_option_list, (reduced_obj, first_node, last_node, output_path))
                else
                    compute_time, obj, opt_x, s = solve_path_iterative(stop_distance_matrix, stop_list, first_node, last_node);
                    next = [findfirst(x -> x > 0.5, opt_x[i, :]) for i = 1:size(opt_x,1)]
                    output_seq = [s]
                    while length(output_seq) < length(stop_list)
                        next_visit_node = next[output_seq[end]]
                        push!(output_seq, next_visit_node)
                    end
                    output_path = [stop_list[i] for i in output_seq]

                    @assert first_node == output_path[1]
                    @assert last_node == output_path[end]

                    push!(inner_option_list, (obj, first_node, last_node, output_path)) # output_path
                end
            end
        end

        inner_zone_tsp_path = sort(inner_option_list)[1][4]
        first_node = inner_zone_tsp_path[1]
        last_node = inner_zone_tsp_path[end]

        output_complete_sequence = vcat(output_complete_sequence, inner_zone_tsp_path)

        prev_node = last_node
    end

    return output_complete_sequence
end

function adjust_zone_dist_matrix(zone_distance_matrix, zone_list, route_min_tt_nearest, route_ecu_tt_nearest, β, γ, α)

    adj_zone_dist_matrix = Dict(i => Dict() for i in zone_list)

    for i in zone_list, j in zone_list

        try
            global big_zone_i = split(i, ".")[1]
            global big_zone_j = split(j, ".")[1]
        catch
            global big_zone_i = i
            global big_zone_j = j
        end

        if big_zone_i != "INIT" && big_zone_j != "INIT" && big_zone_i != big_zone_j
            adj_zone_dist_matrix[i][j] = zone_distance_matrix[i][j] * β
        else
            adj_zone_dist_matrix[i][j] = zone_distance_matrix[i][j]
        end

        if i == "INIT"
            if j ∉ route_ecu_tt_nearest[route_ecu_tt_nearest.from_zone .== i, :].to_zone && j ∉ route_min_tt_nearest[route_min_tt_nearest.from_zone .== i, :].to_zone
                adj_zone_dist_matrix[i][j] = zone_distance_matrix[i][j] * α
            end
        end

        # if i == "INIT" && j ∉ route_zone_min_tt_df[route_zone_min_tt_df.from_zone .== i, :].to_zone
        #     adj_zone_dist_matrix[i][j] = zone_distance_matrix[i][j] * α
        # end

        if i == "INIT" || j == "INIT"
            continue
        end

        try
            global after_dot_i1 = split(i, ".")[2][1]
            global after_dot_i2 = split(i, ".")[2][2]
            global after_dot_j1 = split(j, ".")[2][1]
            global after_dot_j2 = split(j, ".")[2][2]
            global first_diff = abs(Int(after_dot_i1) - Int(after_dot_j1))
            global second_diff = abs(Int(after_dot_i2) - Int(after_dot_j2))
            global sum_diff = first_diff + second_diff
        catch
            global sum_diff = 0
        end

        if big_zone_i == big_zone_j && sum_diff != 1
            adj_zone_dist_matrix[i][j] = adj_zone_dist_matrix[i][j] * γ
        else
            adj_zone_dist_matrix[i][j] = adj_zone_dist_matrix[i][j]
        end
    end

    return adj_zone_dist_matrix
end


text:
#importing the module
using annsp

function test_annsp()
        train_data = {{5.1,3.5,1.4}, {4.9,3.0,1.4}, {4.7,3.2,1.3}, {4.6,3.1,1.5}, {5.0,3.6,1.4}, {5.4,3.9,1.7}, {4.6,3.4,1.4}, {5.0,3.4,1.5}, {4.4,2.9,1.4}, {4.9,3.1,1.5}, {5.4,3.7,1.5}, {4.8,3.4,1.6}, {4.8,3.0,1.4}, {4.3,3.0,1.1}, {5.8,4.0,1.2}, {5.7,4.4,1.5}, {5.4,3.9,1.3}, {5.1,3.5,1.4}, {5.7,3.8,1.7}, {5.1,3.8,1.5}, {5.4,3.4,1.7}, {5.1,3.7,1.5}, {4.6,3.6,1.0}, {5.1,3.3,1.7}, {4.8,3.4,1.9}, {5.0,3.0,1.6}, {5.0,3.4,1.6}, {5.2,3.5,1.5}, {5.2,3.4,1.4}, {4.7,3.2,1.6}, {4.8,3.1,1.6}, {5.4,3.4,1.5}, {5.2,4.1,1.5}, {5.5,4.2,1.4}, {4.9,3.1,1.5}, {5.0,3.2,1.2}, {5.5,3.5,1.3}, {4.9,3.1,1.5}, {4.4,3.0,1.3}, {5.1,3.4,1.5}, {5.0,3.5,1.3}, {4.5,2.3,1.3}, {4.4,3.2,1.3}, {5.0,3.5,1.6}, {5.1,3.8,1.9}, {4.8,3.0,1.4}, {5.1,3.8,1.6}, {4.6,3.2,1.4}, {5.3,3.7,1.5}, {5.0,3.3,1.4}, {7.0,3.2,4.7}, {6.4,3.2,4.5}, {6.9,3.1,4.9}, {5.5,2.3,4.0}, {6.5,2.8,4.6}, {5.7,2.8,4.5}, {6.3,3.3,4.7}, {4.9,2.4,3.3}, {6.6,2.9,4.6}, {5.2,2.7,3.9}, {5.0,2.0,3.5}, {5.9,3.0,4.2}, {6.0,2.2,4.0}, {6.1,2.9,4.7}, {5.6,2.9,3.6}, {6.7,3.1,4.4}, {5.6,3.0,4.5}, {5.8,2.7,4.1}, {6.2,2.2,4.5}, {5.6,2.5,3.9}, {5.9,3.2,4.8}, {6.1,2.8,4.0}, {6.3,2.5,4.9}, {6.1,2.8,4.7}, {6.4,2.9,4.3}, {6.6,3.0,4.4}, {6.8,2.8,4.8}, {6.7,3.0,5.0}, {6.0,2.9,4.5}, {5.7,2.6,3.5}, {5.5,2.4,3.8}, {5.5,2.4,3.7}, {5.8,2.7,3.9}, {6.0,2.7,5.1}, {5.4,3.0,4.5}, {6.0,3.4,4.5}, {6.7,3.1,4.7}, {6.3,2.3,4.4}, {5.6,3.0,4.1}, {5.5,2.5,4.0}, {5.5,2.6,4.4}, {6.1,3.0,4.6}, {5.8,2.6,4.0}, {5.0,2.3,3.3}, {5.6,2.7,4.2}, {5.7,3.0,4.2}, {5.7,2.9,4.2}, {6.2,2.9,4.3}, {5.1,2.5,3.0}, {5.7,2.8,4.1}}
        train_labels = {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1}
        test_data = {{5.1,3.5,1.4}, {4.9,3.0,1.4}, {4.7,3.2,1.3}, {4.6,3.1,1.5}, {5.0,3.6,1.4}, {5.4,3.9,1.7}, {4.6,3.4,1.4}, {5.7,3.0,4.2}, {5.7,2.9,4.2}, {6.2,2.9,4.3}, {5.1,2.5,3.0}, {5.7,2.8,4.1}}
        test_labels = {1,1,1,1,1,1,1,-1,-1,-1,-1,-1}
        test_n = length(test_labels)
        #new instance of annsp
        k = annsp
        #train
        k.__init__(train_data, train_labels)
        #test
        #using linspace right now, because it seems that the recent push to julia repo has
        #screwed up the range method somehow.
        for i in linspace(1,test_n, test_n)
            print(k.classify(test_data[int(i)]), test_labels[int(i)], '\n')
        end
    end

#call
test_annsp()

ext: jl


text:
#-----------------------------------------------------------------------------# common
function mergestats(partition)
    parts = partition.parts
    o = copy(parts[1][2])
    for p in parts[2:end]
        merge!(o, p[2])
    end
    o
end

#-----------------------------------------------------------------------# Partition
"""
    Partition(stat)
    Partition(stat, nparts)

Split a data stream into `nparts` (default 100) where each part is summarized by `stat`.

# Example

    o = Partition(Extrema())
    fit!(o, cumsum(randn(10^5)))

    using Plots
    plot(o)
"""
mutable struct Partition{T, O <: OnlineStat{T}} <: OnlineStat{T}
    parts::Vector{Pair{Tuple{Int,Int}, O}} # (a,b) => stat
    b::Int
    init::O
    n::Int
end
Partition(o::OnlineStat, b::Int=100) = Partition([(1,1) => copy(o)], b, o, 0)

function _fit!(o::Partition, y)
    n = o.n += 1
    parts = o.parts
    (a, b), stat = parts[end]
    a ≤ n ≤ b ?
        _fit!(stat, y) :
        push!(parts, (n, n + nobs(stat) - 1) => fit!(copy(o.init), y))
    length(parts) > o.b && nobs(parts[end][2]) == (b - a + 1) && merge_next!(parts)
end

function merge_next!(parts)
    n = nobs(parts[1][2])
    i = 1
    for (j, p) in enumerate(parts)
        nobs(p[2]) < n && (i = j; break;)
    end
    a, b = parts[i], parts[i+1]
    parts[i] = (min(a[1][1], b[1][1]), max(a[1][2], b[1][2])) => merge!(a[2], b[2])
    deleteat!(parts, i + 1)
end

# Assumes `a` goes before `b`
function _merge!(a::Partition, b::Partition)
    n = nobs(a)
    a.n += b.n
    for p in b.parts
        push!(a.parts, (p[1][1] + n, p[1][2] + n) => copy(p[2]))
    end
    while length(a.parts) > a.b
        merge_next!(a.parts)
    end
    a
end

#-----------------------------------------------------------------------# IndexedPartition
"""
    IndexedPartition(T, stat, b=100)

Summarize data with `stat` over a partition of size `b` where the data is indexed by a
variable of type `T`.

# Example

    x, y = randn(10^5), randn(10^6)
    o = IndexedPartition(Float64, KHist(10))
    fit!(o, zip(x, y))

    using Plots
    plot(o)
"""
mutable struct IndexedPartition{I, T, O <: OnlineStat{T}, P <: Pair{<:Tuple, O}} <: OnlineStat{TwoThings{I,T}}
    parts::Vector{P}
    b::Int
    init::O
    method::Symbol
    n::Int
end
function IndexedPartition(I::Type, o::O, b::Int=100; method=:weighted_nearest) where {T, O<:OnlineStat{T}}
    IndexedPartition{I,T, O,Pair{Tuple{I,I}, O}}(Pair{Tuple{I,I}, O}[], b, o, method, 0)
end

function _fit!(o::IndexedPartition{I}, xy) where {I}
    x, y = xy
    n = o.n += 1
    isempty(o.parts) && push!(o.parts, (I(x), I(x)) => copy(o.init))
    i = findfirst(p -> p[1][1] ≤ x ≤ p[1][2], o.parts)
    if isnothing(i)
        push!(o.parts, (I(x), I(x)) => fit!(copy(o.init), y))
    else
        _fit!(o.parts[i][2], y)
    end
    length(o.parts) > o.b && indexed_merge_next!(sort!(o.parts, by = x -> x[1][1]), o.method)
end

function indexed_merge_next!(parts, method)
    if method === :weighted_nearest
        diffs = map(neighbors(parts)) do (a, b)
            (b[1][1] - a[1][2]) * round(Int, middle(nobs(a[2]), nobs(b[2])))
        end
        _, i = findmin(diffs)
        parts[i] = (parts[i][1][1], parts[i+1][1][2]) => merge!(parts[i][2], parts[i + 1][2])
        deleteat!(parts, i + 1)
    else
        error("method not recognized")
    end
end

# function _merge!(o::IndexedPartition, o2::IndexedPartition)
#     # If there's any overlap, merge
#     for p2 in o2.parts
#         pushpart = true
#         for p in o.parts
#             if (p[1][1] ≤ p2[1][1] ≤ p[1][2]) || (p[1][1] ≤ p2[1][2] ≤ p[1][2])
#                 pushpart = false
#                 merge!(p, p2)
#                 break
#             end
#         end
#         pushpart && push!(o.parts, p2)
#     end
#     # merge parts that overlap
#     for i in reverse(2:length(o.parts))
#         p1, p2 = o.parts[i-1], o.parts[i]
#         if p1.domain.first > p2.domain.last
#             # info("hey I deleted something at $i")
#             merge!(p1, p2)
#             deleteat!(o.parts, i)
#         end
#     end
#     # merge until there's b left
#     sort!(o.parts)
#     while length(o.parts) > o.b
#         indexed_merge_next!(o.parts, o.method)
#     end
#     o
# end

#-----------------------------------------------------------------------------# KIndexedPartition
"""
    KIndexedPartition(T, stat_init, k=100)

Similar to [`IndexedPartition`](@ref), but indexes the first variable by centroids (like [`KHist`](@ref))
rather than intervals.

- Note: `stat_init` must be a function e.g. `() -> Mean()`

# Example
    using Plots

    o = KIndexedPartition(Float64, () -> KHist(10))

    fit!(o, zip(randn(10^6), randn(10^6)))

    plot(o)
"""
mutable struct KIndexedPartition{I,T,O<:OnlineStat{T},F} <: OnlineStat{TwoThings{I,T}}
    parts::Vector{Pair{I,O}}
    k::Int
    init::F
    n::Int
end
function KIndexedPartition(I::Type, init::Base.Callable, k::Int=100)
    o = init()
    T, O = OnlineStatsBase.input(o), typeof(o)
    KIndexedPartition{I,T,O,typeof(init)}(Pair{I,O}[], k, init, 0)
end

function _fit!(o::KIndexedPartition, xy)
    x, y = xy
    parts = o.parts
    newpart = x => fit!(o.init(), y)
    insert!(parts, searchsortedfirst(parts, newpart; by=first), newpart)
    if length(parts) > o.k
        mindiff = Inf
        i = 0
        for (j, (a,b)) in enumerate(neighbors(parts))
            d = first(b) - first(a)
            if d < mindiff && 1 < j < (length(parts) - 1)
                mindiff = d
                i = j
            end
        end
        a = parts[i]
        b = parts[i + 1]
        n = nobs(last(a)) + nobs(last(b))
        parts[i] = smooth(first(a), first(b), nobs(last(b)) / n) => merge!(a[2], b[2])
        deleteat!(parts, i + 1)
    end
end


text:
function to_rpr_object(context, matsys, scene, plot::Makie.Volume)
    volume = plot.volume[]
    xyz = plot.x[], plot.y[], plot.z[]
    mini_maxi = extrema.(xyz)
    mini = first.(mini_maxi)
    maxi = last.(mini_maxi)

    vol_cube = RadeonProRender.Shape(context, Rect3f(mini, maxi .- mini))
    color_lookup = to_colormap(plot.colormap[])
    density_lookup = [Vec3f(plot.absorption[])]

    mini, maxi = extrema(volume)
    grid = RPR.VoxelGrid(context, (volume .- mini) ./ (maxi - mini))
    rpr_vol = RPR.HeteroVolume(context)

    RPR.set_albedo_grid!(rpr_vol, grid)
    RPR.set_albedo_lookup!(rpr_vol, color_lookup)

    RPR.set_density_grid!(rpr_vol, grid)
    RPR.set_density_lookup!(rpr_vol, density_lookup)

    mat = RPR.TransparentMaterial(matsys)
    mat.color = Vec4f(1)

    set!(vol_cube, rpr_vol)
    set!(vol_cube, mat)

    return [vol_cube, rpr_vol]
end


text:
##################################################
## evaluate binary operations of symbolic objects
## XXX -- this may prove too narryw with the use of ::Sym

+(x::SymbolicObject, y::SymbolicObject) = pycall(sympy.Add, Sym, x, y)
*(x::SymbolicObject, y::SymbolicObject) = pycall(sympy.Mul, Sym, x, y)
-(x::SymbolicObject, y::SymbolicObject) = x + (-y)
-(x::SymbolicObject)                    =  (-1) * x
/(x::SymbolicObject, y::SymbolicObject) = x * inv(y)
^(x::SymbolicObject, y::SymbolicObject) = pycall(sympy.Pow, Sym, x, y)::Sym
^(x::SymbolicObject, y::Rational) = x^convert(Sym,y)
#^(x::SymbolicObject, y::Integer) = x^convert(Sym,y) # no Union{Integer, Rational}, as that has ambiguity
//(x::SymbolicObject, y::Int) = x / Sym(y)
//(x::SymbolicObject, y::Rational) = x / Sym(y)
//(x::SymbolicObject, y::SymbolicObject) = x / y

\(x::SymbolicObject, y::SymbolicObject) = (y'/x')' # ?

#inv(x::Sym) = x\one(x)
#inv(x::Sym) = x^(-1)
inv(x::Sym) = pycall(sympy.Pow, Sym, x, -1)


text:
time_modeling(model::ModelAll, srcGeometry::Geometry, srcData, recGeometry::Geometry, recData, perturbation, srcnum::UnitRange{Int64}, op::Char, mode::Int64) =
	time_modeling(model, srcGeometry, srcData, recGeometry, recData, perturbation, srcnum, op, mode, Options())

function time_modeling(model::ModelAll, srcGeometry::Geometry, srcData, recGeometry::Geometry, recData, perturbation, srcnum::UnitRange{Int64}, op::Char, mode::Int64, options)
# time_modeling function for multiple sources. Depending on the operator and mode, this function distributes the sources
# and if applicable the input data amongst the available workers.

	p = default_worker_pool()
	time_modeling = remote(TimeModeling.time_modeling)
	# time_modeling = wrap_retry(time_modeling, options.retry_n)

	numSources = length(srcnum)
	results = Array{Any}(numSources)

	# Process shots from source channel asynchronously
	@sync begin
		for j=1:numSources

			# local geometry for current position
			srcGeometryLocal = subsample(srcGeometry,j)
			recGeometryLocal = subsample(recGeometry,j)

			# Parallelization
			if op=='F' && mode==1
				srcDataLocal = Array{Any}(1)
				srcDataLocal[1] = srcData[j]
				@async results[j] = time_modeling(model, srcGeometryLocal, srcDataLocal, recGeometryLocal, [], [], j, op, mode, options)
			elseif op=='F' && mode==-1
				recDataLocal = Array{Any}(1)
				recDataLocal[1] = recData[j]
				@async results[j] = time_modeling(model, srcGeometryLocal, [], recGeometryLocal, recDataLocal, [], j, op, mode, options)
			elseif op=='J' && mode==1
				srcDataLocal = Array{Any}(1)
				srcDataLocal[1] = srcData[j]
				@async results[j] = time_modeling(model, srcGeometryLocal, srcDataLocal, recGeometryLocal, [], perturbation, j, op, mode, options)
			elseif op=='J' && mode==-1
				srcDataLocal = Array{Any}(1)
				srcDataLocal[1] = srcData[j]
				recDataLocal = Array{Any}(1)
				recDataLocal[1] = recData[j]
				@async results[j] = time_modeling(model, srcGeometryLocal, srcDataLocal, recGeometryLocal, recDataLocal, [], j, op, mode, options)
			end
		end
	end

	if op=='F' || (op=='J' && mode==1)
		argout1 = results[1]
		for j=2:numSources
			argout1 = [argout1; results[j]]
		end
	elseif op=='J' && mode==-1
		argout1 = results[1]
		for j=2:numSources
			argout1 += results[j]
		end
	else
		error("operation no defined")
	end
	return argout1
end

max_stars_repo_name: SINBADconsortium/opesciSLIM

max_issues_repo_name: SINBADconsortium/opesciSLIM

max_forks_repo_name: SINBADconsortium/opesciSLIM


text:
""" Convert wiring diagrams to and from syntactic expressions.

Wiring diagrams are a graphical syntax for morphisms in a monoidal category. As
mathematical objects, they are intermediate between the morphisms themselves and
expressions in the textual syntax: a single morphism may correspond to many
wiring diagrams, and a single wiring diagram may correspond to many syntactic
expressions.

It is trivial to convert a morphism expression to a wiring diagram, but not to
go the other way around.
"""
module WiringDiagramExpressions
export to_ob_expr, to_hom_expr, to_wiring_diagram, to_undirected_wiring_diagram

using ...Syntax, ...Theories, ...Permutations, ...CategoricalAlgebra
using ...Syntax: syntax_module
using ...Graphs, ..DirectedWiringDiagrams, ..UndirectedWiringDiagrams,
  ..MonoidalDirectedWiringDiagrams, ..MonoidalUndirectedWiringDiagrams
using ..WiringDiagramAlgorithms: crossing_minimization_by_sort
using ..DirectedWiringDiagrams: WiringDiagramGraph

# Expression -> Diagram
#######################

""" Convert a morphism expression into a wiring diagram.
"""
function to_wiring_diagram(expr::GATExpr, args...)
  T = syntax_module(expr).theory()
  to_wiring_diagram(T, expr, args...)
end
function to_wiring_diagram(T::Type, expr::GATExpr)
  to_wiring_diagram(T, expr, first, first)
end
function to_wiring_diagram(T::Type, expr::GATExpr, ob_map, hom_map)
  to_wd = expr -> to_wiring_diagram(T, expr, ob_map, hom_map)
  functor((Ports, WiringDiagram), expr; terms=Dict(
    :Ob => expr -> Ports{T}([ob_map(expr)]),
    :Hom => expr -> begin
      A, B = to_wd(dom(expr)), to_wd(codom(expr))
      singleton_diagram(T, Box(hom_map(expr), A, B))
    end
  ))
end

""" Convert a morphism expression into an undirected wiring diagram.

Returns a `HomUWD`, or a UWD with outer ports partitioned domain and codomain.
"""
function to_undirected_wiring_diagram(expr::GATExpr)
  to_undirected_wiring_diagram(expr, Symbol, Symbol, first, first)
end
function to_undirected_wiring_diagram(expr::GATExpr, Ob::Type, Hom::Type,
                                      ob_map, hom_map)
  to_uwd = expr -> to_undirected_wiring_diagram(expr, Ob, Hom, ob_map, hom_map)
  HDOb, HDHom = HypergraphDiagramOb{Ob,Hom}, HypergraphDiagramHom{Ob,Hom}
  functor((HDOb, HDHom), expr; terms=Dict(
    :Ob => expr -> HDOb([ob_map(expr)]),
    :Hom => expr -> begin
      A, B = to_uwd(dom(expr)), to_uwd(codom(expr))
      HDHom(singleton_diagram(A⊗B, name=hom_map(expr)),
            dom=[trues(length(A)); falses(length(B))])
    end
  ))
end

# Diagram -> Expression
#######################

""" Convert a wiring diagram into a morphism expression.
"""
function to_hom_expr(Ob::Type, Hom::Type, d::WiringDiagram)
  box_to_expr(v::Int) = to_hom_expr(Ob, Hom, box(d,v))

  # Initial reduction: Add junction nodes to ensure any wiring layer between two
  # boxes is a permutation.
  d = add_junctions(d)

  # Dispatch special case: no boxes.
  if nboxes(d) == 0
    return hom_expr_between(Ob, d, input_id(d), output_id(d))
  end

  # Main loop: reduce until no further reduction is possible.
  n, changed = nboxes(d), true
  while changed
    # Transtive reduction: can only increase box count.
    d = transitive_reduction!(Ob, d)
    n, changed = nboxes(d), false
    while true
      # Alternating series/parallel reduction: can only decrease box count.
      d = series_reduction(Ob, Hom, d)
      d = parallel_reduction(Ob, Hom, d)
      d = input_parallel_reduction(Ob, Hom, d)
      d = output_parallel_reduction(Ob, Hom, d)
      nboxes(d) == n && break
      n, changed = nboxes(d), true
    end
  end

  # Termination: process all remaining boxes (always at least one, and exactly
  # one if there is no creation or deletion).
  if nboxes(d) > 1
    @assert all((wire.source.box == input_id(d) ||
                 wire.target.box == output_id(d)) for wire in wires(d))
    d = encapsulate_parallel(Ob, Hom, d,
      [(box_ids(d), [input_id(d)], [output_id(d)])])
  end
  v = first(box_ids(d))
  foldl(compose_simplify_id, [
    hom_expr_between(Ob, d, input_id(d), v),
    box_to_expr(v),
    hom_expr_between(Ob, d, v, output_id(d)),
  ])
end
function to_hom_expr(Syntax::Module, d::WiringDiagram)
  to_hom_expr(Syntax.Ob, Syntax.Hom, d)
end

to_hom_expr(::Type{Ob}, ::Type{Hom}, box::Box{<:Hom}) where {Ob,Hom} = box.value

function to_hom_expr(Ob::Type, Hom::Type, box::Box)
  dom = otimes(to_ob_exprs(Ob, input_ports(box)))
  codom = otimes(to_ob_exprs(Ob, output_ports(box)))
  Hom(box.value, dom, codom)
end
function to_hom_expr(Ob::Type, Hom::Type, op::BoxOp)
  invoke_term(parentmodule(Hom), head(op), to_hom_expr(Ob, Hom, op.box))
end
function to_hom_expr(Ob::Type, Hom::Type, junction::Junction)
  junction_to_expr(Ob, junction)
end

""" All possible parallel reductions of a wiring diagram.
"""
function parallel_reduction(Ob::Type, Hom::Type, d::WiringDiagram)
  g, outer_vs = graph_with_outer_ids(d)
  parallel = find_parallel(g, skip=outer_vs)
  parallel = Dict([((box_id(g, src), box_id(g, tgt)), box_id(g,vs)) for ((src, tgt), vs) in parallel])
  encapsulate_parallel(Ob, Hom, d, [
    (vs, [src], [tgt]) for ((src, tgt), vs) in parallel
  ])
end

""" Input-sided parallel reduction of a wiring diagram.

Because these reductions are not necessarily unique, only one is performed,
the first one in topological sort order.
"""
function input_parallel_reduction(Ob::Type, Hom::Type, d::WiringDiagram)
  g, outer_vs = graph_with_outer_ids(d)
  parallel = find_one_sided_parallel(g, input=true, skip=outer_vs)
  if isempty(parallel)
    return d
  end
  vs = collect(keys(parallel))
  src = vs[first(topological_sort(induced_subgraph(g, vs)))]
  encapsulate_parallel(Ob, Hom, d, [ (box_id(g, parallel[src]), [box_id(g, src)], Int[]) ])
end

""" Output-sided parallel reduction of a wiring diagram.

Because these reductions are not necessarily unique, only one is performed,
the last one in topological sort order.
"""
function output_parallel_reduction(Ob::Type, Hom::Type, d::WiringDiagram)
  g, outer_vs = graph_with_outer_ids(d)
  parallel = find_one_sided_parallel(g, input=false, skip=outer_vs)
  if isempty(parallel)
    return d
  end
  vs = collect(keys(parallel))
  tgt = vs[last(topological_sort(induced_subgraph(g, vs)))]
  encapsulate_parallel(Ob, Hom, d, [ (box_id(g, parallel[tgt]), Int[], [box_id(g, tgt)]) ])
end

function encapsulate_parallel(Ob::Type, Hom::Type, d::WiringDiagram, parallel)
  parallel = Vector{Int}[
    vs[crossing_minimization_by_sort(d, vs, sources=sources, targets=targets)]
    for (vs, sources, targets) in parallel
  ]
  products = map(parallel) do vs
    exprs = Hom[ to_hom_expr(Ob, Hom, box(d,v)) for v in vs ]
    otimes(exprs)
  end
  encapsulate(d, parallel, discard_boxes=true, values=products)
end

""" All possible series reductions of a wiring diagram.
"""
function series_reduction(Ob::Type, Hom::Type, d::WiringDiagram)
  box_to_expr(v::Int) = to_hom_expr(Ob, Hom, box(d,v))

  g, (input_v, output_v) = graph_with_outer_ids(d)
  series = find_series(g, source=input_v, sink=output_v)
  composites = map(series) do vs
    exprs = Hom[ box_to_expr(vs[1]) ]
    for i in 2:length(vs)
      layer_expr = hom_expr_between(Ob, d, vs[i-1], vs[i])
      append!(exprs, [layer_expr, box_to_expr(vs[i])])
    end
    foldl(compose_simplify_id, exprs)
  end
  encapsulate(d, series, discard_boxes=true, values=composites)
end

""" All possible transitive reductions of a wiring diagram.
"""
function transitive_reduction!(Ob::Type, d::WiringDiagram)
  # Compute transitive reduction of underlying graph.
  # First add extra edges for the "invisible wires" corresponding to monoidal
  # units, since transitive reduction can be needed even in this case.
  g, (input_v, output_v) = graph_with_outer_ids(d)
  reduced = copy(g)
  for v in box_ids(d)
    if isempty(input_ports(d, v))
      add_edge!(reduced, input_v, v)
    end
    if isempty(output_ports(d, v))
      add_edge!(reduced, v, output_v)
    end
  end
  reduced = Graphs.transitive_reduction!(reduced)

  # Add junction node for each wire removed by transitive reduction.
  for edge in collect(edges(g))
    s, t = src(g, edge), tgt(g, edge)
    if !has_edge(reduced, s, t)
      for wire in wires(d, box_id(g,s), box_id(g,t))
        value = port_value(d, wire.source) # =?= port_value(d, wire.target)
        v = add_box!(d, Junction(value, 1, 1))
        add_wire!(d, Wire(wire.source => Port(v,InputPort,1)))
        add_wire!(d, Wire(Port(v,OutputPort,1) => wire.target))
        rem_wire!(d, wire)
      end
    end
  end
  return d
end

""" Morphism expression for wires between two boxes.

Assumes that the wires form a permutation morphism.
"""
function hom_expr_between(Ob::Type, diagram::WiringDiagram, v1::Int, v2::Int)
  inputs, outputs = output_ports(diagram, v1), input_ports(diagram, v2)
  σ = zeros(Int, length(inputs))
  for wire in wires(diagram, v1, v2)
    src, tgt = wire.source.port, wire.target.port
    @assert σ[src] == 0 "Wire mapping must be single-valued"
    σ[src] = tgt
  end
  @assert 0 ∉ σ "Wire mapping must be total"
  @assert allunique(σ) "Wiring mapping must be a permutation"

  inputs, outputs = to_ob_exprs(Ob, inputs), to_ob_exprs(Ob, outputs)
  @assert inputs == outputs[σ]
  permutation_to_expr(σ, inputs)
end

""" Convert a port value into an object expression.
"""
to_ob_expr(Syntax::Module, x) = to_ob_expr(Syntax.Ob, x)
to_ob_expr(::Type{Ob}, ob::Ob) where Ob = ob
to_ob_expr(Ob::Type, value) = Ob(Ob, value)

to_ob_expr(Ob::Type, ports::Ports) = otimes(to_ob_exprs(Ob, ports))
to_ob_expr(Ob::Type, op::PortOp) =
  invoke_term(parentmodule(Ob), head(op), to_ob_expr(Ob, op.value))

to_ob_exprs(Ob::Type, values) = Ob[ to_ob_expr(Ob, value) for value in values ]

""" Compose morphism expressions, eliminating identities.
"""
function compose_simplify_id(f::GATExpr, g::GATExpr)
  if head(f) == :id; g
  elseif head(g) == :id; f
  else compose(f,g) end
end

""" Convert a junction node to a morphism expression.
"""
function junction_to_expr(Ob::Type, junction::Junction)
  ob = to_ob_expr(Ob, junction.value)
  nin, nout = length(input_ports(junction)), length(output_ports(junction))
  if (nin == 2 && nout == 0) dcounit(ob)
  elseif (nin == 0 && nout == 2) dunit(ob)
  else compose_simplify_id(mmerge_foldl(ob, nin), mcopy_foldl(ob, nout)) end
end

# FIXME: These functions belong elsewhere, probably in the standard library.
function mcopy_foldl(A, n::Int)
  @assert n >= 0
  if (n > 2) compose(mcopy(A), otimes(mcopy_foldl(A, n-1), id(A)))
  elseif (n == 2) mcopy(A)
  elseif (n == 1) id(A)
  else delete(A) end
end
function mmerge_foldl(A, n::Int)
  @assert n >= 0
  if (n > 2) compose(otimes(mmerge_foldl(A, n-1), id(A)), mmerge(A))
  elseif (n == 2) mmerge(A)
  elseif (n == 1) id(A)
  else create(A) end
end

function graph_with_outer_ids(d::WiringDiagram)
  g = graph(d)
  outer = (incident(g, input_id(d), :box), incident(g, output_id(d), :box))
  (g, outer)
end

box_id(g::WiringDiagramGraph, v) = subpart(g, v, :box)

# Graph operations
##################

""" Find parallel compositions in a directed graph.

This two-sided notion of parallel composition is standard in the literature on
series-parallel digraphs.
"""
function find_parallel(g::AbstractGraph; skip=())
  parallel = Dict{Pair{Int,Int},Vector{Int}}()
  for v in 1:nv(g)
    if v in skip; continue end
    v_in, v_out = unique(inneighbors(g,v)), unique(outneighbors(g,v))
    if length(v_in) == 1 && length(v_out) == 1
      u, w = first(v_in), first(v_out)
      push!(get!(parallel, u => w, Int[]), v)
    end
  end
  filter(pair -> length(last(pair)) > 1, parallel)
end

""" Find one-sided parallel compositions in a directed graph.

Finds either input-sided or output-sided compositions. This weaker notion of
parallel composition seems to be nonstandard.
"""
function find_one_sided_parallel(g::AbstractGraph; input::Bool=true, skip=())
  parallel = Dict{Int,Vector{Int}}()
  for v in 1:nv(g)
    if v in skip; continue end
    vs = unique(input ? inneighbors(g,v) : outneighbors(g,v))
    if length(vs) == 1
      push!(get!(parallel, first(vs), Int[]), v)
    end
  end
  filter(pair -> length(last(pair)) > 1, parallel)
end

""" Find series compositions in a directed graph.
"""
function find_series(g::AbstractGraph; source=nothing, sink=nothing)::Vector{Vector{Int}}
  series_graph = Graph(nv(g))
  for edge in edges(g)
    s, t = src(g, edge), tgt(g, edge)
    if (length(unique(outneighbors(g,s))) == 1 &&
        length(unique(inneighbors(g,t))) == 1 &&
        s != source && t != sink)
      add_edge!(series_graph, s, t)
    end
  end
  series = Vector{Int}[]
  for component in connected_components(series_graph)
    if length(component) > 1
      sub = induced_subgraph(series_graph, component)
      push!(series, Int[ component[v] for v in topological_sort(sub) ])
    end
  end
  series
end

end


text: Dict{String,Any}("a" => Any[Dict{String,Any}("b" => Any[Dict{String,Any}("c" => Dict{String,Any}("d" => Dict{String,Any}("value" => "val0","type" => "string"))), Dict{String,Any}("c" => Dict{String,Any}("d" => Dict{String,Any}("value" => "val1","type" => "string")))])])

max_stars_repo_name: vanillajonathan/julia

max_issues_repo_name: vanillajonathan/julia

max_forks_repo_name: vanillajonathan/julia


text:
"""
    prior_probs(model::Model, parm_grid)

Computes a grid of prior probabilities over model parameters.

- `model::Model`: a model object
- `parm_grid`: a grid of parameter values
"""
function prior_probs(model::Model, parm_grid)
    return prior_probs(model.prior, parm_grid)
end

"""
    prior_probs(prior, parm_grid)

Computes a grid of prior probabilities over model parameters.

- `prior`: a vector of `Distribution` objects, one for each parameter
- `parm_grid`: a grid of parameter values
"""
function prior_probs(prior, parm_grid)
    dens = [mapreduce((θ,d)->pdf(d, θ), *, g, prior) for g in parm_grid]
    return dens/sum(dens)
end

"""
    prior_probs(prior, parm_grid)

Computes a grid of uniform prior probabilities over model parameters.

- `prior::Nothing`: a vector of `Distribution` objects, one for each parameter
- `parm_grid`: a grid of parameter values
"""
function prior_probs(prior::Nothing, parm_grid)
    return fill(1/length(parm_grid), size(parm_grid))
end

"""
    loglikelihood(model::Model, design_grid, parm_grid, data_grid, model_type, model_state)

Computes a three dimensional grid of log likelihoods where the first dimension is the model parameters, the second
dimension is the design parameters, and the third dimension is the data values.

- `model::Model`: a model object containing a log likelihood function and prior distributions
- `parm_grid`: a grid of model parameters
- `design_grid`: a grid of design parameters
- `data_grid`: a grid of data values
- `model_type`: model type can be `Static` or `Dynamic`
- `model_state`: model state variables that are updated for `Dynamic` models
"""
function loglikelihood(model::Model, design_grid, parm_grid, data_grid, model_type, model_state)
    return loglikelihood(model.loglike, design_grid, parm_grid, data_grid)
end

"""
    loglikelihood(model::Model, design_grid, parm_grid, data_grid, model_type, model_state)

Computes a three dimensional grid of log likelihoods where the first dimension is the model parameters, the second
dimension is the design parameters, and the third dimension is the data values.

- `model::Model`: a model object containing a log likelihood function and prior distributions
- `parm_grid`: a grid of model parameters
- `design_grid`: a grid of design parameters
- `data_grid`: a grid of data values
- `model_type`: model type can be `Static` or `Dynamic`
- `model_state`: model state variables that are updated for `Dynamic` models
"""
function loglikelihood(model::Model, design_grid, parm_grid, data_grid, model_type::Dyn, model_state)
    return loglikelihood(model.loglike, design_grid, parm_grid, data_grid, model_state)
end

"""
    loglikelihood(model::Model, design_grid, parm_grid, data_grid, model_type, model_state)

Computes a three dimensional grid of log likelihoods where the first dimension is the model parameters, the second
dimension is the design parameters, and the third dimension is the data values.

- `loglike`: a model object containing a log likelihood function and prior distributions
- `parm_grid`: a grid of model parameters
- `design_grid`: a grid of design parameters
- `data_grid`: a grid of data values
- `model_type`: model type can be `Static` or `Dynamic`
- `model_state`: model state variables that are updated for `Dynamic` models
"""
function loglikelihood(loglike, design_grid, parm_grid, data_grid)
    LLs = zeros(length(parm_grid), length(design_grid), length(data_grid))
    for (d, data) in enumerate(data_grid)
        for (k,design) in enumerate(design_grid)
            for (p,parms) in enumerate(parm_grid)
                LLs[p,k,d] = loglike(parms..., design..., data...)
            end
        end
    end
    return LLs
end

"""
    loglikelihood(model::Model, design_grid, parm_grid, data_grid, model_type, model_state)

Computes a three dimensional grid of log likelihoods for a dynamic model where the first dimension is the model parameters,
the second dimension is the design parameters, and the third dimension is the data values.

- `loglike`: a model object containing a log likelihood function and prior distributions
- `parm_grid`: a grid of model parameters
- `design_grid`: a grid of design parameters
- `data_grid`: a grid of data values
- `model_type`: model type can be `Static` or `Dynamic`
- `model_state`: model state variables that are updated for `Dynamic` models
"""
function loglikelihood(loglike, design_grid, parm_grid, data_grid, model_state)
    LLs = zeros(length(parm_grid), length(design_grid), length(data_grid))
    i = 0
    for (d, data) in enumerate(data_grid)
        for (k,design) in enumerate(design_grid)
            for (p,parms) in enumerate(parm_grid)
                i += 1
                state = deepcopy(model_state[i])
                LLs[p,k,d] = loglike(parms..., design..., data..., state)
            end
        end
    end
    return LLs
end

"""
    loglikelihood(optimizer::Optimizer)

Computes a three dimensional grid of log likelihoods for a dynamic model where the first dimension is the model parameters,
the second dimension is the design parameters, and the third dimension is the data values.

- `optimizer`: an optimizer object

"""
function loglikelihood!(optimizer)
    @unpack model, log_like, design_grid, parm_grid, data_grid, model_state = optimizer
    return loglikelihood!(model.loglike, log_like, design_grid, parm_grid, data_grid, model_state)
end

"""
    loglikelihood(optimizer::Optimizer)

Computes a three dimensional grid of log likelihoods where the first dimension is the model parameters, the second
dimension is the design parameters, and the third dimension is the data values.

- `loglike`: a model object containing a log likelihood function and prior distributions
- `parm_grid`: a grid of model parameters
- `design_grid`: a grid of design parameters
- `data_grid`: a grid of data values
- `model_type`: model type can be `Static` or `Dynamic`
- `model_state`: model state variables that are updated for `Dynamic` models
"""
function loglikelihood!(loglike, log_like, design_grid, parm_grid, data_grid, model_state)
    i = 0
    for (d, data) in enumerate(data_grid)
        for (k,design) in enumerate(design_grid)
            for (p,parms) in enumerate(parm_grid)
                i += 1
                state = deepcopy(model_state[i])
                log_like[p,k,d] = loglike(parms..., design..., data..., state)
            end
        end
    end
    return nothing
end

"""
    marginal_log_like!(optimizer)

Marginalizes the log likelihoods over model parameters, resulting in a three dimensional array where the first dimension is length 1,
the second dimension is the design parameters, and the third dimension is the data values.

- `optimizer`: an optimizer object
"""
function marginal_log_like!(optimizer)
    @unpack marg_log_like,log_like,log_post = optimizer
    marg_log_like .= marginal_log_like(log_post, log_like)
end

"""
    marginal_log_like(optimizer)

Marginalizes the log likelihoods over model parameters, resulting in a three dimensional array where the first dimension is length 1,
the second dimension is the design parameters, and the third dimension is the data values.

- `log_post`: a vector of posterior log likelihood values for the model parameters
- `log_like`: a three dimensional vector of log likelihood values where the first dimension corresponds to the model parameters,
the second dimension corresponds to the design parameters, and the third dimension corresponds to the data values
"""
function marginal_log_like(log_post, log_like)
    return logsumexp(log_post .+ log_like, dims=1)
end

# """
#     marginal_posterior(optimizer)

# Computes a three dimensional grid of log likelihoods where the first dimension is the model parameters, the second
# dimension is the design parameters, and the third dimension is the data values.

# - `optimizer`: an optimizer object
# """
# function marginal_posterior(optimizer)
#     @unpack posteriors = optimizer
#     return map(d->sum(posterior, dims=d), ndims(posterior):-1:1)
# end

function compute_entropy(log_like)
    return -1*sum(exp.(log_like) .* log_like, dims=3)[:,:]
end

function conditional_entropy(entropy, post)
    return entropy'*post
end

function conditional_entropy!(optimizer)
    @unpack cond_entropy, entropy,log_post = optimizer
    post = exp.(log_post)
    cond_entropy .= conditional_entropy(entropy, post)
end

function marginal_entropy!(optimizer::Optimizer)
    @unpack marg_entropy,marg_log_like = optimizer
    marg_entropy .= marginal_entropy(marg_log_like)
end

function marginal_entropy(marg_log_like)
    return -sum(exp.(marg_log_like).*marg_log_like, dims=3)[:]
end

"""
    mutual_information(optimizer)

Computes a vector of the mutual information values where each element corresponds to a design parameter.

- `marg_entropy`: a vector of marginal entropy values where each element corresponds to a design parameter
- `cond_entropy`: a vector of conditional entropy values where each element corresponds to a design parameter
"""
function mutual_information(marg_entropy, cond_entropy)
    return marg_entropy .- cond_entropy
end

"""
    mutual_information!(optimizer)

Computes a vector of the mutual information values where each element corresponds to a design parameter.

- `optimizer`: an optimizer object
"""
function mutual_information!(optimizer)
    @unpack mutual_info,marg_entropy,cond_entropy = optimizer
    mutual_info .= mutual_information(marg_entropy, cond_entropy)
    return nothing
end

function get_best_design(optimizer)
    @unpack best_design,design_names = optimizer
    return NamedTuple{design_names}(best_design)
end

function find_best_design!(optimizer)
    @unpack design_names,mutual_info,design_grid = optimizer
    best_design = find_best_design(mutual_info, design_grid, design_names)
    optimizer.best_design = best_design
    return best_design
end

function find_best_design(mutual_info, design_grid, design_names)
    _,best = findmax(mutual_info)
    best_design = design_grid[best]
    return best_design
end

function find_best_design!(optimizer::Optimizer{A}) where {A<:Rand}
    @unpack design_grid,design_names = optimizer
    best_design = rand(design_grid)
    optimizer.best_design = best_design
    return best_design
end

function update_posterior!(optimizer, data)
    @unpack log_post,design_grid,data_grid,log_like,best_design = optimizer
    dn = find_index(design_grid, best_design)
    for datum in data
        da = find_index(data_grid, (datum,))
        log_post .+= log_like[:,dn,da]
    end
    log_post .-= logsumexp(log_post)
    return nothing
end

function update!(optimizer, data)
    update_posterior!(optimizer, data)
    marginal_log_like!(optimizer)
    marginal_entropy!(optimizer)
    conditional_entropy!(optimizer)
    mutual_information!(optimizer)
    best_design = find_best_design!(optimizer)
    return best_design
end

function update!(optimizer::Optimizer{A,MT}, data, args...; kwargs...) where {A,MT<:Dyn}
    update_posterior!(optimizer, data)
    update_states!(optimizer, data, args...; kwargs...)
    loglikelihood!(optimizer)
    marginal_log_like!(optimizer)
    marginal_entropy!(optimizer)
    conditional_entropy!(optimizer)
    mutual_information!(optimizer)
    best_design = find_best_design!(optimizer)
    return best_design
end

function update!(optimizer::Optimizer{A,MT}, data, args...; kwargs...) where {A<:Rand,MT<:Dyn}
    update_posterior!(optimizer, data)
    update_states!(optimizer, data, args...; kwargs...)
    loglikelihood!(optimizer)
    best_design = find_best_design!(optimizer)
    return best_design
end

function update!(optimizer::Optimizer{A}, data) where {A<:Rand}
    update_posterior!(optimizer, data)
    best_design = find_best_design!(optimizer)
    return best_design
end

function update_states!(optimizer, obs_data, args...; kwargs...)
    @unpack model_state, data_grid, design_grid, parm_grid, update_state! = optimizer
    i = 0
    for data in data_grid
        for design in design_grid
            for parms in parm_grid
                i += 1
                update_state!(model_state[i], parms, design, data, obs_data, args...; kwargs...)
            end
        end
    end
    return nothing
end

function to_grid(vals::NamedTuple)
    k = keys(vals)
    v = product(vals...) |> collect
    return k,v
end

function to_grid(vals)
    k = [Symbol(string("v",i)) for i in 1:length(vals[1])]
    return (k...,),vals
end

to_grid(vals::Tuple) = vals

function find_index(grid, val)
    i = 0
    for g in grid
        i += 1
        if g == val
            return i
        end
    end
    return i
end

function mean_post(optimizer)
    @unpack log_post,parm_grid = optimizer
    post = exp.(log_post)
    return mean_post(post, parm_grid)
end

function mean_post(post, parm_grid)
    return mapreduce((p,v)->p.*v, .+, post, parm_grid)
end

function std_post(post, parm_grid)
    mu = mean_post(post, parm_grid)
    return mapreduce((p,v)->p.*(v .- mu).^2, .+, post, parm_grid) .|> sqrt
end

function std_post(optimizer)
    @unpack log_post,parm_grid = optimizer
    post = exp.(log_post)
    return std_post(post, parm_grid)
end

function create_state(model_type::Dyn, T, dims, args...; kwargs...)
    state = fill(T(args...; kwargs...), dims)
    return state .= deepcopy.(state)
end

function create_state(model_type::Stat, T, dims, args...; kwargs...)
    return nothing
end

# function update(θ1, θ2, p, data)
#     #like = pdf.(Binomial.(n, x), k)
#     like = pdf.(Bernoulli.(x), d)
#     for d in data
#         θ1 .*= like
#         θ2 .*= like
#     end

#     # e1 = θ1'*like
#     # e2 = θ2'*like

#     # θ1 .= (θ1 .* like)/e1
#     # θ2 .= (θ2 .* like)/e2
#     e1 = sum(θ1)
#     e2 = sum(θ2)
#     θ1 ./= e1
#     θ2 ./= e2

#     post_odds = (p * e1) / ((1 - p) * e2)
#     p = post_odds/(1 + post_odds)
#     return p, θ1, θ2
# end

# # function update_model_posterior(ps, es)
# #     map(i->ps[i]./(sum(ps.*(es./es[i]))),1:length(ps))
# # end

# using Distributions
# p = .5
# x = .001:.001:.999
# θ1 = pdf.(Beta(10,1), x)
# θ1 /= sum(θ1)

# θ2 = pdf.(Beta(1,1), x)
# θ2 /= sum(θ2)

# p, θ1, θ2 = update(θ1, θ2, p, [1])


text:
@doc """
    abs_difference(edge)

Return the difference between two points in euclidean space, given an edge between those points.

# Example
```jldoctest
julia> abs_difference( (5,1) )
4

julia> abs_difference( ((2,2), (5,-5)) )
(3, 7)
```
"""
abs_difference(edge::Tuple{T,T}) where T<:Number = abs(edge[1] - edge[2])
# FIXME should really be L2 norm
abs_difference(edge::Tuple{Tup,Tup}) where {Nminusone,T,Tup<:Tuple{T,Vararg{T,Nminusone}}} = abs.(edge[1] .- edge[2])

"""
    abs_difference_periodic(edge, period)

Return the difference between two points in euclidean space as in abs_difference, but let the space wrap with period.

# Example
```jldoctest
julia> abs_difference_periodic( (5,1), 3 )
-1

julia> abs_difference_periodic( ((2,2), (5,-5)), (3,4) )
(0, -3)
```
"""
function abs_difference_periodic(edge::Tuple{T,T}, (period,)::Tuple{T}) where T<:Number
    diff = abs_difference(edge)
    if diff >= period / 2
        return period - diff
    else
        return diff
    end
end
function abs_difference_periodic(edge::Tuple{Tup,Tup}, periods::Tup) where {Nminusone,T,Tup<:Tuple{T,Vararg{T,Nminusone}}}
    diffs = abs_difference(edge)
    diffs = map(zip(diffs, periods)) do (diff, period)
        if diff > period / 2
            return period - diff
        else
            return diff
        end
    end
    return Tup(diffs)
end


text: Dict{String, Any}("~!@\$^&*()_+-`1234567890[]|/?><.,;:'" => Dict{String, Any}("value" => "1", "type" => "integer"))


text:
# This file is a part of BAT.jl, licensed under the MIT License (MIT).

bat_default(::typeof(bat_findmode), ::Val{:algorithm}, ::DensitySampleVector) = MaxDensitySampleSearch()

bat_default(::typeof(bat_findmode), ::Val{:algorithm}, ::AbstractDensity) = MaxDensityNelderMead()

bat_default(::typeof(bat_findmode), ::Val{:algorithm}, ::Distribution) = ModeAsDefined()

bat_default(::typeof(bat_findmode), ::Val{:algorithm}, ::DistLikeDensity) = ModeAsDefined()


text:
using OctreeBH
using StaticArrays
using .Threads #for parallel treewalk
using PyPlot
plt.style.use("dark_background")

const N = 3 #spatial dimension
const T = Float64

const Npart = 2000 #number of particles
const ANGLE = 0.6 #node opening criteria
const FAC_GRAV = 0.05 #timestep control
const boxsizes = @SVector(ones(N)) * Inf  #open B.C.
const hsml = zeros(Npart) #dummy

const R0 = 1.0 #scale radius
const Rmax = 5.0 #cutoff radius
const softening = 5*R0 / sqrt(Npart) #gravitational softening length
const M_tot = 1.0 #total mass
const m_gas = M_tot / Npart
const mass = ones(Npart) .* m_gas #particle mass
const G_const = 1.0 #in code units

const tend = 10. #end time
const dt_dump = 0.02 #time between snapshots

using Random

function set_spherical_cloud(Npart, R0, Rmax)
    Random.seed!(1114);
    X = Vector{SVector{3,T}}(undef,Npart)
    V = [zero(SVector{N,T}) for _ in 1:Npart]
    for i in 1:Npart
        x,y,z = 0.,0.,0.
        r = 0.
        while(true)
            m = rand();
            r = R0 * sqrt(m) / (1. - sqrt(m)) #importance sampling
            if(r<Rmax) break end
        end
        phi = ( 2 * rand() - 1 ) * pi;
        cos_theta = 2 * rand() - 1;

        x = r * sqrt( 1. - cos_theta^2 ) * cos(phi);
        y = r * sqrt( 1. - cos_theta^2 ) * sin(phi);
        z = r * cos_theta;

        X[i] = SVector{3,T}(x,y,z)
    end
    return X,V
end

function calc_grav_acc(X)
    part = [Data{N,T}(SVector(X[i]), i, hsml[i], mass[i]) for i in eachindex(X)]

    Xmin = @SVector [minimum(getindex.(X,i)) for i in 1:N]
    Xmax = @SVector [maximum(getindex.(X,i)) for i in 1:N]

    topnode_length = (Xmax - Xmin) * 1.01  #actual length of tree
    center = 0.5 * (Xmax + Xmin)

    #build the tree
    tree = buildtree(part, center, topnode_length);

    acc = Vector{SVector{N,T}}(undef,Npart)

    @threads for i in eachindex(X)
        ga = GravTreeGather{N,T}()
        gravity_treewalk!(ga,X[i],tree,ANGLE,softening,boxsizes)
        acc[i] = G_const .* ga.acc
    end
    return acc
end

function plot_particle_configuration(X,count_dump)
    clf()
    fig, ax = PyPlot.subplots(1, 1, figsize=(7,6))
    ax.plot(getindex.(X,1), getindex.(X,2), ".", ms=1, c="white")
    ax.axis(R0.*[-1,1,-1,1])
    if count_dump < 10
        num = "000"*string(count_dump)
    elseif count_dump < 100
        num = "00"*string(count_dump)
    elseif count_dump < 1000
        num = "0"*string(count_dump)
    else
        num = string(count_dump)
    end
    savefig("plot_"*num*".png")
end

norm(a::SVector{N,T}) where{N,T} = sqrt(sum(a.^2))

function run()

    #randomly distributing particles
    #X = [@SVector rand(N) for _ in 1:Npart]
    #V = [zero(SVector{N,T}) for _ in 1:Npart]
    X,V = set_spherical_cloud(Npart, R0, Rmax)

    acc = calc_grav_acc(X)

    count_dump::Int64 = 0
    t = 0.
    t_output = 0.
    println("Start the calculation...")
    @time while t < tend
        if t >= t_output
            @show t
            plot_particle_configuration(X,count_dump)
            t_output += dt_dump
            count_dump += 1
        end

        dt = minimum( @. FAC_GRAV * sqrt(softening / norm(acc) ) )
        dt = tend - t < dt ? tend - t : dt

        #use the kick-drift-kick integration (i.e. midpoint method)
        V = V + 0.5 * acc * dt #1st kick
        X = X + V * dt #drift
        acc = calc_grav_acc(X) #calculate force using X^(n+1)
        V = V + 0.5 * acc * dt #2nd kick
        t += dt
    end
    return X,V
end

X,V = run();
0


text:
#module Visualize
#using Bridge, PyPlot, StaticVector
#end
using JLD2

simid = 3
sim = [:lorenz1, :lorenz2, :lorenz3][simid]
simname = String(sim)

mkpath(joinpath("output", simname))
try # save cp of this file as documentation
    cp(@__FILE__(), joinpath("output",simname,"$simname.jl"); force=true)
catch
end

iterations = 100000
saveit = 5000
alpha = 0.5 # 1 - exp(-alpha Exp) is AR(1) coefficient of Brownian motion valued random walk
independent = false # independent proposals
adaptit = 1000 # adapt every `it`th step
adaptmax = iterations

if simid == 4
    partial = true
    adaptive = true
    initnu = :backward
elseif simid == 3
    partial = false
    adaptive = true
    initnu = :brown
    alpha = 0.5
elseif simid == 2
    partial = false
    adaptive = false
    initnu = :backward
    alpha = 5.0
elseif simid == 1
    partial = false
    adaptive = false
    initnu = :brown
    alpha = 5.0
else
    error("provide `simid in 1:3`")
end

include("lorenz.jl")

#M = div(M,10)
#M = M*4

XX = Vector{typeof(X)}(m)
XXmean = Vector{typeof(X)}(m)
XXᵒ = Vector{typeof(X)}(m)
WW = Vector{typeof(W)}(m)
WWᵒ = Vector{typeof(W)}(m)

# Create linear noise approximations

if initnu == :backward
    TPt = Bridge.LinearNoiseAppr{Bridge.Models.Lorenz,StaticArrays.SDiagonal{3,Float64},SVector{3,Float64}}
else
    TPt = Bridge.LinearAppr{StaticArrays.SDiagonal{3,Float64},StaticArrays.SArray{Tuple{3,3},Float64,2,9},SVector{3,Float64}}
end
TPᵒ = Bridge.GuidedBridge{SVector{3,Float64},StaticArrays.SArray{Tuple{3,3},Float64,2,9},TPt,Bridge.Models.Lorenz}

#Pt = Vector{TPt}(m)
#Pᵒ = Vector{TPᵒ}(m)
Pt = Vector(m)
Pᵒ = Vector(m)
H♢, v = Bridge.gpupdate(πH*one(SM), zero(SV), L, Σ, V.yy[end])

for i in m:-1:1
    tt_ = range(V.tt[i], stop=V.tt[i+1], length=M+1)
    XX[i] = Bridge.samplepath(tt_, zero(SV))
    WW[i] = Bridge.samplepath(tt_, zero(RV))

    a_ = Bridge.a(XX[i].tt[end], v, P)

     if initnu == :brown
        Pt[i] = Bridge.linearappr(SamplePath(XX[i].tt, zeros(ℝ{3}, M+1)), Bridge.NoDrift(P))
    elseif initnu == :backward
        Pt[i] = Bridge.LinearNoiseAppr(XX[i].tt, P, v, a_, :backward)
    elseif initnu == :foci
        if partial
            x_ = Models.foci(P)[1 + ( V.yy[i][1] > 0)]
        else
            x_ = Models.foci(P)[1 + ( V.yy[i][2] + V.yy[i][3] > 0)]
        end
        Pt[i] = Bridge.linearappr(SamplePath(XX[i].tt, fill(x_, M+1)), P)
    end

    Pᵒ[i] = Bridge.GuidedBridge(XX[i].tt, P, Pt[i], v, H♢)
    H♢, v = Bridge.gpupdate(Pᵒ[i], L, Σ, V.yy[i])
end

π0 = Bridge.Gaussian(v, Hermitian(H♢))

# initialize

y = π0.μ
for i in 1:m
    sample!(WW[i], Wiener{ℝ{3}}())
    y = Bridge.bridge!(XX[i], y, WW[i], Pᵒ[i])
end
XXmean = [zero(XX[i]) for i in 1:m]

function smooth(π0, XX, WW, P, Pᵒ, iterations, alpha; verbose = true,adaptive = true, adaptmax = iterations, adaptit = 5000, saveit = 500, independent = false,
    smoothmean = false, hwindow=20)
    Paths = []

    m = length(XX)

    # create workspace
    XXᵒ = deepcopy(XX)
    WWᵒ = deepcopy(WW)
    Wnr = Wiener{valtype(WW[1])}()

    # initialize
    mcstate = [mcstart(XX[i].yy) for i in 1:m]
    acc = 0
    y0 = π0.μ
    X0 = typeof(y0)[]
    Xt = typeof(y0)[]
    newblock = false
    for it in 1:iterations
        doaccept = false
        if it % saveit == 0
            push!(Paths, collect(Iterators.flatten(XX[i].yy[1:end-1] for i in 1:m)))
        end

        if adaptive && it < adaptmax && it % adaptit == 0 # adaptive smoothing
            H♢, v = Bridge.gpupdate(πH*one(SM), zero(SV), L, Σ, V.yy[end])
            for i in m:-1:1
                xx = mcstate[i][1]
                if isa(Pᵒ[i].Pt, Bridge.LinearNoiseAppr)
                    if smoothmean
                        Pᵒ[i].Pt.Y.yy[:] = [mean(xx[max(1, j-hwindow):min(end, j+hwindow)]) for j in 1:length(xx)]
                    else
                        Pᵒ[i].Pt.Y.yy[:] = xx
                    end
                elseif isa(Pᵒ[i].Pt, Bridge.LinearAppr)
                    if smoothmean
                        Y = SamplePath(XX[i].tt, [mean(xx[max(1, j-hwindow):min(end, j+hwindow)]) for j in 1:length(xx)])
                    else
                        Y = SamplePath(XX[i].tt, xx)
                    end
                    Bridge.linearappr!(Pᵒ[i].Pt, Y, P)
                else
                    error("Adaption not implemented for proposals Pᵒ.Pt")
                end
                Pᵒ[i] = Bridge.GuidedBridge(XX[i].tt, P, Pᵒ[i].Pt, v, H♢)
                H♢, v = Bridge.gpupdate(Pᵒ[i], L, Σ, V.yy[i])
            end
            π0 = Bridge.Gaussian(v, Hermitian(H♢))
            #y0ᵒ = rand(π0)
            newblock = true
            if it == adaptit # first adaptive proposal needs some help sometimes
                doaccept = true
            end

        end

        push!(X0, y0)
        push!(Xt, XX[div(end,2)].yy[1])
        rho_ = exp(-alpha*randexp())
        if newblock
            #y0ᵒ = rand(π0)
            y0ᵒ = y0
        elseif !independent

            y0ᵒ = π0.μ + sqrt(rho_)*(rand(π0) - π0.μ) + sqrt(1-rho_)*(y0 - π0.μ)
        else
            y0ᵒ = rand(π0)
        end
        y = y0ᵒ
        for i in 1:m
            sample!(WWᵒ[i], Wnr)
            if !independent
                WWᵒ[i].yy[:] = sqrt(rho_)*WWᵒ[i].yy + sqrt(1-rho_)*WW[i].yy
            end
            y = Bridge.bridge!(XXᵒ[i], y, WWᵒ[i], Pᵒ[i])
        end

        ll = 0.0
        for i in 1:m
            ll += llikelihood(LeftRule(), XXᵒ[i],  Pᵒ[i]) - llikelihood(LeftRule(), XX[i],  Pᵒ[i])
        end
        verbose && print(@sprintf("%7i", it), " ll ",  @sprintf("%10.3f", ll))

        #if true
        if doaccept || rand() < exp(ll)
            acc += 1
            verbose && print("\t ✓")
            y0 = y0ᵒ
            for i in 1:m
                XX[i], XXᵒ[i] = XXᵒ[i], XX[i]
                WW[i], WWᵒ[i] = WWᵒ[i], WW[i]
            end
            newblock = false
        else
            verbose && print("\t .")
        end
        verbose && println([" ", "N"][1 + newblock], "\t", round(acc/it, digits=2), "\t", round.(y0 - x0, 2))
        for i in 1:m
            mcstate[i] = Bridge.mcnext!(mcstate[i],XX[i].yy)
        end
    end
    Paths, X0, Xt, mcstate, acc
end

Paths, X0, Xt, mcstates, acc = smooth(π0, XX, WW, P, Pᵒ, iterations, alpha;
adaptive = adaptive,
adaptit = adaptit,
adaptmax = adaptmax,
saveit = saveit,
verbose = true, independent = independent)

open(joinpath("output", simname,"info.txt"), "w") do f
    println(f, "acc")
    println(f, round(acc/iterations, digits=3))
end

writedlm(joinpath("output", simname, "x0n$simid.csv"), [1:iterations Bridge.mat(X0)'])
writedlm(joinpath("output", simname, "xtn$simid.csv"), [1:iterations Bridge.mat(Xt)'])

# Plot result
include(joinpath(Pkg.dir("Bridge"),"extra/makie.jl"))

Xmeanm, Xstdm = Bridge.mcmarginalstats(mcstates)

Xmean, Xrot, Xscal = mcsvd3(mcstates)

save(joinpath(ENV["BRIDGE_OUTDIR"],  "$simname$(simid)paths.jld2"), "Path", Paths)
save(joinpath(ENV["BRIDGE_OUTDIR"],  "$simname$(simid)states.jld2"), "mcstates", mcstates)


text:
using JLD
using LinearAlgebra
using Distributions
using Plots
pyplot()

# Dimensionalities
DX = 1
DY = 1

# Length of time-series
T = 100

# Process noise precision
τ = 2.

# Measurement noise precision
γ = 10.

# Preallocate data arrays
X = zeros(T,)
Y = zeros(T,)

# Initial state
X_0 = 0.0

# Keep previous state in memory
X_tmin1 = X_0

# Generate data for entire time-series
for t = 1:T

    # Transition from previous state
    X[t] = X_tmin1 + sqrt(inv(τ))*randn(1,)[1]

    # Emission of current state
    Y[t] = X[t] + sqrt(inv(γ))*randn(1,)[1]

    # Update previous state
    X_tmin1 = X[t]

end

# Write data to file
save("source/probprog4/signal.jld",
     "Y", Y,
     "γ", γ,
     "T", T)

# Visualization of states
plot(1:T, X, color="red", label="states", grid=false, xlabel="time (t)", ylabel="signal")
scatter!(1:T, Y, color="black", label="observations")

max_stars_repo_name: Yikeru/BMLIP

max_issues_repo_name: Yikeru/BMLIP

max_forks_repo_name: Yikeru/BMLIP


text:
module PhaseBases
using  RecursiveArrayTools
using LinearAlgebra
using SampledDomains: CartesianDomain2D

export  ZernikeBW, compose, decompose
export zernike

include("types.jl")
include("utils.jl")

include("Zernike/Zernike.jl")

end


text:
export FFMLP, reset!

immutable FFMLP{C} <: MLP
    layers::Vector{Int}
    data::MLPData{C}
end

function FFMLP{C}(ctx::C, layers::Vector{Int})
    numLayers = length(layers)
    dims = Vector{WeightDim}()
    for i in 1:numLayers - 1
        rows = (layers[i] + 1)
        cols = layers[i + 1]
        push!(dims, WeightDim(rows, cols))
    end
    FFMLP{C}(layers, MLPData(ctx, dims))
end

makeInputVector(mlp::FFMLP, index) = addBias!(mlp.data, mlp.data.signal[index])

reset!(mlp::FFMLP) = nothing

Base.show(io::IO, mlp::FFMLP) = print(io, "Feed Forward Multiplayer Perceptron\nLayers: $(mlp.layers)\nData: $(mlp.data)")

max_stars_repo_name: unbornchikken/julia-ml-proto

max_issues_repo_name: unbornchikken/julia-ml-proto

max_forks_repo_name: unbornchikken/julia-ml-proto


text:
include("gamess_custom_basis_sets/631+gd_split.jl")
include("gamess_custom_basis_sets/6311++g2d2p_split.jl")
include("gamess_custom_basis_sets/6311++gdp-J.jl")

const atom_to_atomic_number_mapping = Dict(
  "H" => 1.0,
  "He" => 2.0,
  "Li" => 3.0,
  "Be" => 4.0,
  "B" => 5.0,
  "C" => 6.0,
  "N" => 7.0,
  "O" => 8.0,
  "F" => 9.0,
  "Ne" => 10.0,
  "Na" => 11.0,
  "Mg" => 12.0,
  "Al" => 13.0,
  "Si" => 14.0,
  "P" => 15.0,
  "S" => 16.0,
  "Cl" => 17.0,
  "Ar" => 18.0,
  "K" => 19.0,
  "Ca" => 20.0,
  "Sc" => 21.0,
  "Ti" => 22.0,
  "V" => 23.0,
  "Cr" => 24.0,
  "Mn" => 25.0,
  "Fe" => 26.0,
  "Co" => 27.0,
  "Ni" => 28.0,
  "Cu" => 29.0,
  "Zn" => 30.0,
  "Ga" => 31.0,
  "Ge" => 32.0,
  "As" => 33.0,
  "Se" => 34.0,
  "Br" => 35.0,
  "Kr" => 36.0,
  "Rb" => 37.0,
  "Sr" => 38.0,
  "Y" => 39.0,
  "Zr" => 40.0,
  "Nb" => 41.0,
  "Mo" => 42.0,
  "Tc" => 43.0,
  "Ru" => 44.0,
  "Rh" => 45.0,
  "Pd" => 46.0,
  "Ag" => 47.0,
  "Cd" => 48.0,
  "In" => 49.0,
  "Sn" => 50.0,
  "Sb" => 51.0,
  "Te" => 52.0,
  "I" => 53.0,
  "Xe" => 54.0
)

function create_input(input)
  xyz_array::Vector{String} = []
  xyz = open(input, "r") do file
    fragment = readlines(file)
    xyz_array = fragment[3:end]
  end
  natoms = length(xyz_array)

  output = open(joinpath(@__DIR__, input[1:(end-4)]*".inp"), "w") do file
    write(file, " \$scf dirscf=.t. damp=.t. shift=.f.\n")
    write(file, "   soscf=.f. diis=.t. fdiff=.f. extrap=.f. \$end\n")
    write(file, " \$system mwords=500 memddi=500 \$end\n")
    write(file, " \$intgrl schwrz=.t. intomp=1 shfock=.f. \$end\n")
    write(file, " \$guess guess=hcore \$end\n")
    write(file, " \$contrl scftyp=rhf runtyp=energy nprint=-5\n")
    write(file, "     icharg=0 ispher=0 inttyp=eric \n")
    write(file, "     maxit=100 \$end\n")
    write(file, " \$end\n")
    write(file, " \$data\n")
    write(file, "$(splitdir(input[1:(end-4)])[end]) \n")
    write(file, "C1\n")
    for line in xyz_array
      atom = rstrip(line[1:2])
      newline = line[1:2]*"   $(atom_to_atomic_number_mapping[atom])   "*
        line[3:end]
      write(file, "$newline\n")

      newline = six_three_one_one_2plus_2d_2p_split[atom]
      #newline = six_three_one_one_2plus_d_p_J[atom]
      #newline = six_three_one_plus_g_d_split[atom]
      write(file, "$newline\n")
    end
    write(file, " \$end\n")
  end
end


text:
moves = []
function mark_angles(r::Robot)
    go_to_corner!(r, moves)
    print_markers(r)
    return_back!(r::Robot, moves)
end

function go_to_corner!(r, moves)
    while !(isborder(r, West) && isborder(r, Sud))
        if !isborder(r, West)
            move!(r, West)
            push!(moves, West)
        end
        if !isborder(r, Sud)
            move!(r, Sud)
            push!(moves, Sud)
        end
    end
end

function moves!(r, side)
    while !isborder(r, side)
        move!(r, side)
    end
end

function print_markers(r)
    for side in (Nord, Ost, Sud, West)
        moves!(r,side)
        putmarker!(r)
    end
end

inverse(side::HorizonSide) = HorizonSide(mod(Int(side) + 2, 4))

function return_back!(r::Robot, moves)
    while length(moves) > 0
        side = pop!(moves)
        move!(r, inverse(side))
    end
end

max_stars_repo_name: VladislavBukin/mirea-progs

max_issues_repo_name: VladislavBukin/mirea-progs

max_forks_repo_name: VladislavBukin/mirea-progs


text:
const WRDS_FLDS = [:gvkey, :datadate, :fyear, :tic, :cusip, :conm, :indfmt,
  :curcd, :curncd, :costat, :loc, :mkvalt, :prcc_f,
  :stalt, :stko, :idbflag, :revt, :adrr, :seq, :naics, :at]
const WRDS_OK_MISSING_FLDS = [:adrr, :stalt]

function wrdsAggDF(; wrdsPath::String = WRDS_PATH, compustatName::String = COMPUSTAT_NAME,
    outcols::Vector{Symbol} = [:naics2name, :conm, :revt, :at, :seq],
    datasuffix::String=DATA_SUFFIX)::DataFrame

  local outdf::DataFrame
  local compdf::DataFrame
  #now work on the compustat data
  compdf = CSV.read("$wrdsPath\\$compustatName.csv")
  compdf = compdf[WRDS_FLDS]
  compdf = compdf[completecases(compdf[setdiff(WRDS_FLDS, WRDS_OK_MISSING_FLDS)]), :]

  filterfunc(r::DataFrameRow) = (
    r[:indfmt] == "INDL" &&
    r[:curcd] == "USD" && #native currency and reported currency in USD
    r[:curncd] == "USD" && #ibid
    #r[:costat] == "A" && #active in db
    r[:loc] == "USA" && #incorporated or registered in US
    r[:mkvalt] > 0.1 && #has a market value > 100k
    (ismissing(r[:stalt]) ||  r[:stalt] == "AW") &&#not in reogranization/bankruptcy
    r[:stko] < 2 && #ignore very thinly traded securities (no stale quotes)
    #r[:idbflag] ≠ "I" && #incorporated or registered in US
    r[:revt] > 0.1 && #revenue is reported and > 100k
    (ismissing(r[:adrr])) && #not an adr
    r[:seq] > 0.1 #book equity reported and > 100k
  )

  compdf = compdf[((r::DataFrameRow)->filterfunc(r)).(eachrow(compdf)),:]

  #make sure we don't have dupe companies
  compdf[:tokeep] = trues(size(compdf,1))
  compdfBygvkey::GroupedDataFrame = groupby(compdf, :gvkey)
  for subdf::SubDataFrame ∈ compdfBygvkey
    if size(subdf,1) > 1
      subdf[:tokeep] = subdf[:datadate] .== maximum(subdf[:datadate])
      (sum(subdf[:tokeep]) > 1) && error("Duplicate firm remains in compustat data, add new filtering rule")
    end
  end
  compdf = compdf[compdf[:tokeep],:]
  compdf[:naics2] = (s::Int->Symbol(("$s")[1:2])).(compdf[:naics])
  compdf[:naics2name] = (s::Symbol->(naicsCodes[s])).(compdf[:naics2])

  #now write out the largest firms
  sort!(compdf, (:naics2name, order(:at,rev=true)))
  outdf = DataFrame(deepcopy(compdf[[],outcols]))

  for subdf::SubDataFrame ∈ groupby(compdf, :naics2name)
    numentries::Int = min(3,size(subdf,1))
    append!(outdf, DataFrame(deepcopy(subdf[1:numentries,outcols])))
  end
  outdf[:firmtype] = :nonprofit

  #housekeeping to match the output dataframe
  rename!(outdf, [:revt=>:revenue, :conm=>:name, :at=>:assets, :seq=>:netrentassets])
  CSV.write("$WORKING_PATH\\largestPublicForProfits_$datasuffix.csv", outdf)

  #this aggregates the industry data using new data frames 15.2+ functions
  compaggdf::DataFrame = by(
    compdf, :naics2name,
      [:naics2name,:mkvalt, :seq, :revt, :at] =>
        x -> (n = length(x.naics2name),
        revenue=sum(skipmissing(x.revt)) * 1_000_000.,
        assets=sum(skipmissing(x.at)) * 1_000_000.,
        netassets=sum(skipmissing(x.seq)) * 1_000_000.))

  #create normalized versions
  compaggdf[:drevenue] = compaggdf[:revenue] ./ sum(compaggdf[:revenue])
  compaggdf[:dassets] = compaggdf[:assets] ./ sum(compaggdf[:assets])
  compaggdf[:dnetassets] = compaggdf[:netassets] ./ sum(compaggdf[:netassets])

  compaggdf[:firmtype] = :publicforprofit

  return compaggdf
end

const SOI_FLDS = [:naics2, :n,	:assets, :netassets,	:revenue]

#this grabs the pre-aggregated SOI data (already aggregated by IRS)
function soiAggDF(irsSOIPath::String = IRS_SOI_PATH, soiName::String = SOI_NAME)::DataFrame

  compaggdf::DataFrame = CSV.read("$irsSOIPath\\$(soiName).csv")
  compaggdf = compaggdf[SOI_FLDS]

  compaggdf[:naics2name] = (i::Int->(naicsCodes[Symbol("$i")])).(compaggdf[:naics2])
  deletecols!(compaggdf, :naics2)

  #create normalized versions
  compaggdf[:drevenue] = compaggdf[:revenue] ./ sum(compaggdf[:revenue])
  compaggdf[:dassets] = compaggdf[:assets] ./ sum(compaggdf[:assets])
  compaggdf[:dnetassets] = compaggdf[:netassets] ./ sum(compaggdf[:netassets])

  compaggdf[:firmtype] = :soiforprofit

  return compaggdf
end

#make the industry distribution plots
function industryCrossDF(data::NCCSData; assetYear::Int = INDUSTRY_CROSS_YEAR,
    industryComparisonSource::Symbol=:wrds,
    outcols::Vector{Symbol} = [:naics2name, :name, :ein, :totrev, :bookassets, :bookliabilities],
    datasuffix::String=DATA_SUFFIX
    )::DataFrame
  local presubdf::SubDataFrame = view(data.df, data.df[:fisyr].==assetYear,:)
  local comparisondf::DataFrame
  local outdf::DataFrame

  #write out the largest firms by assets
  #need to copy it to use the sort command
  local preoutdf::DataFrame = DataFrame(deepcopy(presubdf[:,outcols]))
  sort!(preoutdf, (:naics2name, order(:bookassets, rev=true)))

  outdf = deepcopy(preoutdf[[],:]) #create a stub df for collecting
  for subdf::SubDataFrame ∈ groupby(preoutdf, :naics2name)
    numentries::Int = min(3,size(subdf,1))
    append!(outdf, DataFrame(deepcopy(subdf[1:numentries,outcols])))
  end
  outdf[:netassets] = outdf[:bookassets] .- outdf[:bookliabilities]
  outdf[:firmtype] = :nonprofit

  #housekeeping to match the output dataframe
  deletecols!(outdf, :bookliabilities)
  rename!(outdf, [:totrev=>:revenue, :bookassets=>:assets])
  CSV.write("$WORKING_PATH\\largestNonProfits_$datasuffix.csv", outdf)

  #this aggregates the industry data using new data frames 15.2+ functions
  local aggdf::DataFrame = by(
    presubdf[:,[:naics2, :naics2name,:totrev, :totexp, :netassets, :bookassets, :bookliabilities]],
      :naics2name,
      [:naics2name,:totrev,:totexp,:netassets, :bookassets, :bookliabilities] =>
        x -> (n = length(x.naics2name),
        revenue=sum(skipmissing(x.totrev)),
        #profits = sum(skipmissing(x.totrev .- x.totexp)),
        assets=sum(skipmissing(x.bookassets)),
        netassets=sum(skipmissing(x.bookassets .- x.bookliabilities))))
  #create normalized versions
  (assetYear ≥ 2000) && (aggdf[:drevenue] = aggdf[:revenue] ./ sum(aggdf[:revenue]))
  aggdf[:dassets] = aggdf[:assets] ./ sum(aggdf[:assets])
  aggdf[:dnetassets] = aggdf[:netassets] ./ sum(aggdf[:netassets])

  aggdf[:firmtype] = :nonprofit

  #load the comparison df based on the requested source
  if industryComparisonSource == :wrds
    comparisondf = wrdsAggDF()
  elseif industryComparisonSource == :soi
    comparisondf = soiAggDF()
  end

  assetYear < 2000 && deletecols!(comparisondf, :revenue)

  aggdf = vcat(aggdf, comparisondf)

  #missing values really correspond to 0
  for s::Symbol ∈ [:revenue, :assets, :netassets, :drevenue, :dassets, :dnetassets]
    aggdf[:,s] .= (f::MFloat64->ismissing(f) ? 0 : f).(aggdf[:,s])
  end

  #this snippit creates a 0-value row if an industry is missing
  #the purpose is to provide the placeholder value to the bar charts
  for s::Symbol ∈ unique(aggdf[:naics2name])
    for t::Symbol ∈ unique(aggdf[:firmtype])
      if t ∉ aggdf[aggdf[:naics2name] .== s, :firmtype]
        aggdf = vcat(aggdf, DataFrame(naics2name=[s], firmtype=[t],n=[0.],
          revenue=[0.], assets=[0.], netassets=[0.],
          drevenue=[0.], dassets=[0.], dnetassets=[0.]))
      end
    end
  end

  #create proportional logged results
  #each value is the corresponding percentage of the total unlogged value in each industry
  # times the logged value
  aggdf[:prevenue] = similar(aggdf[:revenue])
  aggdf[:passets] = similar(aggdf[:assets])
  aggdf[:pnetassets] = similar(aggdf[:netassets])
  aggdfBynaics2name::GroupedDataFrame = groupby(aggdf, :naics2name)
  for subdf::SubDataFrame ∈ aggdfBynaics2name

    subdf[:prevenue] = log(sum(subdf[:revenue])) .* (subdf[:revenue] ./ sum(subdf[:revenue]))
    subdf[:passets] = log(sum(subdf[:assets])) .* (subdf[:assets] ./ sum(subdf[:assets]))
    subdf[:pnetassets] = log(sum(subdf[:netassets])) .* (subdf[:netassets] ./ sum(subdf[:netassets]))
  end

  return aggdf
end

function industryCross(data::NCCSData;
        outputPath::String = GRAPH_PATH,
        industryComparisonSource::Symbol = INDUSTRY_COMPARISON_SOURCE)::Nothing

  #runs the routine for multiple comparison sources if desired
  if industryComparisonSource == :soiwrds
    industryCross(data, outputPath=outputPath, industryComparisonSource=:soi)
    industryComparisonSource = :wrds
  end

  local aggdf::DataFrame = industryCrossDF(data, industryComparisonSource=industryComparisonSource)
  local plotNames::Vector{String} = Vector{String}()
  local plots::Vector{PlotContainer} = Vector{PlotContainer}()

  sort!(aggdf, [:naics2name, :firmtype])

  #first make the concentration plots
  push!(plotNames, "Concentration_of_Revenue_by_Industry")
  push!(plots, plot(aggdf,
    x=:naics2name, y=:drevenue, color=:firmtype,
    Guide.title("Distribution of Revenue by Industry"),
    Guide.xlabel("Sector (NAICS-2)"), Guide.ylabel("Percentage of Revenue"),
     Geom.bar(position=:dodge)))

   push!(plotNames, "Concentration_of_Assets_by_Industry")
   push!(plots, plot(aggdf,
     x=:naics2name, y=:dassets, color=:firmtype,
     Guide.title("Distribution of Assets by Industry"),
     Guide.xlabel("Sector (NAICS-2)"), Guide.ylabel("Percentage of Total Assets"),
      Geom.bar(position=:dodge),
      Coord.cartesian(ymin=0.0, ymax=0.4)))

   push!(plotNames, "Concentration_of_Book_Equity_by_Industry")
   push!(plots, plot(aggdf,
     x=:naics2name, y=:dnetassets, color=:firmtype,
     Guide.title("Distribution of Equity by Industry"),
     Guide.xlabel("Sector (NAICS-2)"), Guide.ylabel("Percentage of Total Book Equity"),
      Geom.bar(position=:dodge)))

  #lets make some stacked plots
  push!(plotNames, "Revenue_by_Industry")
  push!(plots, plot(aggdf,
    x=:naics2name, y=:prevenue, color=:firmtype,
    Guide.title("Revenue by Industry"),
    Guide.xlabel("Sector (NAICS-2)"), Guide.ylabel("Log Revenue (Color areas are % of total bar)"),
     Geom.bar(position=:stack)))

   push!(plotNames, "Assets_by_Industry")
   push!(plots, plot(aggdf,
     x=:naics2name, y=:passets, color=:firmtype,
     Guide.title("Assets by Industry"),
     Guide.xlabel("Sector (NAICS-2)"), Guide.ylabel("Log Assets (Color areas are % of total bar)"),
      Geom.bar(position=:stack)))

   push!(plotNames, "Book_Equity_by_Industry")
   push!(plots, plot(aggdf,
     x=:naics2name, y=:pnetassets, color=:firmtype,
     Guide.title("Equity by Industry"),
     Guide.xlabel("Sector (NAICS-2)"), Guide.ylabel("Log Book Equity (Color areas are % of total bar)"),
      Geom.bar(position=:stack)))

  for i ∈ 1:length(plots) #write the graphs
    draw(PDF("$outputPath\\$(plotNames[i])_$(industryComparisonSource)_$INDUSTRY_CROSS_YEAR.pdf", 9inch, 7inch), plots[i])
    #println("Graph $(plotNames[i]) written.")
  end

  println("Cross-secitonal industry graphs drawn.")

  return nothing
end

max_stars_repo_name: clintonTE/CCA

max_issues_repo_name: clintonTE/CCA

max_forks_repo_name: clintonTE/CCA


text:
@recipe(Violin, x, y) do scene
    Theme(;
        default_theme(scene, Poly)...,
        side = :both,
        width = 0.8,
        trim = false,
        strokecolor = :white,
        show_median = false,
        mediancolor = automatic,
        medianlinewidth = 1.0,
    )
end

conversion_trait(x::Type{<:Violin}) = SampleBased()

function plot!(plot::Violin)
    width, side, trim, show_median = plot[:width], plot[:side], plot[:trim], plot[:show_median]

    signals = lift(plot[1], plot[2], width, side, trim, show_median) do x, y, bw, vside, trim, show_median
        vertices = Vector{Point2f0}[]
        lines = Pair{Point2f0, Point2f0}[]
        for (key, idxs) in finduniquesorted(x)
            v = view(y, idxs)
            spec = (x = key, kde = density(v; trim = trim), median = median(v))
            min, max = extrema_nan(spec.kde.density)
            scale = 0.5*bw/max
            xl = reverse(spec.x .- spec.kde.density .* scale)
            xr = spec.x .+ spec.kde.density .* scale
            yl = reverse(spec.kde.x)
            yr = spec.kde.x

            x_coord, y_coord = if vside == :left
                [spec.x; xl; spec.x], [yl[1]; yl; yl[end]]
            elseif vside == :right
                [spec.x; xr; spec.x], [yr[1]; yr; yr[end]]
            else
                [spec.x; xr; spec.x; xl], [yr[1]; yr; yl[1]; yl]
            end
            verts = Point2f0.(x_coord, y_coord)
            push!(vertices, verts)

            if show_median
                # interpolate median bounds between corresponding points
                xm = spec.median
                ip = findfirst(>(xm), spec.kde.x)
                ym₋, ym₊ = spec.kde.density[ip-1], spec.kde.density[ip]
                xm₋, xm₊ = spec.kde.x[ip-1], spec.kde.x[ip]
                ym = (xm * (ym₊ - ym₋) + xm₊ * ym₋ - xm₋ * ym₊) / (xm₊ - xm₋)
                median_left = Point2f0(vside == :right ? spec.x : spec.x - ym * scale, xm)
                median_right = Point2f0(vside == :left ? spec.x : spec.x + ym * scale, xm)
                push!(lines, median_left => median_right)
            end
        end
        return vertices, lines
    end
    t = copy(Theme(plot))
    mediancolor = pop!(t, :mediancolor)
    poly!(plot, t, lift(first, signals))
    linesegments!(
        plot,
        lift(last, signals),
        color = lift(
            (mc, sc) -> mc === automatic ? sc : mc,
            mediancolor,
            plot[:strokecolor],
        ),
        linewidth = plot[:medianlinewidth],
        visible = plot[:show_median],
    )
end

max_stars_repo_name: piever/StatMakie

max_issues_repo_name: piever/StatMakie

max_forks_repo_name: piever/StatMakie


text:
# Data is from  ``Data analysis recipes: Fitting a model to data'' by Hogg, Bovy & Lang
# available on https://arxiv.org/abs/1008.4686.

if !isdef(:penalty_names); penalty_names = ["huber"] end
if !isdef(:ovf_formulations); ovf_formulations = ["equilibrium"] end

@testset "Equilibrium loss test: penalty = $penalty_name; ovf_formulation = $ovf_formulation" for penalty_name in penalty_names, ovf_formulation in ovf_formulations

if ovf_formulation == "dual"
  if penalty_name == "hubnik"
    n_MP = 2
  else
    n_MP = 3
  end
else
  n_MP = 5
end

m = direct_model(ReSHOP.Optimizer(;ovf_formulation=ovf_formulation))

fit_model = EMP.Model(m)

mps = [MathPrgm(fit_model) for i=1:n_MP]

EquilibriumProblem(fit_model, mps)

# primary unknonwn
@variableMP(mps[1], c1[j=1:1])
@variableMP(mps[1], d1)

# OVF variable
@variableMP(mps[1], loss_var1)

# OVF arguments
@variableMP(mps[1], fit1[i=1:N]);

@constraintMP(mps[1], fiteqn1[i=1:N], (y[i] - sum(c1[j]*x[i, j] for j=1:M) - d1) / sigma_y[i] == fit1[i]);

addovf!(fit_model, loss_var1, fit1, penalty_name, all_params[penalty_name])

@objectiveMP(mps[1], Min, loss_var1)

# primary unknonwn
@variableMP(mps[2], c2[j=1:1])
@variableMP(mps[2], d2)

# OVF variable
@variableMP(mps[2], loss_var2)

# OVF arguments
@variableMP(mps[2], fit2[i=1:N]);

@NLconstraintMP(mps[2], fiteqn2[i=1:N], log(exp((y[i] - sum(c2[j]*x[i, j] for j=1:M) - d2) / sigma_y[i])) == fit2[i]);

addovf!(fit_model, loss_var2, fit2, penalty_name, all_params[penalty_name])

@objectiveMP(mps[2], Min, loss_var2)

if n_MP >= 3
# primary unknonwn
@variableMP(mps[3], c3[j=1:1])
@variableMP(mps[3], d3)

# OVF variable
@variableMP(mps[3], loss_var3)

# OVF arguments
@variableMP(mps[3], fit3[i=1:N]);

@constraintMP(mps[3], fiteqn3[i=1:N], fit3[i] == (y[i] - sum(c3[j]*x[i, j] for j=1:M) - d3) / sigma_y[i]);

addovf!(fit_model, loss_var3, fit3, penalty_name, all_params[penalty_name])

@objectiveMP(mps[3], Min, loss_var3)

elseif n_MP >= 4
  # primary unknonwn
  @variableMP(mps[4], c4[j=1:1])
  @variableMP(mps[4], d4)

  # OVF variable
  @variableMP(mps[4], loss_var4)

  # OVF arguments
  @variableMP(mps[4], fit4[i=1:N]);

  @NLconstraintMP(mps[4], fiteqn4[i=1:N], fit4[i] == log(exp((y[i] - sum(c4[j]*x[i, j] for j=1:M) - d4) / sigma_y[i])));

  addovf!(fit_model, loss_var4, fit4, penalty_name, all_params[penalty_name])

  @objectiveMP(mps[4], Min, loss_var4)

# primary unknonwn
  @variableMP(mps[5], c5[j=1:1])
  @variableMP(mps[5], d5)

  # OVF variable
  @variableMP(mps[5], loss_var5)

  # OVF arguments
  @variableMP(mps[5], fit5[i=1:N]);

  @constraintMP(mps[5], fiteqn5[i=1:N], pi*(y[i] - sum(c5[j]*x[i, j] for j=1:M) - d5) / sigma_y[i] == pi*fit5[i]);

  addovf!(fit_model, loss_var5, fit5, penalty_name, all_params[penalty_name])

  @objectiveMP(mps[5], Min, loss_var5)
end

cmp_dir = joinpath(dirname(@__FILE__), "test_res")
unknown_ref = readdlm(joinpath(cmp_dir, "mcp_" * penalty_name * "_v1.out"))
file_ref = readdlm(joinpath(cmp_dir, "mcp_" * penalty_name * "_v1_fit.out"))

solveEMP(fit_model)
@test EMP.termination_status(fit_model) == MOI.LOCALLY_SOLVED

occursin("hinge", penalty_name) && continue

@test isapprox(value.(c1), [unknown_ref[1]], rtol=1e-4)
@test isapprox(value(d1), unknown_ref[2], rtol=1e-4)
@test isapprox(value.(fit1), file_ref[1:N], rtol=1e-4)
@test isapprox(value(loss_var1), unknown_ref[3], rtol=1e-4)

@test isapprox(value.(c2), [unknown_ref[1]], rtol=1e-4)
@test isapprox(value(d2), unknown_ref[2], rtol=1e-4)
@test isapprox(value.(fit2), file_ref[1:N], rtol=1e-4)
@test isapprox(value(loss_var2), unknown_ref[3], rtol=1e-4)

if n_MP >= 3
@test isapprox(value.(c3), [unknown_ref[1]], rtol=1e-4)
@test isapprox(value(d3), unknown_ref[2], rtol=1e-4)
@test isapprox(value.(fit3), file_ref[1:N], rtol=1e-4)
@test isapprox(value(loss_var3), unknown_ref[3], rtol=1e-4)

elseif n_MP >= 4
  @test isapprox(value.(c4), [unknown_ref[1]], rtol=1e-4)
  @test isapprox(value(d4), unknown_ref[2], rtol=1e-4)
  @test isapprox(value.(fit4), file_ref[1:N], rtol=1e-4)
  @test isapprox(value(loss_var4), unknown_ref[3], rtol=1e-4)

  @test isapprox(value.(c5), [unknown_ref[1]], rtol=1e-4)
  @test isapprox(value(d5), unknown_ref[2], rtol=1e-4)
  @test isapprox(value.(fit5), file_ref[1:N], rtol=1e-4)
  @test isapprox(value(loss_var5), unknown_ref[3], rtol=1e-4)
end

end