text:
# Parts of this code were taken / derived from Graphs.jl. See LICENSE for
# licensing details.

# Depth-first visit / traversal

#################################################
#
#  Depth-first visit
#
#################################################
"""
**Conventions in Breadth First Search and Depth First Search**
VertexColorMap :
- color == 0    => unseen
- color < 0     => examined but not closed
- color > 0     => examined and closed

EdgeColorMap :
- color == 0    => unseen
- color == 1     => examined
"""

type DepthFirst <: SimpleGraphVisitAlgorithm
end

function depth_first_visit_impl!(
    graph::SimpleGraph,      # the graph
    stack,                          # an (initialized) stack of vertex
    vertexcolormap::AbstractVertexMap,    # an (initialized) color-map to indicate status of vertices
    edgecolormap::AbstractEdgeMap,      # an (initialized) color-map to indicate status of edges
    visitor::SimpleGraphVisitor)  # the visitor

    while !isempty(stack)
        u, udsts, tstate = pop!(stack)
        found_new_vertex = false

        while !done(udsts, tstate) && !found_new_vertex
            v, tstate = next(udsts, tstate)
            u_color = get(vertexcolormap, u, 0)
            v_color = get(vertexcolormap, v, 0)
            v_edge = Edge(u,v)
            e_color = get(edgecolormap, v_edge, 0)
            examine_neighbor!(visitor, u, v, u_color, v_color, e_color) #no return here

            edgecolormap[v_edge] = 1

            if v_color == 0
                found_new_vertex = true
                vertexcolormap[v] = vertexcolormap[u] - 1 #negative numbers
                discover_vertex!(visitor, v) || return
                push!(stack, (u, udsts, tstate))

                open_vertex!(visitor, v)
                vdsts = fadj(graph, v)
                push!(stack, (v, vdsts, start(vdsts)))
            end
        end

        if !found_new_vertex
            close_vertex!(visitor, u)
            vertexcolormap[u] *= -1
        end
    end
end

function traverse_graph!(
    graph::SimpleGraph,
    alg::DepthFirst,
    s::Int,
    visitor::SimpleGraphVisitor;
    vertexcolormap = Dict{Int, Int}(),
    edgecolormap = DummyEdgeMap())

    vertexcolormap[s] = -1
    discover_vertex!(visitor, s) || return

    sdsts = fadj(graph, s)
    sstate = start(sdsts)
    stack = [(s, sdsts, sstate)]

    depth_first_visit_impl!(graph, stack, vertexcolormap, edgecolormap, visitor)
end

#################################################
#
#  Useful applications
#
#################################################

# Test whether a graph is cyclic

type DFSCyclicTestVisitor <: SimpleGraphVisitor
    found_cycle::Bool

    DFSCyclicTestVisitor() = new(false)
end

function examine_neighbor!(
    vis::DFSCyclicTestVisitor,
    u::Int,
    v::Int,
    ucolor::Int,
    vcolor::Int,
    ecolor::Int)

    if vcolor < 0 && ecolor == 0
        vis.found_cycle = true
    end
end

discover_vertex!(vis::DFSCyclicTestVisitor, v) = !vis.found_cycle

"""
    is_cyclic(g)

Tests whether a graph contains a cycle through depth-first search. It
returns `true` when it finds a cycle, otherwise `false`.
"""
function is_cyclic(g::SimpleGraph)
    cmap = zeros(Int, nv(g))
    visitor = DFSCyclicTestVisitor()

    for s in vertices(g)
        if cmap[s] == 0
            traverse_graph!(g, DepthFirst(), s, visitor, vertexcolormap=cmap)
        end
        visitor.found_cycle && return true
    end
    return false
end

# Topological sort using DFS

type TopologicalSortVisitor <: SimpleGraphVisitor
    vertices::Vector{Int}

    function TopologicalSortVisitor(n::Int)
        vs = Array(Int, 0)
        sizehint!(vs, n)
        new(vs)
    end
end

function examine_neighbor!(visitor::TopologicalSortVisitor, u::Int, v::Int, ucolor::Int, vcolor::Int, ecolor::Int)
    (vcolor < 0 && ecolor == 0) && error("The input graph contains at least one loop.")
end

function close_vertex!(visitor::TopologicalSortVisitor, v::Int)
    push!(visitor.vertices, v)
end

function topological_sort_by_dfs(graph::SimpleGraph)
    nvg = nv(graph)
    cmap = zeros(Int, nvg)
    visitor = TopologicalSortVisitor(nvg)

    for s in vertices(graph)
        if cmap[s] == 0
            traverse_graph!(graph, DepthFirst(), s, visitor, vertexcolormap=cmap)
        end
    end

    reverse(visitor.vertices)
end

type TreeDFSVisitor <:SimpleGraphVisitor
    tree::DiGraph
    predecessor::Vector{Int}
end

TreeDFSVisitor(n::Int) = TreeDFSVisitor(DiGraph(n), zeros(Int,n))

function examine_neighbor!(visitor::TreeDFSVisitor, u::Int, v::Int, ucolor::Int, vcolor::Int, ecolor::Int)
    if (vcolor == 0)
        visitor.predecessor[v] = u
    end
    return true
end

"""
    dfs_tree(g, s::Int)

Provides a depth-first traversal of the graph `g` starting with source vertex `s`,
and returns a directed acyclic graph of vertices in the order they were discovered.
"""
function dfs_tree(g::SimpleGraph, s::Int)
    nvg = nv(g)
    visitor = TreeDFSVisitor(nvg)
    traverse_graph!(g, DepthFirst(), s, visitor)
    # visitor = traverse_dfs(g, s, TreeDFSVisitor(nvg))
    h = DiGraph(nvg)
    for (v, u) in enumerate(visitor.predecessor)
        if u != 0
            add_edge!(h, u, v)
        end
    end
    return h
end


text:
using Statistics
using PyPlot
using DataFrames
using Distributed

# Experimental Prediction 1 (section 2.4)
#
# Sgm stands for Bayes Factor Surprise
# Ssh stands for Shannon Surprise
function runExperimentsSingleLearner_Second(;
                        nsteps = [1000, 2000],
                        p_centers = collect(0.025:0.025:0.4),
                        Δp = [0.0125],
                        seeds = [3, 4],
                        sigma = 0.5, changeprobability = 0.05,
                        learnertype = GNas12,
                        Npar = 1, etaleak = 0.1, m = 0.1)

    settings = collect(Iterators.product(nsteps,
                                        seeds,
                                        Δp
                                        )
                                    )
    setting = settings[1]
    df = @distributed vcat for setting in settings
        experiment = SurpriseExperiment(sigma = sigma,
                                    changeprobability = changeprobability,
                                    nsteps = setting[1], seed = [setting[2]],
                                    learnertype = learnertype,
                                    Npar = Npar, etaleak = etaleak, m = m
                                    )
        callback = runSingleExperiment(experiment)

        mean_Sgm, se_Sgm, mean_Ssh, se_Ssh = meanSurpriseOverp(callback;
                                p_centers = p_centers, Δp = Δp)

        DataFrame(learnertype = experiment.learnertype,
                learnerparam = experiment.learnerparam,
                sigma = sigma,
                changeprobability = changeprobability,
                nsteps = setting[1],
                seedenvpolicy = experiment.specs.seedenvpolicy[1],
                seedlearner = experiment.specs.seedlearner[1],
                Δp = setting[3],
                p_centers = [p_centers],
                mean_Sgm = [mean_Sgm],
                se_Sgm = [se_Sgm],
                mean_Ssh = [mean_Ssh],
                se_Ssh = [se_Ssh],
                )
            end
    df
end
function getconditionedSurprise_second(callback;
                        p_center = 0.025,
                        Δp = 0.0125,
                        )

    Sgm_t, Ssh_t, Py_0_t, Py_n_t = getExperimentQuantities_second(callback)

    indices = ((abs.(Py_n_t .- p_center) .<= Δp)
                .&
                (abs.(Py_0_t .- p_center) .<= Δp)
                )
    mean_Sgm_s, se_Sgm_s = getmeanSurprise_second(indices, Sgm_t)
    mean_Ssh_s, se_Ssh_s = getmeanSurprise_second(indices, Ssh_t)

    mean_Sgm_s, se_Sgm_s, mean_Ssh_s, se_Ssh_s
end
function getmeanSurprise_second(indices, Surprise_t)
    Surprise = Surprise_t[indices]

    mean_Surprise = mean(Surprise)
    se_Surprise = std(Surprise) / sqrt(length(Surprise))

    mean_Surprise, se_Surprise
end
function getExperimentQuantities_second(callback)
    Sgm_t = callback.Sgm_history[3:end]
    Ssh_t = callback.Ssh_history[3:end]
    Py_0_t = callback.Py_0_history[3:end]

    Py_n_t = Py_0_t./Sgm_t

    Sgm_t, Ssh_t, Py_0_t, Py_n_t
end
function meanSurpriseOverp(callback;
                        p_centers = collect(0.025:0.025:0.4),
                        Δp = 0.0125
                        )
    mean_Sgm = zeros(length(p_centers))
    se_Sgm = zeros(length(p_centers))
    mean_Ssh = zeros(length(p_centers))
    se_Ssh = zeros(length(p_centers))

    for i in 1:length(p_centers)
        mean_Sgm[i], se_Sgm[i],
        mean_Ssh[i], se_Ssh[i] = getconditionedSurprise_second(callback,
                                                p_center = p_centers[i],
                                                Δp = Δp)
    end
    mean_Sgm, se_Sgm, mean_Ssh, se_Ssh
end

ext: jl


text:
module FactoredValueMCTS

using Random
using LinearAlgebra

using POMDPs
using MultiAgentPOMDPs
using POMDPPolicies
using POMDPLinter: @req, @subreq, @POMDP_require
using MCTS
using LightGraphs
using BeliefUpdaters

using MCTS: convert_estimator
import POMDPModelTools

using POMDPSimulators: RolloutSimulator
import POMDPs

# Patch simulate to support vector of rewards
function POMDPs.simulate(sim::RolloutSimulator, mdp::JointMDP, policy::Policy, initialstate::S) where {S}

    if sim.eps === nothing
        eps = 0.0
    else
        eps = sim.eps
    end

    if sim.max_steps === nothing
        max_steps = typemax(Int)
    else
        max_steps = sim.max_steps
    end

    s = initialstate

    # TODO: doesn't this add unnecessary action search?
    r = @gen(:r)(mdp, s, action(policy, s), sim.rng)
    if r isa AbstractVector
        r_total = zeros(n_agents(mdp))
    else
        r_total = 0.0
    end
    sim_helper!(r_total, sim, mdp, policy, s, max_steps, eps)
    return r_total
end

function sim_helper!(r_total::AbstractVector{F}, sim, mdp, policy, s, max_steps, eps) where {F}
    step = 1
    disc = 1.0
    while disc > eps && !isterminal(mdp, s) && step <= max_steps
        a = action(policy, s)

        sp, r = @gen(:sp, :r)(mdp, s, a, sim.rng)

        r_total .+= disc.*r

        s = sp

        disc *= discount(mdp)
        step += 1
    end

end

function sim_helper!(r_total::AbstractFloat, sim, mdp, policy, s, max_steps, eps)
    step = 1
    disc = 1.0
    while disc > eps && !isterminal(mdp, s) && step <= max_steps
        a = action(policy, s)

        sp, r = @gen(:sp, :r)(mdp, s, a, sim.rng)

        r_total += disc.*r

        s = sp

        disc *= discount(mdp)
        step += 1
    end

end

###
# Factored Value MCTS
#

abstract type CoordinationStatistics end

include(joinpath("fvmcts", "factoredpolicy.jl"))
include(joinpath("fvmcts", "fv_mcts_vanilla.jl"))
include(joinpath("fvmcts", "action_coordination", "varel.jl"))
include(joinpath("fvmcts", "action_coordination", "maxplus.jl"))

export
    FVMCTSSolver,
    MaxPlus,
    VarEl

end


text:
"""
Creates Ohms constraints for DC branches

```
p[f_idx] + p[t_idx] == p * g[l] * (wdc[f_bus] - wdcr[f_bus,t_bus])
```
"""
function constraint_ohms_dc_branch(pm::AbstractBFQPModel, n::Int, cnd::Int, f_bus, t_bus, f_idx, t_idx, r, p)
    l = f_idx[1];
    p_dc_fr = PowerModels.var(pm, n, cnd, :p_dcgrid, f_idx)
    p_dc_to = PowerModels.var(pm, n, cnd, :p_dcgrid, t_idx)
    ccm_dcgrid = PowerModels.var(pm, n, cnd, :ccm_dcgrid, l)
    wdc_fr = PowerModels.var(pm, n, cnd, :wdc, f_bus)
    wdc_to = PowerModels.var(pm, n, cnd, :wdc, t_bus)

    @constraint(pm.model, p_dc_fr + p_dc_to ==  r * p * ccm_dcgrid)
    @constraint(pm.model, p_dc_fr^2 <= p^2 * wdc_fr * ccm_dcgrid)
    @constraint(pm.model, wdc_to == wdc_fr - 2 * r * (p_dc_fr/p) + (r)^2 * ccm_dcgrid)
end

function constraint_ohms_dc_branch(pm::AbstractBFConicModel, n::Int, cnd::Int, f_bus, t_bus, f_idx, t_idx, r, p)
    l = f_idx[1];
    p_dc_fr = PowerModels.var(pm, n, cnd, :p_dcgrid, f_idx)
    p_dc_to = PowerModels.var(pm, n, cnd, :p_dcgrid, t_idx)
    ccm_dcgrid = PowerModels.var(pm, n, cnd, :ccm_dcgrid, l)
    wdc_fr = PowerModels.var(pm, n, cnd, :wdc, f_bus)
    wdc_to = PowerModels.var(pm, n, cnd, :wdc, t_bus)

    @constraint(pm.model, p_dc_fr + p_dc_to ==  r * p * ccm_dcgrid)
    @constraint(pm.model, [p*wdc_fr/sqrt(2), p*ccm_dcgrid/sqrt(2), p_dc_fr/sqrt(2), p_dc_fr/sqrt(2)] in JuMP.RotatedSecondOrderCone())
    @constraint(pm.model, wdc_to == wdc_fr - 2 * r * (p_dc_fr/p) + (r)^2 * ccm_dcgrid)
end
"""
Model to approximate cross products of node voltages
```
wdcr[(i,j)] <= wdc[i]*wdc[j]
```
"""
function constraint_voltage_dc(pm::AbstractBFModel, n::Int = pm.cnw, cnd::Int = pm.ccnd)
# do nothing
end

function variable_dcbranch_current(pm::AbstractBFModel; kwargs...)
    variable_dcbranch_current_sqr(pm; kwargs...)
end
function constraint_dc_branch_current(pm::AbstractBFModel, n::Int, cnd::Int, f_bus, f_idx, ccm_max, p)
# do nothing
end


text:
module NNlib

using Pkg
using Requires
using ChainRulesCore
import ChainRulesCore: rrule
using Base.Broadcast: broadcasted
using Base.Threads
using Statistics
using Statistics: mean
using LinearAlgebra
using LinearAlgebra: BlasFloat, Transpose, Adjoint, AdjOrTransAbsMat
using LinearAlgebra.BLAS: BlasInt, @blasfunc

const libblas = Base.libblas_name

const IntOrIntTuple = Union{Integer, NTuple{N,<:Integer} where N}
const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}

# Include APIs
include("dim_helpers.jl")
export ConvDims, DenseConvDims, PoolDims, DepthwiseConvDims

is_nnpack_available() = false

@init @require NNPACK_jll="a6bfbf70-4841-5cb9-aa18-3a8ad3c413ee"  begin
  if isdefined(NNPACK_jll, :libnnpack)
    include("nnpack/NNPACK.jl")
  else
    @warn "NNPACK not available for your platform: " *
          "$( Pkg.BinaryPlatforms.platform_name(Pkg.BinaryPlatforms.platform_key_abi()))" *
          "($( Pkg.BinaryPlatforms.triplet(Pkg.BinaryPlatforms.platform_key_abi())))
          You will be able to use only the default Julia NNlib backend"
  end
end

include("activations.jl")
for f in ACTIVATIONS
    @eval export $(f)
end
export sigmoid, hardsigmoid, logsigmoid, thresholdrelu # Aliases

include("softmax.jl")
export softmax, softmax!, ∇softmax, ∇softmax!, logsoftmax,
    logsoftmax!, ∇logsoftmax, ∇logsoftmax!, logsumexp

include("batched/batchedadjtrans.jl")
include("batched/batchedmul.jl")
export batched_mul, batched_mul!, ⊠,  batched_vec,
    batched_transpose, batched_adjoint

include("gemm.jl")
export grid_sample, ∇grid_sample

include("conv.jl")
export conv, conv!, ∇conv_data, ∇conv_data!, ∇conv_filter,
    ∇conv_filter!, depthwiseconv, depthwiseconv!,
    ∇depthwiseconv_data, ∇depthwiseconv_data!,
    ∇depthwiseconv_filter, ∇depthwiseconv_filter!

include("conv_bias_act.jl")
export conv_bias_act, conv_bias_act!

include("pooling.jl")
export maxpool, maxpool!, meanpool, meanpool!,
    ∇maxpool, ∇maxpool!, ∇meanpool, ∇meanpool!

include("padding.jl")
export pad_constant, pad_repeat, pad_reflect, pad_zeros

include("upsample.jl")
export upsample_nearest, ∇upsample_nearest,
    upsample_linear, ∇upsample_linear,
    upsample_bilinear, ∇upsample_bilinear,
    upsample_trilinear, ∇upsample_trilinear,
    pixel_shuffle

include("gather.jl")
include("scatter.jl")
include("utils.jl")
include("sampling.jl")
include("functions.jl")

## Include implementations
include("impl/padding_edges.jl")

# Direct implementations of convolutional and depthwise-convolutional algorithms
include("impl/conv_direct.jl")
include("impl/depthwiseconv_direct.jl")
# im2col implementations of convolutional and depthwise-convolutional algorithms
include("impl/conv_im2col.jl")
include("impl/depthwiseconv_im2col.jl")

# Direct implementations of pooling
include("impl/pooling_direct.jl")
include("deprecations.jl")

end # module NNlib


text:
#=
# Description of the algorithm and definition of used symbols

## 0. Input embedding
s is the input timeseries or Dataset.
On s one performs a d-dimensional embedding, that uses a combination of d timeseries
(with arbitrary amount of repetitions) and d delay times.

I.e. input is `s, τs, js`, with `js, τs` being tuples.

## 1. Core loop given a specific input embedding
We have a d-dimensional embedding.
Define a radius δ around point v in ℝᵈ space (d-dimensional embedding). k points are inside
the δ-ball (with respect to some metric) around v. For simplicity, the time index of
v is t0. The other points inside the δ_ball have indices i (with several i).

We want to check if we can add an additional dimension to the embedding, using the j-th
timeseries. We check with continuity statistic of Pecora et al.

let x(t+τ) ≡ s_j(t+τ) be this extra dimension we add into the embedding. Out of the
k points in the δ-ball, we count l of them that land into a range ε around x.  Notice that
"points" is a confusing term that should not be used interchange-bly. Here in truth we
refer to **indices** not points. Because of delay embedding, all points are mapped 1-to-1
to a unique time index. We count the points with the same time indices ti, if they
are around the original x point with index t0.

Now, if l ≥ δ_to_ε_amount[k] (where δ_to_ε_amount a dictionary defined below), we can
reject the null hypothesis (that the point mapping was by chance) at a significance
level α, and thus we satisfy the continuity criterion.

## 2. Finding minimum ε

Notice that in the current code implementation, δ (which is a real number)
is never given/existing in code.
Let there be a fiducial point v, and we find the k nearest neighbors (say Vs)
and then map their indices to the ε-space. The map of the fiducial point in ε-space
is a and the map of the neighbors is As.

In the ε-space we calculate the distances of As from a. We then sort these distances.

We want the minimum range ε★ within which there are at least l (l = δ_to_ε_amount[k])
neighbors of a. This is simply the l-th maximum distance of As from a.
Why? Because if ε★ was any smaller, one neighbor wouldn't be a neighbor anymore and
we would have 1 less l.

δ_to_ε_amount is a dictionary, which stores the number of points of the δ
neighborhood as keys, and the corresponding maximum numbers of points in the ε
neighborhood as values, for rejecting the Null-Hypothesis (that there is a
functional relationship) under the significance-level `α`. The input parameter
`p` controls the shape of the binomial distribution. This is by default set to
0.5, which and has been proposed by Pecora et al. . The locations of the maxima
of the resulting ⟨ε★⟩ stay the same for a wide range of `p`-values.

## 3. Averaging ε
We repeat step 1 and 2 for several different input points v, and possibly several
input `k ∈ K`, and average the result in ⟨ε★⟩. The number `K` of considered points
in the δ-neighborhood is by default 13. In a range from `8:K`, the corresponding
ε★'s get computed (see description above) and finally the minimum of all these
values is chosen. This is due to the fact that there is no preferable size
of the δ-neighborhood a priori and it also depends on the quality of the data. So
we can not wheight any of the choices of `k`, which is why we take the minimum of
all considered choices.

The larger ⟨ε★⟩, the more functionaly independent is the new d+1 entry to the rest
d entries of the embedding.

## Creating a proper embedding
The Pecora embedding is a sequential process. This means that one should start with
a 1-dimensional embedding, with delay time 0 (i.e. a single timeseries).
Then, one performs steps 1-3 for a
choice of one more embedded dimension, i.e. the j-th timeseries and a delay τ.
The optimal choice for the second dimension of the embedding is the j entry with highest
⟨ε★⟩ and τ in a local maximum of ⟨ε★⟩.

Then these two dimensions are used again as input to the algorithm, and sequentially
the third optimal entry for the embedding is chosen. Each added entry successfully
reduces ⟨ε★⟩ for the next entry.

This process continues until ε cannot be reduced further, in which scenario the
process terminates and we have found an optimal embedding that maximizes
functional independence among the dimensions of the embedding.

## Perforamance notes
`samplesize` controls the number of considered fiducial points v. It determines
the fraction of all available points in the trajectory. This is just for
performance reasons and one would get the "very best" result for setting
`samplesize = 1`, obviously, i.e. considering every single point of the trajectory.
Default is `samplesize = 0.1`.

## The undersampling statistic
Because real world data are finite, the aforementioned process (of seeing when ⟨ε★⟩
will saturate) isn't very accurate because as the dimension of v increases, and
as τ increases, we are undersampling a high-dimensional object.

To counter this, one calcualtes the undersampling statistic ⟨Γ⟩ like so:
For each timeseries calculate the (normalized) histogram, i.e. pdf function
ρ, by using e.g. `db` bins.
Then, for each timeseries x prepare the function σ(ζ), defined as
σ(ζ) = ∫ρ(x) ρ(ζ-x) dx
where ζ is positive, running from 0 to max(x)-min(x).

The following part of the algorithm demands using Chebyshev metric for nearest
neighbors (see performance notes for why)!
Let v be a fiducial point. Create a point q = (v..., a), i.e. the fiducial point
in the product of the δ- and ε- spaces. Now, let qNN be the nearest neigbhor of q.
For every entry of q, calculate ξᵢ=|qᵢ - qNNᵢ|. Then, calculate the integral
Γᵢ = ∫_0 ^(ξᵢ) σ_{J[i]}(ζ) dζ
Then Γ is the maximum of Γᵢ . The undersampling statistic ⟨Γ⟩ is the average of
Γ over all fiducial points.

If Γ>β we are undersampling and thus we don't have statistical confidence
to continue computing ε★.

## Perforamance notes
An optimized σ function is necessary. As many σ functions as timeseries existing
both in the embedding as well as in J are neccessary, irrespectively of the size
of `js`, the input embedding (because the same σ is reused if the same timeseries is
delayed). The optimized calculation of σ is done by assuming that ξ is always at a
bin edge. Thus, the higher `db` the more accurate the computation.

Now, finding the neighbor qNN is extremely tricky. If we were to do it in Euclidean
norm, then for every τ and for every timeseries x we would have to create a _new_
tree and find the nearest neigghbors for all fiducial points. This is insane.

Instead, we use Chebyshev (maximum) norm to simplify the problem. The Chebyshev
distance is the greatest distance along each dimension. Thus, the nearest neighbor
of a fiducial point v0 with index n0 will be the one whose Chebyshev distance
with v0 is the smallest. Let this distance be d0 and this index ni.

We now look at the "ε space", where the x timeseries lives.
We calculate d = |x(ni+τ) - x(n0+τ)|. If d ≤ d0, nothing changes,
and q_ni is the nearest neighbor of q_n0 (just as v_ni is the nearest neighbor to v_n0).
This happens because any other point with index nj would have a Chebyshev distance from
q_n0 that is greater than the existing d0 (since we know guaranteed that the existing
dimensions in v space have _more_ difference along each dimension, since the Chebyshev
is the maximum of those).

Now, if d is greater than d0, then we have to actually find the nearest neighbor,
since there is no guarantee that q_ni is still the nearest neighbor of q_n0.
The extra dimension has _increased_ the Chebyshev distance between them, and we
don't know whether there exists some new point q_nj whose Chebyshev distance from q_n0
is between d0 < dj ≤ d. Here is how we find the nearest neighbor:

1. We collect all indices i that satisfy d0 < |x(ni+τ) - x(n0+τ)| ≤ d, and sort them
   by their distance
2. We start iterating over these indices i, checking the Chebyshev distance of the
   points in v space with index i, with respect to the point with index n0
3. If we find distance that is less than `d`, we terminate the algorithm, and that
   i is the true nearest neighbor in q-space.
4. If not, and we have exhausted all i, then indeed the original ni
   remains the true neighbor.

=#

using Distances, Statistics, Distributions
using StatsBase: fit, Histogram
export pecora

"""
    pecora(s, τs, js; kwargs...) → ⟨ε★⟩, ⟨Γ⟩
Compute the (average) continuity statistic `⟨ε★⟩` and undersampling statistic `⟨Γ⟩`
according to Pecora et al.[^Pecoral2007] (A unified approach to attractor reconstruction),
for a given input
`s` (timeseries or `Dataset`) and input generalized embedding defined by `(τs, js)`,
according to [`genembed`](@ref). The continuity statistic represents functional independence
between the components of the existing embedding and one additional timeseries.
The returned results are *matrices* with size `T`x`J`.

## Keyword arguments
* `delays = 0:50`: Possible time delay values `delays` (in sampling time units).
  For each of the `τ`'s in `delays` the continuity-statistic `⟨ε★⟩` gets computed.
  If `undersampling = true` (see further down), also the undersampling statistic
  `⟨Γ⟩` gets returned for all considered delay values.
* `J = 1:dimension(s)`: calculate for all timeseries indices in `J`.
  If input `s` is a timeseries, this is always just 1.
* `samplesize::Real = 0.1`: determine the fraction of all phase space points (=`length(s)`)
  to be considered (fiducial points v) to average ε★ to produce `⟨ε★⟩, ⟨Γ⟩`
* `K::Int = 13`: the amount of nearest neighbors in the δ-ball (read algorithm description).
  Must be at least 8 (in order to gurantee a valid statistic). `⟨ε★⟩` is computed
  taking the minimum result over all `k ∈ K`.
* `metric = Chebyshev()`: metrix with which to find nearest neigbhors in the input
  embedding (ℝᵈ space, `d = length(τs)`).
* `w = 1`: Theiler window (neighbors in time with index `w` close to the point, that
  are excluded from being true neighbors). `w=0` means to exclude only the
  point itself, and no temporal neighbors.
* `undersampling = false` : whether to calculate the undersampling statistic or not
  (if not, zeros are returned for `⟨Γ⟩`). Calculating `⟨Γ⟩` is thousands of times
  slower than `⟨ε★⟩`.
* `db::Int = 100`: Amount of bins used into calculating the histograms of
  each timeseries (for the undersampling statistic).
* `α::Real = 0.05`: The significance level for obtaining the continuity statistic
* `p::Real = 0.5`: The p-parameter for the binomial distribution used for the
  computation of the continuity statistic.

## Description
Notice that the full algorithm is too large to discuss here, and is
written in detail (several pages!) in the source code of `pecora`.

[^Pecora2007]: Pecora, L. M., Moniz, L., Nichols, J., & Carroll, T. L. (2007). [A unified approach to attractor reconstruction. Chaos 17(1)](https://doi.org/10.1063/1.2430294).
"""
function pecora(
        s, τs::NTuple{D, Int}, js::NTuple{D, Int} = Tuple(ones(Int, D));
        delays = 0:50 , J=maxdimspan(s), samplesize::Real = 0.5, K::Int = 13, w::Int = 1,
        db = 250, undersampling = false, metric = Chebyshev(), α::T = 0.05,
        p::T = 0.5) where {D, T<:Real}

    @assert K ≥ 8 "You must provide a δ-neighborhood size consisting of at least 8 neighbors."
    @assert all(x -> x ≥ 0, js) "j's for generalized embedding must be positive integers"
    @assert 0 < samplesize ≤ 1 "`samplesize` must be ∈ (0,1]"

    if undersampling
        error("Undersampling statistic is not yet accurate for production use.")
    end
    undersampling && metric ≠ Chebyshev() && error("Chebyshev metric required for undersampling")
    vspace = genembed(s, τs, js)
    vtree = KDTree(vspace.data[1:end-maximum(delays)], metric)
    all_ε★ = zeros(length(delays), length(J))
    all_Γ = copy(all_ε★)
    allts = columns(s)
    # indices of random fiducial points (with valid time range w.r.t. T)
    L = length(vspace)
    if samplesize==1
        ns = vec(1:length(vec(max(1, (-minimum(delays) + 1)):min(L, L - maximum(delays)))))
    else
        ns = sample(vec(max(1, (-minimum(delays) + 1)):min(L, L - maximum(delays))),
        floor(Int,samplesize*length(vec(max(1, (-minimum(delays) + 1)):min(L, L - maximum(delays))))),
        replace = false)
    end

    vs = vspace[ns]
    allNNidxs, allNNdist = all_neighbors(vtree, vs, ns, K, w)
    # prepare things for undersampling statistic
    if undersampling
        # TODO: Allow `db` to be vector
        ρs = histograms(s, unique([js..., J...]), db)
        uidxs = [n[1] for n in allNNidxs]
        udist = [d[1] for d in allNNdist]
    end
    # Loop over potential timeseries to use in new embedding
    for i in 1:length(J)
        x = allts[J[i]]
        x = (x .- mean(x))./std(x) # so that different timeseries can be compared
        all_ε★[:, i] .= continuity_per_timeseries(x, ns, allNNidxs, delays, K, α, p)
        if undersampling
            println("Started undersampling for $(J[i]) timeseries")
            all_Γ[:, i] .= undersampling_per_timeseries(
                x, vspace, ns, uidxs, udist, [js..., J[i]], delays, ρs
            )
        end
    end
    return all_ε★, all_Γ
end

maxdimspan(s) = 1:size(s)[2]
maxdimspan(s::AbstractVector) = 1
columns(s::AbstractVector) = (s, )

function continuity_per_timeseries(x::AbstractVector, ns, allNNidxs, delays, K, α, p)
    avrg_ε★ = zeros(size(delays))
    Ks = [k for k in 8:K]
    δ_to_ε_amount = get_binomial_table(p, α; trial_range = length(Ks))
    for (ι, τ) in enumerate(delays) # Loop over the different delays
        c = 0
        for (i, n) in enumerate(ns) # Loop over fiducial points
            NNidxs = allNNidxs[i] # indices of k nearest neighbors to v
            # Check if any of the indices of the neighbors falls out of temporal range
            any(j -> (j+τ < 1), NNidxs) && continue
            # If not, calculate minimum ε
            avrg_ε★[ι] += ε★(x, n, τ, NNidxs, δ_to_ε_amount, Ks)
            c += 1
        end
        c == 0 && error("Encountered astronomically small chance of all neighbors having "*
                        "invalid temporal range... Run the function again or decrease `w`.")
        avrg_ε★[ι] /= c
    end
    return avrg_ε★
end

function ε★(x, n, τ, NNidxs, δ_to_ε_amount::Dict, Ks::AbstractVector)
    a = x[n+τ] # fiducial point in ε-space
    @inbounds dis = [abs(a - x[i+τ]) for i in NNidxs]
    ε = zeros(length(Ks))
    for (i, k) in enumerate(Ks)
        sortedds = sort!(dis[1:k]; alg = QuickSort)
        l = δ_to_ε_amount[k]
        ε[i] = sortedds[l]
    end
    return minimum(ε)
end

##########################################################################################
# Undersampling statistic code
##########################################################################################
function undersampling_per_timeseries(x, vspace, ns, uidxs, udist, JS, delays, ρs)
    avrg_Γ = zeros(size(delays))
    for (ι, τ) in enumerate(delays) # loop over delay times
        println("τ = $τ")
        c = 0
        τr = max(1, (-τ + 1)):min(length(vspace), length(vspace) - τ) # valid time range
        for (i, n0) in enumerate(ns) # Loop over fiducial points
            ni, d0 = uidxs[i], udist[i] # distances in vspace
            # Check if the neighbor + τ falls out of temporal range
            (ni+τ > length(x)) | (ni+τ < 1) && continue
            nℓ = true_neighbor_index(ni, n0, vspace, d0, x, τ, τr)
            q = (vspace[n0]..., x[n0+τ])
            qNN = (vspace[nℓ]..., x[nℓ+τ])
            avrg_Γ[ι] += Γ(q, qNN, JS, ρs)
            c += 1
        end
        avrg_Γ[ι] /= c
    end
    return avrg_Γ
end

function Γ(q, qNN, JS, ρs)
    Γmax = 0.0
    for i in 1:length(q)
        ξᵢ = abs(q[i] - qNN[i])
        ρᵢ = ρs[JS[i]]
        Γᵢ = integral_σ(ρᵢ, ξᵢ)
        Γmax = Γᵢ > Γmax ? Γᵢ : Γmax
    end
    return Γmax
end

# this function assumes Chebyshev norm
function true_neighbor_index(ni, n0, vspace, d0, x, τ, τr)
    d = abs(x[ni+τ] - x[n0+τ])
    if d ≤ d0
        return ni
    else # we now have to find a different point
        # TODO: This can be optimized with a dedicated function
        # that doesn't allocate 4 extra arrays...
        tocheck = τr[findall(i -> d0 < abs(x[i+τ] - x[n0+τ]) ≤ d, τr)]
        dist = map(i -> abs(x[i+τ] - x[n0+τ]), tocheck)
        _si = sortperm(dist)
        tocheck = tocheck[_si]
        dist = dist[_si]
        true_ni = ni
        for i in tocheck
            dj = evaluate(Chebyshev(), vspace[n0], vspace[i])
            if dj < d
                true_ni = i
                break
            end
        end
        return true_ni
    end
end

function histograms(s, J, db)
    sort!(J); hs = []
    for j in J
        x = s[:, j]
        h = normalize(fit(Histogram, x; nbins = db); mode = :pdf)
        push!(hs, h)
    end
    H = typeof(hs[1])
    return Dict{Int, H}((J[i]=>hs[i]) for i in 1:length(J))
end

"calculate σ(ζ) = ∫ρ(x) ρ(ζ-x) dx , approximating ζ to be in bin start (whole bin counts)."
function difference_distribution(ρ::Histogram, ζ)
    e = ρ.edges[1]
    dx = Float64(e.step)
    ∫ = 0.0
    vmin = max(1, searchsortedfirst(e, ζ-maximum(e)))
    vmax = min(length(e)-1, searchsortedfirst(e, ζ-minimum(e)))
    valid = vmin:vmax # all indices where (ζ-emax, ζ-emin) ∈ e
    # @show e, valid, ζ
    for i in 1:length(valid)
        # x = e[valid[i]]
        # ζmx = e[valid[end]-i+1] # ζ - x
        # @show (valid[i], valid[end]-i, x, ζmx, ζ-x)
        # we know the bin of x and ζ-x
        ∫ += (ρ.weights[valid[i]] * ρ.weights[valid[end]-i+1])*dx
    end
    return ∫ # I don't need trapezoid rule, histogram is PDF normalized
end

"approximately calculate ∫_0 ^(ξ) σ(ζ) dζ."
function integral_σ(ρ, ξ)
    e = ρ.edges[1]
    dζ = Float64(e.step)
    ∫ = 0.0
    for ζ in 0.0:dζ:ξ
        σ = difference_distribution(ρ, ζ)
        ∫ += σ*dζ
    end
    return ∫
end

"""
    get_binomial_table(p, α; trial_range::Int=8) -> `δ_to_ε_amount`, Dict(δ_points => ϵ_points)
compute the numbers of points from the δ-neighborhood, which need to fall outside
the ϵ-neighborhood, in order to reject the Null Hypothesis at a significance
level `α`. One parameter of the binomial distribution is `p`, the other one would
be the number of trials, i.e. the considered number of points of the δ-neighborhood.
`trial_range` determines the number of considered δ-neighborhood-points, always
starting from 8. For instance, if `trial_range=8` (Default), then δ-neighborhood
sizes from 8 up to 15 are considered.
Return `δ_to_ε_amount`, a dictionary with `δ_points` as keys and the corresponding number of
points in order to reject the Null, `ϵ_points`, constitute the values.
"""
function get_binomial_table(p::T, α::T; trial_range::Int=8) where {T<:Real}
    @assert trial_range ≥ 1 "`trial_range` must be an integer ≥ 1"
    δ_to_ε_amount = Dict{Int, Int}()
    @inbounds for key = 8:(7+trial_range)
        δ_to_ε_amount[key] = quantile(Distributions.Binomial(key,p), 1-α)
    end
    return δ_to_ε_amount
end


text:
mutable struct MeanNormResDense <: Block
    blocks::Vector
    function MeanNormResDense(ichannels,  # input #channels, equals to output #channels
                              mchannels;  # intermediary #channels (or dimension)
                              type::Type=Array{Float32})

        linear1 = Affine(ichannels, mchannels, type=type)
        normer1 = MeanNorm(ndims=2, keptdims=1, keptsize=mchannels, type=type)
        linear2 = Affine(mchannels, ichannels, type=type)
        normer2 = MeanNorm(ndims=2, keptdims=1, keptsize=ichannels, type=type)
        new([linear1, normer1, linear2, normer2])
    end
    function MeanNormResDense()
        new(Vector(undef,4))
    end
end

Base.length(m::MeanNormResDense)     = 4
Base.lastindex(m::MeanNormResDense)  = 4
Base.firstindex(m::MeanNormResDense) = 1
Base.getindex(m::MeanNormResDense, k...)     =  m.blocks[k...]
Base.setindex!(m::MeanNormResDense, v, k...) = (m.blocks[k...] = v)
Base.iterate(m::MeanNormResDense, i=1) = i>4 ? nothing : (m[i], i+1)

@basic(MeanNormResDense, blocks)

function clone(this::MeanNormResDense; type::Type=Array{Float32})
    cloned = MeanNormResDense()
    for i = 1:4
        cloned[i] = clone(this[i], type=type)
    end
    return cloned
end

function forward(m::MeanNormResDense, x0::Variable{T}) where T
    x =      forward(m[1], x0)
    x = relu(forward(m[2], x))
    x =      forward(m[3], x)
    x = relu(forward(m[4], x) + x0)
    return x
end

function predict(m::MeanNormResDense, x0::AbstractArray)
    x =      predict(m[1], x0)
    x = relu(predict(m[2], x))
    x =      predict(m[3], x)
    x = relu(predict(m[4], x) + x0)
    return x
end

max_stars_repo_name: sonosole/Delta

max_issues_repo_name: sonosole/Delta

max_forks_repo_name: sonosole/Delta


text: dice_metric(ŷ, y) = 2 * sum((ŷ .& y)) / (sum(ŷ) + sum(y))


text:
module OceanDistributions

using CSV, DataFrames
export readoceandistribution

"""
    readoceandistribution(fil::String)

Read a distribution from csv file

```
Tcensus=readoceandistribution("examples/Tcensus.txt")
```
"""
function readoceandistribution(file::String)
    return CSV.read(file,DataFrame)
end

end # module


text:
using chitchat
using Base.Test

# write your own tests here
@test 1 == 1
@test reduce(add_ab, [1 2 3 4 5]) == 15


text:
# __precompile__(true)
module Microwave
# Microwave.jl
# Author: Eunjong Kim
import Base: +, -, *, ^, /, convert, show, getindex, setindex!
# export setMicrowaveFloat

abstract type AbstractParams end
abstract type AbstractData end

export AbstractParams, AbstractData
# const MFloat = BigFloat
#using Touchstone
include(joinpath("Circuit", "Circuit.jl"))

import .Circuit: CircuitParams, Impedance, Admittance, CircuitData, ∥,
    resistor, inductor, capacitor
export Circuit
export CircuitParams, Impedance, Admittance, CircuitData, ∥,
    resistor, inductor, capacitor

include(joinpath("Network", "Network.jl"))
import .Network: NetworkParams, Sparams, Yparams, Zparams, TwoPortParams, ABCDparams,
      Port, NetworkData, impedances, swap_ports!, permute_ports!, connect_ports,
      innerconnect_ports, cascade, reflection_coefficient, transmission_coefficient,
      impedance_step, series_network, parallel_network, terminated_network,
      π_network, t_network, check_frequency_identical
export Network
export NetworkParams, Sparams, Yparams, Zparams, TwoPortParams, ABCDparams,
      Port, NetworkData, impedances, swap_ports!, permute_ports!, connect_ports,
      innerconnect_ports, cascade, reflection_coefficient, transmission_coefficient,
      impedance_step, series_network, parallel_network, terminated_network,
      π_network, t_network, check_frequency_identical

include(joinpath("Touchstone", "Touchstone.jl"))
import .Touchstone: TouchstoneData, read_touchstone
export TouchstoneData, read_touchstone

# interpolate data
include("interpolate.jl")

end # module


text:
"
    Represents a periodic (in 2 directions) statmech mpo
"
struct PeriodicMPO{O<:MPOTensor} <: Operator
    opp::PeriodicArray{O,2}
end

PeriodicMPO(t::AbstractTensorMap) = PeriodicMPO(fill(t,1,1));
PeriodicMPO(t::AbstractArray{T,2}) where T<:MPOTensor = PeriodicMPO(PeriodicArray(t));

Base.getindex(o::PeriodicMPO,i,j) = o.opp[i,j]
Base.size(o::PeriodicMPO,i) = size(o.opp,i);
Base.size(o::PeriodicMPO) = size(o.opp);


text:
module UnscentedKalmanFilter
export UKFState, predict!, update!

using LinearAlgebra: I, dot, diagm

include("Helpers.jl")
using ..SigmaPoints
using .Helpers

"""
# State of the Unscented Kalman Filter

- `dim_x::Int64`: State dimension
- `dim_z::Int64`: Measurement dimension
- `x::AbstractArray{Float64, 1}`: State of the filter
- `P::AbstractArray{Float64, 2}`: Covariance matrix of the filter.
- `K::AbstractArray{Float64, 2}`: Kalman gain of the filter.
- `S::AbstractArray{Float64, 2}`: System uncertainty.
- `S_inv::AbstractArray{Float64, 2}`: Inverse of system uncertainty.
- `σ_parameters::MerweScaled`: Parameters for Van der Merwe sigma points.
- `σ_fx::AbstractArray{Float64, 2}`: Sigma points after `fx`.
- `σ_hx::AbstractArray{Float64, 2}`: Sigma points after `hx`.
- `add_x::Function`: Function to calculate addition between two states.
- `residual_x::Function`: Function to calculate residual between two states.
- `residual_z::Function`: Function to calculate residual between state and measurement.
- `mean_x::Function`: Function to calculate mean between Σ-points and weights
    after passing them through `fx`.
- `mean_z::Function`: Function to calculate mean between Σ-points and weights
    after passing them through `hx`.
"""
struct UKFState
    dim_x::Int64
    dim_z::Int64

    x::AbstractArray{Float64, 1}
    P::AbstractArray{Float64, 2}

    K::AbstractArray{Float64, 2}
    y::AbstractArray{Float64, 1}

    S::AbstractArray{Float64, 2}
    S_inv::AbstractArray{Float64, 2}

    σ_parameters::MerweScaled
    σ_fx::AbstractArray{Float64, 2}
    σ_hx::AbstractArray{Float64, 2}

    add_x::Function
    residual_x::Function
    residual_z::Function
    mean_x::Function
    mean_z::Function
end

"""
```julia
UKFState(
    ;dim_x::Int64, dim_z::Int64, σ_parameters::MerweScaled,
    add_x::Function = default_add_x,
    residual_x::Function = default_residual,
    residual_z::Function = default_residual,
    mean_x::Function = default_mean_f,
    mean_z::Function = default_mean_f,
)
```

# Arguments
- `dim_x::Int64`: State dimension
- `dim_z::Int64`: Measurement dimension
- `σ_parameters::MerweScaled`: Parameters for Van der Merwe sigma points.
- `add_x::Function`:
    Function to calculate addition between two states.
    Should accept two positional arguments `y` and `x`
    both of type `AbstractArray{Float64, 1}` and return `AbstractArray{Float64, 1}` as well.
- `residual_x::Function`: Function to calculate residual between two states.
    Should accept two positional arguments `y` and `x`,
    but should have methods both for `AbstractArray{Float64, 1}`, `AbstractArray{Float64, 1}`
    and for `AbstractArray{Float64, 2}`, `AbstractArray{Float64, 1}` types of arguments.
    In second case, residual will be calculated between `Σ`-points
    and state `x` and return matrix of the same shape as `Σ`.
- `residual_z::Function`: Function to calculate residual between state and measurement.
    Same as for `residual_x`, you should define methods for both kinds
    of arguments.
- `mean_x::Function`: Function to calculate mean between `Σ`-points
    and weights after passing them through `fx`.
    Arguments are positional.
- `mean_z::Function`: Function to calculate mean between `Σ`-points
    and weights after passing them through `hx`.
    Arguments are positional.

!!! note
    Do not forget to implement methods for both kinds of arguments for
    `residual_x` and/or `residual_z` as described above.
"""
function UKFState(
    ;dim_x::Int64, dim_z::Int64, σ_parameters::MerweScaled,
    add_x::Function = default_add_x,
    residual_x::Function = default_residual,
    residual_z::Function = default_residual,
    mean_x::Function = default_mean_f,
    mean_z::Function = default_mean_f,
)
    x = zeros(Float64, dim_x)
    P = Matrix{Float64}(I, dim_x, dim_x)
    K = zeros(Float64, dim_x, dim_z)
    y = zeros(Float64, dim_z)
    σ_fx = zeros(Float64, num_σ_points(σ_parameters), dim_x)
    σ_hx = zeros(Float64, num_σ_points(σ_parameters), dim_z)
    S = zeros(Float64, dim_z, dim_z)
    S_inv = zeros(Float64, dim_z, dim_z)

    UKFState(
        dim_x, dim_z,
        x, P, K, y, S, S_inv,
        σ_parameters, σ_fx, σ_hx,
        add_x, residual_x, residual_z, mean_x, mean_z
    )
end

function Base.show(io::IO, ukf::UKFState)
    print(
        io, "UKF State:\n",
        "dim_x: ", repr(MIME("text/plain"), ukf.dim_x, context=io), "\n",
        "dim_z: ", repr(MIME("text/plain"), ukf.dim_z, context=io), "\n",
        "x: ", repr(MIME("text/plain"), ukf.x, context=io), "\n",
        "P: ", repr(MIME("text/plain"), ukf.P, context=io), "\n",
        "K: ", repr(MIME("text/plain"), ukf.K, context=io), "\n",
        "y: ", repr(MIME("text/plain"), ukf.y, context=io), "\n",
        "S: ", repr(MIME("text/plain"), ukf.S, context=io), "\n",
        "S_inv: ", repr(MIME("text/plain"), ukf.S_inv, context=io), "\n",
        "σ_parameters: ", repr(MIME("text/plain"), ukf.σ_parameters, context=io), "\n",
        "σ_fx: ", repr(MIME("text/plain"), ukf.σ_fx, context=io), "\n",
        "σ_hx: ", repr(MIME("text/plain"), ukf.σ_hx, context=io), "\n",
        "add_x: ", repr(MIME("text/plain"), ukf.add_x, context=io), "\n",
        "residual_x: ", repr(MIME("text/plain"), ukf.residual_x, context=io), "\n",
        "residual_z: ", repr(MIME("text/plain"), ukf.residual_z, context=io), "\n",
        "mean_x: ", repr(MIME("text/plain"), ukf.mean_x, context=io), "\n",
        "mean_z: ", repr(MIME("text/plain"), ukf.mean_z, context=io), "\n",
    )
end

"""
Pass σ-points for the current UKF state `(x, P)` through `fx` function.
"""
function calculate_process_σ!(
    ;ukf_state::UKFState, δt::Float64, fx::Function, fx_args...
)
    σ_points = calculate_σ_points(
        σ_parameters=ukf_state.σ_parameters, x=ukf_state.x, P=ukf_state.P
    )
    @inbounds for i = 1:size(σ_points, 1)
        ukf_state.σ_fx[i, :] = fx(x=σ_points[i, :], δt=δt; fx_args...)
    end
end

"""
Calculate cross-variance between state and measurement.
"""
function cross_variance(
    ;ukf::UKFState, x::AbstractArray{Float64, 1}, z::AbstractArray{Float64, 1}
)::AbstractArray{Float64, 2}
    Δx = ukf.residual_x(ukf.σ_fx, x)
    Δz = ukf.residual_z(ukf.σ_hx, z)
    Δx' * diagm(ukf.σ_parameters.Σ_w) * Δz
end

"""
```julia
unscented_transform(
    ;ukf=UKFState, Σ::AbstractArray{Float64, 2}, residual_f::Function,
    noise_cov::Union{Nothing, AbstractArray{Float64, 2}} = nothing,
)::Tuple{AbstractArray{Float64, 1}, AbstractArray{Float64, 2}}
```

Unscented transformation of Σ-points and their weights.
This is used to calculate new mean `x` and covariance `P`
for the state of Unscented Kalman Filter.

# Arguments
- `ukf::UKFState`: State of the UKF filter which contains weights for `Σ`.
- `Σ::AbstractArray{Float64, 2}`: `[2x+1, n]` matrix of Σ-points
- `residual_f::Function`: Function to calculate residual between `y` and `x`.
- `mean_f::Function`: Function to calculate mean between Σ points and weights.
- `noise_cov::Union{Nothing, AbstractArray{Float64, 2}} = nothing`:
    Optional noise, that will be added to final covariance matrix.
    Its size should match that of a covariance matrix.
"""
function unscented_transform(
    ;ukf=UKFState, Σ::AbstractArray{Float64, 2},
    residual_f::Function, mean_f::Function,
    noise_cov::Union{Nothing, AbstractArray{Float64, 2}} = nothing,
)::Tuple{AbstractArray{Float64, 1}, AbstractArray{Float64, 2}}
    x = mean_f(Σ, ukf.σ_parameters.m_w)
    Y = residual_f(Σ, x)
    P = Y' * diagm(ukf.σ_parameters.Σ_w) * Y
    if isa(noise_cov, AbstractArray{Float64, 2})
        P .+= noise_cov
    end
    x, P
end

"""
```julia
predict!(
    ;ukf::UKFState, δt::Float64, Q::Union{Nothing, AbstractArray{Float64, 2}},
    fx::Function, fx_args...
)
```

Perform prediction of the Unscented Kalman Filter (UKF).
This updates `x` and `P` of the `ukf`
to contain predicted state and covariance.

# Arguments
- `ukf::UKFState`: State of the UKF on which to perform prediction.
- `δt::Float64`: Time delta.
- `Q::Union{Nothing, AbstractArray{Float64, 2}}`: Process noise.
- `fx::Function(;x::AbstractArray{Float64, 1}, δt::Float64, ...)`:
    State transition function.
    Should accept keyword-only arguments, with mandatory `x` and `δt`.
    All other arguments may serve as control input or provide additional info.
- `fx_args...`: Arguments to pass to `fx` function.
    Must contain at least `x` and `δt` arguments.

!!! note
    Just to emphasize, `fx` should accept keyword-only arguments,
    with mandatory `x` and `δt`. All other arguments may serve
    as control input or provide additional info.
"""
function predict!(
    ;ukf::UKFState, δt::Float64, Q::Union{Nothing, AbstractArray{Float64, 2}},
    fx::Function, fx_args...
)
    calculate_process_σ!(ukf_state=ukf, δt=δt, fx=fx; fx_args...)
    ukf.x[:], ukf.P[:] = unscented_transform(
        ukf=ukf, Σ=ukf.σ_fx, noise_cov=Q,
        residual_f=ukf.residual_x, mean_f=ukf.mean_x
    )
    ukf.σ_fx .= calculate_σ_points(
        σ_parameters=ukf.σ_parameters, x=ukf.x, P=ukf.P
    )
end

"""
```julia
update!(
    ;ukf::UKFState, z::AbstractArray{Float64, 1}, R::Union{Nothing, AbstractArray{Float64, 2}},
    hx::Function, hx_args...
)
```

Perform update step of the Unscented Kalman Filter (UKF)
given measurements.
This updates `x` and `P` of the `ukf_state`
to contain updated state and covariance.

# Arguments
- `ukf::UKFState`: State of the UKF on which to perform update step.
- `z::AbstractArray{Float64, 1}`: Measurements.
- `R::Union{Nothing, AbstractArray{Float64, 2}}`: Measurement noise.
- `hx::Function`:
    Measurement function.
    Should accept keyword-only arguments with mandatory `x` argument for state.
    Given state `x` of UKF it must transform it to the respective measurement.
- `hx_args...`: Arguments that will be passed to `hx` function.
    Should contain at least `x` argument.

!!! note
    Just to emphasize, `hx` should accept keyword-only arguments
    with mandatory `x` argument for state.
    Given state `x` of UKF it must transform it to the respective measurement.
"""
function update!(
    ;ukf::UKFState, z::AbstractArray{Float64, 1},
    R::Union{Nothing, AbstractArray{Float64, 2}},
    hx::Function, hx_args...
)
    @inbounds for i = 1:size(ukf.σ_fx, 1)
        ukf.σ_hx[i, :] = hx(x=ukf.σ_fx[i, :]; hx_args...)
    end

    zp, ukf.S[:] = unscented_transform(
        ukf=ukf, Σ=ukf.σ_hx, noise_cov=R,
        residual_f=ukf.residual_z, mean_f=ukf.mean_z
    )
    ukf.S_inv .= inv(ukf.S)

    V = cross_variance(ukf=ukf, x=ukf.x, z=zp)
    ukf.K .= V * ukf.S_inv
    ukf.y .= ukf.residual_z(z, zp)

    ukf.x .= ukf.add_x(ukf.x, ukf.K * ukf.y)
    ukf.P .-= ukf.K * ukf.S * ukf.K'
end

end

max_stars_repo_name: pxl-th/KalmanFilters

max_issues_repo_name: pxl-th/KalmanFilters

max_forks_repo_name: pxl-th/KalmanFilters


text:
using Printf

using StructuralIdentifiability
using StructuralIdentifiability: ODE, print_for_DAISY, print_for_COMBOS, print_for_maple

include("benchmarks.jl")

formats = [
    Dict(
        :name => "SIAN",
        :function => ode -> print_for_maple(ode, :SIAN),
        :extension => ".mpl"
    ),
    Dict(
        :name => "RosenfeldGroebner",
        :function => ode -> print_for_maple(ode, :DifferentialAlgebra),
        :extension => ".mpl"
    ),
    Dict(
        :name => "DifferentialThomas",
        :function => ode -> print_for_maple(ode, :DifferentialThomas),
        :extension => ".mpl"
    ),
    Dict(
        :name => "DAISY",
        :function => print_for_DAISY,
        :extension => ".txt"
    ),
    Dict(
        :name => "COMBOS",
        :function => print_for_COMBOS,
        :extension => ".txt"
    )
]

for frmt in formats
    mkpath(frmt[:name])
    cd(frmt[:name])
    for bnchmrk in benchmarks
        as_text = frmt[:function](bnchmrk[:ode])
        fname = replace(bnchmrk[:name], " " => "-") * frmt[:extension]
        open(fname, "w") do io
            write(io, as_text)
        end
    end
    cd("..")
end


text:
using Completion
using Rewrite
using Test

f(xs...) = Fn(:f, xs...)
g(xs...) = Fn(:g, xs...)
x, y, z = Variable.([:x, :y, :z])

include("unify.jl")
include("orders.jl")
include("complete.jl")


text:
#Show how to display plots using EasyPlotPlots/Plots.jl
#-------------------------------------------------------------------------------
module CMDimData_SampleUsage

using CMDimData
using CMDimData.EasyPlot
CMDimData.@includepkg EasyPlotPlots

#==Constants
===============================================================================#
demolist = EasyPlot.demofilelist()
backend = :gr
	#no display: {:hdf5}
	#Text-based {:unicodeplots}
	#GUI: {:gr, :inspectdr, :pgfplotsx}
	#Python-based {:pyplot}
	#Browser-based? {:plotly, :plotlyjs}
	##########deprecated: {:gadfly, :bokeh, :winston}

#==Helper functions
===============================================================================#
printsep(label, sep="-") = println("\n", label, "\n", repeat(sep, 80))
printheader(label) = printsep(label, "=")

#==Write EasyPlot plots to file
===============================================================================#
printsep("Write EasyPlot.Plot to file using `backend=:$backend`...")
plot = evalfile(demolist[1])
bld_headless = EasyPlot.getbuilder(:image, :PlotsJl, backend=backend)
	EasyPlot._write(:png, "sample_PlotsJl_$backend.png", bld_headless, plot)
	EasyPlot._write(:svg, "sample_PlotsJl_$backend.svg", bld_headless, plot)

#==Render sample EasyPlot plots
===============================================================================#
disp = EasyPlot.GUIDisplay(:PlotsJl, backend=backend)
for demofile in demolist
	fileshort = basename(demofile)
	printsep("Display $fileshort...")
	plot = evalfile(demofile)
	display(disp, plot)
end

end #module
:SampleCode_Executed


text:
const TO = TrajectoryOptimization

altro = ALTROSolver(Problems.Cartpole()...)
solver = altro.solver_al.solver_uncon
Altro.initialize!(solver)

Z = solver.Z
TO.state_diff_jacobian!(solver.G, solver.model, Z)
TO.dynamics_expansion!(TO.integration(solver), solver.D, solver.model, solver.Z)
TO.error_expansion!(solver.D, solver.model, solver.G)
TO.cost_expansion!(solver.quad_obj, solver.obj, solver.Z, true, true)
TO.error_expansion!(solver.E, solver.quad_obj, solver.model, Z, solver.G)

ΔV1 = Altro.static_backwardpass!(solver)
ΔV2 = Altro.backwardpass!(solver)
@test ΔV1 ≈ ΔV2

# Make sure the backward pass doesn't modify it's state
ΔV = Altro.static_backwardpass!(solver)
ΔV2 = Altro.backwardpass!(solver)
@test ΔV ≈ ΔV1 ≈ ΔV2

# Test allocations
@test (@ballocated Altro.static_backwardpass!($solver) samples=2 evals=2) == 0
@test (@ballocated Altro.backwardpass!($solver) samples=2 evals=2) == 0


text:
using PowerSystems
using NLsolve
const PSY = PowerSystems

include(joinpath(dirname(@__FILE__), "dynamic_test_data.jl"))
include(joinpath(dirname(@__FILE__), "data_utils.jl"))
############### Data Network ########################
sys_dir = joinpath(dirname(@__FILE__), "OMIB.raw")
sys = System(sys_dir, runchecks = false)
add_source_to_ref(sys)
############### Data Dynamic devices ########################

function pvs_simple(source)
    return PeriodicVariableSource(
        name = get_name(source),
        R_th = get_R_th(source),
        X_th = get_X_th(source),
        internal_voltage_bias = 1.0,
        internal_voltage_frequencies = [2 * pi],
        internal_voltage_coefficients = [(1.0, 0.0)],
        internal_angle_bias = 0.0,
        internal_angle_frequencies = [2 * pi],
        internal_angle_coefficients = [(0.0, 1.0)],
    )
end

#Attach dynamic generator
gen = [g for g in get_components(Generator, sys)][1]
case_gen = dyn_gen_second_order(gen)
add_component!(sys, case_gen, gen)

#Attach periodic variable source
source = [s for s in get_components(Source, sys)][1]
pvs = pvs_simple(source)
add_component!(sys, pvs, source)


text:
#############################################################################
# Joulia
# A Large-Scale Spatial Power System Model for Julia
# See https://github.com/JuliaEnergy/Joulia.jl
#############################################################################
# This file contains functions and structs for renewables

"""
    get_res_dict(res_df::DataFrame)

Load renewables.
"""
function get_res_dict(res_df::DataFrame)
	stacked_res = melt(res_df, :Node, variable_name=:Technology,
			           value_name=:Capacity)

	stacked_res[:id] = Symbol.(string.(stacked_res[:Technology]) .* "_"
		                      .* string.(stacked_res[:Node]))

	stacked_res = stacked_res[[:id, :Node, :Technology, :Capacity]]
	stacked_res = stacked_res[stacked_res[:Capacity] .> 0, :]
	res_dict = df_to_dict_with_id(stacked_res)

	return stacked_res[:id], res_dict
end

"""
    get_avail_dict(avail::Dict, res_dict)

Determine renewables availablity.
"""
function get_avail_dict(avail::Dict, res_dict)

	tech = res_dict[:Technology]
	node = res_dict[:Node]

	avail_dict = Dict{Symbol,Array{Float64, 1}}()
	for id in keys(tech)

		t = tech[id]
		n = Symbol(node[id])

		if haskey(avail, t)
			if Symbol(node[id]) in names(avail[t])
				avail_dict[id] = avail[t][n]
			else
				avail_dict[id] = fill(0.0, 8760)
			end
		else
			avail_dict[id] = avail[:global][t]
		end
	end

	return avail_dict
end

"""
    calc_infeed(cap_dict, avail_dict)

Determine renewables infeed.
"""
function calc_infeed(cap_dict, avail_dict)

	infeed = map(collect(keys(cap_dict))) do id
		cap = cap_dict[id]
		avail = avail_dict[id]

		id => avail * cap
	end |> Dict

	return infeed
end

mutable struct RenewableEnergySource
	id
	technology
	node
	infeed

	function RenewableEnergySource(res_df::DataFrame, avail::Dict)

		ids, res_dict = get_res_dict(res_df)
		avail_dict = get_avail_dict(avail,res_dict)
		infeed = calc_infeed(res_dict[:Capacity], avail_dict)

		return new(ids, res_dict[:Technology], res_dict[:Node], infeed)
	end
end


text:
using CUDA.CUTENSOR
using CUDA
using LinearAlgebra

@test has_cutensor()
@test CUTENSOR.version() isa VersionNumber

@testset "CuTensor type basics" begin
    N = 2
    dmax = 2^div(18,N)
    dims = rand(2:dmax, N)
    p = randperm(N)
    indsA = collect(('a':'z')[1:N])
    dimsA = dims
    A = rand(Float64, dimsA...)
    dA = CuArray(A)
    p = randperm(N)
    indsA = collect(('a':'z')[1:N])
    ctA = CuTensor(dA, indsA)
    @test length(ctA) == length(A)
    @test size(ctA) == size(A)
    @test size(ctA, 1) == size(A, 1)
    @test ndims(ctA) == ndims(A)
    @test strides(ctA) == strides(A)
    @test eltype(ctA) == eltype(A)
end

@testset "Elementwise binary" begin
    eltypes = ((Float16, Float16),
               #(Float16, Float32),
               (Float32, Float32),
               #(Float32, Float64),
               (Float64, Float64),
               #(ComplexF16, ComplexF16), (ComplexF16, ComplexF32),
               (ComplexF32, ComplexF32), #(ComplexF32, ComplexF64),
               (ComplexF64, ComplexF64))
    @testset for N=2:5
        @testset for (eltyA, eltyC) in eltypes
            # setup
            eltyD = eltyC
            dmax  = 2^div(18,N)
            dims  = rand(2:dmax, N)
            p     = randperm(N)
            indsA = collect(('a':'z')[1:N])
            indsC = indsA[p]
            dimsA = dims
            dimsC = dims[p]
            A     = rand(eltyA, dimsA...)
            dA    = CuArray(A)
            C     = rand(eltyC, dimsC...)
            dC    = CuArray(C)
            Mem.pin(A)
            Mem.pin(C)

            # simple case
            opA   = CUTENSOR.CUTENSOR_OP_IDENTITY
            opC   = CUTENSOR.CUTENSOR_OP_IDENTITY
            dD    = similar(dC, eltyD)
            opAC  = CUTENSOR.CUTENSOR_OP_ADD
            dD    = CUTENSOR.elementwiseBinary!(1, dA, indsA, opA, 1, dC, indsC, opC,
                                                dD, indsC, opAC)
            D = collect(dD)
            @test D ≈ permutedims(A, p) .+ C
            Dsimple = similar(C)
            Mem.pin(Dsimple)
            Dsimple = CUTENSOR.elementwiseBinary!(1, A, indsA, opA, 1, C, indsC, opC,
                                                  Dsimple, indsC, opAC)
            synchronize()
            @test Dsimple ≈ permutedims(A, p) .+ C

            # using integers as indices
            dD = CUTENSOR.elementwiseBinary!(1, dA, 1:N, opA, 1, dC, p, opC, dD, p, opAC)
            D = collect(dD)
            @test D ≈ permutedims(A, p) .+ C
            Dint = zeros(eltyC, dimsC...)
            Mem.pin(Dint)
            Dint = CUTENSOR.elementwiseBinary!(1, A, 1:N, opA, 1, C, p, opC, Dint, p, opAC)
            synchronize()
            @test Dint ≈ permutedims(A, p) .+ C

            # multiplication as binary operator
            opAC = CUTENSOR.CUTENSOR_OP_MUL
            dD = CUTENSOR.elementwiseBinary!(1, dA, indsA, opA, 1, dC, indsC, opC,
                                                dD, indsC, opAC)
            D = collect(dD)
            @test D ≈ permutedims(A, p) .* C
            Dmult = zeros(eltyC, dimsC...)
            Mem.pin(Dmult)
            Dmult = CUTENSOR.elementwiseBinary!(1, A, indsA, opA, 1, C, indsC, opC,
                                                Dmult, indsC, opAC)
            synchronize()
            @test Dmult ≈ permutedims(A, p) .* C

            # with non-trivial coefficients and conjugation
            opA = eltyA <: Complex ? CUTENSOR.CUTENSOR_OP_CONJ :
                                    CUTENSOR.CUTENSOR_OP_IDENTITY
            opC = CUTENSOR.CUTENSOR_OP_IDENTITY
            opAC = CUTENSOR.CUTENSOR_OP_ADD
            α = rand(eltyD)
            γ = rand(eltyD)
            dD = CUTENSOR.elementwiseBinary!(α, dA, indsA, opA, γ, dC, indsC, opC,
                                                dD, indsC, opAC)
            D = collect(dD)
            @test D ≈ α .* conj.(permutedims(A, p)) .+ γ .* C
            Dnontrivial = similar(C)
            Mem.pin(Dnontrivial)
            Dnontrivial = CUTENSOR.elementwiseBinary!(α, A, indsA, opA, γ, C, indsC, opC,
                                                Dnontrivial, indsC, opAC)
            synchronize()
            @test Dnontrivial ≈ α .* conj.(permutedims(A, p)) .+ γ .* C

            # test in-place, and more complicated unary and binary operations
            opA = eltyA <: Complex ? CUTENSOR.CUTENSOR_OP_IDENTITY :
                                    CUTENSOR.CUTENSOR_OP_SQRT
            # because we use rand, entries of A will be positive when elty is real
            opC = eltyC <: Complex ? CUTENSOR.CUTENSOR_OP_CONJ :
                                    CUTENSOR.CUTENSOR_OP_IDENTITY
            opAC = eltyD <: Complex ? CUTENSOR.CUTENSOR_OP_ADD :
                                    CUTENSOR.CUTENSOR_OP_MAX
            α = rand(eltyD)
            γ = rand(eltyD)
            dD = CUTENSOR.elementwiseBinary!(α, dA, indsA, opA, γ, dC, indsC, opC,
                                                dC, indsC, opAC)
            D = collect(dC)
            if eltyD <: Complex
                if eltyA <: Complex
                    @test D ≈ α .* permutedims(A, p) .+ γ .* conj.(C)
                else
                    @test D ≈ α .* sqrt.(convert.(eltyD, permutedims(A, p))) .+
                                γ .* conj.(C)
                end
            else
                @test D ≈ max.(α .* sqrt.(convert.(eltyD, permutedims(A, p))), γ .* C)
            end            # # using host memory

            # using CuTensor type
            dA = CuArray(A)
            dC = CuArray(C)
            ctA = CuTensor(dA, indsA)
            ctC = CuTensor(dC, indsC)
            ctD = ctA + ctC
            hD = collect(ctD.data)
            @test hD ≈ permutedims(A, p) .+ C
            ctD = ctA - ctC
            hD = collect(ctD.data)
            @test hD ≈ permutedims(A, p) .- C

            α = rand(eltyD)
            ctD = LinearAlgebra.axpy!(α, ctA, ctC)
            hD = collect(ctD.data)
            @test hD ≈ α.*permutedims(A, p) .+ C

            γ = rand(eltyD)
            ctD = LinearAlgebra.axpby!(α, ctA, γ, ctC)
            hD = collect(ctD.data)
            @test hD ≈ α.*permutedims(A, p) .+ γ.*C
        end
    end
end

@testset "Permutations" begin
    eltypes = ((Float16, Float16),
               #(Float16, Float32),
               (Float32, Float32),
               #(Float32, Float64),
               (Float64, Float64),
               #(ComplexF16, ComplexF16),
               #(ComplexF16, ComplexF32),
               (ComplexF32, ComplexF32),
               #(ComplexF32, ComplexF64),
               (ComplexF64, ComplexF64))
    @testset for N=2:5
        @testset for (eltyA, eltyC) in eltypes
            # setup
            dmax = 2^div(18,N)
            dims = rand(2:dmax, N)
            p = randperm(N)
            indsA = collect(('a':'z')[1:N])
            indsC = indsA[p]
            dimsA = dims
            dimsC = dims[p]
            A = rand(eltyA, dimsA...)
            Mem.pin(A)
            dA = CuArray(A)
            dC = similar(dA, eltyC, dimsC...)

            # simple case
            dC = CUTENSOR.permutation!(one(eltyA), dA, indsA, dC, indsC)
            C  = collect(dC)
            @test C == permutedims(A, p) # exact equality
            Csimple = zeros(eltyC, dimsC...)
            Mem.pin(Csimple)
            Csimple = CUTENSOR.permutation!(one(eltyA), A, indsA, Csimple, indsC)
            synchronize()
            @test Csimple == permutedims(A, p) # exact equality

            # with scalar
            α  = rand(eltyA)
            dC = CUTENSOR.permutation!(α, dA, indsA, dC, indsC)
            C  = collect(dC)
            @test C ≈ α * permutedims(A, p) # approximate, floating point rounding

            Cscalar = zeros(eltyC, dimsC...)
            Mem.pin(Cscalar)
            Cscalar = CUTENSOR.permutation!(α, A, indsA, Cscalar, indsC)
            synchronize()
            @test Cscalar ≈ α * permutedims(A, p) # approximate, floating point rounding
        end
    end
end

@testset "Elementwise trinary" begin
    eltypes = ((Float16, Float16, Float16),
               #(Float16, Float32, Float32),
               # (Float32, Float16, Float32),
                (Float32, Float32, Float32),
               # (Float32, Float32, Float64),
               # (Float32, Float64, Float64),
               # (Float64, Float32, Float64),
                (Float64, Float64, Float64),
                (ComplexF32, ComplexF32, ComplexF32),
                (ComplexF64, ComplexF64, ComplexF64))
    @testset for N=2:5
        @testset for (eltyA, eltyB, eltyC) in eltypes
            # setup
            eltyD = eltyC
            dmax = 2^div(18,N)
            dims = rand(2:dmax, N)
            pA = randperm(N)
            ipA = invperm(pA)
            pB = randperm(N)
            ipB = invperm(pB)
            indsC = collect(('a':'z')[1:N])
            dimsC = dims
            indsA = indsC[ipA]
            dimsA = dims[ipA]
            indsB = indsC[ipB]
            dimsB = dims[ipB]
            A = rand(eltyA, dimsA...)
            dA = CuArray(A)
            B = rand(eltyB, dimsB...)
            dB = CuArray(B)
            C = rand(eltyC, dimsC...)
            dC = CuArray(C)
            dD = similar(dC)

            # simple case
            opA = CUTENSOR.CUTENSOR_OP_IDENTITY
            opB = CUTENSOR.CUTENSOR_OP_IDENTITY
            opC = CUTENSOR.CUTENSOR_OP_IDENTITY
            opAB = CUTENSOR.CUTENSOR_OP_ADD
            opABC = CUTENSOR.CUTENSOR_OP_ADD
            dD = CUTENSOR.elementwiseTrinary!(1, dA, indsA, opA, 1, dB, indsB, opB,
                                                1, dC, indsC, opC, dD, indsC, opAB, opABC)
            D = collect(dD)
            @test D ≈ permutedims(A, pA) .+ permutedims(B, pB) .+ C

            # test with Mem.pinned host memory
            Mem.pin(A)
            Mem.pin(B)
            Mem.pin(C)
            Dc = similar(C)
            Mem.pin(Dc)
            Dc = CUTENSOR.elementwiseTrinary!(1, A, indsA, opA, 1, B, indsB, opB,
                                                1, C, indsC, opC, Dc, indsC, opAB, opABC)
            synchronize()
            @test Dc ≈ permutedims(A, pA) .+ permutedims(B, pB) .+ C

            # using integers as indices
            dD = CUTENSOR.elementwiseTrinary!(1, dA, ipA, opA, 1, dB, ipB, opB,
                                                1, dC, 1:N, opC, dD, 1:N, opAB, opABC)
            D = collect(dD)
            @test D ≈ permutedims(A, pA) .+ permutedims(B, pB) .+ C
            Dd = similar(C)
            Mem.pin(Dd)
            Dd = CUTENSOR.elementwiseTrinary!(1, A, ipA, opA, 1, B, ipB, opB,
                                                1, C, 1:N, opC, Dd, 1:N, opAB, opABC)
            synchronize()
            @test Dd ≈ permutedims(A, pA) .+ permutedims(B, pB) .+ C

            # multiplication as binary operator
            opAB = CUTENSOR.CUTENSOR_OP_MUL
            opABC = CUTENSOR.CUTENSOR_OP_ADD
            dD = CUTENSOR.elementwiseTrinary!(1, dA, indsA, opA, 1, dB, indsB, opB,
                                                1, dC, indsC, opC, dD, indsC, opAB, opABC)
            D = collect(dD)
            @test D ≈ (convert.(eltyD, permutedims(A, pA)) .* convert.(eltyD, permutedims(B, pB))) .+ C
            De = similar(C)
            Mem.pin(De)
            De = CUTENSOR.elementwiseTrinary!(1, A, indsA, opA, 1, B, indsB, opB,
                                                1, C, indsC, opC, De, indsC, opAB, opABC)
            synchronize()
            @test De ≈ (convert.(eltyD, permutedims(A, pA)) .* convert.(eltyD, permutedims(B, pB))) .+ C

            opAB = CUTENSOR.CUTENSOR_OP_ADD
            opABC = CUTENSOR.CUTENSOR_OP_MUL
            dD = CUTENSOR.elementwiseTrinary!(1, dA, indsA, opA, 1, dB, indsB, opB,
                                                1, dC, indsC, opC, dD, indsC, opAB, opABC)
            D = collect(dD)
            @test D ≈ (convert.(eltyD, permutedims(A, pA)) .+ convert.(eltyD, permutedims(B, pB))) .* C
            Df = similar(C)
            Mem.pin(Df)
            Df = CUTENSOR.elementwiseTrinary!(1, A, indsA, opA, 1, B, indsB, opB,
                                                1, C, indsC, opC, Df, indsC, opAB, opABC)
            synchronize()
            @test Df ≈ (convert.(eltyD, permutedims(A, pA)) .+ convert.(eltyD, permutedims(B, pB))) .* C

            opAB = CUTENSOR.CUTENSOR_OP_MUL
            opABC = CUTENSOR.CUTENSOR_OP_MUL
            dD = CUTENSOR.elementwiseTrinary!(1, dA, indsA, opA, 1, dB, indsB, opB,
                                                1, dC, indsC, opC, dD, indsC, opAB, opABC)
            D = collect(dD)
            @test D ≈ convert.(eltyD, permutedims(A, pA)) .*
                        convert.(eltyD, permutedims(B, pB)) .* C
            Dg = similar(C)
            Mem.pin(Dg)
            Dg = CUTENSOR.elementwiseTrinary!(1, A, indsA, opA, 1, B, indsB, opB,
                                                1, C, indsC, opC, Dg, indsC, opAB, opABC)
            synchronize()
            @test Dg ≈ convert.(eltyD, permutedims(A, pA)) .*
                        convert.(eltyD, permutedims(B, pB)) .* C

            # with non-trivial coefficients and conjugation
            α = rand(eltyD)
            β = rand(eltyD)
            γ = rand(eltyD)
            opA = eltyA <: Complex ? CUTENSOR.CUTENSOR_OP_CONJ :
                                    CUTENSOR.CUTENSOR_OP_IDENTITY
            opAB = CUTENSOR.CUTENSOR_OP_ADD
            opABC = CUTENSOR.CUTENSOR_OP_ADD
            dD = CUTENSOR.elementwiseTrinary!(α, dA, indsA, opA, β, dB, indsB, opB,
                                                γ, dC, indsC, opC, dD, indsC, opAB, opABC)
            D = collect(dD)
            @test D ≈ α .* conj.(permutedims(A, pA)) .+ β .* permutedims(B, pB) .+ γ .* C
            Dh = similar(C)
            Mem.pin(Dh)
            Dh = CUTENSOR.elementwiseTrinary!(α, A, indsA, opA, β, B, indsB, opB,
                                                γ, C, indsC, opC, Dh, indsC, opAB, opABC)
            synchronize()
            @test Dh ≈ α .* conj.(permutedims(A, pA)) .+ β .* permutedims(B, pB) .+ γ .* C

            opB = eltyB <: Complex ? CUTENSOR.CUTENSOR_OP_CONJ :
                                    CUTENSOR.CUTENSOR_OP_IDENTITY
            opAB = CUTENSOR.CUTENSOR_OP_ADD
            opABC = CUTENSOR.CUTENSOR_OP_ADD
            dD = CUTENSOR.elementwiseTrinary!(α, dA, indsA, opA, β, dB, indsB, opB,
                                                γ, dC, indsC, opC, dD, indsC, opAB, opABC)
            D = collect(dD)
            @test D ≈ α .* conj.(permutedims(A, pA)) .+
                        β .* conj.(permutedims(B, pB)) .+ γ .* C
            Di = similar(C)
            Mem.pin(Di)
            Di = CUTENSOR.elementwiseTrinary!(α, A, indsA, opA, β, B, indsB, opB,
                                                γ, C, indsC, opC, Di, indsC, opAB, opABC)
            synchronize()
            @test Di ≈ α .* conj.(permutedims(A, pA)) .+
                        β .* conj.(permutedims(B, pB)) .+ γ .* C

            opA = CUTENSOR.CUTENSOR_OP_IDENTITY
            opAB = CUTENSOR.CUTENSOR_OP_MUL
            opABC = CUTENSOR.CUTENSOR_OP_ADD
            dD = CUTENSOR.elementwiseTrinary!(α, dA, indsA, opA, β, dB, indsB, opB,
                                                γ, dC, indsC, opC, dD, indsC, opAB, opABC)
            D = collect(dD)
            @test D ≈ α .* permutedims(A, pA) .* β .* conj.(permutedims(B, pB)) .+ γ .* C
            Dj = similar(C)
            Mem.pin(Dj)
            Dj = CUTENSOR.elementwiseTrinary!(α, A, indsA, opA, β, B, indsB, opB,
                                                γ, C, indsC, opC, Dj, indsC, opAB, opABC)
            synchronize()
            @test Dj ≈ α .* permutedims(A, pA) .* β .* conj.(permutedims(B, pB)) .+ γ .* C

            # test in-place, and more complicated unary and binary operations
            opA = eltyA <: Complex ? CUTENSOR.CUTENSOR_OP_IDENTITY :
                                    CUTENSOR.CUTENSOR_OP_SQRT
            opB = eltyB <: Complex ? CUTENSOR.CUTENSOR_OP_IDENTITY :
                                    CUTENSOR.CUTENSOR_OP_SQRT
            # because we use rand, entries of A will be positive when elty is real
            opC = eltyC <: Complex ? CUTENSOR.CUTENSOR_OP_CONJ :
                                    CUTENSOR.CUTENSOR_OP_IDENTITY
            opAB = eltyD <: Complex ? CUTENSOR.CUTENSOR_OP_MUL :
                                    CUTENSOR.CUTENSOR_OP_MIN
            opABC = eltyD <: Complex ? CUTENSOR.CUTENSOR_OP_ADD :
                                    CUTENSOR.CUTENSOR_OP_MAX
            α = rand(eltyD)
            β = rand(eltyD)
            γ = rand(eltyD)
            dD = CUTENSOR.elementwiseTrinary!(α, dA, indsA, opA, β, dB, indsB, opB,
                                                γ, dC, indsC, opC, dC, indsC, opAB, opABC)
            D = collect(dD)
            if eltyD <: Complex
                if eltyA <: Complex && eltyB <: Complex
                    @test D ≈ α .* permutedims(A, pA) .* β .* permutedims(B, pB) .+
                                γ .* conj.(C)
                elseif eltyB <: Complex
                    @test D ≈ α .* sqrt.(convert.(eltyD, permutedims(A, pA))) .*
                                β .* permutedims(B, pB) .+ γ .* conj.(C)
                elseif eltyB <: Complex
                    @test D ≈ α .* permutedims(A, pA) .*
                                β .* sqrt.(convert.(eltyD, permutedims(B, pB))) .+
                                γ .* conj.(C)
                else
                    @test D ≈ α .* sqrt.(convert.(eltyD, permutedims(A, pA))) .*
                                β .* sqrt.(convert.(eltyD, permutedims(B, pB))) .+
                                γ .* conj.(C)
                end
            else
                @test D ≈ max.(min.(α .* sqrt.(convert.(eltyD, permutedims(A, pA))),
                                    β .* sqrt.(convert.(eltyD, permutedims(B, pB)))),
                                γ .* C)
            end
        end
    end
end

@testset "Reduction" begin
    eltypes = (#(Float16, Float16), #(Float16, Float32),
                            (Float32, Float32), #(Float32, Float64),
                            (Float64, Float64),
                            #(ComplexF16, ComplexF16), (ComplexF16, ComplexF32),
                            (ComplexF32, ComplexF32), #(ComplexF32, ComplexF64),
                            (ComplexF64, ComplexF64))
    @testset for NA=2:5, NC = 1:NA-1
        @testset for (eltyA, eltyC) in eltypes
            # setup
            eltyD = eltyC
            dmax = 2^div(18,NA)
            dims = rand(2:dmax, NA)
            p = randperm(NA)
            indsA = collect(('a':'z')[1:NA])
            indsC = indsA[p][1:NC]
            dimsA = dims
            dimsC = dims[p][1:NC]
            A = rand(eltyA, (dimsA...,))
            dA = CuArray(A)
            C = rand(eltyC, (dimsC...,))
            dC = CuArray(C)
            # setup
            Mem.pin(A)

            opA = CUTENSOR.CUTENSOR_OP_IDENTITY
            opC = CUTENSOR.CUTENSOR_OP_IDENTITY
            opReduce = CUTENSOR.CUTENSOR_OP_ADD
            # simple case
            dC = CUTENSOR.reduction!(1, dA, indsA, opA, 0, dC, indsC, opC, opReduce)
            C = collect(dC)
            @test reshape(C, (dimsC..., ones(Int,NA-NC)...)) ≈
                sum(permutedims(A, p); dims = ((NC+1:NA)...,))
            Csimple = zeros(eltyC, dimsC...)
            Mem.pin(Csimple)
            Csimple = CUTENSOR.reduction!(1, A, indsA, opA, 0, Csimple, indsC, opC, opReduce)
            synchronize()
            @test reshape(Csimple, (dimsC..., ones(Int,NA-NC)...)) ≈
                sum(permutedims(A, p); dims = ((NC+1:NA)...,))

            # using integers as indices
            Cinteger = zeros(eltyC, dimsC...)
            Mem.pin(Cinteger)
            dC = CUTENSOR.reduction!(1, dA, collect(1:NA), opA, 0, dC, p[1:NC], opC, opReduce)
            synchronize()
            C = collect(dC)
            @test reshape(C, (dimsC..., ones(Int,NA-NC)...)) ≈
                sum(permutedims(A, p); dims = ((NC+1:NA)...,))
            Cinteger = CUTENSOR.reduction!(1, A, collect(1:NA), opA, 0, Cinteger, p[1:NC], opC, opReduce)
            synchronize()
            @test reshape(Cinteger, (dimsC..., ones(Int,NA-NC)...)) ≈
                sum(permutedims(A, p); dims = ((NC+1:NA)...,))

            # multiplication as reduction operator
            opReduce = CUTENSOR.CUTENSOR_OP_MUL
            Cmult = zeros(eltyC, dimsC...)
            Mem.pin(Cmult)
            dC = CUTENSOR.reduction!(1, dA, indsA, opA, 0, dC, indsC, opC, opReduce)
            synchronize()
            C = collect(dC)
            @test reshape(C, (dimsC..., ones(Int,NA-NC)...)) ≈
                prod(permutedims(A, p); dims = ((NC+1:NA)...,)) atol=eps(Float16) rtol=Base.rtoldefault(Float16)
            Cmult = CUTENSOR.reduction!(1, A, indsA, opA, 0, Cmult, indsC, opC, opReduce)
            synchronize()
            @test reshape(Cmult, (dimsC..., ones(Int,NA-NC)...)) ≈
                prod(permutedims(A, p); dims = ((NC+1:NA)...,)) atol=eps(Float16) rtol=Base.rtoldefault(Float16)
            # NOTE: this test often yields values close to 0 that do not compare approximately

            # with non-trivial coefficients and conjugation
            opA = eltyA <: Complex ? CUTENSOR.CUTENSOR_OP_CONJ :
                                    CUTENSOR.CUTENSOR_OP_IDENTITY
            opC = CUTENSOR.CUTENSOR_OP_IDENTITY
            opReduce = CUTENSOR.CUTENSOR_OP_ADD
            C = rand(eltyC, (dimsC...,))
            dC = CuArray(C)
            α = rand(eltyC)
            γ = rand(eltyC)
            dC = CUTENSOR.reduction!(α, dA, indsA, opA, γ, dC, indsC, opC, opReduce)
            @test reshape(collect(dC), (dimsC..., ones(Int,NA-NC)...)) ≈
                α .* conj.(sum(permutedims(A, p); dims = ((NC+1:NA)...,))) .+ γ .* C
        end
    end
end

@testset "Contraction" begin
    eltypes = ( (Float32, Float32, Float32, Float32),
                (Float32, Float32, Float32, Float16),
                (ComplexF32, ComplexF32, ComplexF32, ComplexF32),
                (Float64, Float64, Float64, Float64),
                (Float64, Float64, Float64, Float32),
                (ComplexF64, ComplexF64, ComplexF64, ComplexF64),
                (ComplexF64, ComplexF64, ComplexF64, ComplexF32)
                )

    @testset for NoA=1:3, NoB=1:3, Nc=1:3
        @testset for (eltyA, eltyB, eltyC, eltyCompute) in eltypes
            # setup
            dmax = 2^div(18, max(NoA+Nc, NoB+Nc, NoA+NoB))
            dimsoA = rand(2:dmax, NoA)
            loA = prod(dimsoA)
            dimsoB = rand(2:dmax, NoB)
            loB = prod(dimsoB)
            dimsc = rand(2:dmax, Nc)
            lc = prod(dimsc)
            allinds = collect('a':'z')
            indsoA = allinds[1:NoA]
            indsoB = allinds[NoA .+ (1:NoB)]
            indsc = allinds[NoA .+ NoB .+ (1:Nc)]
            pA = randperm(NoA + Nc)
            ipA = invperm(pA)
            pB = randperm(Nc + NoB)
            ipB = invperm(pB)
            pC = randperm(NoA + NoB)
            ipC = invperm(pC)
            compute_rtol = (real(eltyCompute) == Float16 || real(eltyC) == Float16) ? 1e-2 : (real(eltyCompute) == Float32 ? 1e-4 : 1e-6)
            dimsA = [dimsoA; dimsc][pA]
            indsA = [indsoA; indsc][pA]
            dimsB = [dimsc; dimsoB][pB]
            indsB = [indsc; indsoB][pB]
            dimsC = [dimsoA; dimsoB][pC]
            indsC = [indsoA; indsoB][pC]

            A = rand(eltyA, (dimsA...,))
            Mem.pin(A)
            dA = CuArray(A)
            mA = reshape(permutedims(A, ipA), (loA, lc))
            B = rand(eltyB, (dimsB...,))
            Mem.pin(B)
            dB = CuArray(B)
            mB = reshape(permutedims(B, ipB), (lc, loB))
            C = zeros(eltyC, (dimsC...,))
            dC = CuArray(C)

            # simple case host side
            Cpin = zeros(eltyC, (dimsC...,))
            Mem.pin(Cpin)
            @test !any(isnan.(Cpin))
            opA = CUTENSOR.CUTENSOR_OP_IDENTITY
            opB = CUTENSOR.CUTENSOR_OP_IDENTITY
            opC = CUTENSOR.CUTENSOR_OP_IDENTITY
            opOut = CUTENSOR.CUTENSOR_OP_IDENTITY
            Cpin  = CUTENSOR.contraction!(1, A, indsA, opA, B, indsB, opB, 0, Cpin, indsC, opC, opOut)
            synchronize()
            mCpin = reshape(permutedims(Cpin, ipC), (loA, loB))
            @test !any(isnan.(A))
            @test !any(isnan.(B))
            @test !any(isnan.(mCpin))
            @test mCpin ≈ mA * mB rtol=compute_rtol

            # simple case
            opA = CUTENSOR.CUTENSOR_OP_IDENTITY
            opB = CUTENSOR.CUTENSOR_OP_IDENTITY
            opC = CUTENSOR.CUTENSOR_OP_IDENTITY
            opOut = CUTENSOR.CUTENSOR_OP_IDENTITY
            dC = CUTENSOR.contraction!(1, dA, indsA, opA, dB, indsB, opB, 0, dC, indsC, opC, opOut, compute_type=eltyCompute)
            C = collect(dC)
            mC = reshape(permutedims(C, ipC), (loA, loB))
            @test mC ≈ mA * mB rtol=compute_rtol

            # simple case with plan storage
            opA = CUTENSOR.CUTENSOR_OP_IDENTITY
            opB = CUTENSOR.CUTENSOR_OP_IDENTITY
            opC = CUTENSOR.CUTENSOR_OP_IDENTITY
            opOut = CUTENSOR.CUTENSOR_OP_IDENTITY
            plan  = CUTENSOR.plan_contraction(dA, indsA, opA, dB, indsB, opB, dC, indsC, opC, opOut)
            dC = CUTENSOR.contraction!(1, dA, indsA, opA, dB, indsB, opB, 0, dC, indsC, opC, opOut, plan=plan)
            C = collect(dC)
            mC = reshape(permutedims(C, ipC), (loA, loB))
            @test mC ≈ mA * mB

            # simple case with plan storage host-side
            opA = CUTENSOR.CUTENSOR_OP_IDENTITY
            opB = CUTENSOR.CUTENSOR_OP_IDENTITY
            opC = CUTENSOR.CUTENSOR_OP_IDENTITY
            opOut = CUTENSOR.CUTENSOR_OP_IDENTITY
            plan  = CUTENSOR.plan_contraction(A, indsA, opA, B, indsB, opB, C, indsC, opC, opOut)
            Cpin = CUTENSOR.contraction!(1, A, indsA, opA, B, indsB, opB, 0, Cpin, indsC, opC, opOut, plan=plan)
            synchronize()
            mC = reshape(permutedims(Cpin, ipC), (loA, loB))
            @test !any(isnan.(mC))
            @test mC ≈ mA * mB

            # simple case with plan storage and compute type
            opA = CUTENSOR.CUTENSOR_OP_IDENTITY
            opB = CUTENSOR.CUTENSOR_OP_IDENTITY
            opC = CUTENSOR.CUTENSOR_OP_IDENTITY
            opOut = CUTENSOR.CUTENSOR_OP_IDENTITY
            plan  = CUTENSOR.plan_contraction(dA, indsA, opA, dB, indsB, opB, dC, indsC, opC, opOut, compute_type=eltyCompute)
            dC = CUTENSOR.contraction!(1, dA, indsA, opA, dB, indsB, opB,
                                        0, dC, indsC, opC, opOut, plan=plan, compute_type=eltyCompute)
            C = collect(dC)
            mC = reshape(permutedims(C, ipC), (loA, loB))
            @test mC ≈ mA * mB rtol=compute_rtol

            # with non-trivial α
            α = rand(eltyCompute)
            dC = CUTENSOR.contraction!(α, dA, indsA, opA, dB, indsB, opB, zero(eltyCompute), dC, indsC, opC, opOut, compute_type=eltyCompute)
            C = collect(dC)
            mC = reshape(permutedims(C, ipC), (loA, loB))
            @test mC ≈ α * mA * mB rtol=compute_rtol

            α = rand(eltyCompute)
            Calpha = zeros(eltyC, (dimsC...,))
            Mem.pin(Calpha)
            @test !any(isnan.(Calpha))
            Calpha = CUTENSOR.contraction!(α, A, indsA, opA, B, indsB, opB, 0, Calpha, indsC, opC, opOut)
            synchronize()
            mCalpha = reshape(permutedims(collect(Calpha), ipC), (loA, loB))
            @test !any(isnan.(mCalpha))
            @test mCalpha ≈ α * mA * mB rtol=compute_rtol

            # with non-trivial β
            C = rand(eltyC, (dimsC...,))
            dC = CuArray(C)
            α = rand(eltyCompute)
            β = rand(eltyCompute)
            copyto!(dC, C)
            dD = CUTENSOR.contraction!(α, dA, indsA, opA, dB, indsB, opB, β, dC, indsC, opC, opOut, compute_type=eltyCompute)
            D = collect(dD)
            mC = reshape(permutedims(C, ipC), (loA, loB))
            mD = reshape(permutedims(D, ipC), (loA, loB))
            @test mD ≈ α * mA * mB + β * mC rtol=compute_rtol
            # with CuTensor objects
            if eltyCompute != Float32 && eltyC != Float16
                ctA = CuTensor(dA, indsA)
                ctB = CuTensor(dB, indsB)
                ctC = CuTensor(dC, indsC)
                ctC = LinearAlgebra.mul!(ctC, ctA, ctB)
                C2, C2inds = collect(ctC)
                mC = reshape(permutedims(C2, ipC), (loA, loB))
                @test mC ≈ mA * mB
                ctC = ctA * ctB
                C2, C2inds = collect(ctC)
                pC2 = convert.(Int, indexin(convert.(Char, C2inds), [indsoA; indsoB]))
                mC = reshape(permutedims(C2, invperm(pC2)), (loA, loB))
                @test mC ≈ mA * mB
            end
            # with conjugation flag for complex arguments
            if !((NoA, NoB, Nc) in ((1,1,3), (1,2,3), (3,1,2)))
            # not supported for these specific cases for unknown reason
                if eltyA <: Complex
                    opA   = CUTENSOR.CUTENSOR_OP_CONJ
                    opB   = CUTENSOR.CUTENSOR_OP_IDENTITY
                    opOut = CUTENSOR.CUTENSOR_OP_IDENTITY
                    dC    = CUTENSOR.contraction!(complex(1.0, 0.0), dA, indsA, opA, dB, indsB, opB,
                                                  0, dC, indsC, opC, opOut, compute_type=eltyCompute)
                    C     = collect(dC)
                    mC    = reshape(permutedims(C, ipC), (loA, loB))
                    @test mC ≈ conj(mA) * mB rtol=compute_rtol
                end
                if eltyB <: Complex
                    opA = CUTENSOR.CUTENSOR_OP_IDENTITY
                    opB = CUTENSOR.CUTENSOR_OP_CONJ
                    opOut = CUTENSOR.CUTENSOR_OP_IDENTITY
                    dC = CUTENSOR.contraction!(complex(1.0, 0.0), dA, indsA, opA, dB, indsB, opB,
                                               complex(0.0, 0.0), dC, indsC, opC, opOut, compute_type=eltyCompute)
                    C = collect(dC)
                    mC = reshape(permutedims(C, ipC), (loA, loB))
                    @test mC ≈ mA*conj(mB) rtol=compute_rtol
                end
                if eltyA <: Complex && eltyB <: Complex
                    opA = CUTENSOR.CUTENSOR_OP_CONJ
                    opB = CUTENSOR.CUTENSOR_OP_CONJ
                    opOut = CUTENSOR.CUTENSOR_OP_IDENTITY
                    dC = CUTENSOR.contraction!(one(eltyCompute), dA, indsA, opA, dB, indsB, opB,
                           zero(eltyCompute), dC, indsC, opC, opOut, compute_type=eltyCompute)
                    C = collect(dC)
                    mC = reshape(permutedims(C, ipC), (loA, loB))
                    @test mC ≈ conj(mA)*conj(mB) rtol=compute_rtol
                end
            end
        end
    end
end