text:
"""测试不同的粒度，bubble的区别"""

using Plots
using MARY_fRG.Basics
using MARY_fRG.Drawers
using MARY_fRG.Fermi
using MARY_fRG.Fermi.Patch
using MARY_fRG.Fermi.Kagome
using MARY_fRG.FlowEquation

"""筛选tris"""
function fliter(Γ4::Gamma4{T, P}, lval) where {T, P}
    lamb = Γ4.λ_0 * exp(-lval)
    resltris::Vector{P} = []
    reslpats::Vector{Int64} = []
    cert = 1#length(Γ4.ltris) > 2000000 ? 5 : 15
    mintri = missing
    minpat = missing
    mineng = 10.
    for (tri, pat) in zip(Γ4.ltris, Γ4.lpats)
        if pat != 1
            continue
        end
        engs = [Γ4.model.dispersion[bidx](
            tri.center.x, tri.center.y) / lamb for bidx in 1:1:Γ4.model.bandnum
        ]
        #如果还有可能有贡献
        if minimum(abs.(engs)) < cert #&& minimum(abs.(engs2)) > cert # && pat == 1# && minimum(abs.(engs)) > 0.1
            push!(resltris, tri)
            push!(reslpats, pat)
        end
        #if minimum(abs.(engs)) < mineng
        #    mintri = tri
        #    minpat = pat
        #    mineng = minimum(abs.(engs))
        #end
    end
    #
    #resltris = [mintri]#resltris[1:1]
    #reslpats = [minpat]#reslpats[1:1]
    #
    ltris_pat = Vector{typeof(resltris)}(undef, Γ4.patchnum)
    for idx in 1:1:Γ4.patchnum
        ltris_pat[idx] = []
    end
    #
    for (tri, pat) in zip(resltris, reslpats)
        push!(ltris_pat[pat], tri)
    end
    #
    return Gamma4(
        Γ4.model,
        Γ4.λ_0,
        Γ4.V,
        Γ4.k4tab,
        Γ4.patchnum,
        Γ4.patches,
        resltris, reslpats, nothing, ltris_pat
    )
end

function refine(Γ4::Gamma4{T, P}) where {T, P}
    resltris::Vector{P} = []
    reslpats::Vector{Int64} = []
    resltris = refine_list_triangle(Γ4.ltris, 5)
    reslpats = group_ltris_into_patches_mt(resltris, Γ4.model.brillouin, Γ4.patchnum)
    #
    #resltris = [mintri]#resltris[1:1]
    #reslpats = [minpat]#reslpats[1:1]
    #
    ltris_pat = Vector{typeof(resltris)}(undef, Γ4.patchnum)
    for idx in 1:1:Γ4.patchnum
        ltris_pat[idx] = []
    end
    #
    for (tri, pat) in zip(resltris, reslpats)
        push!(ltris_pat[pat], tri)
    end
    #
    return Gamma4(
        Γ4.model,
        Γ4.λ_0,
        Γ4.V,
        Γ4.k4tab,
        Γ4.patchnum,
        Γ4.patches,
        resltris, reslpats, nothing, ltris_pat
    )
end

function run_rftf()
    model = common_square_lattice(0.20)
    disp = model.dispersion[1]
    lval = 15.
    #500
    Γ4_5 = TFGamma4(
        model, 8.0, 16, 500
    )
    Γ4_5.V .+= 1.0
    Γ4_5_1 = fliter(Γ4_5, 11.0)
    Γ4_5_2 = fliter(Γ4_5, 12.0)
    println(length(Γ4_5_1.ltris))
    println(length(Γ4_5_2.ltris))
    #fliter2
    bubbval_fs = Array{Float64, 2}(undef, Γ4_5_2.patchnum, Γ4_5_2.patchnum)
    lamb = Γ4_5_2.λ_0 * exp(-lval)
    println("lamb ", lamb)
    Threads.@threads for idxs in CartesianIndices(bubbval_fs)
        i2, i3 = Tuple(idxs)
        k2, k3 = Γ4_5_2.patches[1][i2], Γ4_5_2.patches[1][i3]
        q_fs = Γ4_5_2.model.kadd(k3, -k2)
        bubbres = pi_αβ_plus_tf(
            Γ4_5_2.ltris, area(Γ4_5_2.model.brillouin),
            lamb, q_fs, Γ4_5_2.model.dispersion[1],
            Γ4_5_2.model.dispersion[1], Γ4_5_2.model.kadd
        )
        bubbval_fs[i2, i3] = bubbres
    end
    plt = heatmap(1e5*bubbval_fs)
    @savepng plt "bubb_5_2_2"
    #fliter1
    bubbval_fs = Array{Float64, 2}(undef, Γ4_5_1.patchnum, Γ4_5_1.patchnum)
    lamb = Γ4_5_1.λ_0 * exp(-lval)
    Threads.@threads for idxs in CartesianIndices(bubbval_fs)
        i2, i3 = Tuple(idxs)
        k2, k3 = Γ4_5_1.patches[1][i2], Γ4_5_1.patches[1][i3]
        q_fs = Γ4_5_1.model.kadd(k3, -k2)
        bubbres = pi_αβ_plus_tf(
            Γ4_5_1.ltris, area(Γ4_5_1.model.brillouin),
            lamb, q_fs, Γ4_5_1.model.dispersion[1],
            Γ4_5_1.model.dispersion[1], Γ4_5_1.model.kadd
        )
        bubbval_fs[i2, i3] = bubbres
    end
    plt = heatmap(1e5*bubbval_fs)
    @savepng plt "bubb_5_2_1"
    #fliter3
    ltris3::Vector{Basics.AbstractTriangle{:RT}} = []
    ltris4::Vector{Basics.AbstractTriangle{:RT}} = []
    for tri in Γ4_5_1.ltris
        if !any(map((x)->x==tri, Γ4_5_2.ltris))
            push!(ltris4, tri)
        end
    end
    #push!(ltris3, Γ4_5_2.ltris[1])
    #println(Γ4_5_2.ltris[1])
    #println(disp(Γ4_5_2.ltris[1].center.x, Γ4_5_2.ltris[1].center.y))
    #push!(ltris3, Γ4_5_2.ltris[2])
    push!(ltris3, Γ4_5_2.ltris[3])
    #push!(ltris3, ltris4[1])
    push!(ltris3, ltris4[2])
    println(disp(ltris4[2].center.x, ltris4[2].center.y))
    #push!(ltris3, ltris4[3])
    println(length(ltris3))
    bubbval_fs = Array{Float64, 2}(undef, Γ4_5_1.patchnum, Γ4_5_1.patchnum)
    lamb = Γ4_5_1.λ_0 * exp(-lval)
    Threads.@threads for idxs in CartesianIndices(bubbval_fs)
        i2, i3 = Tuple(idxs)
        k2, k3 = Γ4_5_1.patches[1][i2], Γ4_5_1.patches[1][i3]
        q_fs = Γ4_5_1.model.kadd(k3, -k2)
        bubbres = pi_αβ_plus_tf(
            ltris3, area(Γ4_5_1.model.brillouin),
            lamb, q_fs, Γ4_5_1.model.dispersion[1],
            Γ4_5_1.model.dispersion[1], Γ4_5_1.model.kadd
        )
        bubbval_fs[i2, i3] = bubbres
    end
    plt = heatmap(1e5*bubbval_fs)
    @savepng plt "bubb_5_2_3"
    return
    #plt = draw_points([tri.center for tri in Γ4_5.ltris])
    #abses = [abs(disp(tri.center.x, tri.center.y)) for tri in Γ4_5.ltris]
    #println(minimum(abses))
    #abscount = 0
    #for ae in abses
    #    if ae < 0.01
    #        abscount += 1
    #    end
    #end
    #println(abscount)
    #@savepng plt "diffe1"
    println(Γ4_5.ltris[1])
    bubb_pp_5, bubb_fs_5, bubb_ex_5 = all_bubble_tf_mt(Γ4_5, lval)
    println("500时的贡献 ", sum(abs.(bubb_fs_5.V)))
    println("500时最大", findmax(abs.(bubb_fs_5.V)))
    k1, k2 = Γ4_5.patches[1][12], Γ4_5.patches[1][4]
    println(k1)
    println(k2)
    kprim = Γ4_5.model.kadd(k1, k2)
    println(kprim)
    kprim = Γ4_5.model.kadd(kprim, -Γ4_5.ltris[1].center)
    println(kprim)
    println("kprim disp ", Γ4_5.model.dispersion[1](kprim.x, kprim.y))

    return
    ###
    #100
    Γ4_1 = TFGamma4(
        model, 8.0, 16, 800
    )
    Γ4_1.V .+= 1.0
    #Γ4_1 = fliter_away_surface(Γ4_1, 10.0)
    #plt = draw_points([tri.center for tri in Γ4_1.ltris])
    #@savepng plt "diffe3"
    #Γ4_1 = TFGamma4_addition_ltris_mt(Γ4_1, 10.0)
    Γ4_1 = fliter(Γ4_1, 12.0)
    #println(Γ4_1.ltris)
    #plt = draw_polygon(Γ4_1.ltris)
    #plot!(plt, aspect_ratio=:equal)
    #@savepng plt "diffe4"
    Γ4_1 = refine(Γ4_1)
    #println(Γ4_1.ltris)
    plt = draw_polygon(Γ4_1.ltris)
    #plot!(plt, aspect_ratio=:equal)
    #@savepng plt "diffe5"
    Γ4_1 = fliter(Γ4_1, 12.0)
    abses = [abs(disp(tri.center.x, tri.center.y)) for tri in Γ4_1.ltris]
    println(minimum(abses))
    abscount = 0
    for ae in abses
        if ae < 0.01
            abscount += 1
        end
    end
    println(abscount)
    #println(minimum([abs(disp(tri.center.x, tri.center.y)) for tri in Γ4_1.ltris]))
    plt = draw_points([tri.center for tri in Γ4_1.ltris])
    @savepng plt "diffe2"
    println(Γ4_1.ltris[1])
    k1, k2 = Γ4_1.patches[1][11], Γ4_1.patches[1][3]
    println(k1)
    println(k2)
    kprim = Γ4_1.model.kadd(k1, k2)
    println(kprim)
    kprim = Γ4_1.model.kadd(kprim, -Γ4_1.ltris[1].center)
    println(kprim)
    println("kprim disp ", Γ4_1.model.dispersion[1](kprim.x, kprim.y))
    bubb_pp_1, bubb_fs_1, bubb_ex_1 = all_bubble_tf_mt(Γ4_1, lval)
    println("200时的贡献 ", sum(abs.(bubb_pp_1.V)))
    println("200时最大", findmax(abs.(bubb_pp_1.V)))
    #500和100的差
    diffe = bubb_pp_5.V - bubb_pp_1.V
    diffe = abs.(diffe)
    println(sum(diffe))
    println(sum(abs.(bubb_pp_1.V)))
    #refine
    #Γ4_r = TFGamma4_addition_ltris_mt(Γ4_1, lval)
    #bubb_pp_r, bubb_fs_r, bubb_ex_r = all_bubble_tf_mt(Γ4_r, lval)
    ##500和100的差
    #diffe = bubb_pp_5.V - bubb_pp_r.V
    #diffe = abs.(diffe)
    #println(sum(diffe))
end

#run_rftf()

function run_rf_ult()
    #model = common_square_lattice(0.20)
    model = upperband_kagome_lattice(-0.2)
    disp = model.dispersion[1]
    lval = 17.
    #500
    Γ4_5 = TFGamma4(
        model, 8.0, 24, 500
    )
    Γ4_5.V .+= 1.0
    #Γ4_5 = fliter(Γ4_5, 12.0)
    #plt = draw_points([tri.center for tri in Γ4_5.ltris])
    #abses = [abs(disp(tri.center.x, tri.center.y)) for tri in Γ4_5.ltris]
    #println(minimum(abses))
    #abscount = 0
    #for ae in abses
    #    if ae < 0.01
    #        abscount += 1
    #    end
    #end
    #println(abscount)
    #@savepng plt "diffe1"
    println(Γ4_5.ltris[1])
    bubb_pp_5, bubb_fs_5, bubb_ex_5 = all_bubble_tf_mt(Γ4_5, lval)
    println("500时的贡献 ", sum(abs.(bubb_pp_5.V)))
    println("500时最大", findmax(abs.(bubb_pp_5.V)))
    ###
    #100
    Γ4_1 = TFGamma4(
        model, 8.0, 24, 200
    )
    Γ4_1.V .+= 1.0
    bubb_pp_1, bubb_fs_1, bubb_ex_1 = all_bubble_tf_mt(Γ4_1, lval)
    println("200时的贡献 ", sum(abs.(bubb_pp_1.V)))
    println("200时最大", findmax(abs.(bubb_pp_1.V)))
    ###
    #ult
    Γ4_1 = fliter_away_surface(Γ4_1, 6.0)
    Γ4_1 = refine_to_surface(Γ4_1)
    while length(Γ4_1.ltris) < 1000000
        Γ4_1 = refine_to_surface(Γ4_1)
    end
    plt = draw_points([tri.center for tri in Γ4_1.ltris])
    @savepng plt "diffe6"
    println("tri大小", length(Γ4_1.ltris))
    maxi, mini = engpeak_to_surface(Γ4_1)
    println(maxi, " ", mini)
    bubb_pp_u, bubb_fs_u, bubb_ex_u = all_bubble_tf_ult_mt(Γ4_1, maxi)
    println("ult时的贡献 ", sum(abs.(bubb_pp_u.V)))
    println("ult时最大", findmax(abs.(bubb_pp_u.V)))
    #500和100的差
    diffe = bubb_pp_5.V - bubb_pp_1.V
    diffe = abs.(diffe)
    println(sum(diffe))
    println(sum(abs.(bubb_pp_1.V)))
    #500个ult的差
    diffe = bubb_pp_5.V - bubb_pp_u.V
    diffe = abs.(diffe)
    println(sum(diffe))
    println(sum(abs.(bubb_pp_u.V)))
    #ult和500的画图
    plt = heatmap(1e5*bubb_pp_u.V[1, 1, 1, 1, 1, :, :])
    png(plt, "bubb_u1")
    plt = heatmap(1e5*bubb_pp_u.V[1, 1, 1, 1, 2, :, :])
    png(plt, "bubb_u2")
    plt = heatmap(1e5*bubb_pp_u.V[1, 1, 1, 1, 3, :, :])
    png(plt, "bubb_u3")
    plt = heatmap(1e5*bubb_pp_u.V[1, 1, 1, 1, 4, :, :])
    png(plt, "bubb_u4")
    plt = heatmap(1e5*bubb_pp_5.V[1, 1, 1, 1, 1, :, :])
    png(plt, "bubb_51")
    plt = heatmap(1e5*bubb_pp_5.V[1, 1, 1, 1, 2, :, :])
    png(plt, "bubb_52")
    plt = heatmap(1e5*bubb_pp_5.V[1, 1, 1, 1, 3, :, :])
    png(plt, "bubb_53")
    plt = heatmap(1e5*bubb_pp_5.V[1, 1, 1, 1, 4, :, :])
    png(plt, "bubb_54")
    println(Γ4_1.patches[1][1], Γ4_1.patches[1][4], Γ4_1.patches[1][9])
end

run_rf_ult()


text:
using Sockets

function matgen(n, seed)
    tmp = seed / n / n
    [tmp * (i - j) * (i + j - 2) for i = 1:n, j = 1:n]
end

function calc(n)
    n = round(Int, n / 2) * 2
    a = matgen(n, 1.0)
    b = matgen(n, 2.0)
    c = a * b
    c[Int(n / 2)+1, Int(n / 2)+1]
end

function notify(msg)
    try
        socket = connect("localhost", 9001)
        write(socket, msg)
        close(socket)
    catch
        # standalone usage
    end
end

if abspath(PROGRAM_FILE) == @__FILE__
    n = length(ARGS) > 0 ? parse(Int, ARGS[1]) : 100

    left = calc(101)
    right = -18.67
    if abs(left - right) > 0.1
        println(stderr, "$(left) != $(right)")
        exit(1)
    end

    # Assuming openblas64, may not work for all environments
    num_threads = ccall((:openblas_get_num_threads64_, Base.libblas_name), Int32, ())
    notify("Julia (threads: $(num_threads))\t$(getpid())")

    t = time()
    results = calc(n)
    elapsed = time() - t

    notify("stop")

    println(results)
    println("time: $(elapsed) s")
end

ext: jl

max_stars_repo_name: clemenswasser/benchmarks

max_issues_repo_name: clemenswasser/benchmarks

max_forks_repo_name: clemenswasser/benchmarks


text:
using Documenter, MultivariateFunctions

makedocs(
    format = Documenter.HTML(),
    sitename = "MultivariateFunctions",
    modules = [MultivariateFunctions],
    pages = ["index.md",
             "1_structs_and_limitations.md",
             "2_interpolation_methods.md",
             "3_approximation_methods.md",
             "4_examples_algebra.md",
             "5_examples_interpolation.md",
             "6_examples_approximation.md",
             "99_refs.md"]
)

deploydocs(
    repo   = "github.com/s-baumann/MultivariateFunctions.jl.git",
    target = "build",
    deps   = nothing,
    make   = nothing
)


text:
"""

`setdiff(dim, arr1, arr2)`

Take difference of two arrays of type `LabeledArray`/`DictArray`/`AbstractArrayWrapper`.
The order is preserved in the sense that the first array order is kept except that only duplicate elements in the first array have been removed.

##### Arguments

* `dim` : direction along which to concat.
* `arr1` : array 1.
* `arr2` : array 2. The return value is array1 - array2, where arrays are treated as sets.

##### Examples

```julia
julia> setdiff(1, nalift([1 2 3;4 5 6]), nalift([1 2 3;5 5 6]))
1x3 DataCubes.AbstractArrayWrapper{Nullable{Int64},2,Array{Nullable{Int64},2}}:
 Nullable(4)  Nullable(5)  Nullable(6)

julia> setdiff(2, nalift([1 2 3;4 5 6]), nalift([1 2 3;5 5 6]))
2x1 DataCubes.AbstractArrayWrapper{Nullable{Int64},2,Array{Nullable{Int64},2}}:
 Nullable(1)
 Nullable(4)

julia> setdiff(1, darr(a=[:x,:y,:z]), darr(a=[:x,:x,:u]))
2 DictArray

a
--
y
z

julia> setdiff(1, larr(a=[1 2 3;4 5 6], axis1=[:X,:Y]), larr(a=[1 2 3;4 3 2], axis1=[:X,:Y]))
1 x 3 LabeledArray

  |1 |2 |3
--+--+--+--
  |a |a |a
--+--+--+--
Y |4 |5 |6

julia> setdiff(1, larr(a=[1 2 3;4 5 6], axis1=[:X,:Y]), larr(a=[1 2 3;4 3 2], axis1=[:Z,:Y]))
2 x 3 LabeledArray

  |1 |2 |3
--+--+--+--
  |a |a |a
--+--+--+--
X |1 |2 |3
Y |4 |5 |6
```

"""
Base.setdiff(dim::Int, arr1::LabeledArray, arr2::LabeledArray) = begin
  # first rearrange the axes so that the calculatioin can be more convenient.
  # this might be later optimized.
  # all arrays will be of the same dimensions.
  arrs = (arr1, arr2)
  ndimsarr = ndims(arr1)
  remainingdirs = filter(x->x != dim, 1:ndimsarr)
  remainingcoords = ntuple(d->Colon(), length(remainingdirs))
  newdirs = (remainingdirs...,dim)
  permuted_arrs = map(arr->permutedims_if_necessary(arr, newdirs), arrs)
  elems = Dict{Any,Int}()
  isdefaxis = isa(permuted_arrs[1].axes[ndimsarr], DefaultAxis)
  permuted_arrs1 = permuted_arrs[1]
  lastsize = size(permuted_arrs[1], ndimsarr)
  if isdefaxis
    for i in 1:lastsize
      full_coords = (remainingcoords..., i)
      oneslice = getindexvalue(permuted_arrs1.data, full_coords...)
      push!(elems, oneslice => i)
    end
  else
    for i in 1:lastsize
      full_coords = (remainingcoords..., i)
      oneslice = getindexvalue(permuted_arrs1.data, full_coords...)
      labels = permuted_arrs1.axes[ndimsarr][i]
      push!(elems, (oneslice,labels) => i)
    end
  end
  if isdefaxis
    parrs2 = permuted_arrs[2]
    for i in 1:size(parrs2,ndimsarr)
      slice = getindexvalue(parrs2.data, remainingcoords..., i)
      if haskey(elems, slice)
        delete!(elems, slice)
      end
    end
  else
    parrs2 = permuted_arrs[2]
    for i in 1:size(parrs2,ndimsarr)
      slice = getindexvalue(parrs2.data, remainingcoords..., i)
      labels = parrs2.axes[ndimsarr][i]
      pair = (slice, labels)
      if haskey(elems, pair)
        delete!(elems, pair)
      end
    end
  end
  #setdiff_slices = setdiff(result_arrs...)
  coords_toshow = sort(collect(values(elems))) #[elems[slice] for slice in setdiff_slices]
  result = permuted_arrs[1][remainingcoords..., coords_toshow]
  ipermutedims_if_necessary(result, newdirs)
end

Base.setdiff(dim::Int, arr1::DictArray, arr2::DictArray) = begin
  arrs = (arr1, arr2)
  # first rearrange the axes so that the calculatioin can be more convenient.
  # this might be later optimized.
  # all arrays will be of the same dimensions.
  ndimsarr = ndims(arrs[1])
  remainingdirs = filter(x->x != dim, 1:ndimsarr)
  remainingcoords = ntuple(d->Colon(), length(remainingdirs))
  newdirs = (remainingdirs...,dim)
  permuted_arrs = map(arr->permutedims_if_necessary(arr, newdirs), arrs)
  result_arrs = map(permuted_arrs) do arr
    map(1:size(arr, ndimsarr)) do i
      full_coords = (remainingcoords..., i)
      getindex(arr, full_coords...)
    end
  end
  setdiff_slices = setdiff(result_arrs...)
  result = similar(permuted_arrs[1],
                   (size(permuted_arrs[1])[1:end-1]...,
                    length(setdiff_slices)))
  for (i,arr) in enumerate(setdiff_slices)
    result[remainingcoords...,i] = arr
  end
  ipermutedims_if_necessary(result, newdirs)
end

Base.setdiff{T,U}(dim::Int,
                  arr1::AbstractArrayWrapper{Nullable{T}},
                  arr2::AbstractArrayWrapper{Nullable{U}}) = begin
  arrs = (arr1, arr2)
  # first rearrange the axes so that the calculatioin can be more convenient.
  # this might be later optimized.
  # all arrays will be of the same dimensions.
  ndimsarr = ndims(arrs[1])
  remainingdirs = filter(x->x != dim, 1:ndimsarr)
  remainingcoords = ntuple(d->Colon(), length(remainingdirs))
  newdirs = (remainingdirs...,dim)
  permuted_arrs = map(arr->permutedims_if_necessary(arr, newdirs), arrs)
  result_arrs = map(permuted_arrs) do arr
    map(1:size(arr, ndimsarr)) do i
      full_coords = (remainingcoords..., i)
      getindex(arr, full_coords...)
    end
  end
  setdiff_slices = setdiff(result_arrs...)
  result = similar(permuted_arrs[1],
                   (size(permuted_arrs[1])[1:end-1]...,
                    length(setdiff_slices)))
  for (i,arr) in enumerate(setdiff_slices)
    result[remainingcoords...,i] = arr
  end
  ipermutedims_if_necessary(result, newdirs)
end


text:
module BayesianRegression

#=

### Model and prior

Assume a linear model
$$
f(\mathbf{x}) = \sum_{i=1}^M c_i \phi_i(\mathbf{x}) = \mathbf{c}^T \boldsymbol{\phi}(\mathbf{x} )
$$
where the elements of the $M$-dimensional coefficient vector have independent priors $p(c_i)=\mathcal{N}(c_i | 0,\sigma_i^2)$. Suppose $N$ noisy input-target pairs are available, $\{\mathbf{x}_i, y_i\}_{i=1}^N$, such that
$$
y_i = f(\mathbf{x}_i) + \epsilon_i,
$$
with noise prior $p(\epsilon_i) = \mathcal{N}(\epsilon_i | 0,\sigma_\epsilon^2)$.

### Posterior

For fixed values of $\{\sigma_i^2\}_{i=1}^M$ and $\sigma_\epsilon^2$, the posterior mean used for predictions is
$$
\boldsymbol{\mu}_c = \sigma_\epsilon^{-2} \boldsymbol{\Sigma}_c \boldsymbol{\Phi}^T \mathbf{y}
$$
with associated covariance
$$
\begin{aligned}
\boldsymbol{\Sigma}_c
&= \left[ \sigma_\epsilon^{-2} \boldsymbol{\Phi}^T \boldsymbol{\Phi} + \boldsymbol{\Sigma}_0^{-1} \right]^{-1} \\
&= \boldsymbol{\Sigma}_0 - \boldsymbol{\Sigma}_0 \boldsymbol{\Phi}^T \boldsymbol{\Sigma}_y^{-1} \boldsymbol{\Phi} \boldsymbol{\Sigma}_0
\end{aligned} ,
$$
where $\boldsymbol{\Phi}$ is the design matrix
$$
\boldsymbol{\Phi} = \left[
\begin{matrix}
\text{---} \boldsymbol{\phi}(\mathbf{x}_1) \text{---} \\
\text{---} \boldsymbol{\phi}(\mathbf{x}_2) \text{---} \\
\vdots \\
\text{---} \boldsymbol{\phi}(\mathbf{x}_3) \text{---}
\end{matrix}
\right] ,
$$
$\boldsymbol{\Sigma}_0$ is diagonal with $[\boldsymbol{\Sigma}_0]_{ii} = \sigma_i^2$, and $\boldsymbol{\Sigma}_y$ is defined below.

### Likelihood and log likelihood

With flat hyperpriors, the marginal likelihood---the "evidence" for the model---is
$$
p(\mathbf{y} | \boldsymbol{\Sigma}_0, \sigma_\epsilon^2) = \mathcal{N}(\mathbf{y} | \mathbf{0}, \boldsymbol{\Sigma}_y)
$$
with
$$
\begin{aligned}
\boldsymbol{\Sigma}_y
&= \boldsymbol{\Phi} \boldsymbol{\Sigma}_0 \boldsymbol{\Phi}^T + \sigma_\epsilon^2 \mathbf{I} \\
&= \left[ \sigma_\epsilon^{-2} \mathbf{I} - \sigma_\epsilon^{-2} \boldsymbol{\Phi} \boldsymbol{\Sigma}_c \boldsymbol{\Phi}^T \sigma_\epsilon^{-2}  \right]^{-1}
\end{aligned}
$$
When maximizing this quantity, one typically works with its log,
$$
\begin{aligned}
\mathrm{\mathcal{L}}
&= \ln p(\mathbf{y} | \boldsymbol{\Sigma}_0, \sigma_\epsilon^2) \\
&= -\frac{1}{2} \mathbf{y}^T \boldsymbol{\Sigma}_y^{-1} \mathbf{y}
    - \frac{1}{2} \ln \left| \boldsymbol{\Sigma}_y \right|
    - \frac{N}{2} \ln (2\pi) \\
&= -\frac{1}{2} \sigma_\epsilon^{-2} \mathbf{y}^T \left[  \mathbf{y} - \boldsymbol{\Phi} \boldsymbol{\mu}_c \right]
    + \frac{1}{2} \ln \left| \boldsymbol{\Sigma}_c \right|
    - \frac{1}{2} \ln \left| \boldsymbol{\Sigma}_0 \right|
    - \frac{N}{2} \ln (\sigma_\epsilon^2)
    - \frac{N}{2} \ln (2\pi) \\
\end{aligned} .
$$
The relevant derivatives with respect to the hyperparameters are
$$
\begin{aligned}
\frac{\partial \mathcal{L}}{\partial \sigma_i^2}
&= \frac{1}{2} [\boldsymbol{\Phi}^T \boldsymbol{\Sigma}_y^{-1} \mathbf{y}]_i^2
    - \frac{1}{2} [\boldsymbol{\Phi}^T \boldsymbol{\Sigma}_y^{-1} \boldsymbol{\Phi}]_{ii} \\
&= \frac{1}{2} \sigma_i^{-4} \left[ \mu_i^2 + [\boldsymbol{\Sigma}_c]_{ii} - \sigma_i^2 \right]
\end{aligned}
$$
and
$$
\begin{aligned}
\frac{\partial \mathcal{L}}{\partial \sigma_\epsilon^2}
& = \frac{1}{2} \left\| \boldsymbol{\Sigma}_y^{-1} \mathbf{y} \right\|^2 - \frac{1}{2} \mathrm{tr}\left(\boldsymbol{\Sigma}_y^{-1}\right) \\
&= \frac{1}{2} \sigma_\epsilon^{-4} \left[
    \left\| \mathbf{y} - \boldsymbol{\Phi} \boldsymbol{\mu}_c \right\|^2
    + \mathrm{tr}\left( \boldsymbol{\Phi} \boldsymbol{\Sigma}_c \boldsymbol{\Phi}^T \right)
    - N \sigma_\epsilon^2
\right]
\end{aligned} .
$$

### References

* M. E. Tipping, "Sparse Bayesian learning and the relevance vector machine," Journal of Machine Learning Research 1, 211 (2001)
* C. K. Williams and C. E. Rasmussen, _Gaussian Processes for Machine Learning_, MIT Press (2006)

=#

using LinearAlgebra
using Optim

function solve(
    y::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    var_c::AbstractFloat,
    var_e::AbstractFloat,
)
    return solve(y, X, var_c*ones(size(X,2)), var_e)
end

function solve(
    y::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    var_c::Vector{<:AbstractFloat},
    var_e::AbstractFloat,
)
    M = size(X,2)
    XTX = X'*X
    Σ_c = Array{Float64}(undef,M,M)
    BLAS.blascopy!(length(Σ_c), XTX, stride(XTX,1), Σ_c, stride(Σ_c,1))
    BLAS.scal!(1.0/var_e, Σ_c)
    for i in 1:M; Σ_c[i,i] += 1.0/var_c[i]; end
    C = cholesky!(Symmetric(Σ_c))
    return 1.0/var_e*(C \ (X'*y))
end

function log_marginal_likelihood_overdetermined!(
    lml::AbstractFloat,
    grad::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    y::Vector{<:AbstractFloat},
    var_c::AbstractFloat,
    var_e::AbstractFloat,
    XTX::Matrix{<:AbstractFloat},
)
    var_c_vec = var_c*ones(size(X,2))
    grad_vec = zeros(size(X,2)+1)
    lml = log_marginal_likelihood_overdetermined!(lml, grad_vec, X, y, var_c_vec, var_e, XTX)
    grad[1] = sum(grad_vec[1:end-1])
    grad[2] = grad_vec[end]
    return lml
end

function log_marginal_likelihood_overdetermined!(
    lml::AbstractFloat,
    grad::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    y::Vector{<:AbstractFloat},
    var_c::Vector{<:AbstractFloat},
    var_e::AbstractFloat,
    XTX::Matrix{<:AbstractFloat},
)
    N = size(X,1)
    M = size(X,2)
    Σ_c = Array{Float64}(undef,M,M)
    BLAS.blascopy!(length(Σ_c), XTX, stride(XTX,1), Σ_c, stride(Σ_c,1))
    BLAS.scal!(1.0/var_e, Σ_c)
    for i in 1:M; Σ_c[i,i] += 1.0/var_c[i]; end
    C = cholesky!(Symmetric(Σ_c))
    Σ_c = C \ I(M)
    μ_c = 1.0/var_e*(C \ (X'*y))
    lml = - 0.5*logdet(C) - 0.5*sum(log.(var_c)) - 0.5*N*log(var_e) - 0.5*N*log(2*π)
    lml -= 0.5/var_e*y'*(y-X*μ_c)
    grad[1:M] .= 0.5*(μ_c.^2 .+ diag(Σ_c) .- var_c)./var_c.^2
    grad[M+1] = 0.5/var_e^2*(sum((y-X*μ_c).^2) + dot(XTX,Σ_c) - N*var_e)
    return lml
end

function log_marginal_likelihood_underdetermined!(
    lml::AbstractFloat,
    grad::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    y::Vector{<:AbstractFloat},
    var_c::AbstractFloat,
    var_e::AbstractFloat,
)
    var_c_vec = var_c*ones(size(X,2))
    grad_vec = zeros(size(X,2)+1)
    lml = log_marginal_likelihood_underdetermined!(lml, grad_vec, X, y, var_c_vec, var_e)
    grad[1] = sum(grad_vec[1:end-1])
    grad[2] = grad_vec[end]
    return lml
end

function log_marginal_likelihood_underdetermined!(
    lml::AbstractFloat,
    grad::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    y::Vector{<:AbstractFloat},
    var_c::Vector{<:AbstractFloat},
    var_e::AbstractFloat,
)
    N = size(X,1)
    M = size(X,2)
    Σ_y = X*Diagonal(var_c)*X'
    for i=1:N; Σ_y[i,i] += var_e; end
    C = cholesky!(Symmetric(Σ_y))
    invΣy_y = C \ y
    lml = -0.5*y'*invΣy_y - 0.5*logdet(C) - 0.5*N*log(2*π)
    grad[1:M] .= 0.5*(X'*invΣy_y).^2
    W = C \ X
    @views for i=1:M; grad[i] -= 0.5*dot(X[:,i], W[:,i]); end
    grad[M+1] = 0.5*dot(invΣy_y,invΣy_y) - 0.5*tr(C\I(N))
    return lml
end

function bayesian_fit(
    y::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat};
    variance_floor::AbstractFloat=1e-8,
    verbose::Bool=false,
)
    if size(X,1) > size(X,2)
        XTX = X'*X  # advantageous to precompute for overdetermined case
    end

    function fg!(f, g, x)
        var_c = variance_floor + x[1]*x[1]
        var_e = variance_floor + x[2]*x[2]
        if size(X,1) >= size(X,2)
            f = log_marginal_likelihood_overdetermined!(f, g, X, y, var_c, var_e, XTX)
        else
            f = log_marginal_likelihood_underdetermined!(f, g, X, y, var_c, var_e)
        end
        if f != nothing
            f = -f
        end
        if g != nothing
            g .*= -2*x
        end
        return f
    end

    res = optimize(Optim.only_fg!(fg!),
                   ones(2),
                   Optim.LBFGS(),
                   Optim.Options(x_tol=1e-4, g_tol=0.0, show_trace=verbose))
    verbose && println(res)

    lml = -Optim.minimum(res)
    var_c, var_e = Optim.minimizer(res)
    var_c = variance_floor + var_c*var_c
    var_e = variance_floor + var_e*var_e

    return solve(y, X, var_c, var_e), var_c, var_e, lml
end

function ard_fit(
    y::Vector{<:AbstractFloat},
    X::Matrix{<:AbstractFloat},
    variance_floor::AbstractFloat=1e-8;
    verbose::Bool=false,
)
    if size(X,1) >= size(X,2)
        XTX = X'*X  # advantageous to precompute for overdetermined case
    end

    function fg!(f, g, x)
        var_c = variance_floor .+ x[1:end-1].*x[1:end-1]
        var_e = variance_floor + x[end]*x[end]
        if size(X,1) >= size(X,2)
            f = log_marginal_likelihood_overdetermined!(f, g, X, y, var_c, var_e, XTX)
        else
            f = log_marginal_likelihood_underdetermined!(f, g, X, y, var_c, var_e)
        end
        if f != nothing
            f = -f
        end
        if g != nothing
            g .*= -2*x
        end
        return f
    end

    res = optimize(Optim.only_fg!(fg!),
                   ones(size(X,2)+1),
                   Optim.LBFGS(),
                   Optim.Options(x_tol=1e-4, g_tol=0.0, show_trace=verbose))
    verbose && println(res)

    lml = -Optim.minimum(res)
    x = Optim.minimizer(res)
    var_c = variance_floor .+ x[1:end-1].*x[1:end-1]
    var_e = variance_floor + x[end]*x[end]

    mask = var_c .> 10*variance_floor
    var_c[.~mask] .= 0
    c_mask = solve(y, X[:,mask], var_c[mask], var_e)
    c = zeros(length(var_c))
    c[mask] .= c_mask

    return c, var_c, var_e, lml, mask
end

end


text:
module NMarkov

using LinearAlgebra: qr
# using Printf

using Origin: @origin
using SparseArrays: SparseMatrixCSC, nnz
using SparseMatrix: SparseCSR, SparseCSC, SparseCOO, spdiag
using Distributions: UnivariateDistribution, pdf, cquantile, Normal
using Deformula: deint

export eye
export unif
export gth!, gth, stsen
export stsenguess, stsengs, stsenpower
export stguess, stgs, stpower
export qstgs, qstpower
export mexp, mexpc
export mexp, mexpmix, mexpc, mexpcmix
export tran

include("_common.jl")
include("_special_matrix.jl")
include("_unif.jl")
include("_gth.jl")
include("_stsen_dense.jl")
include("_gsstep.jl")
include("_stationary_iterative.jl")
include("_sensitivity_iterative.jl")
include("_quasistationary_iterative.jl")

include("_poisson.jl")
include("_forward_backward.jl")
include("_mexp.jl")
include("_transient.jl")

include("_mix.jl")

include("_conv.jl")

end # module


text:
abstract type SpecEnv end #global settings for spectral calculations

#=
    Spectral environment
    λmin: lower wavelength boundary for all spectral calculations
            default: 380.0 nm, suitable for the 2012 CIE CMFs
    Δλ : wavelength stepwidth
            default: 1.0 nm, suitable for the 2012 CIE CMFs
    λmax: upper wavelength boundery for all spectral calculations
            default: 830.0 nm, suitable for the 2012 CIE CMFs
    cmf: color matching function used for all spectral calculations
            default: CMF2012_10 (CIE2012 color matching function for a 10° observer)
    ex: extrapolation mode for spectra that have a smaller range than the environment settings
            default: :linear (linear extrapolation)
=#
mutable struct SpecEnvironment <: SpecEnv
    λmin::Real
    Δλ::Real
    λmax::Real
    cmf
    ex::Symbol
    function SpecEnvironment(λmin,Δλ,λmax,cmf,ex)
        new(λmin,Δλ,λmax,cmf,ex)
    end
end

"""
`set_specenv(λmin=390.0,Δλ=1.0,λmax=830.0,extrapolate= :zero)`

Initializes the spectral environment to standard values:

λmin = 380.0 nm
λmax = 830.0 nm
Δλ   =   1.0 nm
cmf  = cmfunc(:cie12_10)
ex = :linear
"""
function set_specenv(λmin = 390.0, Δλ = 1.0, λmax = 830.0, cmf = cmfunc(:cie12_10), ex = :linear)
    SpecEnvironment(λmin,Δλ,λmax,cmf,ex)
end

# Set spectral environment to default settings
SPECENV=set_specenv()

"""
`set_extrap(env::SpecEnvironment,extrapolation::Symbol)`

Sets the extrapolation mode of the environment `env` to one of the available extrapolation settings:

`:none` or `:zero`: out of range values are zero.

`:constant`: out of range values are set to λmin and λmax, respectively.

`:linear`: out of range values are linearly extrapolated through the first/last two data points.

`:parabolic`: out fo range values are extrapolated with a parabola through the first/last three datapoints. *Recommended extrapolation method according to D.L. MacAdams* in "Color Measurement", chapter 5.4 *"Truncations"*.
"""
function set_extrap(env::SpecEnvironment,extrapolation::Symbol)
    extrapolation in (:none, :boundary, :linear, :parabolic) ? env.ex = extrapolation :
    extrapolation == :zero ? set_extrap(env, :none) : throw(DomainError(extrapolation,"Extrapolation mode does not exist!"))
end

"""
`set_extrap(extrapolation::Symbol)`

Sets the extrapolation mode of the environment `env` to one of the available extrapolation settings:

`:none` or `:zero`: out of range values are zero.

`:constant`: out of range values are set to λmin and λmax, respectively.

`:linear`: out of range values are linearly extrapolated through the first/last two data points.

`:parabolic`: out fo range values are extrapolated with a parabola through the first/last three datapoints. *Recommended extrapolation method according to D.L. MacAdams* in "Color Measurement", chapter 5.4 *"Truncations"*.
"""
function set_extrap(extrapolation::Symbol)
    set_extrap(SPECENV, extrapolation)
end

"""
`set_limits(env::SpecEnvironment, λmin::Real,Δλ::Real,λmax::Real)`

sets the wavelength limits and resolution of the spectral environment
"""
function set_limits(env::SpecEnvironment, λmin::Real,Δλ::Real,λmax::Real)
    env.λmin = λmin
    env.Δλ = Δλ
    env.λmax = λmax
end

"""
`set_limits(λmin::Real,Δλ::Real,λmax::Real)`

sets the wavelength limits and resolution of the spectral environment
"""
function set_limits(λmin::Real,Δλ::Real,λmax::Real)
    set_limits(SPECENV, λmin, Δλ, λmax)
end

"""
`set_colormatch(env::SpecEnvironment, colormatch::CMatch)`

Sets the color matching function (CMF) used in a defined spectral environment.

Available CMFs:

`:cie31`, `:cie31_j`, `:cie31_jv`, `:cie64`, `:cie12_2`, `:cie12_10`, `:lms06_2`, `:lms06_10`
"""
function set_colormatch(env::SpecEnvironment, colormatch::Symbol)
    env.cmf = create_cmf(colormatch)
end

"""
`set_colormatch(colormatch::CMatch)`

Sets the color matching function (CMF) used in the spectral environment.

Available CMFs:

`:cie31`, `:cie31_j`, `:cie31_jv`, `:cie64`, `:cie12_2`, `:cie12_10`, `:lms06_2`, `:lms06_10`
"""
function set_colormatch(colormatch::Symbol)
    set_colormatch(SPECENV, colormatch)
end


text:
require("../src/CXX.jl")
RequireCompleteType(d::pcpp"clang::Type") = ccall(:RequireCompleteType,Cint,(Ptr{Void},),d.ptr) > 0
function cxxsizeof(d::pcpp"clang::CXXRecordDecl")
    executionEngine = pcpp"llvm::ExecutionEngine"(ccall(:jl_get_llvm_ee,Ptr{Void},()))
    cgt = pcpp"clang::CodeGen::CodeGenTypes"(ccall(:clang_get_cgt,Ptr{Void},()))
    dl = @cxx executionEngine->getDataLayout()
    RequireCompleteType(@cxx d->getTypeForDecl())
    t = @cxx cgt->ConvertRecordDeclType(d)
    @assert @cxx t->isSized()
    div((@cxx dl->getTypeSizeInBits(t)),8)
end
@assert cxxsizeof(pcpp"clang::CXXRecordDecl"(lookup_name(["llvm","ExecutionEngine"]).ptr)) == 152

code_llvmf(f,t) = pcpp"llvm::Function"(ccall(:jl_get_llvmf, Ptr{Void}, (Any,Any,Bool), f, t, false))
function code_graph(f,args)
    v = @cxx std::string()
    os = @cxx llvm::raw_string_ostream(v)
    graphf = code_llvmf(f,args)
    @cxx llvm::WriteGraph(os,graphf)
    @cxx os->flush()
    bytestring((@cxx v->data()), (@cxx v->length()))
end

gt = code_graph(factorize,(typeof(rand(4,4)),))

@assert sizeof(gt) > 0

# LLDB test

initd() = @cxx lldb_private::Debugger::Initialize(cast(C_NULL,vcpp"lldb_private::Debugger::LoadPluginCallbackType"))
initd()

debugger() = @cxx lldb_private::Debugger::CreateInstance()
const dbg = @cxx (debugger())->get()
ci(dbg::pcpp"lldb_private::Debugger") = @cxx dbg->GetCommandInterpreter()
function hc(ci::rcpp"lldb_private::CommandInterpreter",cmd)
    cro = @cxx lldb_private::CommandReturnObject()
    @cxx ci->HandleCommand(pointer(cmd),0,cro)
    @cxx cro->GetOutputData()
end

stdout = pcpp"FILE"(ccall(:fdopen,Ptr{Void},(Int32,Ptr{Uint8}),1,"a"))

function setSTDOUT(dbg,stdout)
    @cxx dbg->SetOutputFileHandle(stdout,false)
end

function setSTDERR(dbg,stdout)
    @cxx dbg->SetErrorFileHandle(stdout,false)
end

setSTDOUT(dbg,stdout)
setSTDERR(dbg,stdout)


text:
using VDPTag
using Plots
using POMDPToolbox
using Reel
using ProgressMeter
using PmapProgressMeter
using ParticleFilters
using ContinuousPOMDPTreeSearchExperiments
using Plots

@show N = 2

@everywhere begin
    using POMDPs
    using POMDPToolbox
    using ContinuousPOMDPTreeSearchExperiments
    using ParticleFilters
    using POMCPOW
    using BasicPOMCP
    using VDPTag

    pomdp = VDPTagPOMDP()

    n = 10000
    seed = 42

    planners = Dict{String, Policy}(

        "pomcpow" => begin
            ro = ToNextML(mdp(pomdp))
            solver = POMCPOWSolver(tree_queries=n,
                                   criterion=MaxUCB(40.0),
                                   final_criterion=MaxTries(),
                                   max_depth=20,
                                   k_action=8.0,
                                   alpha_action=1/20,
                                   k_observation=4.0,
                                   alpha_observation=1/20,
                                   estimate_value=FORollout(ro),
                                   check_repeat_act=false,
                                   check_repeat_obs=false,
                                   rng=MersenneTwister(13)
                                  )
            solve(solver, pomdp)
        end,

        # "discrete_pomcpow" => begin
        #     dpomdp = AODiscreteVDPTagPOMDP(n_angles=1000, n_obs_angles=1000)
        #     ro = translate_policy(ToNextML(mdp(pomdp)), mdp(pomdp), dpomdp, dpomdp)
        #     solver = POMCPOWSolver(tree_queries=10_000,
        #                            criterion=MaxUCB(40.0),
        #                            final_criterion=MaxTries(),
        #                            max_depth=20,
        #                            k_action=8.0,
        #                            alpha_action=1/20,
        #                            k_observation=4.0,
        #                            alpha_observation=1/20,
        #                            estimate_value=FORollout(ro),
        #                            check_repeat_act=false,
        #                            check_repeat_obs=false,
        #                            rng=MersenneTwister(13)
        #                           )
        #     translate_policy(solve(solver, dpomdp), dpomdp, pomdp, dpomdp)
        # end,

        "discrete_pomcp" => begin
            dpomdp = AODiscreteVDPTagPOMDP()
            ro = translate_policy(ToNextML(mdp(pomdp)), mdp(pomdp), dpomdp, dpomdp)
            solver = POMCPSolver(tree_queries=n,
                                   c=40.0,
                                   max_depth=20,
                                   estimate_value=FORollout(ro),
                                   rng=MersenneTwister(13)
                                  )
            dpolicy = solve(solver, dpomdp)
            translate_policy(dpolicy, dpomdp, pomdp, dpomdp)
        end,
    )
end

for (k,p) in planners
    prog = Progress(N, desc="Simulating...")
    # rewards = pmap(prog, 1:N) do i
    rewards = map(1:N) do i
        hr = HistoryRecorder(max_steps=100, rng=MersenneTwister(i))
        up_rng = MersenneTwister(i+100_000)
        up = ObsAdaptiveParticleFilter(deepcopy(pomdp), LowVarianceResampler(10_000), 0.05, up_rng)
        hist = simulate(hr, pomdp, p, up)
        discounted_reward(hist)
    end
    @show k
    @show mean(rewards)
end


text:
#Wrappers for functions directly in the rospy namespace
export init_node, is_shutdown, spin,
       get_param, has_param, set_param, delete_param,
       logdebug, loginfo, logwarn, logerr, logfatal

#General rospy functions
init_node(name::AbstractString; args...) =
    __rospy__[:init_node](ascii(name); args...)
spin()                 = __rospy__[:spin]()
is_shutdown()          = __rospy__[:is_shutdown]()
get_published_topics() = __rospy__[:get_published_topics]()
get_ros_root()         = __rospy__[:get_ros_root]()

#Parameter server API
function get_param(param_name::AbstractString, def=nothing)
    try
        if def == nothing
            __rospy__[:get_param](ascii(param_name))
        else
            __rospy__[:get_param](ascii(param_name), def)
        end
    catch ex
        throw(KeyError(pycall(pybuiltin("str"), PyAny, ex.val)[2:end-1]))
    end
end
set_param(param_name::AbstractString, val) =
    __rospy__[:set_param](ascii(param_name), val)
has_param(param_name::AbstractString) =
    __rospy__[:has_param](ascii(param_name))
function delete_param(param_name::AbstractString)
    try
        __rospy__[:delete_param](ascii(param_name))
    catch ex
        throw(KeyError(pycall(pybuiltin("str"), PyAny, ex.val)[2:end-1]))
    end
end
#Doesn't work for some reason
#rospy_search_param(param_name::AbstractString) =
#    __rospy__[:rospy_search_param](ascii(param_name))
get_param_names() = __rospy__[:get_param_names]()

#Logging API
logdebug(msg, args...) = __rospy__[:logdebug](msg, args...)
loginfo(msg, args...)  = __rospy__[:loginfo](msg, args...)
logwarn(msg, args...)  = __rospy__[:logwarn](msg, args...)
logerr(msg, args...)   = __rospy__[:logerr](msg, args...)
logfatal(msg, args...) = __rospy__[:logfatal](msg, args...)

#Node information
get_name()             = __rospy__[:get_name]()
get_namespace()        = __rospy__[:get_namespace]()
get_node_uri()         = __rospy__[:get_node_uri]()
get_caller_id()        = __rospy__[:get_caller_id]()


text:
function tprun(t::Tuple)
  data = t[1]
  ProjDir = data[:pdir]
  wob = t[2]
  incl = t[3]
  holediameter = t[4]
  heading = t[5]
  dalpha = 0.0

  tptable = Vector{TpRecord}()
  printvec = [1; 100; 1000; 5000:5000:data[:maxiters]]

  local fstream
  cd(joinpath(ProjDir, "tp", "$(string(Int(wob)))", "$(string(Int(incl)))")) do
    fstream = open("tp.log", "w")
  end
  ptup = (printvec, fstream)

  write(fstream, "$(ProjDir)\n")
  write(fstream, "WOB: $(wob), Incl: $(incl)\n")
  write(fstream, "Dalpha: $(data[:welltrajectory][:dalpha])\n")

  data[:welltrajectory][:dalpha] = 0.0
  data[:wob] = wob * 1000.
  data[:welltrajectory][:inclination] = incl
  survey_df = trajectory2df(data)
	params = Parameters(data)
  createmesh!(data, params)
  interpolate!(survey_df, params)

  #println(elements2df(params.elements))

  local tpstring
  count = 1
  while true && count < 10
    dalphainit!(data[:welltrajectory][:dalpha], params)
    xyinit!(data, params)
    createstiffnessmatrices!(params)
    finalinit!(params)
    fem!(params, ptup)
    df = sol2df(params)
    if data[:welltrajectory][:dalpha] == 0.0
      cd(joinpath(ProjDir, "tp0", "$(string(Int(wob)))", "$(string(Int(incl)))")) do
        CSV.write("solution.csv", df)
      end
    end
    ft, f_ft = findfirsttouch(df)
    tprec = TpRecord(data[:welltrajectory][:dalpha], data[:wob],
      data[:welltrajectory][:inclination], df[1, :fx], ft, f_ft, df)
    push!(tptable, tprec)
    if abs(df[1, :fx]) < 1.0
      break
    else
      data[:welltrajectory][:dalpha] = nexttprun(tptable)
      count += 1
    end
    write(fstream, "\nWOB: $(data[:wob]), Incl: $(data[:welltrajectory][:inclination])\n")
    write(fstream, "Force on bit: $(df[1, :fx])\n")
    write(fstream, "New dalpha: $(data[:welltrajectory][:dalpha])\n")
  end
  cd(joinpath(ProjDir, "tp", "$(string(Int(wob)))", "$(string(Int(incl)))")) do
    #tpdf = tp2df(tptable)
    CSV.write("solution.csv", sol2df(params))
    CSV.write("tp.csv", tp2df(tptable))
  end
  close(fstream)
  println("Completed $wob and $incl.")
end

function nexttprun(tptab)
  cnt = length(tptab)
  if tptab[cnt].dalpha == 0.0
    dalpha = tptab[1].fbit > 10.0 ? 0.5 : -0.5
  else
    dfbit = tptab[cnt].fbit - tptab[cnt-1].fbit
    dalpha = (tptab[cnt].fbit*tptab[cnt-1].dalpha-tptab[cnt-1].fbit*tptab[cnt].dalpha)/dfbit
  end
  dalpha
end

function findfirsttouch(df::DataFrame)
  nrows = size(df, 1)
  ft = df[!, :dtb][nrows]
  f_ft = df[!, :fx][nrows]
  for i in 2:nrows
    if abs(df[!, :fx][i]) > 0.0
      ft = df[!, :dtb][i]
      f_ft = df[!, :fx][i]
      break
    end
  end
  (ft, f_ft)
end

function trajectory2df(data::Dict{Symbol, Any})
  df = BHAtp.segments2df(data)
  DataFrame(
    survey_type = [BHAtp.HoleSurvey, BHAtp.BitSurvey],
    depth=[0.0, sum(df[:, :length])],
    inclination = data[:welltrajectory][:inclination]*ones(2),
    heading = data[:welltrajectory][:heading]*ones(2),
    diameter = data[:welltrajectory][:diameter]*ones(2)
  )
end


text:
import PerronFrobenius: AbstractTriangulationInvariantMeasure
import CausalityToolsBase: RectangularBinning, CustomReconstruction
import StateSpaceReconstruction: Simplex, generate_interior_points
import StaticArrays: SVector

abstract type TransferEntropyEstimator end

struct TransferOperatorGrid <: TransferEntropyEstimator end
struct VisitationFrequency <: TransferEntropyEstimator end

export transferentropy, TransferEntropyEstimator, TransferOperatorGrid, VisitationFrequency

"""
    transferentropy(μ::AbstractTriangulationInvariantMeasure, vars::TEVars,
        binning_scheme::RectangularBinning;
        estimator = VisitationFrequency(), n::Int = 10000, b = 2)

#### Transfer entropy using a precomputed invariant measure over a triangulated partition

Estimate transfer entropy from an invariant measure over a triangulation
that has been precomputed either as

1. `μ = invariantmeasure(pts, TriangulationBinning(), ApproximateIntersection())`, or
2. `μ = invariantmeasure(pts, TriangulationBinning(), ExactIntersection())`

where the first method uses approximate simplex intersections (faster) and
the second method uses exact simplex intersections (slow). `μ` contains
all the information needed to compute transfer entropy.

Note: `pts` must be a vector of states, not a vector of
variables/(time series). Wrap your time series in a `Dataset`
first if the latter is the case.

#### Computing transfer entropy (triangulation -> rectangular partition)

Because we need to compute marginals, we need a rectangular grid. To do so,
transfer entropy is computed by sampling the simplices of the
triangulation according to their measure with a total of approximately
`n` points. Introducing multiple points as representatives for the partition
elements does not introduce any bias, because we in computing the
invariant measure, we use no more information than what is encoded in the
dynamics of the original data points. However, from the invariant measure,
we can get a practically infinite amount of points to estimate transfer
entropy from.

Then, transfer entropy is estimated using the visitation
frequency estimator on those points (see docs for `transferentropy_visitfreq`
for more information), on a rectangular grid specified by `binning_scheme`.

#### Common use case

This method is good to use if you want to explore the sensitivity
of transfer entropy to the bin size in the final rectangular grid,
when you have few observations in the time series. The invariant
measure, which encodes the dynamical information, is slow to compute over
the triangulation, but only needs to be computed once.
After that, transfer entropy may be estimated at multiple scales very quickly.

### Example

```julia
# Assume these points are an appropriate delay embedding {(x(t), y(t), y(t+1))} and
# that we're measure transfer entropy from x -> y.
pts = invariantize([rand(3) for i = 1:30])

v = TEVars(Tf = [3], Tpp = [2], Spp = [1])

# Compute invariant measure over a triangulation using approximate
# simplex intersections. This is relatively slow.
μ = invariantmeasure(pts, TriangulationBinning(), ApproximateIntersection())

# Compute transfer entropy from the invariant measure over multiple
# bin sizes. This is fast, because the measure has been precomputed.
tes = map(ϵ -> transferentropy(μ, v, RectangularBinning(ϵ)), 2:50)
```
"""
function transferentropy(μ::AbstractTriangulationInvariantMeasure, vars::TEVars,
        binning_scheme::RectangularBinning;
        estimator = VisitationFrequency(), n::Int = 20000, b = 2)

    dim = length(μ.points[1])
    triang = μ.triangulation.simplexindices
    n_simplices = length(triang)

    simplices = [Simplex(μ.points[triang[i]]) for i = 1:n_simplices]

    # Find a number of points to fill each simplex with so that we
    # obey the measure over the simplices of the triangulation.
    # The total number of points will be roughly `n`, but slightly
    # higher because we need integer numbers of points and use `ceil`
    # for this.
    n_fillpts_persimplex = ceil.(Int, μ.measure.dist .* n)

    # Array to store the points filling the simplices
    fillpts = Vector{SVector{dim, Float64}}()
    sizehint!(fillpts, sum(n_fillpts_persimplex))

    for i = 1:n_simplices
        sᵢ = simplices[i]

        if n_fillpts_persimplex[i] > 0
            pts = generate_interior_points(sᵢ, n_fillpts_persimplex[i])
            append!(fillpts, [SVector{dim, Float64}(pt) for pt in pts])
        end
    end

    sizehint!(fillpts, length(fillpts))

    transferentropy(fillpts, vars, binning_scheme, estimator, b = b)
end

"""
    transferentropy(pts, vars::TEVars, ϵ::RectangularBinning,
        estimator::VisitationFrequency; b = 2)

Compute transfer entropy for a set of ordered points representing
an appropriate embedding of some time series. See documentation for
`TEVars` for info on how to specify the marginals (i.e. which variables
of the embedding are treated as what).

`b` sets the base of the logarithm (e.g `b = 2` gives the transfer
entropy in bits).
"""
function transferentropy(pts, vars::TEVars, ϵ::RectangularBinning,
        estimator::VisitationFrequency; b = 2)

    # Collect variables for the marginals
    C = vars.conditioned_presentpast
    XY = [vars.target_future;      vars.target_presentpast; C]
    YZ = [vars.target_presentpast; vars.source_presentpast; C]
    Y =  [vars.target_presentpast;                          C]

    # Find the bins visited by the joint system (and then get
    # the marginal visits from that, so we don't have to encode
    # bins multiple times).
    joint_bin_visits = joint_visits(pts, ϵ)

    # Compute visitation frequencies for nonempty bi
    p_Y = non0hist(marginal_visits(joint_bin_visits, Y))
    p_XY = non0hist(marginal_visits(joint_bin_visits, XY))
    p_YZ = non0hist(marginal_visits(joint_bin_visits, YZ))
    p_joint = non0hist(joint_bin_visits)

    te = StatsBase.entropy(p_YZ, b) +
            StatsBase.entropy(p_XY, b) -
            StatsBase.entropy(p_Y, b) -
            StatsBase.entropy(p_joint, b)
end

"""
    transferentropy(pts, vars::TEVars, ϵ::RectangularBinning,
        estimator::TransferOperatorGrid; b = 2)

Compute transfer entropy for a set of ordered points representing
an appropriate embedding of some time series. See documentation for
`TEVars` for info on how to specify the marginals (i.e. which variables
of the embedding are treated as what).

`b` sets the base of the logarithm (e.g `b = 2` gives the transfer
entropy in bits).
"""
function transferentropy(pts, vars::TEVars, ϵ::RectangularBinning,
        estimator::TransferOperatorGrid; b = 2)

    # Collect variables for the marginals
    C = vars.conditioned_presentpast
    XY = [vars.target_future;      vars.target_presentpast; C]
    YZ = [vars.target_presentpast; vars.source_presentpast; C]
    Y =  [vars.target_presentpast;                          C]

    # Calculate the invariant distribution over the bins.
    μ = invariantmeasure(pts, ϵ)

    # Find the unique visited bins, then find the subset of those bins
    # with nonzero measure.
    positive_measure_bins = unique(μ.encoded_points)[μ.measure.nonzero_inds]

    p_Y  = marginal(Y, Array(transpose(hcat(positive_measure_bins...,))), μ.measure)
    p_XY = marginal(XY, Array(transpose(hcat(positive_measure_bins...,))), μ.measure)
    p_YZ = marginal(YZ, Array(transpose(hcat(positive_measure_bins...,))), μ.measure)

    te = StatsBase.entropy(p_YZ, b) +
            StatsBase.entropy(p_XY, b) -
            StatsBase.entropy(p_Y, b) -
            StatsBase.entropy(μ.measure.dist[μ.measure.nonzero_inds], b)
end

# Allow CustomReconstructions to be used.
transferentropy(pts::CustomReconstruction, vars::TEVars, ϵ::RectangularBinning,
    estimator::VisitationFrequency; b = 2) =
    transferentropy(pts.reconstructed_pts, vars, ϵ, estimator, b = b)

transferentropy(pts::CustomReconstruction, vars::TEVars, ϵ::RectangularBinning,
    estimator::TransferOperatorGrid; b = 2) =
    transferentropy(pts.reconstructed_pts, vars, ϵ, estimator, b = b)

"""
    transferentropy(pts, vars::TEVars, ϵ::RectangularBinning,
        estimator::TransferEntropyEstimator; b = 2)

#### Transfer entropy using a rectangular partition

Estimate transfer entropy on a rectangular partition over
a set of points `pts`, which represent a generalised embedding of
data series ``x``, ``y`` and (potentially) ``z`` of the form
outlined below.

Note: `pts` must be a vector of states, not a vector of
variables/(time series). Wrap your time series in a `Dataset`
first if the latter is the case.

#### Estimators (and their acronyms)

- Visitation frequency estimator: `VisitationFrequency`; an instance must be provided.
- Transfer operator grid estimator: `TransferOperatorGrid`; an instance must be provided.

#### Relationship between `pts` and `vars`

`pts` should be an embedding of the form
``(y(t + \\eta)^{k}, y(t)^{l}, x(t)^{m}, z(t)^{n}``.
Here, ``y(t + \\eta)^{k})`` indicates that ``k`` future states of `y`
should be included, ``y(t)^{l}`` indicates that ``l`` present/past
states of ``y`` should be included, ``x(t)^{m}`` indicates that ``m``
present/past states of ``x`` should be included, and ``z(t)^{n}`` indicates
that ``n`` present/past states of the variable we're conditioning on should
be included in the embedding vectors. Thus, the total dimension
of the embedding space will be ``k + l + m + n``.

`vars` is a `TEVars` instance contain the instruction on which
variables of the embedding will be treated as part of which marginals
during transfer entropy computation. Check the documentation of
`TEVars` for more details.

### Example

#### 1. Time series

We'll generate two 80-point long realizations, ``x`` and ``y``, of two 1D
logistic maps starting at different initial conditions.

```julia
sys1 = DynamicalSystems.Systems.logistic()
sys2 = DynamicalSystems.Systems.logistic()
x = trajectory(sys1, 80, Ttr = 1000);
y = trajectory(sys1, 80, Ttr = 1000);

# Wrap the time series in a dataset containing the states of the
# composite system.
D = Dataset(x, y)
```

#### 2. Generalised embedding
Say we want to compute transfer entropy from ``x`` to ``y``, and that we
require a 4-dimensional embedding. We do an appropriate delay reconstruction of the data
(``E = \\{S_{pp}, T_{pp}, T_f \\}= \\{x_t, (y_t, y_{t-\\tau}), y_{t+\\eta} \\}``), so that
we're computing the following TE

```math
TE_{x \\to y} =  \\int_E P(x_t, y_{t-\\tau} y_t, y_{t + \\eta}) \\log{\\left( \\dfrac{P(y_{t + \\eta} | (y_t, y_{t - \\tau}, x_t)}{P(y_{t + \\eta} | y_t, y_{t-\\tau})} \\right)}.
```

To create the embedding, we'll use the `customembed` function (check its
documentation for a detailed explanation on how it works).

```julia
# Embed the data, putting time series in the 2nd column (y) of `data` in the
# first three embedding columns, lagging them with lags (η, 0, -τ), and
# putting the 1st column of `data` (x) in the last position of the embedding,
# not lagging it.
τ = optimal_delay(y) # embedding lag
η = 2 # prediction lag

pts = customembed(D, Positions(2, 2, 2, 1), Lags(η, 0, -τ, 0))
```

#### 3. Instructions to the estimator

Now, tell the estimator how to relative the dynamical variables of the
generalised embedding to the marginals in the transfer entropy computation.

```julia
vars = TEVars(Tf = [1], Tpp = [2, 3], Spp = [4])
```

#### 4. Rectangular grid specification

We'll compute transfer entropy using the visitation frequency estimator over
a rectangular partition where the box sizes are determined by
splitting each coordinate axis into ``6`` equally spaced intervals each.

```julia
binning = RectangularBinning(6)
```

#### 5. Compute transfer entropy

```julia
# Over a single rectangular grid
transferentropy(pts, vars, binning, VisitationFrequency()) #, or
transferentropy(pts, vars, binning, TransferOperatorGrid())

# Over multiple cubic grids with differing box sizes
# logarithmically spaced from edge length 0.001 to 0.3
ϵs = 10 .^ range(log(10, 0.001), log10(0.3), length = 15)
map(ϵ -> transferentropy(pts, vars, RectangularBinning(ϵ), VisitationFrequency())) #, or
map(ϵ -> transferentropy(pts, vars, RectangularBinning(ϵ), TransferOperatorGrid()))
```
"""
transferentropy(pts, vars::TEVars, ϵ::RectangularBinning,
    estimator::TransferEntropyEstimator; b = 2)


text:
# This file includes various datasets and dataset generators used in testing.
using Distances
using LinearAlgebra
using SparseArrays
using Random

"""
    rand_dist_matrix(n, [sparse])

Construct a random distance matrix with `n` rows and columns.
"""
function rand_dist_matrix(n)
    A = rand(n, n)
    A .+= A'
    A -= Diagonal(A)
    return A
end
function rand_dist_matrix(n, sparse)
    A = sprand(n, n, sparse / 2)
    A .+= A'
    A -= Diagonal(A)
    dropzeros!(A)
    return A
end

"Distances on an icosahedron graph with edge length 1."
icosahedron = Float64[
    0 1 2 2 1 2 1 1 2 2 1 3
    1 0 3 2 1 1 2 1 2 1 2 2
    2 3 0 1 2 2 1 2 1 2 1 1
    2 2 1 0 3 2 1 1 2 1 2 1
    1 1 2 3 0 1 2 2 1 2 1 2
    2 1 2 2 1 0 3 2 1 1 2 1
    1 2 1 1 2 3 0 1 2 2 1 2
    1 1 2 1 2 2 1 0 3 1 2 2
    2 2 1 2 1 1 2 3 0 2 1 1
    2 1 2 1 2 1 2 1 2 0 3 1
    1 2 1 2 1 2 1 2 1 3 0 2
    3 2 1 1 2 1 2 2 1 1 2 0
]

"Distances on a cycle graph with edge length 1."
cycle = [
    0 1 2 3 4 5 6 7 8 9 8 7 6 5 4 3 2 1
    1 0 1 2 3 4 5 6 7 8 9 8 7 6 5 4 3 2
    2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5 4 3
    3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5 4
    4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5
    5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6
    6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7
    7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8
    8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9
    9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8
    8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7
    7 8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6
    6 7 8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5
    5 6 7 8 9 8 7 6 5 4 3 2 1 0 1 2 3 4
    4 5 6 7 8 9 8 7 6 5 4 3 2 1 0 1 2 3
    3 4 5 6 7 8 9 8 7 6 5 4 3 2 1 0 1 2
    2 3 4 5 6 7 8 9 8 7 6 5 4 3 2 1 0 1
    1 2 3 4 5 6 7 8 9 8 7 6 5 4 3 2 1 0
]

"Projective plane, taken from ripser/examples."
projective_plane = [
    0 1 1 1 1 1 1 1 1 2 2 2 2
    1 0 2 2 2 1 2 1 2 1 2 2 2
    1 2 0 2 2 2 1 2 1 1 2 2 2
    1 2 2 0 2 1 2 2 1 2 2 2 1
    1 2 2 2 0 2 1 1 2 2 2 2 1
    1 1 2 1 2 0 2 2 2 1 1 2 1
    1 2 1 2 1 2 0 2 2 1 1 2 1
    1 1 2 2 1 2 2 0 2 1 2 1 1
    1 2 1 1 2 2 2 2 0 1 2 1 1
    2 1 1 2 2 1 1 1 1 0 1 1 2
    2 2 2 2 2 1 1 2 2 1 0 2 1
    2 2 2 2 2 2 2 1 1 1 2 0 1
    2 2 2 1 1 1 1 1 1 2 1 1 0
]

"""
    rand_n_sphere(n, dim)

Randomly sample `n` points from `dim`-dimensional sphere. Return euclidean distance matrix.
"""
function rand_n_sphere(n, dim)
    points = mapslices(normalize, randn(dim + 1, n); dims=1)
    return pairwise(Euclidean(), points)
end

"""
    torus_dist(pts)

Calculate euclidean distances between points on [-1,1]×[-1,1] as if they were on a flat
torus.
"""
function torus_dist(pts)
    n = size(pts, 2)
    dist = fill(Inf, (n, n))
    for i in 1:n
        for j in (i + 1):n
            for x_offset in -2:2:2, y_offset in -2:2:2
                off = [x_offset, y_offset]
                dst = evaluate(Euclidean(), pts[:, i], pts[:, j] .+ off)
                dist[i, j] = dist[j, i] = min(dist[i, j], dst)
            end
        end
        dist[i, i] = 0
    end
    return dist
end

"""
    torus_points(n, R=2, r=1)

Get `n` points from torus with external radius `R` and internal radius `r` embedded in 3d
euclidean space.
"""
function torus_points(n, R=2, r=1)
    n = floor(Int, sqrt(n))
    return [
        ((R + r * cos(θ)) * cos(φ), (R + r * cos(θ)) * sin(φ), r * sin(θ)) for
        θ in range(0, 2π; length=n + 1)[1:(end - 1)] for
        φ in range(0, 2π; length=n + 1)[1:(end - 1)]
    ]
end

"""
    torus(n)

Construct a torus distance matrix with `n` points. The points are equidistant. `n` must be a
square number.
"""
function torus(n)
    sqn = √n
    @assert sqn == floor(Int, sqn)
    n = floor(Int, sqn)
    r = range(-1, 1; length=n + 1)[1:(end - 1)]
    pts = fill(0.0, (2, n * n))
    i = 1
    for x in r, y in r
        pts[1, i] = x
        pts[2, i] = y
        i += 1
    end
    return torus_dist(pts)
end

"""
    rand_torus(n)

Construct a random torus distance matrix with `n` points.
"""
function rand_torus(n)
    pts = rand(2, n) .* 2 .- 1
    return torus_dist(pts)
end

"""
    disconnected_tori(n, m)

Construct sparse distance matrix of `m` random toruses with `n` points each.
"""
function disconnected_tori(n, m)
    dist = zeros(n * m, n * m)
    for i in 1:n:(n * m)
        dist[i:(i + n - 1), i:(i + n - 1)] .= rand_torus(n)
    end
    perm = shuffle(1:(n * m))
    return sparse(dist[perm, perm])
end


text:
using DocThemePC2
using Test

@testset "install" begin
    DocThemePC2.install(@__DIR__)
    files = ["src/assets/themes/documenter-light.css", "src/assets/themes/documenter-dark.css"]
    relfiles = [joinpath(@__DIR__, f) for f in files]
    for f in relfiles
        @test isfile(f)
    end
    rm.(relfiles)
    rm("src"; recursive = true, force = true)
end


text:
using BiMaps
using Base.Test

let
    bimap = BiMap{Int,ASCIIString}()
    @test isempty(bimap)
    @test length(bimap) == 0

    bimap.left[1] = "one"
    @test bimap.left[1] == "one"
    @test bimap.right["one"] == 1
    @test !isempty(bimap)
    @test length(bimap) == 1

    bimap.right["two"] = 2
    @test bimap.left[2] == "two"
    @test bimap.right["two"] == 2
    @test !isempty(bimap)
    @test length(bimap) == 2

    for (l, r) in bimap
        @test l == 1 ? r == "one" : true
        @test l == 2 ? r == "two" : true
    end

    println(bimap)
end

let
    bimap = BiMap([
        (1, "one"),
        (2, "two"),
    ])

    bimap.left[1] = "1st"
    @test bimap.left[1] == "1st"
    @test bimap.right["1st"] == 1
    @test !haskey(bimap.right, "one")

    bimap.left[10] = "1st"
    @test bimap.left[10] == "1st"
    @test bimap.right["1st"] == 10
    @test !haskey(bimap.left, 1)

    bimap.right["2nd"] = 2
    @test bimap.left[2] == "2nd"
    @test bimap.right["2nd"] == 2
    @test !haskey(bimap.right, "two")

    bimap.right["2nd"] = 20
    @test bimap.left[20] == "2nd"
    @test bimap.right["2nd"] == 20
    @test !haskey(bimap.left, 2)
end


text:
import Pkg

function my_run(cmd)
    @info("Attempting to run command: $(cmd)")
    run(cmd)
    @info("Command ran successfully")
end

function apply_template(original::String,
                        templates::Vector{Pair{String,String}})::String
    content = original
    for pair in templates
        content = replace(content, pair)
    end
    return content
end

function generate_dockerfiles(; directory::String,
                                image_name_prefix::String,
                                image_owner::String,
                                image_list::Vector{String})::Nothing
    original_directory::String = pwd()
    templates = ["{{image_owner}}" => "$(image_owner)",
                 "{{image_name_prefix}}" => "$(image_name_prefix)"]
    for image_name in image_list
        image_directory = joinpath(directory, "images", image_name)
        builddir = joinpath(image_directory, "builddir")
        if isfile(joinpath(builddir, "before_build.jl"))
            cd(builddir)
            include(joinpath(builddir, "before_build.jl"))
            cd(original_directory)
        end
        build_in_filename = joinpath(builddir, "Dockerfile.template")
        build_out_filename = joinpath(builddir, "Dockerfile")
        build_in = read(build_in_filename, String)
        build_out = apply_template(build_in, templates)
        rm(build_out_filename; force = true, recursive = true)
        mkpath(dirname(build_out_filename))
        open(build_out_filename, "w") do f
            write(f, build_out)
        end
        @info("Wrote file $(build_out_filename)")
        testdir = joinpath(image_directory, "testdir")
        test_in_filename = joinpath(testdir, "Dockerfile.template")
        test_out_filename = joinpath(testdir, "Dockerfile")
        test_in = read(test_in_filename, String)
        test_out = apply_template(test_in, templates)
        rm(test_out_filename; force = true, recursive = true)
        mkpath(dirname(test_out_filename))
        open(test_out_filename, "w") do f
            write(f, test_out)
        end
        @info("Wrote file $(test_out_filename)")
    end
    cd(original_directory)
    return nothing
end

function build(; image_name::String,
                 image_name_prefix::String,
                 image_owner::String,
                 image_directory::String,
                 docker_command::String)::Nothing
    original_directory::String = pwd()
    cd(image_directory)
    cd("builddir")
    my_run(`$(docker_command) build -t $(image_owner)/$(image_name_prefix)$(image_name) .`)
    cd(image_directory)
    cd("testdir")
    my_run(`$(docker_command) build -t $(image_owner)/$(image_name_prefix)test-$(image_name) .`)
    cd(original_directory)
    return nothing
end

function test(; image_name::String,
                image_name_prefix::String,
                image_owner::String,
                image_directory::String,
                docker_command::String)::Nothing
    original_directory::String = pwd()
    cd(image_directory)
    cd("testdir")
    test_args::Vector{String} = Vector{String}(undef, 0)
    if ispath("get-test-args.sh")
        for x in split(strip(read(`./get-test-args.sh`, String)))
            if length(strip(x)) > 0
                push!(test_args, strip(x))
            end
        end
    end
    runtests = join(vcat(["/bin/runtests.sh"], test_args), " ")
    my_run(`$(docker_command) run --user predictmdtestuser --network none -it $(image_owner)/$(image_name_prefix)test-$(image_name) /bin/bash -c "$(runtests)"`)
    cd(original_directory)
    return nothing
end

function login(; image_name::String,
                image_name_prefix::String,
                image_owner::String,
                image_directory::String,
                docker_command::String)::Nothing
    original_directory::String = pwd()
    my_run(`$(docker_command) login`)
    cd(original_directory)
    return nothing
end

function push(; image_name::String,
                image_name_prefix::String,
                image_owner::String,
                image_directory::String,
                docker_command::String)::Nothing
    original_directory::String = pwd()
    cd(image_directory)
    cd("builddir")
    my_run(`$(docker_command) push $(image_owner)/$(image_name_prefix)$(image_name)`)
    cd(image_directory)
    cd("testdir")
    my_run(`$(docker_command) push $(image_owner)/$(image_name_prefix)test-$(image_name)`)
    cd(original_directory)
    return nothing
end

function main()::Nothing
    DIR::String = @__DIR__
    DOCKER_COMMAND::String = "docker"
    CONFIG_FILE_NAME::String = "config.toml"
    CONFIG_FILE_PATH::String = joinpath(splitpath(DIR)...,CONFIG_FILE_NAME)
    original_directory::String = pwd()
    if length(ARGS) < 2
        error(string("Syntax: julia make.jl COMMAND IMAGE [IMAGE_NAME_PREFIX]\n",
                     "Examples: \n",
                     "julia make.jl build julia-1\n",
                     "julia make.jl test julia-1\n",
                     "julia make.jl push julia-1\n",
                     "julia make.jl build julia-1 staging-\n",
                     "julia make.jl test julia-1 staging-\n",
                     "julia make.jl push julia-1 staging-\n"))
    end
    command::Symbol = Symbol(strip(ARGS[1]))
    image_name::String = strip(ARGS[2])
    image_name_prefix::String = strip(length(ARGS) == 2 ? "" : ARGS[3])
    config = try
        Pkg.TOML.parsefile(CONFIG_FILE_PATH)
    catch
        Dict()
    end
    _image_owner::String = strip(get(config, "image_owner", ""))
    image_owner::String = strip(length(_image_owner) == 0 ? "dilumaluthge" : _image_owner)
    image_list::Vector{String} = strip.(get(config, "image_list", String[]))
    if !(image_name in image_list)
        error(string("image $(image_name) is not in list of images"))
    end

    generate_dockerfiles(; directory = DIR,
                           image_name_prefix = image_name_prefix,
                           image_owner = image_owner,
                           image_list = [image_name])
    image_directory::String = joinpath(DIR, "images", image_name)
    if command == :build
        build(; image_name = image_name,
                image_name_prefix = image_name_prefix,
                image_owner = image_owner,
                image_directory = image_directory,
                docker_command = DOCKER_COMMAND)
    elseif command == :test
        test(; image_name = image_name,
                image_name_prefix = image_name_prefix,
                image_owner = image_owner,
                image_directory = image_directory,
                docker_command = DOCKER_COMMAND)
    elseif command == :login
        login(; image_name = image_name,
                image_name_prefix = image_name_prefix,
                image_owner = image_owner,
                image_directory = image_directory,
                docker_command = DOCKER_COMMAND)
    elseif command == :push
        push(; image_name = image_name,
                image_name_prefix = image_name_prefix,
                image_owner = image_owner,
                image_directory = image_directory,
                docker_command = DOCKER_COMMAND)
    else
        error(string("command $(string(command)) is not one of build, test, login, push"))
    end
    cd(original_directory)
    return nothing
end

main()


text:
# Company Overview (https://www.alphavantage.co/documentation/#company-overview)
# ex: https://www.alphavantage.co/query?function=OVERVIEW&symbol=IBM&apikey=demo

function company_overview(symbol::String; outputsize::String="compact", datatype::String="json")
    @argcheck in(datatype, ["json", "csv"])
    uri = _form_uri_head("OVERVIEW") * "&symbol=$symbol" * _form_uri_tail(outputsize, datatype)
    data = _get_request(uri)
    return _parse_response(data, datatype)
end
export company_overview

# Income Statement (https://www.alphavantage.co/documentation/#income-statement)
# ex: https://www.alphavantage.co/query?function=INCOME_STATEMENT&symbol=IBM&apikey=demo

function income_statement(symbol::String; outputsize::String="compact", datatype::String="json")
    @argcheck in(datatype, ["json", "csv"])
    uri = _form_uri_head("INCOME_STATEMENT") * "&symbol=$symbol" * _form_uri_tail(outputsize, datatype)
    data = _get_request(uri)
    return _parse_response(data, datatype)
end
export income_statement

# Balance Sheet (https://www.alphavantage.co/documentation/#balance-sheet)
# ex: https://www.alphavantage.co/query?function=BALANCE_SHEET&symbol=IBM&apikey=demo

function balance_sheet(symbol::String; outputsize::String="compact", datatype::String="json")
    @argcheck in(datatype, ["json", "csv"])
    uri = _form_uri_head("BALANCE_SHEET") * "&symbol=$symbol" * _form_uri_tail(outputsize, datatype)
    data = _get_request(uri)
    return _parse_response(data, datatype)
end
export balance_sheet

# Cash Flow (https://www.alphavantage.co/documentation/#cash-flow)
# ex: https://www.alphavantage.co/query?function=CASH_FLOW&symbol=IBM&apikey=demo

function cash_flow(symbol::String; outputsize::String="compact", datatype::String="json")
    @argcheck in(datatype, ["json", "csv"])
    uri = _form_uri_head("CASH_FLOW") * "&symbol=$symbol" * _form_uri_tail(outputsize, datatype)
    data = _get_request(uri)
    return _parse_response(data, datatype)
end
export cash_flow

# Earnings (https://www.alphavantage.co/documentation/#earnings)
# ex: https://www.alphavantage.co/query?function=EARNINGS&symbol=IBM&apikey=demo

function earnings(symbol::String; outputsize::String="compact", datatype::String="json")
    @argcheck in(datatype, ["json", "csv"])
    uri = _form_uri_head("EARNINGS") * "&symbol=$symbol" * _form_uri_tail(outputsize, datatype)
    data = _get_request(uri)
    return _parse_response(data, datatype)
end
export earnings

function listingstatus()
    uri = _form_uri_head("LISTING_STATUS") * "&apikey=" * ENV["ALPHA_VANTAGE_API_KEY"]
    data = _get_request(uri)
    return _parse_response(data, "csv")
end

function listingstatus(state::String, date::String)
    @argcheck in(state, ["active", "delisted"])
    uri = _form_uri_head("LISTING_STATUS") * "&state=$(state)&date=$(date)&apikey=" * ENV["ALPHA_VANTAGE_API_KEY"]
    data = _get_request(uri)
    return _parse_response(data, "csv")
end

export listingstatus


text:
export RStepAdapter, SelectiveRStepAdapter

"""
$(TYPEDEF)

Simple adapter based on reachability. Whenever the current FSP solution error is found to exceed the acceptable tolerance, the adapter will expand the state space by exploring all states that could be reached from the current state space within a set number of reaction events.

This adapter only works with `StateSpaceSparse`.
"""
struct RStepAdapter <: AbstractSpaceAdapterSparse
    initial_step_count::Int
    max_step_count::Int
    dropstates::Bool
end

"""
$(TYPEDSIGNATURES)

Make the state space and solution vector ready for FSP integration.
"""
function init!(statespace::StateSpaceSparse, adapter::RStepAdapter, p::Vector{<:AbstractFloat}, t::AbstractFloat, fsptol::AbstractFloat)
    nold = get_state_count(statespace)
    expand!(statespace, adapter.initial_step_count)
    append!(p, zeros(typeof(p[1]), get_state_count(statespace) - nold))
    nothing
end

"""
$(TYPEDSIGNATURES)

Adapt the state space and probability vector based on current error recorded in the `sinks` vector.
"""
function adapt!(statespace::StateSpaceSparse, adapter::RStepAdapter, p::Vector{RealT},
    sinks::Vector{RealT},
    t::RealT, tend::RealT,
    fsptol::RealT; integrator::Union{DEIntegrator,Nothing} = nothing) where {RealT<:AbstractFloat}

    if adapter.dropstates
        pids = sortperm(p)
        dropcount = sum((sum(p) .- cumsum(p[pids])) .>= (1.0 - t * fsptol / tend))
        dropids = sort(pids[1:dropcount])
        deleteat!(statespace, dropids)
        deleteat!(p, dropids)
    end

    nold = get_state_count(statespace)
    expand!(statespace, adapter.max_step_count)
    append!(p, zeros(RealT, get_state_count(statespace) - nold))
    nothing
end

"""
$(TYPEDEF)

Adapter based on reachability. This adapter only explore new states through reaction channels with positive derivatives.

This adapter only works with `StateSpaceSparse`.
"""
struct SelectiveRStepAdapter <: AbstractSpaceAdapterSparse
    initial_step_count::Int
    max_step_count::Int
    dropstates::Bool
end

"""
$(TYPEDSIGNATURES)

`init!(statespace::StateSpaceSparse, adapter::SelectiveRStepAdapter, p::Vector{RealT}, t::RealT, fsptol::RealT) where {RealT <: AbstractFloat}`

Make the state space and solution vector ready for FSP integration.
"""
function init!(statespace::StateSpaceSparse, adapter::SelectiveRStepAdapter, p::Vector{RealT}, t::RealT, fsptol::RealT) where {RealT<:AbstractFloat}
    nold = get_state_count(statespace)
    expand!(statespace, adapter.initial_step_count)
    append!(p, zeros(RealT, get_state_count(statespace) - nold))
    nothing
end

"""
$(TYPEDSIGNATURES)

Adapt the state space and probability vector based on current error recorded in the `sinks` vector.
"""
function adapt!(statespace::StateSpaceSparse, adapter::SelectiveRStepAdapter,
    p::Vector{<:AbstractFloat}, sinks::Vector{<:AbstractFloat},
    t::AbstractFloat, tend::AbstractFloat, fsptol::AbstractFloat;
    integrator::Union{DEIntegrator,Nothing} = nothing)
    if isempty(p)
        throw(ArgumentError("Empty `p` input in `adapt!`."))
    end
    if adapter.dropstates
        pids = sortperm(p)
        dropcount = sum((sum(p) .- cumsum(p[pids])) .> (1.0 - t * fsptol / tend))
        dropids = sort(pids[1:dropcount])
        deleteat!(statespace, dropids)
        deleteat!(p, dropids)
    end
    sink_count = get_sink_count(statespace)
    du = similar(integrator.u)
    integrator.f(du, integrator.u, [], t)
    dsinks = du[end-sink_count+1:end]
    expandreactions = findall(dsinks .> 0)
    nold = get_state_count(statespace)
    expand!(statespace, adapter.max_step_count; onlyreactions = expandreactions)
    append!(p, zeros(typeof(p[1]), get_state_count(statespace) - nold))
    nothing
end


text:
versioninfo()
cd(expanduser("~/Code/GitHub/adventofcode2015"))
input = open(readall, "1input.txt")

pos = count(e -> isequal('(', e), input)
neg = count(e -> isequal(')', e), input)

res = pos-neg