text:
using Sockets: connect

# Script to launch Julia worker process
# NOTE: this file should be run as a standalone script so that all included definitions
# are available in module Main

import Spark: readint, writeint, readobj, writeobj, load_stream, dump_stream
import Spark: END_OF_DATA_SECTION, END_OF_STREAM, JULIA_EXCEPTION_THROWN

# if there are any attached files in the worker directory, include them
for filename in readdir()
    if occursin(r"attached_.{8}\.jl", filename)
        include(joinpath(pwd(), filename))
    end
end

function main()
    port = parse(Int, readline(stdin))
    sock = connect("127.0.0.1", port)
    try
        split = readint(sock)
        func = readobj(sock)[2]
        itc = load_stream(sock)             # return chain representing partition iterator
        dump_stream(sock, Base.invokelatest(func,split, itc))
        writeint(sock, END_OF_DATA_SECTION)
        writeint(sock, END_OF_STREAM)
    catch e
        # TODO: handle the case when JVM closes connection
        io = IOBuffer()
        Base.show_backtrace(io, catch_backtrace())
        seekstart(io)
        bt = read(io, String)
        @info(bt)
        writeint(sock, JULIA_EXCEPTION_THROWN)
        writeobj(sock, string(e) * bt)
    end
end

main()


text:
push!(LOAD_PATH,"../src/")

# using Pkg

# pkg"activate .."

using Documenter, Streamfall

makedocs(sitename="Streamfall Documentation",
    format = Documenter.HTML(
        prettyurls = get(ENV, "CI", nothing) == "true"
    ),
    pages = [
        "index.md",
        "primer.md",
        "network.md",
        "Nodes" => [
            "Node.md",
            "IHACRES.md",
            "HyMod.md",
            "GR4J.md",
            "Dam.md"
        ],
        "use_methods.md",
        "metrics.md",
        "Examples" => [
            "Calibration" => [
                "calibration_setup.md",
                "calibration.md",
            ],
            "simple_showcase.md",
            "model_comparison.md",
            # "multisystem_showcase.md",
        ]

    ]
)

deploydocs(
    repo = "github.com/ConnectedSystems/Streamfall.jl.git",
    devbranch = "main",
    target="build",
    deps=nothing,
    make=nothing
)

ext: jl


text:
module _SampEn
export SampEn
using Statistics: mean, std
using LinearAlgebra: UpperTriangular, I
"""
    Samp, A, B = SampEn(Sig)

 Returns the sample entropy estimates `Samp` and the number of matched state
 vectors (`m`:B, `m+1`:A) for `m` = [0,1,2] estimated from the data sequence `Sig`
 using the default parameters: embedding dimension = 2, time delay = 1,
 radius threshold = 0.2*SD(`Sig`), logarithm = natural

    Samp, A, B = SampEn(Sig::AbstractArray{T,1} where T<:Real; m::Int=2, tau::Int=1, r::Real=0.2*std(Sig,corrected=false), Logx::Real=exp(1))

 Returns the sample entropy estimates `Samp` for dimensions = [0,1,...,`m`]
 estimated from the data sequence `Sig` using the specified keyword arguments:

# Arguments:
 `m`     - Embedding Dimension, a positive integer\n
 `tau`   - Time Delay, a positive integer\n
 `r`     - Radius Distance Threshold, a positive scalar  \n
 `Logx`  - Logarithm base, a positive scalar \n

# See also `ApEn`, `FuzzEn`, `PermEn`, `CondEn`, `XSampEn`, `SampEn2D`, `MSEn`

# References:
    [1] Joshua S Richman and J. Randall Moorman.
        "Physiological time-series analysis using approximate entropy
        and sample entropy."
        American Journal of Physiology-Heart and Circulatory Physiology (2000).

"""
    function SampEn(Sig::AbstractArray{T,1} where T<:Real; m::Int=2, tau::Int=1, r::Real=0.2*std(Sig,corrected=false), Logx::Real=exp(1))

    N = length(Sig)
    (N>10) ? nothing : error("Sig:   must be a numeric vector")
    (m > 0) ? nothing : error("m:     must be an integer > 0")
    (tau > 0) ? nothing : error("tau:   must be an integer > 0")
    (r>=0) ? nothing : error("r:     must be a positive scalar value")
    (Logx>0) ? nothing : error("Logx:     must be a positive number > 0")

    Counter = 1*(abs.(Sig .- transpose(Sig)) .<= r).*UpperTriangular(ones(N,N)) - I(N)
    M = Int.([m*ones(N-m*tau); repeat(collect(m-1:-1:1),inner=tau)])
    A = zeros(m + 1)
    B = zeros(m + 1)
    A[1] = sum(Counter)
    B[1] = N*(N-1)/2

    for n = 1:N-tau
        ix = findall(Counter[n, :] .== 1)
        for k = 1:M[n]
            ix = ix[ix .+ (k*tau) .<= N]
            p1 = repeat(transpose(Sig[n:tau:n+(tau*k)]), length(ix))
            p2 = Sig[ix .+ transpose(collect(0:tau:(k*tau)))]
            ix = ix[findall(maximum(abs.(p1 - p2),dims=2) .<= r)]
            if length(ix)>0
                Counter[n, ix] .+= 1
            else
                break
            end
        end
    end

    for k = 1:m
        A[k+1] = sum(Counter.>k)
        B[k+1] = sum(Counter[:,1:N-(k*tau)].>=k)
    end

    Samp = -log.(Logx, A./B)
    return Samp, A, B
    end

end

"""Copyright 2021 Matthew W. Flood, EntropyHub

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

For Terms of Use see https://github.com/MattWillFlood/EntropyHub"""


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

if !isdefined(@__MODULE__, :DefsOptionalAppId)
    const DefsOptionalAppId = String
else
    @warn("Skipping redefinition of DefsOptionalAppId to String")
end


text:
using Tar, Inflate, SHA

download_urls = [
    "http://dalexander.github.io/admixture/binaries/admixture_linux-1.3.0.tar.gz",
    "http://dalexander.github.io/admixture/binaries/admixture_macosx-1.3.0.tar.gz"
    ]

for download_url in download_urls
    filename = download(download_url)
    println("sha256: ", bytes2hex(open(sha256, filename)))
    println("git-tree-sha1: ", Tar.tree_hash(IOBuffer(inflate_gzip(filename))))
end


text:
# Include this file to conditionally include the helpers.

# Pull in helpers if they're not already there. Including this file in
# test files allows the test files to be run separately, without causing
# duplicate definition or "replacing docs" warnings when the whole suite is
# run.
try
    _HELPERS_INCLUDED_
catch UndefVarError
    include("helpers.jl")
end


text:
using MagneticReadHead
using Test

"""
    DemoException()
A custom exception type for testing that exceptions are thrown
"""
struct DemoException <: Exception
end

@testset "Should not block throwing errors" begin
    foo() = throw(DemoException())
    @test_throws(DemoException, (@run foo()))

    bar() = map(_->foo(), 1:10)
    @test_throws(DemoException, (@run bar()))
end


text:
# A quadratic segment in 3D space that passes through three points: x⃗₁, x⃗₂, and x⃗₃.
# The assumed relation of the points may be seen in the diagram below.
#                 ___x⃗₃___
#            ____/        \____
#        ___/                  \
#     __/                       x⃗₂
#   _/
#  /
# x⃗₁
#
# NOTE: x⃗₃ is between x⃗₁ and x⃗₂, but not necessarily the midpoint.
# q(r) = (2r-1)(r-1)x⃗₁ + r(2r-1)x⃗₂ + 4r(1-r)x⃗₃
# See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
# Chapter 8, Advanced Data Representation, in the interpolation functions section
struct QuadraticSegment_3D{T <: AbstractFloat}
    points::NTuple{3, Point_3D{T}}
end

# Constructors
# -------------------------------------------------------------------------------------------------
QuadraticSegment_3D(p₁::Point_3D{T},
                    p₂::Point_3D{T},
                    p₃::Point_3D{T}) where {T <: AbstractFloat} = QuadraticSegment_3D((p₁, p₂, p₃))

# Methods
# -------------------------------------------------------------------------------------------------
function (q::QuadraticSegment_3D{T})(r::R) where {T <: AbstractFloat, R <: Real}
    # See The Visualization Toolkit: An Object-Oriented Approach to 3D Graphics, 4th Edition
    # Chapter 8, Advanced Data Representation, in the interpolation functions section
    rₜ = T(r)
    return (2rₜ-1)*(rₜ-1)*q.points[1] + rₜ*(2rₜ-1)*q.points[2] + 4rₜ*(1-rₜ)*q.points[3]
end

#function derivative(q::QuadraticSegment_3D{T}, r::R) where {T <: AbstractFloat, R <: Real}
#    # dq⃗/dr
#    rₜ = T(r)
#    return (4rₜ - 3)*q.points[1] + (4rₜ - 1)*q.points[2] + (4 - 8rₜ)*q.points[3]
#end
#
#function arc_length(q::QuadraticSegment_3D{T}; N::Int64=20) where {T <: AbstractFloat}
#    # Mathematica solution is pages long and can produce NaN results when the segment is
#    # straight, so numerical integration is used. (Gauss-Legengre quadrature)
#    #     1                  N
#    # L = ∫ ||q⃗'(r)||dr  ≈   ∑ wᵢ||q⃗'(rᵢ)||
#    #     0                 i=1
#    # The default number of points is N = 20, since the timing difference is very small for
#    # additional accuracy when compared with N = 15, and small N give poor accuracy.
#    w, r = gauss_legendre_quadrature(T, N)
#    return sum(norm.(w .* derivative.(q, r)))
#end
#
#function intersect(l::LineSegment_3D{T}, q::QuadraticSegment_3D{T}) where {T <: AbstractFloat}
#    # q(r) = (2r-1)(r-1)x⃗₁ + r(2r-1)x⃗₂ + 4r(1-r)x⃗₃
#    # q(r) = 2r²(x⃗₁ + x⃗₂ - 2x⃗₃) + r(-3x⃗₁ - x⃗₂ + 4x⃗₃) + x⃗₁
#    # Let D⃗ = 2(x⃗₁ + x⃗₂ - 2x⃗₃), E⃗ = (-3x⃗₁ - x⃗₂ + 4x⃗₃), F⃗ = x₁
#    # q(r) = r²D⃗ + rE⃗ + F⃗
#    # l(s) = x⃗₄ + sw⃗
#    # If D⃗ × w⃗ ≠ 0⃗
#    #   x⃗₄ + sw⃗ = r²D⃗ + rE⃗ + F⃗
#    #   sw⃗ = r²D⃗ + rE⃗ + (F⃗ - x⃗₄)
#    #   0⃗ = r²(D⃗ × w⃗) + r(E⃗ × w⃗) + (F⃗ - x⃗₄) × w⃗
#    #   Let A⃗ = (D⃗ × w⃗), B⃗ = (E⃗ × w⃗), C⃗ = (F⃗ - x⃗₄) × w⃗
#    #   0⃗ = r²A⃗ + rB⃗ + C⃗
#    #   0 = (A⃗ ⋅ A⃗)r² + (B⃗ ⋅ A⃗)r + (C⃗ ⋅ A⃗)
#    #   A = (A⃗ ⋅ A⃗), B = (B⃗ ⋅ A⃗), C = (C⃗ ⋅ A⃗)
#    #   0 = Ar² + Br + C
#    #   r = (-B - √(B²-4AC))/2A, -B + √(B²-4AC))/2A)
#    #   s = ((q(r) - p₄)⋅w⃗/(w⃗ ⋅ w⃗)
#    #   r is invalid if:
#    #     1) A = 0
#    #     2) B² < 4AC
#    #     3) r < 0 or 1 < r   (Curve intersects, segment doesn't)
#    #   s is invalid if:
#    #     1) s < 0 or 1 < s   (Line intersects, segment doesn't)
#    # If D⃗ × w⃗ = 0, we need to use line intersection instead.
#    bool = false
#    npoints = 0
#    points = [Point_3D(T, 0), Point_3D(T, 0)]
#    D⃗ = 2*(q.points[1] + q.points[2] - 2*q.points[3])
#    E⃗ = 4*q.points[3] - 3*q.points[1] - q.points[2]
#    w⃗ = l.points[2] - l.points[1]
#    A⃗ = D⃗ × w⃗
#    B⃗ = E⃗ × w⃗
#    C⃗ = (q.points[1] - l.points[1]) × w⃗
#    A = A⃗ ⋅ A⃗
#    B = B⃗ ⋅ A⃗
#    C = C⃗ ⋅ A⃗
#    if abs(A) < 1.0e-6
#        # Line intersection
#        r = (-C⃗ ⋅ B⃗)/(B⃗ ⋅ B⃗)
#        s = (q(r)- l.points[1]) ⋅ w⃗/(w⃗ ⋅ w⃗)
#        points[1] = q(r)
#        if (0 ≤ s ≤ 1) && (0 ≤ r ≤ 1)
#            bool = true
#            npoints = 1
#        end
#    elseif B^2 ≥ 4A*C
#        # Quadratic intersection
#        r₁ = (-B - √(B^2-4A*C))/2A
#        r₂ = (-B + √(B^2-4A*C))/2A
#        points[1] = q(r₁)
#        points[2] = q(r₂)
#        s₁ = (points[1] - l.points[1]) ⋅ w⃗/(w⃗ ⋅ w⃗)
#        s₂ = (points[2] - l.points[1]) ⋅ w⃗/(w⃗ ⋅ w⃗)
#
#        # Check points to see if they are valid intersections.
#        # First r,s valid?
#        if (0 ≤ r₁ ≤ 1) && (0 ≤ s₁ ≤ 1) && (points[1] ≈ l(s₁))
#            npoints += 1
#        end
#        # Second r,s valid?
#        if (0 ≤ r₂ ≤ 1) && (0 ≤ s₂ ≤ 1) && (points[2] ≈ l(s₂))
#            npoints += 1
#            # If only point 2 is valid, return it in index 1 of points
#            if npoints === 1
#                points[1] = points[2]
#            end
#        end
#        bool = npoints > 0
#    end
#    return bool, npoints, points
#end
#intersect(q::QuadraticSegment_3D, l::LineSegment_3D) = intersect(l, q)

max_stars_repo_name: KyleVaughn/MOCNeutronTransport

max_issues_repo_name: KyleVaughn/MOCNeutronTransport

max_forks_repo_name: KyleVaughn/MOCNeutronTransport


text:
function generate_traindataset2d(datatyp, n_meshes, polydeg, data_input_size)

    n_dims = 2
    X = zeros(data_input_size, 0)
    if datatyp == "NNPP" || datatyp == "NNRH"
        X = zeros(data_input_size, 0)
    elseif datatyp == "CNN"
        X = zeros(data_input_size, data_input_size, 0)
    end
    Y = zeros(2,0)

    n_traindata = zeros(19)
    n_troubledcells = zeros(4)

    # loop over meshs
    for i in 1:n_meshes
        println("Mesh $i")
        #include("meshes/$(n_dims)d/mesh$(i).jl")
        mesh = get_mesh_2d(i)

        leaf_cell_ids = leaf_cells(mesh.tree)
        n_elements = length(leaf_cell_ids)

        #loop over polydeg
        for r in polydeg
            n_nodes = r + 1
            nodes, _ = gauss_lobatto_nodes_weights(n_nodes)
            node_coordinates = zeros(2,n_nodes, n_nodes, n_elements)
            c2e = zeros(Int, length(mesh.tree))

            for element_id in 1:n_elements
                cell_id = leaf_cell_ids[element_id]
                dx = length_at_cell(mesh.tree, cell_id)
                # Calculate node coordinates
                for j in 1:n_nodes
                    for i in 1:n_nodes
                        node_coordinates[1, i, j, element_id] = (mesh.tree.coordinates[1, cell_id] + dx/2 * nodes[i])
                        node_coordinates[2, i, j, element_id] = (mesh.tree.coordinates[2, cell_id] + dx/2 * nodes[j])
                    end
                end
                c2e[leaf_cell_ids[element_id]] = element_id
            end

            #loop over functions
            for func in 1:19
                #println("Funktion $func")
                u(x,y) = trainfunction2d(func,x,y)

                if datatyp == "NNPP"
                    Xi = zeros(4,1)
                    input_data = zeros(4, n_elements)
                    for element_id in 1:n_elements
                        data = node_coordinates[:, :, :, element_id]
                        input_data[:, element_id] .= generate_inputdata2d(u, data, r, datatyp)
                    end
                elseif datatyp == "NNRH"
                    Xi = zeros(15,1)
                    input_data = zeros(3, n_elements)
                    for element_id in 1:n_elements
                        data = node_coordinates[:, :, :, element_id]
                        input_data[:, element_id] .= generate_inputdata2d(u, data, r, datatyp)
                    end
                elseif datatyp == "CNN"
                    Xi = zeros(4,4,1)
                    input_data = zeros(4, 4, n_elements)
                    for element_id in 1:n_elements
                        data = node_coordinates[:, :, :, element_id]
                        input_data[:, :, element_id] .= generate_inputdata2d(u, data, r, datatyp)
                    end
                end

                # for element_id in 1:n_elements
                #     data = node_coordinates[:, :, :, element_id]
                #     input_data[:, element_id] .= generate_inputdata2d(u, data, r, datatyp)
                # end

                for element_id in 1:n_elements
                    cell_id = leaf_cell_ids[element_id]
                    n_traindata[func] += 1
                    Yi = zeros(2,1)

                    # Create Xi
                    Xi .= generate_inputXi2d(element_id, input_data, c2e, mesh, datatyp)

                    # Create label Yi
                    Yi = [0,1]

                    # append to X and Y
                    if datatyp == "NNPP" || datatyp == "NNRH"
                        X = cat(X, Xi, dims=2)
                    elseif datatyp == "CNN"
                        X = cat(X, Xi, dims=3)
                    end

                    Y = cat(Y, Yi, dims=2)
                end

            end

            # troubled cells
            for t in 1:530
                if t < 75
                    func = 1
                    a = rand(Uniform(-100, 100))
                    m = rand(Uniform(-1,1))
                    x0 = rand(Uniform(-0.5, 0.5))
                    y0 = rand(Uniform(-0.5, 0.5))
                    u1(x,y) = troubledcellfunctionabs2d(x, y, a, m, x0, y0)
                elseif t >= 75 && t<200
                    func = 2
                    ui = rand(Uniform(-1,1),4)
                    m = rand(Uniform(0,20))
                    x0 = rand(Uniform(-0.5, 0.5))
                    y0 = rand(Uniform(-0.5, 0.5))
                    u2(x,y) = troubledcellfunctionstep2d(x, y, ui, m, x0, y0)
                elseif t >= 200 && t<400
                    func = 3
                    ui = rand(Uniform(-100,100),4)
                    m = rand(Uniform(0,20))
                    x0 = rand(Uniform(-0.5, 0.5))
                    y0 = rand(Uniform(-0.5, 0.5))
                    u3(x,y) = troubledcellfunctionstep2d(x, y, ui, m, x0, y0)
                elseif t >= 400
                    func = 4
                    ui = rand(Uniform(-10,10),2)
                    m = rand(Uniform(0,0.5))
                    x0 = rand(Uniform(-0.5, 0.5))
                    y0 = rand(Uniform(-0.5, 0.5))
                    u4(x,y)= roundstep2d(x, y, ui, m, x0, y0)
                end

                if datatyp == "NNPP"
                    Xi = zeros(4,1)
                    input_data = zeros(4, n_elements)
                elseif datatyp == "NNRH"
                    Xi = zeros(15,1)
                    input_data = zeros(3, n_elements)
                elseif datatyp == "CNN"
                    Xi = zeros(4,4,1)
                    input_data = zeros(4, 4, n_elements)
                end
                Yi = zeros(2,1)
                if datatyp == "NNPP" || datatyp == "NNRH"
                    for element_id in 1:n_elements
                        data = node_coordinates[:, :, :, element_id]
                        if func == 1
                            input_data[:, element_id] .= generate_inputdata2d(u1, data, r, datatyp)
                        elseif func == 2
                            input_data[:, element_id] .= generate_inputdata2d(u2, data, r, datatyp)
                        elseif func == 3
                            input_data[:, element_id] .= generate_inputdata2d(u3, data, r, datatyp)
                        elseif func == 4
                            input_data[:, element_id] .= generate_inputdata2d(u4, data, r, datatyp)
                        end
                    end
                elseif datatyp == "CNN"
                    for element_id in 1:n_elements
                        data = node_coordinates[:, :, :, element_id]
                        if func == 1
                            input_data[:, :, element_id] .= generate_inputdata2d(u1, data, r, datatyp)
                        elseif func == 2
                            input_data[:, :, element_id] .= generate_inputdata2d(u2, data, r, datatyp)
                        elseif func == 3
                            input_data[:, :, element_id] .= generate_inputdata2d(u3, data, r, datatyp)
                        elseif func == 4
                            input_data[:, :, element_id] .= generate_inputdata2d(u4, data, r, datatyp)
                        end
                    end
                end

                for element_id in 1:n_elements
                    cell_id = leaf_cell_ids[element_id]
                    # Get cell length
                    dx = length_at_cell(mesh.tree, cell_id)
                    # for last functions just troubled cells
                    if good_cell2d(node_coordinates[:, 1, 1, element_id], dx, func, m, x0, y0)
                        continue
                    end
                    n_traindata[func] += 1
                    Yi = zeros(2,1)

                    # Create Xi
                    Xi .= generate_inputXi2d(element_id, input_data, c2e, mesh, datatyp)

                    # Create label Yi
                    Yi = [1,0]

                    # append to X and Y
                    if datatyp == "NNPP" || datatyp == "NNRH"
                        X = cat(X, Xi, dims=2)
                    elseif datatyp == "CNN"
                        X = cat(X, Xi, dims=3)
                    end
                    Y = cat(Y, Yi, dims=2)

                end
            end
            println(size(X))
        end
        println(size(X))
    end

    println(size(X))

    #TODO CNN
    # Scale data
    println("Scale data")
    if datatyp == "NNPP" || datatyp == "NNRH"
        for i in 1:size(X)[2]
            X[:,i]=X[:,i]./max(maximum(abs.(X[:,i])),1)
        end
    elseif datatyp == "CNN"
        for i in 1:size(X)[3]
            X[:,:,i]=X[:,:,i]./max(maximum(abs.(X[:,:,i])),1)
        end
    end

    println("Safe data")
    h5open("datasets/$(n_dims)d/traindata2d$(datatyp).h5", "w") do file
        write(file, "X", X)
        write(file, "Y", Y)
    end

end


text:
module DOFBasesTests

using Test
using Gridap

fun(x) = x[1]+x[2]

D = 2
T = Float64
field = AnalyticalField(fun,D)

basis = MonomialBasis(T,(1,1))

polytope = Polytope(HEX_AXIS,HEX_AXIS)
nodes, nfacenodes = Gridap.RefFEs._high_order_lagrangian_nodes_polytope(polytope,[1,1])
nodes
dofbasis = LagrangianDOFBasis(T,nodes)

field_dofs = [0.0, 1.0, 1.0, 2.0]

basis_dofs = [
 1.0 1.0 1.0 1.0;
 0.0 1.0 0.0 1.0;
 0.0 0.0 1.0 1.0;
 0.0 0.0 0.0 1.0]

test_dof_basis(dofbasis,field,basis,field_dofs,basis_dofs)

T = VectorValue{2,Float64}

fun(x) = VectorValue(x[1]+x[2],x[1])

field = AnalyticalField(fun,D)
basis = MonomialBasis(T,(1,1))

nodes, nfacenodes = Gridap.RefFEs._high_order_lagrangian_nodes_polytope(polytope,[1,1])
dofbasis = LagrangianDOFBasis(T,nodes)

# # Component major
#field_dofs = [0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 2.0, 1.0]

# Node major
field_dofs = [0.0, 1.0, 1.0, 2.0, 0.0, 1.0, 0.0, 1.0]

# Node major
basis_dofs = [
  1.0 1.0 1.0 1.0 0.0 0.0 0.0 0.0;
  0.0 0.0 0.0 0.0 1.0 1.0 1.0 1.0;
  0.0 1.0 0.0 1.0 0.0 0.0 0.0 0.0;
  0.0 0.0 0.0 0.0 0.0 1.0 0.0 1.0;
  0.0 0.0 1.0 1.0 0.0 0.0 0.0 0.0;
  0.0 0.0 0.0 0.0 0.0 0.0 1.0 1.0;
  0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0;
  0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0]

test_dof_basis(dofbasis,field,basis,field_dofs,basis_dofs)

@test dofbasis.dof_to_node == [1, 2, 3, 4, 1, 2, 3, 4]
@test dofbasis.dof_to_comp == [1, 1, 1, 1, 2, 2, 2, 2]
@test dofbasis.node_and_comp_to_dof == [1 5; 2 6; 3 7; 4 8]

end # module


text:
using Revise
using TSExplore
using Plots

plot(hcat([simulate(AR1Model(1.0), zeros(1), randn(101), randn(100).*0.01)[1] for _ in 1:10]...))
plot(hcat([simulate(ARMA11Model(1.0, 1.0), zeros(2), randn(101), randn(100).*0.01)[1] for _ in 1:10]...))

max_stars_repo_name: venuur/TSExplore

max_issues_repo_name: venuur/TSExplore

max_forks_repo_name: venuur/TSExplore


text:
using LogDensityProblems
using Distributions, Parameters
using Calculus, StatsPlots, PositiveFactorizations

include(string(pwd(),"/src/simulation.jl"))
include(string(pwd(),"/gali_bayesian.jl"))
include(string(pwd(),"/mcmc/priors.jl"))
include(string(pwd(),"/mcmc/foos.jl"))

yy,shocks = simulate_dsge(GAMMA_0,GAMMA_1,PSI,PI,500)

prob = DSGE_Model(yy[:,2],1/3)

prob((bet = 0.99,epsilon = 6,theta=2/3,sig=1,s2=1,phi=1,phi_pi=1.5,phi_y=0.5/4,rho_v=0.5))

t = problem_transform(prob)

P = TransformedLogDensity(t,prob)

unif = Uniform(0,1)

adapt = 50000
adapt_warm_up = 5000

num_iter = 100000

npar = 9

pars_aceitos = zeros(num_iter,npar+1)
pars_aceitos[:,1] .= 1/3

true_vals = TransformVariables.inverse(t,(bet = 0.99,epsilon = 6,theta = 2/3,sig = 1,s2 = 1, phi = 1,phi_pi = 1.5,phi_y = 0.5/4, rho_v = 0.5))#

#[rand(prior_bet),rand(prior_epsilon),rand(prior_theta),rand(prior_sig),rand(prior_s2),rand(prior_phi),rand(prior_phi_pi),rand(prior_phi_y),rand(prior_rho_v)]

hes = -Calculus.hessian(x->LogDensityProblems.logdensity(P,x),true_vals)

hes_inv = inv(hes)

########################################################
### Adaptative phase: finding optimal c and hes_inv ###
########################################################

warm_up_pars = zeros(adapt_warm_up,npar+1)
warm_up_pars[:,1] .= 1/3

hes_inv = inv(hes)

hes_inv = (hes_inv + hes_inv')/2

hes_inv = cholesky(Positive,hes_inv)

hes_inv = hes_inv.L*hes_inv.L'

isposdef(hes_inv)

warm_up_pars[1,2:10] = rand(MvNormal(hes_inv),1)

coef_escala = 2.4/sqrt(npar)

rejec = 0

k = 2

while k <= adapt_warm_up
    kernel_velho = MvNormal(warm_up_pars[k-1,2:10],coef_escala*hes_inv)
    novo_par = rand(kernel_velho)
            #pars_aceitos[j,2:10] = novo_par
    kernel_novo = MvNormal(novo_par,coef_escala*hes_inv)

    teste = LogDensityProblems.logdensity(P,novo_par)
    if isnan(teste)
            #pars_aceitos[j,2:10] = pars_aceitos[j-1,2:10]
            #global rejec += 1
            #global j += 1
        continue
    else
        num = teste + logpdf(kernel_novo,warm_up_pars[k-1,2:10])
        #println(string(teste," ",logpdf(kernel_novo, pars_aceitos[j-1,2:10])))
        dem = LogDensityProblems.logdensity(P,warm_up_pars[k-1,2:10]) + logpdf(kernel_velho,novo_par)
        alpha = exp(min(0,num - dem))
        p = rand(unif)
        if alpha < p
            warm_up_pars[k,2:10] = warm_up_pars[k-1,2:10]
            global rejec += 1
        else
            warm_up_pars[k,2:10] = novo_par
        end
    end

    if k % 500 == 0
        acc = 1 - rejec/k
        println("Warm up phase, iteration " ,k, " acceptance ", acc*100,"%")
    end

    if k % 1000 == 0
        pars_space = warm_up_pars[(k-499):k,2:10]
        global hes_inv = cov(pars_space)
        global hes_inv = (hes_inv + hes_inv' )/2
        global hes_inv = cholesky(Positive,hes_inv)
        global hes_inv = hes_inv.L*hes_inv.L'
        global rejec = 0
    end
    global k += 1
end

hes_inv_backup = hes_inv

j = 2
rejec = 0
test = true

coef_escala = 2.4/sqrt(npar)

final_countdown = 0

pars_aceitos[1,2:10] = warm_up_pars[adapt_warm_up,2:10]

while j <= adapt && test
    kernel_velho = MvNormal(pars_aceitos[j-1,2:10],coef_escala*hes_inv)
    novo_par = rand(kernel_velho)
            #pars_aceitos[j,2:10] = novo_par
    kernel_novo = MvNormal(novo_par,coef_escala*hes_inv)

    teste = LogDensityProblems.logdensity(P,novo_par)
    if isnan(teste)
            #pars_aceitos[j,2:10] = pars_aceitos[j-1,2:10]
            #global rejec += 1
            #global j += 1
        continue
    else
        num = teste + logpdf(kernel_novo,pars_aceitos[j-1,2:10])
        #println(string(teste," ",logpdf(kernel_novo, pars_aceitos[j-1,2:10])))
        dem = LogDensityProblems.logdensity(P,pars_aceitos[j-1,2:10]) + logpdf(kernel_velho,novo_par)
        alpha = exp(min(0,num - dem))
        p = rand(unif)
        if alpha < p
            pars_aceitos[j,2:10] = pars_aceitos[j-1,2:10]
            if j % 1200 > 900
                global rejec += 1
            end
        else
            pars_aceitos[j,2:10] = novo_par
        end
    end
    if j % 1200 == 0
        acc = 1 - rejec/300
        #pars_space = pars_aceitos[(j-299):j,2:10]
        # global hes_inv = cov(pars_space)
        # global hes_inv = (hes_inv + hes_inv' )/2
        # global hes_inv = cholesky(Positive,hes_inv)
        # global hes_inv = hes_inv.L*hes_inv.L'
        if round(acc;digits=4) > 0.3
            rescale = 1.1#acc/0.23
            global coef_escala = coef_escala*rescale
            println("Adaptação: iteração ", j, " taxa de aceitação ", acc*100, "% new c = ", coef_escala)
            sleep(0.4)
            #pars_aceitos[j,2:10] = pars_aceitos[1,2:10]
            global j += 1
            global rejec = 0
        elseif round(acc;digits=4) < 0.18
            rescale = 0.8#acc/0.23
            global coef_escala = coef_escala*rescale
            println("Adaptação: iteração ", j, " taxa de aceitação ", acc*100, "% new c = ", coef_escala)
            sleep(0.4)
            #pars_aceitos[j,2:10] = pars_aceitos[1,2:10]
            global j += 1
            global rejec = 0
        else
            acc = 1 - rejec/300
            println("Adaptação: iteração ", j, " taxa de aceitação ", acc*100, "% c = ", coef_escala)
            global rejec = 0
            while final_countdown <= 500
                kernel_velho = MvNormal(pars_aceitos[j+final_countdown-1,2:10],coef_escala*hes_inv)
                novo_par = rand(kernel_velho)
                        #pars_aceitos[j,2:10] = novo_par
                kernel_novo = MvNormal(novo_par,coef_escala*hes_inv)

                teste = LogDensityProblems.logdensity(P,novo_par)
                if isnan(teste)
                        #pars_aceitos[j,2:10] = pars_aceitos[j-1,2:10]
                        #global rejec += 1
                        #global j += 1
                    continue
                else
                    num = teste + logpdf(kernel_novo,pars_aceitos[j+final_countdown-1,2:10])
                    #println(string(teste," ",logpdf(kernel_novo, pars_aceitos[j-1,2:10])))
                    dem = LogDensityProblems.logdensity(P,pars_aceitos[j-1,2:10]) + logpdf(kernel_velho,novo_par)
                    alpha = exp(min(0,num - dem))
                    p = rand(unif)
                    if alpha < p
                        pars_aceitos[j+final_countdown,2:10] = pars_aceitos[j+final_countdown-1,2:10]
                        global rejec += 1
                    else
                        pars_aceitos[j+final_countdown,2:10] = novo_par
                    end
                    acc = 1 - rejec/final_countdown
                end
                println("Adaptação: iteração ", final_countdown, " taxa de aceitação ", acc*100, "% c = ", coef_escala)
                global final_countdown += 1
            end
            if 0.3 > round(acc;digits=4) > 0.18
                global test = false
            else
                global final_countdown = 0
                if round(acc;digits=4) > 0.3
                    rescale = 1.1#acc/0.23
                    global coef_escala = coef_escala*rescale
                    println("Adaptação: iteração ", j, " taxa de aceitação ", acc*100, "% new c = ", coef_escala)
                    sleep(0.4)
                    #pars_aceitos[j,2:10] = pars_aceitos[1,2:10]
                    global j += 1
                    global rejec = 0
                elseif round(acc;digits=4) < 0.18
                    rescale = 0.8#acc/0.23
                    global coef_escala = coef_escala*rescale
                    println("Adaptação: iteração ", j, " taxa de aceitação ", acc*100, "% new c = ", coef_escala)
                    sleep(0.4)
                    #pars_aceitos[j,2:10] = pars_aceitos[1,2:10]
                end

                global rejec = 0
                global j += 1
            end
        end
    else
        global j += 1
    end
end

####################################################
#### RWMC for sampling the true distribution ######
###################################################

j = 2
rejec = 0

pars_aceitos[1,2:10] = rand(MvNormal(hes_inv),1)

#coef_escala = 0.5

while j <= num_iter
    kernel_velho = MvNormal(pars_aceitos[j-1,2:10],coef_escala*hes_inv)
    novo_par = rand(kernel_velho)
        #pars_aceitos[j,2:10] = novo_par
    kernel_novo = MvNormal(novo_par,coef_escala*hes_inv)

    teste = LogDensityProblems.logdensity(P,novo_par)
    if isnan(teste)
        #pars_aceitos[j,2:10] = pars_aceitos[j-1,2:10]
        #global rejec += 1
        #global j += 1
        continue
    else
        num = teste + logpdf(kernel_novo,pars_aceitos[j-1,2:10])
        #println(string(teste," ",logpdf(kernel_novo, pars_aceitos[j-1,2:10])))
        dem = LogDensityProblems.logdensity(P,pars_aceitos[j-1,2:10]) + logpdf(kernel_velho,novo_par)
        alpha = exp(min(0,num - dem))
        p = rand(unif)
        if alpha < p
            pars_aceitos[j,2:10] = pars_aceitos[j-1,2:10]
            global rejec += 1
        else
            pars_aceitos[j,2:10] = novo_par
        end
        acc = 1 - rejec/j
        if j % 50 == 0
            println("Iteração ", j, " taxa de aceitação ", acc)
            sleep(0.4)
        end
        global j += 1
    end
end

max_stars_repo_name: danmrc/DSGE-HMC

max_issues_repo_name: danmrc/DSGE-HMC

max_forks_repo_name: danmrc/DSGE-HMC


text:
module Algorithms

using Nonconvex, ..Functions, Parameters, Requires
using Nonconvex.NonconvexCore: AbstractModel
using ..TopOpt: whichdevice, AbstractDevice, CPU, GPU, PENALTY_BEFORE_INTERPOLATION
using TimerOutputs, Setfield, StaticArrays
using Parameters: @unpack, @pack!
using ..Utilities, Ferrite
using LinearAlgebra, Zygote

export Optimizer,
    SIMP,
    ExponentialContinuation,
    ContinuationSIMP,
    AdaptiveSIMP,
    MMAOptionsGen,
    CSIMPOptions,
    BESO,
    GESO,
    Continuation,
    PowerContinuation

const to = TimerOutput()

abstract type TopOptAlgorithm end

# Solid isotropic material with penalisation
include(joinpath("SIMP", "SIMP.jl"))

# Bidirectional evolutionary strctural optimisation
include("beso.jl")

# Genetic evolutionary structural optimisation
include("geso.jl")

end


text:
# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    Segment(p1, p2)

A line segment with points `p1`, `p2`.
"""
struct Segment{Dim,T,V<:AbstractVector{Point{Dim,T}}} <: Polytope{1,Dim,T}
  vertices::V
end

measure(s::Segment) = norm(s.vertices[2] - s.vertices[1])

facets(s::Segment) = (v for v in s.vertices)

"""
    s1 ∩ s2

Compute the intersection of two line segments `s1` and `s2`.
"""
Base.intersect(s1::Segment, s2::Segment) = get(intersecttype(s1, s2))

abstract type Intersection end

Base.get(I::Intersection) = I.value

struct CrossingSegments{P<:Point} <: Intersection
  value::P
end

struct MidTouchingSegments{P<:Point} <: Intersection
  value::P
end

struct CornerTouchingSegments{P<:Point} <: Intersection
  value::P
end

struct OverlappingSegments{S<:Segment} <: Intersection
  value::S
end

struct NonIntersectingSegments <: Intersection end

Base.get(::NonIntersectingSegments) = nothing

"""
    intersecttype(s1, s2)

Compute the intersection type of two line segments `s1` and `s2`.

The intersection type can be one of five types according to
Balbes, R. and Siegel, J. 1990:

1. intersect at point interior to both segments
2. intersect at end point of one and an interior point of the other
3. intersect at end point of both segments
4. overlap at more than one point
5. do not overlap nor intersect

## References

* Balbes, R. and Siegel, J. 1990. [A robust method for calculating
  the simplicity and orientation of planar polygons]
  (https://www.sciencedirect.com/science/article/abs/pii/0167839691900198)
"""
function intersecttype(s1::Segment{2,T}, s2::Segment{2,T}) where {T}
  x̄ = center(s1)
  ȳ = center(s2)
  x1, x2 = vertices(s1)
  y1, y2 = vertices(s2)
  Q = [x1, y1, x2, y2, x1]

  # winding number of x̄
  windingx = zero(T)
  determinatex = true
  for j in 1:length(Q)-1
    α = ∠(Q[j], x̄, Q[j+1])
    if α ≈ 0 || α ≈ π || α ≈ -π
      determinatex = false
    end
    windingx += α
  end

  # winding number of ȳ
  windingy = zero(T)
  determinatey = true
  for j in 1:length(Q)-1
    β = ∠(Q[j], ȳ, Q[j+1])
    if β ≈ 0 || β ≈ π || β ≈ -π
      determinatey = false
    end
    windingy += β
  end

  # winding number ≈ 0?
  wxzero = isapprox(windingx, 0, atol=atol(T))
  wyzero = isapprox(windingy, 0, atol=atol(T))

  # there are three cases to consider
  # Balbes, R. and Siegel, J. 1990.
  if determinatex && determinatey # CASE (I)
    if !wxzero && !wyzero
      # configuration (1)
      CrossingSegments(intersectpoint(s1, s2))
    else
      # configuration (5)
      NonIntersectingSegments()
    end
  elseif determinatex || determinatey # CASE (II)
    if !(determinatex ? wxzero : wyzero)
      if x1 ≈ y1 || x1 ≈ y2 || x2 ≈ y1 || x2 ≈ y2
        # configuration (3)
        CornerTouchingSegments(intersectpoint(s1, s2))
      else
        # configuration (2)
        MidTouchingSegments(intersectpoint(s1, s2))
      end
    else
      # configuration (5)
      NonIntersectingSegments()
    end
  elseif !determinatex && !determinatey # CASE (III)
    if !isapprox((x2 - x1) × (y2 - y1), 0, atol=atol(T))
      # configuration (3)
      CornerTouchingSegments(intersectpoint(s1, s2))
    else
      # configuration (3), (4) or (5)
      intersectcolinear(s1, s2)
    end
  end
end

# compute the intersection of two line segments assuming that it is a point
function intersectpoint(s1::Segment{Dim,T}, s2::Segment{Dim,T}) where {Dim,T}
  x1, x2 = vertices(s1)
  y1, y2 = vertices(s2)

  v1 = x1 - x2
  v2 = y1 - y2
  v12 = x1 - y1

  # the intersection point lies in between x1 and x2 at a fraction s
  # (https://en.wikipedia.org/wiki/Line-line_intersection#Formulas)
  s = (v12[1] * v2[2] - v12[2] * v2[1]) /
      (v1[1] * v2[2] - v1[2] * v2[1])

  x1 - s*v1
end

# intersection of two line segments assuming that they are colinear
function intersectcolinear(s1::Segment{Dim,T}, s2::Segment{Dim,T}) where {Dim,T}
  # make sure the first segment is larger than the second
  sa, sb = measure(s1) < measure(s2) ? (s2, s1) : (s1, s2)
  x1, x2 = vertices(sa)
  y1, y2 = vertices(sb)

  # given that the segments are colinear we can simply
  # operate on the scalar parameters of the vertices
  # along the line defined by the vector --x1--x2->
  Δ = x2 - x1
  i = findfirst(!iszero, Δ)
  t1 = (y1 - x1)[i] / Δ[i]
  t2 = (y2 - x1)[i] / Δ[i]

  # fix direction of second segment to match the first
  t1 > t2 && ((t1, t2) = (t2, t1))

  if t1 < zero(T) && t2 < zero(T)
    NonIntersectingSegments()
  elseif t1 < zero(T) && t2 ≈ zero(T)
    CornerTouchingSegments(x1)
  elseif t1 < zero(T) && t2 > zero(T)
    OverlappingSegments(Segment(x1, x1 + t2*Δ))
  elseif t1 ≈ zero(T) && t2 > zero(T)
    OverlappingSegments(Segment(x1, x1 + t2*Δ))
  elseif t1 > zero(T) && t2 < one(T)
    OverlappingSegments(Segment(x1 + t1*Δ, x1 + t2*Δ))
  elseif t1 < one(T) && t2 ≈ one(T)
    OverlappingSegments(Segment(x1 + t1*Δ, x2))
  elseif t1 < one(T) && t2 > one(T)
    OverlappingSegments(Segment(x1 + t1*Δ, x2))
  elseif t1 ≈ one(T) && t2 > one(T)
    CornerTouchingSegments(x2)
  elseif t1 > one(T) && t2 > one(T)
    NonIntersectingSegments()
  else
    @error "please report bug"
  end
end


text:
module FFAST

using Requires
using DataFrames
using CSV

# This type identifies the class of MAC to compute
abstract type MassAbsorptionCoefficientType end
# photoelectric mass absorption coefficient
struct PhotoElectricMAC <: MassAbsorptionCoefficientType end
# coherent/incoherent mass attenuation coefficient
struct CoherentIncoherentMAC <: MassAbsorptionCoefficientType end
# sum mass attenuation coefficient
struct TotalMAC <: MassAbsorptionCoefficientType end
# K-shell contribution to the total mass attenuation coefficient
struct KMAC <: MassAbsorptionCoefficientType end

# Used by the fractionaluncertainty function
abstract type MACUncertaintyType end
# The uncertainties associated with the MAC due to a solid or liquid
struct SolidLiquid <: MACUncertaintyType end
# The uncertainties associated with the MAC due to a monatomic gas (typically smaller)
struct MonatomicGas <: MACUncertaintyType end

include("base.jl")

export eachelement
export eachedge
export hasedge
export edgeenergy
export atomicweight
export crosssectionfactor
export density
export EV
export relativisticcorrection
export nuclearthompsoncorrection
export formfactors # Tuple of two form factors in e/atom
export jumpratio # Ratio of the height after the edge over before the edge

export mac
export MassAbsorptionCoefficientType, PhotoElectricMAC, CoherentIncoherentMAC, TotalMAC, KMAC

include("uncertainties.jl")
export MACUncertaintyType, SolidLiquid, MonatomicGas
export fractionaluncertainty # Fractional uncertainties for a SolidLiquid or MonatomicGas

function __init__()
    @require Gadfly = "c91e804a-d5a3-530f-b6f0-dfbca275c004" include("gadflysupport.jl")
end

export plotMACs # Plot [μ/ρ].  Only available if Gadfly is loaded.
export plotFormFactors # Plot form factors.  Only available if Gadfly is loaded.
export plotEdgeCount # Plot edge/shell count.  Only available if Gadfly is loaded.
export plotAtomicWeight # Plot nominal atomic weight.  Only available if Gadfly is loaded.
export plotCrossSectionFactor # Plot crossection factor.  Only available if Gadfly is loaded.
export plotDensity # Plot nominal density.  Only available if Gadfly is loaded.
export plotRelativisticCorrection # Plot the relativistic corretion factors.  Only available if Gadfly is loaded.
export plotNuclearThompsonCorrection # Plot the nuclear Thompson correction.  Only available if Gadfly is loaded.
export plotJumpRatios # Plot the jump ratios for the specified shell

end


text:
using Test
using QuantumOptics
using Random, LinearAlgebra

@testset "phasespace" begin

Random.seed!(0)

D(op1::AbstractOperator, op2::AbstractOperator) = abs(tracedistance_nh(dense(op1), dense(op2)))

# Test quasi-probability functions
b = FockBasis(100)
alpha = complex(0.3, 0.7)
nfock = 3

X = [-2.1:1.5:2;]
Y = [0.13:1.4:3;]
psi_coherent = coherentstate(b, alpha)
rho_coherent = dm(psi_coherent)
psi_fock = fockstate(b, nfock)
rho_fock = dm(psi_fock)

Qpsi_coherent = qfunc(psi_coherent, X, Y)
Qrho_coherent = qfunc(rho_coherent, X, Y)
Qpsi_fock = qfunc(psi_fock, X, Y)
Qrho_fock = qfunc(rho_fock, X, Y)

Wpsi_coherent = wigner(psi_coherent, X, Y)
Wrho_coherent = wigner(rho_coherent, X, Y)
Wpsi_fock = wigner(psi_fock, X, Y)
Wrho_fock = wigner(rho_fock, X, Y)

laguerre3(x) = (-x^3+9x^2-18x+6)/6

for (i,x)=enumerate(X), (j,y)=enumerate(Y)
    beta = 1.0/sqrt(2)*complex(x, y)
    betastate = coherentstate(b, beta)

    q_coherent = 1/pi*exp(-abs2(alpha-beta))
    @test Qpsi_coherent[i, j] ≈ q_coherent
    @test Qrho_coherent[i, j] ≈ q_coherent
    @test qfunc(psi_coherent, beta) ≈ q_coherent
    @test qfunc(rho_coherent, beta) ≈ q_coherent
    @test qfunc(psi_coherent, x, y) ≈ q_coherent
    @test abs2(dagger(betastate) * psi_coherent)/pi ≈ q_coherent
    @test dagger(betastate) * rho_coherent * betastate/pi ≈ q_coherent

    q_fock = 1/pi*exp(-abs2(beta))*abs2(beta)^nfock/factorial(nfock)
    @test Qpsi_fock[i, j] ≈ q_fock
    @test Qrho_fock[i, j] ≈ q_fock
    @test qfunc(psi_fock, beta) ≈ q_fock
    @test qfunc(rho_fock, beta) ≈ q_fock
    @test abs2(dagger(betastate) * psi_fock)/pi ≈ q_fock
    @test dagger(betastate) * rho_fock * betastate/pi ≈ q_fock

    w_coherent = 1/pi*exp(-2*abs2(alpha-beta))
    @test Wpsi_coherent[i, j] ≈ w_coherent
    @test Wrho_coherent[i, j] ≈ w_coherent
    @test wigner(psi_coherent, beta) ≈ w_coherent
    @test wigner(rho_coherent, beta) ≈ w_coherent

    w_fock = 1/pi*(-1)^nfock*laguerre3(4*abs2(beta))*exp(-2*abs2(beta))
    @test Wpsi_fock[i, j] ≈ w_fock
    @test Wrho_fock[i, j] ≈ w_fock
    @test wigner(psi_fock, beta) ≈ w_fock
    @test wigner(rho_fock, beta) ≈ w_fock
end

# Test with offset
alpha = complex(4.3, 2.7)
nfock = 3

X = [0.1:0.1:4;] .+ real(alpha)
Y = [0.13:0.1:3;] .+ imag(alpha)
psi_coherent = coherentstate(b, alpha)
rho_coherent = dm(psi_coherent)
psi_fock = fockstate(b, nfock)
rho_fock = dm(psi_fock)

Qpsi_coherent = qfunc(psi_coherent, X, Y)
Qrho_coherent = qfunc(rho_coherent, X, Y)
Qpsi_fock = qfunc(psi_fock, X, Y)
Qrho_fock = qfunc(rho_fock, X, Y)

Wpsi_coherent = wigner(psi_coherent, X, Y)
Wrho_coherent = wigner(rho_coherent, X, Y)
Wpsi_fock = wigner(psi_fock, X, Y)
Wrho_fock = wigner(rho_fock, X, Y)

b_off = FockBasis(100,3)
psi_coherent_off = coherentstate(b_off, alpha)
rho_coherent_off = dm(psi_coherent_off)
psi_fock_off = fockstate(b_off, nfock)
rho_fock_off = dm(psi_fock)

@test isapprox(qfunc(psi_coherent, alpha), qfunc(psi_coherent_off, alpha))
@test isapprox(Qpsi_coherent, qfunc(psi_coherent_off, X, Y), atol=1e-6)
@test isapprox(Qrho_coherent, qfunc(rho_coherent_off, X, Y), atol=1e-6)
@test isapprox(qfunc(psi_fock, alpha), qfunc(psi_fock_off, alpha))
@test isapprox(Qpsi_fock, qfunc(psi_fock_off, X, Y), atol=1e-6)
@test isapprox(Qrho_fock, qfunc(rho_fock_off, X, Y), atol=1e-6)

@test isapprox(wigner(psi_coherent, alpha), wigner(psi_coherent_off, alpha), atol=1e-6)
@test isapprox(Wpsi_coherent, wigner(psi_coherent_off, X, Y), atol=1e-5, rtol=1e-5)
@test isapprox(Wrho_coherent, wigner(rho_coherent_off, X, Y), atol=1e-5, rtol=1e-5)
@test isapprox(wigner(psi_fock, alpha), wigner(psi_fock_off, alpha), atol=1e-5, rtol=1e-5)
@test isapprox(Wpsi_fock, wigner(psi_fock_off, X, Y), atol=1e-5, rtol=1e-5)
@test isapprox(Wrho_fock, wigner(rho_fock_off, X, Y), atol=1e-5, rtol=1e-5)

# Test qfunc with rand operators
b = FockBasis(50)
psi = randstate(b)
rho = randoperator(b)
X = [-2.1:1.5:2;]
Y = [-0.5:.8:3;]

Qpsi = qfunc(psi, X, Y)
Qrho = qfunc(rho, X, Y)
for (i,x)=enumerate(X), (j,y)=enumerate(Y)
    c = complex(x, y)/sqrt(2)
    state = coherentstate(b, c)
    q_rho = dagger(state) * rho * state/pi
    q_psi = abs2(dagger(state) *psi)/pi
    @test 1e-14 > abs(Qpsi[i,j] - q_psi)
    @test 1e-14 > abs(Qrho[i,j] - q_rho)
    @test 1e-14 > abs(qfunc(psi, c) - q_psi)
    @test 1e-14 > abs(qfunc(rho, c) - q_rho)
end

# Test SU(2) phasespace
b = SpinBasis(5)
theta = π*rand()
phi =2π*rand()
css = coherentspinstate(b,theta,phi)
dmcss = dm(css)
csssx = coherentspinstate(b,π/2,0)
dmcsssx = dm(csssx)
rs = randstate(b)
dmrs = dm(rs)
sx = expect(sigmax(b)/2,css); # eigenstate of jx operator
sy = expect(sigmay(b)/2,css); # eigenstate of jy
sz = expect(sigmaz(b)/2,css); # eigenstate of jz operator
ssq = sx^2 + sy^2 + sz^2

qsu2sx = qfuncsu2(csssx,theta,phi)
qsu2sxdm = qfuncsu2(dmcsssx,theta,phi)
res = 250
costhetam = Array{Float64}(undef,res,2*res)
for i = 1:res, j = 1:2*res
    costhetam[i,j] = sin(i*1pi/(res-1))
end
wsu2 = sum(wignersu2(rs,res).*costhetam)*(π/res)^2
wsu2dm = sum(wignersu2(dmrs,res).*costhetam)*(π/res)^2
qsu2 = sum(qfuncsu2(rs,res).*costhetam)*(π/res)^2
qsu2dm = sum(qfuncsu2(dmrs,res).*costhetam)*(π/res)^2

@test ssq ≈ float(b.spinnumber)^2

@test isapprox(qsu2sxdm, (2*float(b.spinnumber)+1)/(4pi)*(0.5*(sin(theta)cos(phi)+1))^(2*float(b.spinnumber)),atol=1e-2)
@test isapprox(qsu2sx, (2*float(b.spinnumber)+1)/(4pi)*(0.5*(sin(theta)cos(phi)+1))^(2*float(b.spinnumber)),atol=1e-2)
@test isapprox(qsu2, 1.0, atol=1e-2)
@test isapprox(qsu2dm, 1.0, atol=1e-2)

@test isapprox(wignersu2(csssx,π/2,0), (4*float(b.spinnumber)+1)/(4pi), atol=1e-2)
@test isapprox(wignersu2(dmcsssx,π/2,0), (4*float(b.spinnumber)+1)/(4pi),atol=1e-2)
@test isapprox(wsu2, 1.0, atol=1e-2)
@test isapprox(wsu2dm, 1.0, atol=1e-2)

# Test CSS for large spin number (#326)
b = SpinBasis(35)  # S=35 overflows `binomial` function, e.g. binomial(70, 26)
theta = π*rand()
phi =2π*rand()
css = coherentspinstate(b,theta,phi)
css0 = coherentspinstate(b,0,0)
css2 = exp(-0.5im * phi * dense(sigmaz(b))) * exp(-0.5im * theta * dense(sigmay(b))) * css0
@test norm(css - css2) < 1e-12
b_big = SpinBasis(1000)  # > 2000 overflows recursive binomial un-moderated with α^n
css_big = coherentspinstate(b_big,theta,phi)
@test abs(1.0 - norm(css_big)) < 1e-12
# we can go bigger than 10000 without overflowing (≈21000 is the limit), but
# at some point we lose more precision to floating-point errors than we'd
# normally want tolerate.
# @test abs(1 - norm(coherentspinstate(SpinBasis(10000),theta,phi))) < 1e-10

########### YLM test #############
res = 1000
int = 0
l = rand(4:33)
m = rand(0:l-1)
for i = 0:2pi/res:2pi, j = 0:pi/res:pi
    int += sin(j)*abs2(ylm(l,m,j,i))
end
t1 = abs(int*2*(pi/res)^2)
@test isapprox(t1, 1.00, atol=1e-2)

l1 = rand(33:40)
m1 = rand(8:30)
l2 = rand(77:80)
m2 = rand(0:10)
int = 0
for i = 0:2pi/res:2pi, j = 0:pi/res:pi
    int += sin(j)*ylm(l1,m1,j,i)*conj(ylm(l2,m2,j,i))
end
t2 = abs(int*2*(pi/res)^2)
@test isapprox(t2, 0, atol=1e-2)

end # testset


text:
module TestFolds

using CUDA
using Folds
using FoldsCUDA
using Referenceables
using Test

function test_sum_pairs()
    xs = CUDA.rand(Int32, 100)
    @test Folds.sum(last, pairs(xs); init = Int32(0)) == sum(xs)
    VERSION >= v"1.6-" || return
    # Non-type-stable reduction:
    @test Folds.sum(last, pairs(xs)) == sum(xs)
end

function inc!(xs, ex = CUDAEx())
    Folds.foreach(referenceable(xs), ex) do ref
        ptr = pointer(ref.x, LinearIndices(ref.x)[ref.i...])
        CUDA.atomic_add!(ptr, one(ref[]))
    end
    return xs
end

""" Test that each side effect is executed exactly once.  """
function test_side_effects()
    @testset for ex in [CUDAEx(), CoalescedCUDAEx()]
        test_side_effects(ex)
    end
end

function test_side_effects(ex)
    # Using an array size large enough s.t. it uses two blocks but not fully
    # occupy all threads.
    xs = inc!(CUDA.ones(1500), ex)

    @test all(==(2), xs)
end

end  # module


text:
@with_kw mutable struct ZWrapperPOMDP{S, A, O} <: POMDP{S,A,O}
    pomdp::POMDP{S,A,O}
    z
end

POMDPs.initialstate(mdp::ZWrapperPOMDP) = initialstate(mdp.pomdp)
POMDPs.initialobs(mdp::ZWrapperPOMDP, s) = initialobs(mdp.pomdp, s)

POMDPs.actions(mdp::ZWrapperPOMDP) = actions(mdp.pomdp)
POMDPs.actionindex(mdp::ZWrapperPOMDP, a) = actionindex(mdp.pomdp)

POMDPs.isterminal(mdp::ZWrapperPOMDP, s) = isterminal(mdp.pomdp, s)
POMDPs.discount(mdp::ZWrapperPOMDP) = discount(mdp.pomdp)

POMDPs.observation(mdp::ZWrapperPOMDP, s) = observation(mdp.pomdp, s)

render(mdp::ZWrapperPOMDP, s, a = nothing; kwargs...) = render(mdp.pomdp, s, a; kwargs...)
render(mdp::ZWrapperPOMDP; kwargs...) = render(mdp.pomdp; kwargs...)

function POMDPs.gen(mdp::ZWrapperPOMDP, s, a, rng::AbstractRNG=Random.GLOBAL_RNG; info=Dict())
    info["z"] = mdp.z
    gen(mdp.pomdp, s, a, rng, info=info)
    # sp, o, r = gen(mdp.pomdp, s, a, rng, info=info)
    # return (sp=sp, o=vcat(mdp.z, o), r=r)
end

@with_kw mutable struct ZWrapperMDP{S, A} <: MDP{S,A}
    mdp::MDP{S,A}
    z
end

POMDPs.initialstate(mdp::ZWrapperMDP) = initialstate(mdp.mdp)

POMDPs.actions(mdp::ZWrapperMDP) = actions(mdp.mdp)
POMDPs.actionindex(mdp::ZWrapperMDP, a) = actionindex(mdp.mdp)

POMDPs.isterminal(mdp::ZWrapperMDP, s) = isterminal(mdp.mdp, s)
POMDPs.discount(mdp::ZWrapperMDP) = discount(mdp.mdp)

render(mdp::ZWrapperMDP, s, a = nothing; kwargs...) = render(mdp.mdp, s, a; kwargs...)
render(mdp::ZWrapperMDP; kwargs...) = render(mdp.mdp; kwargs...)

function POMDPs.gen(mdp::ZWrapperMDP, s, a, rng::AbstractRNG = Random.GLOBAL_RNG; info=Dict())
    info["z"] = mdp.z
    gen(mdp.mdp, s, a, rng, info=info)
end

ZWrap(mdp::T, z) where T <: MDP = ZWrapperMDP(mdp, z)
ZWrap(pomdp::T, z) where T <: POMDP = ZWrapperPOMDP(pomdp, z)


text:
using SpreadsheetFunctions

sheet_1_a2 = 1
sheet_1_a3 = 2
sheet_1_a4 = 3
sheet_1_a2a4 = [sheet_1_a2; sheet_1_a3; sheet_1_a4]
sheet_1_a5 = SUM(sheet_1_a2a4)
sheet_1_b2 = 4
sheet_1_b3 = 5
sheet_1_b4 = 6
sheet_1_b2b4 = [sheet_1_b2; sheet_1_b3; sheet_1_b4]
sheet_1_b5 = SUM(sheet_1_b2b4)
sheet_1_c2 = 7
sheet_1_c3 = 8
sheet_1_c4 = 9
sheet_1_c2c4 = [sheet_1_c2; sheet_1_c3; sheet_1_c4]
sheet_1_c5 = SUM(sheet_1_c2c4)
sheet_1_a3a5 = [sheet_1_a3; sheet_1_a4; sheet_1_a5]
sheet_1_a6 = SUM(sheet_1_a3a5)
sheet_1_b11 = 1
sheet_1_c11 = 2
sheet_1_d11 = 3
sheet_1_b11d11 = [sheet_1_b11 sheet_1_c11 sheet_1_d11]
sheet_1_e11 = SUM(sheet_1_b11d11)
sheet_1_b12 = 4
sheet_1_c12 = 5
sheet_1_d12 = 6
sheet_1_b12d12 = [sheet_1_b12 sheet_1_c12 sheet_1_d12]
sheet_1_e12 = SUM(sheet_1_b12d12)
sheet_1_b13 = 7
sheet_1_c13 = 8
sheet_1_d13 = 9
sheet_1_b13d13 = [sheet_1_b13 sheet_1_c13 sheet_1_d13]
sheet_1_e13 = SUM(sheet_1_b13d13)
sheet_1_b11b13 = [sheet_1_b11; sheet_1_b12; sheet_1_b13]
sheet_1_b14 = SUM(sheet_1_b11b13)
sheet_1_c11c13 = [sheet_1_c11; sheet_1_c12; sheet_1_c13]
sheet_1_c14 = SUM(sheet_1_c11c13)
sheet_1_d11d13 = [sheet_1_d11; sheet_1_d12; sheet_1_d13]
sheet_1_d14 = SUM(sheet_1_d11d13)
sheet_1_e11e13 = [sheet_1_e11; sheet_1_e12; sheet_1_e13]
sheet_1_e14 = SUM(sheet_1_e11e13)
sheet_1_b14d14 = [sheet_1_b14 sheet_1_c14 sheet_1_d14]
sheet_1_e15 = SUM(sheet_1_b14d14)
sheet_1_b1 = "B"
sheet_1_b6 = nothing
sheet_1_b7 = nothing
sheet_1_b8 = nothing
sheet_1_b9 = nothing
sheet_1_b10 = "Widget"
sheet_1_b15 = nothing
sheet_1_b16 = nothing
sheet_1_bb = [sheet_1_b1; sheet_1_b2; sheet_1_b3; sheet_1_b4; sheet_1_b5; sheet_1_b6; sheet_1_b7; sheet_1_b8; sheet_1_b9; sheet_1_b10; sheet_1_b11; sheet_1_b12; sheet_1_b13; sheet_1_b14; sheet_1_b15; sheet_1_b16]
sheet_1_a16 = SUM(sheet_1_bb)+5
sheet_1_a1 = "A"
sheet_2_d3 = sheet_1_a1
sheet_2_e3 = sheet_1_b1
sheet_2_a2 = 1
sheet_2_a3 = 2
sheet_2_a4 = 6
sheet_2_a2a4 = [sheet_2_a2; sheet_2_a3; sheet_2_a4]
sheet_2_a5 = SUM(sheet_2_a2a4)
sheet_2_b2 = 4
sheet_2_b3 = 5
sheet_2_b4 = 6
sheet_2_b2b4 = [sheet_2_b2; sheet_2_b3; sheet_2_b4]
sheet_2_b5 = SUM(sheet_2_b2b4)
sheet_2_c2 = 7
sheet_2_c3 = 8
sheet_2_c4 = 9
sheet_2_c2c4 = [sheet_2_c2; sheet_2_c3; sheet_2_c4]
sheet_2_c5 = PRODUCT(sheet_2_c2c4)
sheet_2_a1 = "A"
sheet_2_b1 = "B"
sheet_2_c1 = "C"
sheet_2_d1 = 1
sheet_2_e1 = nothing
sheet_2_11 = [sheet_2_a1 sheet_2_b1 sheet_2_c1 sheet_2_d1 sheet_2_e1]
sheet_2_c6 = CONCATENATE(sheet_2_11)
sheet_2_d5 = nothing
sheet_2_e5 = 3
sheet_2_d6 = nothing
sheet_2_e6 = nothing
sheet_2_d7 = nothing
sheet_2_e7 = nothing
sheet_2_d5f7 = [sheet_2_d5 sheet_2_e5; sheet_2_d6 sheet_2_e6; sheet_2_d7 sheet_2_e7]
sheet_2_d8 = PRODUCT(sheet_2_d5f7)


text:
struct ClusterInfo
    data::Matrix{Float64}
    label::Int
    index::Vector{Int}
    size::Int
    df::Float64
    center::Vector{Float64}
    cov::Matrix{Float64}
    function ClusterInfo(x, index, kmeans, label::Int)
        c = findall(kmeans.labels.==label)
        data = x[c, :]
        index = index[c]
        cluster_size = size(data, 1)
        features = size(data, 2)
        df = features*(features+3)/2
        centers = kmeans.μ[label, :]
        data_cov = cov(data)
        return new(data, label, index, cluster_size, df, centers, data_cov)
    end
end

function build_cluster_info(x, k_means::Kmeans; index=nothing)
    if index === nothing
        index = collect(1:size(x, 1))
    end
    labels = 1:n_clusters(k_means)

    return [ClusterInfo(x, index, k_means, label) for label in labels]
end

function log_likehood(c::ClusterInfo)
    return sum([logpdf(c.data[i, :], c.center, c.cov) for i in 1 : size(c.data, 1)])
end
function cluster_bic(c::ClusterInfo)
    return -2*log_likehood(c)+c.df*log(c.size)
end

"""
    Xmeans(; kinit=2, max=300, th=1e-4)
Xmeans method. This algorithm uses BIC to determine if a cluster should be split in two.
In other words, you don't need to give the number of clusters, you can cluster only with the data.

# Parameters:
- `kinit`: initial number of class
- `max`: maximum number of repitition
- `th`: converge threshold

# Example
```jldoctest
julia> model = Xmenas()
Kmeans{3}(Matrix{Float64}(undef, 0, 0), 100000000, 0.0001)

julia> using HorseML.Clustering: fit!

julia> fit!(model, x)

julia> model.labels |> size
(100, 3)
```
"""
mutable struct Xmeans
    kinit::Int
    max::Int64
    th::Float64
    labels::Vector{Int}
    centers::Vector{Vector{Float64}}
    log_likehoods::Vector{Float64}
    sizes::Vector{Int}
    Xmeans(; kinit=2, max=300, th=1e-4) = new(kinit, max, th, [], [], [], [])
end

function fit!(model::Xmeans, x)
    xmeans_clusters = ClusterInfo[]

    kmeans = Kmeans(model.kinit, max=model.max, th=model.th)
    fit!(kmeans, x)
    clusters = build_cluster_info(x, kmeans)
    recursively_split!(model, clusters, xmeans_clusters)

    labels = Array{Int}(undef, size(x, 1))
    for (i, c) in enumerate(xmeans_clusters)
        labels[c.index] .= i
    end
    model.labels = labels
    model.centers = [cluster.center for cluster in xmeans_clusters]
    model.log_likehoods = [log_likehood(cluster) for cluster in xmeans_clusters]
    model.sizes = [cluster.size for cluster in xmeans_clusters]
end

function recursively_split!(
        model::Xmeans,
        clusters::Vector{ClusterInfo},
        xmeans_clusters::Vector{ClusterInfo}
    )
    for cluster in clusters
        if cluster.size <= 3
            push!(xmeans_clusters)
            continue
        end
        kmeans = Kmeans(2, max=model.max, th=model.th)
        fit!(kmeans, cluster.data)
        c1, c2 = build_cluster_info(cluster.data, kmeans, index=cluster.index)

        β = norm(c1.center-c2.center) / sqrt(det(c1.cov) + det(c2.cov))
        α = 0.5 / norm_cdf(β)
        bic = -2*(cluster.size * log(α) + log_likehood(c1) + log_likehood(c2)) + 2*cluster.df*log(cluster.size)

        if bic < cluster_bic(cluster)
            recursively_split!(model, [c1, c2], xmeans_clusters)
        else
            push!(xmeans_clusters, cluster)
        end
    end
end


text:
import Base:   ==
"""
    Bounds(Real2D,Real2D,Real2D,Real2D)

    ....
"""
struct Bounds{T<:Real, P<:Position, F<:Field}
    nw::P
    ne::P
    sw::P
    se::P
    f::F
    function Bounds(f::Field2D{T}, bagID::Real2D{Int}) where {T<:Real}
        nw = Real2D{T}(bagID.x * f.discretization, bagID.y * f.discretization)
        ne = Real2D{T}(nw.x, min(nw.y + f.discretization, f.height))
        sw = Real2D{T}(min(nw.x+f.discretization, f.width), nw.y)
        se = Real2D{T}(sw.x, ne.y)
        bounds = new{T,Real2D{T},Field2D{T}}(nw, ne, sw, se, f)
        return bounds
    end
end

function ==(b1::Bounds{T,P,F}, b2::Bounds{T,P,F}) where {T<:Real, P<:Position, F<:Field}
        b1.nw == b2.nw &&
            b1.ne == b2.ne &&
                 b1.sw == b2.sw &&
                    b1.sw == b2.sw
end

function checkBoundCircle(b::Bounds{T,P,F}, center::P, radius::T, toroidal::Bool) where {T<:Real, P<:Position, F<:Field}
        if distance(b.nw, center, b.f.width, b.f.height,toroidal) <= radius &&
                distance(b.ne, center, b.f.width, b.f.height,toroidal) <= radius &&
                    distance(b.sw, center, b.f.width, b.f.height,toroidal) <= radius &&
                        distance(b.se, center, b.f.width, b.f.height,toroidal) <= radius
                    return 1
        elseif distance(b.nw, center, b.f.width, b.f.height,toroidal) > radius &&
                distance(b.ne, center, b.f.width, b.f.height,toroidal) > radius &&
                    distance(b.sw, center, b.f.width, b.f.height,toroidal) > radius &&
                        distance(b.se, center, b.f.width, b.f.height,toroidal) > radius
                            return -1

        end
        return 0

end


text:
abstract type AbstractShootingDataManager end

# ShootingDataManager: Standard shooting data manager for solving an
# indirect trajectory optimization problem. Does not include any data
# for continuation
struct ShootingDataManager <: AbstractShootingDataManager
    # Vector containing initial guess data
    initGuessData::Vector{Float64}

    # Initial and final boundary conditions
    ics::Vector{Float64}
    fcs::Vector{Float64}

    # Preallocated function input vectors
    inVec::Vector{Float64}
    outVec::Vector{Float64}

    # Solution data
    sol::Vector{Float64}
    cflag::Bool
end

function ShootingDataManager()
    ShootingDataManager(Vector{Float64}(undef, 0),
                        Vector{Float64}(undef, 0),
                        Vector{Float64}(undef, 0),
                        Vector{Float64}(undef, 0),
                        Vector{Float64}(undef, 0),
                        Vector{Float64}(undef, 0),
                        false)
end

# ShootingHomotopyDataManager: Shooting data manager when performing
# homotopic continuation
struct ShootingHomotopyDataManager <: AbstractShootingDataManager
    # Vector containing initial guess data
    initGuessData::Vector{Float64}

    # Initial and final boundary conditions
    ics::Vector{Float64}
    fcs::Vector{Float64}

    # Preallocated function input vectors
    inVec::Vector{Float64}
    outVec::Vector{Float64}

    # Continuation parameter vector
    ϵs::Vector{Float64}

    # continuation solution vector
    sols::Vector{Vector{Float64}}
    cflags::Vector{Bool}
end

function ShootingHomotopyDataManager()
    ShootingHomotopyDataManager(Vector{Float64}(undef, 0),
                        Vector{Float64}(undef, 0),
                        Vector{Float64}(undef, 0),
                        Vector{Float64}(undef, 0),
                        Vector{Float64}(undef, 0),
                        Vector{Float64}(undef, 0),
                        Vector{Vector{Float64}}(undef, 0),
                        Vector{Bool}(undef, 0))
end

function SetInitGuessData!(sdm::AbstractShootingDataManager, igv::AbstractVector)
    resize!(sdm.initGuessData, length(igv))
    sdm.initGuessData .= igv
    return nothing
end

function SetInitConditions!(sdm::AbstractShootingDataManager, icv::AbstractVector)
    resize!(sdm.ics, length(icv))
    sdm.ics .= icv
    return nothing
end

function SetFinConditions!(sdm::AbstractShootingDataManager, fcv::AbstractVector)
    resize!(sdm.fcs, length(fcv))
    sdm.fcs .= fcv
    return nothing
end

function SetPreallocatedVecs!(sdm::AbstractShootingDataManager, size::Int)
    resize!(sdm.inVec, size)
    resize!(sdm.outVec, size)
    return nothing
end

function SetPreallocatedVecs!(sdm::AbstractShootingDataManager, inVec::AbstractVector)
    n = length(inVec)
    resize!(sdm.inVec, n)
    resize!(sdm.outVec, n)
    sdm.inVec .= inVec
    return nothing
end

function GetInitGuessData(sdm::AbstractShootingDataManager)
    return sdm.initGuessData
end

function GetInitConditions(sdm::AbstractShootingDataManager)
    return sdm.ics
end

function GetFinConditions(sdm::AbstractShootingDataManager)
    return sdm.fcs
end

function GetInputVec(sdm::AbstractShootingDataManager)
    return sdm.inVec
end

function GetOutputVec(sdm::AbstractShootingDataManager)
    return sdm.outVec
end

function SetHomotopyParams!(sdm::ShootingHomotopyDataManager, ϵs::AbstractVector)
    n = length(ϵs)
    resize!(sdm.ϵs, n)
    sdm.ϵs .= ϵs
    return nothing
end

function GetHomotopyParams(sdm::ShootingHomotopyDataManager)
    return sdm.ϵs
end

function Initialize!(sdm::AbstractShootingDataManager)
    return nothing
end

function Initialize!(sdm::ShootingDataManager)
    m = length(sdm.initGuessData)
    resize!(sdm.sol, m)
end

function Initialize!(sdm::ShootingHomotopyDataManager)
    # Need to size sols vector of vectors
    n = length(sdm.ϵs)
    m = length(sdm.initGuessData)
    resize!(sdm.cflags, n)
    resize!(sdm.sols, n)
    for i in 1:n
        sdm.sols[i] = Vector{Float64}(undef, m)
    end
    #for i in 1:n
    #    push!(sdm.sols, Vector{Float64}(undef, m))
    #end
    return nothing
end

function GetSolution(sdm::ShootingDataManager)
    return sdm.sol
end

function GetHomotopySolutionVector(sdm::ShootingHomotopyDataManager)
    return sdm.sols
end

function GetHomotopyConvergenceFlags(sdm::ShootingHomotopyDataManager)
    return sdm.cflags
end

function GetInitialGuessConverged(sdm::ShootingDataManager)
    return sdm.cflag
end

function GetInitialGuessConverged(sdm::ShootingHomotopyDataManager)
    return sdm.cflags[1]
end

function GetHomotopyConverged(sdm::ShootingHomotopyDataManager)
    return sdm.cflags[end]
end


text:
using Base: @pure, @propagate_inbounds

@inline function fllog2(x::T) where {T <: Integer}
    return (sizeof(T) * 8 - 1) - leading_zeros(x)
end
@inline function cllog2(x::T) where {T <: Integer}
    return fllog2(x - 1) + 1
end

@inline function fld2(x::T) where {T <: Integer}
    return x >>> true
end
@inline function cld2(x::T) where {T <: Integer}
    return (x + true) >>> true
end

@pure nbits(::Type{T}) where {T} = sizeof(T) * 8
@pure log2nbits(::Type{T}) where {T} = fllog2(nbits(T))

@inline function flpow1m(x::T) where {T <: Integer}
    return (1 << x) - 1
end

@inline undefs(T::Type, dims::Vararg{Any, N}) where {N} = Array{T, N}(undef, dims...)

zero!(arr) = fill!(arr, zero(eltype(arr)))
one!(arr) = fill!(arr, one(eltype(arr)))

@inline maptuple(f::F) where {F} = ()
@inline maptuple(f::F, x) where {F} = (f(x),)
@inline maptuple(f::F, x, y, z...) where {F} = (f(x), maptuple(f, y, z...)...)

function pattern(A::SparseMatrixCSC{Tv, Ti}) where {Tv, Ti}
    return SparseMatrixCSC{Bool, Ti}(size(A)..., A.colptr, A.rowval, ones(Bool, nnz(A)))
end

macro stabilize(args...)
    @assert all(map(arg -> arg isa Symbol, args[1:end-1]))
    syms = args[1:end-1]
    body = args[end]
    name = gensym(:stabilized)
    name_inbounds = gensym(:stabilized_inbounds)
    Base.eval(__module__, quote
        function $name($(syms...))
            begin
                $body
            end
        end
        function $name_inbounds($(syms...))
            @inbounds begin
                $body
            end
        end
    end)
    inbounds = gensym(:inbounds)
    return esc(quote
        $inbounds = true
        @boundscheck($inbounds = false)
        if $inbounds
            $name_inbounds($(syms...))
        else
            $name($(syms...))
        end
    end)
end

function adjointpattern(A::SparseMatrixCSC{Tv, Ti}) where {Tv, Ti}
    @inbounds begin
        m, n = size(A)
        N = nnz(A)
        pos = A.colptr
        idx = A.rowval
        pos′ = zeros(Ti, m + 1)
        idx′ = undefs(Ti, N)
        for j = 1:n
            for q = pos[j]:pos[j + 1]-1
                i = idx[q]
                pos′[i+1] += 1
            end
        end
        tmp = 1
        for i = 1:m + 1
            (pos′[i], tmp) = (tmp, tmp + pos′[i])
        end
        for j = 1:n
            for q = pos[j]:pos[j + 1]-1
                i = idx[q]
                q′ = pos′[i + 1]
                idx′[q′] = j
                pos′[i + 1] = q′ + 1
            end
        end
        return SparseMatrixCSC(n, m, pos′, idx′, A.nzval)
    end
end


text:
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

module TestBind
using MXNet
using Test

using ..Main: rand_dims

################################################################################
# Test Implementations
################################################################################
function test_arithmetic(::Type{T}, uf, gf) where T <: mx.DType
  shape = rand_dims()
  @info "Bind::arithmetic::$T::$uf::dims = $shape"

  lhs = mx.Variable(:lhs)
  rhs = mx.Variable(:rhs)
  ret = uf(lhs, rhs)
  @test mx.list_arguments(ret) == [:lhs, :rhs]

  lhs_arr  = NDArray(rand(T, shape))
  rhs_arr  = NDArray(rand(T, shape))
  lhs_grad = NDArray{T}(undef, shape)
  rhs_grad = NDArray{T}(undef, shape)

  exec2 = mx.bind(ret, mx.Context(mx.CPU), [lhs_arr, rhs_arr], args_grad=[lhs_grad, rhs_grad])
  exec3 = mx.bind(ret, mx.Context(mx.CPU), [lhs_arr, rhs_arr])
  exec4 = mx.bind(ret, mx.Context(mx.CPU), Dict(:lhs=>lhs_arr, :rhs=>rhs_arr),
                  args_grad=Dict(:rhs=>rhs_grad, :lhs=>lhs_grad))

  mx.forward(exec2)
  mx.forward(exec3)
  mx.forward(exec4)

  out1 = uf(copy(lhs_arr), copy(rhs_arr))
  out2 = copy(exec2.outputs[1])
  out3 = copy(exec3.outputs[1])
  out4 = copy(exec4.outputs[1])
  @test isapprox(out1, out2)
  @test isapprox(out1, out3)
  @test isapprox(out1, out4)

  # test gradients
  out_grad = mx.NDArray(ones(T, shape))
  lhs_grad2, rhs_grad2 = gf(copy(out_grad), copy(lhs_arr), copy(rhs_arr))
  mx.backward(exec2, out_grad)
  @test isapprox(copy(lhs_grad), lhs_grad2)
  @test isapprox(copy(rhs_grad), rhs_grad2)

  # reset grads
  lhs_grad[:] = 0
  rhs_grad[:] = 0
  # compute using another binding
  mx.backward(exec4, out_grad)
  @test isapprox(copy(lhs_grad), lhs_grad2)
  @test isapprox(copy(rhs_grad), rhs_grad2)
end

function test_arithmetic()
  for T in [mx.fromTypeFlag(TF) for TF in instances(mx.TypeFlag)]
    test_arithmetic(T, (x,y) -> x .+ y, (g,x,y) -> (g,g))
    test_arithmetic(T, (x,y) -> x .- y, (g,x,y) -> (g,-g))
    test_arithmetic(T, (x,y) -> x .* y, (g,x,y) -> (y.*g, x.*g))
    if T <: Integer || T == Float16
      @warn "Not running division test for $T"
    else
      test_arithmetic(T, (x,y) -> x ./ y, (g,x,y) -> (g ./ y, -x .* g ./ (y.^2)))
    end
  end
end

function test_forward()
  # forward with data keyword argument
  x = @var x
  y = x .+ 42

  A = 1:5
  B = A .+ 42

  e = bind(y, args = Dict(:x => NDArray(24:28)))
  z = forward(e, x = NDArray(A))[1]

  @test copy(z) == collect(B)
end

################################################################################
# Run tests
################################################################################
@testset "Bind Test" begin
  test_arithmetic()
  test_forward()
end

end


text:
# cache

struct SinkhornBarycenterGibbsCache{U,V,KT,A}
    u::U
    v::V
    K::KT
    Kv::U
    a::A
end

# solver cache
function build_cache(
    ::Type{T}, ::SinkhornGibbs, size2::Tuple, μ::AbstractMatrix, C::AbstractMatrix, ε::Real
) where {T}
    # compute Gibbs kernel (has to be mutable for ε-scaling algorithm)
    K = similar(C, T)
    @. K = exp(-C / ε)

    # create and initialize dual potentials
    u = similar(μ, T, size(μ, 1), size2...)
    v = similar(μ, T, size(μ, 1), size2...)
    a = similar(μ, T, size(μ, 1), 1)
    fill!(u, one(T))
    fill!(v, one(T))
    fill!(a, one(T))

    # cache for next iterate of `u`
    Kv = similar(u)

    return SinkhornBarycenterGibbsCache(u, v, K, Kv, a)
end

# Sinkhorn algorithm steps (see solve!)
prestep!(::SinkhornBarycenterSolver{SinkhornGibbs}, ::Int) = nothing

function init_step!(solver::SinkhornBarycenterSolver{SinkhornGibbs})
    return A_batched_mul_B!(solver.cache.Kv, solver.cache.K, solver.cache.v)
end

function step!(solver::SinkhornBarycenterSolver{SinkhornGibbs}, iter::Int)
    μ = solver.source
    w = solver.w
    cache = solver.cache
    u = cache.u
    v = cache.v
    Kv = cache.Kv
    K = cache.K
    a = cache.a

    a .= prod(Kv' .^ w; dims=1)'  # TODO: optimise
    u .= a ./ Kv
    At_batched_mul_B!(v, K, u)
    v .= μ ./ v
    return A_batched_mul_B!(Kv, K, v)
end

function check_convergence(solver::SinkhornBarycenterSolver{SinkhornGibbs})
    return OptimalTransport.check_convergence(
        solver.cache.a,
        solver.cache.u,
        solver.cache.Kv,
        solver.convergence_cache,
        solver.atol,
        solver.rtol,
    )
end

function solution(solver::SinkhornBarycenterSolver{SinkhornGibbs})
    cache = solver.cache
    return cache.u[:, 1] .* cache.Kv[:, 1]
end


text:
using Base.Test
using Compat
using SparseVectors

import SparseVectors: exact_equal

### Data

spv_x1 = SparseVector(8, [2, 5, 6], [1.25, -0.75, 3.5])
_x2 = SparseVector(8, [1, 2, 6, 7], [3.25, 4.0, -5.5, -6.0])
spv_x2 = view(_x2)

@test isa(spv_x1, SparseVector{Float64,Int})
@test isa(spv_x2, SparseVectorView{Float64,Int})

x1_full = zeros(length(spv_x1))
x1_full[nonzeroinds(spv_x1)] = nonzeros(spv_x1)

x2_full = zeros(length(spv_x2))
x2_full[nonzeroinds(spv_x2)] = nonzeros(spv_x2)

### Basic Properties

let x = spv_x1
    @test eltype(x) == Float64
    @test ndims(x) == 1
    @test length(x) == 8
    @test size(x) == (8,)
    @test size(x,1) == 8
    @test size(x,2) == 1
    @test !isempty(x)

    @test countnz(x) == 3
    @test nnz(x) == 3
    @test nonzeroinds(x) == [2, 5, 6]
    @test nonzeros(x) == [1.25, -0.75, 3.5]
end

let x = spv_x2
    @test eltype(x) == Float64
    @test ndims(x) == 1
    @test length(x) == 8
    @test size(x) == (8,)
    @test size(x,1) == 8
    @test size(x,2) == 1
    @test !isempty(x)

    @test countnz(x) == 4
    @test nnz(x) == 4
    @test nonzeroinds(x) == [1, 2, 6, 7]
    @test nonzeros(x) == [3.25, 4.0, -5.5, -6.0]
end

# full

for (x, xf) in [(spv_x1, x1_full), (spv_x2, x2_full)]
    @test isa(full(x), Vector{Float64})
    @test full(x) == xf
end

### Show

@test string(spv_x1) == "Sparse vector, length = 8, with 3 Float64 entries:\n" *
"  [2]  =  1.25\n" *
"  [5]  =  -0.75\n" *
"  [6]  =  3.5\n"

### Other Constructors

# construct empty sparse vector

@test exact_equal(sparsevector(Float64, 8), SparseVector(8, Int[], Float64[]))

# from list of indices and values

@test exact_equal(
    sparsevector(Int[], Float64[], 8),
    SparseVector(8, Int[], Float64[]))

@test exact_equal(
    sparsevector(Int[], Float64[]),
    SparseVector(0, Int[], Float64[]))

@test exact_equal(
    sparsevector([3, 3], [5.0, -5.0], 8),
    sparsevector(Float64, 8))

@test exact_equal(
    sparsevector([2, 3, 6], [12.0, 18.0, 25.0]),
    SparseVector(6, [2, 3, 6], [12.0, 18.0, 25.0]))

let x0 = SparseVector(8, [2, 3, 6], [12.0, 18.0, 25.0])
    @test exact_equal(
        sparsevector([2, 3, 6], [12.0, 18.0, 25.0], 8), x0)

    @test exact_equal(
        sparsevector([3, 6, 2], [18.0, 25.0, 12.0], 8), x0)

    @test exact_equal(
        sparsevector([2, 3, 4, 4, 6], [12.0, 18.0, 5.0, -5.0, 25.0], 8),
        x0)

    @test exact_equal(
        sparsevector([1, 1, 1, 2, 3, 3, 6], [2.0, 3.0, -5.0, 12.0, 10.0, 8.0, 25.0], 8),
        x0)

    @test exact_equal(
        sparsevector([2, 3, 6, 7, 7], [12.0, 18.0, 25.0, 5.0, -5.0], 8), x0)
end

# from dictionary

function my_intmap(x)
    a = Dict{Int,eltype(x)}()
    for i in nonzeroinds(x)
        a[i] = x[i]
    end
    return a
end

let x = spv_x1
    a = my_intmap(x)
    xc = sparsevector(a, 8)
    @test exact_equal(x, xc)

    xc = sparsevector(a)
    @test exact_equal(xc, SparseVector(6, [2, 5, 6], [1.25, -0.75, 3.5]))
end

# sprand & sprandn

let xr = sprand(1000, 0.3)
    @test isa(xr, SparseVector{Float64,Int})
    @test length(xr) == 1000
    @test all(nonzeros(xr) .>= 0.0)
end

let xr = sprand(1000, 0.3, Float32)
    @test isa(xr, SparseVector{Float32,Int})
    @test length(xr) == 1000
    @test all(nonzeros(xr) .>= 0.0)
end

let xr = sprandn(1000, 0.3)
    @test isa(xr, SparseVector{Float64,Int})
    @test length(xr) == 1000
    @test any(nonzeros(xr) .> 0.0) && any(nonzeros(xr) .< 0.0)
end

### Element access

# getindex

# single integer index
for (x, xf) in [(spv_x1, x1_full), (spv_x2, x2_full)]
    for i = 1:length(x)
        @test x[i] == xf[i]
    end
end

# range index
let x = spv_x2
    # range that contains no non-zeros
    @test exact_equal(x[3:2], sparsevector(Float64, 0))
    @test exact_equal(x[3:3], sparsevector(Float64, 1))
    @test exact_equal(x[3:5], sparsevector(Float64, 3))

    # range with non-zeros
    @test exact_equal(x[1:length(x)], x)
    @test exact_equal(x[1:5], SparseVector(5, [1,2], [3.25, 4.0]))
    @test exact_equal(x[2:6], SparseVector(5, [1,5], [4.0, -5.5]))
    @test exact_equal(x[2:8], SparseVector(7, [1,5,6], [4.0, -5.5, -6.0]))
end

# generic array index
let x = sprand(100, 0.5)
    I = rand(1:length(x), 20)
    @which x[I]
    r = x[I]
    @test isa(r, SparseVector{Float64,Int})
    @test all(nonzeros(r) .!= 0.0)
    @test full(r) == full(x)[I]
end

# setindex

let xc = sparsevector(Float64, 8)
    xc[3] = 2.0
    @test exact_equal(xc, SparseVector(8, [3], [2.0]))
end

let xc = copy(spv_x1)
    xc[5] = 2.0
    @test exact_equal(xc, SparseVector(8, [2, 5, 6], [1.25, 2.0, 3.5]))
end

let xc = copy(spv_x1)
    xc[3] = 4.0
    @test exact_equal(xc, SparseVector(8, [2, 3, 5, 6], [1.25, 4.0, -0.75, 3.5]))

    xc[1] = 6.0
    @test exact_equal(xc, SparseVector(8, [1, 2, 3, 5, 6], [6.0, 1.25, 4.0, -0.75, 3.5]))

    xc[8] = -1.5
    @test exact_equal(xc, SparseVector(8, [1, 2, 3, 5, 6, 8], [6.0, 1.25, 4.0, -0.75, 3.5, -1.5]))
end

let xc = copy(spv_x1)
    xc[5] = 0.0
    @test exact_equal(xc, SparseVector(8, [2, 6], [1.25, 3.5]))

    xc[6] = 0.0
    @test exact_equal(xc, SparseVector(8, [2], [1.25]))

    xc[2] = 0.0
    @test exact_equal(xc, SparseVector(8, Int[], Float64[]))
end

### Array manipulation

# copy

let x = spv_x1
    xc = copy(x)
    @test isa(xc, SparseVector{Float64,Int})
    @test !is(x.nzind, xc.nzval)
    @test !is(x.nzval, xc.nzval)
    @test exact_equal(x, xc)
end

let a = SparseVector(8, [2, 5, 6], Int32[12, 35, 72])
    # reinterpret
    au = reinterpret(UInt32, a)
    @test isa(au, SparseVector{UInt32,Int})
    @test exact_equal(au, SparseVector(8, [2, 5, 6], UInt32[12, 35, 72]))

    # float
    af = float(a)
    @test isa(af, SparseVector{Float64,Int})
    @test exact_equal(af, SparseVector(8, [2, 5, 6], [12., 35., 72.]))

    # complex
    acp = complex(af)
    @test isa(acp, SparseVector{Complex128,Int})
    @test exact_equal(acp, SparseVector(8, [2, 5, 6], complex([12., 35., 72.])))
end

### Type conversion

let x = convert(SparseVector, sparse([2, 5, 6], [1, 1, 1], [1.25, -0.75, 3.5], 8, 1))
    @test isa(x, SparseVector{Float64,Int})
    @test exact_equal(x, spv_x1)
end

let x = spv_x1, xf = x1_full
    xc = convert(SparseVector, xf)
    @test isa(xc, SparseVector{Float64,Int})
    @test exact_equal(xc, x)

    xc = convert(SparseVector{Float32,Int}, x)
    xf32 = SparseVector(8, [2, 5, 6], [1.25f0, -0.75f0, 3.5f0])
    @test isa(xc, SparseVector{Float32,Int})
    @test exact_equal(xc, xf32)

    xc = convert(SparseVector{Float32}, x)
    @test isa(xc, SparseVector{Float32,Int})
    @test exact_equal(xc, xf32)

    xm = convert(SparseMatrixCSC, x)
    @test isa(xm, SparseMatrixCSC{Float64,Int})
    @test full(xm) == reshape(xf, 8, 1)

    xm = convert(SparseMatrixCSC{Float32}, x)
    @test isa(xm, SparseMatrixCSC{Float32,Int})
    @test full(xm) == reshape(convert(Vector{Float32}, xf), 8, 1)
end

### Concatenation

let m = 80, n = 100
    A = Array(SparseVector{Float64,Int}, n)
    tnnz = 0
    for i = 1:length(A)
        A[i] = sprand(m, 0.3)
        tnnz += nnz(A[i])
    end

    H = hcat(A...)
    @test isa(H, SparseMatrixCSC{Float64,Int})
    @test size(H) == (m, n)
    @test nnz(H) == tnnz
    Hr = zeros(m, n)
    for j = 1:n
        Hr[:,j] = full(A[j])
    end
    @test full(H) == Hr

    V = vcat(A...)
    @test isa(V, SparseVector{Float64,Int})
    @test length(V) == m * n
    Vr = vec(Hr)
    @test full(V) == Vr
end


text:
#! /usr/bin/julia

# Rosetta Code, Loops/Do-while

i = 0
while true
    println(i)
    i += 1
    if i%6 == 0
        break
    end
end

max_stars_repo_name: MichaeLeroy/rosetta-code

max_issues_repo_name: MichaeLeroy/rosetta-code

max_forks_repo_name: MichaeLeroy/rosetta-code


text:
export SinkConfig, SINKVECTORS

# sinks — general interface

# A *sink* is a container that collects elements, expanding as necessary. When the
# collection is finished, sinks are *finalized* into *columns*. Details are governed by sink
# configuration objects (see [`SinkConfig`](@ref).
#
# Interface for sinks
#
# 1. [`make_sink`](@ref) for creating a sink for a single element.
#
# 2. [`store_or_reallocate!`](@ref) for saving another element. Either the result is `≡` to
# the sink in the argument, or a new sink was reallocated (potentially changing type).
#
# 3. [`finalize_sink`](@ref) turns the sink into a *column*, which is an iterable with a
# length and element type, supports `iterate,` but is not necessarily optimized for random
# access.

"""
$(TYPEDEF)
"""
struct SinkConfig{useRLE, M}
    missingvalue::M
    function SinkConfig{useRLE}(missingvalue::M) where {useRLE, M}
        @argcheck useRLE isa Bool
        @argcheck issingletontype(M)
        new{useRLE, M}(missingvalue)
    end
end

use_rle(cfg::SinkConfig{useRLE}) where useRLE = useRLE

"""
$(SIGNATURES)

Make a sink configuration, using defaults.
"""
SinkConfig(; useRLE = true, missingvalue = missing) = SinkConfig{useRLE}(missingvalue)

"Default sink configuration."
const SINKCONFIG = SinkConfig{true}(missing)

"Sink configuration that collects to vectors."
const SINKVECTORS = SinkConfig{false}(missing)

####
#### Reference implementation for sinks: `Vector`
####

"""
$(SIGNATURES)

Create and return a sink using configuration `cfg` that stores elements of type `T`. When
`T` is unkown, use `Base.Bottom`.
"""
make_sink(cfg::SinkConfig{true,M}, ::Type{T}) where {M, T} = RLEVector{M}(Int8, T)

make_sink(cfg::SinkConfig{false,M}, ::Type{T}) where {M, T} = Vector{T}()

"""
$(SIGNATURES)

Either store `elt` in `sink` (in which case the returned value is `≡ sink`), or
allocate a new sink that can do this, copy the contents, save `elt` and return that (then
the returned value is `≢ sink`).
"""
function store!_or_reallocate(::SinkConfig, sink::Vector{T}, elt) where T
    if cancontain(T, elt)
        (push!(sink, elt); sink)
    else
        append1(sink, narrow(elt))
    end
end

"""
$(SIGNATURES)

Convert `sink` to a *column*.

`sink` may share structure with the result and is not supposed to be used for saving any
more elements.
"""
finalize_sink(::SinkConfig, sink::Vector) = sink

####
#### RepeatValue type
####

"""
RepeatValue(value, len)

Equivalent to a vector containing `len` instances of `value`. Used *internally*.
"""
struct RepeatValue{T} <: AbstractVector{T}
    value::T
    len::Int
end

Base.size(s::RepeatValue) = (s.len, )

Base.IndexStyle(::Type{<:RepeatValue}) = Base.IndexLinear()

function Base.getindex(s::RepeatValue, i::Integer)
    @boundscheck checkbounds(s, i)
    s.value
end

####
#### RLE compressed vector
####

"""
$(TYPEDEF)

An RLE encoded vector, using negative lengths for missing values. Use the
`RLEVector{S}(C, T)` constructor for creating an empty one.

When an elemenet in `counts` is positive, it encodes that many of the corresponding element
in `data`.

Negative `counts` encode values of type `S` (has to be a concrete singleton). In this case
there is no corresponding value in `data`, ie `data` may have *fewer elements* than
`counts`. Note that `0` values in count are reserved, and currently should not happen.

The flag `anyS::Bool` is `true` iff there are *any* values of type `S` in object.

An RLEVector is iterable.
"""
struct RLEVector{C,T,S,anyS}
    counts::Vector{C}
    data::Vector{T}
    function RLEVector{S,anyS}(counts::Vector{C}, data::Vector{T}
                               ) where {C <: Signed, T, S, anyS}
        @argcheck issingletontype(S) "$(S) is not a concrete singleton type."
        @argcheck anyS isa Bool
        @argcheck length(counts) ≥ length(data)
        new{eltype(counts), eltype(data), S, anyS}(counts, data)
    end
end

"""
$(SIGNATURES)

Create an empty RLEVector for `Union{T,S}`, with special-casing the singleton type `S`. RLE
counts are stored in type `C`.
"""
RLEVector{S}(C::Type{<:Signed}, ::Type{S}) where {S} =
    RLEVector{S, true}(Vector{C}(), Vector{Union{}}())

RLEVector{S}(C::Type{<:Signed}, ::Type{T}) where {S, T} =
    RLEVector{S, false}(Vector{C}(), Vector{T}())

RLEVector{S}(C::Type{<:Signed}, ::Type{Union{S,T}}) where {S, T} =
    RLEVector{S, true}(Vector{C}(), Vector{T}())

function store!_or_reallocate(::SinkConfig, sink::RLEVector{C,T,S,anyS}, elt) where {C,T,S,anyS}
    @unpack counts, data = sink
    if cancontain(T, elt)       # can accommodate elt, same sink
        if isempty(data)
            push!(data, elt)
            push!(counts, one(C))
        elseif data[end] == elt && 0 < counts[end] < typemax(C)
            counts[end] += one(C) # increment existing count
        else
            push!(counts, one(C)) # start new RLE run
            push!(data, elt)
        end
        sink
    else                        # can't accommodate elt, allocate new sink
        RLEVector{S,anyS}(append1(counts, one(C)), append1(data, narrow(elt)))
    end
end

function store!_or_reallocate(cfg::SinkConfig, sink::RLEVector{C, T, S, false}, elt::S
                              ) where {C, T, S}
    @unpack counts, data = sink
    # simply flip the flag anyS
    store!_or_reallocate(cfg, RLEVector{S, true}(counts, data), elt)
end

function store!_or_reallocate(::SinkConfig, sink::RLEVector{C,T,S,true}, elt::S) where {C,T,S}
    @unpack counts = sink
    if isempty(counts)                  # no elements yet
        push!(counts, -one(C))
    elseif 0 > counts[end] > typemin(C) # ongoing RLE run with S
        counts[end] -= one(C)
    else
        push!(counts, -one(C))  # start new RLE runx
    end
    sink
end

finalize_sink(::SinkConfig, rle::RLEVector) = rle

function Base.eltype(::Type{RLEVector{C,T,S,anyS}}) where {C,T,S,anyS}
    anyS ? Base.promote_typejoin(T,S) : T
end

Base.length(rle::RLEVector) = isempty(rle.counts) ? 0 : sum(abs ∘ Int, rle.counts)

function Base.iterate(rle::RLEVector{C,T,S},
                 (countsindex, dataindex, remaining) = (0, 0, zero(C))) where {C,T,S}
    @unpack counts, data = rle
    if remaining < 0
        (S(), (countsindex, dataindex, remaining + one(C)))
    elseif remaining > 0
        (data[dataindex], (countsindex, dataindex, remaining - one(C)))
    else
        countsindex += 1
        countsindex > length(counts) && return nothing
        remaining = counts[countsindex]
        if remaining > 0
            dataindex += 1
        end
        iterate(rle, (countsindex, dataindex, remaining))
    end
end

####
#### Collecting named tuples
####

"""
$(TYPEDEF)

Wrapper type to indicate that the length should not be checked.

!!! note
    The perfect footgun. Only use when the lengths are known and verified by construction.
"""
struct TrustLength
    len::Int
end

"""
$(SIGNATURES)

Collect results from `itr` into a sink (using config `cfg`), then finalize and return the
column.
"""
function collect_column(cfg::SinkConfig, itr)
    y = iterate(itr)
    y ≡ nothing && return nothing
    elt, state = y
    narrow_elt = narrow(elt)
    sink = make_sink(cfg, typeof(narrow_elt))
    collect_column!(store!_or_reallocate(cfg, sink, narrow_elt), cfg, itr, state)
end

function collect_column!(sink, cfg::SinkConfig, itr, state)
    while true
        y = iterate(itr, state)
        y ≡ nothing && return finalize_sink(cfg, sink)
        elt, state = y
        newsink = store!_or_reallocate(cfg, sink, elt)
        sink ≡ newsink || return collect_column!(newsink, cfg, itr, state)
    end
end

"""
$(SIGNATURES)

Empty sinks for a named tuple of elements, using a type.
"""
function empty_sinks(cfg, ::Type{NamedTuple{N,T}}) where {N,T}
    NamedTuple{N}(map(S -> make_sink(cfg, S), fieldtypes(T)))
end

"""
$(SIGNATURES)

Start sinks using row, using the default `known_types` when available.
"""
function start_sinks(cfg, row::T, known_types) where {T}
    rowtypes = merge_default_types(T, known_types)
    sinks = empty_sinks(cfg, rowtypes)
    store!_or_reallocate_row(cfg, sinks, row)
end

"""
$(SIGNATURES)

Finalize a (named) tuple of sinks.
"""
finalize_sinks(cfg, sinks::NamedTuple) = map(sink -> finalize_sink(cfg, sink), sinks)

"""
$(SIGNATURES)

Broadcast `store!_or_rellocate` for a compatible (named) tuple of `sinks` and `elts`. Return
the (potentially) new sinks.
"""
@inline store!_or_reallocate_row(cfg, sinks, elts) =
    map((sink, elt) -> store!_or_reallocate(cfg, sink, elt), sinks, elts)

"""
len, columns, ordering_rule = $(SIGNATURES)

Collect results from `itr`, which are supposed to be `NamedTuple`s with the same names, into
sinks (using config `cfg`), then finalize and return

1. the length,

2. the `NamedTuple` of the columns, and

3. the ordering rule (which is always `::TrustOrdering`, by construction).

Determine the names and types from the first named tuple, using `known_types` as the
narrowest types for the given columns.

# Special rules for empty iterators

When `itr` is empty, use a `known_types` will be used to create empty columns, and only the
`TryOrdering` rule will be narrowed to these. Other rule with more column names may cause an
error in the callee, which is intentional.
"""
function collect_columns(cfg::SinkConfig, itr, ordering_rule::OrderingRule{R},
                         known_types::Type{<:NamedTuple{N}} = NamedTuple{(), Tuple{}}
                         ) where {R,N}
    y = iterate(itr)
    if y ≡ nothing
        columns = finalize_sinks(cfg, empty_sinks(cfg, known_types))
        return 0, columns, mask_try_ordering(ordering_rule, N)
    end
    elts, state = y
    sinks = start_sinks(cfg, elts, known_types)
    collect_columns!(sinks, 1, cfg, itr, mask_try_ordering(ordering_rule, keys(elts)),
                     # :trust, we don't need the last element for comparison, hence the ()
                     R ≡ :trust ? () : elts, state)
end

function collect_columns!(sinks::NamedTuple, len::Int, cfg::SinkConfig, itr,
                          ordering_rule::OrderingRule{R}, lastelts, state) where R
    @unpack ordering = ordering_rule
    while true
        y = iterate(itr, state)
        if y ≡ nothing
            return TrustLength(len), finalize_sinks(cfg, sinks), TrustOrdering(ordering_rule)
        end
        elts, state = y
        newsinks = store!_or_reallocate_row(cfg, sinks, elts)
        len += 1
        if R ≢ :trust
            if cmp_ordering(ordering, lastelts, elts) > 0
                if R ≡ :verify
                    error("Sorting $(sorting) violated: $(lastelts) ≰ $(elts).")
                else # R ≡ :try
                    new_ordering = retained_ordering(ordering, lastelts, elts)
                    return collect_columns!(newsinks, len, cfg, itr,
                                            OrderingRule{R}(new_ordering), elts, state)
                end
            end
            lastelts = elts
        end
        sinks ≡ newsinks || return collect_columns!(newsinks, len, cfg, itr, ordering_rule,
                                                    lastelts, state)
    end
end


text:
const EXPRStack = Tuple{EXPR,Int,Int}

function find_enclosing_expr(cst, offset, pos = 0, stack = EXPRStack[])
    pos1 = pos
    for (i, a) in enumerate(cst)
        if pos < first(offset) < pos + a.span && pos < last(offset) < pos + a.span
            push!(stack, (cst, pos1, i))
            return find_enclosing_expr(a, offset, pos, stack)
        else
            pos += a.fullspan
        end
    end
    push!(stack, (cst, pos1, 0))
    return stack
end

containing_expr(stack) = first(last(stack))
parent_expr(stack) = first(stack[end - 1])
insert_size(inserttext, insertrange) = sizeof(inserttext) - max(last(insertrange) - first(insertrange), 0)
enclosing_expr_range(stack, insertsize) = last(stack)[2] .+ (1:first(last(stack)).fullspan + insertsize)

function reparse(edittedtext::String, inserttext::String, insertrange::UnitRange{Int}, oldcst)
    #need to handle empty replacement case, i.e. inserttext = ""
    stack = find_enclosing_expr(oldcst, insertrange)
    insertsize = insert_size(inserttext, insertrange)
    reparsed = reparse(stack, edittedtext, insertsize, oldcst)
    return reparsed, oldcst
end

function replace_args!(replacement_args, stack)
    pexpr, ppos, pi = stack[end - 1]
    oldspan, oldfullspan = pexpr.args[pi].span, pexpr.args[pi].fullspan
    if length(replacement_args) > 1
        deleteat!(pexpr.args, pi)
        for i = 0:length(replacement_args) - 1
            insert!(pexpr.args, pi + i, replacement_args[1 + i])
        end
    else
        pexpr.args[pi] = replacement_args[1]
        dfullspan = replacement_args[1].fullspan - oldfullspan
        @info "Replacing $(typeof(pexpr)) at $(pi)"
        fix_stack_span(stack, dfullspan, replacement_args[1].span - replacement_args[1].fullspan)
    end
end

function reparse(stack::Array{EXPRStack}, edittedtext::String, insertsize::Int, oldcst)
    # Assume no existing error in CST
    # need to update (full)spans
    if length(stack) == 1
        return false
    else
        if parent_expr(stack) isa EXPR{FileH}
            replacement_args = parse(edittedtext[enclosing_expr_range(stack, insertsize)], true).args
        elseif parent_expr(stack) isa EXPR{Block} && !(length(stack) > 2 && stack[end - 2] isa EXPR{If})
            replacement_args = let
                ps = ParseState(edittedtext[enclosing_expr_range(stack, insertsize)])
                newblockargs = Any[]
                CSTParser.parse_block(ps, newblockargs)
            end
        else
            pop!(stack)
            return reparse(stack, edittedtext, insertsize, oldcst)
        end
        replace_args!(replacement_args, stack)
        return true
    end
    return false
end

function fixlastchild(x, pi, dspan)
    nx = length(x)
    if pi == nx
        x[pi] = x.fullspan - dspan
        return true
    elseif x[nx].fullspan == 0
        for i = nx - 1:-1:1
            if x[i].fullspan > 0
                x[pi] = x.fullspan - dspan
                return true
            end
        end
    end
    return false
end

function fix_stack_span(stack::Array{EXPRStack}, dfullspan, dspan)
    islast = true
    for i = length(stack) - 1:-1:1
        stack[i][1].fullspan += dfullspan
        if islast && stack[i][3] == length(stack[i][1])
            stack[i][1].span = stack[i][1].fullspan + dspan
        else
            stack[i][1].span += dfullspan
            islast = false
        end
    end
end

function reparse_test(text, insertrange, inserttext)
    cst = parse(text, true)
    cst0 = deepcopy(cst)
    edittedtext = edit_string(text, insertrange, inserttext)
    reparsed, reparsed_cst = reparse(edittedtext, inserttext, insertrange, cst)
    new_cst = CSTParser.parse(edittedtext, true)
    return reparsed, isequiv(new_cst, reparsed_cst), text, edittedtext, cst0, new_cst, reparsed_cst
end

function edit_string(text, insertrange, inserttext)
    if first(insertrange) == last(insertrange) == 0
        text = string(inserttext, text)
    elseif first(insertrange) == 0 && last(insertrange) == sizeof(text)
        text = inserttext
    elseif first(insertrange) == 0
        text = string(inserttext, text[nextind(text, last(insertrange)):end])
    elseif first(insertrange) == last(insertrange) == sizeof(text)
        text = string(text, inserttext)
    elseif last(insertrange) == sizeof(text)
        text = string(text[1:first(insertrange)], inserttext)
    else
        text = string(text[1:first(insertrange)], inserttext, text[nextind(text, last(insertrange)):end])
    end
end

spanequiv(a::EXPR, b::EXPR) = a.span == b.span && a.fullspan == b.fullspan

isequiv(a, b; span = true) = false

function isequiv(a::EXPR, b::EXPR; span = true)
    t = typof(a) === typof(b)
    typeof(a.args) != typeof(b.args) && return false
    if a.args isa Vector
        length(a.args) != length(b.args) && return false
        for i = 1:length(a.args)
            t = t && isequiv(a.args[i], b.args[i], span = span)
            t || return false
        end
    end
    return (!span || spanequiv(a, b))
end


text:
# ncg.jl

using MIRT: ncg

using LinearAlgebra: I, norm, opnorm
using Plots: Plot
using Random: seed!
using Plots: plot, plot!, scatter!, gui
using LaTeXStrings
using Test: @test, @test_throws, @inferred
using LinearMapsAA

odim = (10,4); idim = (5,2)
M = prod(odim)
N = prod(idim)
seed!(0)
A = randn(M,N)
a2 = opnorm(A)^2
y = randn(M)
reg = 0.1 * a2
xh = (A'A + reg*I) \ A'y
if true # stress test with AO
	A = LinearMapAA(A ; odim=odim, idim=idim)
	xh = reshape(xh, idim)
	y = reshape(y, odim)
end
cost = (x) -> 1/2 * norm(A * x - y)^2 + reg/2 * norm(x)^2
fun = (x,iter) -> (cost(x) - cost(xh), norm(x - xh) / norm(xh), time())
grad1 = (x) -> A' * (A * x - y) + reg * x
curv1 = (x) -> a2 + reg

B = [A, I] # matrix blocks
gradf = [u -> u - y, v -> reg * v] # f functions gradients
curvf = [v -> 1, v -> reg]

#	todo-i: these next two fail:
#	@inferred ncg([ones(3,3)], [x -> x], [v -> 1], zeros(Float32, 3), niter=2)
#	@inferred ncg(x -> x, v -> 1, zeros(Float32, 3), niter=2)
#	error("todo")

niter = 40
x0 = zeros(size(xh))
x1, out1 = ncg(   grad1, curv1, x0, niter=niter, fun=fun)
x2, out2 = ncg(B, gradf, curvf, x0, niter=niter, fun=fun)
@test x1 ≈ xh
@test x2 ≈ xh

@test_throws ArgumentError ncg(grad1, curv1, x0, betahow=:test)

lf = x -> log10(max(x,1e-17))
costk = out -> lf.([out[k][1] for k=1:niter+1])
errk = out -> lf.([out[k][2] for k=1:niter+1])
allk = out -> (costk(out), errk(out))
cost1, err1 = allk(out1)
cost2, err2 = allk(out2)

if isinteractive()
	k = 0:niter
	plot(xlabel="k", ylabel=L"\log(\Psi(x_k) - \Psi(x_*))")
	scatter!(k, cost1, color=:blue, label="cost1")
	scatter!(k, cost2, color=:red, marker=:x, label="cost2")
	p1 = plot!()

	plot(xlabel="k", ylabel=L"\log(\|x_k - x_*\|/\|x_*\|)")
	scatter!(k, err1, color=:blue, label="NRMSD1")
	scatter!(k, err2, color=:red, marker=:x, label="NRMSD2")
	p2 = plot!()
	plot(p1, p2)
	gui()
end

#=
todo
f = x -> x + one(eltype(x))
f = x -> eltype(x).(x)

@inferred f(3)
@inferred f(3.0f0)

function g(x0::AbstractVector{<:Number} ; f::Function = x -> 2*x)
	x = copy(x0)
#	return x .+ ones(eltype(x0), size(x0))
#	return 2*x
	return f(x)
end

g(1:3)
@inferred g(1:3)
@code_warntype g(1:3)

#	the following test returns nothing!?
#	@code_warntype ncg([x -> 2x], [v -> 1], zeros(Float32, 3), niter=2)
=#