text:
#Start real-time evolution from initial variational parameters WRBM_Init and with "method" integration scheme
function runDynamics(method,DYNHp::DYN_hp,WRBM_Init::Array{Complex{Float64},1})

    println("# Starting dynamics for $nspins = $L_x x $L_y spins and α=$alpha")
    println("# Time evolution = [0.,$(DYNHp.tot_time)], time-step = $(DYNHp.step_size). ")
    println("# Number of sweeps = $(DYNHp.samples)")
    println("# Number of workers = $(nprocs())")

    #Define array storing independent variational parameters at each time step
    #starting from the initial parameters WRBM_Init
    WRBM_t = Array{Array{Complex{Float64}}}(undef,0)
    push!(WRBM_t, WRBM_Init)

    #Iniziate the unitary time evolution
    @time method(RHS,WRBM_t,DYNHp)

end

#Heun integration scheme
function heun(GenRHS,WRBMt::Array{Array{Complex{Float64}}},DYNHp::DYN_hp)

    #initialized the time-interval of the simulation
    time = collect(0.:DYNHp.step_size:DYNHp.tot_time)

    #Assign independent variational parameters at t=0 to new variable W_
    W_ = WRBMt[1][:]

    #Initialize network parameters
    RBM_Par = init_RBM_par(alpha,nspins)
    #Sort network parameters to W_RBM independent parameters
    WToPar(RBM_Par,W_)

    #If the number of workers exceeds the number of variational parameters,
    #only a number of workers equals to the number of variational parameters will be used.
    #The exceeding workers will be removed.
    if nprocs() > dim
        println("WARNING: number of workers exceeds number of viariational parameters. Only $dim workers will be used.")
        diff = nprocs()-dim
        rmprocs(nprocs()-diff+1:nprocs())
    end

    #Define matrix storing energy and wavefunction gradients for each sampled state
    #@everywhere ensure a copy of ENWfDer is created for each worker
    @everywhere ENWfDer = Array{Complex{Float64}}(undef,@getfrom 1 NsampleDYN,dim + 1)

    save_ind = 1
    for t in time
            RHS_t = GenRHS(WRBMt[end],t,RBM_Par,true,DYN_HP.minres_)
            W_ .+=  DYNHp.step_size*(RHS_t .+ GenRHS(WRBMt[end] .+ DYNHp.step_size*RHS_t,t+DYNHp.step_size,RBM_Par,false,DYN_HP.minres_))/2
            push!(WRBMt, W_[:] ); save_ind += 1

            #Backup save of independent variational parameters every ten time steps
            if save_ind == 10
                writedlm("W_RBM_t_$(nspins)_$(alpha)_real.jl",real(WRBMt));writedlm("W_RBM_t_$(nspins)_$(alpha)_imag.jl",imag(WRBMt))
                save_ind = 1
            end

    end

    #Save indipendent variational parameters at each time step t
    writedlm("W_RBM_t_$(nspins)_$(alpha)_real.jl",real(WRBMt));writedlm("W_RBM_t_$(nspins)_$(alpha)_imag.jl",imag(WRBMt))

end

#Generate the right hand side of the TDVP equations  d/dt W(t) = -i*S^-1 F(W(t))
function RHS(W_t::Array{Complex{Float64},1},t::Float64,RBM_Par::RBM_par,time_flag::Bool,minres_::Bool)

    #Store number of workers for later use
    NW = Int64[]; push!(NW, NumberWorkers())

    #Define tt to pass current time to each worker
    @everywhere tt = Array{Float64}(undef,1)
    for i in procs()
    fetch(@spawnat i tt[1] = t )
    end

    #perturbation of J_ex along y
    @everywhere J = [J_x + DeltaJ(tt[1]); J_y]

    #Update neural network parameters at the current time t
    for i in procs()
        fetch(@spawnat i WToPar(RBM_Par,W_t))
    end

    #Sampling of states
    #At the end of the loop, the energy sampled are returned in EnergyStates
    #The derivatives of the wavefunction are stored in ENWFDer
     EnergyStates = @distributed (vcat) for i=1:NW[1]
          sampling(RBM_Par,J,i)
     end

     #Evaluation entries of S-matrix stored in the vector SVector
     SVector = @distributed (+) for i=1:NW[1]
        S_gen(dim)./NW[1]
     end

     #Sort entries SVector into the S matrix
     SMatrix = Gen_S(SVector)

     #Store the total energy for each state
     energy = EnergyStates[:]
     #Calculate the mean value of the energy and energy variance
     mean_E = mean(energy); var_E = var(energy)

     #Calculate the derivative of the energy respect to the variational parameters and assign them to the array dF
     EGrad = @distributed (+) for i=1:NW[1]
        Egrad(dim)./NW[1]
     end

     #Print the energy only in the first step of the Heun scheme
      if time_flag == true
      println("time $t- total energy: $(real(mean_E)) +- $(real(var_E)/length(energy))")
      end

      #Return the rhs of the TDVP equation
      if minres_ == true
          return rhsMinres(SMatrix,EGrad)
      else
          return rhsPinv(SMatrix,EGrad)
      end

end

#Generate rhs of TDVP equation by using the iterative solver MINRES
function rhsMinres(S_::Matrix{Complex{Float64}}, eGrad::Array{Complex{Float64},1})

    return minres(S_,-im*eGrad)

end

#Generate rhs of TDVP equation by using singular value decomposition (pseudo inverse)
function rhsPinv(S_::Matrix{Complex{Float64}}, eGrad::Array{Complex{Float64},1})

    #Regularization of the S-matrix with a small shift of the diagonal element of S
    S_ .= S_ .+ Matrix{Float64}(I,dim,dim).*0.0001

    #Calculate the pseudo inverse of S_
    invS = pinv(S_)

    return -im*invS*eGrad

end

#Tamed Euler integration scheme
function tamedEuler(GenRHS,WRBMt::Array{Array{Complex{Float64}}},DYNHp::DYN_hp)

    #initialized the time-interval of the simulation
    time = collect(0.:DYNHp.step_size:DYNHp.tot_time)

    #Assign independent variational parameters at t=0 to new variable W_
    W_ = WRBMt[1][:]

    #Initialize network parameters
    RBM_Par = init_RBM_par(alpha,nspins)
    #Sort network parameters to W_RBM independent parameters
    WToPar(RBM_Par,W_)

    #If the number of workers exceeds the number of variational parameters,
    #only a number of workers equals to the number of variational parameters will be used.
    #The exceeding workers will be removed.
    if nprocs() > dim
        println("WARNING: number of workers exceeds number of viariational parameters. Only $dim workers will be used.")
        diff = nprocs()-dim
        rmprocs(nprocs()-diff+1:nprocs())
    end

    #Define matrix storing energy and wavefunction gradients for each sampled state
    #@everywhere ensure a copy of ENWfDer is created for each worker
    @everywhere ENWfDer = Array{Complex{Float64}}(undef,@getfrom 1 NsampleDYN,dim + 1)

    save_ind = 1
    for t in time
            RHS_t = GenRHS(WRBMt[end],t,RBM_Par,true,DYNHp.minres_)
            W_ .+= (RHS_t * DYNHp.step_size)./(1 .+ DYNHp.step_size*norm(RHS_t))
            push!(WRBMt, W_[:] ); save_ind += 1

            #Backup save of independent variational parameters every ten time steps
            if save_ind == 10
                writedlm("W_RBM_t_$(nspins)_$(alpha)_real.jl",real(WRBMt));writedlm("W_RBM_t_$(nspins)_$(alpha)_imag.jl",imag(WRBMt))
                save_ind = 1
            end

    end

    #Save independent variational parameters at each time step t
    writedlm("W_RBM_t_$(nspins)_$(alpha)_real.jl",real(WRBMt));writedlm("W_RBM_t_$(nspins)_$(alpha)_imag.jl",imag(WRBMt))

end

#Tamed Heun integration scheme
function tamedHeun(GenRHS,WRBMt::Array{Array{Complex{Float64}}},DYNHp::DYN_hp)

    #initialized the time-interval of the simulation
    time = collect(0.:DYNHp.step_size:DYNHp.tot_time)

    #Assign independent variational parameters at t=0 to new variable W_
    W_ = WRBMt[1][:]

    #Initialize network parameters
    RBM_Par = init_RBM_par(alpha,nspins)
    #Sort network parameters to W_RBM independent parameters
    WToPar(RBM_Par,W_)

    #If the number of workers exceeds the number of variational parameters,
    #only a number of workers equals to the number of variational parameters will be used.
    #The exceeding workers will be removed.
    if nprocs() > dim
        println("WARNING: number of workers exceeds number of viariational parameters. Only $dim workers will be used.")
        diff = nprocs()-dim
        rmprocs(nprocs()-diff+1:nprocs())
    end

    #Define matrix storing energy and wavefunction gradients for each sampled state
    #@everywhere ensure a copy of ENWfDer is created for each worker
    @everywhere ENWfDer = Array{Complex{Float64}}(undef,@getfrom 1 NsampleDYN,dim + 1)

    save_ind = 1
    for t in time
            RHS_t1 = GenRHS(WRBMt[end],t,RBM_Par,true,DYNHp.minres_)
            RHS_t2 = GenRHS(WRBMt[end] .+ DYNHp.step_size*RHS_t1,t+DYNHp.step_size,RBM_Par,false,DYN_HP.minres_)
            W_ .+=  0.5*DYNHp.step_size*(RHS_t1/(1 + DYNHp.step_size*norm(RHS_t1) ) .+  RHS_t2/(1 + DYNHp.step_size*norm(RHS_t2)))
            push!(WRBMt, W_[:] ); save_ind += 1

            #Backup save of independent variational parameters every ten time steps
            if save_ind == 10
                writedlm("W_RBM_t_$(nspins)_$(alpha)_real.jl",real(WRBMt));writedlm("W_RBM_t_$(nspins)_$(alpha)_imag.jl",imag(WRBMt))
                save_ind = 1
            end

    end

    #Save independent variational parameters at each time step t
    writedlm("W_RBM_t_$(nspins)_$(alpha)_real.jl",real(WRBMt));writedlm("W_RBM_t_$(nspins)_$(alpha)_imag.jl",imag(WRBMt))

end

max_stars_repo_name: GiamFbn/MyULTRAFAST

max_issues_repo_name: GiamFbn/MyULTRAFAST

max_forks_repo_name: GiamFbn/MyULTRAFAST


text:
@testset "Z Scoring" begin
    data = (x=collect(1:10),
            y=rand(10) .+ 3,
            z=Symbol.(repeat('a':'e', 2)))

    @testset "Automatic scaling" begin
        xc = concrete_term(term(:x), data, ZScore())
        @test xc isa ZScoredTerm
        @test xc.center ≈ mean(data.x)
        @test xc.scale ≈ std(data.x)
        @test modelcols(xc, data) ≈ zscore(data.x) ≈ zscore(data.x, xc.center, xc.scale)

        yc = concrete_term(term(:y), data, ZScore())
        @test yc.center ≈ mean(data.y)
        @test yc.scale ≈ std(data.y)
        @test modelcols(yc, data) ≈ zscore(data.y) ≈ zscore(data.y, yc.center, yc.scale)

        @testset "alternative center and scale functions" begin
            f, g = first, mad
            xc = concrete_term(term(:x), data, ZScore(f, g))
            @test xc isa ZScoredTerm
            @test xc.center ≈ f(data.x)
            @test xc.scale ≈ g(data.x)
            @test modelcols(xc, data) ≈ zscore(data.x, xc.center, xc.scale)
            # why test this? well this makes sure that our tests
            # wouldn't pass if we were using the default scale function
            # in other words, this tests we're actually hitting a different codepath
            @test !isapprox(mean(data.x), f(data.x))
            @test !isapprox(std(data.x), g(data.x))
        end
    end

    @testset "Manual scaling" begin
        xc = concrete_term(term(:x), data, ZScore(; center=5))
        @test xc isa ZScoredTerm
        @test xc.center ≈ 5 != mean(data.x)
        @test xc.scale ≈ std(data.x)
        @test modelcols(xc, data) ≈ zscore(data.x, xc.center, xc.scale)

        yc = concrete_term(term(:y), data, ZScore(; scale=3))
        @test yc isa ZScoredTerm
        @test yc.center ≈ mean(data.y)
        @test yc.scale == 3 != std(data.y)
        @test modelcols(yc, data) ≈ zscore(data.y, yc.center, yc.scale)

        yc = concrete_term(term(:y), data, ZScore(; scale=3, center=1))
        @test yc isa ZScoredTerm
        @test yc.center ≈ 1 != mean(data.y)
        @test yc.scale == 3 != std(data.y)
        @test modelcols(yc, data) ≈ zscore(data.y, yc.center, yc.scale)

        @test zscore(concrete_term(term(:x), data, nothing); center=3, scale=5) ==
              concrete_term(term(:x), data, ZScore(; center=3, scale=5))
    end

    @testset "Schema hints dict" begin
        sch = schema(data, Dict(:x => ZScore(), :y => ZScore(; center=2.2, scale=5)))
        xc = sch[term(:x)]
        @test xc isa ZScoredTerm
        @test !StatsModels.needs_schema(xc)
        @test xc.center == mean(data.x)
        @test xc.scale == std(data.x)

        yc = sch[term(:y)]
        @test yc isa ZScoredTerm
        @test yc.center == 2.2
        @test yc.scale == 5

        @test modelcols(xc, data) ≈ zscore(data.x)
        @test modelcols(yc, data) ≈ zscore(data.y, 2.2, 5)
    end

    # no need to test "zscore function"
    # taken care of by StatsBase 😎

    @testset "plays nicely with formula" begin
        f = @formula(0 ~ x * y)

        sch = schema(f, data)

        sch_c = schema(f, data, Dict(:x => ZScore(), :y => ZScore(; center=2.2, scale=5)))
        ff_c = apply_schema(f, sch_c)

        mm_c = modelcols(ff_c.rhs, data)
        @test mm_c ≈ hcat(zscore(data.x),
                          zscore(data.y, 2.2, 5),
                          zscore(data.x) .* zscore(data.y, 2.2, 5))

        @test coefnames(ff_c.rhs) == ["x(centered: 5.5 scaled: 3.03)",
                                      "y(centered: 2.2 scaled: 5)",
                                      "x(centered: 5.5 scaled: 3.03) & y(centered: 2.2 scaled: 5)"]

        # round-trip schema is empty since needs_schema is false
        sch_2 = schema(ff_c, data)
        @test isempty(sch_2.schema)

        @testset "missing omit" begin
            d = (x=[1, 2, missing, 3], y=[1, missing, 2, 3], z=[missing, 1, 2, 3])
            dd, nonmissing = StatsModels.missing_omit(d, ff_c)
            # z is ignored:
            @test findall(nonmissing) == [1, 4]
            # x and y are handled correclty
            @test dd.x == d.x[nonmissing]
            @test dd.y == d.y[nonmissing]
        end
    end

    @testset "printing" begin
        xc = concrete_term(term(:x), data, ZScore())
        @test StatsModels.termsyms(xc) == Set([:x])
        @test "$(xc)" == string_mime(MIME("text/plain"), xc) ==
              "x(centered: 5.5 scaled: 3.03)"
        @test coefnames(xc) == "x(centered: 5.5 scaled: 3.03)"
    end

    @testset "categorical term" begin
        z = concrete_term(term(:z), data)
        @test_throws ArgumentError zscore(z)
        @test_throws ArgumentError zscore(z, ZScore())
        zc = zscore(z, ZScore(; center=0.5, scale=3))
        @test modelcols(zc, data) ≈ (modelcols(z, data) .- 0.5) ./ 3
        @test StatsModels.width(zc) == StatsModels.width(z)
        @test coefnames(zc) == coefnames(z) .* "(centered: 0.5 scaled: 3)"

        zc2 = zscore(z, ZScore(; center=[1 2 3 4], scale=[5 6 7 8]))
        @test string_mime(MIME("text/plain"), zc2) ==
              "z(centered: [1 2 3 4] scaled: [5 6 7 8])"
        @test modelcols(zc2, data) ≈ (modelcols(z, data) .- [1 2 3 4]) ./ [5 6 7 8]
        @test coefnames(zc2) ==
              coefnames(z) .* "(centered: " .* string.([1, 2, 3, 4]) .* " scaled: " .*
              string.([5, 6, 7, 8]) .* ")"
    end
end

ext: jl


text:
module CPython

using Libdl, Dates, Compat, Requires
import ..PythonCall:
    CONFIG,
    isnull,
    ism1,
    PYERR,
    NOTIMPLEMENTED,
    _typeintersect,
    _type_lb,
    _type_ub,
    _type_union_split,
    tryconvert,
    ispyreftype,
    ispyref,
    pyptr,
    putresult,
    takeresult,
    CACHE,
    PythonCall
using Base: @kwdef
using UnsafePointers: UnsafePtr

pyglobal(name) = dlsym(CONFIG.libptr, name)
pyloadglobal(::Type{T}, name) where {T} = unsafe_load(Ptr{T}(pyglobal(name)))
pyglobal(r::Ref{Ptr{T}}, name) where {T} = begin
    p = r[]
    if isnull(p)
        p = r[] = Ptr{T}(pyglobal(name))
    end
    p
end
pyloadglobal(r::Ref{Ptr{T}}, name) where {T} = begin
    p = r[]
    if isnull(p)
        p = r[] = unsafe_load(Ptr{Ptr{T}}(pyglobal(name)))
    end
    p
end

include("consts.jl")
include("pointers.jl")
include("dsl.jl")
include("fundamentals.jl")
include("none.jl")
include("object.jl")
include("sequence.jl")
include("mapping.jl")
include("method.jl")
include("str.jl")
include("bytes.jl")
include("tuple.jl")
include("type.jl")
include("number.jl")
include("iter.jl")
include("bool.jl")
include("int.jl")
include("float.jl")
include("complex.jl")
include("list.jl")
include("dict.jl")
include("set.jl")
include("buffer.jl")
include("stdlib.jl")
include("collections.jl")
include("io.jl")
include("range.jl")
include("ctypes.jl")
include("numpy.jl")
include("slice.jl")
include("fraction.jl")
include("datetime.jl")
include("newtype.jl")
include("juliaerror.jl")
include("juliabase.jl")
include("juliaraw.jl")
include("juliaany.jl")
include("juliaiterator.jl")
include("juliatype.jl")
include("juliadict.jl")
include("juliaset.jl")
include("juliaarray.jl")
include("juliavector.jl")
include("juliamodule.jl")
include("julianumber.jl")
include("juliaio.jl")
include("as.jl")
include("arg.jl")

init() = begin
    PyObject_TryConvert_AddRule("builtins.NoneType", Nothing, PyNone_TryConvertRule_nothing, 100)
    PyObject_TryConvert_AddRule("builtins.NoneType", Missing, PyNone_TryConvertRule_missing)
    PyObject_TryConvert_AddRule("builtins.bool", Bool, PyBool_TryConvertRule_bool, 100)
    PyObject_TryConvert_AddRule("numbers.Integral", Integer, PyLongable_TryConvertRule_integer, 100)
    PyObject_TryConvert_AddRule("numbers.Integral", Rational, PyLongable_TryConvertRule_tryconvert)
    PyObject_TryConvert_AddRule("numbers.Integral", Real, PyLongable_TryConvertRule_tryconvert)
    PyObject_TryConvert_AddRule("numbers.Integral", Number, PyLongable_TryConvertRule_tryconvert)
    PyObject_TryConvert_AddRule("numbers.Integral", Any, PyLongable_TryConvertRule_tryconvert)
    PyObject_TryConvert_AddRule("builtins.float", Float64, PyFloatable_TryConvertRule_convert, 100)
    PyObject_TryConvert_AddRule("numbers.Real", Float64, PyFloatable_TryConvertRule_convert)
    PyObject_TryConvert_AddRule("numbers.Real", BigFloat, PyFloatable_TryConvertRule_convert)
    PyObject_TryConvert_AddRule("numbers.Real", Float32, PyFloatable_TryConvertRule_convert)
    PyObject_TryConvert_AddRule("numbers.Real", Float16, PyFloatable_TryConvertRule_convert)
    PyObject_TryConvert_AddRule("numbers.Real", AbstractFloat, PyFloatable_TryConvertRule_tryconvert)
    PyObject_TryConvert_AddRule("numbers.Real", Real, PyFloatable_TryConvertRule_tryconvert)
    PyObject_TryConvert_AddRule("numbers.Real", Number, PyFloatable_TryConvertRule_tryconvert)
    PyObject_TryConvert_AddRule("numbers.Real", Nothing, PyFloatable_TryConvertRule_nothing)
    PyObject_TryConvert_AddRule("numbers.Real", Missing, PyFloatable_TryConvertRule_missing)
    PyObject_TryConvert_AddRule("builtins.complex", Complex{Float64}, PyComplexable_TryConvertRule_convert, 100)
    PyObject_TryConvert_AddRule("numbers.Complex", Complex{Float64}, PyComplexable_TryConvertRule_convert)
    PyObject_TryConvert_AddRule("numbers.Complex", Complex{BigFloat}, PyComplexable_TryConvertRule_convert)
    PyObject_TryConvert_AddRule("numbers.Complex", Complex{Float32}, PyComplexable_TryConvertRule_convert)
    PyObject_TryConvert_AddRule("numbers.Complex", Complex{Float16}, PyComplexable_TryConvertRule_convert)
    PyObject_TryConvert_AddRule("numbers.Complex", Complex{T} where {T<:AbstractFloat}, PyComplexable_TryConvertRule_tryconvert)
    PyObject_TryConvert_AddRule("numbers.Complex", Complex{T} where {T<:Real}, PyComplexable_TryConvertRule_tryconvert)
    PyObject_TryConvert_AddRule("numbers.Complex", Number, PyComplexable_TryConvertRule_tryconvert)
    PyObject_TryConvert_AddRule("builtins.bytes", Vector{UInt8}, PyBytes_TryConvertRule_vector)
    PyObject_TryConvert_AddRule("builtins.bytes", Vector{Int8}, PyBytes_TryConvertRule_vector)
    PyObject_TryConvert_AddRule("builtins.bytes", String, PyBytes_TryConvertRule_string)
    PyObject_TryConvert_AddRule("builtins.str", String, PyUnicode_TryConvertRule_string, 100)
    PyObject_TryConvert_AddRule("builtins.str", Symbol, PyUnicode_TryConvertRule_symbol)
    PyObject_TryConvert_AddRule("builtins.str", Char, PyUnicode_TryConvertRule_char)
    PyObject_TryConvert_AddRule("builtins.str", Vector{UInt8}, PyUnicode_TryConvertRule_vector)
    PyObject_TryConvert_AddRule("builtins.str", Vector{Int8}, PyUnicode_TryConvertRule_vector)
    PyObject_TryConvert_AddRule("builtins.tuple", Tuple, PyIterable_ConvertRule_tuple, 100)
    PyObject_TryConvert_AddRule("builtins.range", StepRange{T,S} where {T<:Integer,S<:Integer}, PyRange_TryConvertRule_steprange, 100)
    PyObject_TryConvert_AddRule("builtins.range", UnitRange{T} where {T<:Integer}, PyRange_TryConvertRule_unitrange)
    PyObject_TryConvert_AddRule("collections.abc.Iterable", Vector, PyIterable_ConvertRule_vecorset)
    PyObject_TryConvert_AddRule("collections.abc.Iterable", Set, PyIterable_ConvertRule_vecorset)
    PyObject_TryConvert_AddRule("collections.abc.Iterable", Tuple, PyIterable_ConvertRule_tuple)
    PyObject_TryConvert_AddRule("collections.abc.Iterable", NamedTuple, PyIterable_ConvertRule_namedtuple)
    PyObject_TryConvert_AddRule("collections.abc.Iterable", Pair, PyIterable_ConvertRule_pair)
    PyObject_TryConvert_AddRule("collections.abc.Sequence", Vector, PyIterable_ConvertRule_vecorset)
    PyObject_TryConvert_AddRule("collections.abc.Set", Set, PyIterable_ConvertRule_vecorset)
    PyObject_TryConvert_AddRule("collections.abc.Mapping", Dict, PyMapping_ConvertRule_dict)
    PyObject_TryConvert_AddRule("datetime.time", Time, PyTime_TryConvertRule_time, 100)
    PyObject_TryConvert_AddRule("datetime.date", Date, PyDate_TryConvertRule_date, 100)
    PyObject_TryConvert_AddRule("datetime.datetime", DateTime, PyDateTime_TryConvertRule_datetime, 100)
    PyObject_TryConvert_AddRule("datetime.timedelta", Dates.Period, PyTimeDelta_TryConvertRule_period, 100)
    PyObject_TryConvert_AddRule("datetime.timedelta", Dates.CompoundPeriod, PyTimeDelta_TryConvertRule_compoundperiod)
    PyObject_TryConvert_AddRule("juliacall.ValueBase", Any, PyJuliaValue_TryConvert_any, 1000)
    PyObject_TryConvert_AddRule("juliacall.As", Any, PyAs_ConvertRule_tryconvert, 1000)
    PyObject_TryConvert_AddExtraType(PyIterableABC_Type)
    PyObject_TryConvert_AddExtraType(PyCallableABC_Type)
    PyObject_TryConvert_AddExtraType(PySequenceABC_Type)
    PyObject_TryConvert_AddExtraType(PyMappingABC_Type)
    PyObject_TryConvert_AddExtraType(PySetABC_Type)
    PyObject_TryConvert_AddExtraType(PyNumberABC_Type)
    PyObject_TryConvert_AddExtraType(PyComplexABC_Type)
    PyObject_TryConvert_AddExtraType(PyRealABC_Type)
    PyObject_TryConvert_AddExtraType(PyRationalABC_Type)
    PyObject_TryConvert_AddExtraType(PyIntegralABC_Type)

    ### ctypes
    for (p, T) in [
        ("char", Cchar),
        ("wchar", Cwchar_t),
        ("byte", Cchar),
        ("ubyte", Cuchar),
        ("short", Cshort),
        ("ushort", Cushort),
        ("int", Cint),
        ("uint", Cuint),
        ("long", Clong),
        ("ulong", Culong),
        ("longlong", Clonglong),
        ("ulonglong", Culonglong),
        ("size_t", Csize_t),
        ("ssize_t", Cssize_t),
        ("float", Cfloat),
        ("double", Cdouble), #=("longdouble", ???),=#
        ("char_p", Ptr{Cchar}),
        ("wchar_p", Ptr{Cwchar_t}),
        ("void_p", Ptr{Cvoid}),
    ]
        isptr = occursin("_p", p)
        isfloat = occursin("float", p) || occursin("double", p)
        isint = !(isfloat || isptr)
        isreal = isint || isfloat
        p == "char_p" && PyObject_TryConvert_AddRule("ctypes.c_$p", Cstring, PySimpleCData_TryConvert_value{T,false}())
        p == "wchar_p" && PyObject_TryConvert_AddRule("ctypes.c_$p", Cwstring, PySimpleCData_TryConvert_value{T,false}())
        PyObject_TryConvert_AddRule("ctypes.c_$p", T, PySimpleCData_TryConvert_value{T,false}())
        isint && PyObject_TryConvert_AddRule("ctypes.c_$p", Integer, PySimpleCData_TryConvert_value{T,true}())
        isint && PyObject_TryConvert_AddRule("ctypes.c_$p", Rational, PySimpleCData_TryConvert_value{T,true}())
        isreal && PyObject_TryConvert_AddRule("ctypes.c_$p", Float64, PySimpleCData_TryConvert_value{T,false}())
        isreal && PyObject_TryConvert_AddRule("ctypes.c_$p", BigFloat, PySimpleCData_TryConvert_value{T,false}())
        isreal && PyObject_TryConvert_AddRule("ctypes.c_$p", Float32, PySimpleCData_TryConvert_value{T,false}())
        isreal && PyObject_TryConvert_AddRule("ctypes.c_$p", Float16, PySimpleCData_TryConvert_value{T,false}())
        isreal && PyObject_TryConvert_AddRule("ctypes.c_$p", AbstractFloat, PySimpleCData_TryConvert_value{T,true}())
        isreal && PyObject_TryConvert_AddRule("ctypes.c_$p", Real, PySimpleCData_TryConvert_value{T,true}())
        isreal && PyObject_TryConvert_AddRule("ctypes.c_$p", Number, PySimpleCData_TryConvert_value{T,true}())
        isptr && PyObject_TryConvert_AddRule("ctypes.c_$p", Ptr, PySimpleCData_TryConvert_value{T,false}())
        PyObject_TryConvert_AddRule("ctypes.c_$p", Any, PySimpleCData_TryConvert_value{T,true}())
    end

    ### numpy
    # TODO: Compound types
    # TODO: datetime64, timedelta64
    for (p, T) in [
        ("int8", Int8),
        ("int16", Int16),
        ("int32", Int32),
        ("int64", Int64),
        ("int128", Int128),
        ("uint8", UInt8),
        ("uint16", UInt16),
        ("uint32", UInt32),
        ("uint64", UInt64),
        ("uint128", UInt128),
        ("float16", Float16),
        ("float32", Float32),
        ("float64", Float64),
        ("complex32", Complex{Float16}),
        ("complex64", Complex{Float32}),
        ("complex128", Complex{Float64}),
    ]
        isint = occursin("int", p)
        isfloat = occursin("float", p)
        iscomplex = occursin("complex", p)
        isreal = isint || isfloat
        PyObject_TryConvert_AddRule("numpy.$p", T, PyNumpySimpleData_TryConvert_value{T,false}(), 100)
        isint && PyObject_TryConvert_AddRule("numpy.$p", Integer, PyNumpySimpleData_TryConvert_value{T,true}())
        isint && PyObject_TryConvert_AddRule("numpy.$p", Rational, PyNumpySimpleData_TryConvert_value{T,true}())
        isreal && PyObject_TryConvert_AddRule("numpy.$p", Float64, PyNumpySimpleData_TryConvert_value{T,false}())
        isreal && PyObject_TryConvert_AddRule("numpy.$p", BigFloat, PyNumpySimpleData_TryConvert_value{T,false}())
        isreal && PyObject_TryConvert_AddRule("numpy.$p", Float32, PyNumpySimpleData_TryConvert_value{T,false}())
        isreal && PyObject_TryConvert_AddRule("numpy.$p", Float16, PyNumpySimpleData_TryConvert_value{T,false}())
        isreal && PyObject_TryConvert_AddRule("numpy.$p", AbstractFloat, PyNumpySimpleData_TryConvert_value{T,true}())
        isreal && PyObject_TryConvert_AddRule("numpy.$p", Real, PyNumpySimpleData_TryConvert_value{T,true}())
        iscomplex && PyObject_TryConvert_AddRule("numpy.$p", Complex{Float64}, PyNumpySimpleData_TryConvert_value{T,false}())
        iscomplex && PyObject_TryConvert_AddRule("numpy.$p", Complex{BigFloat}, PyNumpySimpleData_TryConvert_value{T,false}())
        iscomplex && PyObject_TryConvert_AddRule("numpy.$p", Complex{Float32}, PyNumpySimpleData_TryConvert_value{T,false}())
        iscomplex && PyObject_TryConvert_AddRule("numpy.$p", Complex{Float16}, PyNumpySimpleData_TryConvert_value{T,false}())
        iscomplex && PyObject_TryConvert_AddRule("numpy.$p", Complex{T} where {T<:AbstractFloat}, PyNumpySimpleData_TryConvert_value{T,true}())
        iscomplex && PyObject_TryConvert_AddRule("numpy.$p", Complex{T} where {T<:Real}, PyNumpySimpleData_TryConvert_value{T,true}())
        PyObject_TryConvert_AddRule("numpy.$p", Number, PyNumpySimpleData_TryConvert_value{T,true}())
        PyObject_TryConvert_AddRule("numpy.$p", Any, PyNumpySimpleData_TryConvert_value{T,true}())
    end
end
precompile(init, ())
@init init()

precompile(PyJuliaValue_New, (PyPtr, Function))
precompile(PyMappingABC_Type, ())

end


text:
setfont() = GR.settextfontprec(232, 3) # CM Serif Roman

function search(s::AbstractString, c::Char)
    result = findfirst(isequal(c), s)
    result != nothing ? result : 0
end

function inqtext(x, y, s)
    if length(s) >= 2 && s[1] == '$' && s[end] == '$'
        tbx, tby = GR.inqmathtex(x, y, s[2:end-1])
    elseif search(s, '\\') != 0 || search(s, '_') != 0 || search(s, '^') != 0
        tbx, tby = GR.inqtextext(x, y, s)
    else
        tbx, tby = GR.inqtext(x, y, s)
    end
    tbx, tby
end

"""
    stringsize(s, charsize [, wc=false])

Return the width and height of the string `s`, written with characters
with height equal to `charsize`.

By default, the dimensions are returned in NDC;
use `wc=true` to return them in world coordinates.
"""
function stringsize(s, charsize, wc=false)
    GR.savestate()
    wc ? GR.selntran(1) : GR.selntran(0)
    setfont() # needed for GR.inqmathtex
    GR.setcharheight(charsize)
    tbx, tby = inqtext(0, 0, s)
    GR.restorestate()
    w = tbx[2] - tbx[1]
    h = tby[4] - tby[1]
    w, h
end

"""
    text(x, y, s [, wc=false])

Write the string `s` at the coordinates `x`, `y`.
NDC coordinates are assumed by default.
use `wc=true` to use world coordinates.
"""
function text(x, y, s, wc = false)
    if wc
        win = GR.inqwindow()
        vp = GR.inqviewport()
        x = (vp[2] - vp[1]) / (win[2] - win[1]) * x + (vp[1] * win[2] - vp[2] * win[1]) / (win[2] - win[1])
        y = (vp[4] - vp[3]) / (win[4] - win[3]) * y + (vp[3] * win[4] - vp[4] * win[3]) / (win[4] - win[3])
    end
    GR.savestate() # to avoid side-effect of GR.mathtex
    # setfont() # needed for GR.mathtex
    if length(s) >= 2 && s[1] == '$' && s[end] == '$'
        GR.mathtex(x, y, s[2:end-1])
    elseif search(s, '\\') != 0 || search(s, '_') != 0 || search(s, '^') != 0
        GR.textext(x, y, s)
    else
        GR.text(x, y, s)
    end
    GR.restorestate()
end


text:
import DifferentialDynamicsModels: SteeringBVP, propagate, state_dim, control_dim

export SimpleCarDynamics
export SE2State, SE2vState, SE2κState, SE2vκState
export VelocityCurvatureControl, AccelerationCurvatureControl, VelocityCurvRateControl, AccelerationCurvRateControl
export VelocityCurvatureStep, VelocityCurvRateStep
export DubinsSteering, ReedsSheppSteering, DubinsCCSteering

# Simple Car Dynamics (with integrators in speed v and curvature κ)
struct SimpleCarDynamics{Dv,Dκ} <: DifferentialDynamics end

state_dim(::SimpleCarDynamics{Dv,Dκ}) where {Dv,Dκ} = Dv + Dκ + 3
control_dim(::SimpleCarDynamics{Dv,Dκ}) where {Dv,Dκ} = 2

@maintain_type struct SE2State{T}   <: FieldVector{3,T} x::T; y::T; θ::T end
@maintain_type struct SE2vState{T}  <: FieldVector{4,T} x::T; y::T; θ::T; v::T end
@maintain_type struct SE2κState{T}  <: FieldVector{4,T} x::T; y::T; θ::T; κ::T end
@maintain_type struct SE2vκState{T} <: FieldVector{5,T} x::T; y::T; θ::T; v::T; κ::T end

@maintain_type struct VelocityCurvatureControl{T}     <: FieldVector{2,T}; v::T; κ::T end
@maintain_type struct AccelerationCurvatureControl{T} <: FieldVector{2,T}; a::T; κ::T end
@maintain_type struct VelocityCurvRateControl{T}      <: FieldVector{2,T}; v::T; σ::T end
@maintain_type struct AccelerationCurvRateControl{T}  <: FieldVector{2,T}; a::T; σ::T end

const VelocityCurvatureStep{T} = StepControl{2,T,VelocityCurvatureControl{T}}    # alias used for dubins, reedsshepp
const VelocityCurvRateStep{T} = StepControl{2,T,VelocityCurvRateControl{T}}      # alias used for dubinsCC, elementary

# Dv = 0, Dκ = 0 (Dubins Car, Reeds-Shepp Car)
(::SimpleCarDynamics{0,0})((x, y, θ)::StaticVector{3}, (v, κ)::StaticVector{2}) = SVector(v*cos(θ), v*sin(θ), v*κ)
function propagate(f::SimpleCarDynamics{0,0}, q::StaticVector{3,X}, c::StepControl{2,U}) where {X,U}
    t, v, κ = c.t, c.u[1], c.u[2]
    x, y, θ = q[1], q[2], q[3]
    T = promote_type(X, U)
    s, c = sincos(θ)
    if abs(κ) > sqrt(eps(T))
        st, ct = sincos(θ + v*κ*t)
        similar_type(q, T)(x + (st - s)/κ,
                           y + (c - ct)/κ,
                           mod2piF(θ + v*κ*t))
    else
        similar_type(q, T)(x + v*c*t - v*v*s*κ*t*t/2,
                           y + v*s*t + v*v*c*κ*t*t/2,
                           mod2piF(θ + v*κ*t))
    end
end

# Dv = 1, Dκ = 0 (Dubins Car, Reeds-Shepp Car with Acceleration)
(::SimpleCarDynamics{1,0})((x, y, θ, v)::StaticVector{4}, (a, κ)::StaticVector{2}) = SVector(v*cos(θ), v*sin(θ), v*κ, a)
function propagate(f::SimpleCarDynamics{1,0}, q::StaticVector{4,X}, c::StepControl{2,U}) where {X,U}
    t, a, κ = c.t, c.u[1], c.u[2]
    x, y, θ, v = q[1], q[2], q[3], q[4]
    T = promote_type(X, U)
    s = v*t + a*t*t/2
    xyθ = propagate(SimpleCarDynamics{0,0}(), SE2State(x, y, θ), StepControl(s, VelocityCurvatureControl(T(1), κ)))
    similar_type(q, T)(xyθ[1], xyθ[2], xyθ[3], v + a*t)
end

# Dv = 0, Dκ = 1 (Continuous Curvature Car)
(::SimpleCarDynamics{0,1})((x, y, θ, κ)::StaticVector{4}, (v, σ)::StaticVector{2}) = SVector(v*cos(θ), v*sin(θ), v*κ, σ)
function propagate(f::SimpleCarDynamics{0,1}, q::StaticVector{4,X}, c::StepControl{2,U}) where {X,U}
    t, v, σ = c.t, c.u[1], c.u[2]
    x, y, θ, κ = q[1], q[2], q[3], q[4]
    T = promote_type(X, U)
    if abs(σ) > sqrt(eps(T))
        spi = sqrt(T(pi))
        θ_ = θ - v*κ*κ/σ/2
        s, c = sincos(θ_)
        sv_σ = sqrt(abs(v/σ))
        PK = sv_σ*κ/spi
        PT = sv_σ*σ*t/spi
        FCK = flipsign(flipsign(fresnelC(PK), σ), v)
        FCT = flipsign(flipsign(fresnelC(PK + PT), σ), v)
        FSK = fresnelS(PK)
        FST = fresnelS(PK + PT)
        similar_type(q, T)(x + spi*sv_σ*(c*(FCT-FCK) + s*(FSK-FST)),
                           y + spi*sv_σ*(c*(FST-FSK) + s*(FCT-FCK)),
                           mod2piF(θ + v*κ*t + v*σ*t*t/2),
                           κ + σ*t)
    elseif abs(κ) > sqrt(eps(T))
        s, c = sincos(θ)
        st, ct = sincos(θ + v*κ*t)
        similar_type(q, T)(x + (st - s)/κ,    # perhaps include higher order σ terms
                           y + (c - ct)/κ,
                           mod2piF(θ + v*κ*t + v*σ*t*t/2),
                           κ + σ*t)
    else
        s, c = sincos(θ)
        similar_type(q, T)(x + v*c*t,         # perhaps include higher order κ, σ terms
                           y + v*s*t,
                           mod2piF(θ + v*κ*t + v*σ*t*t/2),
                           κ + σ*t)
    end
end

# General Simple Car Dynamics
(::SimpleCarDynamics{1,1})((x, y, θ, v, κ)::StaticVector{5}, (a, σ)::StaticVector{2}) = SVector(v*cos(θ), v*sin(θ), v*κ, a, σ)
@generated function (::SimpleCarDynamics{Dv,Dκ})(q::StaticVector{N}, u::StaticVector{2}) where {Dv,Dκ,N}
    @assert Dv + Dκ + 3 == N
    θ = :(q[3])
    v = :(q[4])
    κ = :(q[$(4+Dv)])
    qdot = [:($v*cos($θ)); :($v*sin($θ)); :($v*$κ);
            [:(q[$(4+i)]) for i in 1:Dv-1]; :(u[1]);
            [:(q[$(4+Dv+i)]) for i in 1:Dκ-1]; :(u[2])]
    :(SVector{N}(tuple($(qdot...))))
end

# Steering
struct DubinsConstraints{T} <: SteeringConstraints
    v_max::T
    r::T
end
struct ReedsSheppConstraints{T} <: SteeringConstraints
    v_max::T
    r::T
end
struct DubinsCCConstraints{T} <: SteeringConstraints
    v_max::T
    κ_max::T
    σ_max::T
end
const DubinsSteering{T} = SteeringBVP{SimpleCarDynamics{0,0},Time,DubinsConstraints{T}}
const ReedsSheppSteering{T} = SteeringBVP{SimpleCarDynamics{0,0},Time,ReedsSheppConstraints{T}}
const DubinsCCSteering{T} = SteeringBVP{SimpleCarDynamics{0,1},Time,DubinsCCConstraints{T}}
DubinsSteering(; v=1, r=1) = SteeringBVP(SimpleCarDynamics{0,0}(), Time(), constraints=DubinsConstraints(promote(v, r)...))
ReedsSheppSteering(; v=1, r=1) = SteeringBVP(SimpleCarDynamics{0,0}(), Time(), constraints=ReedsSheppConstraints(promote(v, r)...))
DubinsCCSteering(; v=1, κ_max=1, σ_max=1) = SteeringBVP(SimpleCarDynamics{0,1}(), Time(), constraints=DubinsCCConstraints(promote(v, κ_max, σ_max)...))
function (bvp::DubinsSteering)(q0::StaticVector{3}, qf::StaticVector{3})
    dubins(q0, qf, v=bvp.constraints.v_max, r=bvp.constraints.r)
end
function (bvp::ReedsSheppSteering)(q0::StaticVector{3}, qf::StaticVector{3})
    reedsshepp(q0, qf, v=bvp.constraints.v_max, r=bvp.constraints.r)
end
function (bvp::DubinsCCSteering)(q0::StaticVector{4}, qf::StaticVector{4})
    dubinsCC(q0, qf, v=bvp.constraints.v_max, κ_max=bvp.constraints.κ_max, σ_max=bvp.constraints.σ_max)
end

## Common Steering Utilities
@inline scaleradius(c::VelocityCurvatureStep, α) = StepControl(c.t*α, VelocityCurvatureControl(c.u.v, c.u.κ/α))
@inline scalespeed(c::VelocityCurvatureStep, λ)  = StepControl(c.t/λ, VelocityCurvatureControl(c.u.v*λ, c.u.κ))
@inline scalespeed(c::VelocityCurvRateStep, λ)   = StepControl(c.t/λ, c.u*λ)
@inline carsegment2stepcontrol(t::Int, d::T) where {T} = StepControl(abs(d), VelocityCurvatureControl(T(sign(d)), T(t)))
@inline timeflip(c::VelocityCurvatureStep) = StepControl(c.t, VelocityCurvatureControl(-c.u.v, c.u.κ))
@inline reflect(c::VelocityCurvatureStep)  = StepControl(c.t, VelocityCurvatureControl(c.u.v, -c.u.κ))
@inline reverse(ctrl::StaticVector{5}) = ctrl[SVector(5,4,3,2,1)]    # TODO: implement reverse in StaticArrays

SE2State(q::StaticVector{4}) = SE2State(q[1], q[2], q[3])
SE2κState(q::StaticVector{3,T}) where {T} = SE2κState(q[1], q[2], q[3], T(0))


text:
using ITensors
using Zygote

s = siteind("Qubit")

f(x) = op("Ry", s; θ=x)[1, 1]

x = 0.2
@show f(x), cos(x / 2)
@show f'(x), -sin(x / 2) / 2

# Simple gate optimization
ψ0 = state(s, "0")
ψp = state(s, "+")

function loss(x)
  U = op("Ry", s; θ=x)
  Uψ0 = replaceprime(U * ψ0, 1 => 0)
  return -(dag(ψp) * Uψ0)[]
end

# Extremely simple gradient descent implementation,
# where gradients are computing with automatic differentiation
# using Zygote.
function gradient_descent(f, x0; γ, nsteps, grad_tol)
  @show γ, nsteps
  x = x0
  f_x = f(x)
  ∇f_x = f'(x)
  step = 0
  @show step, x, f_x, ∇f_x
  for step in 1:nsteps
    x -= γ * ∇f_x
    f_x = f(x)
    ∇f_x = f'(x)
    @show step, x, f_x, ∇f_x
    if norm(∇f_x) ≤ grad_tol
      break
    end
  end
  return x, f_x, ∇f_x
end

x0 = 0
γ = 2.0 # Learning rate
nsteps = 30 # Number of steps of gradient descent
grad_tol = 1e-4 # Stop if gradient falls below this value
x, loss_x, ∇loss_x = gradient_descent(loss, x0; γ=γ, nsteps=nsteps, grad_tol=grad_tol)

@show x0, loss(x0)
@show x, loss(x)
@show π / 2, loss(π / 2)


text:
# ======================================================
# Testing continuous-time homogeneous system (no input)
# ======================================================
A = sparse([1, 1, 2, 3, 4], [1, 2, 2, 4, 3], [1., 2., 3., 4., 5.], 4, 4)
X0 = BallInf(zeros(4), 0.1)
cont_sys_homog = ContinuousSystem(A, X0)

# Check if the input is constant
@test inputdim(cont_sys_homog) == 0
# Check if the input is empty
#@test isa(next_set(cont_sys_homog.U), ZeroSet) # there is no input set
# Check data fields
@test cont_sys_homog.s.A == A
@test cont_sys_homog.x0.center == zeros(4) && cont_sys_homog.x0.radius == 0.1

# ===================================================
# Testing continuous-time system with constant input
# ===================================================
U = Ball2(ones(4), 0.5)
cont_sys = ContinuousSystem(A, X0, U)

# check initial state
@test cont_sys.x0.center ≈ zeros(4) && cont_sys.x0.radius ≈ 0.1

# recover input
inputs = next_set(cont_sys.s.U) #collect(nextinput(cont_sys.s.U))[1]

@test inputs.center ≈ ones(4) && inputs.radius ≈ 0.5

# ========================================================
# Testing continuous-time system with time-varying input
# ========================================================
Ui = [Ball2(0.01*i*ones(4), i*0.2) for i in 1:3]
cont_sys = ContinuousSystem(A, X0, Ui)

input_state = start(cont_sys.s.U)
inputs, input_state = next(cont_sys.s.U, input_state)
@test inputs.center ≈ 0.01*ones(4)
@test inputs.radius ≈ 0.2

inputs, input_state = next(cont_sys.s.U, input_state)
@test inputs.center ≈ 0.02*ones(4)
@test inputs.radius ≈ 0.4

inputs, input_state = next(cont_sys.s.U, input_state)
@test inputs.center ≈ 0.03*ones(4)
@test inputs.radius ≈ 0.6

# =========================================
# Testing discrete-time homogeneous system
# =========================================
δ = 0.01
discr_sys_homog = DiscreteSystem(A, X0)

# Check if the input is constant
#@test isa(discr_sys_homog.s.U, ConstantInput) # no field U
# Check if the input is empty
#@test isa(next_set(discr_sys_homog.s.U), ZeroSet)
# Check data fields
@test discr_sys_homog.s.A == A
@test discr_sys_homog.x0.center == zeros(4) && discr_sys_homog.x0.radius == 0.1

# =================================================
# Testing discrete-time system with constant input
# =================================================
U = Ball2(ones(4), 0.5)
discr_sys = DiscreteSystem(A, X0, U)

# check initial state
@test discr_sys.x0.center ≈ zeros(4) && discr_sys.x0.radius ≈ 0.1

# recover input
inputs = collect(nextinput(discr_sys.s.U))[1]

@test inputs.center ≈ ones(4) && inputs.radius ≈ 0.5

# =====================================================
# Testing discrete-time system with time-varying input
# =====================================================
Ui = [Ball2(0.01*i*ones(4), i*0.2) for i in 1:3]
discr_sys = DiscreteSystem(A, X0, Ui)
Uset = inputset(discr_sys.s)

input_state = start(Uset)
inputs, input_state = next(Uset, input_state)
@test inputs.center ≈ 0.01*ones(4)
@test inputs.radius ≈ 0.2

inputs, input_state = next(Uset, input_state)
@test inputs.center ≈ 0.02*ones(4)
@test inputs.radius ≈ 0.4

inputs, input_state = next(Uset, input_state)
@test inputs.center ≈ 0.03*ones(4)
@test inputs.radius ≈ 0.6


text:
module LearnBase

import StatsBase
using StatsBase: nobs

# AGGREGATION MODES
include("aggmode.jl")

# VIEW AND ITERATORS
include("iteration.jl")

# OBSERVATION DIMENSIONS
include("observation.jl")

# LEARNING COSTS (e.g. loss & penalty)
include("costs.jl")

# LABELS
include("labels.jl")

# OTHER CONCEPTS
include("other.jl")

end # module


text:
module DA_Nlayer_functionTests

using Test
using LightPropagation

# CW comparison of N-layer cylinder to semi-infinite
# note that for large layer thicknesses and radii many more besselroots are needed

### test at multiple SDS

# ρ = 0.0 cm
cylinder_data = Nlayer_cylinder(ρ = 0.0, a = 100.0, l = (10.0, 10.0, 10.0, 10.0), N_J0Roots = 10000)
@test fluence_DA_semiinf_CW(0.0, 0.1, 10.0) ≈ fluence_DA_Nlay_cylinder_CW(cylinder_data)

# ρ = 1.0 cm
cylinder_data = Nlayer_cylinder(ρ = 1.0, a = 100.0, l = (10.0, 10.0, 10.0, 10.0), N_J0Roots = 10000)
@test fluence_DA_semiinf_CW(1.0, 0.1, 10.0) ≈ fluence_DA_Nlay_cylinder_CW(cylinder_data)

# ρ = 4.0 cm
cylinder_data = Nlayer_cylinder(ρ = 4.0, a = 100.0, l = (10.0, 10.0, 10.0, 10.0), N_J0Roots = 10000)
@test fluence_DA_semiinf_CW(4.0, 0.1, 10.0) ≈ fluence_DA_Nlay_cylinder_CW(cylinder_data)

# ρ = 8.0 cm; for long SDS you need more besselroots or decrease the dimensions of cylinder
cylinder_data = Nlayer_cylinder(ρ = 8.0, a = 100.0, l = (10.0, 10.0, 10.0, 10.0),N_J0Roots = 10000)
@test isapprox(fluence_DA_semiinf_CW(8.0, 0.1, 10.0), fluence_DA_Nlay_cylinder_CW(cylinder_data), rtol=1e-6)

### test that 2, 3, 4 layers equal with same optical properties
cylinder_4 = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0, 10.0, 10.0), μa = (0.1, 0.1, 0.1, 0.1), l = (10.0, 10.0, 10.0, 10.0), n_med = (1.0, 1.0, 1.0, 1.0), N_J0Roots = 10000)
cylinder_3 = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0, 10.0), μa = (0.1, 0.1, 0.1), l = (10.0, 10.0, 10.0), n_med = (1.0, 1.0, 1.0), N_J0Roots = 10000)
cylinder_2 = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0), μa = (0.1, 0.1), l = (10.0, 10.0), n_med = (1.0, 1.0), N_J0Roots = 10000)
@test fluence_DA_Nlay_cylinder_CW(cylinder_4) ≈  fluence_DA_Nlay_cylinder_CW(cylinder_3)  ≈ fluence_DA_Nlay_cylinder_CW(cylinder_2)

### test that arbitrary layers are equal
cylinder_5 = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0, 10.0, 10.0, 10.0), μa = (0.1, 0.1, 0.1, 0.1, 0.1), l = (2.0, 2.0, 2.0, 5.0, 5.0), n_med = (1.0, 1.0, 1.0, 1.0, 1.0), N_J0Roots = 10000)
cylinder_8 = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0), μa = (0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1), l = (2.0, 2.0, 2.0, 5.0, 5.0, 3.0, 3.0, 3.0), n_med = (1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0), N_J0Roots = 10000)

@test fluence_DA_Nlay_cylinder_CW(cylinder_4) ≈  fluence_DA_Nlay_cylinder_CW(cylinder_5) ≈  fluence_DA_Nlay_cylinder_CW(cylinder_8)
### test high scatterring
# SDS = 1.0, μsp = 60.0, once again need more besselroots if you want higher scattering coefficients
cylinder_data = Nlayer_cylinder(ρ = 1.0, μsp = (60.0, 60.0, 60.0, 60.0), μa = (0.1, 0.1, 0.1, 0.1), l = (2.0, 2.0, 2.0, 2.0), n_med = (1.0, 1.0, 1.0, 1.0), a = 20.0, N_J0Roots = 10000)
@test fluence_DA_semiinf_CW(1.0, 0.1, 60.0) ≈ fluence_DA_Nlay_cylinder_CW(cylinder_data)

### test for varying index of refractions
cyl_2 = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0), μa = (0.1, 0.1), l = (10.0, 10.0), n_med = (1.5, 1.5), n_ext = 1.4, a = 20.0, N_J0Roots = 10000)
cyl_3 = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0, 10.0), μa = (0.1, 0.1, 0.1), l = (10.0, 10.0, 10.0), n_med = (1.5, 1.5, 1.5), n_ext = 1.4, a = 20.0, N_J0Roots = 10000)
cyl_4 = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0, 10.0, 10.0), μa = (0.1, 0.1, 0.1, 0.1), l = (10.0, 10.0, 10.0, 10.0), n_med = (1.5, 1.5, 1.5, 1.5), n_ext = 1.4, a = 20.0, N_J0Roots = 10000)
@test fluence_DA_Nlay_cylinder_CW(cyl_4) ≈ fluence_DA_Nlay_cylinder_CW(cyl_3) ≈ fluence_DA_Nlay_cylinder_CW(cyl_2) ≈ fluence_DA_semiinf_CW(1.0, 0.1, 10.0, n_med = 1.5, n_ext = 1.4)

#### tests for CW fluence in the bottom layer GN
@test fluence_DA_Nlay_cylinder_CW(0.0, (0.1, 0.1), (10.0, 10.0), 1.0, (1.0, 1.0), (1.0, 1.0), 10.0, 2.0, 10000) ≈ fluence_DA_slab_CW(0.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0, s = 2.0, z = 2.0, xs = 50)
@test fluence_DA_Nlay_cylinder_CW(1.0, (0.1, 0.1), (10.0, 10.0), 1.0, (1.0, 1.0), (1.0, 1.0), 10.0, 2.0, 10000) ≈ fluence_DA_slab_CW(1.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0, s = 2.0, z = 2.0, xs = 50)
@test fluence_DA_Nlay_cylinder_CW(1.6, (0.21, 0.21), (12.1, 12.1), 1.0, (1.0, 1.0), (1.0, 1.0), 10.0, 2.0, 10000) ≈ fluence_DA_slab_CW(1.6, 0.21, 12.1; n_ext = 1.0, n_med = 1.0, s = 2.0, z = 2.0, xs = 50)
@test fluence_DA_Nlay_cylinder_CW(2.2, (0.31, 0.31, 0.31), (12.1, 12.1, 12.1), 1.0, (1.0, 1.0, 1.0), (1.0, 1.0, 1.0), 10.0, 3.0, 10000) ≈ fluence_DA_slab_CW(2.2, 0.31, 12.1; n_ext = 1.0, n_med = 1.0, s = 3.0, z = 3.0, xs = 50)
@test fluence_DA_Nlay_cylinder_CW(2.8, (0.41, 0.41, 0.41, 0.41), (13.1, 13.1, 13.1, 13.1), 1.0, (1.0, 1.0, 1.0, 1.0), (1.0, 1.0, 1.0, 1.0), 10.0, 4.0, 10000) ≈ fluence_DA_slab_CW(2.8, 0.41, 13.1; n_ext = 1.0, n_med = 1.0, s = 4.0, z = 4.0, xs = 50)

@test fluence_DA_Nlay_cylinder_CW(0.0, (0.1, 0.1), (10.0, 10.0), 1.0, (1.0, 1.0), (1.5, 1.5), 10.0, 3.0, 10000) ≈ fluence_DA_slab_CW(0.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0, s = 3.0, z = 3.0, xs = 50)

@test fluence_DA_Nlay_cylinder_CW(0.0, (0.1, 0.1, 0.1, 0.1, 0.1, 0.1), (10.0, 10.0, 10.0, 10.0, 10.0, 10.0), 1.0, (1.0, 1.0, 1.0, 1.0, 1.0, 1.0), (0.5, 0.5, 0.5, 0.5, 0.5, 0.5), 10.0, 3.0, 10000) ≈ fluence_DA_slab_CW(0.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0, s = 3.0, z = 3.0, xs = 50)
@test fluence_DA_Nlay_cylinder_CW(0.0, (0.1, 0.1, 0.1, 0.1, 0.1, 0.1), (10.0, 10.0, 10.0, 10.0, 10.0, 10.0), 1.0, (1.0, 1.0, 1.0, 1.0, 1.0, 1.0), (1.0, 1.0, 1.0, 1.0, 1.0, 1.0), 10.0, 6.0, 10000) ≈ fluence_DA_slab_CW(0.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0, s = 6.0, z = 6.0, xs = 50)

### test that scalar and vector ρ matches (they call different routines)
μsp = (10.0, 10.0, 10.0, 10.0); μa = (0.1, 0.1, 0.1, 0.1); n_ext = 1.0;
n_med = (1.0, 1.0, 1.0, 1.0); l = (0.5, 0.8, 1.0, 5.0); a = 5.0; z = 0.0
N_J0Roots = 1000
ρ1 = (1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0)
@test [fluence_DA_Nlay_cylinder_CW(ρ1, μa, μsp, n_ext, n_med, l, a, z, N_J0Roots)...] ≈ map(ρ -> fluence_DA_Nlay_cylinder_CW(ρ, μa, μsp, n_ext, n_med, l, a, z, N_J0Roots), 1.0:0.1:2.0)

#########
# Time-Domain
##########

## test reflectance in first layer
# test long times at SDS = 0.5 cm
t = range(0.03, 8.0, length = 60)
cylinder_data = Nlayer_cylinder(a = 8.0, ρ = 0.5, N_J0Roots = 600)

si = fluence_DA_semiinf_TD(t, 0.5, 0.1, 10.0)

a72 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 72)
a96 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 96)

@test a72 ≈ si
@test a96 ≈ si

# test long times at SDS = 3.0
t = range(0.04, 8.0, length = 60)
cylinder_data = Nlayer_cylinder(a = 8.0, ρ = 3.0, l = (4.0, 4.0, 4.0, 5.0), N_J0Roots = 600)

si = fluence_DA_semiinf_TD(t, 3.0, 0.1, 10.0)

a72 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 72)
a96 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 96)

@test a72 ≈ si
@test a96 ≈ si

# test high scattering at SDS = 1.5 cm
t = range(0.04, 8.0, length = 60)
cylinder_data = Nlayer_cylinder(a = 8.0, ρ = 1.5, l = (4.0, 4.0, 4.0, 5.0), μsp = (60.0, 60.0, 60.0, 60.0), N_J0Roots = 600)

si = fluence_DA_semiinf_TD(t, 1.5, 0.1, 60.0)

a72 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 72)
a96 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 96)

@test a72 ≈ si
@test a96 ≈ si

## in general the error in the TD is like 1e-14..

## test transmittance in bottom layer
t = range(0.03, 8.0, length = 60)
cylinder_data = Nlayer_cylinder(μa = (0.1, 0.1), μsp = (10.0, 10.0), n_ext = 1.0, n_med = (1.0, 1.0), l = (1.0, 1.0), a = 8.0, ρ = 0.5, z = 2.0, N_J0Roots = 600)

slab = fluence_DA_slab_TD(t, 0.5, 0.1, 10.0; s = 2.0, z = 2.0, xs = 50)

a72 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 72)
a96 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 96)

@test a72 ≈ slab
@test a96 ≈ slab

t = range(0.03, 8.0, length = 60)
cylinder_data = Nlayer_cylinder(μa = (0.1, 0.1), μsp = (10.0, 10.0), n_ext = 1.0, n_med = (1.0, 1.0), l = (1.5, 1.5), a = 8.0, ρ = 0.5, z = 3.0, N_J0Roots = 600)

slab = fluence_DA_slab_TD(t, 0.5, 0.1, 10.0; s = 3.0, z = 3.0, xs = 50)

a72 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 72)
a96 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 96)

@test a72 ≈ slab
@test a96 ≈ slab

t = range(0.03, 8.0, length = 60)
cylinder_data = Nlayer_cylinder(μa = (0.1, 0.1), μsp = (10.0, 10.0), n_ext = 1.0, n_med = (1.0, 1.0), l = (1.5, 1.5), a = 8.0, ρ = 1.5, z = 3.0, N_J0Roots = 600)

slab = fluence_DA_slab_TD(t, 1.5, 0.1, 10.0; s = 3.0, z = 3.0, xs = 50)

a72 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 72)
a96 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 96)

@test a72 ≈ slab
@test a96 ≈ slab

t = range(0.03, 8.0, length = 60)
cylinder_data = Nlayer_cylinder(μa = (0.1, 0.1), μsp = (10.0, 10.0), n_ext = 1.3, n_med = (1.2, 1.2), l = (1.0, 1.0), a = 8.0, ρ = 0.5, z = 2.0, N_J0Roots = 600)

slab = fluence_DA_slab_TD(t, 0.5, 0.1, 10.0; s = 2.0, z = 2.0, xs = 50, n_ext = 1.3, n_med = 1.2)

a72 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 72)
a96 = fluence_DA_Nlay_cylinder_TD(t, cylinder_data, N = 96)

@test a72 ≈ slab
@test a96 ≈ slab

### Flux tests

# CW
cylinder_data = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0), μa = (0.1, 0.1), l = (10.0, 10.0), n_med = (1.0, 1.0), n_ext = 1.0, a = 40.0, N_J0Roots = 10000)
@test flux_DA_Nlay_cylinder_CW(cylinder_data) ≈ flux_DA_semiinf_CW(1.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0)

cylinder_data = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0), μa = (0.1, 0.1), l = (2.0, 2.0), z = 4.0, n_med = (1.0, 1.0), n_ext = 1.0, a = 40.0, N_J0Roots = 10000)
@test flux_DA_Nlay_cylinder_CW(cylinder_data) ≈ flux_DA_slab_CW(1.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0, s = 4.0, z = 4.0, xs = 15)

# TD
t = 1.0
cylinder_data = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0), μa = (0.1, 0.1), l = (10.0, 10.0), n_med = (1.0, 1.0), n_ext = 1.0, a = 40.0, N_J0Roots = 10000)
@test flux_DA_Nlay_cylinder_TD(t, cylinder_data) ≈ flux_DA_semiinf_TD(t, 1.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0)

cylinder_data = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0), μa = (0.1, 0.1), l = (2.0, 2.0), z = 4.0, n_med = (1.0, 1.0), n_ext = 1.0, a = 40.0, N_J0Roots = 10000)
@test flux_DA_Nlay_cylinder_TD(t, cylinder_data) ≈ flux_DA_slab_TD(t, 1.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0, s = 4.0, z = 4.0, xs = 15)

t = 0.5:0.1:1.5
cylinder_data = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0), μa = (0.1, 0.1), l = (10.0, 10.0), n_med = (1.0, 1.0), n_ext = 1.0, a = 40.0, N_J0Roots = 10000)
@test flux_DA_Nlay_cylinder_TD(t, cylinder_data) ≈ flux_DA_semiinf_TD(t, 1.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0)

cylinder_data = Nlayer_cylinder(ρ = 1.0, μsp = (10.0, 10.0), μa = (0.1, 0.1), l = (2.0, 2.0), z = 4.0, n_med = (1.0, 1.0), n_ext = 1.0, a = 40.0, N_J0Roots = 10000)
@test flux_DA_Nlay_cylinder_TD(t, cylinder_data) ≈ flux_DA_slab_TD(t, 1.0, 0.1, 10.0; n_ext = 1.0, n_med = 1.0, s = 4.0, z = 4.0, xs = 15)

end # module


text:
#!/bin/bash
# -*- mode: julia -*-
#=
JULIA="${JULIA:-julia}"
HERE="$(dirname "${BASH_SOURCE[0]}")"
JULIA_CMD="${JULIA_CMD:-$JULIA --color=yes --startup-file=no --project="$HERE"}"
export JULIA_LOAD_PATH=@:@stdlib  # exclude default environment
exec $JULIA_CMD "${BASH_SOURCE[0]}" "$@"
=#

using Pkg

Pkg.develop([
    PackageSpec(
        name = "Try",
        path = dirname(@__DIR__),
        # url = "https://github.com/JuliaPreludes/Try.jl.git",
    ),
    PackageSpec(
        name = "TryExperimental",
        path = dirname(@__DIR__),
        # url = "https://github.com/JuliaPreludes/Try.jl.git",
        subdir = "lib/TryExperimental",
    ),
])

Pkg.instantiate()


text:
# Retrieve filename of literate script
if length(ARGS) != 2
    error("please specify the literate script and the output directory")
end
const SCRIPTJL = ARGS[1]
const OUTDIR = ARGS[2]

# Activate environment
using Pkg: Pkg
Pkg.activate(dirname(SCRIPTJL))
Pkg.instantiate()
using Literate: Literate

# Add link to nbviewer below the first heading of level 1
function preprocess(content)
    sub = SubstitutionString(
        """
#md # ```@meta
#md # EditURL = "@__REPO_ROOT_URL__/examples/@__NAME__/script.jl"
#md # ```
#md #
\\0
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/examples/@__NAME__.ipynb)
#md #
# You are seeing the
#md # HTML output generated by [Documenter.jl](https://github.com/JuliaDocs/Documenter.jl) and
#nb # notebook output generated by
# [Literate.jl](https://github.com/fredrikekre/Literate.jl) from the
# [Julia source file](@__REPO_ROOT_URL__/examples/@__NAME__/script.jl).
# The corresponding
#md # notebook can be viewed in [nbviewer](@__NBVIEWER_ROOT_URL__/examples/@__NAME__.ipynb).
#nb # HTML output can be viewed [here](https://juliagaussianprocesses.github.io/KernelFunctions.jl/dev/examples/@__NAME__/).
#
        """,
    )
    return replace(content, r"^# # [^\n]*"m => sub; count=1)
end

# Convert to markdown and notebook
Literate.markdown(
    SCRIPTJL,
    OUTDIR;
    name=basename(dirname(SCRIPTJL)),
    documenter=false,
    execute=true,
    preprocess=preprocess,
)
Literate.notebook(
    SCRIPTJL,
    OUTDIR;
    name=basename(dirname(SCRIPTJL)),
    documenter=false,
    execute=true,
    preprocess=preprocess,
)


text:
module Turing

##############
# Dependency #
########################################################################
# NOTE: when using anything from external packages,                    #
#       let's keep the practice of explictly writing Package.something #
#       to indicate that's not implemented inside Turing.jl            #
########################################################################

using Requires
using Reexport
@reexport using Distributions
import Flux
using ForwardDiff
using Bijectors
@reexport using MCMCChain
using StatsFuns
using SpecialFunctions
using Statistics
using LinearAlgebra
using ProgressMeter
using Markdown
using Libtask
using MacroTools

function __init__()
    @require CmdStan="593b3428-ca2f-500c-ae53-031589ec8ddd" @eval begin
        using CmdStan
        import CmdStan: Adapt, Hmc
    end

    @require DynamicHMC="bbc10e6e-7c05-544b-b16e-64fede858acb" @eval begin
        using DynamicHMC, LogDensityProblems
        using LogDensityProblems: AbstractLogDensityProblem, ValueGradient

        struct FunctionLogDensity{F} <: AbstractLogDensityProblem
          dimension::Int
          f::F
        end

        LogDensityProblems.dimension(ℓ::FunctionLogDensity) = ℓ.dimension

        LogDensityProblems.logdensity(::Type{ValueGradient}, ℓ::FunctionLogDensity, x) = ℓ.f(x)::ValueGradient
    end
end
import Base: ~, convert, promote_rule, rand, getindex, setindex!
import Distributions: sample
import ForwardDiff: gradient
using Flux.Tracker
import MCMCChain: AbstractChains, Chains

##############################
# Global variables/constants #
##############################

const ADBACKEND = Ref(:forward_diff)
setadbackend(backend_sym) = begin
  @assert backend_sym == :forward_diff || backend_sym == :reverse_diff
  backend_sym == :forward_diff && CHUNKSIZE[] == 0 && setchunksize(40)
  ADBACKEND[] = backend_sym
end

const ADSAFE = Ref(false)
setadsafe(switch::Bool) = begin
  @info("[Turing]: global ADSAFE is set as $switch")
  ADSAFE[] = switch
end

const CHUNKSIZE = Ref(40) # default chunksize used by AD

setchunksize(chunk_size::Int) = begin
  if ~(CHUNKSIZE[] == chunk_size)
    @info("[Turing]: AD chunk size is set as $chunk_size")
    CHUNKSIZE[] = chunk_size
  end
end

const PROGRESS = Ref(true)
turnprogress(switch::Bool) = begin
  @info("[Turing]: global PROGRESS is set as $switch")
  PROGRESS[] = switch
end

# Constants for caching
const CACHERESET  = 0b00
const CACHEIDCS   = 0b10
const CACHERANGES = 0b01

#######################
# Sampler abstraction #
#######################
abstract type AbstractAdapter end
abstract type InferenceAlgorithm end
abstract type Hamiltonian{AD} <: InferenceAlgorithm end
abstract type StaticHamiltonian{AD} <: Hamiltonian{AD} end
abstract type AdaptiveHamiltonian{AD} <: Hamiltonian{AD} end
abstract type AbstractSampler end
"""
    Sampler{T}

Generic interface for implementing inference algorithms.
An implementation of an algorithm should include the following:

1. A type specifying the algorithm and its parameters, derived from InferenceAlgorithm
2. A method of `sample` function that produces results of inference, which is where actual inference happens.

Turing translates models to chunks that call the modelling functions at specified points. The dispatch is based on the value of a `sampler` variable. To include a new inference algorithm implements the requirements mentioned above in a separate file,
then include that file at the end of this one.
"""
mutable struct Sampler{T<:InferenceAlgorithm} <: AbstractSampler
  alg   ::  T
  info  ::  Dict{Symbol, Any}         # sampler infomation
end
Sampler(alg, model) = Sampler(alg)

# mutable struct HMCState{T<:Real}
#     epsilon  :: T
#     std     :: Vector{T}
#     lf_num   :: Integer
#     eval_num :: Integer
# end
#
#  struct Sampler{TH<:Hamiltonian,TA<:AbstractAdapter} <: AbstractSampler
#    alg   :: TH
#    state :: HMCState
#    adapt :: TA
#  end

"""
Robust initialization method for model parameters in Hamiltonian samplers.
"""
struct HamiltonianRobustInit <: AbstractSampler end
struct SampleFromPrior <: AbstractSampler end

# This can be removed when all `spl=nothing` is replaced with
#   `spl=SampleFromPrior`
const AnySampler = Union{Nothing, AbstractSampler}

include("utilities/resample.jl")
@static if isdefined(Turing, :CmdStan)
    include("utilities/stan-interface.jl")
end
include("utilities/helper.jl")
include("utilities/robustinit.jl")
include("utilities/util.jl")         # utility functions
include("utilities/io.jl")           # I/O
include("models/distributions.jl")
include("core/varinfo.jl")  # core internal variable container
include("core/trace.jl")   # to run probabilistic programs as tasks

using Turing.VarReplay

###########
# Exports #
###########

# Turing essentials - modelling macros and inference algorithms
export @model, @~, @VarName                   # modelling
export MH, Gibbs                              # classic sampling
export HMC, SGLD, SGHMC, HMCDA, NUTS          # Hamiltonian-like sampling
export DynamicNUTS
export IS, SMC, CSMC, PG, PIMH, PMMH, IPMCMC  # particle-based sampling
export sample, setchunksize, resume           # inference
export auto_tune_chunk_size!, setadbackend, setadsafe # helper
export turnprogress  # debugging
export consume, produce

# Turing-safe data structures and associated functions
export TArray, tzeros, localcopy, IArray

export @sym_str

export Flat, FlatPos, BinomialLogit, VecBinomialLogit

##################
# Inference code #
##################

include("core/compiler.jl")     # compiler
include("core/container.jl")    # particle container
include("samplers/sampler.jl")  # samplers
include("core/ad.jl")           # Automatic Differentiation

end


text:
### A Pluto.jl notebook ###
# v0.14.8

using Markdown
using InteractiveUtils

# ╔═╡ 6e3ef066-d115-11eb-2338-013a707dfe8a
begin
	using PlutoUI
	using DataFrames
	using Statistics
	using HTTP
	using JSON
	using Chain
	using Tables
	using MLStyle
end

# ╔═╡ 41bd06a2-e7bd-46c6-9249-6e69223b0e11
HASRUA_ENDPOINT_URL = "http://rldl12.eecs.umich.edu:1200/v1/graphql"

# ╔═╡ 94185c13-b6d2-4337-b5ce-336f5e128032
function gql_query(query:: String; variables:: Dict = nothing)
	r = HTTP.request(
		"POST",
		HASRUA_ENDPOINT_URL;
		verbose=3,
		body= JSON.json(Dict("query" => query, "variables" => variables)) ,
		# headers=Dict("x-hasura-admin-secret" => HASURA_ADMIN_SECRET),
	)
	@chain r begin
		_.body
		String
		JSON.parse(_, null=missing)
	    _["data"]
	end
end

# ╔═╡ ed56319c-fd7e-478e-a673-f799debbf7b3
function get_sweep(sweep_ids::AbstractVector{Int}, max_step::Int)
	query = """
		query getSweepRuns(\$ids: [Int!], \$max_step: Int!) {
		  logs_less_than_step(args: {max_step: \$max_step}, where: {run: {sweep_id: {_in: \$ids}}}) {
			log
			run_id
			run {
			  metadata
			  sweep_id
			}
		  }
		}
  	"""
	@chain query begin
		gql_query(_; variables=Dict("ids" => sweep_ids, "max_step" => max_step))
		_["logs_less_than_step"]
		map(_) do d
			Dict(
				"run_id" => d["run_id"],
				"sweep_id" => d["run"]["sweep_id"],
				d["log"]...,
				d["run"]["metadata"]["parameters"]...
			)
		end
		map(_) do d
			Dict(
				d...,
				[k => v for (k1, v1, k2, v2) in [
							(
								"hours", get(d, "time-delta", 0) / 3600,
								"time-delta", get(d, "hours", 0) * 3600,
							),
						]
						for (k, v) in [
								(k1, get(d, k1, v1)),
								(k2, get(d, k2, v2)),
								]]...,
				[name => get(d, name, false) for name in [
							"randomize_parameters"
						]]...,
				[name => get(d, name, nothing) for name in [
							"config",
						]]...
			)
		end
		collect
		vcat(DataFrame.(_)...)
	end
end;

# ╔═╡ da091e45-0d08-490a-ba85-ee9acb6700bc
sweeps = filter(get_sweep([784], 10000000)) do row
				row.env == "Seaquest-v0"
			end

# ╔═╡ 1a6b7810-977d-4a82-ab14-b24e154ad491
EPISODE_RETURN = "episode return"

# ╔═╡ 7e2ac3d5-9b72-4859-96a7-b6b8711e14b3
min_returns = @chain sweeps begin
	dropmissing(_, EPISODE_RETURN)
	groupby(_, [:env])
	combine(_, EPISODE_RETURN => minimum)
	Dict(k=>v for (k,v) in eachrow(_))
end

# ╔═╡ 5a29110f-deca-4812-9240-ab445ed665c8
max_returns = Dict(
	"BeamRider-v0" => 1590,
	"PongNoFrameskip-v0" => 20.7,
	"Seaquest-v0" => 1204.5,
	"Qbert-v0" => 14293.3
) # from PPO paper

# ╔═╡ c838c44c-4bd8-4eb0-a18b-ceaaa5bdce93
dframe = @chain sweeps begin
	filter(:step => >=(8000000), _)
	groupby(_, [:env])
	transform(_, ["env", EPISODE_RETURN] =>
		function (envs, ret)
			@match [Set(envs)...] (
				[env] => (ret .- min_returns[env]) ./ max_returns[env] 							)
		end => [EPISODE_RETURN])
	_[!, filter(names(_)) do name
			!(name in [
				"action loss",
				"config",
				"cuda",
				"entropy",
				"eval_interval",
				"fps",
				"gradient norm",
				"hours",
				"log_interval",
				"log_level",
				"num_env_steps",
				"recurrent_policy",
				"run_id",
				"save_interval",
				"save_path",
				"step",
				"subcommand",
				"sweep_id",
				"time",
				"time-delta",
				"value loss",
			])
		end]
end

# ╔═╡ 2735fbc7-a318-4b4e-8ceb-0e65084cf972
function filter_by_type(ty)
	filter(name -> eltype(dframe[:, name]) == ty, names(dframe))
end

# ╔═╡ 1fe4e26a-2786-4094-a7e5-0154b461e874
df = @chain dframe begin
	groupby(_, "run ID")
	combine(_,
		filter_by_type(Float64) .=> first,
		filter_by_type(Int64) .=> first,
		filter_by_type(Bool) .=> first,
		EPISODE_RETURN .=> mean .=> :episode_return_mean,
		)
	_[!, filter(n -> !(n in  ["run ID", "episode return_first"]), names(_))]
	sort!(_, [:episode_return_mean], rev=true)
	rename(name -> replace(name, "_first" => ""), _)
end

# ╔═╡ 490d6756-f769-4a13-9438-ad724aa013ee
begin
	bool_df = DataFrame()
	for name in names(df)
		if !occursin("episode_return", name)
			for value in df[:, name]
				new_name = string(replace(name, "_first"=>""), " = ", value)
				bool_df[:, :($new_name)] = df[:, name] .== value
			end
		end
	end
    bool_df.episode_return_mean = df.episode_return_mean
	bool_df
end

# ╔═╡ c3e89733-4a83-48d4-bf34-9f4f7895b798
function get_cor_df(df)

	cor_df = DataFrame()

	cor_df.name = names(df)

	cor_mat = cor(Matrix(df))
	return_index = findfirst(name -> name == "episode_return_mean", names(df))
	cor_df.correlation = cor_mat[:, return_index]

	sort!(cor_df, [:correlation], )
	cor_df
end

# ╔═╡ f43f3cf4-b9e4-4c21-bf95-96a0c494c99b
md"# Continuous Correlation"

# ╔═╡ a4cb4dbe-b880-4a4a-9b91-ccf6eaff1b07
get_cor_df(df)

# ╔═╡ 5fb3a66b-49a2-41f6-aea5-65a2bd4d7240
md"# Discrete Correlation"

# ╔═╡ 6c27eed0-e73d-4a44-b286-dca6ac404dbd
get_cor_df(bool_df)

# ╔═╡ 8b15faf7-0d54-481f-a513-311e963dd8c1
Dict(k => unique(df, k)[!, k] for k in names(df))

# ╔═╡ Cell order:
# ╠═6e3ef066-d115-11eb-2338-013a707dfe8a
# ╠═41bd06a2-e7bd-46c6-9249-6e69223b0e11
# ╠═94185c13-b6d2-4337-b5ce-336f5e128032
# ╠═ed56319c-fd7e-478e-a673-f799debbf7b3
# ╠═da091e45-0d08-490a-ba85-ee9acb6700bc
# ╠═1a6b7810-977d-4a82-ab14-b24e154ad491
# ╠═7e2ac3d5-9b72-4859-96a7-b6b8711e14b3
# ╠═5a29110f-deca-4812-9240-ab445ed665c8
# ╠═2735fbc7-a318-4b4e-8ceb-0e65084cf972
# ╠═c838c44c-4bd8-4eb0-a18b-ceaaa5bdce93
# ╠═1fe4e26a-2786-4094-a7e5-0154b461e874
# ╠═490d6756-f769-4a13-9438-ad724aa013ee
# ╟─c3e89733-4a83-48d4-bf34-9f4f7895b798
# ╟─f43f3cf4-b9e4-4c21-bf95-96a0c494c99b
# ╠═a4cb4dbe-b880-4a4a-9b91-ccf6eaff1b07
# ╟─5fb3a66b-49a2-41f6-aea5-65a2bd4d7240
# ╟─6c27eed0-e73d-4a44-b286-dca6ac404dbd
# ╠═8b15faf7-0d54-481f-a513-311e963dd8c1

max_stars_repo_name: GPT-RL/generalization

max_issues_repo_name: GPT-RL/generalization

max_forks_repo_name: GPT-RL/generalization


text:
"""
`ellipse_parameters(dataset, index::Int)`

Gives the entire set of ellipse parameters for a dataset at a given index.
"""
ellipse_parameters(dataset, index::Int) = dataset[index,:]

"""
`ellipse_a(dataset, index::Int)`

Gives the major axis parameter `a` for a dataset at a given index.
"""
function ellipse_a(dataset, index::Int)
    dataset[index,end-1]
end

"""
`ellipse_b(dataset, index::Int)`

Gives the minor axis parameter `b` for a dataset at a given index.
"""
function ellipse_b(dataset, index::Int)
    dataset[index,end]
end

"""
`ellipse_b(dataset, index::Int)`

Gives the axis tilt parameter `θ` for a dataset at a given index.
"""
function ellipse_θ(dataset, index::Int)
    dataset[index,end-2]
end

"""
`ellipse_center(dataset, index::Int)`

Gives the pair of coordinates `x0, y0` of the ellipse center for a dataset at a given index.
"""
function ellipse_center(dataset, index::Int)
    dataset[index,1], dataset[index,2]
end

function ellipse_l0(dataset, index::Int)
    if dataset in (WRB_ellipses, DLM_ellipses, B_ellipses)
        dataset[index,3]
    else
        error("Dataset does not contain info about l0.")
    end
end

"""
`ellipse_g11(dataset, index::Int)`

Gives the ellipse parameter `g11` at a given index.
"""
function ellipse_g11(dataset, index::Int)
    # θ is the tilt of the ellipse
    θ = ellipse_θ(dataset, index)
    cosd(θ)^2/ellipse_a(dataset, index) + sind(θ)^2/ellipse_b(dataset, index)^2
end

"""
`ellipse_g12(dataset, index::Int)`

Gives the ellipse parameter `g12` at a given index.
"""
function ellipse_g12(dataset, index::Int)
    θ = ellipse_θ(dataset, index)
    sind(θ) * cosd(θ) * (1.0/ellipse_a(dataset, index)^2 - 1/ellipse_b(dataset, index)^2)
end

"""
`ellipse_g22(dataset, index::Int)`

Gives the ellipse parameter `g22` at a given index.
"""
function ellipse_g22(dataset, index::Int)
    θ = ellipse_θ(dataset, index)
    sind(θ)^2/ellipse_a(dataset, index)^2 + cosd(θ)/ellipse_b(dataset, index)^2
end

"""
`ellipse_g23(dataset, index::Int)`

Gives the ellipse parameter `g23` at a given index.
"""
function ellipse_g23(dataset, index::Int)

end

"""
`ellipse_g33(dataset, index::Int)`

Gives the ellipse parameter `g33` at a given index.
"""
function ellispe_g33(dataset, index::Int)

end

"""
`ellipse_g13(dataset, index::Int)`

Gives the ellipse parameter `g13` at a given index.
"""
function ellispe_g13(dataset, index::Int)

end

"""
`ellipse_luminance(dataset,index::Int)`

Gives the luminance in cd/m² of a given test color
"""
function ellipse_luminance(dataset, index::Int)
    if dataset in (WRB_ellipses, DLM_ellipses, B_ellipses)
        10^(5 * ellipse_l0(dataset, index))
    else error("Dataset $dataset does not contain l0/luminance information for test colors")
    end
end

"""
`draw_ellipses(dataset)`

Draws all ellipses of a chosen dataset.
"""
function draw_ellipses(dataset,sizefactor)
    ellcoll = []
    @inbounds for ind = 1:size(dataset)[1]
        a = ellipse_a(dataset, ind)
        b = ellipse_b(dataset, ind)
        θ = ellipse_θ(dataset, ind)
        x0, y0 = ellipse_center(dataset, ind)
        ell = ellipse(x0, y0, sizefactor*a, sizefactor*b, θ/180 * π)
        push!(ellcoll, ell)
    end
    ellcoll
end

# 3d ellipse
#  ds^2 = g11(x-x0)^2 + g22(y-y0)^2 + g33(l-l0)^2
# +2 * g12(x-x0)(y-y0) + 2 * g13(x-x0)(l-l0) + 2 * g23(y-y0)(l-l0)
# with ds^2 = 7.81, corresponding to χ^2 value for 3 degrees of freedom,
# ellipsoid contains on average 95% of any random set of color matches

"""
`ellipse(u, v, a, b, theta)`

Create coordinate vectors of an ellipse with center coordinates `u`, `v`, major axis `a`, minor axis `b`, and tilt angle `θ`.

Result: `x::Vector, y::vector`
"""
function ellipse(u, v, a, b, theta)
    n = 40
    x = zeros(n+1)
    y = zeros(n+1)
    @inbounds for i in 1:1:n+1
        t = i * (2π)/n
        e = a * sin(t)
        f = b * cos(t)
        x[i] = e * cos(theta) - f * sin(theta) + u
        y[i] = e * sin(theta) + f * cos(theta) + v
    end
    x, y
end


text:
using MaGe, BenchmarkTools
using Statistics: mean

const eventpath = realpath(joinpath(dirname(pathof(MaGe)), "..", "test", "test.root.hits"))

SUITE = BenchmarkGroup()

SUITE["read_event"] = @benchmarkable read(MaGe.loadstreaming($eventpath), Event{Vector{Hit}})

SUITE["load_file"] = @benchmarkable MaGe.load($eventpath; T=Event{Vector{Hit}})

e = MaGe.loadstreaming(first, eventpath)
SUITE["energy"] = @benchmarkable energy($e)

SUITE["mean_energy"] = @benchmarkable MaGe.loadstreaming($eventpath) do es
    energies = Vector{Float64}()
    while !eof(es)
        e = read(es, Event{MaGe.RootHitIter})
        push!(energies, energy(e))
    end
    return mean(energies)
end

getfirstloc(e::MaGe.AbstractEvent) = location(first(hits(e)))
SUITE["collect_first_locs"] = @benchmarkable map(getfirstloc, MaGe.loadstreaming($eventpath))


text:
# functions to try to debug cui3.out
# converted to test function

PhyloNetworks.CHECKNET || error("need CHECKNET==true in PhyloNetworks to test snaq in test_correctLik.jl")

@testset "test: level-1 partitions" begin

global net, flag, flag2, flag3, nocycle
text = "(Xmayae,((Xhellerii,(((Xclemenciae_F2,Xmonticolus):1.458,(((((Xmontezumae,(Xnezahuacoyotl)#H26:0.247::0.804):0.375,((Xbirchmanni_GARC,Xmalinche_CHIC2):0.997,Xcortezi):0.455):0.63,(#H26:0.0::0.196,((Xcontinens,Xpygmaeus):1.932,(Xnigrensis,Xmultilineatus):1.401):0.042):2.439):2.0)#H7:0.787::0.835,(Xmaculatus,(Xandersi,(Xmilleri,((Xxiphidium,#H7:9.563::0.165):1.409,(Xevelynae,(Xvariatus,(Xcouchianus,(Xgordoni,Xmeyeri):0.263):3.532):0.642):0.411):0.295):0.468):0.654):1.022):0.788):1.917)#H27:0.149::0.572):0.668,Xalvarezi):0.257,(Xsignum,#H27:1.381::0.428):4.669);"

#net = readTopology("cui3.out")
net = readTopology(text)
#printPartitions(net)
#printEdges(net)
cleanBL!(net)
cleanAfterRead!(net,false)
length(net.hybrid)
@test [n.number for n in net.hybrid] == [7,15,25] # or: wrong hybrid nodes"

# ----------------- update everything but partition ------------------
@testset "udpate hybrid 1" begin
n=net.hybrid[1];
flag, nocycle, edgesInCycle, nodesInCycle = updateInCycle!(net,n);
@test flag
@test [n.number for n in nodesInCycle] == [7,-16,-10,-11,-12]
@test [n.number for n in edgesInCycle] == [16,24,15,9,8]
flag2, edgesGammaz = updateGammaz!(net,n,false);
@test flag2
@test isempty([n.number for n in edgesGammaz])
@test !n.isBadDiamondI
@test !n.isBadDiamondII
flag3, edgesRoot = updateContainRoot!(net,n);
@test flag3
@test [n.number for n in edgesRoot] == [7]
end

@testset "udpate hybrid 2" begin
n=net.hybrid[2];
flag, nocycle, edgesInCycle, nodesInCycle = updateInCycle!(net,n);
@test flag
@test [n.number for n in nodesInCycle] == [15,-24,-23,-22,-21,-20,-8]
@test [n.number for n in edgesInCycle] == [31,32,42,43,44,45,26]
flag2, edgesGammaz = updateGammaz!(net,n,false);
@test flag2
@test isempty([n.number for n in edgesGammaz])
@test !n.isBadDiamondI
@test !n.isBadDiamondII
flag3, edgesRoot = updateContainRoot!(net,n);
@test flag3
@test [n.number for n in edgesRoot] == [25,15,9,6,14,12,10,11,13,24,23,19,17,18,22,20,21]
end

@testset "udpate hybrid 3" begin
n=net.hybrid[3];
flag, nocycle, edgesInCycle, nodesInCycle = updateInCycle!(net,n);
@test flag
@test [n.number for n in nodesInCycle] == [25,-29,-2,-3,-4]
@test [n.number for n in edgesInCycle] == [53,54,51,49,48]
flag2, edgesGammaz = updateGammaz!(net,n,false);
@test flag2
@test isempty([n.number for n in edgesGammaz])
@test !n.isBadDiamondI
@test !n.isBadDiamondII
flag3, edgesRoot = updateContainRoot!(net,n);
@test flag3
@test [n.number for n in edgesRoot] == [47,5,3,4,46,45,27,44,28,43,29,42,32,30,41,33,40,34,39,35,38,36,37]

@test isempty(net.partition)
end

# -------------- update partition ---------------------
@testset "partition hybrid 1" begin
n=net.hybrid[1];
nocycle, edgesInCycle, nodesInCycle = identifyInCycle(net,n);
@test [n.number for n in edgesInCycle] == [16,24,15,9,8]
@test [n.number for n in nodesInCycle] == [7,-16,-10,-11,-12]
updatePartition!(net,nodesInCycle)
@test length(net.partition) == 5
@test [n.number for n in net.partition[1].edges] == [7]
@test [n.number for n in net.partition[2].edges] == [23,19,17,18,22,20,21]
@test [n.number for n in net.partition[3].edges] == [25]
@test [n.number for n in net.partition[4].edges] == [14,12,10,11,13]
@test [n.number for n in net.partition[5].edges] == [6]
end

@testset "partition hybrid 2" begin
n=net.hybrid[2];
nocycle, edgesInCycle, nodesInCycle = identifyInCycle(net,n);
@test [n.number for n in nodesInCycle] == [15,-24,-23,-22,-21,-20,-8]
@test [n.number for n in edgesInCycle] == [31,32,42,43,44,45,26]
updatePartition!(net,nodesInCycle)
@test length(net.partition) == 11
@test [n.number for n in net.partition[6].edges] == [30]
@test [n.number for n in net.partition[7].edges] == [41,33,40,34,39,35,38,36,37]
@test [n.number for n in net.partition[8].edges] == [29]
@test [n.number for n in net.partition[9].edges] == [28]
@test [n.number for n in net.partition[10].edges] == [27]
@test [n.number for n in net.partition[11].edges] == [46,5,3,4,47]
end

@testset "partition hybrid 3" begin
n=net.hybrid[3];
nocycle, edgesInCycle, nodesInCycle = identifyInCycle(net,n);
@test [n.number for n in nodesInCycle] == [25,-29,-2,-3,-4]
@test [n.number for n in edgesInCycle] == [53,54,51,49,48]
updatePartition!(net,nodesInCycle)
@test length(net.partition) == 15
@test [n.number for n in net.partition[12].edges] == [52]
@test [n.number for n in net.partition[13].edges] == [1]
@test [n.number for n in net.partition[14].edges] == [50]
@test [n.number for n in net.partition[15].edges] == [2]
end
#printPartitions(net)
end


text:
"""
    Gaussian(θ, γ)

Gaussian layer, with location parameters `θ` and scale parameters `γ`.
"""
struct Gaussian{N,Aθ,Aγ} <: AbstractLayer{N}
    θ::Aθ
    γ::Aγ
    function Gaussian(θ::AbstractArray, γ::AbstractArray)
        @assert size(θ) == size(γ)
        return new{ndims(θ), typeof(θ), typeof(γ)}(θ, γ)
    end
end

Gaussian(::Type{T}, n::Int...) where {T} = Gaussian(zeros(T, n...), ones(T, n...))
Gaussian(n::Int...) = Gaussian(Float64, n...)

Base.repeat(l::Gaussian, n::Int...) = Gaussian(repeat(l.θ, n...), repeat(l.γ, n...))

energies(layer::Gaussian, x::AbstractArray) = gauss_energy.(layer.θ, layer.γ, x)

function free_energies(l::Gaussian, inputs::Union{Real,AbstractArray} = 0)
    return @. -(l.θ .+ inputs)^2 / abs(2l.γ) + log(abs(l.γ)/π/2) / 2
end

function sample_from_inputs(layer::Gaussian, inputs::Union{Real,AbstractArray} = 0)
    μ = mean_from_inputs(layer, inputs)
    σ = sqrt.(var_from_inputs(layer, inputs))
    z = randn!(similar(μ))
    return μ .+ σ .* z
end

mean_from_inputs(l::Gaussian, inputs::Union{Real,AbstractArray} = 0) = (l.θ .+ inputs) ./ abs.(l.γ)
var_from_inputs(l::Gaussian, inputs::Union{Real,AbstractArray} = 0) = inv.(abs.(l.γ .+ zero(inputs)))
std_from_inputs(l::Gaussian, inputs::Union{Real,AbstractArray} = 0) = sqrt.(var_from_inputs(l, inputs))
mode_from_inputs(l::Gaussian, inputs::Union{Real,AbstractArray} = 0) = mean_from_inputs(l, inputs)

function meanvar_from_inputs(l::Gaussian, inputs::Union{Real,AbstractArray} = 0)
    return mean_from_inputs(l, inputs), var_from_inputs(l, inputs)
end

function mean_abs_from_inputs(layer::Gaussian, inputs::Union{Real,AbstractArray} = 0)
    μ = mean_from_inputs(layer, inputs)
    ν = var_from_inputs(layer, inputs)
    return @. √(2ν/π) * exp(-μ^2 / (2ν)) + μ * erf(μ / √(2ν))
end

gauss_energy(θ::Real, γ::Real, x::Real) = (abs(γ) * x / 2 - θ) * x

function ∂free_energies(layer::Gaussian, inputs::Union{Real,AbstractArray} = 0)
    θ = layer.θ .+ inputs
    abs_γ = abs.(layer.γ)
    return (
        θ = -θ ./ abs_γ,
        γ = @. sign(layer.γ) * (abs_γ + θ^2) / abs_γ^2 / 2
    )
end

struct GaussStats{A<:AbstractArray}
    x1::A
    x2::A
    function GaussStats(layer::AbstractLayer, data::AbstractArray; wts = nothing)
        x1 = batchmean(layer, data; wts)
        x2 = batchmean(layer, data.^2; wts)
        return new{typeof(x1)}(x1, x2)
    end
end

Base.size(stats::GaussStats) = size(stats.x1)
suffstats(layer::Gaussian, x::AbstractArray; wts = nothing) = GaussStats(layer, x; wts)

function ∂energy(layer::Gaussian, stats::GaussStats)
    @assert size(layer) == size(stats)
    return (; θ = -stats.x1, γ = sign.(layer.γ) .* stats.x2 / 2)
end


text:
using Documenter

#if haskey(ENV, "GITHUB_ACTIONS")
#    ENV["JULIA_DEBUG"] = "Documenter"
#end

#Documenter.post_status(; type="pending", repo="github.com/ptiede/Comrade.jl.git")
using Comrade

makedocs(;
    modules=[Comrade],
    authors="Paul Tiede",
    repo="https://github.com/ptiede/Comrade.jl/blob/{commit}{path}#L{line}",
    sitename="Comrade.jl",
    format=Documenter.HTML(;
        prettyurls=get(ENV, "CI", "false") == "true",
        canonical="https://ptiede.github.io/Comrade.jl",
        assets=String[],
    ),
    pages=[
        "Home" => "index.md",
        "getting_started.md",
        "api.md"
    ],
)

deploydocs(;
    repo="github.com/ptiede/Comrade.jl",
    push_preview=true,
    forcepush=true,
    devbranch = "main",
    devurl = "latest",
)


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

"""
    GitBlame(repo::GitRepo, path::AbstractString; options::BlameOptions=BlameOptions())

Construct a `GitBlame` object for the file at `path`, using change information gleaned
from the history of `repo`. The `GitBlame` object records who changed which chunks of
the file when, and how. `options` controls how to separate the contents of the file and
which commits to probe - see [`BlameOptions`](@ref) for more information.
"""
function GitBlame(repo::GitRepo, path::AbstractString; options::BlameOptions=BlameOptions())
    blame_ptr_ptr = Ref{Ptr{Cvoid}}(C_NULL)
    @check ccall((:git_blame_file, :libgit2), Cint,
                  (Ptr{Ptr{Cvoid}}, Ptr{Cvoid}, Cstring, Ptr{BlameOptions}),
                   blame_ptr_ptr, repo.ptr, path, Ref(options))
    return GitBlame(repo, blame_ptr_ptr[])
end

"""
    counthunks(blame::GitBlame)

Return the number of distinct "hunks" with a file. A hunk may contain multiple lines.
A hunk is usually a piece of a file that was added/changed/removed together, for example,
a function added to a source file or an inner loop that was optimized out of
that function later.
"""
function counthunks(blame::GitBlame)
    return ccall((:git_blame_get_hunk_count, :libgit2), Int32, (Ptr{Cvoid},), blame.ptr)
end

function Base.getindex(blame::GitBlame, i::Integer)
    if !(1 <= i <= counthunks(blame))
        throw(BoundsError(blame, (i,)))
    end
    GC.@preserve blame begin
        hunk_ptr = ccall((:git_blame_get_hunk_byindex, :libgit2),
                          Ptr{BlameHunk},
                          (Ptr{Cvoid}, Csize_t), blame.ptr, i-1)
        elem = unsafe_load(hunk_ptr)
    end
    return elem
end

function Base.show(io::IO, blame_hunk::BlameHunk)
    println(io, "GitBlameHunk:")
    println(io, "Original path: ", unsafe_string(blame_hunk.orig_path))
    println(io, "Lines in hunk: ", blame_hunk.lines_in_hunk)
    println(io, "Final commit oid: ", blame_hunk.final_commit_id)
    print(io, "Final signature: ")
    show(io, Signature(blame_hunk.final_signature))
    println(io)
    println(io, "Original commit oid: ", blame_hunk.orig_commit_id)
    print(io, "Original signature: ")
    show(io, Signature(blame_hunk.orig_signature))
end

max_stars_repo_name: djsegal/julia-fork

max_issues_repo_name: djsegal/julia-fork

max_forks_repo_name: djsegal/julia-fork


text:
using Distributed
addprocs(3, exeflags="--project")

const jobs    = RemoteChannel(()->Channel{Int}(32))
const results = RemoteChannel(()->Channel{Tuple}(32))
n = 12

function make_jobs(n)
	for i in 1:n
		put!(jobs, i)
	end
end

make_jobs(n) # Feed the jobs channel with "n" jobs.

@everywhere function do_work(jobs, results) # Define work function everywhere.
	while true
       job_id = take!(jobs)
       exec_time = rand()
       sleep(exec_time)
       put!(results, (job_id, exec_time, myid()))
	end
end

for p in workers() # Start tasks on the workers to process requests in parallel.
   @async remote_do(do_work, p, jobs, results) # Similar to remotecall.
end

@elapsed while n > 0 # Print out results.
   job_id, exec_time, location_worker = take!(results)
   println("$job_id finished in $(round(exec_time, digits=2)) seconds on worker $location_worker")
   n = n - 1
end


text:
"""
    check_interval(f, a::Arf, b::Arf; check_unique = true)

Check if the function `f` has a zero on the interval `[a, b]`.

Returns `(maybe, unique)`. If `maybe` is `false` then `f` is proved to
not have a zero on the interval, if it's `true` then `f` might or
might not have a zero on the interval. If `unique` is `true` then `f`
is guaranteed to have exactly one zero on the interval, if it's
`false` then no information is given.

If `check_unique = false` then don't check for uniqueness, `unique`
will always be `false` in this case.

The function `f` should satisfy the same properties as for
[`isolate_roots`](@ref).

FUTURE WORK:
- Take values on endpoints as arguments to avoid having to compute
  them multiple times.

"""
function check_interval(f, a::Arf, b::Arf; check_unique::Bool = true)
    x = Arb((a, b))
    maybe = Arblib.contains_zero(f(x))

    maybe || return false, false

    check_unique || return true, false

    a_sign = Arblib.sgn_nonzero(f(Arb(a)))
    b_sign = Arblib.sgn_nonzero(f(Arb(b)))

    if a_sign * b_sign < 0
        df = f(ArbSeries((x, 1)))[1]
        return true, !Arblib.contains_zero(df)
    else
        return true, false
    end
end

function check_interval(p::ArbPoly, a::Arf, b::Arf; check_unique::Bool = true)
    x = Arb((a, b))

    check_unique || return Arblib.contains_zero(p(x)), false

    px, dpx = Arblib.evaluate2(p, x)

    Arblib.contains_zero(px) || return false, false

    a_sign = Arblib.sgn_nonzero(p(Arb(a)))
    b_sign = Arblib.sgn_nonzero(p(Arb(b)))

    if a_sign * b_sign < 0
        unique = !Arblib.contains_zero(dpx)
        return true, unique
    elseif a_sign * b_sign > 0
        maybe = Arblib.contains_zero(dpx)
        return maybe, false
    else
        return true, false
    end
end

"""
    isolate_roots(f, a::Arf, b::Arf; depth = 10, check_unique = true, verbose = false)

Isolate the roots of `f` on the interval `[a, b]`.

Returns a tuple `(found, flags)` where `found` is a vector of
subintervals of `[a, b]` given by pairs of `Arf` and `flags` is a
vector of booleans. The output has the following properties:
- The function has no roots on interval outside of the subintervals in
  `found`.
- Subintervals are sorted in increasing order (with no overlap except
  possibly starting and ending with the same point).
- Subintervals with a flag `true` has exactly one (simple) root.
- Subintervals with a flag `false` may or may not contain roots.

If all flags are true then all roots of the function on interval have
been isolated. If there are output subintervals on which the existence
or nonexistence of roots could not be determined, the user may attempt
further searches on those subintervals (possibly with increased
precision and/or increased bounds for the breaking criteria). Note
that roots of multiplicity higher than one and roots located exactly
at endpoints cannot be isolated by the algorithm.

This method closely mimics the function `arb_calc_isolate_roots` in
Arb, both in behaviour and in implementation.

The function `f` should support evaluation on both `Arb` and
`ArbSeries` and should return an enclosure of the result in both
cases. The evaluation on `ArbSeries` is done to compute the derivative
to check uniqueness, this can be skipped by setting `check_unique =
false`, see below. Alternatively `f` can be an `ArbPoly`, in which
case some optimizations are done.

If `check_unique = false` then don't check if roots are unique. In
this case `flags` will always be false and it will simply continue to
split the subintervals until it reaches the maximum depth. The only
benefit to this is that the derivative of `f` doesn't have to be
computed so it can be used when the function of interest doesn't
support evaluating derivatives.

If `verbose = true` then output information after each iteration on
the number of found roots and the number of remaining intervals that
needs to be split further.

FUTURE WORK:
- Parallelize.
- Compute values on endpoints of subintervals once to not have to
  compute them multiple times.
- Should it use `BitVector` or `Vector{Bool}` for flags?
- The current implementation does a breadth-first-search. Consider
  rewriting it to do a depth-first-search instead to reduce memory
  usage in case of many splittings. This might however make it harder
  to parallelize.
"""
function isolate_roots(
    f,
    a::Arf,
    b::Arf;
    depth::Integer = 10,
    check_unique::Bool = true,
    verbose = false,
)
    check_interval(a, b)

    if a == b
        fa = f(Arb(a))
        if Arblib.contains_zero(fa)
            return [(a, b)], BitVector((iszero(fa),))
        else
            return Vector{NTuple{2,Arf}}(), BitVector()
        end
    end

    intervals = [(a, b)]

    found = Vector{NTuple{2,Arf}}()
    flags = BitVector()

    iterations = 0
    while !isempty(intervals) && iterations < depth
        iterations += 1

        next_intervals = empty(intervals)

        for interval in intervals
            maybe, unique = check_interval(f, interval...; check_unique)

            if unique
                push!(found, interval)
                push!(flags, true)
            elseif maybe
                if iterations < depth
                    push!(next_intervals, bisect_interval(interval...)...)
                else
                    # If we are on the last iteration don't split the interval
                    push!(next_intervals, interval)
                end
            end
        end

        intervals = next_intervals

        verbose &&
            iterations < depth &&
            @info "isolate_roots iteration: $(lpad(iterations, 2)), " *
                  "found: $(lpad(length(found), 2)), " *
                  "remaining intervals: $(length(intervals) ÷ 2)"
    end

    verbose &&
        iterations == depth &&
        @info "isolate_roots iteration: $(lpad(iterations, 2)), " *
              "found: $(lpad(length(found), 2)), " *
              "remaining intervals: $(length(intervals))"

    append!(found, intervals)
    append!(flags, zeros(Bool, length(intervals)))

    p = sortperm(found, by = interval -> interval[1])

    return found[p], flags[p]
end


text:
using StructJuMP

N = 1000
numScen = 1

m = StructuredModel(numScen)

@defVar(m, x[1:N] >= 0)
@addConstraint(m, sum{x[i], i=1:N} >= -1)
@setObjective(m, Min, sum{x[i], i=1:N})

bl = StructuredModel(parent=m)
@defVar(bl, y[1:N] >= 0)
@addConstraint(bl, constr[r=2:(N-1)], x[r] + x[r-1] + y[r] + y[r+1] >= 0)
@setObjective(bl, Min, sum{y[i], i=1:N})


text:
using Test
using JAC, ..Defaults, ..TestFrames

@testset "Name" begin
    ##x global short = true
    printstyled("\nPerform tests on the JAC program; this may take a while .... \n", color=:cyan)
    ## Defaults.Constants.define("print test: open", pwd() * "/runtests.report")

    ##x include("inc-halfintegers.jl")
    ##x include("inc-wignersymbols.jl")
    ##x include("inc-racahsum.jl")

    ##x redirect_stdout(streamDummy) do
    @testset "JAC methods" begin
        @test TestFrames.testEvaluation_rcsfValues()
        @test TestFrames.testEvaluation_angularCoefientsValues()
    end

    ## Defaults.setDefaults("print test: close", "")
end


text:
module henon_heiles
using ReachabilityAnalysis
@taylorize function henon_heiles!(du, u, p, t)
    p₁, p₂, q₁, q₂ = u[1], u[2], u[3], u[4]
    du[1] = -q₁*(1 + 2q₂)
    du[2] = -q₂-(q₁^2 - q₂^2)
    du[3] = p₁
    du[4] = p₂
end

function model(X0)
    S = @system(x' = henon_heiles!(x), dim:4)
    return IVP(S, X0)
end
end # module

max_stars_repo_name: mforets/ReachabilityLibrary

max_issues_repo_name: mforets/ReachabilityLibrary

max_forks_repo_name: mforets/ReachabilityLibrary


text:
module DataWrapper

using JavaCall

const AbsoluteDate = @jimport org.orekit.time.AbsoluteDate
const Class = @jimport java.lang.Class
const CssiSpaceWeatherData = @jimport org.orekit.models.earth.atmosphere.data.CssiSpaceWeatherData
const CssiSpaceWeatherDataLoader = @jimport org.orekit.models.earth.atmosphere.data.CssiSpaceWeatherDataLoader
const DataProvidersManager = @jimport org.orekit.data.DataProvidersManager
const DateComponents = @jimport org.orekit.time.DateComponents
const InputStream = @jimport java.io.InputStream
const JString = @jimport java.lang.String
const MarshallSolarActivityFutureEstimation = @jimport org.orekit.models.earth.atmosphere.data.MarshallSolarActivityFutureEstimation
const MarshallSolarActivityFutureEstimation_StrengthLevel = @jimport org.orekit.models.earth.atmosphere.data.MarshallSolarActivityFutureEstimation$StrengthLevel
const Object = @jimport java.lang.Object
const SortedSet = @jimport java.util.SortedSet
const TimeScale = @jimport org.orekit.time.TimeScale

export CssiSpaceWeatherData
export CssiSpaceWeatherDataLoader
export MarshallSolarActivityFutureEstimation
export equals
export get24_hours_kp
export get_ap
export get_average_flux
export get_class
export get_daily_flux
export get_data_set
export get_file_date
export get_instant_flux
export get_last_daily_predicted_date
export get_last_observed_date
export get_max_date
export get_mean_flux
export get_min_date
export get_strength_level
export get_supported_names
export get_three_hourly_kp
export hash_code
export load_data
export notify
export notify_all
export still_accepts_data
export to_string
export wait

include("cssi_space_weather_data.jl")
include("cssi_space_weather_data_loader.jl")
include("marshall_solar_activity_future_estimation.jl")

end


text:
function doOptim(vecModelSub, modMas, vecConsRef, vecConsConvex, vecLambda, dualPen,
    dualPenMult, dualPenThreshold, vecMuMinus, vecMuPlus, vecMuMinusConv, vecMuPlusConv, epsilon, whePrint)
    ## Initialize
    numQ = size(vecModelSub[1].mat_e)[1]  # [num of constraints in matA_0]
    numSub = length(vecModelSub)
    extremePoints = [[]]
    extremePointForSub = [-1]
    vecObjCoef = [-1000.0]
    wheDoneStab = false
    iter = 1
    ##
    while !wheDoneStab
        vecPi = []
        vecKappa = []
        wheDoneColGen = false
        while !wheDoneColGen
            (vecPi, vecKappa, obj_master) = solveMas(modMas, vecConsRef, vecConsConvex)
            ## Print the result
            println("$(iter)-th iteration. obj_master = $(obj_master).")
            if whePrint
                println("vecPi = $(vecPi)\n",
                        "vecKappa = $(vecKappa)")
            end
            wheDoneColGen = true
            ## Column Generation
            costReduceBest = -1
            for k = 1:numSub
                (costReduce, vec_x_result) = solveSub(vecModelSub[k], vecPi, vecKappa[k])
                if whePrint
                    println("Reduced cost of $(k)-th sub model is $costReduce.")
                end
                if costReduce > costReduceBest
                    costReduceBest = costReduce
                end
                # remember to look for negative reduced cost if we are minimizing.
                if costReduce > epsilon
                    (lambdaNew, objCoefNew) = addColToMas(modMas, vecModelSub[k], vec_x_result,
                        vecConsRef, vecConsConvex[k])
                    push!(vecLambda, lambdaNew)
                    push!(vecObjCoef, objCoefNew)
                    push!(extremePoints, vec_x_result)
                    push!(extremePointForSub, k)
                    wheDoneColGen = false
                end
            end
            ##
            iter += 1
            println("best reduced cost is $(costReduceBest)")
            println("--------------------------------------------------------------------------------")
        end
        ## See whether we should update stabilization parameters.
        if dualPen > 0
            print("Updating dual penalty. old value = $(dualPen), ")
            dualPen *= dualPenMult
            if dualPen < dualPenThreshold
                dualPen = 0
            end
            println("new value = $dualPen")
            vecDualGuessPi = vecPi
            vecDualGuessKappa = vecKappa
            println("New vecDualGuessPi = $vecDualGuessPi")
            println("New vecDualGuessKappa = $vecDualGuessKappa")
            updateStabInfo(modMas, vecObjCoef, vecLambda, numQ, numSub, dualPen, vecDualGuessPi, vecDualGuessKappa,
                vecMuMinus, vecMuPlus, vecMuMinusConv, vecMuPlusConv)
        else
            wheDoneStab = true
        end
    end
    ## 5,  Print the Result
    println("Optimization Done After $(iter-1) Iterations.\n",
            "objMaster = $(objective_value(modMas))")
    vecLambdaResult = value(vecLambda)
    return (vecLambdaResult, extremePointForSub, extremePoints)
end

function getStrNameVar(numSub, numXPerSub)
    vecStrNameVar = Vector{String}(undef, numSub * numXPerSub)
    idx = 1
    for i = 1:numSub
        for j = 1:numXPerSub
            vecStrNameVar[idx] = "x_($i, $j)"
            idx += 1
        end
    end
    return vecStrNameVar
end

function printResult(vecStrNameVar, indexSub, vecLambda, vecLambdaResult, extremePointForSub, extremePoints)
    # compute values of original variables
    origVarValSub = []
    for s = 1: length(indexSub)
        push!(origVarValSub, zeros(length(indexSub[s])))
    end
    for p = 1: length(vecLambdaResult)
        if vecLambdaResult[p] > 0.000001
            println("lambda_$p = ", vecLambdaResult[p], ", sub = $(extremePointForSub[p]), \n",
                    "extreme point = $(extremePoints[p]).")
            origVarValSub[extremePointForSub[p]] += vecLambdaResult[p] * extremePoints[p]
        end
    end
    for s = 1: length(indexSub)
        #println("var val for sub problem $s: $(origVarValSub[s])")
        for t = 1: length(origVarValSub[s])
            if abs(origVarValSub[s][t]) > 0.000001
                println("$(vecStrNameVar[indexSub[s][t]]) = $(origVarValSub[s][t])")
            end
        end
    end
end


text:
import LinearAlgebra: svd
import Base: getproperty

# Iteration allows destructuring, e.g. U, S, V = svd(A)
# These definitions mirror those defined in the LinearAlgebra module, see
# https://github.com/JuliaLang/julia/blob/master/stdlib/LinearAlgebra/src/svd.jl#L20-L24
Base.iterate(usv::Branch{<:SVD}) = (usv.U, Val(:S))
Base.iterate(usv::Branch{<:SVD}, ::Val{:S}) = (usv.S, Val(:V))
Base.iterate(usv::Branch{<:SVD}, ::Val{:V}) = (usv.V, Val(:done))
Base.iterate(usv::Branch{<:SVD}, ::Val{:done}) = nothing