text:
module HighestDensityRegions

using ArgCheck: @argcheck
using DocStringExtensions: SIGNATURES

export hdr_thresholds

####
#### utilities
####

"""
$(SIGNATURES)

When `i == 0`, return `0` (with the element type of the vector), otherwise
behaves like `getindex`.
"""
@inline getindex_or_zero(x::AbstractVector{T}, i::Integer) where T =
    i == 0 ? zero(T) : x[i]

"""
$(SIGNATURES)

Find `i` in the sorted vector `xs` such that `xs[i-1] < α ≤ xs[i]`, with the
assumed extension ``xs[0] = 0``.

Return `γ, i` such that `interpolate_in(xs, γ, i) ≈ α`. See
[`interpolate_in`](@ref).
"""
function interpolation_coefficient_and_index(xs::AbstractVector{T}, α) where {T <: Real}
    i = searchsortedfirst(xs, α)
    right = xs[i]
    left = getindex_or_zero(xs, i - 1)
    γ = (α - left) / (right - left)
    γ, i
end

"""
$(SIGNATURES)

Calculate the linear interpolation `(1-γ)*xs[i-1] + γ*xs[i]`, with the assumed
extension ``xs[0] = 0``.
"""
function interpolate_in(xs, γ, i)
    left = getindex_or_zero(xs, i - 1)
    right = xs[i]
    (1 - γ) * left + γ * right
end

"""
$(SIGNATURES)

Return the thresholds for highest density regions with probability `αs`, given
the sorted densities and cumulative probabilities (using the same sorting order)
for each bin.
"""
function interpolated_density_thresholds(αs, sorted_densities, cumprobs)
    @argcheck all(0 .≤ αs .≤ 1)
    map(αs) do α
        γ, i = interpolation_coefficient_and_index(cumprobs, 1 - α)
        interpolate_in(sorted_densities, γ, i)
    end
end

"""
$(SIGNATURES)

Cumulative probability.
"""
cumprob(x) = cumsum(x) |> x -> x ./ x[end]

# interface

"""
$(SIGNATURES)

Return thresholds for the highest density regions with coverage probabilities
`αs`.

`counts` is a vector proportional to count data for some bin structure, while
`densities` is a vector of corresponding densities. For each bin, it is assumed
that `counts[i] == densities[i] * bin_size[i]`, where `bin_size[i]` is the
measure of the bin.

A highest density region threshold `t` for a given `α` satisfies
`sum(counts[densities .> t]) ≈ α * sum(counts)`.

When bins have the same size, `counts` may be omitted and assumed to be
proportional to `densities`.
"""
function hdr_thresholds(αs, densities::AbstractVector, counts::AbstractVector)
    order = sortperm(densities)
    sorted_densities = densities[order]
    interpolated_density_thresholds(αs, sorted_densities, cumprob(counts[order]))
end

function hdr_thresholds(αs, densities)
    sorted_densities = sort(densities)
    interpolated_density_thresholds(αs, sorted_densities, cumprob(sorted_densities))
end

end # module


text:
function check_order_conditions_b(tab::Tableau{T}, k; atol=16eps(T), rtol=16eps(T)) where {T}
    b, c, s = tab.b, tab.c, tab.s
    isapprox(mapreduce(i -> b[i] * c[i]^(k-1), +, 1:s), 1/k; atol=atol, rtol=rtol)
end

function satisfies_simplifying_assumption_b(tab::Tableau, σ=tab.s; kwargs...)
    all([check_order_conditions_b(tab, k; kwargs...) for k in 1:σ])
end

function check_order_conditions_c(tab::Tableau{T}, k; atol=16eps(T), rtol=16eps(T)) where {T}
    a, c, s = tab.a, tab.c, tab.s
    [isapprox(mapreduce(j -> a[i,j] * c[j]^(k-1), +, 1:s), c[i]^k/k; atol=atol, rtol=rtol) for i in 1:s]
end

function satisfies_simplifying_assumption_c(tab::Tableau, σ=tab.s; kwargs...)
    all(hcat([check_order_conditions_c(tab, k; kwargs...) for k in 1:σ]...))
end

function check_order_conditions_d(tab::Tableau{T}, k; atol=16eps(T), rtol=16eps(T)) where {T}
    a, b, c, s = tab.a, tab.b, tab.c, tab.s
    [isapprox(mapreduce(j -> b[j] * c[j]^(k-1) * a[j,i], +, 1:s), b[i] * (1 - c[i]^k)/k; atol=atol, rtol=rtol) for i in 1:s]
end

function satisfies_simplifying_assumption_d(tab::Tableau, σ=tab.s; kwargs...)
    all(hcat([check_order_conditions_d(tab, k; kwargs...) for k in 1:σ]...))
end

@doc raw"""
Compute the weights by solving the simplifying assumption $B(s)$:
```math
\sum \limits_{j=1}^{s} b_{j} c_{j}^{k-1} = \frac{1}{k}  \qquad k = 1 , \, ... , \, s .
```
"""
function solve_simplifying_assumption_b(c::AbstractVector{T}) where {T}
    s = length(c)
    M = [ c[j]^(k-1) for k in 1:s, j in 1:s ]
    r = [ 1 / T(k) for k in 1:s ]
    M \ r
end

@doc raw"""
Compute the coefficients by solving the simplifying assumption $C(s)$:
```math
\sum \limits_{j=1}^{s} a_{ij} c_{j}^{k-1} = \frac{c_i^k}{k}  \qquad i = 1 , \, ... , \, s , \; k = 1 , \, ... , \, s .
```
"""
function solve_simplifying_assumption_c(c::AbstractVector{T}) where {T}
    s = length(c)
    M = [ c[j]^(k-1) for k in 1:s, j in 1:s ]

    row(i) = begin
        r = [ c[i]^k / k for k in 1:s ]
        M \ r
    end

    vcat([transpose(row(i)) for i in 1:s]...)
end

@doc raw"""
Compute the coefficients by solving the simplifying assumption $D(s)$:
```math
\sum \limits_{i=1}^{s} b_i c_{i}^{k-1} a_{ij} = \frac{b_j}{k} ( 1 - c_j^k)  \qquad j = 1 , \, ... , \, s , \; k = 1 , \, ... , \, s .
```
"""
function solve_simplifying_assumption_d(b::AbstractVector{T}, c::AbstractVector{T}) where {T}
    @assert axes(b) == axes(c)

    s = length(c)
    M = [ b[i] * c[i]^(k-1) for k in 1:s, i in 1:s ]

    row(j) = begin
        r = [ b[j] / k * (1 - c[j]^k) for k in 1:s ]
        M \ r
    end

    hcat([row(j) for j in 1:s]...)
end

ext: jl


text:
# This file was generated, do not modify it. # hide
for (id, d) in res
    @show id => d
end

max_stars_repo_name: sadit/SimilaritySearchDemos

max_issues_repo_name: sadit/SimilaritySearchDemos

max_forks_repo_name: sadit/SimilaritySearchDemos


text:
#=
  Sum of 2 dice.

  I asked how one should observe a sum of the two dice here:
  https://discourse.julialang.org/t/turing-observing-sum-of-two-random-dice/50486

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function two_dice(ss)
    d1 ~ DiscreteUniform(1, 6)
    d2 ~ DiscreteUniform(1, 6)

    s ~ Dirac(d1 + d2)
    true ~ Dirac(s == ss)

end

# Note: Here we observe c=5 from outside
model = two_dice(5)

# chns = sample(model, Prior(), 10_000)
# chns = sample(model, MH(), 10_000)
# chns = sample(model, IS(), 10_000)
chns = sample(model, PG(20), 1_000)
# chns = sample(model, SMC(1000), 1000)

display(chns)
# display(plot(chns))
# display(histogram(chns))
# display(gelmandiag(chns))

show_var_dist_pct(chns, :d1)
show_var_dist_pct(chns, :d2)
show_var_dist_pct(chns, :s)

max_stars_repo_name: tias/hakank

max_issues_repo_name: tias/hakank

max_forks_repo_name: tias/hakank


text:
immutable NoBackend <: AbstractBackend end

const _backendType = Dict{Symbol, DataType}(:none => NoBackend)
const _backendSymbol = Dict{DataType, Symbol}(NoBackend => :none)
const _backends = Symbol[]
const _initialized_backends = Set{Symbol}()

backends() = _backends
backend_name() = CURRENT_BACKEND.sym
_backend_instance(sym::Symbol) = haskey(_backendType, sym) ? _backendType[sym]() : error("Unsupported backend $sym")

macro init_backend(s)
    str = lowercase(string(s))
    sym = Symbol(str)
    T = Symbol(string(s) * "Backend")
    esc(quote
        immutable $T <: AbstractBackend end
        export $sym
        $sym(; kw...) = (default(; kw...); backend(Symbol($str)))
        backend_name(::$T) = Symbol($str)
        push!(_backends, Symbol($str))
        _backendType[Symbol($str)] = $T
        _backendSymbol[$T] = Symbol($str)
        include("backends/" * $str * ".jl")
    end)
end

@init_backend Immerse
@init_backend Gadfly
@init_backend PyPlot
@init_backend Qwt
@init_backend UnicodePlots
@init_backend Winston
@init_backend Bokeh
@init_backend Plotly
@init_backend PlotlyJS
@init_backend GR
@init_backend GLVisualize
@init_backend PGFPlots

include("backends/web.jl")
# include("backends/supported.jl")

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

# don't do anything as a default
_create_backend_figure(plt::Plot) = nothing
_prepare_plot_object(plt::Plot) = nothing
_initialize_subplot(plt::Plot, sp::Subplot) = nothing

_series_added(plt::Plot, series::Series) = nothing
_series_updated(plt::Plot, series::Series) = nothing

_before_layout_calcs(plt::Plot) = nothing

title_padding(sp::Subplot) = sp[:title] == "" ? 0mm : sp[:titlefont].pointsize * pt
guide_padding(axis::Axis) = axis[:guide] == "" ? 0mm : axis[:guidefont].pointsize * pt

# Set the (left, top, right, bottom) minimum padding around the plot area
# to fit ticks, tick labels, guides, colorbars, etc.
function _update_min_padding!(sp::Subplot)
    leftpad   = 10mm + sp[:left_margin]   + guide_padding(sp[:yaxis])
    toppad    =  2mm + sp[:top_margin]    + title_padding(sp)
    rightpad  =  3mm + sp[:right_margin]
    bottompad =  5mm + sp[:bottom_margin] + guide_padding(sp[:xaxis])
    # @show (leftpad, toppad, rightpad, bottompad)
    sp.minpad = (leftpad, toppad, rightpad, bottompad)
end

_update_plot_object(plt::Plot) = nothing

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

type CurrentBackend
  sym::Symbol
  pkg::AbstractBackend
end
CurrentBackend(sym::Symbol) = CurrentBackend(sym, _backend_instance(sym))

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

function pickDefaultBackend()
    env_default = get(ENV, "PLOTS_DEFAULT_BACKEND", "")
    if env_default != ""
        try
            Pkg.installed(env_default)  # this will error if not installed
            sym = Symbol(lowercase(env_default))
            if haskey(_backendType, sym)
                return backend(sym)
            else
                warn("You have set PLOTS_DEFAULT_BACKEND=$env_default but it is not a valid backend package.  Choose from:\n\t",
                     join(sort(_backends), "\n\t"))
            end
        catch
            warn("You have set PLOTS_DEFAULT_BACKEND=$env_default but it is not installed.")
        end
    end

    # the ordering/inclusion of this package list is my semi-arbitrary guess at
    # which one someone will want to use if they have the package installed...accounting for
    # features, speed, and robustness
    for pkgstr in ("PyPlot", "GR", "PlotlyJS", "Immerse", "Gadfly", "UnicodePlots")
        if Pkg.installed(pkgstr) != nothing
            return backend(Symbol(lowercase(pkgstr)))
        end
    end

    # the default if nothing else is installed
    backend(:plotly)
end

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

"""
Returns the current plotting package name.  Initializes package on first call.
"""
function backend()

  global CURRENT_BACKEND
  if CURRENT_BACKEND.sym == :none
    pickDefaultBackend()
  end

  sym = CURRENT_BACKEND.sym
  if !(sym in _initialized_backends)

    # initialize
    println("[Plots.jl] Initializing backend: ", sym)

    inst = _backend_instance(sym)
    try
      _initialize_backend(inst)
    catch err
      warn("Couldn't initialize $sym.  (might need to install it?)")
      rethrow(err)
    end

    push!(_initialized_backends, sym)

  end
  CURRENT_BACKEND.pkg
end

"""
Set the plot backend.
"""
function backend(pkg::AbstractBackend)
    CURRENT_BACKEND.sym = backend_name(pkg)
    warn_on_deprecated_backend(CURRENT_BACKEND.sym)
    CURRENT_BACKEND.pkg = pkg
end

function backend(modname::Symbol)
    warn_on_deprecated_backend(modname)
    CURRENT_BACKEND.sym = modname
    CURRENT_BACKEND.pkg = _backend_instance(modname)
end

const _deprecated_backends = [:qwt, :winston, :bokeh, :gadfly, :immerse]

function warn_on_deprecated_backend(bsym::Symbol)
    if bsym in _deprecated_backends
        warn("Backend $bsym has been deprecated.  It may not work as originally intended.")
    end
end

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

supported_types(::AbstractBackend) = []
supported_styles(::AbstractBackend) = [:solid]
supported_markers(::AbstractBackend) = [:none]
supported_scales(::AbstractBackend) = [:identity]
is_subplot_supported(::AbstractBackend) = false
is_string_supported(::AbstractBackend) = false
nativeImagesSupported(b::AbstractBackend) = :image in supported_types(b)

supported_types() = supported_types(backend())
supported_styles() = supported_styles(backend())
supported_markers() = supported_markers(backend())
supported_scales() = supported_scales(backend())
is_subplot_supported() = is_subplot_supported(backend())
is_string_supported() = is_string_supported(backend())
nativeImagesSupported() = nativeImagesSupported(backend())

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

# these are args which every backend supports because they're not used in the backend code
const _base_supported_args = [
    :color_palette,
    :background_color, :background_color_subplot,
    :foreground_color, :foreground_color_subplot,
    :group,
    :seriestype,
    :seriescolor, :seriesalpha,
    :smooth,
    :xerror, :yerror,
    :subplot,
    :x, :y, :z,
    :show, :size,
    :margin,
    :left_margin,
    :right_margin,
    :top_margin,
    :bottom_margin,
    :html_output_format,
    :layout,
    :link,
    :primary,
    :series_annotations,
    :subplot_index,
    :discrete_values,
    :projection,

]

function merge_with_base_supported(v::AVec)
    v = vcat(v, _base_supported_args)
    for vi in v
        if haskey(_axis_defaults, vi)
            for letter in (:x,:y,:z)
                push!(v, Symbol(letter,vi))
            end
        end
    end
    v
end


text:
function _up_push_dummy(gl::GitLink)
    dir = global_state_dir(gl)
    mkpath(dir)
    dummy = joinpath(dir, ".gl-push-dummy")
    write(dummy, rand_str())
end

function soft_push(gl::GitLink; verbose = false, commit_msg = "Up at $(time())", tries = 1)
    for t in 1:max(tries, 1)
        # repo info
        gl_repo = repo_dir(gl)
        url = remote_url(gl)

        # check remote
        rhash0 = _remote_HEAD_hash(url)
        isempty(rhash0) && continue

        # enforce repo format (lazy method)
        _format_repo!(gl_repo, url; verbose)

        # check it is sync
        chash0 = _HEAD_hash(gl_repo)
        (rhash0 != chash0) && continue

        # check repodir
        !_check_gitdir(gl_repo) && continue

        # soft push
        _up_push_dummy(gl) # Update push dummy (always push)
        _rm(joinpath(gl_repo, ".gitignore")) # avoid interference
        _run("git -C $(gl_repo) add -A 2>&1"; verbose)
        _run("git -C $(gl_repo) status 2>&1"; verbose)
        user_name = get_global_config("user.name", "GitLink")
        user_email = get_global_config("user.email", "fake@email.com")
        _run("git -C $(gl_repo) -c user.name='$(user_name)' -c user.email='$(user_email)' commit -am '$(commit_msg)' 2>&1"; verbose)
        _run("git -C $(gl_repo) push 2>&1"; verbose)

        # check success
        rhash1 = _remote_HEAD_hash(url)
        chash1 = _HEAD_hash(gl_repo)
        if rhash0 == rhash1 && chash0 == chash1
            _rm(gl_repo) # something fail
            continue
        end
        return true
    end
    return false
end


text:
"""
    plot(net::HybridNetwork, method::Symbol)

Plot a network using R graphics.
`method` should be `:R` (actually, any symbol would do, for now!).

optional arguments, shared with the Gadfly-based plot function:
- useEdgeLength: if true, the tree edges and major hybrid edges are
  drawn proportionally to their length. Minor hybrid edges are not, however.
  Note that edge lengths in coalescent units may scale very poorly with time.
- showTipLabel: if true, taxon labels are shown. You may need to zoom out to see them.
- showNodeNumber: if true, nodes are labelled with the number used internally.
- showEdgeLength: if true, edges are labelled with their length (above)
- showGamma: if true, hybrid edges are labelled with their heritability (below)
- edgeColor: color for tree edges. black by default.
- majorHybridEdgeColor: color for major hybrid edges
- minorHybridEdgeColor: color for minor hybrid edges
- showEdgeNumber: if true, edges are labelled with the number used internally.
- showIntNodeLabel: if true, internal nodes are labelled with their names.
  Useful for hybrid nodes, which do have tags like 'H1'.
- edgeLabel: dataframe with two columns: the first with edge numbers, the second with labels
  (like bootstrap values) to annotate edges. empty by default.
- nodeLabel: dataframe with two columns: the first with node numbers, the second with labels
  (like bootstrap values for hybrid relationships) to annotate nodes. empty by default.

optional arguments specific to this function:
- xlim, ylim: array of 2 values
- tipOffset: to offset tip labels

Note that `plot` actually modifies some (minor) attributes of the network,
as it calls `directEdges!`, `preorder!` and `cladewiseorder!`.

If hybrid edges cross tree and major edges, you may choose to rotate some tree
edges to eliminate crossing edges, using `rotate!`
(in [`PhyloNetworks`](http://crsl4.github.io/PhyloNetworks.jl/latest/lib/public/#PhyloNetworks.rotate!)).

**Alternative**: a tree or network can be exported with [`sexp`](@ref)
and then displayed with R's "plot" and all its options.
"""
function plot(net::HybridNetwork, method::Symbol; useEdgeLength=false::Bool,
    mainTree=false::Bool, showTipLabel=true::Bool, showNodeNumber=false::Bool,
    showEdgeLength=false::Bool, showGamma=false::Bool,
    edgeColor="black"::String,
    majorHybridEdgeColor="deepskyblue4"::String,
    minorHybridEdgeColor="deepskyblue"::String,
    showEdgeNumber=false::Bool, showIntNodeLabel=true::Bool,
    edgeLabel=DataFrame()::DataFrame, nodeLabel=DataFrame()::DataFrame,
    xlim=Float64[]::Array{Float64,1}, ylim=Float64[]::Array{Float64,1},
    tipOffset=0.0::Float64, tipcex=1.0::Float64)

    (edge_xB, edge_xE, edge_yB, edge_yE, node_x, node_y, node_yB, node_yE,
     xmin, xmax, ymin, ymax) = getEdgeNodeCoordinates(net, useEdgeLength)
    labelnodes, nodeLabel = checkNodeDataFrame(net, nodeLabel)
    ndf = prepareNodeDataFrame(net, nodeLabel, showNodeNumber,
            showIntNodeLabel, labelnodes, node_x, node_y)
    if (showTipLabel || showNodeNumber || showIntNodeLabel || labelnodes)
        expfac = 0.1  # force 10% more space to show tip/node/root name
        expfacy = 0.5 # additive expansion for y axis
        xmin -= (xmax-xmin)*expfac
        xmax += (xmax-xmin)*expfac
        ymin -= expfacy
        ymax += expfacy
    end
    xmax += tipOffset
    if length(xlim)==2
        xmin=xlim[1]; xmax=xlim[2]
    end
    if length(ylim)==2
        ymin=ylim[1]; ymax=ylim[2]
    end
    leaves = [n.leaf for n in net.node]
    eCol = fill(edgeColor, length(net.edge))
    eCol[ [ e.hybrid  for e in net.edge] ] .= majorHybridEdgeColor
    eCol[ [!e.isMajor for e in net.edge] ] .= minorHybridEdgeColor

    R"""
    plot($(node_x[leaves]), $(node_y[leaves]), type='n',
         xlim=c($xmin,$xmax), ylim=c($ymin,$ymax),
         axes=FALSE, xlab='', ylab='')
    segments($edge_xB, $edge_yB, $edge_xE, $edge_yE, col=$eCol)
    segments($node_x, $node_yB, $node_x, $node_yE, col=$edgeColor)
    """
    if showTipLabel
      R"text"(node_x[leaves] .+ tipOffset, node_y[leaves],
              tipLabels(net), adj=0, font=3, cex=tipcex)
    end
    if showIntNodeLabel
      R"text"(ndf[.!ndf[!,:lea],:x], ndf[.!ndf[!,:lea],:y],
              ndf[.!ndf[!,:lea],:name], font=3, cex=tipcex, adj=[.5,0])
    end
    if showNodeNumber
      R"text"(ndf[!,:x], ndf[!,:y], ndf[!,:num], adj=1)
    end
    if labelnodes
      R"text"(ndf[!,:x], ndf[!,:y], ndf[!,:lab], adj=1)
    end
    labeledges, edf = prepareEdgeDataFrame(net, edgeLabel, mainTree,
                        edge_xB, edge_xE, edge_yB, edge_yE)
    if labeledges
      R"text"(edf[!,:x], edf[!,:y], edf[!,:lab], adj=[.5,0])
    end
    if showEdgeLength
      R"text"(edf[!,:x], edf[!,:y], edf[!,:len], adj=[.5,1.])
    end
    if (showGamma && net.numHybrids>0)
      im = edf[!,:hyb] .& edf[!,:min]
      iM = edf[!,:hyb] .& .!edf[!,:min]
      R"text"(edf[im,:x], edf[im,:y], edf[im,:gam],
              adj=[.5,1], col=minorHybridEdgeColor)
      R"text"(edf[iM,:x], edf[iM,:y], edf[iM,:gam],
              adj=[.5,1], col=majorHybridEdgeColor)
    end
    if showEdgeNumber
      R"text"(edf[!,:x], edf[!,:y], edf[!,:num], adj=[.5,0])
    end
    return (xmin, xmax, ymin, ymax, node_x, node_y, node_yB, node_yE,
      edge_xB, edge_xE, edge_yB, edge_yE, ndf, edf)
end


text:
# from NLPModels

export DimensionError, @lencheck

struct DimensionError <: Exception
  name :: Union{Symbol,String}
  dim_expected :: Int
  dim_found :: Int
end

function Base.showerror(io::IO, e::DimensionError)
  print(io, "DimensionError: Input $(e.name) should have length $(e.dim_expected) not $(e.dim_found)")
end

# Check that arrays have a prescribed size.
# https://groups.google.com/forum/?fromgroups=#!topic/julia-users/b6RbQ2amKzg
macro lencheck(l, vars...)
  exprs = Expr[]
  for var in vars
    varname = string(var)
    push!(exprs,
          :(if length($(esc(var))) != $(esc(l))
                throw(DimensionError($varname, $(esc(l)), length($(esc(var)))))
            end))
  end
  Expr(:block, exprs...)
end

macro rangecheck(lo, hi, vars...)
  exprs = Expr[]
  for var in vars
    varname = string(var)
    push!(exprs,
          :(if (length($(esc(var))) > 0 && (any(broadcast(<, $(esc(var)), $(esc(lo)))) || any(broadcast(>, $(esc(var)), $(esc(hi))))))
            error(string($varname, " elements must be between ", $(esc(lo)), " and ", $(esc(hi))))
            end))
  end
  Expr(:block, exprs...)
end


text:
@testset "Transformed cubes" begin
    @testset "Simple map" begin
        data = collect(reshape(1:20, 4, 5))
        axlist = [RangeAxis("XVals", 1.0:4.0), CategoricalAxis("YVals", [1, 2, 3, 4, 5])]
        props = Dict("att1" => 5, "att2" => "Hallo")
        a = YAXArray(axlist, data, props)
        a2 = map(i -> i + 1, a)
        @test all(a2.data .== a.data .+ 1)
        @test a2.axes == a.axes
        @test a2.properties == a.properties
        #TODO: Test that cleaners get promoted properly
        a3 = map(+, a, a2)
        @test all(a3.data .== a.data .+ a2.data)
        @test a3.axes == a.axes
        @test a3.properties == a.properties
    end
    @testset "Cube concatenation" begin
        data = [rand(4, 5) for i = 1:3]
        axlist = [RangeAxis("XVals", 1.0:4.0), CategoricalAxis("YVals", [1, 2, 3, 4, 5])]
        props = [Dict("att$i" => i) for i = 1:3]
        singlecubes = [YAXArray(axlist, data[i], props[i]) for i = 1:3]
        newcube = concatenatecubes(singlecubes, CategoricalAxis("ZVals", ["A", "B", "C"]))
        @test caxes(newcube) == [axlist; CategoricalAxis("ZVals", ["A", "B", "C"])]
        @test ndims(newcube) == 3
        @test size(newcube) == (4, 5, 3)
        @test newcube[:, :, :] == cat(data..., dims = 3)
        @test getattributes(newcube) == reduce(merge, props)
    end
end

max_forks_repo_name: meggart/YAXArrays


text:
using JLD2
using Distributions
using StatsBase
using LinearAlgebra
using SparseArrays

"""
    COOccurence(coo_matrix, labels)
"""
mutable struct COOccurence
    coo_matrix::SparseMatrixCSC{Float64,Int64}
    ind2labels::Dict{Int, String}
end

# Note: Sample input occurrance files are those generated with PubMedMiner.jl
# https://github.com/bcbi/PubMedMiner.jl.git
function occur2coo(occur, label2ind)
    #invert map
    ind2label = Dict()
    for key  in keys(label2ind)
        val = label2ind[key]
        ind2label[val] = key
    end
    # labels = Array{String}(length(ind2label))
    # for i in range(1,length(ind2label))
    #     labels[i] = ind2label[i]
    # end

    coo_m = occur * transpose(occur)
    COO = COOccurence(coo_m, ind2label)
    return COO
end

function read_sp_occur(file="./occur_sp.jdl")
    fin  = jldopen(file, "r")
    occur = read(fin, "occur")
    return occur
end

function read_occur_label_dict(file="/mesh2ind.jdl")
    file  = jldopen(file, "r")
    label2ind = read(file, "mesh2ind")
    return label2ind
end

"""
    corrcoef(occur)

Compute correlation coeffiecients for a occurence data matrix
###Arguments

* `occur`: Occurrence data matrix
"""
function corrcoef(occur; vardim = 1)
    # compute the covariance matrix
    C = cov(occur, dims=vardim)

    # the correlation coefficients are given by
    # C_{i,j} / sqrt(C_{ii} * C_{jj})
    d = diag(C)
    CiiCjj = zeros(size(C))
    BLAS.ger!(1.0, d, d, CiiCjj) #CiiCjj += d*d'

    #should we devectorize?
    coeffs = C./sqrt.(CiiCjj)

    return coeffs

end

"""
    pmi_mat(coo_matrix, N)

Point-wise mutual information from co-occurrence  matrix

Point-wise mutual information
log p(x,y)/p(x)p(y)

A measure the takes into account chance by comparing the probability of
joint event to the product of the individual events

###Outcome

* `pmi`:PMI matrix (symmetric matrix)
"""
function pmi_mat(coo_matrix, N)
    d = diag(coo_matrix)

    #Compute conditionals P(X,Y)/P(X)
    cooccurrence_diagonal = diag(coo_matrix)
    conditional = coo_matrix./cooccurrence_diagonal

    # PMI (point-wise mutual information) P(X,Y)/P(X)P(Y)
    pmi_full = log.((N .* conditional') ./ cooccurrence_diagonal)
    # pmi = LowerTriangular(pmi_full)
end

"""
    chi2(coo_matrix)

Point-wise Chi square statistic from co-occurrence  matrix

Chi square statistic to test whether the occurence of one
feature is independent of another, e.i
Null hypothesis = p(x,y) = p(x)p(y)

###Arguments
* `occur`: Occurrance data matrix: Observations are rows, columns are features

###Outcome

* `chi2`:Matrix of Chi Square Statistics  (symmetric matrix)
"""
function chi2_mat(occur; min_freq=5)

    nvars = size(occur, 2)
    chi2 = LowerTriangular(zeros(nvars, nvars))
    for j1=1:nvars
        x = occur[:,j1]
        for j2=j1+1:nvars
            y = occur[:,j2]
            chi2[j2, j1] = chi2_statistic(x,y,span(occur); min_freq=min_freq)
        end
    end
    return chi2
end

max_stars_repo_name: bcbi/BCBIStats

max_issues_repo_name: bcbi/BCBIStats

max_forks_repo_name: bcbi/BCBIStats


text:
module Model

using LightGraphs
using GraphIO

using DataStructures
using SparseArrays
using Optim
using LineSearches
using StatsBase
using Statistics
using JSON
using Dates
using Serialization
import Optim: converged

include("./structures.jl")
include("./nodemap.jl")

setprecision(BigFloat, 64)

mutable struct GraphModel
    dataset::String # constant
    structure_alphabet_size::Int64 # constant
    G::LightGraphs.SimpleGraph{Int64} # constant
    A::SparseArrays.SparseMatrixCSC{Int64,Int64} # constant
    n::Int64 # constant
    m::Int64 # constant
    maximum_m_simple::Int64 # constant
    maximum_m_loopy::Int64 # constant
    lambdas::Array{BigFloat,1} # growing at end (and values changing by optimization)
    constraint_values::Array{BigFloat,1} # growing at end
    micro_structures::Array{Union{Bool, Set{Tuple{Int64,Int64}}},1} # growing at end
    edge_lambda_indices::SortedDict{Tuple{Int64,Int64},Array{Int64,1}} # new edges inserted, growing at selected value ends
    edge_equivalence_classes::Dict{Set{Int64},Array{Tuple{Int64,Int64},1}} # currently reconstructed but could be smarter
    backup_optimization_result::Optim.MultivariateOptimizationResults # reconstructed
    last_optimization_result::Optim.MultivariateOptimizationResults # reconstructed
    entropy::BigFloat # changing
    equivalence_class_edge_probabilities::Dict{Set{Int64},BigFloat} # reconstructed
    surprise_matrix::SparseArrays.SparseMatrixCSC{BigFloat,Int64} # reconstructed
    surprise_matrix_transpose::SparseArrays.SparseMatrixCSC{BigFloat,Int64} # reconstructed
    macro_structures::Array{Structures.Structure} # growing at end
    macro_structure_description_lengths::Array{BigFloat} # growing at end
    edge_probability_description_length_maxent::BigFloat # changing
    structure_type_description_length::BigFloat # changing
    node_surprises::Array{Tuple{Int64,BigFloat},1} # reconstructed
    surprising_edges::Array{Tuple{Int64,Int64},1} # reconstructed
    start_description_length_maxent::BigFloat # constant
    total_description_length_maxent::BigFloat # changing
    description_lengths_over_time::Array{BigFloat,1} # growing at end
    parameters::Dict{String,Float64} # constant resp. set only upon saving
    io::IO # constant
    hub_spoke_edges::Set{Tuple{Int64,Int64}} # changing
    star_spoke_edges::Set{Tuple{Int64,Int64}} # changing
    zero_edges::Set{Tuple{Int64,Int64}} # changing
    last_hub_spoke_edges::Set{Tuple{Int64,Int64}} # changing
    last_star_spoke_edges::Set{Tuple{Int64,Int64}} # changing
    last_zero_edges::Set{Tuple{Int64,Int64}} # changing - note that we'll be a bit imprecise here: if a zero edge is discovered after already having been covered, it will remain also in its previous equivalence class, if any
    covered_nodes::Set{Int64} # growing
    function GraphModel(dataset::String, io=stdout)
        setup_graph!(new(), dataset, io)
    end
    function GraphModel(kwargs::Dict)
        M = new()
        for (k,v) in kwargs
            setproperty!(M, k, v)
        end
        return M
    end
end

Base.show(io::IO, M::GraphModel) = print(io, "Model for $(M.dataset):\n$(typeof(M.G)) with n = $(M.n), m=$(M.m)")
JSON.lower(M::GraphModel) = Dict(key=> getfield(M, key) for key in fieldnames(GraphModel) if key in
            (:dataset, :n, :m, :macro_structures, :macro_structure_description_lengths,
                :edge_probability_description_length_maxent,
                :start_description_length_maxent,
                :total_description_length_maxent,
                :description_lengths_over_time,
                :lambdas,
                :parameters,
                :structure_alphabet_size,
                :structure_type_description_length
            )
)

"""
helper providing basic logging
"""
function log_progress!(M::GraphModel, msg::String)
    println(M.io, now(), " - ", msg)
    flush(M.io)
end

"""
save model as jldb (serialization format) or json
"""
function save_model(M::GraphModel, directory::String; final=true, overwrite=true)
    path, file = splitdir(M.dataset)
    last_folder = splitdir(path)[end]
    file_base, ext = splitext(file)
    size_threshold = M.parameters["size_threshold"]
    stop_after_n_structures = M.parameters["stop_after_n_structures"]
    filepath = "$(directory)/$(last_folder)/$(file_base)_size-$(size_threshold)_max-$(stop_after_n_structures)"
    mkpath(splitdir(filepath)[1])
    if !final
        open(filepath*".jldb", "w") do io
            serialize(io, M)
        end
    else
        open(filepath*".json", "w") do io
            write(io, json(M, 4))
        end
        log_progress!(M, "Saved model to $(filepath).json")
    end
    mv("$(path)/$(file_base)-nodemap.csv", "$(directory)/$(last_folder)/$(file_base)-nodemap.csv", force=overwrite)
end

"""
restore model from json
"""
function restore_model(path_to_file::String)
    model_dict = JSON.parsefile(path_to_file)
    M = Model.GraphModel(model_dict["dataset"])
    structures = Array{Structures.Structure,1}()
    for sd in model_dict["macro_structures"]
        if sd["structure_type"] == "clique"
            structure = Structures.Clique(sd["n_edges_total"], Array{Int64,1}(sd["nodes"]))
        elseif sd["structure_type"] == "star"
            structure = Structures.Star(sd["n_edges_total"], sd["hub"], Array{Int64,1}(sd["spokes"]))
        elseif sd["structure_type"] == "biclique"
            structure = Structures.Biclique(sd["n_edges_total"], Array{Int64,1}(sd["left_nodes"]), Array{Int64,1}(sd["right_nodes"]), sd["n_edges_in_left"], sd["n_edges_in_right"])
        elseif sd["structure_type"] == "starclique"
            structure = Structures.Starclique(sd["n_edges_total"], Array{Int64,1}(sd["left_nodes"]), Array{Int64,1}(sd["right_nodes"]), sd["n_edges_in_left"], sd["n_edges_in_right"])
        # generic structures are already present
        else
            continue
        end
        push!(structures, structure)
    end
    add_macro_structures!(M, structures; lambdas=get!(model_dict, "lambdas", []))
    if model_dict["lambdas"] == []
        println("WARNING: Recomputing lambdas when restoring model - true edge probability description length for computed lambdas might differ from restored edge probability description length due to different optimization process...")
    end
    M.parameters["surprise_threshold"] = model_dict["parameters"]["surprise_threshold"]  === nothing ? Inf : model_dict["parameters"]["surprise_threshold"]
    M.parameters["size_threshold"] = model_dict["parameters"]["size_threshold"]  === nothing ? Inf : model_dict["parameters"]["size_threshold"]
    M.description_lengths_over_time = model_dict["description_lengths_over_time"]
    M.macro_structure_description_lengths = model_dict["macro_structure_description_lengths"]
    M.edge_probability_description_length_maxent = model_dict["edge_probability_description_length_maxent"]
    M.total_description_length_maxent = M.edge_probability_description_length_maxent + M.structure_type_description_length + sum(M.macro_structure_description_lengths)
    return M
end

"""
initial graph setup with first optimization
"""
function setup_graph!(M::GraphModel, dataset::String, io::IO)
    set_constant_terms!(M, dataset, io)
    M.lambdas = zeros(2)
    M.constraint_values = [M.m, 0]
    M.micro_structures = [true, Set([(i,i) for i=1:M.n])]
    M.edge_lambda_indices = SortedDict((i,i)=>[1,2] for i=1:M.n)
    M.edge_equivalence_classes = Dict(Set([1,2])=>[(i,i) for i=1:M.n])
    M.macro_structures = [Structures.GenericStructure("Overall number of edges", M.m), Structures.GenericStructure("Loop-freeness", 0)]
    # below, we add the cost of transmitting n to the cost of transmitting m (since it is not a proper structure but necessary to specify m)
    M.macro_structure_description_lengths = [Structures.MDL.universal_integer(M.n) + Structures.MDL.log2(M.maximum_m_loopy), 1.]
    M.structure_type_description_length = 0
    M.description_lengths_over_time = []
    M.star_spoke_edges = Set()
    M.hub_spoke_edges = Set()
    M.zero_edges = Set()
    M.last_star_spoke_edges = Set()
    M.last_hub_spoke_edges = Set()
    M.last_zero_edges = Set()
    M.covered_nodes = Set()
    update_variable_terms!(M)
    M.backup_optimization_result = deepcopy(M.last_optimization_result)
    M.start_description_length_maxent = M.total_description_length_maxent
    # the below is to account for the fact that we have two initial constraints but perform only one optimization
    append!(M.description_lengths_over_time, [M.start_description_length_maxent])
    M.parameters = Dict("surprise_threshold" => Inf, "size_threshold" => Inf)
    return M
end

"""
initial graph setup: set constant terms - NB: dataset now needs to be full path with extension
"""
function set_constant_terms!(M::GraphModel, dataset::String, io::IO)
    M.dataset = dataset
    M.structure_alphabet_size = 4 # cliques, stars, bicliques, starcliques
    M.io = io
    M.G = SimpleGraph(NodeMap.loadgraph(dataset, GraphIO.EdgeList.EdgeListFormat())[1])
    NodeMap.create_node_map(dataset)
    remove_loops!(M.G)
    M.A = LightGraphs.LinAlg.adjacency_matrix(M.G)
    M.n = nv(M.G)
    M.m = ne(M.G)
    M.maximum_m_simple = Structures.MDL.choose(M.n, 2)
    M.maximum_m_loopy = M.maximum_m_simple + M.n
end

"""
make graph loop-free
"""
function remove_loops!(G::SimpleGraph)
    for e in [e for e in edges(G) if src(e) == dst(e)]
        rem_edge!(G, e)
    end
end

"""
each step: update model (re-optimize, update description length, update surprises)
"""
function update_variable_terms!(M::GraphModel)
    log_progress!(M, "starting optimization")
    optimize_maxent!(M)
    log_progress!(M, "finished optimization")
    log_progress!(M, "updating description length")
    update_description_length!(M)
end

"""
optimization wrapper function
"""
function optimize_maxent!(M::Model.GraphModel)
    n_iter = 50
    td = TwiceDifferentiable(x -> lagrange_dual_equiv(x, M), M.lambdas; autodiff = :forward)
        result = optimize(
        td, M.lambdas,
        Newton(linesearch=LineSearches.BackTracking()), # the solver - using BackTracking instead of the default HagerZhang because that sometimes throws a weird assertion error
        Optim.Options(allow_f_increases=true,
            show_trace=false, show_every=100, iterations=n_iter, x_tol=1e-8, f_tol=1e-8, g_tol=1e-6)
    )
    M.last_optimization_result = result
    M.lambdas = M.last_optimization_result.minimizer
    M.entropy = M.last_optimization_result.minimum
    if !converged(result)
        log_progress!(M, "NO CONVERGENCE AFTER $(n_iter) ITERATIONS")
    end
end

"""
optimization computation: wrapper
"""
function lagrange_dual_equiv(lambdas, M::Model.GraphModel)
    return (get_first_term(lambdas, M) - get_second_term(lambdas, M))
end

"""
optimization computation: first term
"""
function get_first_term(lambdas, M::Model.GraphModel)
    #log_progress!(M, "computing first term in optimization")
    rest_lambda = lambdas[1]
    if isempty(M.edge_equivalence_classes)
        return M.maximum_m_loopy * log(get_normalizer(rest_lambda))
    else
        lambda_total_contribution = BigFloat(0)
        lambda_total_sizes = 0
        for equivalence_class_key in keys(M.edge_equivalence_classes)
            size = get_equiv_size(equivalence_class_key, M.edge_equivalence_classes)
            lambda_sum = get_lambda_sum_from_equiv(equivalence_class_key, lambdas)
            normalized_sum = get_normalizer(lambda_sum)
            contribution = log(normalized_sum)
            lambda_total_contribution += size * contribution
            lambda_total_sizes += size
        end
        rest_class_n = M.maximum_m_loopy - sum(lambda_total_sizes)
        rest_contribution = rest_class_n * log(get_normalizer(rest_lambda))
        return lambda_total_contribution + rest_contribution
    end
end

"""
optimization helper (normalization)
"""
function get_normalizer(lambda_sum)
    return (1 + exp(lambda_sum))
end

"""
optimization helper (size of equivalence class)
"""
function get_equiv_size(equivalence_class_key::Set{Int64}, equivalence_classes::Dict{Set{Int64},Array{Tuple{Int64,Int64},1}})
    return length(equivalence_classes[equivalence_class_key])
end

"""
optimization helper (lambda sum for equivalence class)
"""
function get_lambda_sum_from_equiv(equivalence_class_key::Set{Int64}, lambdas)
    lambda_sum = BigFloat(0)
    for idx in equivalence_class_key
        lambda_sum += lambdas[idx]
    end
    return lambda_sum
end

"""
optimization computation: second term
"""
function get_second_term(lambdas, M::Model.GraphModel)
    #log_progress!(M, "computing second term in optimization")
    second_term = BigFloat(0)
    for (lambda, constraint_value) in zip(lambdas, M.constraint_values)
        second_term += lambda * constraint_value
    end
    return second_term
end

"""
description length computation wrapper
"""
function update_description_length!(M::Model.GraphModel)
    update_equivalence_class_edge_probabilities!(M)
    update_edge_probability_description_length!(M)
    update_structure_type_description_length!(M)
    update_total_description_length!(M)
    # NB: the below gives misleading results when multiple structures are added at a time, e.g., during replay -> we have taken care of that in restore_model
    append!(M.description_lengths_over_time, M.total_description_length_maxent)
end

"""
description length helper: structure type encoding costs
"""
function update_structure_type_description_length!(M::Model.GraphModel)
    n_different_structures = length(M.macro_structures) - 2 # we don't count the generic macro structures as their type is equivalence_class_edge_probabilities
    M.structure_type_description_length = (Structures.MDL.universal_integer(n_different_structures + 1) + Structures.MDL.log2_choose(M.structure_alphabet_size + n_different_structures - 1, n_different_structures - 1)
    )
end

"""
description length helper: edge probabilities for equivalence classes
"""
function update_equivalence_class_edge_probabilities!(M::Model.GraphModel)
    probabilities = Dict{Array{Int64,1},BigFloat}(
        equivalence_class=>edge_probability_for_equiv(equivalence_class, M.lambdas) for equivalence_class in keys(M.edge_equivalence_classes))
    probabilities[[1,]] = edge_probability_for_equiv(Set([1,]), M.lambdas)
    M.equivalence_class_edge_probabilities = probabilities
end

"""
description length helper: edge probability for one equivalence class
"""
function edge_probability_for_equiv(equivalence_class, lambdas)
    if equivalence_class == Set([1,])
        lambda_sum = lambdas[1]
    else
        lambda_sum = get_lambda_sum_from_equiv(equivalence_class, lambdas)
    end
    edge_prop = exp(lambda_sum) / get_normalizer(lambda_sum)
    return edge_prop
end

"""
description length helper: surprise matrices update
"""
function update_surprise_matrices!(M::Model.GraphModel)
    # NB: we currently keep track only of surprising EDGES (else we get a non-sparse matrix :( )
    surprise_matrix = SparseMatrixCSC{BigFloat,Int64}(M.A)
    for idx in filter!(x->(x[1] < x[2]), findall(isone, M.A))
        edge = (idx[1],idx[2])
        # due to symmetry
        eq_class = Set(get(M.edge_lambda_indices, edge, [1,]))
        rest_class_probability = M.equivalence_class_edge_probabilities[Set([1,])]
        if eq_class == Set([1,])
            surprise_matrix[edge...] = 1 - rest_class_probability
        else
            surprise_matrix[edge...] = 0
        end
        surprise_matrix[reverse(edge)...] = 0
    end
    M.surprise_matrix = surprise_matrix
    M.surprise_matrix_transpose = sparse(transpose(M.surprise_matrix))
end

"""
description length helper: total description length of edge probabilities (error in MDL, L(D|M))
"""
function update_edge_probability_description_length!(M::Model.GraphModel)
    dl = 0.
    rest_size = M.maximum_m_loopy # because we have loop-freeness as a separate constraint
    n_guaranteed_edges = length(M.hub_spoke_edges) # guaranteed to be there and not appearing in any equivalence class
    n_zero_edges = length(M.zero_edges) # guaranteed to be not there and not appearing in any equivalence class
    rest_size -=  (n_guaranteed_edges + n_zero_edges)
    rest_edges = M.m - n_guaranteed_edges
    for (eq_class, edges) in M.edge_equivalence_classes
        size = length(edges)
        size_ones = length(filter(edge -> has_edge(M.G, edge...), edges))
        size_zeros = size - size_ones
        eq_class_probability = M.equivalence_class_edge_probabilities[eq_class]
        contribution = size_ones * -Structures.MDL.log2_zero(eq_class_probability) + size_zeros * -Structures.MDL.log2_zero(1 - eq_class_probability)
        dl += contribution
        rest_size -= size
        rest_edges -= size_ones
    end
    rest_probability = M.equivalence_class_edge_probabilities[Set([1,])]
    rest_contribution = rest_edges * -Structures.MDL.log2_zero(rest_probability) + (rest_size - rest_edges) * -Structures.MDL.log2_zero(1 - rest_probability)
    dl += rest_contribution
    M.edge_probability_description_length_maxent = dl
end

"""
description length helper: total description length (L(M) + L(D|M))
"""
function update_total_description_length!(M::Model.GraphModel)
    M.total_description_length_maxent = sum(M.macro_structure_description_lengths) + M.edge_probability_description_length_maxent + M.structure_type_description_length
end

"""
sanity check whether probability mass approximately sums to 1
"""
function sanity_check(M::Model.GraphModel)
    if !isempty(M.edge_equivalence_classes)
        equiv = sum([edge_probability_for_equiv(equivalence_class, M.lambdas) * get_equiv_size(equivalence_class, M.edge_equivalence_classes)
            for equivalence_class in keys(M.edge_equivalence_classes)])
        rest_size = M.maximum_m_loopy - sum([get_equiv_size(equivalence_class, M.edge_equivalence_classes)
            for equivalence_class in keys(M.edge_equivalence_classes)])
    else
        equiv = 0
        rest_size = M.maximum_m_loopy
    end
    rest = edge_probability_for_equiv([1,], M.lambdas)
    probability_mass_sum = equiv + rest * rest_size
    print("Sanity Check: probability mass sums to $(probability_mass_sum), expected $(M.m)\n")
    return abs(probability_mass_sum-M.m) <= 1.0e-2
end

"""
generate micro structures and constraints from clique candidate
"""
function generate_micro_structures(structure::Structures.Clique, M::Model.GraphModel)
    micro_structures = [Set([(n1,n2) for n1 in structure.nodes for n2 in structure.nodes if n1 < n2])]
    clique_edges_in_hs_edges = intersect(micro_structures[1], M.hub_spoke_edges)
    setdiff!(micro_structures[1], clique_edges_in_hs_edges)
    constraint_values = [structure.n_edges_total - length(clique_edges_in_hs_edges)]
    return micro_structures, constraint_values
end

"""
generate micro structures and constraints from star candidate
"""
function generate_micro_structures(structure::Structures.Star, M::Model.GraphModel)
    hub = structure.hub
    micro_structures = [Set([(hub < n ? (hub,n) : (n, hub)) for n in structure.spokes]),
        Set([(n1,n2) for n1 in structure.spokes for n2 in structure.spokes if (n1 < n2) && (n1,n2) ∉ M.star_spoke_edges])]
    hs_edges_in_hs_edges = intersect(micro_structures[1], M.hub_spoke_edges)
    setdiff!(micro_structures[1],hs_edges_in_hs_edges)
    ss_edges_in_hs_edges = intersect(micro_structures[end], M.hub_spoke_edges)
    setdiff!(micro_structures[end],ss_edges_in_hs_edges)
    constraint_values = [structure.n_nodes_in_spokes-length(hs_edges_in_hs_edges), structure.n_edges_in_spokes-length(ss_edges_in_hs_edges)] # note that this is theoretically wrong - we do it for performance reasons
    return micro_structures, constraint_values
end

"""
generate micro structures and constraints from biclique candidate
"""
function generate_micro_structures(structure::Union{Structures.Biclique,Structures.Starclique}, M::Model.GraphModel)
    if structure isa Structures.Biclique
        left = Set([(n1, n2) for n1 in structure.left_nodes for n2 in structure.left_nodes if n1 < n2 && (n1,n2) ∉ M.star_spoke_edges])
    else # assuming Starclique
        left = Set([(n1, n2) for n1 in structure.left_nodes for n2 in structure.left_nodes if n1 < n2])
    end
    left_edges_in_hs_edges = intersect(left, M.hub_spoke_edges)
    setdiff!(left,left_edges_in_hs_edges)
    right = Set([(n1, n2) for n1 in structure.right_nodes for n2 in structure.right_nodes if n1 < n2 && (n1,n2) ∉ M.star_spoke_edges])
    right_edges_in_hs_edges = intersect(right, M.hub_spoke_edges)
    setdiff!(right,right_edges_in_hs_edges)
    across = Set([(n1 < n2) ? (n1, n2) : (n2, n1) for n1 in structure.left_nodes for n2 in structure.right_nodes])
    across_edges_in_hs_edges = intersect(across, M.hub_spoke_edges)
    setdiff!(across, across_edges_in_hs_edges)
    micro_structures = [left, right, across]
    constraint_values = [   structure.n_edges_in_left - length(left_edges_in_hs_edges),
                            structure.n_edges_in_right - length(right_edges_in_hs_edges),
                            structure.n_edges_across - length(across_edges_in_hs_edges)
                        ]
    return micro_structures, constraint_values
end

"""
try adding structure to model, keeping the result only if it reduces the description length
"""
function test_adding_structure!(M::Model.GraphModel, structure::Structures.Structure, size_threshold::Float64=1.)
    log_progress!(M, "saving previous state...")
    M.backup_optimization_result = deepcopy(M.last_optimization_result)
    log_progress!(M, "testing structure $(structure)...")
    if structure.n_nodes_total < size_threshold
        log_progress!(M, "failure: structure too small")
        return false
    end
    add_macro_structures!(M, Array{Structures.Structure,1}([structure]))
    if converged(M.last_optimization_result) && M.description_lengths_over_time[end] <= M.description_lengths_over_time[end-1]
        log_progress!(M, "success: new description length $(M.total_description_length_maxent)")
        return true
    else
        roll_back_model!(M)
        log_progress!(M, "failure: no decrease in description length")
        return false
    end
end

"""
add macro structures to model (assuming no connector knowledge)
we allow for adding multiple structures at once to enable faster replay of macro_structure_description_lengths
(in the model building phase, we test adding one structure at a time only)
"""
function add_macro_structures!(M::Model.GraphModel, structures::Array{Structures.Structure,1}; lambdas=[])
    log_progress!(M, "updating micro structures et al")
    for structure in structures
        micro_structures, constraint_values = generate_micro_structures(structure, M)
        if structure isa Structures.Star
            M.last_hub_spoke_edges = deepcopy(M.hub_spoke_edges)
            M.last_star_spoke_edges = deepcopy(M.star_spoke_edges)
            n_edges_in_spokes = structure.n_edges_in_spokes
            union!(M.star_spoke_edges, micro_structures[end])
            union!(M.hub_spoke_edges, micro_structures[1])
            if n_edges_in_spokes > 0 # if the star is perfect, we know exactly what its associated probabilities are
                append!(M.micro_structures, [micro_structures[end]])
                append!(M.constraint_values, [constraint_values[end]])
                append!(M.lambdas, zeros(1))
                update_edge_lambda_indices!(M, [micro_structures[end]], [constraint_values[end]])
                update_equivalence_classes_from_edge_lambda_indices!(M)
            else
                M.last_zero_edges = deepcopy(M.last_zero_edges)
                union!(M.zero_edges, micro_structures[1])
            end
        elseif structure isa Structures.Starclique || structure isa Structures.Biclique
            M.last_star_spoke_edges = deepcopy(M.star_spoke_edges)
            union!(M.star_spoke_edges, micro_structures[2])
            if structure isa Structures.Biclique
                union!(M.star_spoke_edges, micro_structures[1])
            end
            append!(M.micro_structures, micro_structures)
            append!(M.constraint_values, constraint_values)
            append!(M.lambdas, zeros(length(micro_structures)))
            update_edge_lambda_indices!(M, micro_structures, constraint_values)
            update_equivalence_classes_from_edge_lambda_indices!(M)
        else
            append!(M.micro_structures, micro_structures)
            append!(M.constraint_values, constraint_values)
            append!(M.lambdas, zeros(length(micro_structures)))
            update_edge_lambda_indices!(M, micro_structures, constraint_values)
            update_equivalence_classes_from_edge_lambda_indices!(M)
        end
    end
    if !isempty(lambdas)
        M.lambdas = lambdas
    end
    update_variable_terms!(M)
    for structure in structures
        append!(M.macro_structures, [structure])
        append!(M.macro_structure_description_lengths, [Structures.compute_description_length(structure, M.n, false)])
    end
end

function update_edge_lambda_indices!(M::Model.GraphModel, new_micro_structures, new_constraint_values)
    log_progress!(M, "updating edge lambda indices")
    start_lambda_index = length(M.lambdas) - length(new_constraint_values)
    for (idx, structure) in enumerate(new_micro_structures)
        new_idx = start_lambda_index+idx
        present_edges = intersect(keys(M.edge_lambda_indices), structure)
        for edge in present_edges
            append!(M.edge_lambda_indices[edge], new_idx)
        end
        absent_edge_dict = Dict(e=>[1,new_idx] for e in setdiff(structure,present_edges))
        merge!(M.edge_lambda_indices,absent_edge_dict)
    end
end

"""
roll back the model if the added structure is no good
eliminates last macro structure added and all its consequences
"""
function roll_back_model!(M::Model.GraphModel)
    structure_to_remove = M.macro_structures[end]
    micros_to_remove = determine_micros_to_remove(structure_to_remove)
    if structure_to_remove isa Structures.Star
        M.hub_spoke_edges = deepcopy(M.last_hub_spoke_edges)
        M.star_spoke_edges = deepcopy(M.last_star_spoke_edges)
        if structure_to_remove.n_edges_in_spokes == 0
            micros_to_remove = 0
            M.last_zero_edges = deepcopy(M.last_zero_edges)
        end
    elseif structure_to_remove isa Structures.Starclique || structure_to_remove isa Structures.Biclique
        M.star_spoke_edges = deepcopy(M.last_star_spoke_edges)
    end
    if micros_to_remove > 0
        roll_back_edge_lambda_indices!(M, micros_to_remove)
        roll_back_lambda_related!(M, micros_to_remove)
        update_equivalence_classes_from_edge_lambda_indices!(M)
    end
    update_equivalence_class_edge_probabilities!(M)
    update_edge_probability_description_length!(M)
    pop!(M.macro_structures)
    pop!(M.macro_structure_description_lengths)
    pop!(M.description_lengths_over_time)
    roll_back_total_description_length!(M)
    roll_back_optimization_result!(M)
end

function determine_micros_to_remove(structure_to_remove::Structures.Structure)
    if structure_to_remove isa Structures.Clique || structure_to_remove isa Structures.Star
        micros_to_remove = 1
    elseif structure_to_remove isa Structures.Biclique || structure_to_remove isa Structures.Starclique
        micros_to_remove = 3
    else
        throw(ErrorException("Cannot roll back for structure of type $(typeof(structure_to_remove))!"))
    end
    return micros_to_remove
end

function roll_back_edge_lambda_indices!(M::Model.GraphModel, micros_to_remove::Int64)
    for micro_structure in M.micro_structures[end-micros_to_remove+1:end]
        for edge in micro_structure
            pop!(M.edge_lambda_indices[edge])
            if M.edge_lambda_indices[edge] == [1]
                delete!(M.edge_lambda_indices, edge)
            end
        end
    end
end

function roll_back_lambda_related!(M::Model.GraphModel, micros_to_remove::Int64)
    first_to_delete = length(M.micro_structures)-micros_to_remove+1
    last_to_delete = length(M.micro_structures)
    deleteat!(M.lambdas, first_to_delete:last_to_delete)
    deleteat!(M.constraint_values, first_to_delete:last_to_delete)
    deleteat!(M.micro_structures, first_to_delete:last_to_delete)
end

function update_equivalence_classes_from_edge_lambda_indices!(M::Model.GraphModel)
    equivalence_classes = Dict{Set{Int64},Array{Tuple{Int64,Int64},1}}()
    for (k,v) in M.edge_lambda_indices
        append!(get!(equivalence_classes, Set(v), Array{Int64,1}()), [k])
    end
    M.edge_equivalence_classes = equivalence_classes
end

function update_equivalence_class_edge_probabilities!(M::Model.GraphModel)
    probabilities = Dict{Set{Int64},BigFloat}(
        equivalence_class=>edge_probability_for_equiv(equivalence_class, M.lambdas) for equivalence_class in keys(M.edge_equivalence_classes))
    probabilities[Set([1,])] = edge_probability_for_equiv(Set([1,]), M.lambdas)
    M.equivalence_class_edge_probabilities = probabilities
end

function roll_back_total_description_length!(M::Model.GraphModel)
    M.total_description_length_maxent = M.description_lengths_over_time[end]
end

function roll_back_optimization_result!(M::Model.GraphModel)
    M.last_optimization_result = deepcopy(M.backup_optimization_result)
    M.entropy = M.last_optimization_result.minimum
end

end

max_stars_repo_name: dataspider/momo

max_issues_repo_name: dataspider/momo

max_forks_repo_name: dataspider/momo


text:
# PROJECT_SDC    Project a point onto the semidefinite cone.
#
#   project_sdc(V), where V is a symmetric n x n matrix, is the
#   projection of V onto the semidefinite cone.
function project_sdc(V)
    W, D = eig(V)
    X = W * max(0, D) * W'
    return X
end


text:
function AcelerarRuido(carretera; vmax::Array{Int8, 1} = Int8[3, 5], R::Float32 = 0.2f0, frontera_izq::Int64 = 5)

    #frontera_izq = 5
    #vmax = [3, 5]
    #R = 0.2
    for v in carretera[frontera_izq+1:end]
        if v.tipo == 1 || v.tipo == 2

            #Acelerar
            if v.velocidad != -1
                #if (i > 349) && ( i < 620 && i > 663) # Zona de Curvas
                #    Cambiar_Vehiculo!(v, min(v.velocidad+1, vmax[v.tipo]), v.posicion, v.tipo, v.cambio, v.num)
                #else
                    Cambiar_Vehiculo!(v, min(v.velocidad+1, vmax[v.tipo]), v.posicion, v.tipo, v.cambio, v.num)
                #end
            end

            #Ruido
            if v.velocidad > 0
                if rand() <= R
                    Cambiar_Vehiculo!(v, max(v.velocidad-1, 0), v.posicion, v.tipo, v.cambio, v.num)
                end
            end
        end
    end
end

function DesacelerarMover(C::Carretera1D)

    alfa = 0.75f0
    frontera_izq = int8(5)

    i = Pos_izq(C.carretera)
    i_1 = ( i > 1 ? Pos_izq(C.carretera[1:i-1]) : 0 )

    # Este while es para los pasos siguientes
    while (i > 0)

        if i_1 != 0
        ############################ cambio ###################################
          distancia = i - i_1 - 1 # distancia = i - i_1 - C.carretera[i_1].longitud

          # Desaceleramos el automovil i-1 a partir del automovil i
          velocidad_estimada = round(Int8, floor( C.carretera[i].velocidad * (1.0 - alfa) + 0.5 ) )
          Cambiar_Vehiculo!(C.carretera[i_1], min( (velocidad_estimada+distancia), C.carretera[i_1].velocidad),
                              C.carretera[i_1].posicion, C.carretera[i_1].tipo, C.carretera[i_1].cambio,
                                C.carretera[i_1].num)
        end

        # Movemos el automovil i

        pos_sig = C.carretera[i].velocidad + C.carretera[i].posicion
        if C.carretera[i].posicion != pos_sig #&& C.carretera[pos_sig].tipo == -1

            if pos_sig > C.N
                Celda_Vacia!(C.carretera[i])
                C.cuenta -= 1
            end

            if pos_sig <= C.N && C.carretera[pos_sig].tipo == -1

                if C.carretera[i].posicion <= frontera_izq && pos_sig > frontera_izq
                    C.cuenta += 1
                end

                Cambiar_Vehiculo!(C.carretera[pos_sig], C.carretera[i].velocidad, pos_sig, C.carretera[i].tipo, 0
                                    , C.carretera[i].num)

                Celda_Vacia!(C.carretera[i])
            elseif pos_sig <= C.N && C.carretera[pos_sig].tipo == -2
                while C.carretera[pos_sig].tipo == -2
                   pos_sig -= 1
                end
               Cambiar_Vehiculo!(C.carretera[pos_sig], C.carretera[i].velocidad, pos_sig, C.carretera[i].tipo, 0,
                                C.carretera[i].num)
               Celda_Vacia!(C.carretera[i])
            end
        end
        # Reindexamos el automovil i-1 como el automovil i

        i = i_1
        i_1 = ( i > 1 ? Pos_izq(C.carretera[1:i-1]) : 0 )
    end

    alfa = frontera_izq = i = i_1 = 0
end

################################## Funciones especiales para las simulaciones en la zona de curvas ############################

######### Sentido 1 ###########

function AcelerarRuidoS1(carretera, vmax::Array{Int8, 1} = Int8[3, 5], R::Float32 = 0.2f0, frontera_izq::Int64 = 5)

    for v in carretera[frontera_izq+1:end]
        if v.tipo == 1 || v.tipo == 2

            #Acelerar
            if v.velocidad != -1
                if (v.posicion < 349) || ( v.posicion > 620 && v.posicion < 663) # Zona de Curvas
                    vmax[2] = int8(4)
                end
                Cambiar_Vehiculo!(v, min(v.velocidad+1, vmax[v.tipo]), v.posicion, v.tipo, v.cambio, v.num)

            end

            #Ruido
            if v.velocidad > 0
                if rand() <= R
                    Cambiar_Vehiculo!(v, max(v.velocidad-1, 0), v.posicion, v.tipo, v.cambio, v.num)
                end
            end
            vmax[2] = int8(5)
        end
    end
  vmax = 0
end

##### Sentido 2 #######

function AcelerarRuidoS2(carretera, vmax::Array{Int8, 1} = Int8[3, 5], R::Float32 = 0.2f0, frontera_izq::Int64 = 5)

    for v in carretera[frontera_izq+1:end]
        if v.tipo == 1 || v.tipo == 2

            #Acelerar
            if v.velocidad != -1
                if (v.posicion > 3518) || ( v.posicion < 3408 && v.posicion > 3265) # Zona de Curvas
                    vmax[2] = int8(4)
                end

                if (v.posicion > 2500) && v.tipo == 1 # Pendiente
                    vmax[1] = int8(2)
                end
                Cambiar_Vehiculo!(v, min(v.velocidad+1, vmax[v.tipo]), v.posicion, v.tipo, v.cambio, v.num)
            end

            #Ruido
            if v.velocidad > 0
                if rand() <= R
                    Cambiar_Vehiculo!(v, max(v.velocidad-1, 0), v.posicion, v.tipo, v.cambio, v.num)
                end
            end
        vmax = Int8[3, 5]
        end
    end
  vmax = 0
end


text: include("hello.jl")


text:
using Pkg
using Documenter, BED

format = Documenter.HTML(
    edit_link = "develop"
)

makedocs(
    format = format,
    checkdocs = :all,
    linkcheck = true,
    modules = [BED],
    sitename = "BED.jl",
    pages = [
        "Home" => "index.md",
        "BED" => "man/bed.md",
        "API Reference" => "man/api.md"
    ],
    authors = replace(join(Pkg.TOML.parsefile("Project.toml")["authors"], ", "), r" <.*?>" => "" ) * ", The BioJulia Organisation, and other contributors."
)

deploydocs(
    repo = "github.com/BioJulia/BED.jl.git",
    devbranch = "develop",
    push_preview = true
)


text:
"""
Direct Stiffness Method
=======================

Plane truss/frame analysis using Direct stiffness method. It currently only supports two types of
beams: truss (both ends are pinned joints) and frame (both ends are fixed joints).

References
----------

- [hibbeler]: Hibbeler, R. C. "Structural Analysis 8th edition." (2012).

"""
module DirectStiffnessMethod

export
    Node, Member,
    JointType, FixedJoint, PinnedJoint, RollerJoint, RollerYJoint, SimpleJoint, FreeJoint,
    EndType, FixedConnection, PinnedConnection, RollerConnection, RollerYConnection,
        SimpleConnection, SimpleYConnection,
    Beam,
    Structure, num_nodes, num_members, node_indices,
    solve

using LinearAlgebra
using SparseArrays

#===== Supports & Connections =====#

"The joint (support) type."
@enum JointType begin
    FixedJoint
    PinnedJoint
    RollerJoint
    RollerYJoint
    SimpleJoint
    FreeJoint
end

"The end (connection) type."
@enum EndType begin
    FixedConnection
    PinnedConnection
    RollerConnection
    RollerYConnection
    SimpleConnection
    SimpleYConnection
end

function join_free_dims(j::JointType)::BitVector
    # The case [1, 1, 0] and [0, 1, 1] (SimpleY) seem very rare in real life.
    if j == FixedJoint
        [0, 0, 0]
    elseif j == PinnedJoint
        [0, 0, 1]
    elseif j == RollerJoint
        [1, 0, 0]
    elseif j == RollerYJoint
        [0, 1, 0]
    elseif j == SimpleJoint
        [1, 0, 1]
    elseif j == FreeJoint
        [1, 1, 1]
    end |> BitVector
end

function end_effective_dims(e::EndType)::BitVector
    # The "Free" connection makes no sense, so it is excluded. The real life connections are mostly
    # "Fixed" and "Pinned".
    if e == FixedConnection
        [1, 1, 1]
    elseif e == PinnedConnection
        [1, 1, 0]
    elseif e == RollerConnection
        [0, 1, 1]
    elseif e == RollerYConnection
        [1, 0, 1]
    elseif e == SimpleConnection
        [0, 1, 0]
    elseif e == SimpleYConnection
        [1, 0, 0]
    end |> BitVector
end

# Both ends must have the same connection type. If one end is pinned and the other is fixed, we
# should use the "condensation method", which has not been implemented here.
function beam_effective_dims(e1::EndType, e2::EndType)::BitVector
    if (e1, e2) == (FixedConnection, FixedConnection)
        [1, 1, 1, 1, 1, 1]
    elseif (e1, e2) == (PinnedConnection, PinnedConnection)
        [1, 0, 0, 1, 0, 0]
    else
        error("unimplemented")
    end |> BitVector
end

#===== Node & Member =====#

"Represents joints, supports, or ends of a member."
struct Node
    id::Int
    x::Float64
    y::Float64
    joint_type::JointType
end

struct Member
    id::Int
    from::Int
    to::Int

    near_end::EndType
    far_end::EndType

    "Young's modulus"
    E::Float64

    "The area"
    A::Float64

    "The area moment of inertia"
    I::Float64
end

#===== Beam =====#

struct Beam
    "The length"
    l::Float64

    "Cosine of the sloping angle"
    λ_x::Float64

    "Sine of the sloping angle"
    λ_y::Float64

    "E * A"
    EA::Float64

    "E * I"
    EI::Float64

    "Effective indices of connection"
    conn_indices::BitVector
end

# Truss beam is treated as a special case of frame beam.
#
#     k_m = EA / l *
#         [
#              1.0    -1.0
#             -1.0     1.0
#         ]
#
#     T =
#     [
#         c       s       0.0     0.0
#         0.0     0.0     c       s
#     ]
#

raw"""
Computes member stiffness matrix $k_m$.

In the local coordinate system, the origin is at the "near" node and X axis extends toward the "far"
node.
"""
function compute_member_k(beam::Beam)::Matrix{Float64}
    l = beam.l
    EA = beam.EA
    EI = beam.EI
    l2 = l * l
    l3 = l * l2

    Symmetric(
        [
            EA / l      0.0         0.0         -EA / l     0.0         0.0
            0.0         12EI / l3   6EI / l2    0.0         -12EI / l3  6EI / l2
            0.0         0.0         4EI / l     0.0         -6EI / l2   2EI / l
            0.0         0.0         0.0         EA / l      0.0         0.0
            0.0         0.0         0.0         0.0         12EI / l3   -6EI / l2
            0.0         0.0         0.0         0.0         0.0         4EI / l
        ],
        :U,
    ) |> Matrix
end

function compute_transform_mat(beam::Beam)::Matrix{Float64}
    (c, s) = (beam.λ_x, beam.λ_y)

    [
        c       s       0.0     0.0     0.0     0.0
        -s      c       0.0     0.0     0.0     0.0
        0.0     0.0     1.0     0.0     0.0     0.0
        0.0     0.0     0.0     c       s       0.0
        0.0     0.0     0.0     -s      c       0.0
        0.0     0.0     0.0     0.0     0.0     1.0
    ]
end

#===== Structure =====#

"Represents a planar structure."
mutable struct Structure
    nodes::Vector{Node}
    members::Vector{Member}
end

@inline function num_nodes(s::Structure)::Int
    length(s.nodes)
end

@inline function num_members(s::Structure)::Int
    length(s.members)
end

function node_indices(node_id::Int)::UnitRange{Int}
    (3 * (node_id - 1) + 1):(3 * node_id)
end

"Creates beams from nodes and members."
function assemble_beam(s::Structure, member_id::Int)::Beam
    mem = s.members[member_id]
    from_node = s.nodes[mem.from]
    to_node = s.nodes[mem.to]

    lx = to_node.x - from_node.x
    ly = to_node.y - from_node.y
    l = sqrt(lx ^ 2 + ly ^ 2)
    λ_x = lx / l
    λ_y = ly / l

    EA = mem.E * mem.A
    EI = mem.E * mem.I

    conn_indices = vcat(end_effective_dims(mem.near_end), end_effective_dims(mem.far_end))
    Beam(l, λ_x, λ_y, EA, EI, conn_indices)
end

raw"""
Computes member global stiffness matrix $k$.
"""

function compute_global_k(
    beam::Beam,
    near_end::EndType=FixedConnection,
    far_end::EndType=FixedConnection,
)::Matrix{Float64}
    # - Displacement transformation matrix `T`, which transforms global displacements to local
    #   displacements.
    # - Member stiffness matrix `k_m`, which transforms local displacements to local forces.
    # - Displacement transformation matrix `T'`, which transforms local forces to global forces.
    #
    # k = T' * k_m * T

    k_m = compute_member_k(beam)
    T = compute_transform_mat(beam)
    selection = beam_effective_dims(near_end, far_end)
    T' * (selection .* k_m * T)
end

raw"""
Assembles structure stiffness matrix $K$.
"""
function assemble_struct_k(s::Structure)::SparseMatrixCSC{Float64}
    n_nodes = num_nodes(s)
    n_members = num_members(s)
    K = sparse([], [], Float64[], 3n_nodes, 3n_nodes)
    for i in 1:n_members
        mem = s.members[i]

        b = assemble_beam(s, i)
        k = compute_global_k(b, mem.near_end, mem.far_end)

        from_node = mem.from
        to_node = mem.to
        indices = [node_indices(from_node); node_indices(to_node)]

        K_view = @view K[indices, indices]
        ix = b.conn_indices
        K_view[ix, ix] .+= k[ix, ix]
    end
    K
end

function free_displacement_indices(s::Structure)::BitVector
    n_nodes = num_nodes(s)
    free_ix = BitVector(undef, 3 * n_nodes)
    for i in 1:n_nodes
        ix = node_indices(i)
        free_ix[ix] = join_free_dims(s.nodes[i].joint_type)
    end
    free_ix
end

#=

The structure stiffness equation (Subscript `k`: known variable; `u`: unknown variable):

    [ f_k     [ K_11 K_12     [ d_u
      f_u ] =   K_21 K_22 ] *   d_k ]

`d_k` is usually 0.

    d_u = K_11^{-1} f_k
    f_u = K_21 d_u

=#

function solve(
    s::Structure,
    loads::AbstractVector{Float64}
)::Tuple{SparseVector{Float64}, SparseVector{Float64}}
    free_disp_ix = free_displacement_indices(s)
    unknown_force_ix = .!free_disp_ix

    K = assemble_struct_k(s)
    K_11 = K[free_disp_ix, free_disp_ix]
    d_u = qr(K_11) \ Vector(loads[free_disp_ix])
    K_21 = K[unknown_force_ix, free_disp_ix]
    f_u = K_21 * d_u

    n_dims = length(free_disp_ix)
    d = sparsevec((1:n_dims)[free_disp_ix], d_u, n_dims)
    f = sparsevec((1:n_dims)[unknown_force_ix], f_u, n_dims)
    (d, f)
end

end # module

max_stars_repo_name: gyk/TrivialSolutions

max_issues_repo_name: gyk/TrivialSolutions

max_forks_repo_name: gyk/TrivialSolutions


text:
# --------------------------------------------------------------------------
# ACE.jl and SHIPs.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

using ACE, JuLIP, BenchmarkTools
using JuLIP: alloc_temp, alloc_temp_d

trans = PolyTransform(2, 1.0)
fcut = PolyCutoff2s(Val(2), 0.5, 3.0)
ships = [SHIPBasis(SparseSHIP(n, 15, 1.5), trans, fcut)
         for n = 2:4]

NR = 50
Rs = 1.0 .+ 2*(rand(JVecF, NR) .- 0.5)
Zs = zeros(Int16, NR)
z0 = 0

@info("profile `precompute_A!`")
tmp = alloc_temp(ships[1], NR)
@btime ACE.precompute_A!($tmp, $(ships[1]), $Rs, $Zs)
@btime ACE.precompute_A!($tmp, $(ships[1]), $Rs, $Zs)
@info("profile `precompute_grads!`")
tmpd = alloc_temp_d(ships[1], NR)
@btime ACE.precompute_grads!($tmpd, $(ships[1]), $Rs, $Zs)
@btime ACE.precompute_grads!($tmpd, $(ships[1]), $Rs, $Zs)

@info("profile basis computation")
for n = 2:4
   @info("  body-order $(n+1):")
   🚢 = ships[n-1]
   B = ACE.alloc_B(🚢)
   @info("     evaluate:")
   @btime evaluate!($B, $tmp, $🚢, $Rs, $Zs, $z0)
   @btime evaluate!($B, $tmp, $🚢, $Rs, $Zs, $z0)
   @info("     evaluate_d:")
   dB = ACE.alloc_dB(🚢, Rs)
   @btime evaluate_d!($B, $dB, $tmpd, $🚢, $Rs, $Zs, $z0)
   @btime evaluate_d!($B, $dB, $tmpd, $🚢, $Rs, $Zs, $z0)
end

# ##
# using Profile
# 🚢 = ships[2]
# B = ACE.alloc_B(🚢)
# dB = ACE.alloc_dB(🚢, Rs)
# @btime evaluate!($B, $tmp, $🚢, $Rs, $Zs, $z0)
# @btime evaluate_d!($B, $dB, $tmpd, $🚢, $Rs, $Zs, $z0)
#
# @code_warntype ACE._evaluate!(B, tmp, 🚢, Val{3}(), 1, 🚢.NuZ[3,1])
#
# ACE._evaluate!(B, tmp, 🚢, Val{3}(), 1, 🚢.NuZ[3,1])
#
# ##
#
# function runn(N, f, args...)
#    for n = 1:N
#       f(args...)
#    end
# end
#
# runn(10, evaluate!, B, tmp, 🚢, Rs, Zs, z0)
# runn(10, evaluate_d!, B, dB, tmpd, 🚢, Rs, Zs, z0)
#
# ##
#
# Profile.clear()
# @profile runn(100, evaluate!, B, tmp, 🚢, Rs, Zs, z0)
# Profile.print()


text:
export Sequence,
	Quality,
	qual2num,
	num2qual,
	qual_str,
	SeqRead,
	FastaRead,
	FastqRead,
	FastqPair,
	dna,
	rna,
	aa,
	reverse,
	complement,
	reverse_complement,
	transcribe,
	display_sequence,
	display_fastq,
	display_fasta,
	display_fastq_pair,
	Interval,
	BED,
	Variant,
	version,
	vcf_intersect,
	vcf_minus,
	vcf_merge,
	vcf_samples,
	vcf_filter,
	vcf_diff_genotype,
	vcf_split_sample,
	vcf_filter_region,
	parse_genotype,
	==,
	-,
	!,
	~,
	>,
	<,
	sort,
	sort!,
	+,
	*,
	-

export VcfHeader,
	Vcf,
	GtfHeader,
	GtfItem,
	GtfData,
	Gtf

import Base: convert,
	==,
	-,
	!,
	~,
	>,
	<,
	sort,
	sort!,
	+,
	*,
	-

if VERSION < v"0.5.0-"
    import Base:complement
end

export DNA_SEQ,
	RNA_SEQ,
	AA_SEQ

include("sequence.jl")
include("quality.jl")
include("read.jl")
include("interval.jl")
include("variant.jl")
include("gtf.jl")


text:
mutable struct DpPolicyAgent <: AbstractMDPAgent
    v_::Float64
    ω_::Float64
    dt::Float64
    prev_v_::Float64
    prev_ω_::Float64
    estimator_::AbstractEstimator
    puddle_coeff::Float64
    puddle_depth_::Float64
    total_reward_::Float64
    in_goal_::Bool
    final_value_::Float64
    goal::Goal
    reso::Vector{Float64}
    pose_min::Vector{Float64}
    pose_max::Vector{Float64}
    index_nums::Vector{Int64}
    policy_data::AbstractArray{Float64,4}
end

function DpPolicyAgent(
    v::Float64,
    ω::Float64,
    dt::Float64,
    estimator::AbstractEstimator,
    goal::Goal,
    reso::Vector{Float64};
    lowerleft = [-4.0, -4.0],
    upperright = [4.0, 4.0],
    puddle_coeff = 100,
)
    pose_min = vcat(lowerleft, [0.0])
    pose_max = vcat(upperright, [2pi])
    index_nums = [convert(Int64, round((pose_max[i] - pose_min[i]) / reso[i])) for i = 1:3]
    policy_data = zeros(Float64, index_nums..., 2)
    return DpPolicyAgent(
        v,
        ω,
        dt,
        0.0,
        0.0,
        estimator,
        puddle_coeff,
        0.0,
        0.0,
        false,
        0.0,
        goal,
        reso,
        pose_min,
        pose_max,
        index_nums,
        policy_data,
    )
end

function init_policy(agent::DpPolicyAgent, fname = "policy.txt")
    policy_data = agent.policy_data
    fd = open(fname, "r")
    lines = readlines(fd)
    for line in lines
        tokens = split(line, " ")
        tokens = map(x -> parse(Float64, x), tokens)
        id1, id2, id3 = map(x -> convert(Int64, x), tokens[1:3])
        v, ω = tokens[4], tokens[5]
        policy_data[id1, id2, id3, :] = [v, ω]
    end
end

function policy(agent::DpPolicyAgent)
    if agent.in_goal_
        return [0.0, 0.0]
    end
    reso = agent.reso
    pose = agent.estimator_.pose_
    pose_min = agent.pose_min
    index_nums = agent.index_nums
    index = [convert(Int64, round((pose[i] - pose_min[i]) / reso[i])) for i = 1:3]
    index[3] = (index[3] + 10 * index_nums[3]) % index_nums[3]
    if index[3] == 0
        index[3] = index_nums[3]
    end
    for i in [1, 2]
        if index[i] < 1
            index[i] = 1
        end
        if index[i] > index_nums[i]
            index[i] = index_nums[i]
        end
    end
    return agent.policy_data[index..., :]
end

mutable struct QMDPAgent <: AbstractMDPAgent
    v_::Float64
    ω_::Float64
    dt::Float64
    prev_v_::Float64
    prev_ω_::Float64
    estimator_::AbstractEstimator
    puddle_coeff::Float64
    puddle_depth_::Float64
    total_reward_::Float64
    in_goal_::Bool
    final_value_::Float64
    goal::Goal
    pose_min::Vector{Float64}
    pose_max::Vector{Float64}
    dp::PolicyEvaluator
    current_value_::Float64
end

function QMDPAgent(
    v::Float64,
    ω::Float64,
    dt::Float64,
    estimator::AbstractEstimator,
    goal::Goal,
    reso::Vector{Float64};
    lowerleft = [-4.0, -4.0],
    upperright = [4.0, 4.0],
    puddle_coeff = 100,
)
    pose_min = vcat(lowerleft, [0.0])
    pose_max = vcat(upperright, [2pi])
    dp = PolicyEvaluator(
        reso,
        goal;
        lowerleft = lowerleft,
        upperright = upperright,
        dt = dt,
        puddle_coeff = puddle_coeff,
    )
    return QMDPAgent(
        v,
        ω,
        dt,
        0.0,
        0.0,
        estimator,
        puddle_coeff,
        0.0,
        0.0,
        false,
        0.0,
        goal,
        pose_min,
        pose_max,
        dp,
        0.0,
    )
end

function init_policy(agent::QMDPAgent, fname = "policy.txt")
    # need this to set `indices` in PolicyEvaluator
    init_value(agent.dp)
    policy_data = agent.dp.policy_
    fd = open(fname, "r")
    lines = readlines(fd)
    for line in lines
        tokens = split(line, " ")
        tokens = map(x -> parse(Float64, x), tokens)
        id1, id2, id3 = map(x -> convert(Int64, x), tokens[1:3])
        v, ω = tokens[4], tokens[5]
        policy_data[id1, id2, id3, :] = [v, ω]
    end
    init_state_transition_probs(agent.dp, agent.dp.dt, 10)
end

function init_value(agent::QMDPAgent, fname = "value.txt")
    value_data = agent.dp.value_function_
    fd = open(fname, "r")
    lines = readlines(fd)
    for line in lines
        tokens = split(line, " ")
        tokens = map(x -> parse(Float64, x), tokens)
        id1, id2, id3 = map(x -> convert(Int64, x), tokens[1:3])
        v = tokens[4]
        value_data[id1, id2, id3] = v
    end
end

function correct_index(index_::Vector{Int64}, index_nums::Vector{Int64})
    index = copy(index_)
    while index[3] > index_nums[3]
        index[3] -= index_nums[3]
    end
    while index[3] < 0
        index[3] += index_nums[3]
    end
    if index[3] == 0
        index[3] = index_nums[3]
    end
    for i in [1, 2]
        if index[i] < 1
            index[i] = 1
        end
        if index[i] > index_nums[i]
            index[i] = index_nums[i]
        end
    end
    return index
end

function to_index(
    pose::Vector{Float64},
    pose_min::Vector{Float64},
    reso::Vector{Float64},
    index_nums::Vector{Int64},
)
    index = [convert(Int64, round((pose[i] - pose_min[i]) / reso[i])) for i = 1:3]
    return correct_index(index, index_nums)
end

function action_value(agent::QMDPAgent, action::Vector{Float64}, index::Vector{Int64})
    v, ω = action[1], action[2]
    value = 0.0
    dp = agent.dp
    state_transition_probs = dp.state_transition_probs
    value_function = dp.value_function_
    index_nums = agent.dp.index_nums
    for trans_probs in state_transition_probs[(v, ω, index[3])]
        trans_ind = trans_probs[1]
        prob = trans_probs[2]
        after_ = [index...] .+ trans_ind
        after = correct_index(after_, index_nums)
        reward = -dp.dt * dp.depth[after[1], after[2]] * dp.puddle_coeff - dp.dt
        value += (value_function[after...] + reward) * prob
    end
    return value
end

function evaluation(
    agent::QMDPAgent,
    action::Vector{Float64},
    indices::Vector{Vector{Int64}},
)
    return sum([action_value(agent, action, index) for index in indices]) / length(indices)
end

function policy(agent::QMDPAgent)
    particles = agent.estimator_.particles_
    reso = agent.dp.reso
    pose_min = agent.pose_min
    index_nums = agent.dp.index_nums
    indices =
        [to_index(particle.pose_, pose_min, reso, index_nums) for particle in particles]
    value_function = agent.dp.value_function_
    agent.current_value_ =
        sum([value_function[index...] for index in indices]) / length(indices)

    max_val = -1e100
    a = nothing
    for action in agent.dp.actions
        val = evaluation(agent, action, indices)
        if val > max_val
            max_val = val
            a = copy(action)
        end
    end
    return a
end

mutable struct BeliefDP
    pose_min::Vector{Float64}
    pose_max::Vector{Float64}
    reso::Vector{Float64}
    goal::Goal
    index_nums::Vector{Int64}
    indices::Vector{Tuple{Int64,Int64,Int64,Int64}}
    value_function_::AbstractArray{Float64,4}
    final_state_flags_::AbstractArray{Float64,4}
    policy_::AbstractArray{Float64,5}
    actions::Set{Vector{Float64}}
    state_transition_probs::Dict{
        Tuple{Float64,Float64,Int64}, # key is (v::Float64, ω::Float64, θ_index::Int64,)
        Vector{Tuple{Vector{Int64},Float64}}, # value is vector of ([tran_x_id, tran_y_id, tran_z_id], prob,)
    }
    depth::AbstractArray{Float64,2}
    dt::Float64
    puddle_coeff::Float64
    dev_borders::Vector{Float64}
    dev_borders_side::Vector{Float64}
    motion_sigma_transition_probs::Dict{
        Tuple{Int64,Vector{Float64}},
        Vector{Tuple{Int64,Float64}},
    }
    obs_sigma_transition_probs::Dict{Tuple{Int64,Int64,Int64,Int64},Tuple{Int64,Float64}}
    depths::AbstractArray{Float64,3}
end

# constructor
function BeliefDP(
    reso::Vector{Float64},
    goal::Goal;
    lowerleft = [-4.0, -4.0],
    upperright = [4.0, 4.0],
    dt = 0.1,
    puddle_coeff = 100.0,
    dev_borders = [0.1, 0.2, 0.4, 0.8],
)
    pose_min = vcat(lowerleft, [0.0])
    pose_max = vcat(upperright, [2pi])

    index_nums = [convert(Int64, round((pose_max[i] - pose_min[i]) / reso[i])) for i = 1:3]
    push!(index_nums, length(dev_borders) + 1)
    v = zeros(Float64, index_nums...)
    f = zeros(Float64, index_nums...)
    indices = Vector{Tuple{Int64,Int64,Int64,Int64}}(undef, 0)
    policy_ = zeros(Float64, index_nums..., 2)
    actions = Set{Vector{Float64}}()
    state_transition_probs =
        Dict{Tuple{Float64,Float64,Int64},Vector{Tuple{Vector{Int64},Float64}}}()
    depth = zeros(Float64, index_nums[1], index_nums[2])
    dev_borders_side = [dev_borders[1] / 10, dev_borders..., dev_borders[end] * 10]
    motion_sigma_transition_probs =
        Dict{Tuple{Int64,Vector{Float64}},Vector{Tuple{Int64,Float64}}}()
    obs_sigma_transition_probs = Dict{Tuple{Int64,Int64,Int64,Int64},Tuple{Int64,Float64}}()
    depths = zeros(Float64, index_nums[1], index_nums[2], index_nums[3])

    return BeliefDP(
        pose_min,
        pose_max,
        reso,
        goal,
        index_nums,
        indices,
        v,
        f,
        policy_,
        actions,
        state_transition_probs,
        depth,
        dt,
        puddle_coeff,
        dev_borders,
        dev_borders_side,
        motion_sigma_transition_probs,
        obs_sigma_transition_probs,
        depths,
    )
end

function set_belief_final_state(
    reso::Vector{Float64},
    goal::Goal,
    pose_min::Vector{Float64},
    index::Vector{Int64},
)
    x_min, y_min, θ_min = pose_min .+ reso .* (index[1:3] .- 1)
    x_max, y_max, θ_max = pose_min .+ reso .* (index[1:3])

    corners = [[x_min, y_min], [x_max, y_min], [x_max, y_max], [x_min, y_max]]
    ret = convert(Float64, all([inside(goal, corner) for corner in corners]))
    return ret * convert(Float64, index[4] == 1)
end

function init_value(agent::BeliefDP)
    index_nums = agent.index_nums
    goal = agent.goal
    reso = agent.reso
    pose_min = agent.pose_min
    v = agent.value_function_
    f = agent.final_state_flags_
    for id1 = 1:index_nums[1]
        for id2 = 1:index_nums[2]
            for id3 = 1:index_nums[3]
                for id4 = 1:index_nums[4]
                    index = (id1, id2, id3, id4)
                    val = set_belief_final_state(reso, goal, pose_min, [index...])
                    @inbounds f[index...] = val
                    @inbounds v[index...] = (val == 1.0) ? goal.value : (-100.0)
                    # this line is problematic if we use @distributed
                    push!(agent.indices, index)
                end
            end
        end
    end
end

function init_policy(agent::BeliefDP)
    indices = agent.indices
    reso = agent.reso
    pose_min = agent.pose_min
    goal = agent.goal

    for index in indices
        center = pose_min .+ reso .* ([index[1:3]...] * 1.0 .- 0.5)
        agent.policy_[index..., :] = initial_policy(center, goal)
    end
end

function init_state_transition_probs(agent::BeliefDP; sampling_num = 100)
    dt = agent.dt
    reso = agent.reso
    indices = agent.indices
    policy = agent.policy_
    actions = Set{Vector{Float64}}()
    index_nums = agent.index_nums
    pose_min = agent.pose_min
    actions = agent.actions
    for index in indices
        push!(actions, policy[index..., :])
    end

    dxs = collect(range(0.001, reso[1] * 0.999, length = sampling_num))
    dys = collect(range(0.001, reso[2] * 0.999, length = sampling_num))
    dθs = collect(range(0.001, reso[3] * 0.999, length = sampling_num))
    # key is (v::Float64, ω::Float64, θ_index::Int64,)
    # value is vector of ([tran_x_id, tran_y_id, tran_z_id], prob,)
    transition_probs = agent.state_transition_probs

    for a in actions
        for i_θ = 1:index_nums[3]
            transitions = Vector{Vector{Int64}}(undef, 0)
            for dx in dxs
                for dy in dys
                    for dθ in dθs
                        before = [dx, dy, dθ + (i_θ - 1) * reso[3]] .+ pose_min
                        before_index = [1, 1, i_θ]
                        after = state_transition(before, a[1], a[2], dt)
                        after_index_ = (after - pose_min) ./ reso
                        after_index = map(x -> convert(Int64, floor(x) + 1), after_index_)
                        push!(transitions, after_index - before_index)
                    end
                end
            end
            elems = unique(transitions)
            counts = [count(x -> x == elem, transitions) for elem in elems]
            probs = [c / (sampling_num)^3 for c in counts]
            transition_probs[(a[1], a[2], i_θ)] = [iter for iter in zip(elems, probs)]
        end
    end
end

function init_depth(agent::BeliefDP, world::PuddleWorld; sampling_num = 100)
    reso = agent.reso
    dxs = collect(range(0, reso[1], length = sampling_num))
    dys = collect(range(0, reso[2], length = sampling_num))
    index_nums = agent.index_nums
    pose_min = agent.pose_min
    puddles = world.puddles_
    depth = agent.depth

    for x = 0:index_nums[1]-1
        for y = 0:index_nums[2]-1
            for dx in dxs
                for dy in dys
                    pose = pose_min .+ (reso .* [x, y, 0]) .+ [dx, dy, 0]
                    for puddle in puddles
                        depth[x+1, y+1] +=
                            (puddle.depth * (convert(Float64, inside(puddle, pose))))
                    end
                end
            end
            depth[x+1, y+1] /= (sampling_num)^2
        end
    end
end

function correct_index(agent::BeliefDP, index_::Vector{Int64})
    index = copy(index_)
    index_nums = agent.index_nums
    reso = agent.reso
    while index[3] < 1
        index[3] += index_nums[3]
    end
    while index[3] > index_nums[3]
        index[3] -= index_nums[3]
    end
    out_reward = 0.0
    for i in [1, 2]
        if index[i] < 1
            index[i] = 1
            out_reward = -1e100
        elseif index[i] > index_nums[i]
            index[i] = index_nums[i]
            out_reward = -1e100
        end
    end
    return index, out_reward
end

function action_value(
    agent::BeliefDP,
    action::Vector{Float64},
    index::Vector{Int64},
    value_function::AbstractArray{Float64,4},
)
    value = 0.0
    dt = agent.dt
    transition = agent.state_transition_probs[(action[1], action[2], index[3])]
    depths = agent.depths
    puddle_coeff = agent.puddle_coeff
    motion_transition = agent.motion_sigma_transition_probs
    obs_transition = agent.obs_sigma_transition_probs

    for (delta, prob) in transition
        after, out_reward = correct_index(agent, index[1:3] + [delta...])
        push!(after, 1)
        if haskey(motion_transition, (index[4], action))
            for (σ_after, σ_prob) in motion_transition[(index[4], action)]
                after[4] = σ_after
                σ_obs, σ_obs_prob = obs_transition[(after...,)]
                reward = -dt * depths[after[1:2]..., σ_obs] * puddle_coeff - dt + out_reward

                value +=
                    (value_function[after[1:3]..., σ_obs] + reward) *
                    σ_prob *
                    prob *
                    σ_obs_prob
            end
        end
    end
    return value
end

function value_iteration_sweep(agent::BeliefDP; γ = 1.0, snapshot = false)
    max_Δ = 0.0
    indices = agent.indices
    final_state_flags = agent.final_state_flags_
    value_function = nothing
    if snapshot == true
        value_function = copy(agent.value_function_)
    else
        value_function = agent.value_function_
    end
    for index in indices
        if final_state_flags[index...] == 0.0
            max_a = nothing
            max_q = -1e100
            for action in agent.actions
                q = action_value(agent, action, [index...], value_function)
                if q > max_q
                    max_a = copy(action)
                    max_q = q
                end
            end
            Δ = abs(value_function[index...] - max_q)
            max_Δ = max(Δ, max_Δ)
            agent.value_function_[index...] = max_q
            agent.policy_[index..., :] = max_a
        end
    end
    return max_Δ
end

function cov_to_index(dev_borders::Vector{Float64}, cov_::Matrix{Float64})
    σ = det(cov_)^(1.0 / 6)
    for (i, elem) in enumerate(dev_borders)
        if σ < elem
            return i
        end
    end
    return length(dev_borders)
end

function cov_to_index(agent::BeliefDP, cov_::Matrix{Float64})
    σ = det(cov_)^(1.0 / 6)
    for (i, elem) in enumerate(agent.dev_borders)
        if σ < elem
            return i
        end
    end
    return length(agent.dev_borders)
end

function calc_motion_sigma_transition_probs(
    agent::BeliefDP,
    min_σ::Float64,
    max_σ::Float64,
    action::Vector{Float64};
    sampling_num = 100,
)
    dt = agent.dt
    v, ω = action[1], action[2]
    if abs(ω) < 1e-5
        ω = 1e-5
    end
    F = matF(v, ω, dt, 0.0)
    M = matM(v, ω, dt, Dict("vv" => 0.19, "vω" => 0.001, "ωv" => 0.13, "ωω" => 0.2))
    A = matA(v, ω, dt, 0.0)
    indices = Dict{Int64,Float64}()
    for σ in range(min_σ, max_σ * 0.999, length = sampling_num)
        cov_ = σ * σ * F * transpose(F) + A * M * transpose(A)
        index_after = cov_to_index(agent.dev_borders, cov_)
        if !haskey(indices, index_after)
            indices[index_after] = 1
        else
            indices[index_after] += 1
        end
    end
    for (k, v) in indices
        indices[k] /= sampling_num
    end
    return [(k, v) for (k, v) in indices]
end

function init_motion_sigma_transition_probs(agent::BeliefDP)
    probs = agent.motion_sigma_transition_probs
    dev_borders_side = agent.dev_borders_side
    for a in agent.actions
        for i = 1:length(agent.dev_borders)+1
            probs[(i, a)] = calc_motion_sigma_transition_probs(
                agent,
                dev_borders_side[i],
                dev_borders_side[i+1],
                a,
            )
        end
    end
end

function observation_update(
    lm_id::Int64,
    S::Matrix{Float64},
    camera::IdealCamera,
    pose::Vector{Float64},
)
    distance_dev_rate = 0.14
    direction_dev = 0.05

    H = matH(pose, camera.landmarks_[lm_id+1].pos_)
    est_z = observation_function(pose, camera.landmarks_[lm_id+1].pos_)
    Q = matQ(distance_dev_rate * est_z[1], direction_dev)
    K = S * transpose(H) * (inv(Q + H * S * transpose(H)))
    return (Matrix(1.0I, 3, 3) - K * H) * S
end

function init_obs_sigma_transition_probs(agent::BeliefDP, camera::IdealCamera)
    sigma_transition = agent.obs_sigma_transition_probs
    pose_min = agent.pose_min
    reso = agent.reso
    dev_borders_side = agent.dev_borders_side
    for index in agent.indices
        pose = pose_min .+ reso .* ([index[1:3]...] * 1.0 .- 0.5)
        σ = (dev_borders_side[index[4]] + dev_borders_side[index[4]+1]) / 2.0
        S = Matrix(σ^2 * I, 3, 3)
        for d in observations(camera, pose)
            lm_id = convert(Int64, d[3])
            S = observation_update(lm_id, S, camera, pose)
        end
        sigma_transition[index] = (cov_to_index(agent, S), 1.0)
    end
end

function init_expected_depths(agent::BeliefDP, world::PuddleWorld; sampling_num = 100)
    puddles = world.puddles_
    index_nums = agent.index_nums
    reso = agent.reso
    pose_min = agent.pose_min
    dev_borders_side = agent.dev_borders_side
    depths = agent.depths
    for id1 = 1:index_nums[1]
        for id2 = 1:index_nums[2]
            for id3 = 1:index_nums[4]
                index = [id1, id2, id3]
                pose = pose_min[1:2] .+ reso[1:2] .* (index[1:2] .- 0.5)
                σ = (dev_borders_side[id3] + dev_borders_side[id3+1]) / 2.0
                belief = MvNormal(pose, Matrix(σ^2 * I, 2, 2))
                depth_sum = 0.0
                samples = rand(belief, sampling_num)
                for i_pos = 1:sampling_num
                    pos = samples[:, i_pos]
                    depth_sum += sum([
                        puddle.depth * convert(Float64, inside(puddle, pos)) for
                        puddle in puddles
                    ])
                end
                depths[index...] = depth_sum / sampling_num
            end
        end
    end
end

mutable struct AMDPPolicyAgent <: AbstractMDPAgent
    v_::Float64
    ω_::Float64
    dt::Float64
    prev_v_::Float64
    prev_ω_::Float64
    estimator_::AbstractEstimator
    puddle_coeff::Float64
    puddle_depth_::Float64
    total_reward_::Float64
    in_goal_::Bool
    final_value_::Float64
    goal::Goal
    reso::Vector{Float64}
    pose_min::Vector{Float64}
    pose_max::Vector{Float64}
    index_nums::Vector{Int64}
    policy_data::AbstractArray{Float64,5}
    dev_borders::Vector{Float64}
end

function AMDPPolicyAgent(
    v::Float64,
    ω::Float64,
    dt::Float64,
    estimator::AbstractEstimator,
    goal::Goal,
    reso::Vector{Float64};
    lowerleft = [-4.0, -4.0],
    upperright = [4.0, 4.0],
    puddle_coeff = 100,
    dev_borders = [0.1, 0.2, 0.4, 0.8],
)
    pose_min = vcat(lowerleft, [0.0])
    pose_max = vcat(upperright, [2pi])
    index_nums = [convert(Int64, round((pose_max[i] - pose_min[i]) / reso[i])) for i = 1:3]
    push!(index_nums, length(dev_borders) + 1)
    policy_data = zeros(Float64, index_nums..., 2)
    return AMDPPolicyAgent(
        v,
        ω,
        dt,
        0.0,
        0.0,
        estimator,
        puddle_coeff,
        0.0,
        0.0,
        false,
        0.0,
        goal,
        reso,
        pose_min,
        pose_max,
        index_nums,
        policy_data,
        dev_borders,
    )
end

function init_policy(agent::AMDPPolicyAgent, fname = "policy.txt")
    policy_data = agent.policy_data
    fd = open(fname, "r")
    lines = readlines(fd)
    for line in lines
        tokens = split(line, " ")
        tokens = map(x -> parse(Float64, x), tokens)
        id1, id2, id3, id4 = map(x -> convert(Int64, x), tokens[1:4])
        v, ω = tokens[5], tokens[6]
        policy_data[id1, id2, id3, id4, :] = [v, ω]
    end
end

function policy(agent::AMDPPolicyAgent)
    reso = agent.reso
    pose = agent.estimator_.pose_
    pose_min = agent.pose_min
    index_nums = agent.index_nums
    # normalize index
    index = [convert(Int64, round((pose[i] - pose_min[i]) / reso[i])) for i = 1:3]
    while index[3] < 1
        index[3] += index_nums[3]
    end
    while index[3] > index_nums[3]
        index[3] -= index_nums[3]
    end
    for i in [1, 2]
        if index[i] < 1
            index[i] = 1
        end
        if index[i] > index_nums[i]
            index[i] = index_nums[i]
        end
    end
    belief_index = cov_to_index(agent.dev_borders, cov(agent.estimator_.belief_))
    push!(index, belief_index)
    a = agent.policy_data[index..., :]
    return a
end

max_stars_repo_name: soblin/JuliaProbo

max_issues_repo_name: soblin/JuliaProbo

max_forks_repo_name: soblin/JuliaProbo