text:
using Test
using StaticArrays
using Printf
using ClimateMachine.VariableTemplates
using ClimateMachine.VariableTemplates: wrap_val
import ClimateMachine.VariableTemplates
VT = VariableTemplates

@testset "Test complex models" begin
    include("complex_models.jl")

    FT = Float32

    # test getproperty
    m = ScalarModel()
    st = state(m, FT)
    vs = varsize(st)
    a_global = collect(1:vs)
    v = Vars{st}(a_global)
    @test v.x == FT(1)

    Nv = 4
    m = VectorModel{Nv}()
    st = state(m, FT)
    vs = varsize(st)
    a_global = collect(1:vs)
    v = Vars{st}(a_global)
    @test v.x == SVector{Nv, FT}(1:Nv)

    N = 3
    M = 6
    m = MatrixModel{N, M}()
    st = state(m, FT)
    vs = varsize(st)
    a_global = collect(1:vs)
    v = Vars{st}(a_global)
    @test v.x == SHermitianCompact{N, FT, M}(collect(1:(1 + M - 1)))

    Nv = 3
    N = 3
    M = 6
    m = CompositModel(Nv, N, M)
    st = state(m, FT)
    vs = varsize(st)
    a_global = collect(1:vs)
    v = Vars{st}(a_global)

    scalar_model = v.scalar_model
    @test v.scalar_model.x == FT(1)

    vector_model = v.vector_model
    @test v.vector_model.x == SVector{Nv, FT}([2, 3, 4])

    matrix_model = v.matrix_model
    @test v.matrix_model.x ==
          SHermitianCompact{N, FT, M}(collect(5:(5 + M - 1)))

    Nv = 3
    N = 5
    m = NTupleContainingModel(N, Nv)
    st = state(m, FT)
    vs = varsize(st)
    a_global = collect(1:vs)
    v = Vars{st}(a_global)

    offset = (Nv + 1) * N
    @test v.vector_model.x == SVector{Nv, FT}(collect(1:Nv) .+ offset)
    @test v.scalar_model.x == FT(1 + Nv) + offset

    # Make sure we bounds error for NTupleModel's:
    @test_throws BoundsError m.ntuple_model[0]
    @test_throws BoundsError m.ntuple_model[N + 1]

    unval(::Val{i}) where {i} = i
    @unroll_map(N) do i
        @test m.ntuple_model[i] isa NTupleModel
        @test m.ntuple_model[i].scalar_model isa ScalarModel
        @test v.ntuple_model[i].scalar_model.x ==
              FT(unval(i)) + (Nv) * (unval(i) - 1)
        @test v.vector_model.x == SVector{Nv, FT}(1:Nv) .+ offset
        @test v.scalar_model.x == FT(Nv + 1) + offset
    end

    @test vuntuple(x -> x, 5) == ntuple(i -> Val(i), Val(5))

    # test flattenednames
    fn = flattenednames(st)
    j = 1
    for i in 1:N
        @test fn[j] === "ntuple_model[$i].scalar_model.x"
        j += 1
        for k in 1:Nv
            @test fn[j] === "ntuple_model[$i].vector_model.x[$k]"
            j += 1
        end
    end
    for k in 1:Nv
        @test fn[j] === "vector_model.x[$k]"
        j += 1
    end
    @test fn[j] === "scalar_model.x"

    # flattened_tup_chain - empty/generic cases
    struct Foo end
    @test flattened_tup_chain(NamedTuple{(), Tuple{}}) == ()
    @test flattened_tup_chain(Foo, RetainArr()) == ((Symbol(),),)
    @test flattened_tup_chain(Foo, FlattenArr()) == ((Symbol(),),)

    # flattened_tup_chain - SHermitianCompact
    Nv, M = 3, 6
    A = SHermitianCompact{Nv, FT, M}(collect(1:(1 + M - 1)))

    ftc = flattened_tup_chain(typeof(A), FlattenArr())
    @test ftc == ntuple(i -> (Symbol(), i), M)

    ftc = flattened_tup_chain(typeof(A), RetainArr())
    @test ftc == ((Symbol(),),)

    # flattened_tup_chain - Retain arrays

    ftc = flattened_tup_chain(st, RetainArr())
    j = 1
    for i in 1:N
        @test ftc[j] === (:ntuple_model, i, :scalar_model, :x)
        j += 1
        @test ftc[j] === (:ntuple_model, i, :vector_model, :x)
        j += 1
    end
    @test ftc[j] === (:vector_model, :x)
    j += 1
    @test ftc[j] === (:scalar_model, :x)

    # test varsindex
    ntuple(N) do i
        i_val = Val(i)
        i_sm = varsindex(st, :ntuple_model, i_val, :scalar_model, :x)
        i_vm = varsindex(st, :ntuple_model, i_val, :vector_model, :x)
        nt_offset = (Nv + 1) - 1

        i_sm_correct = (i + nt_offset * (i - 1)):(i + nt_offset * (i - 1))
        @test i_sm == i_sm_correct

        offset = 1
        i_start = i + nt_offset * (i - 1) + offset
        i_vm_correct = (i_start):(i_start + Nv - 1)
        @test i_vm == i_vm_correct
    end

    # test that getproperty matches varsindex
    ntuple(N) do i
        i_ϕ = varsindex(st, wrap_val.(ftc[i])...)
        ϕ = getproperty(v, wrap_val.(ftc[i]))
        @test all(parent(v)[i_ϕ] .≈ ϕ)
    end

    # test getproperty with tup-chain
    @unroll_map(N) do i
        @test v.scalar_model.x == getproperty(v, (:scalar_model, :x))
        @test v.vector_model.x == getproperty(v, (:vector_model, :x))
        @test v.ntuple_model[i] == getproperty(v, (:ntuple_model, i))
        @test v.ntuple_model[i].scalar_model ==
              getproperty(v, (:ntuple_model, i, :scalar_model))
        @test v.ntuple_model[i].scalar_model.x ==
              getproperty(v, (:ntuple_model, i, :scalar_model, :x))
    end

    # Test converting to flattened NamedTuple
    fnt = flattened_named_tuple(v, RetainArr())
    @test fnt.ntuple_model_1_scalar_model_x == 1.0f0
    @test fnt.ntuple_model_1_vector_model_x == Float32[2.0, 3.0, 4.0]
    @test fnt.ntuple_model_2_scalar_model_x == 5.0f0
    @test fnt.ntuple_model_2_vector_model_x == Float32[6.0, 7.0, 8.0]
    @test fnt.ntuple_model_3_scalar_model_x == 9.0f0
    @test fnt.ntuple_model_3_vector_model_x == Float32[10.0, 11.0, 12.0]
    @test fnt.ntuple_model_4_scalar_model_x == 13.0f0
    @test fnt.ntuple_model_4_vector_model_x == Float32[14.0, 15.0, 16.0]
    @test fnt.ntuple_model_5_scalar_model_x == 17.0f0
    @test fnt.ntuple_model_5_vector_model_x == Float32[18.0, 19.0, 20.0]
    @test fnt.vector_model_x == Float32[21.0, 22.0, 23.0]
    @test fnt.scalar_model_x == 24.0f0

    # flattened_tup_chain - Flatten arrays

    ftc = flattened_tup_chain(st, FlattenArr())
    j = 1
    for i in 1:N
        @test ftc[j] === (:ntuple_model, i, :scalar_model, :x)
        j += 1
        for k in 1:Nv
            @test ftc[j] === (:ntuple_model, i, :vector_model, :x, k)
            j += 1
        end
    end
    for i in 1:Nv
        @test ftc[j] === (:vector_model, :x, i)
        j += 1
    end
    @test ftc[j] === (:scalar_model, :x)

    # test varsindex (flatten arrays)
    ntuple(N) do i
        i_val = Val(i)
        i_sm = varsindex(st, :ntuple_model, i_val, :scalar_model, :x)
        nt_offset = (Nv + 1) - 1

        i_sm_correct = (i + nt_offset * (i - 1)):(i + nt_offset * (i - 1))
        @test i_sm == i_sm_correct

        for j in 1:Nv
            i_vm =
                varsindex(st, :ntuple_model, i_val, :vector_model, :x, Val(j))
            offset = 1
            i_start = i + nt_offset * (i - 1) + offset
            i_vm_correct = i_start + j - 1
            @test i_vm == i_vm_correct:i_vm_correct
        end
    end

    # test that getproperty matches varsindex
    ntuple(N) do i
        i_ϕ = varsindex(st, wrap_val.(ftc[i])...)
        ϕ = getproperty(v, wrap_val.(ftc[i]))
        @test all(parent(v)[i_ϕ] .≈ ϕ)
    end

    # test getproperty with tup-chain
    for k in 1:Nv
        @test v.vector_model.x[k] == getproperty(v, (:vector_model, :x, Val(k)))
    end

    # test getindex with Val
    @test getindex((1, 2), Val(1)) == 1
    @test getindex((1, 2), Val(2)) == 2
    @test getindex(SVector(1, 2), Val(1)) == 1
    @test getindex(SVector(1, 2), Val(2)) == 2

    nt = (; a = ((; x = 1), (; x = 2)))
    fnt = VT.flattened_tuple(FlattenArr(), nt)
    vg = Grad{typeof(nt)}(zeros(MMatrix{3, length(fnt), FT}))
    parent(vg)[1, :] .= fnt
    parent(vg)[2, :] .= fnt
    parent(vg)[3, :] .= fnt
    for i in 1:2
        @test getindex(vg.a, Val(i)).x[1] == i
        @test getindex(vg.a, Val(i)).x[2] == i
        @test getindex(vg.a, Val(i)).x[3] == i
    end

    # getpropertyorindex
    @test VT.getpropertyorindex((1, 2), Val(1)) == 1
    @test VT.getpropertyorindex((1, 2), Val(2)) == 2
    @test VT.getpropertyorindex([1, 2], Val(1)) == 1
    @test VT.getpropertyorindex([1, 2], Val(2)) == 2
    @test VT.getpropertyorindex(v, :scalar_model) == v.scalar_model
    for i in 1:N
        @test VT.getpropertyorindex(v.ntuple_model, Val(i)) ==
              v.ntuple_model[Val(i)]
        @test VT.getpropertyorindex(v.ntuple_model, (Val(i),)) ==
              v.ntuple_model[Val(i)]
        @test getindex(v.ntuple_model, (Val(i),)) ==
              VT.getpropertyorindex(v.ntuple_model, (Val(i),))
    end

    # Test converting to flattened NamedTuple
    fnt = flattened_named_tuple(v, FlattenArr())
    @test fnt.ntuple_model_1_scalar_model_x == 1.0f0
    @test fnt.ntuple_model_1_vector_model_x_1 == 2.0
    @test fnt.ntuple_model_1_vector_model_x_2 == 3.0
    @test fnt.ntuple_model_1_vector_model_x_3 == 4.0
    @test fnt.ntuple_model_2_scalar_model_x == 5.0f0
    @test fnt.ntuple_model_2_vector_model_x_1 == 6.0
    @test fnt.ntuple_model_2_vector_model_x_2 == 7.0
    @test fnt.ntuple_model_2_vector_model_x_3 == 8.0
    @test fnt.ntuple_model_3_scalar_model_x == 9.0f0
    @test fnt.ntuple_model_3_vector_model_x_1 == 10.0
    @test fnt.ntuple_model_3_vector_model_x_2 == 11.0
    @test fnt.ntuple_model_3_vector_model_x_3 == 12.0
    @test fnt.ntuple_model_4_scalar_model_x == 13.0f0
    @test fnt.ntuple_model_4_vector_model_x_1 == 14.0
    @test fnt.ntuple_model_4_vector_model_x_2 == 15.0
    @test fnt.ntuple_model_4_vector_model_x_3 == 16.0
    @test fnt.ntuple_model_5_scalar_model_x == 17.0f0
    @test fnt.ntuple_model_5_vector_model_x_1 == 18.0
    @test fnt.ntuple_model_5_vector_model_x_2 == 19.0
    @test fnt.ntuple_model_5_vector_model_x_3 == 20.0
    @test fnt.vector_model_x_1 == 21.0
    @test fnt.vector_model_x_2 == 22.0
    @test fnt.vector_model_x_3 == 23.0
    @test fnt.scalar_model_x == 24.0f0

    struct Foo end
    nt = (;
        nest = (;
            v = SVector(1, 2, 3),
            nt = (;
                shc = SHermitianCompact{3, FT, 6}(collect(1:6)),
                f = FT(1.0),
            ),
            d = SDiagonal(collect(1:3)...),
            tt = (Foo(), Foo()),
            t = Foo(),
        ),
    )
    # Test flattened_tuple:

    @test VT.flattened_tuple(RetainArr(), NamedTuple()) == ()
    @test VT.flattened_tuple(FlattenArr(), NamedTuple()) == ()
    @test VT.flattened_tuple(RetainArr(), Tuple(NamedTuple())) == ()
    @test VT.flattened_tuple(FlattenArr(), Tuple(NamedTuple())) == ()

    ft = FlattenArr()
    @test VT.flattened_tuple(ft, nt.nest.nt.f) == (1.0f0,)
    @test VT.flattened_tuple(ft, nt.nest.nt) ==
          (1.0f0, 2.0f0, 3.0f0, 4.0f0, 5.0f0, 6.0f0, 1.0f0)
    @test VT.flattened_tuple(ft, nt.nest.d) == (1, 2, 3)
    @test VT.flattened_tuple(ft, nt.nest.t) == (Foo(),)
    @test VT.flattened_tuple(ft, nt.nest.tt) == (Foo(), Foo())

    ft = RetainArr()
    @test VT.flattened_tuple(ft, nt.nest.nt.f) == (1.0f0,)
    @test VT.flattened_tuple(ft, nt.nest.nt)[1] == nt.nest.nt.shc.lowertriangle
    @test VT.flattened_tuple(ft, nt.nest.nt)[2] == 1.0f0
    @test VT.flattened_tuple(ft, nt.nest.d) == (nt.nest.d.diag,)
    @test VT.flattened_tuple(ft, nt.nest.t) == (Foo(),)
    @test VT.flattened_tuple(ft, nt.nest.tt) == (Foo(), Foo())

    # Test flattened_named_tuple for NamedTuples
    fnt = flattened_named_tuple(nt, FlattenArr())
    @test fnt.nest_v_1 == 1
    @test fnt.nest_v_2 == 2
    @test fnt.nest_v_3 == 3
    @test fnt.nest_nt_shc_1 == 1.0
    @test fnt.nest_nt_shc_2 == 2.0
    @test fnt.nest_nt_shc_3 == 3.0
    @test fnt.nest_nt_shc_4 == 4.0
    @test fnt.nest_nt_shc_5 == 5.0
    @test fnt.nest_nt_shc_6 == 6.0
    @test fnt.nest_nt_f == 1.0
    @test fnt.nest_tt_1 == Foo()
    @test fnt.nest_tt_2 == Foo()
    @test fnt.nest_t == Foo()

    fnt = flattened_named_tuple(nt, RetainArr())
    @test fnt.nest_v == SVector(1, 2, 3)
    @test fnt.nest_nt_shc == nt.nest.nt.shc.lowertriangle
    @test fnt.nest_nt_f == 1.0
    @test fnt.nest_tt_1 == Foo()
    @test fnt.nest_tt_2 == Foo()
    @test fnt.nest_t == Foo()

    # Test that show doesn't break
    sprint(show, v)
    sprint(show, vg)

end


text:
bm_nonsymbolic_small = @benchmark sem_fit(model_ml_small)

BenchmarkTools.Trial: 7037 samples with 1 evaluation.
 Range (min … max):  536.094 μs …  13.142 ms  ┊ GC (min … max): 0.00% … 84.11%
 Time  (median):     623.008 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   703.678 μs ± 577.872 μs  ┊ GC (mean ± σ):  4.18% ±  4.91%

     ▃█▇▅▃▁
  ▃▄███████▇▆▆▄▄▄▃▃▃▃▃▃▂▂▂▂▂▂▂▂▂▂▁▂▂▁▁▂▁▁▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▂
  536 μs           Histogram: frequency by time         1.24 ms <

 Memory estimate: 439.34 KiB, allocs estimate: 696.

bm_nonsymbolic_big = @benchmark sem_fit(model_ml_big)

BenchmarkTools.Trial: 10 samples with 1 evaluation.
 Range (min … max):  377.030 ms … 742.732 ms  ┊ GC (min … max): 0.00% … 0.62%
 Time  (median):     516.519 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   538.650 ms ± 122.110 ms  ┊ GC (mean ± σ):  0.36% ± 0.49%

  ▁ ▁           █ ▁            ▁         ▁   ▁  ▁             ▁
  █▁█▁▁▁▁▁▁▁▁▁▁▁█▁█▁▁▁▁▁▁▁▁▁▁▁▁█▁▁▁▁▁▁▁▁▁█▁▁▁█▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
  377 ms           Histogram: frequency by time          743 ms <

 Memory estimate: 56.59 MiB, allocs estimate: 956.

# MKL

using MKL

bm_nonsymbolic_small_mkl = @benchmark sem_fit(model_ml_small)

BenchmarkTools.Trial: 4948 samples with 1 evaluation.
 Range (min … max):  712.489 μs …  12.225 ms  ┊ GC (min … max): 0.00% … 82.20%
 Time  (median):     899.797 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):     1.003 ms ± 549.112 μs  ┊ GC (mean ± σ):  2.60% ±  4.65%

        ▂▁█▄▁
  ▁▂▂▃▃▄██████▆▅▅▅▄▃▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▂
  712 μs           Histogram: frequency by time         1.78 ms <

 Memory estimate: 439.34 KiB, allocs estimate: 696.

bm_nonsymbolic_big_mkl = @benchmark sem_fit(model_ml_big)

BenchmarkTools.Trial: 21 samples with 1 evaluation.
 Range (min … max):  149.162 ms … 401.825 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     243.611 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   249.148 ms ±  53.616 ms  ┊ GC (mean ± σ):  0.54% ± 0.63%

  ▁   ▁      ▁    ▁█▁█▁ ▁ ▁ ▁▁█▁    ▁▁     ▁                  ▁
  █▁▁▁█▁▁▁▁▁▁█▁▁▁▁█████▁█▁█▁████▁▁▁▁██▁▁▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
  149 ms           Histogram: frequency by time          402 ms <

 Memory estimate: 56.59 MiB, allocs estimate: 956.

# BFGS

algo = BFGS(;linesearch = BackTracking(order=3), alphaguess = InitialHagerZhang())

bm_algo_small = @benchmark sem_fit(model_ml_small)

BenchmarkTools.Trial: 3270 samples with 1 evaluation.
 Range (min … max):  937.452 μs … 21.758 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):       1.255 ms              ┊ GC (median):    0.00%
 Time  (mean ± σ):     1.510 ms ±  1.043 ms  ┊ GC (mean ± σ):  2.31% ± 4.83%

   ▆█▄▂
  ▆█████▇▆▆▅▄▄▅▄▄▄▄▄▄▄▄▄▄▄▃▃▃▃▂▂▂▂▂▂▂▂▂▁▂▂▂▂▁▁▂▁▁▁▂▂▁▁▁▁▂▂▁▁▁▂ ▃
  937 μs          Histogram: frequency by time          4.1 ms <

 Memory estimate: 439.34 KiB, allocs estimate: 696.

bm_algo_big = @benchmark sem_fit(model_ml_big)

BenchmarkTools.Trial: 14 samples with 1 evaluation.
 Range (min … max):  274.281 ms … 578.298 ms  ┊ GC (min … max): 0.00% … 0.58%
 Time  (median):     340.173 ms               ┊ GC (median):    0.47%
 Time  (mean ± σ):   358.096 ms ±  82.599 ms  ┊ GC (mean ± σ):  0.54% ± 0.57%

  ▁ ▁█▁    ▁ ▁  ▁  ▁▁  ▁         ▁   ▁                        ▁
  █▁███▁▁▁▁█▁█▁▁█▁▁██▁▁█▁▁▁▁▁▁▁▁▁█▁▁▁█▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█ ▁
  274 ms           Histogram: frequency by time          578 ms <

 Memory estimate: 56.59 MiB, allocs estimate: 956.

# profiling

ProfileView.@profview profile_test(model_ml_big, 10)

- small model: matrix inversion takes more time

big model takes more time
- SemML: 153

ML.jl, 221, uniformscaling bad performance for small model

- factorize upper triangular

ext: jl


text:
try
    pyimport("sys")
catch e
    using PyCall
    @info "Python executable used by PyCall: $(pyimport("sys").executable)";
    # Append the code directory to python path
    # (see https://github.com/JuliaPy/PyCall.jl#troubleshooting)
    file_dir = @__DIR__;
    pushfirst!(PyVector(pyimport("sys")."path"),
               file_dir);
end

py"""
from tlqg import TLQG_PosRangeLocalization2D, TLQG_Manipulate2D
"""
import SACBP: BelMvNormal, MControl2D
using Distributions
using Printf

mutable struct TLQG_Planner_PosRangeLocalization2D
    planner::PyObject
    μ_array::Union{Nothing, Vector{Vector{Float64}}}
    Σ_array::Union{Nothing, Vector{Matrix{Float64}}}
    l_array::Union{Nothing, Vector{Vector{Float64}}}
    L_array::Union{Nothing, Vector{Matrix{Float64}}}
    kl_threshold::Float64 # kl threshold to trigger replanning
    t_last_plan::Union{Nothing, Float64} # last time when replanning occured
end

function TLQG_Planner_PosRangeLocalization2D(Th::Float64, # planning horizon [s]
                                             dto::Float64,
                                             Nt::Int64, # number of targets
                                             Q::Matrix{Float64},
                                             Cu::Matrix{Float64},
                                             u_param_min::Vector{Float64},
                                             u_param_max::Vector{Float64},
                                             kl_threshold::Float64)
    N = Int64(round(Th/dto, digits=5));
    Cs = diagm(vcat([10.0, 10.0], zeros(2*Nt)))
    Q_tlqg = zeros(2*Nt, 2*Nt)
    for ii = 1:Nt
        Q_tlqg[2*(ii-1)+1:2*ii, 2*(ii-1)+1:2*ii] = Q
    end
    planner = py"TLQG_PosRangeLocalization2D"(N, dto, Nt, Q_tlqg, Cs, Cu,
                                              u_param_min, u_param_max)
    return TLQG_Planner_PosRangeLocalization2D(planner, nothing, nothing, nothing, nothing,
                                               kl_threshold, nothing)
end

function get_kl_divergence(s1::AugState{<:PhysPos,<:VecBelMvNormal},
                           s2::AugState{<:PhysPos,<:VecBelMvNormal})
    Nt = length(s1.b)
    @assert length(s2.b) == Nt
    kl = 0.0;
    for ii = 1:Nt
        b1 = s1.b[ii]
        μ1, Σ1 = b1.params[1:b1.dim], reshape(b1.params[b1.dim+1:end],b1.dim,b1.dim);
        if !isposdef(Σ1)
            λ, V = eigen(Σ1)
            λ .= max.(1e-6, λ)
            Σ1 = Symmetric(V*diagm(λ)*V');
        end
        b2 = s2.b[ii]
        μ2, Σ2 = b2.params[1:b2.dim], reshape(b2.params[b2.dim+1:end],b2.dim,b2.dim);
        if !isposdef(Σ2)
            λ, V = eigen(Σ2)
            λ .= max.(1e-6, λ)
            Σ2 = Symmetric(V*diagm(λ)*V');
        end
        kl += 0.5*(tr(Σ2\Σ1) + (μ2 - μ1)'*(Σ2\(μ2 - μ1))
            - size(Σ2, 1) + logdet(Σ2) - logdet(Σ1))
    end
    return kl
end

function get_symmetric_kl(s1::AugState{<:PhysPos,<:VecBelMvNormal},
                          s2::AugState{<:PhysPos,<:VecBelMvNormal})
    symmetric_kl = 0.5*(get_kl_divergence(s1, s2) + get_kl_divergence(s2, s1))
    return symmetric_kl
end

function tlqgControlUpdate(tlqg::TLQG_Planner_PosRangeLocalization2D,
                           s::AugState{<:PhysPos,<:VecBelMvNormal},
                           UArray::Vector{<:PosControl};
                           verbose=true, online=true)
    dto = tlqg.planner.dto
    dim = s.b[1].dim;
    μ_current = vcat(s.p.pos, [b.params[1:dim] for b in s.b]...);
    Σ_current = zeros(length(μ_current)-s.p.dim, length(μ_current)-s.p.dim)
    for ii = 1:tlqg.planner.Nt
        Σ_current[dim*(ii-1)+1:dim*ii, dim*(ii-1)+1:dim*ii] = reshape(s.b[ii].params[dim+1:end], dim, dim)
    end
    replan_flag = false;
    if isnothing(tlqg.t_last_plan)
        # No planning done yet
        replan_flag = true
    else
        time_idx = Int64(round((s.t - tlqg.t_last_plan)/dto, digits=5)) + 1
        if time_idx > tlqg.planner.N
            # Reached the end of planning horizon
            if verbose
                println("Replanning since the end of planning horizon reached.")
            end
            replan_flag = true
        else
            # Symmetric KL larger than the replanning threshold
            ego = PhysPos(s.t, tlqg.μ_array[time_idx][1:dim]);
            bVec = [BelMvNormal(s.t, tlqg.μ_array[time_idx][dim+dim*(ii-1)+1:dim+dim*ii],
                                tlqg.Σ_array[time_idx][dim*(ii-1)+1:dim*ii, dim*(ii-1)+1:dim*ii])
                    for ii = 1:tlqg.planner.Nt];
            s_model = AugState(ego, bVec);
            kl_val = get_symmetric_kl(s_model, s)
            kl_str = @sprintf "%2.2f" kl_val
            if verbose
                println("KL: $(kl_str)")
            end
            kl_thresh_str = @sprintf "%2.2f" tlqg.kl_threshold
            if kl_val > tlqg.kl_threshold
                if verbose
                    println("Replanning since KL: $(kl_str) > $(kl_thresh_str).")
                end
                replan_flag = true
            end
        end
    end
    if replan_flag
        tlqg.μ_array, tlqg.Σ_array, tlqg.l_array, tlqg.L_array, elapsed =
            tlqg.planner.solve_tlqg(μ_current, Σ_current, verbose=false, online=online)
        for ii = 1:length(tlqg.Σ_array)
            tlqg.Σ_array[ii] = Symmetric(tlqg.Σ_array[ii])
            if !isposdef(tlqg.Σ_array[ii])
                @warn "At T = $(b.t)[s], tlqg.Σ_array[$(ii)] is not posdef. Fixing..."
                λ, V = eigen(tlqg.Σ_array[ii])
                λ .= max.(1e-6, λ)
                tlqg.Σ_array[ii] = Symmetric(V*diagm(λ)*V');
            end
        end
        tlqg.t_last_plan = s.t;
    else
        elapsed = nothing
    end
    time_idx = Int64(round((s.t - tlqg.t_last_plan)/dto, digits=5)) + 1
    if verbose
        println("Executing policy at time step $(time_idx) out of $(tlqg.planner.N)")
    end
    u_val = -tlqg.L_array[time_idx]*(μ_current - tlqg.μ_array[time_idx]) + tlqg.l_array[time_idx];
    # saturate control input
    for ii = 1:length(u_val)
        if u_val[ii] < tlqg.planner.u_param_min[ii]
            u_val[ii] = tlqg.planner.u_param_min[ii];
        elseif u_val[ii] > tlqg.planner.u_param_max[ii]
            u_val[ii] = tlqg.planner.u_param_max[ii];
        end
    end
    if verbose
        println(u_val)
    end
    UArray_new = copy(UArray)
    for ii = 1:length(UArray_new)
        if UArray_new[ii].t < round(s.t + dto,digits=5)
            UArray_new[ii] = PosControl(UArray_new[ii].t,u_val);
        end
    end
    return elapsed, UArray_new
end

mutable struct TLQG_Planner_Manipulate2D
    planner::PyObject
    μ_array::Union{Nothing, Vector{Vector{Float64}}}
    Σ_array::Union{Nothing, Vector{Matrix{Float64}}}
    l_array::Union{Nothing, Vector{Vector{Float64}}}
    L_array::Union{Nothing, Vector{Matrix{Float64}}}
    kl_threshold::Float64 # kl threshold to trigger replanning
    t_last_plan::Union{Nothing, Float64} # last time when replanning occured
end

function TLQG_Planner_Manipulate2D(Th::Float64, # planning horizon [s]
                                   dto::Float64,
                                   Q::Matrix{Float64}, R::Matrix{Float64},
                                   Cs::Matrix{Float64}, Cu::Matrix{Float64},
                                   u_param_min::Vector{Float64},
                                   u_param_max::Vector{Float64},
                                   posGain::Float64, rotGain::Float64,
                                   kl_threshold::Float64)
    x_target = [0.,0.,pi,0.,0.,0.];
    N = Int64(round(Th/dto, digits=5));
    planner = py"TLQG_Manipulate2D"(N, dto, Q, R, Cs, Cu, x_target,
                                    u_param_min, u_param_max, posGain, rotGain)
    return TLQG_Planner_Manipulate2D(planner, nothing, nothing, nothing, nothing,
                                     kl_threshold, nothing)
end

function get_kl_divergence(b1::BelMvNormal{Float64}, b2::BelMvNormal{Float64})
    μ1, Σ1 = b1.params[1:b1.dim], reshape(b1.params[b1.dim+1:end],b1.dim,b1.dim);
    if !isposdef(Σ1)
        λ, V = eigen(Σ1)
        λ .= max.(1e-6, λ)
        Σ1 = Symmetric(V*diagm(λ)*V');
    end
    μ2, Σ2 = b2.params[1:b2.dim], reshape(b2.params[b2.dim+1:end],b2.dim,b2.dim);
    if !isposdef(Σ2)
        λ, V = eigen(Σ2)
        λ .= max.(1e-6, λ)
        Σ2 = Symmetric(V*diagm(λ)*V');
    end
    kl = 0.5*(tr(Σ2\Σ1) + (μ2 - μ1)'*(Σ2\(μ2 - μ1))
       - size(Σ2, 1) + logdet(Σ2) - logdet(Σ1))
    return kl
end

function get_symmetric_kl(b1::BelMvNormal{Float64}, b2::BelMvNormal{Float64})
    symmetric_kl = 0.5*(get_kl_divergence(b1, b2) + get_kl_divergence(b2, b1))
    return symmetric_kl
end

function tlqgControlUpdate(tlqg::TLQG_Planner_Manipulate2D,
                           b::BelMvNormal{Float64},
                           UArray::Vector{<:MControl2D};
                           verbose=true, online=true)
    dto = tlqg.planner.dto
    μ_current, Σ_current = b.params[1:b.dim], reshape(b.params[b.dim+1:end],b.dim,b.dim);
    replan_flag = false;
    if isnothing(tlqg.t_last_plan)
        # No planning done yet
        replan_flag = true
    else
        time_idx = Int64(round((b.t - tlqg.t_last_plan)/dto, digits=5)) + 1
        if time_idx > tlqg.planner.N
            # Reached the end of planning horizon
            if verbose
                println("Replanning since the end of planning horizon reached.")
            end
            replan_flag = true
        else
            # Symmetric KL larger than the replanning threshold
            b_model = BelMvNormal(b.t, tlqg.μ_array[time_idx], tlqg.Σ_array[time_idx])
            kl_val = get_symmetric_kl(b_model, b)
            kl_str = @sprintf "%2.2f" kl_val
            if verbose
                println("KL: $(kl_str)")
            end
            kl_thresh_str = @sprintf "%2.2f" tlqg.kl_threshold
            if kl_val > tlqg.kl_threshold
                if verbose
                    println("Replanning since KL: $(kl_str) > $(kl_thresh_str).")
                end
                replan_flag = true
            end
        end
    end
    if replan_flag
        tlqg.μ_array, tlqg.Σ_array, tlqg.l_array, tlqg.L_array, elapsed =
            tlqg.planner.solve_tlqg(μ_current, Σ_current, verbose=false, online=online)
        for ii = 1:length(tlqg.Σ_array)
            tlqg.Σ_array[ii] = Symmetric(tlqg.Σ_array[ii])
            if !isposdef(tlqg.Σ_array[ii])
                @warn "At T = $(b.t)[s], tlqg.Σ_array[$(ii)] is not posdef. Fixing..."
                λ, V = eigen(tlqg.Σ_array[ii])
                λ .= max.(1e-6, λ)
                tlqg.Σ_array[ii] = Symmetric(V*diagm(λ)*V');
            end
        end
        tlqg.t_last_plan = b.t;
    else
        elapsed = nothing
    end
    time_idx = Int64(round((b.t - tlqg.t_last_plan)/dto, digits=5)) + 1
    if verbose
        println("Executing policy at time step $(time_idx) out of $(tlqg.planner.N)")
    end
    u_val = -tlqg.L_array[time_idx]*(μ_current - tlqg.μ_array[time_idx]) + tlqg.l_array[time_idx];
    # saturate control input
    for ii = 1:length(u_val)
        if u_val[ii] < tlqg.planner.u_param_min[ii]
            u_val[ii] = tlqg.planner.u_param_min[ii];
        elseif u_val[ii] > tlqg.planner.u_param_max[ii]
            u_val[ii] = tlqg.planner.u_param_max[ii];
        end
    end
    UArray_new = copy(UArray)
    for ii = 1:length(UArray_new)
        if UArray_new[ii].t < round(b.t + dto,digits=5)
            UArray_new[ii] = MControl2D(UArray_new[ii].t,u_val);
        end
    end
    return elapsed, UArray_new
end


text:
using DataDrivenDiffEq
using ModelingToolkit
using LinearAlgebra
using OrdinaryDiffEq
using Test

@testset "Basis" begin
    @variables u[1:3]
    @parameters w[1:2]
    h = [1u[1]; 1u[2]; cos(w[1]*u[2]+w[2]*u[3]); 1u[3]+1u[2]]
    basis = Basis(h, u, parameters = w)
    basis_2 = unique(basis)
    @test size(basis) == size(h)
    # TODO This works fine for "manual" execution of the testset but fails
    # with Pkg.test
    @test basis([1.0; 2.0; π], p = [0. 1.]) ≈ [1.0; 2.0; -1.0; π+2.0]
    @test size(basis) == size(basis_2)
    push!(basis_2, sin(u[2]))
    @test size(basis_2)[1] == length(h)+1
    push!(basis, 1u[3]+1u[2])
    unique!(basis)
    @test size(basis) == size(h)
    g = [1.0*u[1]; 1.0*u[3]; 1.0*u[2]]
    basis = Basis(g, u, parameters = w)
    f = jacobian(basis)
    @test f([1;1;1], [0.0 0.0]) ≈ [1.0 0.0 0.0; 0.0 0.0 1.0; 0.0 1.0 0.0]
    @test_nowarn sys = ODESystem(basis)
end

@testset "DMD" begin
    # Create some linear data
    A = [0.9 -0.2; 0.0 0.2]
    y = [[10.; -10.]]
    for i in 1:10
        push!(y, A*y[end])
    end
    X = hcat(y...)
    @test_throws AssertionError ExactDMD(X[:, 1:end-2], dt = -1.0)
    estimator = ExactDMD(X[:,1:end-2])
    @test isstable(estimator)
    @test estimator.Ã ≈ A
    @test eigvals(estimator) ≈ eigvals(A)
    @test eigvecs(estimator) ≈ eigvecs(A)
    @test_nowarn dynamics(estimator)
    @test_throws AssertionError dynamics(estimator, discrete = false)
    @test_nowarn update!(estimator, X[:, end-1], X[:,end])
end

@testset "EDMD" begin
    # Test for linear system
    function linear_sys(u, p, t)
        x = -0.9*u[1]
        y = -0.3*u[2]
        return [x;y]
    end

    u0 = [π; 1.0]
    tspan = (0.0, 20.0)
    prob = DiscreteProblem(linear_sys, u0, tspan)
    sol = solve(prob,FunctionMap())

    @variables u[1:2]
    h = [1u[1]; 1u[2]; sin(u[1]); cos(u[1]); u[1]*u[2]]
    basis = Basis(h, u)

    @test_throws AssertionError ExtendedDMD(sol[:,:], basis, dt = -1.0)
    estimator = ExtendedDMD(sol[:,:], basis)
    @test basis == estimator.basis
    basis_2 = reduce_basis(estimator, threshold = 1e-5)
    @test size(basis_2)[1] < size(basis)[1]
    estimator_2 = ExtendedDMD(sol[:,:], basis_2)
    p1 = DiscreteProblem(dynamics(estimator), u0, tspan, [])
    s1 = solve(p1,FunctionMap())
    p2 = DiscreteProblem(dynamics(estimator_2), u0, tspan, [])
    s2 = solve(p2,FunctionMap())
    p3 = DiscreteProblem(linear_dynamics(estimator_2), estimator_2(u0), tspan, [])
    s3 = solve(p3,FunctionMap())
    @test sol[:,:] ≈ s1[:,:]
    @test sol[:,:] ≈ s2[:,:]
    @test hcat(estimator_2.basis.(eachcol(sol[:,:]))...)≈ s3[:,:]
    @test eigvals(estimator_2) ≈ [-0.9; -0.3]

    # Test for nonlinear system
    function nonlinear_sys(du, u, p, t)
        du[1] = sin(u[1])
        du[2] = -0.3*u[2] -0.9*u[1]
    end

    prob = DiscreteProblem(nonlinear_sys, u0, tspan)
    sol = solve(prob,FunctionMap())
    estimator = ExtendedDMD(sol[:,:], basis)
    p4 = DiscreteProblem(dynamics(estimator), u0, tspan, [])
    s4 = solve(p4,FunctionMap())
    @test sol[:,:] ≈ s4[:,:]
end

@testset "DMDc" begin
    # Define measurements from unstable system with known control input
    X = [4 2 1 0.5 0.25; 7 0.7 0.07 0.007 0.0007]
    U = [-4 -2 -1 -0.5]
    B = Float32[1; 0]

    # But with a little more knowledge
    sys = DMDc(X, U, B = B)
    @test isa(get_dynamics(sys), ExactDMD)
    @test sys.koopman.Ã ≈[1.5 0; 0 0.1]
    @test get_input_map(sys) ≈ [1.0; 0.0]
    @test !isstable(sys)
    @test_nowarn eigen(sys)

    # Check the solution of an unforced and forced system against each other
    dudt_ = dynamics(sys)
    prob = DiscreteProblem(dudt_, X[:, 1], (0., 10.))
    sol_unforced = solve(prob,  FunctionMap())

    dudt_ = dynamics(sys, control = (u, p, t) -> -0.5u[1])
    prob = DiscreteProblem(dudt_, X[:, 1], (0., 10.))
    sol = solve(prob, FunctionMap())

    @test all(abs.(diff(sol[1,:])) .< 1e-5)
    @test sol[2,:] ≈ sol_unforced[2,:]
end

@testset "SInDy" begin
    # Test the pendulum
    function pendulum(u, p, t)
        du1 = u[2]
        du2 = -sin(u[1]) - 0.1*u[2]
        return [du1; du2]
    end
    u0 = [0.99π; 0.3]
    tspan = (0.0, 10.0)
    prob = ODEProblem(pendulum, u0, tspan)
    sol = solve(prob,Tsit5())
    # Create the differential data
    DX = similar(sol[:,:])
    for (i, xi) in enumerate(eachcol(sol[:,:]))
        DX[:,i] = pendulum(xi, [], 0.0)
    end
    # Create a basis
    @variables u[1:2]

    polys = [u[1]^0]
    for i ∈ 1:3
        for j ∈ 1:3
            push!(polys, u[1]^i*u[2]^j)
        end
    end

    h = [1u[1];1u[2]; cos(u[1]); sin(u[1]); u[1]*u[2]; u[1]*sin(u[2]); u[2]*cos(u[2]); polys...]

    basis = Basis(h, u, parameters = [])
    Ψ = SInDy(sol[:,:], DX, basis, ϵ = 1e-2)
    @test size(Ψ)[1] == 2

    # Simulate
    estimator = ODEProblem(dynamics(Ψ), u0, tspan, [])
    sol_ = solve(estimator,Tsit5())
    @test sol[:,:] ≈ sol_[:,:]
end


text:
using Documenter, Maxima

makedocs(
    modules = [Maxima],
    clean = false,
    format = :html,
    sitename = "Maxima.jl",
    pages = Any[
        "Home" => "index.md",
        "Manual" => "man/manual.md",
        "Library" => Any[
            "lib/basics.md",
            "lib/simplify.md",
            "lib/calculus.md",
            "lib/plot.md"]
    ]
)

deploydocs(
    deps = nothing,
	repo = "github.com/nsmith5/Maxima.jl.git",
    julia = "0.5",
    osname = "osx",
    make = nothing,
	target = "build"
)


text:
 using JLD2
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
#using SparseArrays, LinearAlgebra
using ViewerGL
GL = ViewerGL

# created by surface_extraction_parallel_ircad01.jl
@load "test/liver01.jld2" V FV
FW = [[v1,v2,v4,v3] for (v1,v2,v3,v4) in FV]
GL.VIEW([ GL.GLGrid(V,FW,GL.COLORS[2]) ]);

# created by surface_extraction_parallel_ircad01.jl
@load "test/liver02.jld2" V FV
TV = cat([[[v1,v2,v3],[v1,v4,v3]] for (v1,v2,v3,v4) in FV])
TV = convert(Lar.Cells,TV)
GL.VIEW([ GL.GLGrid(V,TV,GL.COLORS[1]) ]);


text:
using Test
using BetweenFlags

@testset "Flags" begin
    for Flag in (StartFlag, StopFlag)
        flag = Flag("flag")
        @test flag.flag == "flag"
        @test flag.flag_boundaries_left == []
        @test flag.flag_boundaries_right == []
        @test flag.trigger == Dict{Tuple,String}(("", "") => "flag")
        @test BetweenFlags.grep_type(flag) == GreedyType

        flag = Flag("flag",
                    ["LB1","LB2"],
                    ["RB1","RB2"];
                    grep_type=ScopeType())
        @test flag.flag == "flag"
        @test flag.flag_boundaries_left == ["LB1","LB2"]
        @test flag.flag_boundaries_right == ["RB1","RB2"]
        @test flag.trigger == Dict{Tuple,String}(
            ("LB2", "RB1") => "LB2flagRB1",
            ("LB2", "RB2") => "LB2flagRB2",
            ("LB1", "RB2") => "LB1flagRB2",
            ("LB1", "RB1") => "LB1flagRB1")
        @test BetweenFlags.grep_type(flag) == ScopeType

        flag = Flag("flag",
                    ["LB1","LB2"],
                    ["RB1","RB2"];
                    grep_type=ScopeType())
        @test flag.flag == "flag"
        @test flag.flag_boundaries_left == ["LB1","LB2"]
        @test flag.flag_boundaries_right == ["RB1","RB2"]
        @test flag.trigger == Dict{Tuple,String}(
            ("LB2", "RB1") => "LB2flagRB1",
            ("LB2", "RB2") => "LB2flagRB2",
            ("LB1", "RB2") => "LB1flagRB2",
            ("LB1", "RB1") => "LB1flagRB1")
        @test BetweenFlags.grep_type(flag) == ScopeType
    end
end


text:
using Multilane
using MCTS
using POMDPToolbox
using POMDPs
# using POMCP
using Missings
using DataFrames
using CSV
using POMCPOW
using Distributions

@everywhere using Missings
@everywhere using Multilane
@everywhere using POMDPToolbox

@show n_iters = 1000
@show max_time = Inf
@show max_depth = 40
@show val = SimpleSolver()
alldata = DataFrame()

function gen_sims(p, n, k, i)
    c = p[1]
    k_state = p[2]
    alpha_state = 1.0/p[3]
    dpws = DPWSolver(depth=max_depth,
                     n_iterations=n_iters,
                     max_time=max_time,
                     exploration_constant=c,
                     k_state=k_state,
                     alpha_state=alpha_state,
                     enable_action_pw=false,
                     check_repeat_state=false,
                     estimate_value=RolloutEstimator(val)
                     # estimate_value=val
                    )

    cor = 1.0
    lambda = 1.0

    behaviors = standard_uniform(correlation=cor)
    pp = PhysicalParam(4, lane_length=100.0)
    dmodel = NoCrashIDMMOBILModel(10, pp,
                                  behaviors=behaviors,
                                  p_appear=1.0,
                                  lane_terminate=true,
                                  max_dist=1000.0
                                 )
    rmodel = SuccessReward(lambda=lambda)
    mdp = NoCrashMDP{typeof(rmodel), typeof(behaviors)}(dmodel, rmodel, 0.95, false)

    sims = []

    for j in 1:n
        rng_seed = j + 40000*i*k
        planner = deepcopy(solve(dpws, mdp))
        srand(planner, rng_seed+50000)
        rng = MersenneTwister(rng_seed + 60000)
        is = initial_state(mdp, rng)
        ips = MLPhysicalState(is)

        sim = Sim(deepcopy(mdp),
                  deepcopy(planner),
                  rng=rng,
                  max_steps=100,
                  metadata=Dict(:i=>i, :k=>k)
                 )
        push!(sims, sim)
    end

    return sims
end

function make_updater(cor, problem, rng_seed)
    wup = WeightUpdateParams(smoothing=0.0, wrong_lane_factor=0.05)
    if cor >= 0.5
        return AggressivenessUpdater(problem, 2000, 0.05, 0.1, wup, MersenneTwister(rng_seed+50000))
    else
        return BehaviorParticleUpdater(problem, 5000, 0.05, 0.2, wup, MersenneTwister(rng_seed+50000))
    end
end

# params are c, k_state, alpha_state
start_mean = [2.0, 4.0, 8.0]
start_cov = diagm([10.0^2, 10.0^2, 30.0^2])
d = MvNormal(start_mean, start_cov)
K = 60 # 60 # number of parameter samples
n = 100 # 100 # number of evaluation simulations
m = 15  # 15 # number of elite samples
max_iters = 100

for i in 1:max_iters
    sims = []
    params = Vector{Vector{Float64}}(undef,K)
    print("creating $K simulation sets")
    for k in 1:K
        p = rand(d)
        p[1] = max(0.0, p[1])
        p[2] = max(1.0, p[2])
        p[3] = max(2.0, p[3])
        params[k] = p
        k_sims = gen_sims(p, n, k, i)
        print(".")
        append!(sims, k_sims)
    end
    println()
    results = run_parallel(sims)
    # results = run(sims)
    combined = by(results, :k) do df
        DataFrame(mean_reward=mean(df[:reward]))
    end
    @show mean(combined[:mean_reward])
    order = sortperm(combined[:mean_reward])
    elite = params[combined[:k][order[K-m+1:end]]]
    elite_matrix = Matrix{Float64}(length(start_mean), m)
    for k in 1:m
        elite_matrix[:,k] = elite[k]
    end
    try
        d = fit(typeof(d), elite_matrix)
    catch ex
        if ex isa Base.LinAlg.PosDefException
            println("pos def exception")
            d = fit(typeof(d), elite_matrix += 0.01*randn(size(elite_matrix)))
        else
            rethrow(ex)
        end
    end
    println("iteration $i")
    @show mean(d)
    @show det(cov(d))
    @show ev = eigvals(cov(d))
    for j in 1:length(ev)
        @show eigvecs(cov(d))[:,j]
    end
end


text:
"""
abstract type that allows users to pass in their preferred Algorithm
"""
abstract type AbstractFlowAlgorithm end

"""
Forces the maximum_flow function to use the Edmonds–Karp algorithm.
"""
struct EdmondsKarpAlgorithm <: AbstractFlowAlgorithm end

"""
Forces the maximum_flow function to use Dinic\'s algorithm.
"""
struct DinicAlgorithm <: AbstractFlowAlgorithm end

"""
Forces the maximum_flow function to use the Boykov-Kolmogorov algorithm.
"""
struct BoykovKolmogorovAlgorithm <: AbstractFlowAlgorithm end

"""
Forces the maximum_flow function to use the Push-Relabel algorithm.
"""
struct PushRelabelAlgorithm <: AbstractFlowAlgorithm end

"""
Type that returns 1 if a forward edge exists, and 0 otherwise
"""
mutable struct DefaultCapacity{G<:ADiGraph, I<:Integer} <: AbstractMatrix{I}
    g::G
    nv::I
end
DefaultCapacity(g::G) where {G<:ADiGraph} = DefaultCapacity(g, signed(nv(g)))

getindex(d::DefaultCapacity, s, t) = has_edge(d.g, s , t) ? 1 : 0
size(d::DefaultCapacity) = (d.nv, d.nv)

function _complete(g::ADiGraph, capacity_matrix::AbstractMatrix{T}) where T
    S = signedtype(T)
    c = Vector{Vector{S}}()
    for i=1:nv(g)
        neigs = neighbors(g, i)
        push!(c, zeros(S, length(neigs)))
        for (k, j) in enumerate(neigs)
            if has_edge(g, i, j)
                val = capacity_matrix[i, j]
                @assert val >= 0 "Capacities should be non-negative"
                c[i][k] = val
            end #else 0
        end
    end
    return c
end

"""
`pl[i][k]` is the position of vertex `i`
in the adjlist of its neighbour `j=fadj[i][k]`,
i.e. `i == adj[j][pl[i][k]]`.
"""
function poslist(g::AGraphOrDiGraph)
    pl = Vector{Vector{Int}}()
    for i=1:nv(g)
        push!(pl, zeros(Int, degree(g, i)))
        for (k,j) in enumerate(neighbors(g, i))
            p = countfirst(neighbors(g,j), i)
            @assert p > 0
            pl[i][k] = p
        end
    end
    return pl
end

"""
    residual_graph{G<:ADiGraph}(g::G, capacity, flow)

Computers the residual graph of `g` associated to given `flow` and `capacity`.
See wikipedia.
"""
function residual_graph(g::G,
        capacity::Vector{Vector{T}}, flow::Vector{Vector{T}}
    ) where {G<:ADiGraph,T<:Number}
    h = G(nv(g))
    for i=1:nv(g)
        c = capacity[i]
        f = flow[i]
        for (k, j) in enumerate(neighbors(g, i))
            if c[k] - f[k] > 0
                add_edge!(h, i, j)
            end
        end
    end
    return h
end

function residual_graph(g::G, capacity::AbstractMatrix, flow::AbstractMatrix) where G<:ADiGraph
    h = G(nv(g))
    for e in edges(g)
        i, j = src(e), dst(e)
        if capacity[i,j] - flow[i,j] > 0
            add_edge!(h, i, j)
        end
    end
    return h
end

function cut_labels(g::G, source, capacity, flow) where G<:ADiGraph
    h = residual_graph(g, capacity, flow)
    svertices = neighborhood(h, source, nv(h))
    labels = fill(2, nv(h))
    for i in svertices
        labels[i] = 1
    end
    return labels
end

"""
    maximum_flow{T<:Number}(
                        g::ADiGraph,
                        source::Int,
                        target::Int,
                        capacity_matrix::AbstractMatrix{T} =
                            DefaultCapacity(g);
                        algorithm::AbstractFlowAlgorithm  =
                            PushRelabelAlgorithm(),
                        restriction::T = zero(T)
                    )

Generic maximum_flow function.
The function defaults to the Push-Relabel (also called Preflow) algorithm. Alternatively, the algorithm
to be used can also be specified through a keyword argument. A default capacity of 1
is assumed for each link if no capacity matrix is provided.
If the restriction is bigger than 0, it is applied to capacity_matrix.

All algorithms return a tuple with
1) the maximum flow
2) the flow matrix
3) the labelling associated to the minimum cut

Available algorithms are `DinicAlgorithm`, `EdmondsKarpAlgorithm`, `BoykovKolmogorovAlgorithm`
and `PushRelabelAlgorithm`.

Time complexity is O(V²√E) for the push relabel algorithm.

### Usage Example:
```julia

# Create a flow-graph and a capacity matrix
g = DiGraph(8)
flow_edges = [
    (1,2,10),(1,3,5),(1,4,15),(2,3,4),(2,5,9),
    (2,6,15),(3,4,4),(3,6,8),(4,7,16),(5,6,15),
    (5,8,10),(6,7,15),(6,8,10),(7,3,6),(7,8,10)
]
capacity_matrix = zeros(Int, 8, 8)
for e in flow_edges
    u, v, f = e
    add_edge!(g, u, v)
    capacity_matrix[u,v] = f
end

# Run default maximum_flow without the capacity_matrix (assumes capacity 1. on each edge).
f, F, labels = maximum_flow(g, 1, 8)

# Run Endmonds-Karp algorithm
f, F, labels = maximum_flow(g,1,8,capacity_matrix,algorithm=EdmondsKarpAlgorithm())
```
"""
function maximum_flow(
        g::G,                   # the input graph
        source,                           # the source vertex
        target,                           # the target vertex
        capacity_matrix::AbstractMatrix{T} =  # edge flow capacities
            DefaultCapacity(g);
        algorithm::AbstractFlowAlgorithm  =    # keyword argument for algorithm
            PushRelabelAlgorithm(),
        restriction::T = zero(T)
    ) where {G<:ADiGraph, T<:Number}
    if restriction > zero(T)
        capacity_matrix = min.(restriction, capacity_matrix)
    end
    flow_graph = complete(g)
    if  algorithm == DinicAlgorithm()
        f, F =  dinic_impl(flow_graph, source, target, capacity_matrix)
    elseif algorithm == EdmondsKarpAlgorithm()
        f, F =  edmonds_karp_impl(flow_graph, source, target, capacity_matrix)
    elseif algorithm == BoykovKolmogorovAlgorithm()
        f, F, labels =  boykov_kolmogorov_impl(flow_graph, source, target, capacity_matrix)
    elseif algorithm == PushRelabelAlgorithm()
        c = _complete(flow_graph, capacity_matrix)
        pos = poslist(flow_graph)
        f, Fvec = push_relabel_impl(flow_graph, source, target, c, pos)
        F = spzeros(T, nv(g), nv(g))
        for i=1:nv(g)
            for (k,j) in enumerate(neighbors(flow_graph, i))
                F[i,j] = Fvec[i][k]
            end
        end
    end
    labels = cut_labels(flow_graph, source, capacity_matrix, F)

    return f, F, labels
end

function maximum_flow(
        flow_graph::G,                   # the input graph
        source,                           # the source vertex
        target,                           # the target vertex
        capacities::Vector{Vector{T}}
    ) where {G<:ADiGraph, T<:Number}
    pos = poslist(flow_graph)
    f, F = push_relabel_impl(flow_graph, source, target, capacities, pos)
    labels = cut_labels(flow_graph, source, capacities, F)
    return f, F, labels
end

"""
    minimum_cut(g, s, t, capacity_matrix=DefaultCapacity(); kws...)

Finds the `s-t cut` of minimal weight according to the `capacities` matrix on
the directed graph `g`.
The solution is found through a maximal flow algorithm.
See [`maximum_flow`](@ref) for the optional arguments.

Returns a triple `(f, cut, labels)`, where `f` is the weight of the cut,
`cut` is a vector of the edges in the cut, and `labels` gives a partitioning
of the vertices in two sets, according to the cut.
"""
function minimum_cut(g::ADiGraph, s::Integer, t::Integer, args...; kws...)
    f, F, labels = maximum_flow(g, s, t, args...; kws...)
    E = edgetype(g)
    cut = Vector{E}()
    for e in edges(g)
        if labels[src(e)] == 1 && labels[dst(e)] == 2
            push!(cut, e)
        end
    end
    f, cut, labels
end


text:
using CSV
using DataFrames
using DataFramesMeta
using DataStructures
using CodecZlib
using XAM
using FASTX
using BioSequences
using BioAlignments
using ArgParse2
using st16SseqJuliaTools

"Parse genus asignment - specific for fasta header"
function parse_genus_info(inf::String)
    out = Dict{String,String}()
    df = DataFrame(CSV.File(inf))
    for r in eachrow(df)
	out[string(r.ID)] = split(split(r.Genus,"-")[1],"_")[1]
    end
    return(out)
end

function main()
    parser = ArgumentParser(prog = "analyse_alignment_on_reference.jl",
                        description = "Extract good contigs based on expected contigs of a standard")
    add_argument!(parser, "--inbam", "-i", help = "input sam ",type = String)
    add_argument!(parser, "--output", "-o", help = "File with names",type = String)
    add_argument!(parser, "--ref","-r", help = "Reference",type = String)
    add_argument!(parser, "--contig","-c", help = "Contig",type = String)
    add_argument!(parser, "--genus","-g", help = "Genus",type = String)
    add_argument!(parser, "--contig-fraction","-f", help = "Fraction of a contig that should be aligned",type = Float64,default = 0.7)
    add_argument!(parser, "--contig-identity","-d", help = "Minimum identity of an aligned contig",type = Float64,default = 0.7)
    args = parse_args(parser)
    #Not command line parameters
    min_al_idfrac = args.contig_identity # length of contig aligned
    min_al_lenfrac = args.contig_fraction  # identity fraction
    write_ref_seq = false # write matching reference sequence instead of a contig - for testing

    #dat of all matches on reference will be collected for filtering
    matches = DataFrame(Contig_ID = Array{String,1}(),
		        Ref_ID = Array{String,1}(),
		        Contig_length = Array{Int64,1}(),
			Match_length = Array{Int64,1}(),
			Identity = Array{Float64,1}(),
			Coverage_contig = Array{Float64,1}(),
			Coverage_ref = Array{Float64,1}(),
			Contig_start = Array{Int64,1}(),
			Contig_end = Array{Int64,1}(),
			Ref_start = Array{Int64,1}(),
			Ref_end = Array{Int64,1}(),
			Ref_genus = Array{String,1}(),
			Expected_genus = Array{Bool,1}())

    out_seqs = OrderedDict{String,LongSequence{DNAAlphabet{4}}}()
    genus_assigns = parse_genus_info(args.genus)
    costmodel = BioAlignments.CostModel(match=0, mismatch=1, insertion=1, deletion=1);
    testseqs =  fasta_to_dict(args.contig)
    ref =  fasta_to_dict(args.ref)
    print(args.inbam)
    reader = open(BAM.Reader,args.inbam)
    ct = 0

    for record in reader
        if BAM.ismapped(record) #&& SAM.isprimary(record)
            ct += 1
	    pos = BAM.position(record)
            seq = BAM.sequence(record)
            rpos = BAM.rightposition(record)
            aln = XAM.BAM.alignment(record)
            s = first(aln.anchors).refpos
            e = last(aln.anchors).refpos
            l = abs(s-e)
            rn = BAM.refname(record)
	    qual = XAM.BAM.mappingquality(record)
	    tempname = BAM.tempname(record)
	    genus = genus_assigns[split(tempname,"_")[1]]
	    real_length = length(testseqs[BAM.tempname(record)])
	    tempname = BAM.tempname(record)
	    t_pos=BioAlignments.ref2seq(aln,pos)
	    t_rpos=BioAlignments.ref2seq(aln,rpos)
	    rtseq = ref[rn][pos:rpos]
	    tseq = testseqs[tempname][t_pos[1]:t_rpos[1]]
	    ff = findfirst(lowercase(genus),lowercase(rn))
	    expected_g = false
	    if ff != nothing && ff.start == 1
		expected_g = true
	    end

	    #println("PAIR: \n>r\n$rtseq\n>t\n$tseq\n---------------")
	    pairalnres = BioAlignments.pairalign(BioAlignments.EditDistance(),rtseq,tseq,costmodel)
	    pairaln = BioAlignments.alignment(pairalnres)
	    matched = BioAlignments.count_matches(pairaln)
	    idfrac = matched/l
	    lfrac_c = l/real_length
	    lfrac_r = l/length(ref[rn])
	    df_row=[tempname,rn,real_length,l,idfrac,lfrac_c,lfrac_r,t_pos[1],t_rpos[1],pos,rpos,genus,expected_g]
	    push!(matches,df_row)

        end
    end
    close(reader)
    matches_expected = @where(matches, :Expected_genus .== 1)
    ref_c_ds = groupby(matches_expected,:Ref_ID)
    lens =Array{Int64,1}()
    chosen_ids = Array{String,1}()
    for ref_c_d in ref_c_ds
	df = DataFrame(ref_c_d)
	df = @where(df,:Coverage_contig .>= min_al_lenfrac, :Identity .>= min_al_idfrac)
	df = sort(df, :Match_length,rev=true)
	s = df[1,:]
	push!(chosen_ids, s.Contig_ID)
	if write_ref_seq
	    out_seqs[s.Contig_ID] = LongSequence{DNAAlphabet{4}}(ref[s.Ref_ID][s.Ref_start:s.Ref_end])
	else
	    out_seqs[s.Contig_ID] = LongSequence{DNAAlphabet{4}}(testseqs[s.Contig_ID])#[s.Contig_start:s.Contig_end]
	end
	push!(lens,length(testseqs[s.Contig_ID]))
    end
    dict_to_fasta(out_seqs,args.output)#,cut_end = minimum(lens))
    Proper = Array{Bool,1}()
    for i in matches.Contig_ID
	push!(Proper,(i in chosen_ids))
    end
    matches.Proper = Proper
    CSV.write(args.output*".info.csv", matches,delim=';')
end

main()


text:
using FileIO,LinearAlgebra,Plots

## Get Test Data
cd(@__DIR__)
displayname = "ROGPG279Q"
datadir = "./$displayname"
figfmt = [".png",".svg"]

dataset = load(joinpath(datadir,"Test_CondTest_Duration.jld2"))

## Duration Plots
title = "From Photodiode Detected Lag-Corrected Signal\n(VSync=On, GSync=On, RefreshRate=144Hz)"
ylabel = "Condition Duration ($(round(Int,dataset["conddur1"]))ms)"
plot(dataset["ctdur1"];xlabel="Condition Tests",ylabel,title,leg=false)

foreach(ext->savefig(joinpath(datadir,"$ylabel$ext")),figfmt)

max_stars_repo_name: VLAB-ExperimentSystem/VLab

max_issues_repo_name: VLAB-ExperimentSystem/VLab

max_forks_repo_name: VLAB-ExperimentSystem/VLab


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

const StdUvDist = Union{StandardUvUniform, StandardUvNormal}
const StdMvDist = Union{StandardMvUniform, StandardMvNormal}

_adignore(f) = f()

function ChainRulesCore.rrule(::typeof(_adignore), f)
    result = _adignore(f)
    _nogradient_pullback(ΔΩ) = (NoTangent(), ZeroTangent())
    return result, _nogradient_pullback
end

macro _adignore(expr)
    :(_adignore(() -> $(esc(expr))))
end

function _pushfront(v::AbstractVector, x)
    T = promote_type(eltype(v), typeof(x))
    r = similar(v, T, length(eachindex(v)) + 1)
    r[firstindex(r)] = x
    r[firstindex(r)+1:lastindex(r)] = v
    r
end

function ChainRulesCore.rrule(::typeof(_pushfront), v::AbstractVector, x)
    result = _pushfront(v, x)
    function _pushfront_pullback(thunked_ΔΩ)
        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
        (NoTangent(), ΔΩ[firstindex(ΔΩ)+1:lastindex(ΔΩ)], ΔΩ[firstindex(ΔΩ)])
    end
    return result, _pushfront_pullback
end

function _pushback(v::AbstractVector, x)
    T = promote_type(eltype(v), typeof(x))
    r = similar(v, T, length(eachindex(v)) + 1)
    r[lastindex(r)] = x
    r[firstindex(r):lastindex(r)-1] = v
    r
end

function ChainRulesCore.rrule(::typeof(_pushback), v::AbstractVector, x)
    result = _pushback(v, x)
    function _pushback_pullback(thunked_ΔΩ)
        ΔΩ = ChainRulesCore.unthunk(thunked_ΔΩ)
        (NoTangent(), ΔΩ[firstindex(ΔΩ):lastindex(ΔΩ)-1], ΔΩ[lastindex(ΔΩ)])
    end
    return result, _pushback_pullback
end

_dropfront(v::AbstractVector) = v[firstindex(v)+1:lastindex(v)]

_dropback(v::AbstractVector) = v[firstindex(v):lastindex(v)-1]

_rev_cumsum(xs::AbstractVector) = reverse(cumsum(reverse(xs)))

function ChainRulesCore.rrule(::typeof(_rev_cumsum), xs::AbstractVector)
    result = _rev_cumsum(xs)
    function _rev_cumsum_pullback(ΔΩ)
        ∂xs = ChainRulesCore.@thunk cumsum(ChainRulesCore.unthunk(ΔΩ))
        (NoTangent(), ∂xs)
    end
    return result, _rev_cumsum_pullback
end

# Equivalent to `cumprod(xs)``:
_exp_cumsum_log(xs::AbstractVector) = exp.(cumsum(log.(xs)))

function ChainRulesCore.rrule(::typeof(_exp_cumsum_log), xs::AbstractVector)
    result = _exp_cumsum_log(xs)
    function _exp_cumsum_log_pullback(ΔΩ)
        ∂xs = inv.(xs) .* _rev_cumsum(exp.(cumsum(log.(xs))) .* ChainRulesCore.unthunk(ΔΩ))
        (NoTangent(), ∂xs)
    end
    return result, _exp_cumsum_log_pullback
end

struct DistributionTransform{
    VT <: AbstractValueShape,
    VF <: AbstractValueShape,
    DT <: ContinuousDistribution,
    DF <: ContinuousDistribution
} <: VariateTransform{VT,VF}
    target_dist::DT
    source_dist::DF
    _valshape::VT
    _varshape::VF
end

# ToDo: Add specialized dist trafo types able to cache relevant quantities, etc.

function _distrafo_ctor_impl(target_dist::Distribution, source_dist::Distribution)
    @argcheck eff_totalndof(target_dist) == eff_totalndof(source_dist)
    _valshape = varshape(target_dist)
    _varshape = varshape(source_dist)
    DistributionTransform(target_dist, source_dist, _valshape, _varshape)
end

DistributionTransform(target_dist::Distribution{VF,Continuous}, source_dist::Distribution{VF,Continuous}) where VF =
    _distrafo_ctor_impl(target_dist, source_dist)

DistributionTransform(target_dist::Distribution{Multivariate,Continuous}, source_dist::Distribution{VF,Continuous}) where VF =
    _distrafo_ctor_impl(target_dist, source_dist)

DistributionTransform(target_dist::Distribution{VF,Continuous}, source_dist::Distribution{Multivariate,Continuous}) where VF =
    _distrafo_ctor_impl(target_dist, source_dist)

DistributionTransform(target_dist::Distribution{Multivariate,Continuous}, source_dist::Distribution{Multivariate,Continuous}) =
    _distrafo_ctor_impl(target_dist, source_dist)

show_distribution(io::IO, d::Distribution) = show(io, d)
function show_distribution(io::IO, d::NamedTupleDist)
    print(io, Base.typename(typeof(d)).name, "{")
    show(io, propertynames(d))
    print(io, "}(…)")
end

function Base.show(io::IO, trafo::DistributionTransform)
    print(io, Base.typename(typeof(trafo)).name, "(")
    show_distribution(io, trafo.target_dist)
    print(io, ", ")
    show_distribution(io, trafo.source_dist)
    print(io, ")")
end

Base.show(io::IO, M::MIME"text/plain", trafo::DistributionTransform) = show(io, trafo)

# apply_dist_trafo(trg_d, src_d, src_v, prev_ladj)
function apply_dist_trafo end

apply_dist_trafo_noladj(trg_d::Distribution, src_d::Distribution, src_v::Any) = apply_dist_trafo(trg_d, src_d, src_v, missing).v

# Use ForwardDiff for univariate distribution transformations:
@inline function ChainRulesCore.rrule(::typeof(apply_dist_trafo), trg_d::Distribution{Univariate}, src_d::Distribution{Univariate}, src_v::Any, prev_ladj::OptionalLADJ)
    ChainRulesCore.rrule(fwddiff(apply_dist_trafo), trg_d, src_d, src_v, prev_ladj)
end
@inline function ChainRulesCore.rrule(::typeof(apply_dist_trafo_noladj), trg_d::Distribution{Univariate}, src_d::Distribution{Univariate}, src_v::Any)
    ChainRulesCore.rrule(fwddiff(apply_dist_trafo_noladj), trg_d, src_d, src_v)
end

import Base.∘
function ∘(a::DistributionTransform, b::DistributionTransform)
    @argcheck a.source_dist == b.target_dist
    DistributionTransform(a.target_dist, b.source_dist)
end

Base.inv(trafo::DistributionTransform) = DistributionTransform(trafo.source_dist, trafo.target_dist)

ValueShapes.varshape(trafo::DistributionTransform) = trafo._varshape
ValueShapes.valshape(trafo::DistributionTransform) = trafo._valshape

function apply_vartrafo_impl(trafo::DistributionTransform, v::Any, prev_ladj::OptionalLADJ)
    apply_dist_trafo(trafo.target_dist, trafo.source_dist, v, prev_ladj)
end

const _StdDistType = Union{Uniform, Normal}

_trg_disttype(::Type{Uniform}, ::Type{Univariate}) = StandardUvUniform
_trg_disttype(::Type{Uniform}, ::Type{Multivariate}) = StandardMvUniform
_trg_disttype(::Type{Normal}, ::Type{Univariate}) = StandardUvNormal
_trg_disttype(::Type{Normal}, ::Type{Multivariate}) = StandardMvNormal

function _trg_dist(disttype::Type{<:_StdDistType}, source_dist::Distribution{Univariate,Continuous})
    trg_dt = _trg_disttype(disttype, Univariate)
    trg_dt()
end

function _trg_dist(disttype::Type{<:_StdDistType}, source_dist::Distribution{Multivariate,Continuous})
    trg_dt = _trg_disttype(disttype, Multivariate)
    trg_dt(eff_totalndof(source_dist))
end

function _trg_dist(disttype::Type{<:_StdDistType}, source_dist::ContinuousDistribution)
    trg_dt = _trg_disttype(disttype, Multivariate)
    trg_dt(eff_totalndof(source_dist))
end

function DistributionTransform(disttype::Type{<:_StdDistType}, source_dist::ContinuousDistribution)
    trg_d = _trg_dist(disttype, source_dist)
    DistributionTransform(trg_d, source_dist)
end

function std_dist_from(src_d::Distribution)
    throw(ArgumentError("No standard intermediate distribution defined to transform from $(typeof(src_d).name)"))
end

function std_dist_to(trg_d::Distribution)
    throw(ArgumentError("No standard intermediate distribution defined to transform into $(typeof(trg_d).name)"))
end

@inline function _intermediate_std_dist(trg_d::Distribution, src_d::Distribution)
    _select_intermediate_dist(std_dist_to(trg_d), std_dist_from(src_d))
end

@inline _intermediate_std_dist(::Union{StdUvDist,StdMvDist}, src_d::Distribution) = std_dist_from(src_d)

@inline _intermediate_std_dist(trg_d::Distribution, ::Union{StdUvDist,StdMvDist}) = std_dist_to(trg_d)

function _intermediate_std_dist(::Union{StdUvDist,StdMvDist}, ::Union{StdUvDist,StdMvDist})
    throw(ArgumentError("Direct conversions must be used between standard intermediate distributions"))
end

@inline _select_intermediate_dist(a::D, ::D) where D<:Union{StdUvDist,StdMvDist} = a
@inline _select_intermediate_dist(a::D, ::D) where D<:Union{StandardUvUniform,StandardMvUniform} = a
@inline _select_intermediate_dist(a::Union{StandardUvUniform,StandardMvUniform}, ::Union{StdUvDist,StdMvDist}) = a
@inline _select_intermediate_dist(::Union{StdUvDist,StdMvDist}, b::Union{StandardUvUniform,StandardMvUniform}) = b

_check_conv_eff_totalndof(trg_d::Uniform, src_d::Uniform) = nothing

function _check_conv_eff_totalndof(trg_d::Distribution, src_d::Distribution)
    trg_d_n = eff_totalndof(trg_d)
    src_d_n = eff_totalndof(src_d)
    if trg_d_n != src_d_n
        throw(ArgumentError("Can't convert to $(typeof(trg_d).name) with $(trg_d_n) eff. DOF from $(typeof(src_d).name) with $(src_d_n) eff. DOF"))
    end
    nothing
end

function apply_dist_trafo(trg_d::Distribution, src_d::Distribution, src_v::Any, prev_ladj::OptionalLADJ)
    _check_conv_eff_totalndof(trg_d, src_d)
    intermediate_d = _intermediate_std_dist(trg_d, src_d)
    intermediate_v, intermediate_ladj = apply_dist_trafo(intermediate_d, src_d, src_v, prev_ladj)
    apply_dist_trafo(trg_d, intermediate_d, intermediate_v, intermediate_ladj)
end

function apply_dist_trafo(trg_d::DT, src_d::DT, src_v::Real, prev_ladj::OptionalLADJ) where {DT <: StdUvDist}
    (v = src_v, ladj = prev_ladj)
end

function apply_dist_trafo(trg_d::DT, src_d::DT, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ) where {DT <: StdMvDist}
    @argcheck length(trg_d) == length(src_d) == length(eachindex(src_v))
    (v = src_v, ladj = prev_ladj)
end

function apply_dist_trafo(trg_d::Distribution{Univariate}, src_d::StdMvDist, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    @_adignore @argcheck length(src_d) == length(eachindex(src_v)) == 1
    apply_dist_trafo(trg_d, view(src_d, 1), first(src_v), prev_ladj)
end

function apply_dist_trafo(trg_d::StdMvDist, src_d::Distribution{Univariate}, src_v::Real, prev_ladj::OptionalLADJ)
    @argcheck length(trg_d) == 1
    r = apply_dist_trafo(view(trg_d, 1), src_d, first(src_v), prev_ladj)
    (v = unshaped(r.v), ladj = r.ladj)
end

@inline _trafo_cdf(d::Distribution{Univariate,Continuous}, x::Real) = _trafo_cdf_impl(params(d), d, x)

@inline _trafo_cdf_impl(::NTuple, d::Distribution{Univariate,Continuous}, x::Real) = cdf(d, x)

@inline function _trafo_cdf_impl(::NTuple{N,Union{Integer,AbstractFloat}}, d::Distribution{Univariate,Continuous}, x::ForwardDiff.Dual{TAG}) where {N,TAG}
    x_v = ForwardDiff.value(x)
    u = cdf(d, x_v)
    dudx = pdf(d, x_v)
    ForwardDiff.Dual{TAG}(u, dudx * ForwardDiff.partials(x))
end

@inline _trafo_quantile(d::Distribution{Univariate,Continuous}, u::Real) = _trafo_quantile_impl(params(d), d, u)

@inline _trafo_quantile_impl(::NTuple, d::Distribution{Univariate,Continuous}, u::Real) = _trafo_quantile_impl_generic(d, u)

@inline function _trafo_quantile_impl(::NTuple{N,Union{Integer,AbstractFloat}}, d::Distribution{Univariate,Continuous}, u::ForwardDiff.Dual{TAG}) where {N,TAG}
    x = _trafo_quantile_impl_generic(d, ForwardDiff.value(u))
    dxdu = inv(pdf(d, x))
    ForwardDiff.Dual{TAG}(x, dxdu * ForwardDiff.partials(u))
end

# Workaround for Beta dist, ForwardDiff doesn't work for parameters:
@inline _trafo_quantile_impl(::NTuple{N,ForwardDiff.Dual}, d::Beta, u::Real) where N = convert(float(typeof(u)), NaN)
# Workaround for Beta dist, current quantile implementation only supports Float64:
@inline _trafo_quantile_impl(::NTuple{N,Real}, d::Beta, u::Float32) where N = quantile(d, convert(promote_type(Float64, typeof(u)), u))

@inline _trafo_quantile_impl_generic(d::Distribution{Univariate,Continuous}, u::Real) = quantile(d, u)

# Workaround for rounding errors that can result in quantile values outside of support of Truncated:
@inline function _trafo_quantile_impl_generic(d::Truncated{<:Distribution{Univariate,Continuous}}, u::Real)
    x = quantile(d, u)
    T = typeof(x)
    min_x = T(minimum(d))
    max_x = T(maximum(d))
    if x < min_x && isapprox(x, min_x, atol = 4 * eps(T))
        min_x
    elseif x > max_x && isapprox(x, max_x, atol = 4 * eps(T))
        max_x
    else
        x
    end
end

@inline function _eval_dist_trafo_func(f::typeof(_trafo_cdf), d::Distribution{Univariate,Continuous}, src_v::Real, prev_ladj::OptionalLADJ)
    R_V = float(promote_type(typeof(src_v), eltype(params(d)), ismissing(prev_ladj) ? Bool : typeof(prev_ladj)))
    R_LADJ = !ismissing(prev_ladj) ? R_V : Missing
    if insupport(d, src_v)
        trg_v = f(d, src_v)
        trafo_ladj = !ismissing(prev_ladj) ? + logpdf(d, src_v) : missing
        var_trafo_result(convert(R_V, trg_v), src_v, convert(R_LADJ, trafo_ladj), prev_ladj)
    else
        var_trafo_result(convert(R_V, NaN), src_v, convert(R_LADJ, !ismissing(prev_ladj) ? NaN : missing), prev_ladj)
    end
end

@inline function _eval_dist_trafo_func(f::typeof(_trafo_quantile), d::Distribution{Univariate,Continuous}, src_v::Real, prev_ladj::OptionalLADJ)
    R_V = float(promote_type(typeof(src_v), eltype(params(d)), ismissing(prev_ladj) ? Bool : typeof(prev_ladj)))
    R_LADJ = !ismissing(prev_ladj) ? R_V : Missing
    if 0 <= src_v <= 1
        trg_v = f(d, src_v)
        trafo_ladj = !ismissing(prev_ladj) ? - logpdf(d, trg_v) : missing
        var_trafo_result(convert(R_V, trg_v), src_v, convert(R_LADJ, trafo_ladj), prev_ladj)
    else
        var_trafo_result(convert(R_V, NaN), src_v, convert(R_LADJ, !ismissing(prev_ladj) ? NaN : missing), prev_ladj)
    end
end

std_dist_from(src_d::Distribution{Univariate,Continuous}) = StandardUvUniform()

function apply_dist_trafo(::StandardUvUniform, src_d::Distribution{Univariate,Continuous}, src_v::Real, prev_ladj::OptionalLADJ)
    _eval_dist_trafo_func(_trafo_cdf, src_d, src_v, prev_ladj)
end

std_dist_to(trg_d::Distribution{Univariate,Continuous}) = StandardUvUniform()

function apply_dist_trafo(trg_d::Distribution{Univariate,Continuous}, ::StandardUvUniform, src_v::Real, prev_ladj::OptionalLADJ)
    TV = float(typeof(src_v))
    # Avoid src_v ≈ 0 and src_v ≈ 1 to avoid infinite variate values for target distributions with infinite support:
    mod_src_v = ifelse(src_v == 0, zero(TV) + eps(TV), ifelse(src_v == 1, one(TV) - eps(TV), convert(TV, src_v)))
    trg_v, ladj = _eval_dist_trafo_func(_trafo_quantile, trg_d, mod_src_v, prev_ladj)
    (v = trg_v, ladj = ladj)
end

function _dist_trafo_rescale_impl(trg_d, src_d, src_v::Real, prev_ladj::OptionalLADJ)
    R = float(typeof(src_v))
    trg_offs, trg_scale = location(trg_d), scale(trg_d)
    src_offs, src_scale = location(src_d), scale(src_d)
    rescale_factor = trg_scale / src_scale
    trg_v = (src_v - src_offs) * rescale_factor + trg_offs
    trafo_ladj = !ismissing(prev_ladj) ? log(rescale_factor) : missing
    var_trafo_result(trg_v, src_v, trafo_ladj, prev_ladj)
end

@inline apply_dist_trafo(trg_d::Uniform, src_d::Uniform, src_v::Real, prev_ladj::OptionalLADJ) = _dist_trafo_rescale_impl(trg_d, src_d, src_v, prev_ladj)
@inline apply_dist_trafo(trg_d::StandardUvUniform, src_d::Uniform, src_v::Real, prev_ladj::OptionalLADJ) = _dist_trafo_rescale_impl(trg_d, src_d, src_v, prev_ladj)
@inline apply_dist_trafo(trg_d::Uniform, src_d::StandardUvUniform, src_v::Real, prev_ladj::OptionalLADJ) = _dist_trafo_rescale_impl(trg_d, src_d, src_v, prev_ladj)

# ToDo: Use StandardUvNormal as standard intermediate dist for Normal? Would
# be useful if StandardUvNormal would be a better standard intermediate than
# StandardUvUniform for some other uniform distributions as well.
#
#     std_dist_from(src_d::Normal) = StandardUvNormal()
#     std_dist_to(trg_d::Normal) = StandardUvNormal()

@inline apply_dist_trafo(trg_d::Normal, src_d::Normal, src_v::Real, prev_ladj::OptionalLADJ) = _dist_trafo_rescale_impl(trg_d, src_d, src_v, prev_ladj)
@inline apply_dist_trafo(trg_d::StandardUvNormal, src_d::Normal, src_v::Real, prev_ladj::OptionalLADJ) = _dist_trafo_rescale_impl(trg_d, src_d, src_v, prev_ladj)
@inline apply_dist_trafo(trg_d::Normal, src_d::StandardUvNormal, src_v::Real, prev_ladj::OptionalLADJ) = _dist_trafo_rescale_impl(trg_d, src_d, src_v, prev_ladj)

# ToDo: Optimized implementation for Distributions.Truncated <-> StandardUvUniform

@inline function apply_dist_trafo(trg_d::StandardMvUniform, src_d::StandardMvNormal, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(StandardUvUniform(), StandardUvNormal(), src_v, prev_ladj)
end

@inline function apply_dist_trafo(trg_d::StandardMvNormal, src_d::StandardMvUniform, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(StandardUvNormal(), StandardUvUniform(), src_v, prev_ladj)
end

std_dist_from(src_d::MvNormal) = StandardMvNormal(length(src_d))

function apply_dist_trafo(trg_d::StandardMvNormal, src_d::MvNormal, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    @argcheck length(trg_d) == length(src_d)
    A = cholesky(src_d.Σ).U
    trg_v = transpose(A) \ (src_v - src_d.μ)
    trafo_ladj = !ismissing(prev_ladj) ? - logabsdet(A)[1] : missing
    var_trafo_result(trg_v, src_v, trafo_ladj, prev_ladj)
end

std_dist_to(trg_d::MvNormal) = StandardMvNormal(length(trg_d))

function apply_dist_trafo(trg_d::MvNormal, src_d::StandardMvNormal, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    @argcheck length(trg_d) == length(src_d)
    A = cholesky(trg_d.Σ).U
    trg_v = transpose(A) * src_v + trg_d.μ
    trafo_ladj = !ismissing(prev_ladj) ? + logabsdet(A)[1] : missing
    var_trafo_result(trg_v, src_v, trafo_ladj, prev_ladj)
end

eff_totalndof(d::Dirichlet) = length(d) - 1
eff_totalndof(d::DistributionsAD.TuringDirichlet) = length(d) - 1

std_dist_to(trg_d::Dirichlet) = StandardMvUniform(eff_totalndof(trg_d))
std_dist_to(trg_d::DistributionsAD.TuringDirichlet) = StandardMvUniform(eff_totalndof(trg_d))

function apply_dist_trafo(trg_d::Dirichlet, src_d::StandardMvUniform, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    apply_dist_trafo(DistributionsAD.TuringDirichlet(trg_d.alpha), src_d, src_v, prev_ladj)
end

function _dirichlet_beta_trafo(α::Real, β::Real, src_v::Real)
    R = float(promote_type(typeof(α), typeof(β), typeof(src_v)))
    convert(R, apply_dist_trafo(Beta(α, β), StandardUvUniform(), src_v, missing).v)::R
end

_a_times_one_minus_b(a::Real, b::Real) = a * (1 - b)

function apply_dist_trafo(trg_d::DistributionsAD.TuringDirichlet, src_d::StandardMvUniform, src_v::AbstractVector{<:Real}, prev_ladj::Missing)
    # See M. J. Betancourt, "Cruising The Simplex: Hamiltonian Monte Carlo and the Dirichlet Distribution",
    # https://arxiv.org/abs/1010.3436

    @_adignore @argcheck length(trg_d) == length(src_d) + 1
    αs = _dropfront(_rev_cumsum(trg_d.alpha))
    βs = _dropback(trg_d.alpha)
    beta_v = fwddiff(_dirichlet_beta_trafo).(αs, βs, src_v)
    beta_v_cp = _exp_cumsum_log(_pushfront(beta_v, 1))
    beta_v_ext = _pushback(beta_v, 0)
    trg_v = fwddiff(_a_times_one_minus_b).(beta_v_cp, beta_v_ext)
    var_trafo_result(trg_v, src_v, missing, prev_ladj)
end

function apply_dist_trafo(trg_d::DistributionsAD.TuringDirichlet, src_d::StandardMvUniform, src_v::AbstractVector{<:Real}, prev_ladj::Real)
    trg_v = apply_dist_trafo(trg_d, src_d, src_v, missing).v
    trafo_ladj = - logpdf(trg_d, trg_v)
    var_trafo_result(trg_v, src_v, trafo_ladj, prev_ladj)
end

g_debug = nothing
function _product_dist_trafo_impl(trg_ds, src_ds, src_v::AbstractVector{<:Real}, prev_ladj::Missing)
    global g_debug = (trg_ds=trg_ds, src_ds=src_ds, src_v=src_v, prev_ladj=missing)
    trg_v = fwddiff(apply_dist_trafo_noladj).(trg_ds, src_ds, src_v)
    var_trafo_result(trg_v, src_v, missing, missing)
end

function _product_dist_trafo_impl(trg_ds, src_ds, src_v::AbstractVector{<:Real}, prev_ladj::Real)
    rs = fwddiff(apply_dist_trafo).(trg_ds, src_ds, src_v, zero(prev_ladj))
    trg_v = map(r -> r.v, rs)
    trafo_ladj = sum(map(r -> r.ladj, rs))
    var_trafo_result(trg_v, src_v, trafo_ladj, prev_ladj)
end

function apply_dist_trafo(trg_d::Distributions.Product, src_d::Distributions.Product, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(trg_d.v, src_d.v, src_v, prev_ladj)
end

function apply_dist_trafo(trg_d::StandardMvUniform, src_d::Distributions.Product, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(StandardUvUniform(), src_d.v, src_v, prev_ladj)
end

function apply_dist_trafo(trg_d::StandardMvNormal, src_d::Distributions.Product, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(StandardUvNormal(), src_d.v, src_v, prev_ladj)
end

function apply_dist_trafo(trg_d::Distributions.Product, src_d::StandardMvUniform, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(trg_d.v, StandardUvUniform(), src_v, prev_ladj)
end

function apply_dist_trafo(trg_d::Distributions.Product, src_d::StandardMvNormal, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    @_adignore @argcheck eff_totalndof(trg_d) == eff_totalndof(src_d)
    _product_dist_trafo_impl(trg_d.v, StandardUvNormal(), src_v, prev_ladj)
end

function _ntdistelem_to_stdmv(trg_d::StdMvDist, sd::Distribution, src_v_unshaped::AbstractVector{<:Real}, trg_acc::ValueAccessor, init_ladj::OptionalLADJ)
    td = view(trg_d, ValueShapes.view_range(Base.OneTo(length(trg_d)), trg_acc))
    sv = stripscalar(view(src_v_unshaped, trg_acc))
    apply_dist_trafo(td, sd, sv, init_ladj)
end

function _ntdistelem_to_stdmv(trg_d::StdMvDist, sd::ConstValueDist, src_v_unshaped::AbstractVector{<:Real}, trg_acc::ValueAccessor, init_ladj::OptionalLADJ)
    (v = Bool[], ladj = init_ladj)
end

_transformed_ntd_elshape(d::Distribution{Univariate}) = varshape(d)
_transformed_ntd_elshape(d::Distribution{Multivariate}) = ArrayShape{Real}(eff_totalndof(d))
function _transformed_ntd_elshape(d::Distribution)
    vs = varshape(d)
    @argcheck totalndof(vs) == eff_totalndof(d)
    vs
end

function _transformed_ntd_accessors(d::NamedTupleDist{names}) where names
    shapes = map(_transformed_ntd_elshape, values(d))
    vs = NamedTupleShape(NamedTuple{names}(shapes))
    values(vs)
end

function apply_dist_trafo(trg_d::StdMvDist, src_d::ValueShapes.UnshapedNTD, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    src_vs = varshape(src_d.shaped)
    @argcheck length(src_d) == length(eachindex(src_v))
    trg_accessors = _transformed_ntd_accessors(src_d.shaped)
    init_ladj = ismissing(prev_ladj) ? missing : zero(Float32)
    rs = map((acc, sd) -> _ntdistelem_to_stdmv(trg_d, sd, src_v, acc, init_ladj), trg_accessors, values(src_d.shaped))
    trg_v = vcat(map(r -> r.v, rs)...)
    trafo_ladj = !ismissing(prev_ladj) ? sum(map(r -> r.ladj, rs)) : missing
    var_trafo_result(trg_v, src_v, trafo_ladj, prev_ladj)
end

function apply_dist_trafo(trg_d::StdMvDist, src_d::NamedTupleDist, src_v::Union{NamedTuple,ShapedAsNT}, prev_ladj::OptionalLADJ)
    src_v_unshaped = unshaped(src_v, varshape(src_d))
    apply_dist_trafo(trg_d, unshaped(src_d), src_v_unshaped, prev_ladj)
end

function _stdmv_to_ntdistelem(td::Distribution, src_d::StdMvDist, src_v::AbstractVector{<:Real}, src_acc::ValueAccessor, init_ladj::OptionalLADJ)
    sd = view(src_d, ValueShapes.view_range(Base.OneTo(length(src_d)), src_acc))
    sv = view(src_v, ValueShapes.view_range(axes(src_v, 1), src_acc))
    apply_dist_trafo(td, sd, sv, init_ladj)
end

function _stdmv_to_ntdistelem(td::ConstValueDist, src_d::StdMvDist, src_v::AbstractVector{<:Real}, src_acc::ValueAccessor, init_ladj::OptionalLADJ)
    (v = Bool[], ladj = init_ladj)
end

function apply_dist_trafo(trg_d::ValueShapes.UnshapedNTD, src_d::StdMvDist, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    trg_vs = varshape(trg_d.shaped)
    @argcheck length(src_d) == length(eachindex(src_v))
    src_accessors = _transformed_ntd_accessors(trg_d.shaped)
    init_ladj = ismissing(prev_ladj) ? missing : zero(Float32)
    rs = map((acc, td) -> _stdmv_to_ntdistelem(td, src_d, src_v, acc, init_ladj), src_accessors, values(trg_d.shaped))
    trg_v_unshaped = vcat(map(r -> unshaped(r.v), rs)...)
    trafo_ladj = !ismissing(prev_ladj) ? sum(map(r -> r.ladj, rs)) : missing
    var_trafo_result(trg_v_unshaped, src_v, trafo_ladj, prev_ladj)
end

function apply_dist_trafo(trg_d::NamedTupleDist, src_d::StdMvDist, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    unshaped_result = apply_dist_trafo(unshaped(trg_d), src_d, src_v, prev_ladj)
    (v = varshape(trg_d)(unshaped_result.v), ladj = unshaped_result.ladj)
end

const AnyReshapedDist = Union{ReshapedDist,MatrixReshaped}

function apply_dist_trafo(trg_d::Distribution{Multivariate}, src_d::AnyReshapedDist, src_v::Any, prev_ladj::OptionalLADJ)
    src_vs = varshape(src_d)
    @argcheck length(trg_d) == totalndof(src_vs)
    apply_dist_trafo(trg_d, unshaped(src_d), unshaped(src_v, src_vs), prev_ladj)
end

function apply_dist_trafo(trg_d::AnyReshapedDist, src_d::Distribution{Multivariate}, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    trg_vs = varshape(trg_d)
    @argcheck totalndof(trg_vs) == length(src_d)
    r = apply_dist_trafo(unshaped(trg_d), src_d, src_v, prev_ladj)
    (v = trg_vs(r.v), ladj = r.ladj)
end

function apply_dist_trafo(trg_d::AnyReshapedDist, src_d::AnyReshapedDist, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    trg_vs = varshape(trg_d)
    src_vs = varshape(src_d)
    @argcheck totalndof(trg_vs) == totalndof(src_vs)
    r = apply_dist_trafo(unshaped(trg_d), unshaped(src_d), unshaped(src_v, src_vs), prev_ladj)
    (v = trg_vs(r.v), ladj = r.ladj)
end

function apply_dist_trafo(trg_d::StdMvDist, src_d::UnshapedHDist, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    src_v_primary, src_v_secondary = _hd_split(src_d, src_v)
    trg_d_primary = typeof(trg_d)(length(eachindex(src_v_primary)))
    trg_d_secondary = typeof(trg_d)(length(eachindex(src_v_secondary)))
    trg_v_primary, ladj_primary = apply_dist_trafo(trg_d_primary, _hd_pridist(src_d), src_v_primary, prev_ladj)
    trg_v_secondary, ladj = apply_dist_trafo(trg_d_secondary, _hd_secdist(src_d, src_v_primary), src_v_secondary, ladj_primary)
    trg_v = vcat(trg_v_primary, trg_v_secondary)
    (v = trg_v, ladj = ladj)
end

function apply_dist_trafo(trg_d::StdMvDist, src_d::HierarchicalDistribution, src_v::Any, prev_ladj::OptionalLADJ)
    src_v_unshaped = unshaped(src_v, varshape(src_d))
    apply_dist_trafo(trg_d, unshaped(src_d), src_v_unshaped, prev_ladj)
end

function apply_dist_trafo(trg_d::UnshapedHDist, src_d::StdMvDist, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    src_v_primary, src_v_secondary = _hd_split(trg_d, src_v)
    src_d_primary = typeof(src_d)(length(eachindex(src_v_primary)))
    src_d_secondary = typeof(src_d)(length(eachindex(src_v_secondary)))
    trg_v_primary, ladj_primary = apply_dist_trafo(_hd_pridist(trg_d), src_d_primary, src_v_primary, prev_ladj)
    trg_v_secondary, ladj = apply_dist_trafo(_hd_secdist(trg_d, trg_v_primary), src_d_secondary, src_v_secondary, ladj_primary)
    trg_v = vcat(trg_v_primary, trg_v_secondary)
    (v = trg_v, ladj = ladj)
end

function apply_dist_trafo(trg_d::HierarchicalDistribution, src_d::StdMvDist, src_v::AbstractVector{<:Real}, prev_ladj::OptionalLADJ)
    unshaped_result = apply_dist_trafo(unshaped(trg_d), src_d, src_v, prev_ladj)
    (v = varshape(trg_d)(unshaped_result.v), ladj = unshaped_result.ladj)
end


text:
"""
Needs to be overloaded.
"""
function update(model::Model, s::T) where {T} end

"""
Needs to be overloaded.
"""
function make_init_state(model::Model) end
make_init_state(gibbs::Gibbs) = make_init_state(gibbs.model)

"""
Subset a named tuple
"""
function subsetnamedtuple(state::NamedTuple, names)
  vs = map(v -> state[v], names)
  return NamedTuple{names}(vs)
end

"""
    mcmc(model::Model, nsamps::Int; init::Union{NamedTuple, Nothing}, nburn::Int=0,
         thin::Int=1, exclude::Vector{Symbol}=[], callback=nothing,
         return_last_state=false)

### Arguments

`model`: is a subtype of MCMC.Model.

`nsamps`: Number of samples to return.

`init`: is the initial state.

`nburn`: Number of burn-in samples to discard at the beginning.

`thin`: Factor by which to thin samples.

`exclude`: Parameters to exclude from final chain.

`callback`: Callback function of the form `callback(chain, state, sample, i, metrics, iterator)`.

**Return**: `(chain::Vector{<:NamedTuple}, metrics::Dict{Symbol, Any})` where `chain` is a Vector of `NamedTuples` of length `nsamps` and `metrics` is additional metrics that would be generated through `callback` if provided.
"""
function mcmc(model::Union{Model,Gibbs}, nsamps::Int; init::Union{Nothing, NamedTuple}=nothing,
              nburn::Int=0, thin::Int=1, exclude::Vector{Symbol}=Symbol[],
              callback=nothing, progress=true, return_last_state=false,
              kwargs...)
  # Create intial state if needed.
  init === nothing && (init = make_init_state(model))

  # Assert that arguments are correctly specified.
  thin >= 1 || error("`thin` should be ≥ 1")
  nburn >= 0 || error("`thin` should be ≥ 0")

  # Total number of iterations.
  niters = nsamps * thin + nburn

  # Index counter
  idx = 0

  # All parameters
  allparams = keys(init)

  # Tracked parameters
  tracked_params = Tuple(filter(s -> !(s in exclude), allparams))

  # Current state
  state = deepcopy(init)

  # Preallocate output.
  chain = [deepcopy(subsetnamedtuple(state, tracked_params)) for _ in 1:nsamps]

  # Additional metrics.
  metrics = Dict{Symbol, Any}()

  iterator = progress ? ProgressBar(1:niters) : 1:niters
  for i in iterator
    # Update current state.
    # NOTE: User needs to imlement MCMC.update(s::typeof(s))
    state = update(model, state; kwargs...)

    # Trim state to get a sample to save.
    sample = subsetnamedtuple(state, tracked_params)

    # Callback function.
    callback === nothing || callback(chain, state, sample, i, metrics, iterator)

    # Save current state.
    (i > nburn) && ((i - nburn) % thin == 0) && setindex!!(chain, sample, idx += 1)
  end

  if return_last_state
    return (chain=chain, metrics=metrics, last_state=state)
  else
    return (chain=chain, metrics=metrics)
  end
end


text:
# missing in FinFld
# base_ring consistently
# rand
# creation with string, symbol, neither
# minpoly, charpoly
# order = size?
# solve behave consistently for rings and fields
#         do pseudo solve for extra den
#       if ring is known to be field, use field version
#       solve for matrix of NOT full rank
#       kernel
#       version that returns false rather than an error
#         alterntively: throw a very precise error object that
#         can be caught
# residueRing creation: rename generator to prefix with r or so
#
# Padic -> pAdic
# gen for ResidueRing(Poly)

#################################################

# Plesken: a ordering on finite field elements and polynomials
# from his lecture notes.

function steinitz(a::nmod_poly)
  p = characteristic(base_ring(a))
  ZZx = PolynomialRing(FlintZZ)[1]
  #  f = lift(ZZx, a)  ## bloody stupid lift for poly uses symmetric residue
  f = [lift(coeff(a, i))::fmpz for i=0:degree(a)]
  return Nemo.evaluate(ZZx(f), p)
end

function steinitz(a::ResElem{fmpz})
  return lift(a)
end

function steinitz(a::ResElem{T}) where T <: Union{nmod_poly, fq_nmod_poly, PolyElem}
  f = [steinitz(coeff(a.data, i))::fmpz for i=0:degree(a.data)]
  ZZx = PolynomialRing(FlintZZ)[1]
  S = base_ring(base_ring(parent(a)))
  return evaluate(ZZx(f), size(S))
end

function steinitz(a::fq)
  st = 0
  p = characteristic(parent(a))
  for i=degree(parent(a))-1:-1:0
    st *=p
    st += coeff(a, i)
  end
  return st
end

function steinitz(a::fq_nmod)
  st = 0
  p = characteristic(parent(a))
  for i=degree(parent(a))-1:-1:0
    st *=p
    st += coeff(a, i)
  end
  return st
end

##############################################################
# this is expensive, but completely generic
# possibly improve by using the fact that aut should be an automorphism
# if the order of aut would be known, one could use this to proceed in layers
function minpoly_aut(a::ResElem{T}, aut :: Function) where T <: Union{fq_nmod_poly, nmod_poly}
  R = parent(a)
  RX, X = PolynomialRing(R)
  o = Set{typeof(X)}()
  push!(o, X-a)
  a = aut(a)
  while !(X-a in o)
    push!(o, X-a)
    a = aut(a)
  end
  f = prod(collect(o))
  # now, the coefficients SHOULD be in a smaller ring, but we
  # don't have that ring
  return f
end

function minpoly_aut(a::ResElem{T}, aut :: Function) where T <: PolyElem
  R = parent(a)
  RX, X = PolynomialRing(R)
  o = Set{typeof(X)}()
  push!(o, X-a)
  a = aut(a)
  while !(X-a in o)
    push!(o, X-a)
    a = aut(a)
  end
  f = prod(collect(o))
  return f
end

function minpoly_pow(a::ResElem{T}, deg::Int) where T <: Union{PolyElem, fq_nmod_poly}
  R = parent(a)
  S = base_ring(base_ring(R))
  M = MatrixSpace(S, deg, degree(R.modulus))()
  B = MatrixSpace(S, 1, degree(R.modulus))()
  elem_to_mat_row!(M, 1, one(R))
  b = a
  for i=1:deg-1
    elem_to_mat_row!(M, i+1, b)
    b *= a
  end
  elem_to_mat_row!(B, 1, -b)
  s = solve_rational(M', B')
  if isa(s, Tuple)
    s = s[1] * inv(s[2])
  end
  ## just to keep the interface...
  Rx,x = PolynomialRing(R)
  arr = Vector{typeof(b)}(deg+1)
  for i=1:deg
    arr[i] = R(s[i, 1])  ## wasteful
  end
  arr[deg+1] = one(R)
  return Rx(arr)
end

###################################################################
###################################################################
###################################################################
##
##
##
##
## Plesken starting
##

function primitive_root_r_div_qm1(R, r::Int)
  #Plesken, 1.39 essentially
  n = size(R)-1
  k, e = remove(n, r)
  @assert k>0
  @assert isprime(r)

  a = rand(R)^e
  while a==0 || a^(r^(k-1)) == 1
    a = rand(R)^e
  end
  #println("found one $r^$k-th root of one: $a")

  #level 1: need a^(r^(k-1))^l minimal, 0<l<r (so a stays primitive)
  b = a^(r^(k-1))
  c = b
  st = steinitz(b)
  opt = 1
  for i=2:r-1
    c *= b
    if steinitz(c) < st
      opt = i
      st = steinitz(c)
    end
  end
  a = a^opt
  #println("adjusted for level 1: ", a, " st: $st")
  for i=2:k
    #println("dealing with level $i")
    ## (a*a^(xr^(i-1)))^r^(k-i) needs to be minimal in the Steinitz sense
    ## a^(r^(k-i) + xr^k-1)
    # use a*a^(xr^(i-1)) as the new root
    b = a^(r^(k-1)) ## basically a r-th root
    c = a^(r^(k-i))
    opt = 0
    st = steinitz(c)
    for ii=1:r
      c *= b
      if steinitz(c) < st
        opt = ii
        st = steinitz(c)
      end
    end
    #println("adjusted for level $i: ", a, " st: $st")
    a = a*a^(opt*r^(i-1))
  end
  return a
end

function get_f(r::Int, p::fmpz, s::Int)
  R = PadicField(r, s)
  return lift(teichmuller(R(p)))
end
# plan
# to get a field p^n, decompose n into prime powers and use compositum (missing)
# to get p^(r^l)
# find minimal k sth. r| p^k-1 and define F_(p^k)
#  in F_p^k find the proper root of 1 (using the primitive_root_r_div_qm1
#    function), cal this mu
# build x^r - mu and iterate l times
# descent: find f using get_f and form the trace relative to f
# to find generators for the degree r extension of F_p, not F_p^k
# needs minpoly in some variant...

function f_tr(a::ResElem, f::fmpz, o::Int)
  s = a
  for i=1:o-1
    a = a^f
    s += a
  end
  return s
end

function plesken_kummer(p::fmpz, r::Int, s::Int)
  #find Plesken rep for F_p^(r^s)
  @assert r%2 == 1
  @assert r!=p
  @assert isprime(r) && isprime(p)

  if (p-1) % r == 0
    R = FiniteField(p)
    descent = false
  else
    f = cyclotomic(r, PolynomialRing(FlintZZ)[2])
    f = PolynomialRing(ResidueRing(FlintZZ, p))[1](f)
    f = factor(f)
    k = keys(f.fac)
    st = start(k)
    f1, st = next(k, st)
    st = steinitz(f1)
    opt = f1
    while !done(k, st)
      fi, st = next(k, st)
      if steinitz(fi) < st
        st = steinitz(fi)
        opt = fi
      end
    end
    descent = true
    ord = degree(opt)
    R = FlintFiniteField(opt, "a")[1]
    T = ResidueRing(FlintZZ, p)
    J = CoerceMap(T, R)
  end
  zeta = primitive_root_r_div_qm1(R, r)

  a = zeta

  if descent
    f = get_f(r, p, s+4)  #+4 is stupid. Need to think
    #println("using $f as f")
    # a -> a^f should be an automorphism of order ord = degree of ext
  end

  S = 1
  U = 1
  for i=1:s ## maybe do one step x^(r^s)-a only?
    #println("doin' stuff")
    Rx, x = PolynomialRing(R, "x_$i")
    S = ResidueRing(Rx, x^r-a)
    I = CoerceMap(R, S)
    a = S(x)
    if descent
      b = f_tr(a, f, ord)
      #println("$i: trace of $a is $b")
#      pol = minpoly_aut(b, x->x^(p^(r^(i-1))))
      pol = minpoly_pow(b, r)  ## does not work: expo too large
      #println(pol)
      arr = Array{typeof(zero(T))}(degree(pol)+1)
      @assert domain(I) == codomain(J)
      @assert parent(coeff(pol, 0)) == codomain(I)
      for j=0:degree(pol)
        arr[j+1] = preimage(J, preimage(I, coeff(pol, j)))
      end
      pol = PolynomialRing(T, "t_$i")[1](arr)
      U = ResidueRing(parent(pol), pol)
      H = ResidueRingPolyMap(U, S, b, J)
      #H.coeff_map = J
      J = H
      T = U
    end
    R = S
  end

  return U, S
end

function plesken_as(p::fmpz, r::Int, s::Int)
  @assert p==r
  R = FiniteField(p)
  g = R(-1)
  t = 1
  while s>1
    Rx,x = PolynomialRing(R, "t_$i")
    R = ResidueRing(Rx, x^r-x-g^(r-1)) ## r==p, but of better type
    g = gen(R)
    s -= 1
    t += 1
  end
  return R
end

function plesken_2(p::fmpz, r::Int, s::Int)
  @assert r==2
  #Plesken, 1.27
  if valuation(p-1, 2) >1
    R = FiniteField(p)
    g = primitive_root_r_div_qm1(p, r)
    t = 1
  else
    @assert valuation(p+1, 2)>1
    R = FiniteField(p)
    Rx,x = PolynomialRing(R, "t_1")
    R = ResidueRing(Rx, x^2+1)
    g = primitive_root_r_div_qm1(R, 2)
    s -= 1
    t = 2
  end
  while s>0
    Rx,x = PolynomialRing(R, "t_$t")
    R = ResidueRing(Rx, x^2-g)
    g = gen(R)
    t += 1
  end
  return R
end

## missing:
# the composition for general r (proper order and such)
# choice or sum or product to generate field uniquely
#
# efficiency: those towers are slow (improved by different minpoly, Plesken
#               suggests a recursion to writem them down, based on the matrix)
#   one should probably collapse them as much as possibly
#   the rep is still unique (in prime power towers) as the
#   fields are cyclic: there are now subfields that will
#   combine to everything
#
#   direct version: the final elt in them kummer case is x^(r^s)-a
#   we can descent on this directly
#   same with 2
#   as is more involved (->Witt ring construction)
#
# maps & coercion
#
# the Log = canonical generator

function h_minus(p::Int, nb::Int)
  #void
  #fmpz_poly_resultant_modular_div(fmpz_t res, const fmpz_poly_t poly1,
  #              const fmpz_poly_t poly2, const fmpz_t divisor, slong nbits)

  #let g be a generator for Z/pZ^*, h = g^-1 mod p
  # F = sum_0^p-2 h^i x^i in Z[x]
  # h^- = Res(F, x^((p-1)/2)+1)/(2p)^((p-3)/2)
  #
  # log(2*p*(p/4/pi^2)*(p-1)/4/log(2)
  #
  # is the asymptotic size, so that's what nb should be

  Zx, x = PolynomialRing(FlintZZ)
  F = Vector{fmpz}(p-1)

  g = rand(1:p-1)
  while modord(g, p) != p-1
    g = rand(1:p-1)
  end
  h = invmod(g, p)
  F[1] = 1
  for i=2:p-1
    F[i] = F[i-1]*h % p
  end
  F = Zx(F)

  f = (2*fmpz(p))^div(p-3, 2)
  @time r1 = resultant(F, x^div(p-1, 2)+1, f, nb)
  @time r2 = resultant(F, x^div(p-1, 2)+1)
  return r1
end


text:
@model Qlearning(na, ns) = begin
    α ~ Beta()
    β ~ Gamma(2,1)
    return (α=α, β=β, Values = fill(1/na,na,ns))
end

mdl1 = Qlearning(2,1)

@evolution mdl1 begin
    delta_rule!(Values, s, a, r, α)
end

@observation mdl1 begin
    V = β .* @views(Values[:,s])
    Categorical(softmax!(V))
end

θ = generated_quantities(mdl1, (α=0.2, β = 2.0))
@test θ == (α = 0.2, β = 2.0, Values = fill(0.5, 2, 1))

AnimalBehavior.evol!(mdl1, 1, 1, 1.0; θ... )
@test θ.Values[1] == 0.6

@test AnimalBehavior.observ(mdl1, 1; θ... ) == Categorical(softmax!([1.2, 1.0]))

max_stars_repo_name: sqwayer/AnimalBehavior

max_issues_repo_name: sqwayer/AnimalBehavior


text:
ClosedPrimitive(p::VP) where {VP <:AbstractVolumePrimitive} = VP(p, COT = ClosedPrimitive)
OpenPrimitive(p::VP) where {VP <: AbstractVolumePrimitive} = VP(p, COT = OpenPrimitive)
switchClosedOpen(p::AbstractVolumePrimitive{<:Any,ClosedPrimitive}) = OpenPrimitive(p)
switchClosedOpen(p::AbstractVolumePrimitive{<:Any,OpenPrimitive}) = ClosedPrimitive(p)

isClosedPrimitive(vp::AbstractVolumePrimitive{<:Any,ClosedPrimitive}) = true
isClosedPrimitive(vp::AbstractVolumePrimitive{<:Any,OpenPrimitive}) = false

distance(pt::CartesianPoint, vp::AbstractVolumePrimitive) =
minimum(map(p -> distance(pt, p), surfaces(vp)))

function sample(vp::AbstractVolumePrimitive)
    surfs = surfaces(vp)
    ls = unique!(collect(Base.Iterators.flatten(lines.(surfs))))
    unique!(mapreduce(sample, vcat, ls))
end

include("Box.jl")
include("Cone.jl")
include("Ellipsoid.jl")
include("RegularPrism.jl")
include("Torus.jl")


text:
function converter(side)
    return HorizonSide((Int(side)+2)%4)
end

function gowhile(r,s)
    a=0
    b=0
    if ismarker(r)
        a+=1
        b+=temperature(r)
    end
    while !isborder(r,s)
        move!(r,s)
        if ismarker(r)
            a+=1
            b+=temperature(r)
        end
    end
    return (a,b)
end

function task_10(r)
    marks=0
    sum=0
    side=Ost
    while true
        mm,ss=gowhile(r, side)
        marks+=mm
        sum+=ss
        if !isborder(r,Nord)
            move!(r,Nord)
            side=converter(side)
        else
            break
        end
    end
    print(sum/marks)
end

task_10(r)


text:
function besselroots(nu::Float64, n::Integer)
#BESSELROOTS    The first N roots of the function J_v(x)

# DEVELOPERS NOTES:
#   V = 0 --> Full double precision for N <= 20 (Wolfram Alpha), and very
#     accurate approximations for N > 20 (McMahon's expansion)
#   -1 <= V <= 5 : V ~= 0 --> 12 decimal figures for the 6 first zeros
#     (Piessens's Chebyshev series approximations), and very accurate
#     approximations for the others (McMahon's expansion)
#   V > 5 --> moderately accurate for the 6 first zeros and good
#     approximations for the others (McMahon's expansion)

# This code was originally written by L. L. Peixoto in MATLAB.
# Later modified by A. Townsend to work in Julia

    if n < 0
        error("Input N must be a positive integer")
    end

    x = zeros(n)
    if n > 0 && nu == 0
        for k in 1:min(n,20)
            x[k] = J0_roots[k]
        end
        for k in min(n,20)+1:n
            x[k] = McMahon(nu, k)
        end
    elseif n>0 && nu >= -1 && nu <= 5
        correctFirstFew = Piessens( nu )
        for k in 1:min(n,6)
            x[k] = correctFirstFew[k]
        end
        for k in min(n,6)+1:n
            x[k] = McMahon(nu, k)
        end
    elseif nu > 5
        for k in 1:n
            x[k] = McMahon(nu, k)
        end
    end
    x
end

function McMahon(nu::Float64, k::Integer)
    # McMahon's expansion. This expansion gives very accurate approximation
    # for the sth zero (s >= 7) in the whole region V >=- 1, and moderate
    # approximation in other cases.
    mu = 4nu^2
    a1 = 1 / 8
    a3 = (7mu-31) / 384
    a5 = 4*(3779+mu*(-982+83mu)) / 61440 # Evaluate via Horner's method.
    a7 = 6*(-6277237+mu*(1585743+mu*(-153855+6949mu))) / 20643840
    a9 = 144*(2092163573+mu*(-512062548+mu*(48010494+mu*(-2479316+70197mu)))) / 11890851840
    a11 = 720 *(-8249725736393+mu*(1982611456181+mu*(-179289628602+mu*(8903961290 +
          mu*(-287149133 + 5592657mu))))) / 10463949619200
    a13 = 576 *(423748443625564327 + mu*(-100847472093088506 + mu*(8929489333108377 +
        mu*(-426353946885548+mu*(13172003634537+mu*(-291245357370 + 4148944183mu)))))) / 13059009124761600
    b = 0.25 * (2nu+4k-1)*pi
    # Evaluate using Horner's scheme:
    x = b - (mu-1)*( ((((((a13/b^2 + a11)/b^2 + a9)/b^2 + a7)/b^2 + a5)/b^2 + a3)/b^2 + a1)/b)
end

const J0_roots =
    [   2.4048255576957728
        5.5200781102863106
        8.6537279129110122
        11.791534439014281
        14.930917708487785
        18.071063967910922
        21.211636629879258
        24.352471530749302
        27.493479132040254
        30.634606468431975
        33.775820213573568
        36.917098353664044
        40.058425764628239
        43.199791713176730
        46.341188371661814
        49.482609897397817
        52.624051841114996
        55.765510755019979
        58.906983926080942
        62.048469190227170  ]

const Piessens_C = [
       2.883975316228  8.263194332307 11.493871452173 14.689036505931 17.866882871378 21.034784308088
       0.767665211539  4.209200330779  4.317988625384  4.387437455306  4.435717974422  4.471319438161
      -0.086538804759 -0.164644722483 -0.130667664397 -0.109469595763 -0.094492317231 -0.083234240394
       0.020433979038  0.039764618826  0.023009510531  0.015359574754  0.011070071951  0.008388073020
      -0.006103761347 -0.011799527177 -0.004987164201 -0.002655024938 -0.001598668225 -0.001042443435
       0.002046841322  0.003893555229  0.001204453026  0.000511852711  0.000257620149  0.000144611721
      -0.000734476579 -0.001369989689 -0.000310786051 -0.000105522473 -0.000044416219 -0.000021469973
       0.000275336751  0.000503054700  0.000083834770  0.000022761626  0.000008016197  0.000003337753
      -0.000106375704 -0.000190381770 -0.000023343325 -0.000005071979 -0.000001495224 -0.000000536428
       0.000042003336  0.000073681222  0.000006655551  0.000001158094  0.000000285903  0.000000088402
      -0.000016858623 -0.000029010830 -0.000001932603 -0.000000269480 -0.000000055734 -0.000000014856
       0.000006852440  0.000011579131  0.000000569367  0.000000063657  0.000000011033  0.000000002536
      -0.000002813300 -0.000004672877 -0.000000169722 -0.000000015222 -0.000000002212 -0.000000000438
       0.000001164419  0.000001903082  0.000000051084  0.000000003677  0.000000000448  0.000000000077
      -0.000000485189 -0.000000781030 -0.000000015501 -0.000000000896 -0.000000000092 -0.000000000014
       0.000000203309  0.000000322648  0.000000004736  0.000000000220  0.000000000019  0.000000000002
      -0.000000085602 -0.000000134047 -0.000000001456 -0.000000000054 -0.000000000004               0.
       0.000000036192  0.000000055969  0.000000000450  0.000000000013               0.               0.
      -0.000000015357 -0.000000023472 -0.000000000140 -0.000000000003               0.               0.
       0.000000006537  0.000000009882  0.000000000043  0.000000000001               0.               0.
      -0.000000002791 -0.000000004175 -0.000000000014               0.               0.               0.
       0.000000001194  0.000000001770  0.000000000004               0.               0.               0.
      -0.000000000512 -0.000000000752              0.               0.              0.               0.
       0.000000000220  0.000000000321               0.               0.               0.               0.
      -0.000000000095 -0.000000000137               0.               0.               0.               0.
       0.000000000041  0.000000000059               0.               0.               0.               0.
      -0.000000000018 -0.000000000025               0.               0.               0.               0.
       0.000000000008  0.000000000011               0.               0.               0.               0.
      -0.000000000003 -0.000000000005               0.               0.               0.               0.
       0.000000000001  0.000000000002               0.               0.               0.               0.]

function Piessens( nu::Float64 )
    # Piessens's Chebyshev series approximations (1984). Calculates the 6 first
    # zeros to at least 12 decimal figures in region -1 <= V <= 5:
    C = Piessens_C
    T = Array{Float64}(undef,size(C,1))
    pt = (nu-2)/3
    T[1], T[2] = 1., pt
    for k = 2:size(C,1)-1
        T[k+1] = 2*pt*T[k] - T[k-1]
    end
    y = C'*T
    y[1] *= sqrt(nu+1)                  # Scale the first root.
    y
end