text:
#---------------------------------------------------------------------
# LABELS
#--------------------------------------------------------------------
"""
$(TYPEDSIGNATURES)

Add `node` to `d[:labels]`, a vector of [`PnmlLabel`](@ref). Return updated `d[:labels]`.
"""
function add_label!(d::PnmlDict, node::XMLNode, pntd; kw...)
    # Pnml considers any "unknown" element to be a label so its key is `:labels`.

    # The value is initialized to `nothing since it is expected that most labels
    # will have defined tags and semantics. And be given a key `:tag`.
    # Will convert value to a vector on first use.
    if d[:labels] === nothing
        d[:labels] = PnmlLabel[]
    end
    add_label!(d[:labels], node, pntd; kw...)
end

function add_label!(v::Vector{PnmlLabel}, node::XMLNode, pntd; kw...)
    #@show "add label! $(nodename(node))"
    haskey(tagmap, node.name) && @info "$(node.name) is known tag being treated as unclaimed."
    label = PnmlLabel(unclaimed_label(node, pntd; kw...), node) #TODO handle types
    haskey(tagmap, node.name) && @info "$(node.name) parsed to type $(typeof(label))."
    push!(v, label)
    return v
end

#---------------------------------------------------------------------
# TOOLINFO
#---------------------------------------------------------------------
"""
$(TYPEDSIGNATURES)

Add [`ToolInfo`](@ref) `node` to `d[:tools]`. Return updated `d[:tools]`.

The UML from the _pnml primer_ (and schemas) use <toolspecific>
as the tag name for instances of the type ToolInfo.
"""
function add_toolinfo!(d::PnmlDict, node, pntd; kw...)
    if d[:tools] === nothing
        d[:tools] = ToolInfo[] #TODO: Pick type based on PNTD/Trait?
        #TODO TokenGraphics is a known flavor.
        #TODO Tools may induce additional subtype, but if is hoped that
        #TODO label based parsing is general & flexible enough to suffice.
    end
    add_toolinfo!(d[:tools], node, pntd; kw...)
end

function add_toolinfo!(v::Vector{ToolInfo}, node, pntd; kw...)
    ti = parse_toolspecific(node, pntd; kw...) # Here is that tag name in use.
    push!(v,ti)
end

"""
$(TYPEDSIGNATURES)

Does any toolinfo attached to `d` have a matching `toolname`.
"""
function has_toolinfo end

# tools vector
function has_toolinfo(v::Vector{PnmlDict},
                      toolname::AbstractString)
    has_toolinfo(v, Regex(toolname))
end
function has_toolinfo(v::Vector{PnmlDict},
                      toolname::AbstractString,
                      version::AbstractString)
    has_toolinfo(v, Regex(toolname), Regex(version))
end
function has_toolinfo(v::Vector{PnmlDict},
                      namerex::Regex,
                      versionrex::Regex=r"^.*$")
    any(v) do tool
       match(namerex, tool.toolname) && match(versionrex, tool.version)
    end
end

"""
$(TYPEDSIGNATURES)

Return first toolinfo having a matching toolname and version.
"""
function get_toolinfo end

get_toolinfo(ti::ToolInfo, name::AbstractString) = get_toolinfo(ti, Regex(name))
get_toolinfo(ti::ToolInfo, name::AbstractString, version::AbstractString) =
    get_toolinfo(ti, Regex(name), Regex(version))
get_toolinfo(ti::ToolInfo, name::AbstractString, versionrex::Regex) =
    get_toolinfo(ti, Regex(name),  versionrex)
get_toolinfo(ti::ToolInfo, namerex::Regex, versionrex::Regex=r"^.*$") =
    get_toolinfo([ti], namerex, versionrex)

get_toolinfo(v::Vector{ToolInfo}, name::AbstractString, version::AbstractString) =
    get_toolinfo(v, Regex(name), Regex(version))
get_toolinfo(v::Vector{ToolInfo}, name::AbstractString, versionrex::Regex) =
    get_toolinfo(v, Regex(name), versionrex)

function get_toolinfo(v::Vector{ToolInfo}, namerex::Regex, versionrex::Regex=r"^.*$")
    first(get_toolinfos(v, namerex, versionrex))
end

function get_toolinfos(v::Vector{ToolInfo}, namerex::Regex, versionrex::Regex=r"^.*$")
    filter(ti -> _match(ti, namerex, versionrex), v)
end

"""
$(TYPEDSIGNATURES)
Match toolname and version.
"""
function _match end
function _match(ti::ToolInfo, name::AbstractString)
    #@show "match toolinfo $name"
    _match(ti.info, Regex(name))
end
function _match(ti::ToolInfo, name::String, version::String)
    #@show "match toolinfo $name, $version"
    _match(ti.inf, Regex(name), Regex(version))
end
function _match(ti::ToolInfo, namerex::Regex, versionrex::Regex=r"^.*$")
    #@show "match toolinfo $namerex, $versionrex"
    match_name = match(namerex, name(ti))
    match_version = match(versionrex, version(ti))
    #@show match_name
    #@show match_version
    !isnothing(match_name) && !isnothing(match_version)
end

#---------------------------------------------------------------------
"""
$(TYPEDSIGNATURES)

Return Dict of tags common to both pnml nodes and pnml labels.
See also: [`pnml_label_defaults`](@ref), [`pnml_node_defaults`](@ref).
"""
function pnml_common_defaults(node)
    PnmlDict(:graphics => nothing, # graphics tag is single despite the 's'.
             :tools => nothing, # Here the 's' indicates multiples are allowed.
             :labels => nothing) # ditto
end

"""
$(TYPEDSIGNATURES)

Merge `xs` into dictonary with default pnml node tags.
Used on: net, page ,place, transition, arc.
Usually default value will be `nothing` or empty vector.
See also: [`pnml_label_defaults`](@ref), [`pnml_common_defaults`](@ref).
"""
function pnml_node_defaults(node, xs...)
    PnmlDict(pnml_common_defaults(node)...,
             :name => nothing,
             xs...)
end

"""
$(TYPEDSIGNATURES)

Merge `xs` into dictonary with default pnml HLannotation label tags.
Used on pnml tags below a [`PnmlNode`](@ref) tag.
Label level tags include: name, inscription, initialMarking.
Notable differences from [`pnml_node_defaults`](@ref): text, structure, no name tag.
See also: [`pnml_common_defaults`](@ref).
"""
function pnml_label_defaults(node, xs...)
    PnmlDict(pnml_common_defaults(node)...,
             :text => nothing,
             :structure => nothing,
             xs...)
end

#---------------------------------------------------------------------
"""
$(TYPEDSIGNATURES)

Update `d` with any graphics, tools, and label child of `node`.
Used by [`parse_pnml_node_common!`](@ref) & [`parse_pnml_label_common!`](@ref).

Note that "labels" are the "everything else" option and this should be called after parsing
any elements that has an expected tag. Any tag that is encountered in an unexpected location
should be treated as an anonymous label for parsing.
"""
function parse_pnml_common!(d::PnmlDict, node, pntd; kw...)
    @match nodename(node) begin
        "graphics"     => (d[:graphics] = parse_graphics(node, pntd; kw...))
        "toolspecific" => add_toolinfo!(d, node, pntd; kw...)
        _ => add_label!(d, node, pntd; kw...) # label with a label allows any node to be attached & parsable.
    end
end

"""
$(TYPEDSIGNATURES)

Update `d` with `name` children, defering other tags to [`parse_pnml_common!`](@ref).
"""
function parse_pnml_node_common!(d::PnmlDict, node, pntd; kw...)
    @match nodename(node) begin
        "name" => (d[:name] = parse_name(node, pntd; kw...))
        _ => parse_pnml_common!(d, node, pntd; kw...)
    end
end

"""
$(TYPEDSIGNATURES)

Update `d` with  'text' and 'structure' children of `node`,
defering other tags to [`parse_pnml_common!`](@ref).
"""
function parse_pnml_label_common!(d::PnmlDict, node, pntd; kw...)
    @match nodename(node) begin
        "text"      => (d[:text] = parse_text(node, pntd; kw...))
        # This is the fallback as "claimed" label's parser
        # should have already consumed the <structure>.
        "structure" => (d[:structure] = parse_structure(node, pntd; kw...))
        _ => parse_pnml_common!(d, node, pntd; kw...)
    end
end

#---------------------------------------------------------------------
"""
$(TYPEDSIGNATURES)

Parse string as a number. First try integer then float.
"""
function number_value(s::AbstractString)
    x = tryparse(Int, s)
    x = isnothing(x) ?  tryparse(Float64, s) : x
end


text:
using Turing, Distributions

# include("../src/sindyc.jl")
# include("sindy_utils.jl")

#####
##### Helper functions for defining a Turing model from a SINDy struct
#####
"""
Takes in a SINDy model and returns a Turing-compatible model that defines
    priors on each parameter with the proper variable names

Note: helpful NN analogue: https://turing.ml/dev/tutorials/3-bayesnn/
"""
function convert_sindy_to_turing(model::sindycModel;
                                dat_noise_prior=Normal(5, 5.0),
                                coef_noise_std=1.0)
    n_in, n_aug = size(model.A)
    total_terms = n_in * n_aug

    prior_coef = sindy_mat2vec(model.A)

    @model turing_model(y, dat) = begin
        # Set the priors using the passed model parameters
        all_coef ~ MvNormal(prior_coef, coef_noise_std .* ones(total_terms))
        A = sindy_vec2mat(model, all_coef)
        noise ~ Truncated(dat_noise_prior, 0, 100)
        # Predict the gradient at each time step (grid collocation)
        for i in 1:size(y,2)
            preds = sindy_predict(model, all_coef, dat[:,i])
            y[:, i] ~ MvNormal(preds, noise.*ones(n_in))
        end
    end;

    return turing_model
end

"""
Takes in a SINDy model and returns a Turing-compatible model that defines
    priors on each parameter with the proper variable names
    Enforces the 0s of the original model

    convert_sindy_to_turing_enforce_zeros(model::sindycModel;
                                    dat_noise_prior=Normal(5, 5.0),
                                    coef_noise_std=1.0)
"""
function convert_sindy_to_turing_enforce_zeros(model::sindycModel;
                                dat_noise_prior=Normal(5, 5.0),
                                coef_noise_std=1.0)
    n_in, n_aug = size(model.A)

    @model turing_model(y, dat) = begin
        term_ind = get_nonzero_terms(model)
        term_val = model.A[term_ind]
        all_coef ~ MvNormal(term_val, coef_noise_std)
        noise ~ Truncated(dat_noise_prior, 0, 100)
        # Predict the gradient at each time step (grid collocation)
        for i in 1:size(y,2)
            preds = sindy_predict(model, all_coef, dat[:,i],
                        nonzero_terms=term_ind)
            y[:, i] ~ MvNormal(preds, noise.*ones(n_in))
            # y[:, i] ~ MvNormal(preds, [noise, noise, noise])
        end
    end;

    return turing_model
end

#####
##### Helper functions for vector and matrix conversion
#####
function sindy_vec2mat(model::sindycModel, coef::AbstractVector)
    return reshape(coef, size(model.A))
end

function sindy_mat2vec(mat::AbstractMatrix)
    return reshape(mat, length(mat))
end

#####
##### Helper functions for readable prediction
#####
"""Skips full matrix multiplication; only does the nonzero coefficients"""
function sindy_predict(model::sindycModel, coef, dat; nonzero_terms=nothing)
    if nonzero_terms==nothing
        nonzero_terms=get_nonzero_terms(model, linear_indices=false)
    end
    dat_aug = augment_data(model, dat)
    original_ind = CartesianIndices(model.A)
    predictions = zeros(eltype(coef), size(dat))
    for (coef_ind,mat_ind) in enumerate(nonzero_terms)
        var, dat_ind = Tuple(mat_ind)
        predictions[var] += coef[coef_ind] * dat_aug[dat_ind]
    end
    return predictions
end

#####
##### Creating model instances from chains
#####
"""
function sindy_from_chain(model_template::sindycModel,
                          chain::MCMCChains.AbstractChains;
                          enforced_zeros=true)

Creates a sindycModel object (nonlinear dynamical system) from
a chain using the template of model terms in 'model_template';
    If enforced_zeros=true, then the zeros of the 'model_template'
    are enforced; the nonzero terms must match the number of
    elements in the chain minus one (the last one is assumed to be noise)
"""
function sindy_from_chain(model_template::sindycModel,
                          chain::MCMCChains.AbstractChains;
                          enforced_zeros=true)
    if enforced_zeros
        nonzero_terms = get_nonzero_terms(model_template)
        # @warn("model_template should have the same
        #         nonzeros structure as the desired model!")
    end

    ts = 1:length(model_template.U) # TODO: multidimensional U

    n_total = length(model_template.A)
    coef_sample = sample(chain, 2) # sample size of 1 gives error
    coef_vec = Array(coef_sample)[1,1:end-1] # Leave out noise
    subs2ind = CartesianIndices(model_template.A)
    if length(coef_vec) < n_total
        # Assume they are a subset of the A matrix, and
        # the proper index is in their name
        #   Add zeros to the other parameters
        var_names = coef_sample.name_map[:parameters][1:end-1]
        coef_tmp = zeros(n_total)
        for i in 1:n_total
            if enforced_zeros
                match_ind = findall(x->subs2ind[i]==x, nonzero_terms)
                found_match = (length(match_ind)==1)
            else
                # If no explicit mapping
                match_ind = occursin.("[$i]", var_names)
                found_match = any(match_ind)
            end
            if found_match
                coef_tmp[i] = coef_vec[match_ind][1]
            end
        end
        coef_vec = coef_tmp
    end
    # Assume these parameters are the A matrix
    A = sindy_vec2mat(model_template, coef_vec)
    new_model = sindycModel(ts, A,
            model_template.B, model_template.U,
            model_template.U_func, model_template.library,
            model_template.variable_names)
    return new_model
end

"""
Generate test realizations of the gradient from the posterior of the system
    parameters. No integration is performed
    Additional function to work with SINDy models
"""
function sample_sindy_posterior_grad(chain, dat, sample_ind, sindy_template;
                                num_samples=100, t = [0], enforced_zeros=true)
    sindy_samples = [sindy_from_chain(sindy_template, chain, enforced_zeros=enforced_zeros)
                        for i in 1:num_samples]
    all_vals = zeros(num_samples, length(sample_ind), size(dat,1))
    for i in 1:num_samples
        this_model = sindy_samples[i]
        for (i_save, i_dat) in enumerate(sample_ind)
            all_vals[i, i_save, :] = this_model(dat[:,i_dat], t)
        end
    end
    all_noise = Array(sample(chain[:noise], num_samples))

    return (sample_trajectories=all_vals, sample_noise=all_noise)
end

##
export convert_sindy_to_turing, convert_sindy_to_turing_enforce_zeros,
        sindy_from_chain, sample_sindy_posterior_grad

ext: jl


text:
function exp(a::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(a) && return a
    isinf(a) && return(signbit(a) ? zero(DoubleFloat{T}) : a)

    if iszero(HI(a))
    return one(DoubleFloat{T})
    elseif isone(abs(HI(a))) && iszero(LO(a))
    if HI(a) >= zero(T)
        return DoubleFloat{T}(2.718281828459045, 1.4456468917292502e-16)
    else # isone(-HI(a)) && iszero(LO(a))
        return DoubleFloat{T}(0.36787944117144233, -1.2428753672788363e-17)
    end
    elseif abs(HI(a)) >= 709.0
      if (HI(a) <= -709.0)
         return zero(DoubleFloat{T})
      else # HI(a) >=  709.0
         return inf(DoubleFloat{T})
      end
    end

    return calc_exp(a)
end

function exp_taylor(a::DoubleFloat{T}) where {T<:IEEEFloat}
    x = a
    x2 = x*x
    x3 = x*x2
    x4 = x2*x2
    x5 = x2*x3
    x10 = x5*x5
    x15 = x5*x10
    x20 = x10*x10
    x25 = x10*x15

    z = x + inv_fact[2]*x2 + inv_fact[3]*x3 + inv_fact[4]*x4
    z2 = x5 * (inv_fact[5] + x*inv_fact[6] + x2*inv_fact[7] +
         x3*inv_fact[8] + x4*inv_fact[9])
    z3 = x10 * (inv_fact[10] + x*inv_fact[11] + x2*inv_fact[12] +
         x3*inv_fact[13] + x4*inv_fact[14])
    z4 = x15 * (inv_fact[15] + x*inv_fact[16] + x2*inv_fact[17] +
         x3*inv_fact[18] + x4*inv_fact[19])
    z5 = x20 * (inv_fact[20] + x*inv_fact[21] + x2*inv_fact[22] +
         x3*inv_fact[23] + x4*inv_fact[24])
    z6 = x25 * (inv_fact[25] + x*inv_fact[26] + x2*inv_fact[27])

    ((((z6+z5)+z4)+z3)+z2)+z + one(DoubleFloat{T})
end

@inline exp_zero_half(a::DoubleFloat{T}) where {T<:IEEEFloat} = exp_taylor(a)

@inline function exp_half_one(a::DoubleFloat{T}) where {T<:IEEEFloat}
    z = mul_by_half(a)
    z = exp_zero_half(z)
    z = square(z)
    return z
end

function mul_by_half(r::DoubleFloat{T}) where {T<:IEEEFloat}
    frhi, xphi = frexp(HI(r))
    frlo, xplo = frexp(LO(r))
    xphi -= 1
    xplo -= 1
    hi = ldexp(frhi, xphi)
    lo = ldexp(frlo, xplo)
    return DoubleFloat{T}(hi, lo)
end

function mul_by_two(r::DoubleFloat{T}) where {T<:IEEEFloat}
    frhi, xphi = frexp(HI(r))
    frlo, xplo = frexp(LO(r))
    xphi += 1
    xplo += 1
    hi = ldexp(frhi, xphi)
    lo = ldexp(frlo, xplo)
    return DoubleFloat{T}(hi, lo)
end

function mul_pow2(r::DoubleFloat{T}, n::Int) where {T<:IEEEFloat}
    frhi, xphi = frexp(HI(r))
    frlo, xplo = frexp(LO(r))
    xphi += n
    xplo += n
    hi = ldexp(frhi, xphi)
    lo = ldexp(frlo, xplo)
    return DoubleFloat{T}(hi, lo)
end

function mul_pwr2(r::DoubleFloat{T}, n::Real) where {T<:IEEEFloat}
    m = 2.0^n
    return DoubleFloat{T}(HI(r)*m, LO(r)*m)
end

function Base.:(^)(r::DoubleFloat{T}, n::Int) where {T<:IEEEFloat}
    if (n == 0)
        return one(DoubleFloat{T})
    end

    s = one(DoubleFloat{T})
    nabs = abs(n)

    if (nabs > 1)
        while nabs > 0
           if nabs % 2 == 1
              s = s*r
           end
           nabs = div(nabs, 2)
           if nabs > 0
               r = square(r)
           end
        end
    else
        s = r
    end

    if (n < 0)
        s = inv(s)
    end
    return s
end

function Base.:(^)(r::DoubleFloat{T}, n::DoubleFloat{T}) where {T<:IEEEFloat}
   if isinteger(n)
      return r^Int(n)
   else
      return exp(n * log(r))
   end
end

function Base.:(^)(r::Int, n::DoubleFloat{T}) where {T<:IEEEFloat}
   if isinteger(n)
      return r^Int(n)
   else
      return exp(n * log(r))
   end
end

function calc_exp(a::DoubleFloat{T}) where {T<:IEEEFloat}
    is_neg = signbit(HI(a))
    xabs = is_neg ? -a : a
    xintpart = modf(xabs)[2]
    xintpart = xintpart.hi + xintpart.lo
    xint = Int64(xintpart)
    xfrac = xabs - T(xint)

    if 0 < xint <= 64
        zint = exp_int[xint]
    elseif xint === zero(Int64)
        zint = zero(DoubleFloat{T})
    else
        dv, rm = divrem(xint, 64)
        zint = exp_int[64]^dv
        if rm > 0
            zint = zint * exp_int[rm]
        end
    end

    # exp(xfrac)
    if HI(xfrac) < 0.5
        zfrac = exp_zero_half(xfrac)
    elseif HI(xfrac) > 0.5
        zfrac = exp_half_one(xfrac)
    else
        if LO(xfrac) == 0.0
            zfrac = DoubleFloat{T}(1.6487212707001282, -4.731568479435833e-17)
        elseif signbit(LO(xfrac))
            zfrac = exp_zero_half(xfrac)
        else
            zfrac = exp_half_one(xfrac)
        end
    end

    z = HI(zint) == zero(T) ? zfrac : zint * zfrac
    if is_neg
        z = inv(z)
    end

    return z
end

function expm1(a::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(a) && return a
    isinf(a) && return(signbit(a) ? zero(DoubleFloat{T}) : a)
    u = exp(a)
    if (u == one(DoubleFloat{T}))
        x
    elseif (u-1.0 == -one(DoubleFloat{T}))
        -one(DoubleFloat{T})
    else
        a*(u-1.0) / log(u)
    end
end

function exp2(a::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(a) && return a
    isinf(a) && return(signbit(a) ? zero(DoubleFloat{T}) : a)
    return DoubleFloat{T}(2)^a
end

function exp10(a::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(a) && return a
    isinf(a) && return(signbit(a) ? zero(DoubleFloat{T}) : a)
    return DoubleFloat{T}(10)^a
end

#=
# ratio of polys from
# https://github.com/sukop/doubledouble/blob/master/doubledouble.py
function calc_exp_frac(x::DoubleFloat{T}) where {T<:AbstractFloat}
  u = (((((((((((x +
                   156.0)*x + 12012.0)*x +
                   600600.0)*x + 21621600.0)*x +
                   588107520.0)*x + 12350257920.0)*x +
                   201132771840.0)*x + 2514159648000.0)*x +
                   23465490048000.0)*x + 154872234316800.0)*x +
                   647647525324800.0)*x + 1295295050649600.0
   v = (((((((((((x -
                   156.0)*x + 12012.0)*x -
                   600600.0)*x + 21621600.0)*x -
                   588107520.0)*x + 12350257920.0)*x -
                   201132771840.0)*x + 2514159648000.0)*x -
                   23465490048000.0)*x + 154872234316800.0)*x -

function exp2(a::DoubleFloat{T}) where {T<:AbstractFloat}
    isnan(x) && return x
    isinf(x) && return(signbit(x) ? zero(DoubleFloat{T}) : x)
    return 2^a
end
647647525324800.0)*x + 1295295050649600.0
  u = u/v
  return u
end
=#

function log(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && !signbit(x) && return x
    x === zero(DoubleFloat{T}) && return neginf(DoubleFloat{T})
    y = DoubleFloat(log(HI(x)), zero(T))
    z = exp(y)
    adj = (z - x) / (z + x)
    adj = mul_by_two(adj)
    y = y - adj
    return y
end

function log1p(x::DoubleFloat{T})  where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && !signbit(x) && return
    u = 1.0 + x
    if u == one(DoubleFloat{T})
        x
    else
        log(u)*x/(u-1.0)
    end
end

logten(::Type{DoubleFloat{Float64}}) = Double64(2.302585092994046, -2.1707562233822494e-16)
logtwo(::Type{DoubleFloat{Float64}}) = Double64(0.6931471805599453, 2.3190468138462996e-17)
logtwo(::Type{DoubleFloat{Float32}}) = Double32(0.6931472, -1.9046542e-9)
logten(::Type{DoubleFloat{Float32}}) = Double32(2.3025851, -3.1975436e-8)
logtwo(::Type{DoubleFloat{Float16}}) = Double16(0.6934, -0.0002122)
logten(::Type{DoubleFloat{Float16}}) = Double16(2.303, -0.0001493)

function log2(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && !signbit(x) && return x
    log(x) / logtwo(DoubleFloat{T})
end

function log10(x::DoubleFloat{T}) where {T<:IEEEFloat}
    isnan(x) && return x
    isinf(x) && !signbit(x) && return x
    log(x) / logten(DoubleFloat{T})
end


text:
module CoupledODETools
    using LabelledArrays
    using MLStyle
    using MacroTools: @capture
    using RecursiveArrayTools
    using CUDA

    export Component, @Component, Network, SharedPar
    export generate, generate_cuarray, generate_ensemble

    include("./utils.jl")
    include("./Component.jl")
    include("./Network.jl")

end # module


text:
# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

export closure!
"""
  `closure!(fst; star=true)`

Same as `closure` but modifies `fst` in-place.

"""
function closure!(fst; star=true)
  if star
    fst = closure_star!(fst)
  else
    fst = closure_plus!(fst)
  end
  return fst
end

export closure
"""
  `closure(fst; star=true)`

Returns a new WFST `cfst` that is the closure of `fst`.
If `fst` transuces `labels` to `olabel` with weight `w`, `cfst` will be able to transduce `repeat(ilabels,n)` to `repeat(olabels,m)` with weight `w^n` for any integer `n`.

If `star` is `true`, `cfst` will transduce the empty string to itself with weight one.

"""
function closure(fst; star=true)
  fst_copy = deepcopy(fst)
  closure!(fst_copy; star=star)
  return fst_copy
end

function closure_plus!(fst::WFST{W,A,D,I,O}) where {W,A,D,I,O}
  ϵi, ϵo = get_eps(I), get_eps(O)
  i = get_initial(fst; single=true)
  for k in keys(get_final(fst; single=false))
    w = get_final_weight(fst,k)
    add_arc!(fst,k,i,ϵi,ϵo,w)
  end
  return fst
end

function closure_star!(fst::WFST{W,A,D,I,O}) where {W,A,D,I,O}
  ϵi, ϵo = get_eps(I), get_eps(O)
  closure_plus!(fst)
  i = get_initial(fst,single=true)
  w = get_initial_weight(fst,i)
  rm_initial!(fst,i)
  ni = size(fst,1)+1
  add_arc!(fst,ni,i,ϵi,ϵo,w)
  initial!(fst,ni)
  final!(fst,ni)
  return fst
end


text:
# Initialize environment in current directory
@info("Ensuring example environment instantiated...")
import Pkg
Pkg.activate(@__DIR__)
Pkg.instantiate()

@info("Loading Zygote...")
using Zygote, LinearAlgebra

# This example will showcase how we do a simple linear fit with Zygote, making
# use of complex datastructures, a home-grown stochastic gradient descent
# optimizer, and some good old-fashioned math.  We start with the problem
# statement:  We wish to learn the mapping `f(X) -> Y`, where `X` is a matrix
# of vector observations, `f()` is a linear mapping function and `Y` is a
# vector of scalar observations.

# Because we like complex objects, we will define our linear regression as the
# following object:
mutable struct LinearRegression
    # These values will be implicitly learned
    weights::Matrix
    bias::Float64

    # These values will not be learned
    name::String
end
LinearRegression(nparams, name) = LinearRegression(randn(1, nparams), 0.0, name)

# Our linear prediction looks very familiar; w*X + b
function predict(model::LinearRegression, X)
    return model.weights * X .+ model.bias
end

# Our "loss" that must be minimized is the l2 norm  between our current
# prediction and our ground-truth Y
function loss(model::LinearRegression, X, Y)
    return norm(predict(model, X) .- Y, 2)
end

# Our "ground truth" values (that we will learn, to prove that this works)
weights_gt = [1.0, 2.7, 0.3, 1.2]'
bias_gt = 0.4

# Generate a dataset of many observations
X = randn(length(weights_gt), 10000)
Y = weights_gt * X .+ bias_gt

# Add a little bit of noise to `X` so that we do not have an exact solution,
# but must instead do a least-squares fit:
X .+= 0.001.*randn(size(X))

# Now we begin our "training loop", where we take examples from `X`,
# calculate loss with respect to the corresponding entry in `Y`, find the
# gradient upon our model, update the model, and continue.  Before we jump
# in, let's look at what `Zygote.gradient()` gives us:
@info("Building model...")
model = LinearRegression(size(X, 1), "Example")

# Calculate gradient upon `model` for the first example in our training set
@info("Calculating gradient (the first time can take a while to compile...)")
grads = Zygote.gradient(model) do m
    return loss(m, X[:,1], Y[1])
end

# The `grads` object is a Tuple containing one element per argument to
# `gradient()`, so we take the first one to get the gradient upon `model`:
grads = grads[1]

# Because our LinearRegression object is mutable, the gradient holds a
# reference to it, which we peel via `grads[]`:
grads = grads[]

# We now get a `NamedTuple` so we can now do things like `grads.weight`. Let's
# print it out, just to see what it looks like.  Note that while `weights` and
# `bias` have gradients, `name` just naturally has a  gradient of `nothing`,
# because it was not involved in the calculation of the output loss.
@info grads

# Let's define an update rule that will allow us to modify the weights
# of our model a tad bit according to the gradients
function sgd_update!(model::LinearRegression, grads, η = 0.001)
    model.weights .-= η .* grads.weights
    model.bias -= η * grads.bias
end

# Now let's do that for each example in our training set:
@info("Running train loop for $(size(X,2)) iterations")
for idx in 1:size(X, 2)
    grads = Zygote.gradient(m -> loss(m, X[:, idx], Y[idx]), model)[1][]
    sgd_update!(model, grads)
end

# Now let's look at how well we've approximated the ground truth weights/bias:
@info("Ground truth weights: $(weights_gt)")
@info("Learned weights: $(round.(model.weights; digits=3))")
@info("Ground truth bias: $(bias_gt)")
@info("Learned bias: $(round(model.bias; digits=3))")


text:
# Make a dense neural network layer
make_layer(num_in, num_out, activation=Id()) = Layer(rand(num_out, num_in), rand(num_out), activation)

collect_params(layer::Layer) = vcat(vec(layer.weights), layer.bias)
collect_params(network::Network) = mapreduce(collect_params, vcat, network.layers)

function set_params!(layer::Layer, values)
    n_weights = length(layer.weights)
    layer.weights[:] .= @view values[1:n_weights]
    layer.bias .= @view values[n_weights+1:end]
    return layer
end
function set_params!(network::Network, values)
    i_begin = 1
    for layer in network.layers
        i_end = i_begin + length(layer.weights) + length(layer.bias) - 1
        set_params!(layer, @view(values[i_begin:i_end]))
        i_begin = i_end+1
    end
    return network
end

Base.@kwdef mutable struct PD{T}
    kp::T = 0.0
    kd::T = 0.0
end

(controller::PD)(x) = controller.kp*x[1] + controller.kd*x[2]

collect_params(controller::PD) = [controller.kp, controller.kd]
function set_params!(controller::PD, values)
    controller.kp = values[1]
    controller.kd = values[2]
    return controller
end


text:
using Test

using POMDPs

mutable struct X <: POMDP{Float64,Bool,Int} end
abstract type Z <: POMDP{Float64,Bool,Int} end
mutable struct Y <: Z end

@test_throws ErrorException statetype(Int)
@test_throws ErrorException actiontype(Int)
@test_throws ErrorException obstype(Int)

@test statetype(X) == Float64
@test statetype(Y) == Float64
@test actiontype(X) == Bool
@test actiontype(Y) == Bool
@test obstype(X) == Int
@test obstype(Y) == Int


text:
### MCRange

abstract MCRange{T<:Integer} <: Range{T}


text:
# Color is from 0.0 to 1.0
type DebugColor
    r::Float32
    g::Float32
    b::Float32
    a::Float32
end

function sfColor(debugcolor)
    SFML.Color(Uint8(floor(debugcolor.r * 255)), Uint8(floor(debugcolor.g * 255)), Uint8(floor(debugcolor.b * 255)), Uint8(floor(debugcolor.a * 255)))
end

function debug_draw_circle(pos::Vect, angle::Cdouble, radius::Cdouble, outlinecolor::DebugColor, fillcolor::DebugColor, data::Ptr{Void})
    SFML.set_radius(debug_circle, radius)
    SFML.set_origin(debug_circle, SFML.Vector2f(radius, radius))
    SFML.set_position(debug_circle, SFML.Vector2f(pos.x, -pos.y))
    # SFML.set_outlinecolor(debug_circle, sfColor(outlinecolor))
    SFML.set_fillcolor(debug_circle, sfColor(fillcolor))
    # SFML.set_outline_thickness(debug_circle, 0)

    SFML.set_points(debug_line, SFML.Vector2f(pos.x, -pos.y), SFML.Vector2f(pos.x + radius * sin(angle), -pos.y + radius * cos(angle)))
    SFML.set_thickness(debug_line, 2)
    SFML.set_fillcolor(debug_line, SFML.Color(128, 85, 85))

    window = SFML.RenderWindow(data)
    SFML.draw(window, debug_circle)
    SFML.draw(window, debug_line)
    return nothing
end

function debug_draw_segment(a::Vect, b::Vect, color::DebugColor, data::Ptr{Void})
    SFML.set_points(debug_line, SFML.Vector2f(a.x, -a.y), SFML.Vector2f(b.x, -b.y))
    SFML.set_thickness(debug_line, 1)
    SFML.set_fillcolor(debug_line, sfColor(color))

    window = SFML.RenderWindow(data)
    SFML.draw(window, debug_line)
    return nothing
end

function debug_draw_fatsegment(a::Vect, b::Vect, radius::Cdouble, outlinecolor::DebugColor, data::Ptr{Void})
    fillcolor = color_for_shape(C_NULL, C_NULL)

    SFML.set_points(debug_line, SFML.Vector2f(a.x, -a.y), SFML.Vector2f(b.x, -b.y))
    SFML.set_thickness(debug_line, radius*2)
    # SFML.set_outlinecolor(debug_line, sfColor(outlinecolor))
    # SFML.set_outline_thickness(debug_line, 0)
    SFML.set_fillcolor(debug_line, sfColor(fillcolor))

    window = SFML.RenderWindow(data)
    SFML.draw(window, debug_line)
    return nothing
end

function debug_draw_polygon(count::Int32, verts::Ptr{Vect}, radius::Cdouble, outlinecolor::DebugColor, fillcolor::DebugColor, data::Ptr{Void})
    SFML.set_pointcount(debug_polygon, count)

    for i = 0:count-1
        vert = unsafe_load(verts, i+1)
        SFML.set_point(debug_polygon, i, SFML.Vector2f(vert.x, -vert.y))
    end

    SFML.set_fillcolor(debug_polygon, sfColor(fillcolor))
    SFML.set_outlinecolor(debug_polygon, sfColor(outlinecolor))
    SFML.set_outline_thickness(debug_polygon, 0)

    window = SFML.RenderWindow(data)
    SFML.draw(window, debug_polygon)

    return nothing
end

function debug_draw_dot(size::Cdouble, pos::Vect, color::DebugColor, data::Ptr{Void})
    SFML.set_radius(debug_dot, size)
    SFML.set_position(debug_dot, SFML.Vector2f(pos.x, -pos.y))
    SFML.set_fillcolor(debug_dot, sfColor(color.r * 255, color.g * 255, color.b * 255, color.a * 255))

    window = SFML.RenderWindow(data)
    SFML.draw(window, debug_dot)
    return nothing
end

function color_for_shape(shape_ptr::Ptr{Void}, data::Ptr{Void})
    if shape_ptr != C_NULL
        shape = CircleShape(shape_ptr)
        body = get_body(shape)

        if is_sleeping(body)
            # Draw in gray
            return DebugColor(0.4, 0.4, 0.4, 0.6)
        else
            # Draw in red
            return DebugColor(0.8, 0.0, 0.0, 0.6)
        end
    else
        return DebugColor(0.8, 0.0, 0.0, 0.6)
    end
end

function DebugDrawOptions(window)
    c_draw_circle = cfunction(debug_draw_circle, Void, (Vect, Cdouble, Cdouble, DebugColor, DebugColor, Ptr{Void}))
    c_draw_segment = cfunction(debug_draw_segment, Void, (Vect, Vect, DebugColor, Ptr{Void}))
    c_draw_fatsegment = cfunction(debug_draw_fatsegment, Void, (Vect, Vect, Cdouble, DebugColor, Ptr{Void}))
    c_draw_polygon = cfunction(debug_draw_polygon, Void, (Int32, Ptr{Vect}, Cdouble, DebugColor, DebugColor, Ptr{Void}))
    c_draw_dot = cfunction(debug_draw_dot, Void, (Cdouble, Vect, DebugColor, Ptr{Void}))
    c_color_for_shape = cfunction(color_for_shape, DebugColor, (Ptr{Void}, Ptr{Void}))

    ccall(dlsym(libchipmunkjl, :cpjlDebugDrawOptions), Ptr{Void},
    (Ptr{Void}, Ptr{Void}, Ptr{Void}, Ptr{Void}, Ptr{Void},
    Cint, DebugColor, Ptr{Void}, DebugColor, DebugColor,
    Ptr{Void},),
    c_draw_circle,
    c_draw_segment,
    c_draw_fatsegment,
    c_draw_polygon,
    c_draw_dot,
    (1<<0) | (1<<1) | (1<<2),
    DebugColor(0.0, 0.0, 1.0, 1.0),
    c_color_for_shape,
    DebugColor(0.0, 1.0, 0.0, 1.0),
    DebugColor(1.0, 1.0, 0.0, 1.0),
    window.ptr)
end

function debug_draw(space::Space, window::SFML.RenderWindow; clear_and_display=false)
    if use_debug_draw
        event = SFML.Event()
        SFML.pollevent(window, event)
        options = DebugDrawOptions(window)

        if clear_and_display
            SFML.clear(window, SFML.white)
        end

        ccall(dlsym(libchipmunk, :cpSpaceDebugDraw), Void, (Ptr{Void}, Ptr{Void},), space.ptr, options)

        if clear_and_display
            SFML.display(window)
        end
    else
        println("Debug draw is unavailable.")
    end
end

export debug_draw


text:
module ArrayInterfaceBandedMatrices

using ArrayInterfaceCore
using BandedMatrices

struct BandedMatrixIndex <: ArrayInterfaceCore.MatrixIndex
    count::Int
    rowsize::Int
    colsize::Int
    bandinds::Array{Int,1}
    bandsizes::Array{Int,1}
    isrow::Bool
end

Base.firstindex(i::BandedMatrixIndex) = 1
Base.lastindex(i::BandedMatrixIndex) = i.count
Base.length(i::BandedMatrixIndex) = lastindex(i)
Base.@propagate_inbounds function Base.getindex(ind::BandedMatrixIndex, i::Int)
    @boundscheck 1 <= i <= ind.count || throw(BoundsError(ind, i))
    _i = i
    p = 1
    while _i - ind.bandsizes[p] > 0
        _i -= ind.bandsizes[p]
        p += 1
    end
    bandind = ind.bandinds[p]
    startfromone = !xor(ind.isrow, (bandind > 0))
    if startfromone
        return _i
    else
        return _i + abs(bandind)
    end
end

function _bandsize(bandind, rowsize, colsize)
    -(rowsize - 1) <= bandind <= colsize - 1 || throw(ErrorException("Invalid Bandind"))
    if (bandind * (colsize - rowsize) > 0) & (abs(bandind) <= abs(colsize - rowsize))
        return min(rowsize, colsize)
    elseif bandind * (colsize - rowsize) <= 0
        return min(rowsize, colsize) - abs(bandind)
    else
        return min(rowsize, colsize) - abs(bandind) + abs(colsize - rowsize)
    end
end

function BandedMatrixIndex(rowsize, colsize, lowerbandwidth, upperbandwidth, isrow)
    upperbandwidth > -lowerbandwidth || throw(ErrorException("Invalid Bandwidths"))
    bandinds = upperbandwidth:-1:-lowerbandwidth
    bandsizes = [_bandsize(band, rowsize, colsize) for band in bandinds]
    BandedMatrixIndex(sum(bandsizes), rowsize, colsize, bandinds, bandsizes, isrow)
end

function ArrayInterfaceCore.findstructralnz(x::BandedMatrices.BandedMatrix)
    l, u = BandedMatrices.bandwidths(x)
    rowsize, colsize = Base.size(x)
    rowind = BandedMatrixIndex(rowsize, colsize, l, u, true)
    colind = BandedMatrixIndex(rowsize, colsize, l, u, false)
    return (rowind, colind)
end

ArrayInterfaceCore.has_sparsestruct(::Type{<:BandedMatrices.BandedMatrix}) = true
ArrayInterfaceCore.isstructured(::Type{<:BandedMatrices.BandedMatrix}) = true
ArrayInterfaceCore.fast_matrix_colors(::Type{<:BandedMatrices.BandedMatrix}) = true

function ArrayInterfaceCore.matrix_colors(A::BandedMatrices.BandedMatrix)
    l, u = BandedMatrices.bandwidths(A)
    width = u + l + 1
    return ArrayInterfaceCore._cycle(1:width, Base.size(A, 2))
end

end # module


text:
using StructuresKit

#Compare Julia solution with LS-DYNA thin-shell FEA
#Example 2 from Plaut, R.H., Moen, C.D.(2020). "Lateral-Torsional Deformations of C-Section and Z-Section Beams with Continuous Bracing".  Proceedings of the Structural Stability Research Council Annual Conference, Atlanta, Georgia.

#Z-section, no slope, simple span
#ax=0, ay=0
#kϕ=0, kx=0, gravity load

#inputs

#Ix Iy Ixy J Cw
section_properties = [(3.230E6,449530,-865760, 397.09, 3.4104E9)]

#E  ν
material_properties = [(200,0.30)]

#ax ay
load_location = [(0.0, 0.0)]

#kx kϕ
spring_stiffness = [(0.0, 0.0)]

#ay_kx
spring_location = [(0.0)]

#member information
#L dL SectionProperties MaterialProperties LoadLocation SpringStiffness RoofSlope
member_definitions = [(7620,7620/12,1,1,1,1,1)]

#end boundary conditions
#type=1 u''=v''=ϕ''=0 (simply supported), type=2 u'=v'=ϕ'=0  (fixed), type=3 u''=v''=ϕ''=u'''=v'''=ϕ'''=0 (free end, e.g., a cantilever)
end_boundary_conditions = [1 1]

#supports
#location where u=v=ϕ=0
supports = [0.0 7620]

load = (0.0, 0.008/1000)  #kN/mm

z, u, v, ϕ, beam_properties = ThinWalledBeam.solve(member_definitions, section_properties, material_properties, spring_stiffness, spring_location, supports, load, load_location, end_boundary_conditions)

umax=maximum(u)
vmax=maximum(v)
ϕmax=maximum(ϕ)

#results from LS-DYNA
umax_LSDYNA = 2.19 #at Node 16556 in LS-DYNA model
vmax_LSDYNA = 1.14 #at Node 16556 in LS-DYNA model

#(y disp. @ Node 21372 - y disp. @ Node 16556)/(z loc @ Node 21372 - zloc @ Node 11740)
ϕmax_LSDYNA = atan((2.26 - 2.19)/(195.936 - 100.851))

#calculate percent error
uError = abs.((umax .- umax_LSDYNA))./ umax_LSDYNA
vError = abs.((vmax .- vmax_LSDYNA))./ vmax_LSDYNA

#accept 1.5% error from numerical solution
@test uError <= 0.015
@test vError <= 0.015

#PlautBeam predicts ϕ=0, judgement call here that 0.0007 rad from LS-DYNA is small
@test ϕmax ≈ ϕmax_LSDYNA atol=ϕmax_LSDYNA


text:
#####
##### Add missing manifolds and methods for common manifolds
#####
##### These are temporary, until Manifolds.jl provides the same functionality
#####

#####
##### Euclidean space with default Euclidean metric
#####

function geodesic_flow!(M::Euclidean, x, v, x₀, v₀, t)
      @. x = x₀ + t * v₀
      copyto!(v, v₀)
      return (x, v)
end

function geodesic_flow(M::Euclidean, x₀, v₀, t)
      return (x₀ .+ t .* v₀, v₀)
end

cogeodesic_flow!(M::Euclidean, args...) = geodesic_flow!(M, args...)

cogeodesic_flow(M::Euclidean, args...) = geodesic_flow(M, args...)

#####
##### Sphere with default ambient Euclidean metric
#####

function geodesic_flow!(M::Sphere, x, v, x₀, v₀, t)
      α = norm(M, x₀, v₀)
      αt = α * t

      sinαt, cosαt = sincos(αt)
      sincαtoα = α == 0 ? one(α) : sinαt / α

      @. x = cosαt * x₀ + sincαtoα * v₀
      @. v = cosαt * v₀ - (α * sinαt) * x₀

      project_point!(M, x)
      project_tangent!(M, v, x, v)

      return (x, v)
end

cogeodesic_flow!(M::Sphere, args...) = geodesic_flow!(M, args...)


text:
using .RCall
# import CorticalSpecrtralTemporalResponses: rplot, raster_plot
export rplot, scale_plot, cohere_plot

R"""
library(ggplot2)
"""

function titlefn(λ,λ_digits,components)
  @assert !any(isnan,λ)
  digits = λ_digits != :automatic ?  λ_digits :
    min(floor(Int,-log10(clamp(λ[minimum(components)],1e-8,1e8)))+2,
        floor(Int,-log10(clamp(λ[maximum(components)],1e-8,1e8)))+1)

  function title(n)
    if n <= length(λ)
      # nstr = @sprintf("%02d",n)
      "lambda[$(n)] == $(round(λ[n],digits=digits))"
      # "Lmb_$nstr = $(round(λ[n],3))"
    else
      "sigma^2 == $(round(sum(var(C)),digits=digits))"
      # "Variance = $(sum(var(C)))"
    end
  end
end

function nearin(xin,xs)
  _,inds = findmin(abs.(vec(xin) .- vec(xs)'),2)
  cols = map(ii -> ii[2],CartesianIndices((length(xin),length(xs))))
  if length(inds) > 0
    cols[inds[:,1]]
  else
    Array{Int}(0)
  end
end

function findnear(x,nearby)
  indices = sort(unique(nearin(filter(!isnan,nearby),
                                filter(!isnan,x))))
  if any(isnan.(nearby))
    [NaN; x[indices]], [find(isnan,x); indices]
  else
    x[indices], indices
  end
end

function cohere_plot(slice::AxisArray{T,3} where T; kwds...)
  @assert axisdim(slice,Axis{:scale}) == 1
  @assert axisdim(slice,Axis{:freq}) == 2
  @assert axisdim(slice,Axis{:component}) == 3

  ii = CartesianIndices(size(slice))
  at(i) = map(ii -> ii[i],ii)

  df = DataFrame(value = vec(slice),
                 scale = vec(round.(ustrip.(uconvert.(cycoct,
                                                      scales(slice)[at(1)])),
                                    digits=2)),
                 freq_bin = vec(at(2)),
                 component = vec(at(3)))

  fbreaks,findices = freq_ticks(slice)
  p = raster_plot(df;value=:value,x=:scale,y=:freq_bin,kwds...)

R"""
  $p + facet_grid(component~.) + #ordered_scales(scale)) +
    scale_y_continuous(breaks=$findices,labels=$fbreaks) +
    ylab('Frequency (Hz)') + xlab('Time (s)')
"""
end

# function rplot(C::Coherence{<:Any,<:Any,<:Any,4} where {M,T};λ_digits=:automatic,
#                kwds...)
#   C = sort_components(C)
#   @assert axisdim(C,Axis{:time}) == 1
#   @assert axisdim(C,Axis{:scale}) == 2
#   @assert axisdim(C,Axis{:freq}) == 3
#   @assert axisdim(C,Axis{:component}) == 4

#   ii = CartesianIndices(size(C))
#   at(i) = map(ii -> ii[i],ii)

#   rowtitle = titlefn(component_means(C),λ_digits,Base.axes(C,4))

#   df = DataFrame(value = vec(C),
#                  time = vec(ustrip.(uconvert.(s,times(C)[at(1)]))),
#                  scale = vec(round.(ustrip.(uconvert.(cycoct,scales(C)[at(2)])),
#                                     digits=2)),
#                  freq_bin = vec(at(3)),
#                  component = vec(at(4)),
#                  component_title = rowtitle.(vec(at(4))))

#   fbreaks,findices = freq_ticks(C)
#   p = raster_plot(df;value=:value,x=:time,y=:freq_bin,kwds...)

# R"""
#   scalevals = $(ustrip.(uconvert.(cycoct,scales(C))))
#   scalestr = function(x){sprintf("'Scale: %3.2f cyc/oct'",x)}
#   ordered_scales = function(x){
#     factor(scalestr(x),levels=scalestr(scalevals))
#   }

#   $p + facet_grid(ordered_scales(scale)~component_title,labeller=label_parsed) +
#   scale_y_continuous(breaks=$findices,labels=$fbreaks) +
#   ylab('Frequency (Hz)') + xlab('Time (s)')
# """
# end

function rplot(S::SourceTracking;λ_digits=:automatic,
               kwds...)
  S = sort_components(S)
  @assert axisdim(S,Axis{:scale}) == 1
  @assert axisdim(S,Axis{:freq}) == 2
  @assert axisdim(S,Axis{:component}) == 3
  @assert axisdim(S,Axis{:time}) == 4

  ii = CartesianIndices(size(S))
  at(i) = map(ii -> ii[i],ii)

  rowtitle = titlefn(component_means(S),λ_digits,Base.axes(S,3))

  df = DataFrame(value = vec(S),
                 time = vec(ustrip.(uconvert.(s,times(S)[at(4)]))),
                 scale = vec(round.(ustrip.(uconvert.(cycoct,scales(S)[at(1)])),
                                    digits=2)),
                 freq_bin = vec(at(2)),
                 component = vec(at(3)),
                 component_title = rowtitle.(vec(at(3))))

  fbreaks,findices = freq_ticks(S)
  p = raster_plot(df;value=:value,x=:time,y=:freq_bin,kwds...)

R"""
  scalevals = $(ustrip.(uconvert.(cycoct,scales(S))))
  scalestr = function(x){sprintf("'Scale: %3.2f cyc/oct'",x)}
  ordered_scales = function(x){
    factor(scalestr(x),levels=scalestr(scalevals))
  }

  $p + facet_grid(ordered_scales(scale)~component_title,labeller=label_parsed) +
  scale_y_continuous(breaks=$findices,labels=$fbreaks) +
  ylab('Frequency (Hz)') + xlab('Time (s)')
"""
end

function rplot(tracks::AxisArray{<:SourceTracking},component=1;kwds...)
  data = AxisArray(cat((sort_components(track)[:,:,component,:]
                        for track in tracks)...,dims=4),
                   AxisArrays.axes(tracks[1])[[1,2,4]]...,
                   Axis{:track}(1:length(tracks)))

  ii = CartesianIndices(size(data))
  at(i) = vec(map(ii -> ii[i],ii))
  track_vals = axisvalues(tracks)[1]

  df = DataFrame(value = vec(data),
                 time = vec(ustrip.(uconvert.(s,times(data)[at(3)]))),
                 scale = vec(round.(ustrip.(uconvert.(cycoct,
                                                      scales(data)[at(1)])),
                                    digits=2)),
                 freq_bin = vec(at(2)),
                 tau = vec([ustrip(τ) for (τ,_,__) in track_vals[at(4)]]),
                 sigma = vec([ustrip(σ) for (_,σ,__) in track_vals[at(4)]]),
                 N = vec([N for (_,__,N) in track_vals[at(4)]]))

  # @show(df[1:5,:])

  fbreaks,findices = freq_ticks(tracks[1])
  p = raster_plot(df;value=:value,x=:time,y=:freq_bin,kwds...)

R"""
  scalevals = $(ustrip.(uconvert.(cycoct,scales(data))))
  scalestr = function(x){sprintf("'Scale: %3.2f cyc/oct'",x)}
  ordered_scales = function(x){
    factor(scalestr(x),levels=scalestr(scalevals))
  }

  $p + facet_grid(ordered_scales(scale)~tau+sigma+N,labeller=label_both) +
  scale_y_continuous(breaks=$findices,labels=$fbreaks) +
  ylab('Frequency (Hz)') + xlab('Time (s)')
"""
end

function rplot(tracks::AxisArray{<:Coherence},component=1;kwds...)
  data = AxisArray(cat(4,(track[:,:,:,component] for track in tracks)...),
                   AxisArrays.axes(tracks[1])[1:3]...,Axis{:track}(1:length(tracks)))

  ii = CartesianIndices(size(data))
  at(i) = vec(map(ii -> ii[i],ii))
  track_vals = axisvalues(tracks)[1]
  # @show track_vals
  # @show unique(at(4))

  df = DataFrame(value = vec(data),
                 time = vec(ustrip.(uconvert.(s,times(data)[at(1)]))),
                 scale = vec(round.(ustrip.(uconvert.(cycoct,
                                                      scales(data)[at(2)])),
                                    digits=2)),
                 freq_bin = vec(at(3)),
                 tau = vec([ustrip(τ) for (τ,_) in track_vals[at(4)]]),
                 sigma = vec([ustrip(σ) for (_,σ) in track_vals[at(4)]]))
  # @show(df[1:5,:])

  fbreaks,findices = freq_ticks(tracks[1])
  p = raster_plot(df;value=:value,x=:time,y=:freq_bin,kwds...)

R"""
  scalevals = $(ustrip.(uconvert.(cycoct,scales(data))))
  scalestr = function(x){sprintf("'Scale: %3.2f cyc/oct'",x)}
  ordered_scales = function(x){
    factor(scalestr(x),levels=scalestr(scalevals))
  }

  $p + facet_grid(ordered_scales(scale)~tau+sigma,labeller=label_parsed) +
  scale_y_continuous(breaks=$findices,labels=$fbreaks) +
  ylab('Frequency (Hz)') + xlab('Time (s)')
"""
end

#=
function rplot(cohere::CoherenceModel,C::Factors;
               components=1:ncomponents(C),λ_digits=:automatic)
  u = factors(C)[:,components]
  u = reshape(u,length(scales(cohere)),:,size(u,2))

  ii = CartesianIndices(size(u))
  at(i) = vec(map(ii -> ii[i],ii))

  title = titlefn(strengths(C),λ_digits,components)

  df = DataFrame(response = vec(u),
                 scale_index = at(1),
                 freq_bin = at(2),
                 component = at(3),
                 component_title = title.(at(3)))

  sindices = 1:2:length(scales(cohere))
  sbreaks = round.(scales(cohere)[sindices],2)
  fbreaks,findices = freq_ticks(cohere)

  p = raster_plot(df,value=:response,x=:scale_index,y=:freq_bin)

R"""

  library(ggplot2)

  $p + facet_wrap(~component_title,labeller=label_parsed) +
    scale_y_continuous(breaks=$findices,labels=$fbreaks) +
    scale_x_continuous(breaks=$sindices,labels=$sbreaks) +
    ylab('Frequency (Hz)') + xlab('Scale (cycles/octave)')
"""

end

function rplot(tempc::CoherenceModel,C::AbstractArray{<:FactorSeries},
               label=:index => 1:length(C))
  x = [fusion_signal(tempc,Ci) for Ci in C]
  df = DataFrame(resp = vcat((real.(xi) for xi in x)...),
                 var = vcat((fill(label[2][i],length(x[i]))
                             for i in eachindex(x))...),
                             time = vcat((ustrip.(uconvert.(s,eachindex(xi) * Δt(tempc)))
                              for xi in x)...))

R"""
  library(ggplot2)

  ggplot($df,aes(x=time,y=resp,group=factor(var),color=factor(var))) +
    geom_line() +
    scale_color_brewer(palette='Set1',name=$(string(label[1]))) +
    coord_cartesian(ylim=c(1.1,0)) + ylab('lambda / var(x)') +
    xlab('time (s)')
"""
end

function strength_plot(cohere::CoherenceModel,C::EigenSeries;n=ncomponents(C))
  λ = eigvals(C)
  λ = λ[:,sortperm(abs.(λ[max(1,end-10),:]),rev=true)]
  ii = CartesianIndices(size(λ))
  at(i) = map(ii -> ii[i],ii)

  prop_complex = mean(imag.(λ))
  if prop_complex > 0
    warn("$(round(100prop_complex,1))% of eigenvalues were "*
         "complex (using absolute value).")
  end

  df = DataFrame(value = vec(abs.(λ) ./ sum.(var.(C))),
  time = vec(ustrip(uconvert.(s,at(1) * Δt(cohere)))),
                 component = vec(at(2)))

  df = df[df[:component] .<= n,:]

R"""
  ggplot($df,aes(x=time,y=value,color=factor(component),group=component)) +
    geom_line() + scale_color_brewer(palette='Set1',name='Component') +
    xlab('Time (s)') + ylab('Value')
"""
end

function rplot(cohere::CoherenceModel,λ::Vector)
  @assert all(imag.(λ) .== 0) "Can't plot complex eigenvalues"
  λ = real.(λ)
  df = DataFrame(value = sort(λ,rev=true),index = collect(eachindex(λ)))

R"""
  ggplot($df,aes(x=index,y=value)) + geom_bar(stat='identity')
"""
end
=#

max_stars_repo_name: JuliaTagBot/AuditoryBistabilityLE

max_issues_repo_name: JuliaTagBot/AuditoryBistabilityLE

max_forks_repo_name: JuliaTagBot/AuditoryBistabilityLE


text:
GetGeodesicODE(Metric::Function, InitialPos::AbstractVector{<:Number}, approx::Bool=false; kwargs...) = GetGeodesicODE(Metric, InitialPos, Val(approx); kwargs...)
function GetGeodesicODE(Metric::Function, InitialPos::AbstractVector{<:Number}, approx::Val{false}; kwargs...)
    function GeodesicODE!(du,u,p,t)
        n = length(u)÷2
        du[1:n] = u[(n+1):end]
        du[(n+1):end] = ChristoffelTerm(ChristoffelSymbol(Metric,u[1:n]; kwargs...),du[1:n])
    end
end
function GetGeodesicODE(Metric::Function, InitialPos::AbstractVector{<:Number}, approx::Val{true}; kwargs...)
    Γ = ChristoffelSymbol(Metric, InitialPos; kwargs...)
    function ApproxGeodesicODE!(du,u,p,t)
        n = length(u)÷2
        du[1:n] = u[(n+1):end]
        du[(n+1):end] = ChristoffelTerm(Γ, du[1:n])
    end
end

# accuracy ≈ 6e-11
function ComputeGeodesic(Metric::Function, InitialPos::AbstractVector, InitialVel::AbstractVector, Endtime::Number=50.;
                        Boundaries::Union{Function,Nothing}=nothing, tol::Real=1e-11, meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), approx::Bool=false, kwargs...)
    @assert length(InitialPos) == length(InitialVel)
    prob = ODEProblem(GetGeodesicODE(Metric, InitialPos, approx), vcat(InitialPos,InitialVel), (0.0,Endtime))
    if isnothing(Boundaries)
        return solve(prob, meth; reltol=tol, abstol=tol, kwargs...)
    else
        return solve(prob, meth; reltol=tol, abstol=tol, callback=DiscreteCallback(Boundaries,terminate!), kwargs...)
    end
end

"""
    ComputeGeodesic(DM::DataModel, InitialPos::AbstractVector, InitialVel::AbstractVector, Endtime::Number=50.;
                                    Boundaries::Union{Function,Nothing}=nothing, tol::Real=1e-11, approx::Bool=false, meth=Tsit5())
Constructs geodesic with given initial position and velocity.
It is possible to specify a boolean-valued function `Boundaries(u,t,int)`, which terminates the integration process when it returns `true`.

By setting the keyword `approx=true`, the ChristoffelSymbols are assumed to be constant and only computed once at the initial position. This simplifies the computation immensely but may also constitute an inaccurate approximation depending on the magnitude of the Ricci curvature.
"""
ComputeGeodesic(DM::AbstractDataModel, args...; kwargs...) = ComputeGeodesic(FisherMetric(DM), args...; kwargs...)

"""
    GeodesicLength(DM::DataModel,sol::AbstractODESolution, Endrange::Number=sol.t[end]; FullSol::Bool=false, tol=1e-14)
    GeodesicLength(Metric::Function,sol::AbstractODESolution, Endrange::Number=sol.t[end]; FullSol::Bool=false, tol=1e-14)
Calculates the length of a geodesic `sol` using the `Metric` up to parameter value `Endrange`.
```math
L[\\gamma] \\coloneqq \\int_a^b \\mathrm{d} t \\, \\sqrt{g_{\\gamma(t)} \\big(\\dot{\\gamma}(t), \\dot{\\gamma}(t)\\big)}
```
"""
GeodesicLength(DM::AbstractDataModel, args...; kwargs...) = GeodesicLength(FisherMetric(DM), args...; kwargs...)
function GeodesicLength(Metric::Function, sol::AbstractODESolution, Endrange::Number=sol.t[end]; FullSol::Bool=false, tol::Real=1e-14, kwargs...)
    @assert length(sol.u[1]) % 2 == 0
    n = length(sol.u[1])÷2
    function Integrand(t)
        FullGamma = sol(t)
        InnerProduct(Metric(FullGamma[1:n]), FullGamma[(n+1):end]) |> sqrt
    end
    return Integrate1D(Integrand, (sol.t[1],Endrange); FullSol=FullSol, tol=tol, kwargs...)
end

"""
    GeodesicCrossing(DM::DataModel, sol::AbstractODESolution, Conf::Real=ConfVol(1); tol=1e-15)
Gives the parameter value of the geodesic `sol` at which the confidence level `Conf` is crossed.
"""
function GeodesicCrossing(DM::AbstractDataModel, sol::AbstractODESolution, Conf::Real=ConfVol(1); tol::Real=1e-15)
    start = sol.t[end]/2
    if (tol < 1e-15)
        start *= one(BigFloat)
        println("GeodesicCrossing: Conf value not programmed as BigFloat yet.")
    end
    A = loglikelihood(DM,sol(0.)[1:2]) - (1/2)*quantile(Chisq(Int(length(sol(0.))/2)),Conf)
    f(t) = A - loglikelihood(DM,sol(t)[1:2])
    find_zero(f,start,Order1B();xatol=tol)
end

"""
    DistanceAlongGeodesic(Metric::Function,sol::AbstractODESolution,L::Number; tol=1e-14)
Calculates at which parameter value of the geodesic `sol` the length `L` is reached.
"""
function DistanceAlongGeodesic(Metric::Function, sol::AbstractODESolution, L::Number; tol::Real=1e-14)
    L < 0 && throw(BoundsError("DistanceAlongGeodesic: L=$L"))
    # Use interpolated Solution of integral for improved accuracy
    GeoLength = GeodesicLength(Metric,sol,sol.t[end], FullSol=true, tol=tol)
    Func(x) = L - GeoLength(x)
    dFunc(x) = ForwardDiff.derivative(Func,x)
    find_zero((Func,dFunc),sol.t[end]/2*one(typeof(L)),Roots.Newton(),xatol=tol)
end

# Input Array of Geodesics, Output Array of its endpoints
function Endpoints(Geodesics::AbstractVector{<:AbstractODESolution})
    Endpoints = Vector{Vector{Float64}}(undef,0)
    Numb = Int(length(Geodesics[1].u[1])/2)
    for Curve in Geodesics
        T = Curve.t[end]
        push!(Endpoints,Curve(T)[1:Numb])
    end;    Endpoints
end

"""
    EvaluateEach(geos::AbstractVector{<:AbstractODESolution}, Ts::AbstractVector) -> Vector
Evalues a family `geos` of geodesics on a set of parameters `Ts`. `geos[1]` is evaluated at `Ts[1]`, `geos[2]` is evaluated at `Ts[2]` and so on.
The second half of the values respresenting the velocities is automatically truncated.
"""
function EvaluateEach(sols::AbstractVector{<:AbstractODESolution}, Ts::AbstractVector{<:Number})
    length(sols) != length(Ts) && throw(ArgumentError("Dimension Mismatch."))
    n = Int(length(sols[1].u[1])/2)
    Res = Vector{Vector{Float64}}(undef,0)
    for i in 1:length(Ts)
        F = sols[i]
        push!(Res,F(Ts[i])[1:n])
    end
    Res
end

# """
#     Truncated(sol::AbstractODESolution) -> Function
# Given a geodesic `sol`, the second half of the components which represent the velocity are truncated off.
# The result gives the position of the geodesic as a function of the parameter.
# However, since it is no longer of type `ODESolution`, one no longer has access to the fields `sol.t`, `sol.u` and so on.
# """
# Truncated(sol::AbstractODESolution) = (t->sol(t)[1:Int(length(sol.u[1])/2)])

################################################################ MERGE THESE FUNCTION VARIATIONS
# function ConstLengthGeodesics(DM::DataModel,Metric::Function,MLE::AbstractVector,Conf::Float64=ConfVol(1),N::Int=100)
#     angles = [2*pi*n/N      for n in 1:N]
#     sols = Vector{ODESolution}(undef,0)
#     for α in angles
#         InitialVelocity = [cos(α),sin(α)]
#         push!(sols,ConfidenceBoundaryViaGeodesic(DM,Metric,[MLE...,InitialVelocity...],Conf))
#         num = α*N/(2*pi) |> round |> Int
#         println("Calculated Geodesic $num of $N.")
#     end
#     Region = Endpoints(sols);   plot(Unpack(Region))
#     sols
# end

# ADAPT FOR PLANES
function ConstLengthGeodesics(DM::AbstractDataModel, Metric::Function, MLE::AbstractVector, Conf::Real=ConfVol(1), N::Int=100; tol::Real=6e-11)
    angles = [2π*n/N      for n in 1:N]
    Initials = [ [MLE...,cos(alpha),sin(alpha)] for alpha in angles]
    solving = 0
    function Constructor(Initial)
        solving += 1
        println("Computing Geodesic $(solving) / $N\t")
        ConfidenceBoundaryViaGeodesic(DM,Metric,Initial,Conf; tol=tol)
    end
    pmap(Constructor,Initials)
end

function BoundaryViaGeodesic(DM::AbstractDataModel, InitialPos::AbstractVector, InitialVel::AbstractVector,
                                    Confnum::Real=1, Endtime::Real=500.0; dof::Int=length(MLE(DM)), kwargs...)
    WilksCond = (1/2)*quantile(Chisq(dof),ConfVol(Confnum))
    BoundaryFunc(u,t,int) = LogLikeMLE(DM) - loglikelihood(DM, u[1:end÷2]) > WilksCond
    ComputeGeodesic(FisherMetric(DM), InitialPos, InitialVel, Endtime; Boundaries=BoundaryFunc, kwargs...)
end

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

function pConstParamGeodesics(Metric::Function,MLE::AbstractVector,Endtime::Number=10.,N::Int=100;
    Boundaries::Union{Function,Nothing}=nothing, tol::Real=1e-13, parallel::Bool=true)
    ConstParamGeodesics(Metric,MLE,Endtime,N;Boundaries=Boundaries, tol=tol, parallel=parallel)
end

function ConstParamGeodesics(Metric::Function,MLE::AbstractVector,Endtime::Number=10.,N::Int=100;
    Boundaries::Union{Function,Nothing}=nothing, tol::Real=1e-13, parallel::Bool=false)
    Initials = [ [cos(alpha),sin(alpha)] for alpha in range(0,2π;length=N)];    solving = 0
    Map = parallel ? pmap : map
    function Constructor(Initial)
        solving += 1
        println("Computing Geodesic $(solving) / $N")
        ComputeGeodesic(Metric,MLE,Initial,Endtime;tol=tol,Boundaries=Boundaries)
    end
    Map(Constructor,Initials)
end

GeodesicBoundaryFunction(Cube::HyperCube) = (u,p,t) -> u[1:end÷2] ∉ Cube
function GeodesicBoundaryFunction(M::ModelMap)
    function ModelMapBoundaries(u,p,t)
        S = !IsInDomain(M, u[1:end÷2])
        S && @warn "Geodesic ran into boundaries specified by ModelMap at $(u[1:end÷2])."
        return S
    end
end

# Also add Plane method!
function RadialGeodesics(DM::AbstractDataModel, Cube::HyperCube; N::Int=50, tol::Real=1e-9, Boundaries::Union{Function,Nothing}=nothing, parallel::Bool=false, kwargs...)
    @assert length(Cube) == 2 && MLE(DM) ∈ Cube
    widths = CubeWidths(Cube);    Map = parallel ? pmap : map;    Metric(x) = FisherMetric(DM, x)
    initialvels = [widths .* [cos(α), sin(α)] for α in range(0, 2π*(1-1/N); length=N)]
    CB = DiscreteCallback(GeodesicBoundaryFunction(Cube))
    CB = !isnothing(Boundaries) ? CallbackSet(CB, DiscreteCallback(Boundaries,terminate!)) : CB
    CB = Predictor(DM) isa ModelMap ? CallbackSet(CB, DiscreteCallback(GeodesicBoundaryFunction(Predictor(DM)),terminate!)) : CB
    solving = 0
    function Constructor(InitialVel)
        solving += 1
        println("Computing Geodesic $(solving) / $N")
        # Already added Boundaries(u,p,t) function to callbacks if any was passed via kwarg
        ComputeGeodesic(Metric, MLE(DM), InitialVel, 10.0; tol=tol, Boundaries=nothing, callback=CB, kwargs...)
    end
    Map(Constructor, initialvels)
end

"""
    GeodesicBetween(DM::DataModel, P::AbstractVector{<:Number}, Q::AbstractVector{<:Number}; tol::Real=1e-10, meth=Tsit5())
    GeodesicBetween(Metric::Function, P::AbstractVector{<:Number}, Q::AbstractVector{<:Number}; tol::Real=1e-10, meth=Tsit5())
Computes a geodesic between two given points on the parameter manifold and an expression for the metric.

By setting the keyword `approx=true`, the ChristoffelSymbols are assumed to be constant and only computed once at the initial position. This simplifies the computation immensely but may also constitute an inaccurate approximation depending on the magnitude of the Ricci curvature.
"""
GeodesicBetween(DM::AbstractDataModel, args...; kwargs...) = GeodesicBetween(FisherMetric(DM), args...; kwargs...)
function GeodesicBetween(Metric::Function, P::AbstractVector{<:Number}, Q::AbstractVector{<:Number}, Endtime::Real=10.0; tol::Real=1e-9, meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), approx::Bool=false, kwargs...)
    length(P) != length(Q) && throw("GeodesicBetween: Points not of same dim.")
    dim = length(P)
    function bc!(resid, u, p, t)
        resid[1:dim] = u[1][1:dim] .- P
        resid[(dim+1):2dim] = u[end][1:dim] .- Q
    end
    # Slightly perturb initial direction:
    initial = vcat(P, ((Q - P) ./ Endtime) .+ 1e-8 .*(rand(dim) .- 0.5))
    BVP = BVProblem(GetGeodesicODE(Metric, P, approx), bc!, initial, (0.0, Endtime))
    solve(BVP, Shooting(meth); reltol=tol, abstol=tol, kwargs...)
end

"""
    GeodesicDistance(DM::DataModel, P::AbstractVector{<:Number}, Q::AbstractVector{<:Number}; tol::Real=1e-10)
    GeodesicDistance(Metric::Function, P::AbstractVector{<:Number}, Q::AbstractVector{<:Number}; tol::Real=1e-10)
Computes the length of a geodesic connecting the points `P` and `Q`.
"""
GeodesicDistance(DM::AbstractDataModel, args...; kwargs...) = GeodesicDistance(FisherMetric(DM), args...; kwargs...)
function GeodesicDistance(Metric::Function, P::AbstractVector{<:Number}, Q::AbstractVector{<:Number}; tol::Real=1e-10, meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), kwargs...)
    GeodesicLength(Metric, GeodesicBetween(Metric, P, Q; tol=tol, meth=meth, kwargs...))
end

ParamVol(sol::AbstractODESolution) = sol.t[end] - sol.t[1]
GeodesicEnergy(DM::DataModel,sol::AbstractODESolution,Endrange::Number=sol.t[end];FullSol::Bool=false,tol::Real=1e-14) = GeodesicEnergy(x->FisherMetric(DM,x),sol,Endrange;tol=tol)
function GeodesicEnergy(Metric::Function,sol::AbstractODESolution,Endrange=sol.t[end]; FullSol::Bool=false,tol::Real=1e-14)
    @assert length(sol.u[1]) % 2 == 0
    n = length(sol.u[1])÷2
    function Integrand(t)
        FullGamma = sol(t)
        InnerProduct(Metric(FullGamma[1:n]), FullGamma[(n+1):end])
    end
    Integrate1D(Integrand, [sol.t[1],Endrange]; FullSol=FullSol, tol=tol)
end

"""
    ExponentialMap(Metric::Function, point::AbstractVector{<:Number}, tangent::AbstractVector{<:Number}; tol::Real=1e-9)
Computes the differential-geometric exponential map ``\\mathrm{exp}_p(v)`` which returns the endpoint that is reached by a geodesic ``\\gamma:[0,1] \\longrightarrow \\mathcal{M}`` with initial direction ``v \\in T_p \\mathcal{M}``.
"""
function ExponentialMap(Metric::Function, point::AbstractVector{<:Number}, tangent::AbstractVector{<:Number}; tol::Real=1e-9, meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), FullSol::Bool=false, kwargs...)
    if FullSol
        ComputeGeodesic(Metric, point, tangent, 1.0; tol=tol, meth=meth, kwargs...)
    else
        ComputeGeodesic(Metric, point, tangent, 1.0; tol=tol, meth=meth, save_everystep=false, save_start=false, save_end=true, kwargs...).u[end][1:end÷2]
    end
end
ExponentialMap(DM::AbstractDataModel, args...; kwargs...) = ExponentialMap(FisherMetric(DM), args...; kwargs...)

"""
    LogarithmicMap(Metric::Function, P::AbstractVector{<:Number}, Q::AbstractVector{<:Number}; tol::Real=1e-9)
Computes the inverse of the differential-geometric exponential map, i.e. ``\\mathrm{ln}_p(q) \\equiv (\\mathrm{exp}^{-1})_p(q)`` which returns a (possibly non-unique!) initial direction ``v \\in T_p \\mathcal{M}`` for a geodesic ``\\gamma:[0,1] \\longrightarrow \\mathcal{M}`` that goes from ``p \\in \\mathcal{M}`` to ``q \\in \\mathcal{M}``.
"""
function LogarithmicMap(Metric::Function, P::AbstractVector{<:Number}, Q::AbstractVector{<:Number}; tol::Real=1e-9, meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), FullSol::Bool=false, kwargs...)
    if FullSol
        GeodesicBetween(Metric, P, Q, 1.0; tol=tol, meth=meth, kwargs...)
    else
        GeodesicBetween(Metric, P, Q, 1.0; tol=tol, meth=meth, save_everystep=false, save_start=true, save_end=false, kwargs...).u[1][((end÷2)+1):end]
    end
end
LogarithmicMap(DM::AbstractDataModel, args...; kwargs...) = LogarithmicMap(FisherMetric(DM), args...; kwargs...)

function KarcherMeanStep(Metric::Function, points::AbstractVector{<:AbstractVector{<:Number}}, initialmean::AbstractVector{<:Number}=sum(points)/length(points); kwargs...)
    dirs = map(x->LogarithmicMap(Metric, initialmean, x; kwargs...), points)
    ExponentialMap(Metric, initialmean, sum(dirs) / length(dirs))
end
function KarcherMean(Metric::Function, points::AbstractVector{<:AbstractVector{<:Number}}, initialmean::AbstractVector{<:Number}=sum(points)/length(points); tol::Real=1e-8, meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), maxiter::Int=10, kwargs...)
    @assert ConsistentElDims(points) == length(initialmean)
    oldmean = initialmean
    for iter in 1:maxiter
        println("Karcher Mean iteration $iter.")
        newmean = KarcherMeanStep(Metric, points, oldmean; tol=tol, meth=meth, kwargs...)
        if norm(newmean-oldmean) < 20tol
            return newmean
        elseif iter < maxiter
            oldmean = newmean
            continue;
        else
            @warn "KarcherMean: Hit maxiter. Returning last result."
            return newmean
        end
    end
end

# function MBAMBoundaries(u, t, int, DM; componentlim::Real=1e4, singularlim::Real=1e-8)::Bool
#     if !all(x->abs(x) < componentlim, u)
#         @warn "Terminated because a position / velocity coordinate > $componentlim at: $u."
#         return true
#     elseif svdvals(FisherMetric(DM,u[1:end÷2]))[end] < singularlim
#         @warn "Terminated because Fisher metric became singular (i.e. < $singularlim) at: $u."
#         return true
#     else
#         return false
#     end
# end

"""
Return `true` when integration of ODE should be terminated.
"""
function MBAMBoundaries(DM::AbstractDataModel; componentlim::Real=1e4, singularlim::Real=1e-8)
    function MBAMBoundary(u,t,int)
        if !all(x->abs(x) < componentlim, u)
            @warn "Terminated because a position / velocity coordinate > $componentlim at: $u."
            return true
        elseif svdvals(FisherMetric(DM, u[1:end÷2]))[end] < singularlim
            @warn "Terminated because Fisher metric became singular (i.e. < $singularlim) at: $u."
            return true
        else
            return false
        end
    end
end

function MBAM(DM::AbstractDataModel; Boundaries::Union{Function,Nothing}=nothing, tol::Real=1e-5, meth::OrdinaryDiffEqAlgorithm=GetMethod(tol), componentlim::Real=1e4, singularlim::Real=1e-8, kwargs...)
    InitialVel = normalize(LeastInformativeDirection(DM,MLE(DM)))
    CB = DiscreteCallback(MBAMBoundaries(DM; componentlim=componentlim, singularlim=singularlim), terminate!)
    CB = !isnothing(Boundaries) ? CallbackSet(CB, DiscreteCallback(Boundaries, terminate!)) : CB
    CB = Predictor(DM) isa ModelMap ? CallbackSet(CB, DiscreteCallback(GeodesicBoundaryFunction(Predictor(DM)),terminate!)) : CB
    # Already added Boundaries(u,p,t) function to callbacks if any was passed via kwarg
    ComputeGeodesic(DM, MLE(DM), InitialVel, 1e4; Boundaries=nothing, callback=CB, tol=tol, meth=meth, kwargs...)
    # if isnothing(Boundaries)
    #     MBAMboundary(u,t,int) = MBAMBoundaries(u,t,int,DM)
    #     return ComputeGeodesic(DM, MLE(DM), InitialVel, 1e4; Boundaries=MBAMboundary, tol=tol, meth=meth, kwargs...)
    # else
    #     CombinedBoundaries(u,t,int)::Bool = Boundaries(u,t,int) || MBAMBoundaries(u,t,int,DM)
    #     return ComputeGeodesic(DM, MLE(DM), InitialVel, 1e4; Boundaries=CombinedBoundaries, tol=tol, meth=meth, kwargs...)
    # end
end


text:
# Example script for basic frame connection class usage.

using CSV
using DataFrames
using BenchmarkTools
using FrameConnection

# Load some data and store it in an smd struct|ure.
filename = "C:\\Users\\David\\Documents\\work_stuff\\frame_connection\\actin_data.csv"
# filename = "C:\\Users\\David\\Documents\\work_stuff\\frame_connection\\uniform_rho10.csv"

data = DataFrames.DataFrame(CSV.File(filename))
smd = FrameConnection.SMD(data)

# Perform frame connection.
params = FrameConnection.ParamStruct()
smd_combined, smd, smd_preclustered = FrameConnection.frameconnect(smd)

# Make some images.
using SMLMData
using Images
using ImageView

circleim_in = SMLMData.makecircleim(Float64.([smd.y smd.x]), Float64.(vec(mean([smd.y_se smd.x_se], dims = 2))), [128; 128])
circleim_out = SMLMData.makecircleim(Float64.([smd_combined.y smd_combined.x]),
    Float64.(vec(mean([smd_combined.y_se smd_combined.x_se], dims = 2))), [128; 128])
circleim_in ./= maximum(circleim_in)
circleim_out ./= maximum(circleim_out)
testim = Images.RGB.(circleim_in, circleim_out, circleim_in)
ImageView.imshow(testim)


text:
#### Operators for a system with moving body ####

export rigid_body_rhs!, rigid_body_rhs, zero_body_state

zero_body_state(b::Body) = zeros(Float64,3*(NDIM-1))
zero_body_state(bl::BodyList) = zeros(Float64,3*length(bl)*(NDIM-1))
zero_body_state(sys::HeatConduction) = zero_body_state(sys.bodies)

## Right-hand sides for rigid-body motion equations ##

function rigid_body_rhs!(dx::Vector{T},x::Vector{T},sys::HeatConduction,t::Real) where {T<:Real}
  # note that x itself is not used here, because sys is presumed to contain
  # updated body data
  length(dx) == 3*length(sys.bodies)*(NDIM-1) || error("Wrong length for vector")
  dx .= rigidbodyvelocity(sys.motions,t)
  return dx
end

rigid_body_rhs(x::Vector{T},sys::HeatConduction,t::Real) where {T<:Real} =
              rigid_body_rhs!(zero_body_state(sys),sys,t)


text:
"""
scipy.interpolate module

- [Interpolation (scipy.interpolate) Reference Guide](https://docs.scipy.org/doc/scipy/reference/interpolate.html)

# Examples

You can interpolate 1D data:

```julia-repl
julia> x = collect(0:10);
julia> y = exp.(-x/3.0);
julia> f = SciPy.interpolate.interp1d(x, y);
julia> f(0.5)
0-dimensional Array{Float64,0}:
0.8582656552868946

```
"""
module interpolate

using PyCall
import PyCall: hasproperty # Base.hasproperty in Julia 1.2

import ..pyinterpolate
import .._generate_docstring
import ..LazyHelp

const _ignore_funcs = ["interpolate"]

# This list can be generated by print_scipy_api_list.py script
all_properties = ["Akima1DInterpolator", "BPoly", "BSpline", "BarycentricInterpolator", "BivariateSpline", "CloughTocher2DInterpolator", "CubicHermiteSpline", "CubicSpline", "InterpolatedUnivaria
teSpline", "KroghInterpolator", "LSQBivariateSpline", "LSQSphereBivariateSpline", "LSQUnivariateSpline", "LinearNDInterpolator", "NdPPoly", "NearestNDInterpolator", "PPoly", "Pch
ipInterpolator", "Rbf", "RectBivariateSpline", "RectSphereBivariateSpline", "RegularGridInterpolator", "SmoothBivariateSpline", "SmoothSphereBivariateSpline", "UnivariateSpline",
 "approximate_taylor_polynomial", "barycentric_interpolate", "bisplev", "bisplrep", "dfitpack", "fitpack", "fitpack2", "griddata", "insert", "interp1d", "interp2d", "interpn", "i
nterpnd", "interpolate", "krogh_interpolate", "lagrange", "make_interp_spline", "make_lsq_spline", "ndgriddata", "pade", "pchip_interpolate", "polyint", "rbf", "spalde", "splanti
der", "splder", "splev", "splint", "splprep", "splrep", "sproot"]

for f in all_properties
    f in _ignore_funcs && continue

    sf = Symbol(f)
    @eval @doc LazyHelp(pyinterpolate, $f) $sf(args...; kws...) = pycall(pyinterpolate.$f, PyAny, args...; kws...)
end

function __init__()
    copy!(pyinterpolate, pyimport_conda("scipy.interpolate", "scipy"))
end

end # module


text:
"""
    linear_search(array, key)
A simple search of `array`, element per element until `key` is found.
"""
function linear_search(array, key)
    for i in 1:length(array)
        if array[i] == key
            return i
        end
    end
    return -1
end

max_stars_repo_name: KohRongSoon/Julia

max_issues_repo_name: KohRongSoon/Julia

max_forks_repo_name: KohRongSoon/Julia


text:
import RayTracing: num_dims, num_cells

"""
    MoCProblem{Dim,NRegions,NGroups,elType,G<:TrackGenerator,Q<:Quadrature,M} <: TransportProblem{Dim,NRegions,NGroups}

Defines a problem that ought to be solved using the Method of Characteristics, with
dimension `Dim`, number of flat source regions `NRegions`, number of energy groups
`NGroups`, ray tracing
"""
struct MoCProblem{Dim,NRegions,NGroups,elType,G<:TrackGenerator,Q<:Quadrature,Xs<:XSs} <: TransportProblem{Dim,NRegions,NGroups}
    nφ::Int
    nψ::Int

    trackgenerator::G
    quadrature::Q

    xss::Vector{Xs}
    fsr_tag::Vector{Int8}  # cell/element/fsr to xs tag (tag is the index in the xss array)

    function MoCProblem{Dim,NRegions,NGroups,elType}(
        nφ, nψ, tg::G, quad::Q, xss::Vector{X}, fsr_tag
    ) where {Dim,NRegions,NGroups,elType,G,Q,X}
        return new{Dim,NRegions,NGroups,elType,G,Q,X}(nφ, nψ, tg, quad, xss, fsr_tag)
    end
end

function MoCProblem(tg::TrackGenerator, polar_quadrature::PolarQuadrature, xss::Vector{<:XSs})
    @unpack mesh, azimuthal_quadrature, n_total_tracks = tg

    Dim = num_dims(mesh)

    NGroups = ngroups(first(xss))
    if !all(g -> isequal(g, NGroups), ngroups.(xss))
        error("all `CrossSections` *must* have the same number of energy groups.")
    end
    elType = promote_type(eltype.(xss)...)

    NRegions = num_cells(mesh)
    n_polar_2 = npolar2(polar_quadrature)

    nφ = NRegions * NGroups
    nψ = n_total_tracks * 2 * n_polar_2 * NGroups # 2 is the number of directions

    quadrature = Quadrature(azimuthal_quadrature, polar_quadrature)

    face_labeling = get_face_labeling(tg.mesh.model)
    tag_to_idx = Dict{Int8,Int8}()
    for (i, xs) in enumerate(xss)
        # given a xs name, find its gridap tag
        gridap_tag = get_tag_from_name(face_labeling, xs.name)

        # map gridap tag to xs idx
        push!(tag_to_idx, gridap_tag => i)
    end

    cell_tag = get_face_tag(face_labeling, Dim)  # gridap cell/element/fsr tags
    fsr_tag = similar(cell_tag)                  # NeutronTransport cell/element/fsr tags
    for (i, gridap_tag) in enumerate(cell_tag)
        fsr_tag[i] = tag_to_idx[gridap_tag]
    end

    return MoCProblem{Dim,NRegions,NGroups,elType}(nφ, nψ, tg, quadrature, xss, fsr_tag)
end

dimension(::MoCProblem{Dim}) where{Dim} = Dim
nregions(::MoCProblem{Dim,NRegions}) where{Dim,NRegions} = NRegions
ngroups(::MoCProblem{Dim,NRegions,NGroups}) where{Dim,NRegions,NGroups} = NGroups
eltype(::MoCProblem{Dim,NRegions,NGroups,elType}) where{Dim,NRegions,NGroups,elType} = elType

function show(io::IO, prob::MoCProblem)
    println(io, "  Problem Dimension: ", dimension(prob))
    println(io, "  Number of regions: ", nregions(prob))
    print(io,   "  Number of energy groups: ", ngroups(prob))
end

@inline function getxs(prob::MoCProblem, region_idx::Integer)
    @unpack xss, fsr_tag = prob
    @inbounds xs_idx = getindex(fsr_tag, region_idx) # get material tag
    @inbounds xs = getindex(xss, xs_idx) # get material XSs
    return xs
end

# si quiero que los primeros NRegions sean para el 1er group, etc...
# macro region_index(i, g)
#     return :((g - 1) * NRegions + i)
# end

# si quiero que esten intercalados
macro region_index(i, g)
    ex = quote
        ($i - 1) * NGroups + $g
    end
    return esc(ex)
end
# puedo hacerlas como function con @inline pero necesitan `groups` o `nfsr`

macro angular_index(t, d, p, g)
    ex = quote
        ($t - 1) * 2 * n_polar_2 * NGroups + $d * n_polar_2 * NGroups + ($p - 1) * NGroups + $g
    end
    return esc(ex)
end

macro reduced_angular_index(p, g)
    ex = quote
        ($p - 1) * NGroups + $g
    end
    return esc(ex)
end


text:
using BasisFunctions
using MoreBasisFunctions
using Test

include("vandermonde_tests.jl")
include("bernstein_tests.jl")
include("lagrange_tests.jl")
include("bspline_periodic_tests.jl")


text:
using SequentialMonteCarlo
using RNGPool
using SMCExamples.MVLinearGaussian
using StaticArrays

setRNGs(0)

model, theta, ys, ko = MVLinearGaussian.defaultMVLGModel(2, 100)

smcio = SMCIO{model.particle, model.pScratch}(2^15, model.maxn,
  Threads.nthreads(), true)
smc!(model, smcio)

predMeanSMC = SequentialMonteCarlo.eta(smcio, p -> p.x, false, smcio.n)
predVarSMC = SequentialMonteCarlo.eta(smcio,
  p -> (p.x - predMeanSMC) * (p.x - predMeanSMC)', false, smcio.n)
println("Predictive mean: $(ko.predictionMeans[smcio.n])")
println("Estimate: $predMeanSMC")
println("Predictive variance: $(ko.predictionVariances[smcio.n])")
println("Estimate: $predVarSMC")

filtMeanSMC = SequentialMonteCarlo.eta(smcio, p -> p.x, true, smcio.n)
filtVarSMC = SequentialMonteCarlo.eta(smcio,
  p -> (p.x - filtMeanSMC) * (p.x - filtMeanSMC)', true, smcio.n)
println("Filtering mean: $(ko.filteringMeans[smcio.n])")
println("Estimate: $filtMeanSMC")
println("Filtering variance: $(ko.filteringVariances[smcio.n])")
println("Estimate: $filtVarSMC")


text:
using CellularInfo, JSON

graphs = Dict{Int,Vector{Tuple{Int,Int}}}()
for rule in 0:255
    @time begin
        @info "Building graph for $(rule)..."
        k = 1
        width = 11
        duration = 2^width

        traj = simulate(rule, width, duration)

        edges = []
        @time for src in 1:width
            for dst in 1:width
                if src == dst
                    continue
                end
                back = filter(f -> f != src && f != dst, 1:width)
                cond = traj[:,back]
                te = transferentropy(traj[:,src], traj[:,dst], cond', k, 1000)
                if te.p < 0.05
                    push!(edges, (src, dst))
                end
            end
        end
        graphs[rule] = sort!(edges)
    end
end

open("graphs.json", "w") do handle
    JSON.print(handle, graphs)
end

max_stars_repo_name: ELIFE-ASU/dglmoore

max_issues_repo_name: ELIFE-ASU/dglmoore

max_forks_repo_name: ELIFE-ASU/dglmoore


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

@doc raw"""
    ObjsExternalOrgMigrations(;
        current=nothing,
        date_updated=nothing,
    )

    - current::Vector{ObjsExternalOrgMigrationsCurrent}
    - date_updated::Int32
"""
mutable struct ObjsExternalOrgMigrations <: SwaggerModel
    current::Any # spec type: Union{ Nothing, Vector{ObjsExternalOrgMigrationsCurrent} } # spec name: current
    date_updated::Any # spec type: Union{ Nothing, Int32 } # spec name: date_updated

    function ObjsExternalOrgMigrations(;current=nothing, date_updated=nothing)
        o = new()
        validate_property(ObjsExternalOrgMigrations, Symbol("current"), current)
        setfield!(o, Symbol("current"), current)
        validate_property(ObjsExternalOrgMigrations, Symbol("date_updated"), date_updated)
        setfield!(o, Symbol("date_updated"), date_updated)
        o
    end
end # type ObjsExternalOrgMigrations

const _property_map_ObjsExternalOrgMigrations = Dict{Symbol,Symbol}(Symbol("current")=>Symbol("current"), Symbol("date_updated")=>Symbol("date_updated"))
const _property_types_ObjsExternalOrgMigrations = Dict{Symbol,String}(Symbol("current")=>"Vector{ObjsExternalOrgMigrationsCurrent}", Symbol("date_updated")=>"Int32")
Base.propertynames(::Type{ ObjsExternalOrgMigrations }) = collect(keys(_property_map_ObjsExternalOrgMigrations))
Swagger.property_type(::Type{ ObjsExternalOrgMigrations }, name::Symbol) = Union{Nothing,eval(Base.Meta.parse(_property_types_ObjsExternalOrgMigrations[name]))}
Swagger.field_name(::Type{ ObjsExternalOrgMigrations }, property_name::Symbol) =  _property_map_ObjsExternalOrgMigrations[property_name]

function check_required(o::ObjsExternalOrgMigrations)
    (getproperty(o, Symbol("current")) === nothing) && (return false)
    (getproperty(o, Symbol("date_updated")) === nothing) && (return false)
    true
end

function validate_property(::Type{ ObjsExternalOrgMigrations }, name::Symbol, val)
end


text:
# Automatically generated using Clang.jl

const CUTENSORNET_MAJOR = 0
const CUTENSORNET_MINOR = 0
const CUTENSORNET_PATCH = 1
const CUTENSORNET_VERSION = CUTENSORNET_MAJOR * 10000 + CUTENSORNET_MINOR * 100 + CUTENSORNET_PATCH

@cenum cutensornetStatus_t::UInt32 begin
    CUTENSORNET_STATUS_SUCCESS = 0
    CUTENSORNET_STATUS_NOT_INITIALIZED = 1
    CUTENSORNET_STATUS_ALLOC_FAILED = 3
    CUTENSORNET_STATUS_INVALID_VALUE = 7
    CUTENSORNET_STATUS_ARCH_MISMATCH = 8
    CUTENSORNET_STATUS_MAPPING_ERROR = 11
    CUTENSORNET_STATUS_EXECUTION_FAILED = 13
    CUTENSORNET_STATUS_INTERNAL_ERROR = 14
    CUTENSORNET_STATUS_NOT_SUPPORTED = 15
    CUTENSORNET_STATUS_LICENSE_ERROR = 16
    CUTENSORNET_STATUS_CUBLAS_ERROR = 17
    CUTENSORNET_STATUS_CUDA_ERROR = 18
    CUTENSORNET_STATUS_INSUFFICIENT_WORKSPACE = 19
    CUTENSORNET_STATUS_INSUFFICIENT_DRIVER = 20
    CUTENSORNET_STATUS_IO_ERROR = 21
    CUTENSORNET_STATUS_CUTENSOR_VERSION_MISMATCH = 22
end

@cenum cutensornetComputeType_t::UInt32 begin
    CUTENSORNET_COMPUTE_16F = 1
    CUTENSORNET_COMPUTE_16BF = 1024
    CUTENSORNET_COMPUTE_TF32 = 4096
    CUTENSORNET_COMPUTE_32F = 4
    CUTENSORNET_COMPUTE_64F = 16
    CUTENSORNET_COMPUTE_8U = 64
    CUTENSORNET_COMPUTE_8I = 256
    CUTENSORNET_COMPUTE_32U = 128
    CUTENSORNET_COMPUTE_32I = 512
end

@cenum cutensornetGraphAlgo_t::UInt32 begin
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_GRAPH_ALGORITHM_RB = 0
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_GRAPH_ALGORITHM_KWAY = 1
end

@cenum cutensornetMemoryModel_t::UInt32 begin
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SLICER_MEMORY_MODEL_HEURISTIC = 0
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SLICER_MEMORY_MODEL_CUTENSOR = 1
end

@cenum cutensornetContractionOptimizerConfigAttributes_t::UInt32 begin
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_GRAPH_NUM_PARTITIONS = 0
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_GRAPH_CUTOFF_SIZE = 1
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_GRAPH_ALGORITHM = 2
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_GRAPH_IMBALANCE_FACTOR = 3
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_GRAPH_NUM_ITERATIONS = 4
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_GRAPH_NUM_CUTS = 5
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_RECONFIG_NUM_ITERATIONS = 6
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_RECONFIG_NUM_LEAVES = 7
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SLICER_DISABLE_SLICING = 8
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SLICER_MEMORY_MODEL = 9
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SLICER_MEMORY_FACTOR = 10
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SLICER_MIN_SLICES = 11
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SLICER_SLICE_FACTOR = 12
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_HYPER_NUM_SAMPLES = 13
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SIMPLIFICATION_DISABLE_DR = 14
    CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SEED = 15
end

@cenum cutensornetContractionOptimizerInfoAttributes_t::UInt32 begin
    CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_NUM_SLICES = 0
    CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_NUM_SLICED_MODES = 1
    CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_SLICED_MODE = 2
    CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_SLICED_EXTENT = 3
    CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_PATH = 4
    CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_PHASE1_FLOP_COUNT = 5
    CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_FLOP_COUNT = 6
    CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_LARGEST_TENSOR = 7
    CUTENSORNET_CONTRACTION_OPTIMIZER_INFO_SLICING_OVERHEAD = 8
end

@cenum cutensornetContractionAutotunePreferenceAttributes_t::UInt32 begin
    CUTENSORNET_CONTRACTION_AUTOTUNE_MAX_ITERATIONS = 0
end

const cutensornetNetworkDescriptor_t = Ptr{Cvoid}
const cutensornetContractionPlan_t = Ptr{Cvoid}
const cutensornetHandle_t = Ptr{Cvoid}
const cutensornetNodePair = Ptr{Cvoid}
const cutensornetNodePair_t = cutensornetNodePair

struct cutensornetContractionPath
    numContractions::Int32
    data::Ptr{cutensornetNodePair_t}
end

const cutensornetContractionPath_t = cutensornetContractionPath
const cutensornetContractionOptimizerConfig_t = Ptr{Cvoid}
const cutensornetContractionOptimizerInfo_t = Ptr{Cvoid}
const cutensornetContractionAutotunePreference_t = Ptr{Cvoid}
const cutensornetLoggerCallback_t = Ptr{Cvoid}


text: Dict{String,Any}("pound" => Dict{String,Any}("value" => "We see no # comments here.","type" => "string"),"poundcomment" => Dict{String,Any}("value" => "But there are # some comments here.","type" => "string"))

max_stars_repo_name: vanillajonathan/julia

max_issues_repo_name: vanillajonathan/julia

max_forks_repo_name: vanillajonathan/julia


text:
"""
    SplitChannels <: Augmentor.Operation

Description
--------------

Splits out the color channels of the given image using the
function `ImageCore.channelview`. This will effectively create a
new array dimension for the colors in the front. In contrast to
`ImageCore.channelview` it will also result in a new dimension
for gray images.

This operation is mainly useful at the end of a pipeline in
combination with [`PermuteDims`](@ref) in order to prepare the
image for the training algorithm, which often requires the color
channels to be separate.

Usage
--------------

    SplitChannels()

See also
--------------

[`PermuteDims`](@ref), [`CombineChannels`](@ref), [`augment`](@ref)

Examples
--------------

```julia-repl
julia> using Augmentor

julia> img = testpattern()
300×400 Array{RGBA{N0f8},2}:
[...]

julia> augment(img, SplitChannels())
4×300×400 Array{N0f8,3}:
[...]

julia> augment(img, SplitChannels() |> PermuteDims(3,2,1))
400×300×4 Array{N0f8,3}:
[...]
```
"""
struct SplitChannels <: Operation end

@inline supports_eager(::Type{SplitChannels}) = false
@inline supports_lazy(::Type{SplitChannels}) = true

function applylazy(op::SplitChannels, img::AbstractArray{<:Colorant}, param)
    plain_indices(channelview(img))
end

function applylazy(op::SplitChannels, img::AbstractArray{<:AbstractGray}, param)
    ns = (1, map(length, indices(img))...)
    reshape(channelview(img), ns)
end

function showconstruction(io::IO, op::SplitChannels)
    print(io, typeof(op).name.name, "()")
end

function Base.show(io::IO, op::SplitChannels)
    if get(io, :compact, false)
        print(io, "Split colorant into its color channels")
    else
        print(io, "Augmentor.")
        showconstruction(io, op)
    end
end

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

"""
    CombineChannels <: Augmentor.Operation

Description
--------------

Combines the first dimension of a given array into a colorant of
type `colortype` using the function `ImageCore.colorview`. The
main difference is that a separate color channel is also expected
for Gray images.

The shape of the input image has to be appropriate for the given
`colortype`, which also means that the separated color channel
has to be the first dimension of the array. See
[`PermuteDims`](@ref) if that is not the case.

Usage
--------------

    CombineChannels(colortype)

Arguments
--------------

- **`colortype`** : The color type of the resulting image. Must
    be a subtype of `ColorTypes.Colorant` and match the color
    channel of the given image.

See also
--------------

[`SplitChannels`](@ref), [`PermuteDims`](@ref), [`augment`](@ref)

Examples
--------------

```julia-repl
julia> using Augmentor, Colors

julia> A = rand(3, 10, 10) # three color channels
3×10×10 Array{Float64,3}:
[...]

julia> augment(A, CombineChannels(RGB))
10×10 Array{RGB{Float64},2}:
[...]

julia> B = rand(1, 10, 10) # singleton color channel
1×10×10 Array{Float64,3}:
[...]

julia> augment(B, CombineChannels(Gray))
10×10 Array{Gray{Float64},2}:
[...]
```
"""
struct CombineChannels{T<:Colorant} <: Operation
    colortype::Type{T}
end

@inline supports_eager(::Type{<:CombineChannels}) = false
@inline supports_lazy(::Type{<:CombineChannels}) = true

function applylazy(op::CombineChannels, img::AbstractArray{<:Number}, param)
    colorview(op.colortype, plain_indices(img))
end

function applylazy(op::CombineChannels{<:AbstractGray}, img::AbstractArray{<:Number}, param)
    length(indices(img,1)) == 1 || throw(ArgumentError("The given image must have a singleton colorchannel in the first dimension in order to combine the channels to a AbstractGray colorant"))
    ns = Base.tail(map(length, indices(img)))
    colorview(op.colortype, reshape(img, ns))
end

function showconstruction(io::IO, op::CombineChannels)
    print(io, typeof(op).name.name, '(')
    _showcolor(io, op.colortype)
    print(io, ')')
end

function Base.show(io::IO, op::CombineChannels)
    if get(io, :compact, false)
        print(io, "Combine color channels into colorant ")
        _showcolor(io, op.colortype)
    else
        print(io, "Augmentor.")
        showconstruction(io, op)
    end
end


text:
module NCP_VB

include("../Misc.jl")
using .Misc
import Base.Iterators: product

using Distributions, SpecialFunctions
import Base.Iterators: product

export standard_VB

function standard_VB(X::Array{ℜ,N}, R::Ƶ; a=1.0, b=a/nansum(X), EPOCHS::Ƶ=1, ϵ=1e-16) where {ℜ<:Real, Ƶ<:Integer, N}
    I = size(X)
    α_I = (a/R) ./ I
    α_R = a/R

    S_R, S₊ = zeros(R), sum(X)
    S_I = [zeros(R,Iₙ) for Iₙ in I]

    log_λ = digamma(S₊ + a) - log(b + 1)
    log_θ_R = log.(rand(Dirichlet(R,α_R + 0.1)))
    log_θ_I = [log.(rand(Dirichlet(Iₙ,α_Iₙ + 0.1),R)) for (Iₙ,α_Iₙ) in zip(I,α_I)]

    log_C = a*log(b) - (a + S₊)*log(b + 1) + lgamma(a + S₊) - lgamma(a) - sum(lgamma, X .+ 1)

    s, log_p = zeros(R), zeros(R)
    ELBO = fill(log_C,EPOCHS)

    for eph=1:EPOCHS
        # initialization
        for n ∈ 1:N
            S_I[n] .= 0
        end
        S_R .= 0

        # expectation of S
        for i ∈ product(map(Iₙ -> 1:Iₙ, I)...)
            log_p .= log_θ_R + sum(n -> log_θ_I[n][i[n],:],1:N)
            log_p .-= logsumexp(log_p)
            s .= X[i...] .* exp.(log_p)

            for (n,iₙ) ∈ enumerate(i)
                S_I[n][:,iₙ] .+= s
            end
            S_R .+= s

            ELBO[eph] -= sum(s .* log_p)
        end

        for (n,Iₙ) ∈ enumerate(I)
            log_θ_I[n] .= digamma.(S_I[n] .+ α_I[n])' .- digamma.(S_R .+ α_R)'
            ELBO[eph] += sum(lbeta(α_I[n] .+ S_I[n];dims=2)) - R*lbeta(α_I[n],Iₙ)
        end
        log_θ_R .= digamma.(S_R .+ α_R) .- digamma(S₊ + R*α_R)
        ELBO[eph] += lbeta(α_R .+ S_R) - lbeta(α_R,R)
    end
    return ELBO, log_λ, (log_θ_R, log_θ_I...)
end

end

max_stars_repo_name: mehmetburakkurutmaz/bam

max_issues_repo_name: mehmetburakkurutmaz/bam

max_forks_repo_name: atcemgil/bam


text:
@inline allequal(a) = isempty(a) ? true : mapfoldl(==(first(a)), &, a)

@inline function same(f, x...)
    fx = f.(x)
    allequal(fx) ||
        throw(DimensionMismatch("$f not equal: $(join(fx, " vs "))"))
    return first(fx)
end