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

mutable struct Process <: AbstractPipe
    cmd::Cmd
    handle::Ptr{Cvoid}
    in::IO
    out::IO
    err::IO
    exitcode::Int64
    termsignal::Int32
    exitnotify::ThreadSynchronizer
    function Process(cmd::Cmd, handle::Ptr{Cvoid})
        this = new(cmd, handle, devnull, devnull, devnull,
                   typemin(fieldtype(Process, :exitcode)),
                   typemin(fieldtype(Process, :termsignal)),
                   ThreadSynchronizer())
        finalizer(uvfinalize, this)
        return this
    end
end
pipe_reader(p::Process) = p.out
pipe_writer(p::Process) = p.in

# Represents a whole pipeline of any number of related processes
# so the entire pipeline can be treated as one entity
mutable struct ProcessChain <: AbstractPipe
    processes::Vector{Process}
    in::IO
    out::IO
    err::IO
    function ProcessChain()
        return new(Process[], devnull, devnull, devnull)
    end
end
pipe_reader(p::ProcessChain) = p.out
pipe_writer(p::ProcessChain) = p.in

# release ownership of the libuv handle
function uvfinalize(proc::Process)
    if proc.handle != C_NULL
        disassociate_julia_struct(proc.handle)
        ccall(:jl_close_uv, Cvoid, (Ptr{Cvoid},), proc.handle)
        proc.handle = C_NULL
    end
    nothing
end

# called when the process dies
function uv_return_spawn(p::Ptr{Cvoid}, exit_status::Int64, termsignal::Int32)
    data = ccall(:jl_uv_process_data, Ptr{Cvoid}, (Ptr{Cvoid},), p)
    data == C_NULL && return
    proc = unsafe_pointer_to_objref(data)::Process
    proc.exitcode = exit_status
    proc.termsignal = termsignal
    ccall(:jl_close_uv, Cvoid, (Ptr{Cvoid},), proc.handle)
    proc.handle = C_NULL
    lock(proc.exitnotify)
    try
        notify(proc.exitnotify)
    finally
        unlock(proc.exitnotify)
    end
    nothing
end

# called when the libuv handle is destroyed
function _uv_hook_close(proc::Process)
    proc.handle = C_NULL
    nothing
end

const SpawnIOs = Vector{Any} # convenience name for readability

function as_cpumask(cpus::Vector{UInt16})
    n = max(Int(maximum(cpus)), Int(ccall(:uv_cpumask_size, Cint, ())))
    cpumask = zeros(Bool, n)
    for i in cpus
        cpumask[i] = true
    end
    return cpumask
end

# handle marshalling of `Cmd` arguments from Julia to C
@noinline function _spawn_primitive(file, cmd::Cmd, stdio::SpawnIOs)
    loop = eventloop()
    cpumask = cmd.cpus
    cpumask === nothing || (cpumask = as_cpumask(cmd.cpus))
    GC.@preserve stdio begin
        iohandles = Tuple{Cint, UInt}[ # assuming little-endian layout
            let h = rawhandle(io)
                h === C_NULL     ? (0x00, UInt(0)) :
                h isa OS_HANDLE  ? (0x02, UInt(cconvert(@static(Sys.iswindows() ? Ptr{Cvoid} : Cint), h))) :
                h isa Ptr{Cvoid} ? (0x04, UInt(h)) :
                error("invalid spawn handle $h from $io")
            end
            for io in stdio]
        handle = Libc.malloc(_sizeof_uv_process)
        disassociate_julia_struct(handle) # ensure that data field is set to C_NULL
        (; exec, flags, env, dir) = cmd
        err = ccall(:jl_spawn, Int32,
                  (Cstring, Ptr{Cstring}, Ptr{Cvoid}, Ptr{Cvoid},
                   Ptr{Tuple{Cint, UInt}}, Int,
                   UInt32, Ptr{Cstring}, Cstring, Ptr{Bool}, Csize_t, Ptr{Cvoid}),
            file, exec, loop, handle,
            iohandles, length(iohandles),
            flags,
            env === nothing ? C_NULL : env,
            isempty(dir) ? C_NULL : dir,
            cpumask === nothing ? C_NULL : cpumask,
            cpumask === nothing ? 0 : length(cpumask),
            @cfunction(uv_return_spawn, Cvoid, (Ptr{Cvoid}, Int64, Int32)))
    end
    if err != 0
        ccall(:jl_forceclose_uv, Cvoid, (Ptr{Cvoid},), handle) # will call free on handle eventually
        throw(_UVError("could not spawn " * repr(cmd), err))
    end
    pp = Process(cmd, handle)
    associate_julia_struct(handle, pp)
    return pp
end

_spawn(cmds::AbstractCmd) = _spawn(cmds, Any[])

# optimization: we can spawn `Cmd` directly without allocating the ProcessChain
function _spawn(cmd::Cmd, stdios::SpawnIOs)
    isempty(cmd.exec) && throw(ArgumentError("cannot spawn empty command"))
    pp = setup_stdios(stdios) do stdios
        return _spawn_primitive(cmd.exec[1], cmd, stdios)
    end
    return pp
end

# assume that having a ProcessChain means that the stdio are setup
function _spawn(cmds::AbstractCmd, stdios::SpawnIOs)
    pp = setup_stdios(stdios) do stdios
        return _spawn(cmds, stdios, ProcessChain())
    end
    return pp
end

# helper function for making a copy of a SpawnIOs, with replacement
function _stdio_copy(stdios::SpawnIOs, fd::Int, @nospecialize replace)
    nio = max(fd, length(stdios))
    new = SpawnIOs(undef, nio)
    copyto!(fill!(new, devnull), stdios)
    new[fd] = replace
    return new
end

function _spawn(redirect::CmdRedirect, stdios::SpawnIOs, args...)
    fdnum = redirect.stream_no + 1
    io, close_io = setup_stdio(redirect.handle, redirect.readable)
    try
        stdios = _stdio_copy(stdios, fdnum, io)
        return _spawn(redirect.cmd, stdios, args...)
    finally
        close_io && close_stdio(io)
    end
end

function _spawn(cmds::OrCmds, stdios::SpawnIOs, chain::ProcessChain)
    in_pipe, out_pipe = link_pipe(false, false)
    try
        stdios_left = _stdio_copy(stdios, 2, out_pipe)
        _spawn(cmds.a, stdios_left, chain)
        stdios_right = _stdio_copy(stdios, 1, in_pipe)
        _spawn(cmds.b, stdios_right, chain)
    finally
        close_pipe_sync(out_pipe)
        close_pipe_sync(in_pipe)
    end
    return chain
end

function _spawn(cmds::ErrOrCmds, stdios::SpawnIOs, chain::ProcessChain)
    in_pipe, out_pipe = link_pipe(false, false)
    try
        stdios_left = _stdio_copy(stdios, 3, out_pipe)
        _spawn(cmds.a, stdios_left, chain)
        stdios_right = _stdio_copy(stdios, 1, in_pipe)
        _spawn(cmds.b, stdios_right, chain)
    finally
        close_pipe_sync(out_pipe)
        close_pipe_sync(in_pipe)
    end
    return chain
end

function _spawn(cmds::AndCmds, stdios::SpawnIOs, chain::ProcessChain)
    _spawn(cmds.a, stdios, chain)
    _spawn(cmds.b, stdios, chain)
    return chain
end

function _spawn(cmd::Cmd, stdios::SpawnIOs, chain::ProcessChain)
    isempty(cmd.exec) && throw(ArgumentError("cannot spawn empty command"))
    pp = _spawn_primitive(cmd.exec[1], cmd, stdios)
    push!(chain.processes, pp)
    return chain
end

# open the child end of each element of `stdios`, and initialize the parent end
function setup_stdios(f, stdios::SpawnIOs)
    nstdio = length(stdios)
    open_io = Vector{Any}(undef, nstdio)
    close_io = falses(nstdio)
    try
        for i in 1:nstdio
            open_io[i], close_io[i] = setup_stdio(stdios[i], i == 1)
        end
        pp = f(open_io)
        return pp
    finally
        for i in 1:nstdio
            close_io[i] && close_stdio(open_io[i])
        end
    end
end

function setup_stdio(stdio::PipeEndpoint, child_readable::Bool)
    if stdio.status == StatusInit
        # if the PipeEndpoint isn't open, set it to the parent end
        # and pass the other end to the child
        rd, wr = link_pipe(!child_readable, child_readable)
        try
            open_pipe!(stdio, child_readable ? wr : rd)
        catch
            close_pipe_sync(rd)
            close_pipe_sync(wr)
            rethrow()
        end
        child = child_readable ? rd : wr
        return (child, true)
    end
    # if it's already open, assume that it's already the child end
    # (since we can't do anything else)
    return (stdio, false)
end

function setup_stdio(stdio::Pipe, child_readable::Bool)
    if stdio.in.status == StatusInit && stdio.out.status == StatusInit
        link_pipe!(stdio)
    end
    io = child_readable ? stdio.out : stdio.in
    return (io, false)
end

setup_stdio(stdio::AbstractPipe, readable::Bool) =
    setup_stdio(readable ? pipe_reader(stdio) : pipe_writer(stdio), readable)

function setup_stdio(stdio::IOStream, child_readable::Bool)
    io = RawFD(fd(stdio))
    return (io, false)
end

function setup_stdio(stdio::FileRedirect, child_readable::Bool)
    if child_readable
        attr = JL_O_RDONLY
        perm = zero(S_IRUSR)
    else
        attr = JL_O_WRONLY | JL_O_CREAT
        attr |= stdio.append ? JL_O_APPEND : JL_O_TRUNC
        perm = S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH
    end
    io = Filesystem.open(stdio.filename, attr, perm)
    return (io, true)
end

# incrementally move data between an arbitrary IO and a system Pipe,
# including copying the EOF (shutdown) when finished
# TODO: probably more efficient (when valid) to use `stdio` directly as the
#       PipeEndpoint buffer field in some cases
function setup_stdio(stdio::IO, child_readable::Bool)
    parent = PipeEndpoint()
    rd, wr = link_pipe(!child_readable, child_readable)
    try
        open_pipe!(parent, child_readable ? wr : rd)
    catch
        close_pipe_sync(rd)
        close_pipe_sync(wr)
        rethrow()
    end
    child = child_readable ? rd : wr
    try
        let in = (child_readable ? parent : stdio),
            out = (child_readable ? stdio : parent)
            @async try
                write(in, out)
            catch ex
                @warn "Process I/O error" exception=(ex, catch_backtrace())
            finally
                close(parent)
                child_readable || closewrite(stdio)
            end
        end
    catch
        close_pipe_sync(child)
        rethrow()
    end
    return (child, true)
end

close_stdio(stdio::OS_HANDLE) = close_pipe_sync(stdio)
close_stdio(stdio) = close(stdio)

# INTERNAL
# pad out stdio to have at least three elements,
# passing either `devnull` or the corresponding `stdio`
# A Redirectable can be any of:
#   - A system IO handle, to be passed to the child
#   - An uninitialized pipe, to be created
#   - devnull (to pass /dev/null for 0-2, or to leave undefined for fd > 2)
#   - An Filesystem.File or IOStream object to redirect the output to
#   - A FileRedirect, containing a string specifying a filename to be opened for the child

spawn_opts_swallow(stdios::StdIOSet) = Any[stdios...]
spawn_opts_inherit(stdios::StdIOSet) = Any[stdios...]
spawn_opts_swallow(in::Redirectable=devnull, out::Redirectable=devnull, err::Redirectable=devnull) =
    Any[in, out, err]
# pass original descriptors to child processes by default, because we might
# have already exhausted and closed the libuv object for our standard streams.
# ref issue #8529
spawn_opts_inherit(in::Redirectable=RawFD(0), out::Redirectable=RawFD(1), err::Redirectable=RawFD(2)) =
    Any[in, out, err]

function eachline(cmd::AbstractCmd; keep::Bool=false)
    out = PipeEndpoint()
    processes = _spawn(cmd, Any[devnull, out, stderr])
    # if the user consumes all the data, also check process exit status for success
    ondone = () -> (success(processes) || pipeline_error(processes); nothing)
    return EachLine(out, keep=keep, ondone=ondone)::EachLine
end

"""
    open(command, mode::AbstractString, stdio=devnull)

Run `command` asynchronously. Like `open(command, stdio; read, write)` except specifying
the read and write flags via a mode string instead of keyword arguments.
Possible mode strings are:

| Mode | Description | Keywords                         |
|:-----|:------------|:---------------------------------|
| `r`  | read        | none                             |
| `w`  | write       | `write = true`                   |
| `r+` | read, write | `read = true, write = true`      |
| `w+` | read, write | `read = true, write = true`      |
"""
function open(cmds::AbstractCmd, mode::AbstractString, stdio::Redirectable=devnull)
    if mode == "r+" || mode == "w+"
        return open(cmds, stdio, read = true, write = true)
    elseif mode == "r"
        return open(cmds, stdio)
    elseif mode == "w"
        return open(cmds, stdio, write = true)
    else
        throw(ArgumentError("mode must be \"r\", \"w\", \"r+\", or \"w+\", not $(repr(mode))"))
    end
end

# return a Process object to read-to/write-from the pipeline
"""
    open(command, stdio=devnull; write::Bool = false, read::Bool = !write)

Start running `command` asynchronously, and return a `process::IO` object.  If `read` is
true, then reads from the process come from the process's standard output and `stdio` optionally
specifies the process's standard input stream.  If `write` is true, then writes go to
the process's standard input and `stdio` optionally specifies the process's standard output
stream.
The process's standard error stream is connected to the current global `stderr`.
"""
function open(cmds::AbstractCmd, stdio::Redirectable=devnull; write::Bool=false, read::Bool=!write)
    if read && write
        stdio === devnull || throw(ArgumentError("no stream can be specified for `stdio` in read-write mode"))
        in = PipeEndpoint()
        out = PipeEndpoint()
        processes = _spawn(cmds, Any[in, out, stderr])
        processes.in = in
        processes.out = out
    elseif read
        out = PipeEndpoint()
        processes = _spawn(cmds, Any[stdio, out, stderr])
        processes.out = out
    elseif write
        in = PipeEndpoint()
        processes = _spawn(cmds, Any[in, stdio, stderr])
        processes.in = in
    else
        stdio === devnull || throw(ArgumentError("no stream can be specified for `stdio` in no-access mode"))
        processes = _spawn(cmds, Any[devnull, devnull, stderr])
    end
    return processes
end

"""
    open(f::Function, command, args...; kwargs...)

Similar to `open(command, args...; kwargs...)`, but calls `f(stream)` on the
resulting process stream, then closes the input stream and waits for the process
to complete. Return the value returned by `f` on success. Throw an error if the
process failed, or if the process attempts to print anything to stdout.
"""
function open(f::Function, cmds::AbstractCmd, args...; kwargs...)
    P = open(cmds, args...; kwargs...)
    ret = try
        f(P)
    catch
        kill(P)
        close(P)
        rethrow()
    end
    close(P.in)
    if !eof(P.out)
        close(P.out)
        throw(_UVError("open(do)", UV_EPIPE))
    end
    success(P) || pipeline_error(P)
    return ret
end

"""
    read(command::Cmd)

Run `command` and return the resulting output as an array of bytes.
"""
function read(cmd::AbstractCmd)
    procs = open(cmd, "r", devnull)
    bytes = read(procs.out)
    success(procs) || pipeline_error(procs)
    return bytes
end

"""
    read(command::Cmd, String)

Run `command` and return the resulting output as a `String`.
"""
read(cmd::AbstractCmd, ::Type{String}) = String(read(cmd))::String

"""
    run(command, args...; wait::Bool = true)

Run a command object, constructed with backticks (see the [Running External Programs](@ref)
section in the manual). Throws an error if anything goes wrong, including the process
exiting with a non-zero status (when `wait` is true).

If `wait` is false, the process runs asynchronously. You can later wait for it and check
its exit status by calling `success` on the returned process object.

When `wait` is false, the process' I/O streams are directed to `devnull`.
When `wait` is true, I/O streams are shared with the parent process.
Use [`pipeline`](@ref) to control I/O redirection.
"""
function run(cmds::AbstractCmd, args...; wait::Bool = true)
    if wait
        ps = _spawn(cmds, spawn_opts_inherit(args...))
        success(ps) || pipeline_error(ps)
    else
        stdios = spawn_opts_swallow(args...)
        ps = _spawn(cmds, stdios)
        # for each stdio input argument, guess whether the user
        # passed a `stdio` placeholder object as input, and thus
        # might be able to use the return AbstractProcess as an IO object
        # (this really only applies to PipeEndpoint, Pipe, TCPSocket, or an AbstractPipe wrapping one of those)
        if length(stdios) > 0
            in = stdios[1]
            isa(in, IO) && (ps.in = in)
            if length(stdios) > 1
                out = stdios[2]
                isa(out, IO) && (ps.out = out)
                if length(stdios) > 2
                    err = stdios[3]
                    isa(err, IO) && (ps.err = err)
                end
            end
        end
    end
    return ps
end

# some common signal numbers that are usually available on all platforms
# and might be useful as arguments to `kill` or testing against `Process.termsignal`
const SIGHUP  = 1
const SIGINT  = 2
const SIGQUIT = 3 # !windows
const SIGKILL = 9
const SIGPIPE = 13 # !windows
const SIGTERM = 15

function test_success(proc::Process)
    @assert process_exited(proc)
    if proc.exitcode < 0
        #TODO: this codepath is not currently tested
        throw(_UVError("could not start process " * repr(proc.cmd), proc.exitcode))
    end
    return proc.exitcode == 0 && proc.termsignal == 0
end

function success(x::Process)
    wait(x)
    return test_success(x)
end
success(procs::Vector{Process}) = mapreduce(success, &, procs)
success(procs::ProcessChain) = success(procs.processes)

"""
    success(command)

Run a command object, constructed with backticks (see the [Running External Programs](@ref)
section in the manual), and tell whether it was successful (exited with a code of 0).
An exception is raised if the process cannot be started.
"""
success(cmd::AbstractCmd) = success(_spawn(cmd))

"""
    ProcessFailedException

Indicates problematic exit status of a process.
When running commands or pipelines, this is thrown to indicate
a nonzero exit code was returned (i.e. that the invoked process failed).
"""
struct ProcessFailedException <: Exception
    procs::Vector{Process}
end
ProcessFailedException(proc::Process) = ProcessFailedException([proc])

function showerror(io::IO, err::ProcessFailedException)
    if length(err.procs) == 1
        proc = err.procs[1]
        println(io, "failed process: ", proc, " [", proc.exitcode, "]")
    else
        println(io, "failed processes:")
        for proc in err.procs
            println(io, "  ", proc, " [", proc.exitcode, "]")
        end
    end
end

function pipeline_error(proc::Process)
    if !proc.cmd.ignorestatus
        throw(ProcessFailedException(proc))
    end
    nothing
end

function pipeline_error(procs::ProcessChain)
    failed = Process[]
    for p = procs.processes
        if !test_success(p) && !p.cmd.ignorestatus
            push!(failed, p)
        end
    end
    isempty(failed) && return nothing
    throw(ProcessFailedException(failed))
end

"""
    kill(p::Process, signum=Base.SIGTERM)

Send a signal to a process. The default is to terminate the process.
Returns successfully if the process has already exited, but throws an
error if killing the process failed for other reasons (e.g. insufficient
permissions).
"""
function kill(p::Process, signum::Integer=SIGTERM)
    iolock_begin()
    if process_running(p)
        @assert p.handle != C_NULL
        err = ccall(:uv_process_kill, Int32, (Ptr{Cvoid}, Int32), p.handle, signum)
        if err != 0 && err != UV_ESRCH
            throw(_UVError("kill", err))
        end
    end
    iolock_end()
    nothing
end
kill(ps::Vector{Process}, signum::Integer=SIGTERM) = for p in ps; kill(p, signum); end
kill(ps::ProcessChain, signum::Integer=SIGTERM) = kill(ps.processes, signum)

"""
    getpid(process) -> Int32

Get the child process ID, if it still exists.

!!! compat "Julia 1.1"
    This function requires at least Julia 1.1.
"""
function Libc.getpid(p::Process)
    # TODO: due to threading, this method is no longer synchronized with the user application
    iolock_begin()
    ppid = Int32(0)
    if p.handle != C_NULL
        ppid = ccall(:jl_uv_process_pid, Int32, (Ptr{Cvoid},), p.handle)
    end
    iolock_end()
    ppid <= 0 && throw(_UVError("getpid", UV_ESRCH))
    return ppid
end

## process status ##

"""
    process_running(p::Process)

Determine whether a process is currently running.
"""
process_running(s::Process) = s.handle != C_NULL
process_running(s::Vector{Process}) = any(process_running, s)
process_running(s::ProcessChain) = process_running(s.processes)

"""
    process_exited(p::Process)

Determine whether a process has exited.
"""
process_exited(s::Process) = !process_running(s)
process_exited(s::Vector{Process}) = all(process_exited, s)
process_exited(s::ProcessChain) = process_exited(s.processes)

process_signaled(s::Process) = (s.termsignal > 0)

function process_status(s::Process)
    return process_running(s) ? "ProcessRunning" :
           process_signaled(s) ? "ProcessSignaled(" * string(s.termsignal) * ")" :
           process_exited(s) ? "ProcessExited(" * string(s.exitcode) * ")" :
           error("process status error")
end

function wait(x::Process)
    process_exited(x) && return
    iolock_begin()
    if !process_exited(x)
        preserve_handle(x)
        lock(x.exitnotify)
        iolock_end()
        try
            wait(x.exitnotify)
        finally
            unlock(x.exitnotify)
            unpreserve_handle(x)
        end
    else
        iolock_end()
    end
    nothing
end

wait(x::ProcessChain) = foreach(wait, x.processes)

show(io::IO, p::Process) = print(io, "Process(", p.cmd, ", ", process_status(p), ")")

# allow the elements of the Cmd to be accessed as an array or iterator
for f in (:length, :firstindex, :lastindex, :keys, :first, :last, :iterate)
    @eval $f(cmd::Cmd) = $f(cmd.exec)
end
eltype(::Type{Cmd}) = eltype(fieldtype(Cmd, :exec))
for f in (:iterate, :getindex)
    @eval $f(cmd::Cmd, i) = $f(cmd.exec, i)
end

max_stars_repo_name: BSnelling/julia

max_issues_repo_name: BSnelling/julia

max_forks_repo_name: BSnelling/julia


text:
# Biryani -- A conversion and validation toolbox
# By: Emmanuel Raviart <emmanuel@raviart.com>
#
# Copyright (C) 2015 Emmanuel Raviart
# https://github.com/eraviart/Biryani.jl
#
# This file is part of Biryani.
#
# 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.

importall Biryani.JsonConverters

# input_to_json
@test Convertible("""{"a": 1, "b": [2, "three"]}""") |> input_to_json |> to_value == Dict("a" => 1, "b" => [2, "three"])
@test Convertible("""      {"a": 1, "b": [2, "three"]}   """) |> input_to_json |> to_value ==
  Dict("a" => 1, "b" => [2, "three"])
@test Convertible("null") |> input_to_json |> to_value == nothing
@test Convertible("   null   ") |> input_to_json |> to_value == nothing
@test Convertible("Hello world!") |> input_to_json |> to_value_error == ("Hello world!", "Invalid JSON")
@test Convertible("   Hello world!   ") |> input_to_json |> to_value_error == ("   Hello world!   ", "Invalid JSON")
@test Convertible("""{"a": 1, "b":""") |> input_to_json |> to_value_error == ("""{"a": 1, "b":""", "Invalid JSON")
@test Convertible("""    {"a": 1, "b":""") |> input_to_json |> to_value_error == ("""    {"a": 1, "b":""",
  "Invalid JSON")
@test Convertible("") |> input_to_json |> to_value == nothing
@test Convertible("   ") |> input_to_json |> to_value_error == ("   ", "Invalid JSON")
@test Convertible(nothing) |> input_to_json |> to_value === nothing
@test Convertible("""{"a": 1, "b": [2, "three"]}""") |> fail("Initial error.") |> input_to_json |> to_value_error == (
  """{"a": 1, "b": [2, "three"]}""", "Initial error.")

ext: jl


text:
module Fields

using Test
using Gridap
using Gridap.Helpers

export FieldLike
export Field
export Basis
export Geomap
export num_dofs
export gradient
export ∇
export test_fieldlike
export test_field
export test_field_without_gradient
export test_basis
import Base: length
import Base: eltype

"""
Umbrella type for Field and Basis
"""
const FieldLike{D,T<:FieldValue,N} = Map{Point{D},1,T,N}

function eltype(::FieldLike{D,T}) where {D,T}
  return T
end

"""
Create the gradient of a `Field` or `Basis`
"""
function gradient(this::FieldLike)
  @abstractmethod
end

const ∇ = gradient

"""
Abstract field of rank `T` (e.g., scalar, vector, tensor) on a manifold of
dimension `D`
"""
const Field{D,T<:FieldValue} = FieldLike{D,T,1}

"""
Abstract basis for a space of fields of rank `T` (e.g., scalar, vector, tensor)
on a manifold of dimension `D`.

A Basis is evaluated at an array of Points and returns a matrix of values.
The first dimension in the returned matrix corresponds to the dofs of the basis,
whereas the second dimension corresponds to the evaluation points.
"""
const Basis{D,T<:FieldValue} = FieldLike{D,T,2}

@inline function num_dofs(b::Basis)
  n, = return_size(b,(1,))
  n
end

@inline length(b::Basis) = num_dofs(b)

"""
Abstract geometry map
"""
const Geomap = Field{D,Point{Z,X}} where {D,Z,X}

# Testers

function test_fieldlike(
  m::FieldLike{D,T,N},
  x::AbstractVector{<:Point{D}},
  v::AbstractArray{T,N},
  g::AbstractArray{G,N}) where {D,T,N,G}
  test_map(m,x,v)
  mg = gradient(m)
  test_map(mg,x,g)
end

function test_field(
  m::Field{D,T},
  x::AbstractVector{<:Point{D}},
  v::AbstractVector{T},
  g::AbstractVector{G}) where {D,T,G}
  test_fieldlike(m,x,v,g)
end

function test_field_without_gradient(
  m::Field{D,T},
  x::AbstractVector{<:Point{D}},
  v::AbstractVector{T}) where {D,T}
  test_map(m,x,v)
end

function test_basis(
  m::Basis{D,T},
  x::AbstractVector{<:Point{D}},
  v::AbstractMatrix{T},
  g::AbstractMatrix{G}) where {D,T,G}
  nd = num_dofs(m)
  @test nd == size(v,1)
  @test nd == size(g,1)
  test_fieldlike(m,x,v,g)
  mg = gradient(m)
  nd = num_dofs(mg)
  @test nd == size(g,1)
  nd = length(mg)
  @test nd == size(g,1)
end

end # module


text:
using ACEds

using StaticArrays
using ACEds.equivRotations3D

"""
Given an l-vector `ll` iterate over all combinations of `mm` vectors  of
the same length such that `sum(mm) == 1`
"""
struct MRange2{N, T2}
   ll::SVector{N, Int}
   cartrg::T2
end

Base.length(mr::MRange2) = sum(_->1, _mrange2(mr.ll))

_mrange2(ll) = MRange2(ll, Iterators.Stateful(
               filter((x) -> abs(sum(x))<= 1, Tuple.(CartesianIndices(ntuple(i -> -ll[i]:ll[i], length(ll)))))
                     ))

function Base.iterate(mr::MRange2, args...)
   while true
      if isempty(mr.cartrg)
         return nothing
      end
      mpre = popfirst!(mr.cartrg)
      return SVector(mpre), nothing
   end
end

"""
Compare behaviour of iterate(mr::MRange) with iterate(mr::MRange2)
"""

function mytest(mr)
   for (i,val) in enumerate(mr)
      print("index:", i, " m-vector: ", val, "\n")
   end
end

ll = @SVector([10,10,10])
mr1 = ACEds.equivRotations3D._mrange(ll)
mr2 = _mrange2(ll)
mytest(mr2)
print("-----------")


text:
#
# Adding constraint handlers and constraints to SCIP.Optimizer.
#

"""
    include_conshdlr(
        o::Optimizer,
        ch::CH;
        name::String,
        description::String,
        enforce_priority::Int,
        check_priority::Int,
        eager_frequency::Int,
        needs_constraints::Bool
    )

Include a user defined constraint handler `ch` to SCIP optimizer instance `o`.

All parameters have default values that can be set as keyword arguments.
In particular, note the boolean `needs_constraints`:
* If set to `true`, then the callbacks are only called for the constraints that
  were added explicitly using `add_constraint`.
* If set to `false`, the callback functions will always be called, even if no
  corresponding constraint was added. It probably makes sense to set
  `misc/allowdualreds` to `FALSE` in this case.
"""
function include_conshdlr(o::Optimizer, ch::CH;
                          name="", description="", enforce_priority=-15,
                          check_priority=-7000000, eager_frequency=100,
                          needs_constraints=true) where CH <: AbstractConstraintHandler
    include_conshdlr(o.mscip, ch, name=name, description=description,
                     enforce_priority=enforce_priority,
                     check_priority=check_priority,
                     eager_frequency=eager_frequency,
                     needs_constraints=needs_constraints)
end

"""
    add_constraint(
        o::Optimizer,
        ch::CH,
        c::C;
        initial=true,
        separate=true,
        enforce=true,
        check=true,
        propagate=true,
        _local=false,
        modifiable=false,
        dynamic=false,
        removable=false,
        stickingatnode=false
    )::ConsRef

Add constraint `c` belonging to user-defined constraint handler `ch` to model.
Returns constraint reference.

All keyword arguments are passed to the `SCIPcreateCons` call.
"""
function add_constraint(o::Optimizer, ch::CH, c::C;
                        initial=true, separate=true, enforce=true, check=true,
                        propagate=true, _local=false, modifiable=false,
                        dynamic=false, removable=false, stickingatnode=false) where {CH <:AbstractConstraintHandler, C <: AbstractConstraint{CH}}
    return add_constraint(o.mscip, ch, c, initial=initial, separate=separate,
                          enforce=enforce, check=check, propagate=propagate,
                          _local=_local, modifiable=modifiable, dynamic=dynamic,
                          removable=removable, stickingatnode=stickingatnode)
end


text:
function no_whitening(
    dataset::DataSet{T, I})::WhiteningResult{T} where {T<:AbstractFloat, I<:Integer}

    datamean = zeros(T, dataset.P)

    transform_data!(dataset, Array{T, 2}(LinearAlgebra.I, dataset.P, dataset.P), datamean)
end

function cholesky_whitening!(
    dataset::DataSet{T, I})::WhiteningResult{T} where {T<:AbstractFloat, I<:Integer}

    datamean = zeros(T, dataset.P)

    for p in eachindex(datamean)
        datamean[p] = mean(view(dataset.data, p, :))
    end

    for n in 1:dataset.N
        buffer = view(dataset.data, :, n) - datamean
        setindex!(dataset.data, buffer, :, n)
    end

    #cov of a Transpose object yields the wrong (0.0) result in 1D, if all weights are equal (e.g. direct sampling)!
    covmatrix = cov(convert(typeof(dataset.data), transpose(dataset.data)), FrequencyWeights(dataset.weights), corrected=true)
    symcovmatrix = Symmetric(covmatrix)
    covmatrix_inv = inv(symcovmatrix)
    w = cholesky(covmatrix_inv).U
    wres = convert(Matrix{T}, w)

    transform_data!(dataset, wres, datamean)
end

function statistical_whitening!(
    dataset::DataSet{T, I})::WhiteningResult{T} where {T<:AbstractFloat, I<:Integer}

    datamean = zeros(T, dataset.P)

    for p in eachindex(datamean)
        datamean[p] = mean(view(dataset.data, p, :))
    end

    dataset.data .-= datamean

    covmatrix = cov(transpose(dataset.data), FrequencyWeights(dataset.weights), corrected=true)

    E = eigfact(covmatrix).vectors
    w_d = transpose(E)
    D = Diagonal(inv(E) * covmatrix * E)
    w = inv(full(sqrt.(D))) * w_d
    wres = convert(Matrix{T}, w)

    transform_data!(dataset, wres, datamean)
end

function transform_data!(
    dataset::DataSet{T, I},
    W::Matrix{T},
    datamean::Vector{T},
    substractmean::Bool = false)::WhiteningResult{T} where {T<:AbstractFloat, I<:Integer}

    local determinant::T

    if substractmean
        dataset.data .-= datamean
    end

    if diag(W) == ones(T, dataset.P)
        determinant = 1.0
    else
        dataset.data = W * dataset.data

        determinant = abs(det(W))
    end

    maxP::T = maximum(dataset.logprob)
    suggTargetProb::T = exp(maxP - partialsort(dataset.logprob, floor(Int64, dataset.N * 0.2)))

    dataset.iswhitened = true

    @debug "Determinant:\t" * string(determinant)
    @debug "Suggested Target Probability Factor:\t" * string(suggTargetProb)

    WhiteningResult(determinant, suggTargetProb, W, datamean)
end


text:
using JuMP, EAGO

                     m = Model()

                     EAGO.register_eago_operators!(m)

                     @variable(m, -1 <= x[i=1:5] <= 1)

                     @variable(m, -6.148474362391325 <= q <= 10.677081718106185)

                     add_NL_constraint(m, :(1/(1 + exp(-(-0.2518902526786948 + 0.9847866884384731*1/(1 + exp(-(0.2752793536861313 + -0.1568397657479923*1/(1 + exp(-(0.860587138591963 + -0.8175452953351838*$(x[1]) + -0.2911974841119127*$(x[2]) + -0.03381280038289569*$(x[3]) + -0.021805451939043596*$(x[4]) + 0.5056820468199099*$(x[5])))) + -0.1062303426643516*1/(1 + exp(-(0.40215662414311426 + 0.47796901525165936*$(x[1]) + -0.8905174133817706*$(x[2]) + 0.3815447257070974*$(x[3]) + -0.7235559998197827*$(x[4]) + -0.43821768660130234*$(x[5])))) + 0.5642457669318222*1/(1 + exp(-(-0.371366893645773 + 0.03574445051584929*$(x[1]) + 0.07916257211879385*$(x[2]) + 0.5096612381208421*$(x[3]) + 0.1651890232922466*$(x[4]) + -0.4691295251335359*$(x[5]))))))) + -0.7649756572019379*1/(1 + exp(-(0.9918790053549298 + 0.85981905523253*1/(1 + exp(-(0.860587138591963 + -0.8175452953351838*$(x[1]) + -0.2911974841119127*$(x[2]) + -0.03381280038289569*$(x[3]) + -0.021805451939043596*$(x[4]) + 0.5056820468199099*$(x[5])))) + 0.9639515148457134*1/(1 + exp(-(0.40215662414311426 + 0.47796901525165936*$(x[1]) + -0.8905174133817706*$(x[2]) + 0.3815447257070974*$(x[3]) + -0.7235559998197827*$(x[4]) + -0.43821768660130234*$(x[5])))) + -0.08816238542180388*1/(1 + exp(-(-0.371366893645773 + 0.03574445051584929*$(x[1]) + 0.07916257211879385*$(x[2]) + 0.5096612381208421*$(x[3]) + 0.1651890232922466*$(x[4]) + -0.4691295251335359*$(x[5]))))))) + 0.10087195663512549*1/(1 + exp(-(-0.8202416682813327 + -0.17766510212211006*1/(1 + exp(-(0.860587138591963 + -0.8175452953351838*$(x[1]) + -0.2911974841119127*$(x[2]) + -0.03381280038289569*$(x[3]) + -0.021805451939043596*$(x[4]) + 0.5056820468199099*$(x[5])))) + -0.1793020696087071*1/(1 + exp(-(0.40215662414311426 + 0.47796901525165936*$(x[1]) + -0.8905174133817706*$(x[2]) + 0.3815447257070974*$(x[3]) + -0.7235559998197827*$(x[4]) + -0.43821768660130234*$(x[5])))) + -0.41892665263312834*1/(1 + exp(-(-0.371366893645773 + 0.03574445051584929*$(x[1]) + 0.07916257211879385*$(x[2]) + 0.5096612381208421*$(x[3]) + 0.1651890232922466*$(x[4]) + -0.4691295251335359*$(x[5])))))))))) + 1/(1 + exp(-(0.638435553115452 + 0.09453389081965424*1/(1 + exp(-(0.2752793536861313 + -0.1568397657479923*1/(1 + exp(-(0.860587138591963 + -0.8175452953351838*$(x[1]) + -0.2911974841119127*$(x[2]) + -0.03381280038289569*$(x[3]) + -0.021805451939043596*$(x[4]) + 0.5056820468199099*$(x[5])))) + -0.1062303426643516*1/(1 + exp(-(0.40215662414311426 + 0.47796901525165936*$(x[1]) + -0.8905174133817706*$(x[2]) + 0.3815447257070974*$(x[3]) + -0.7235559998197827*$(x[4]) + -0.43821768660130234*$(x[5])))) + 0.5642457669318222*1/(1 + exp(-(-0.371366893645773 + 0.03574445051584929*$(x[1]) + 0.07916257211879385*$(x[2]) + 0.5096612381208421*$(x[3]) + 0.1651890232922466*$(x[4]) + -0.4691295251335359*$(x[5]))))))) + -0.07927014075141203*1/(1 + exp(-(0.9918790053549298 + 0.85981905523253*1/(1 + exp(-(0.860587138591963 + -0.8175452953351838*$(x[1]) + -0.2911974841119127*$(x[2]) + -0.03381280038289569*$(x[3]) + -0.021805451939043596*$(x[4]) + 0.5056820468199099*$(x[5])))) + 0.9639515148457134*1/(1 + exp(-(0.40215662414311426 + 0.47796901525165936*$(x[1]) + -0.8905174133817706*$(x[2]) + 0.3815447257070974*$(x[3]) + -0.7235559998197827*$(x[4]) + -0.43821768660130234*$(x[5])))) + -0.08816238542180388*1/(1 + exp(-(-0.371366893645773 + 0.03574445051584929*$(x[1]) + 0.07916257211879385*$(x[2]) + 0.5096612381208421*$(x[3]) + 0.1651890232922466*$(x[4]) + -0.4691295251335359*$(x[5]))))))) + 0.2624391124261729*1/(1 + exp(-(-0.8202416682813327 + -0.17766510212211006*1/(1 + exp(-(0.860587138591963 + -0.8175452953351838*$(x[1]) + -0.2911974841119127*$(x[2]) + -0.03381280038289569*$(x[3]) + -0.021805451939043596*$(x[4]) + 0.5056820468199099*$(x[5])))) + -0.1793020696087071*1/(1 + exp(-(0.40215662414311426 + 0.47796901525165936*$(x[1]) + -0.8905174133817706*$(x[2]) + 0.3815447257070974*$(x[3]) + -0.7235559998197827*$(x[4]) + -0.43821768660130234*$(x[5])))) + -0.41892665263312834*1/(1 + exp(-(-0.371366893645773 + 0.03574445051584929*$(x[1]) + 0.07916257211879385*$(x[2]) + 0.5096612381208421*$(x[3]) + 0.1651890232922466*$(x[4]) + -0.4691295251335359*$(x[5])))))))))) + 1/(1 + exp(-(0.5273095272255199 + 0.4025366346817978*1/(1 + exp(-(0.2752793536861313 + -0.1568397657479923*1/(1 + exp(-(0.860587138591963 + -0.8175452953351838*$(x[1]) + -0.2911974841119127*$(x[2]) + -0.03381280038289569*$(x[3]) + -0.021805451939043596*$(x[4]) + 0.5056820468199099*$(x[5])))) + -0.1062303426643516*1/(1 + exp(-(0.40215662414311426 + 0.47796901525165936*$(x[1]) + -0.8905174133817706*$(x[2]) + 0.3815447257070974*$(x[3]) + -0.7235559998197827*$(x[4]) + -0.43821768660130234*$(x[5])))) + 0.5642457669318222*1/(1 + exp(-(-0.371366893645773 + 0.03574445051584929*$(x[1]) + 0.07916257211879385*$(x[2]) + 0.5096612381208421*$(x[3]) + 0.1651890232922466*$(x[4]) + -0.4691295251335359*$(x[5]))))))) + -0.585290488444234*1/(1 + exp(-(0.9918790053549298 + 0.85981905523253*1/(1 + exp(-(0.860587138591963 + -0.8175452953351838*$(x[1]) + -0.2911974841119127*$(x[2]) + -0.03381280038289569*$(x[3]) + -0.021805451939043596*$(x[4]) + 0.5056820468199099*$(x[5])))) + 0.9639515148457134*1/(1 + exp(-(0.40215662414311426 + 0.47796901525165936*$(x[1]) + -0.8905174133817706*$(x[2]) + 0.3815447257070974*$(x[3]) + -0.7235559998197827*$(x[4]) + -0.43821768660130234*$(x[5])))) + -0.08816238542180388*1/(1 + exp(-(-0.371366893645773 + 0.03574445051584929*$(x[1]) + 0.07916257211879385*$(x[2]) + 0.5096612381208421*$(x[3]) + 0.1651890232922466*$(x[4]) + -0.4691295251335359*$(x[5]))))))) + 0.1647823489958915*1/(1 + exp(-(-0.8202416682813327 + -0.17766510212211006*1/(1 + exp(-(0.860587138591963 + -0.8175452953351838*$(x[1]) + -0.2911974841119127*$(x[2]) + -0.03381280038289569*$(x[3]) + -0.021805451939043596*$(x[4]) + 0.5056820468199099*$(x[5])))) + -0.1793020696087071*1/(1 + exp(-(0.40215662414311426 + 0.47796901525165936*$(x[1]) + -0.8905174133817706*$(x[2]) + 0.3815447257070974*$(x[3]) + -0.7235559998197827*$(x[4]) + -0.43821768660130234*$(x[5])))) + -0.41892665263312834*1/(1 + exp(-(-0.371366893645773 + 0.03574445051584929*$(x[1]) + 0.07916257211879385*$(x[2]) + 0.5096612381208421*$(x[3]) + 0.1651890232922466*$(x[4]) + -0.4691295251335359*$(x[5])))))))))) - $q <= 0.0))

                     @objective(m, Min, q)

                     return m

max_stars_repo_name: PSORLab/RSActivationFunctions

max_issues_repo_name: PSORLab/RSActivationFunctions

max_forks_repo_name: PSORLab/RSActivationFunctions


text:
# This file contains the definition for the Renewables struct.

mutable struct Renewables
    unit
    node
    pmax
	generation
    mc
    fuel
    name

    function Renewables(res_df::DataFrame,
                        avail_solar_df::DataFrame,
                        avail_wind_df::DataFrame)

		R = res_df[!, 1]
        RES = Vector(R)
        res_dict = df_to_dict_with_id(res_df)

		availability = [avail_solar_df avail_wind_df]

		gen_dict = Dict(RES[col] => availability[!, Symbol(res_dict[:node][col])]
			* res_dict[:pmax][col] for col in 1:length(RES))

        return new(R,
                   res_dict[:node],
                   res_dict[:pmax],
				   gen_dict,
				   res_dict[:mc],
				   res_dict[:fuel],
				   res_dict[:name],
                   )
    end
end


text:
A = randn(3,3)

x = fill(1.0, (3))
b = A*x

l,u,p = lu(A)

norm(l*u - A)

norm(l*u - A[p,:])

l,u,p = lu(A, Val{false})

norm(l*u - A)

Alu = lufact(A)

Alu[:P]

Alu[:L]

Alu[:U]

# PA = LU
# A = P'LU
# P'LUx = b
# LUx = Pb
# Ux = L\Pb
# x = U\L\Pb
Alu[:U]\(Alu[:L]\(Alu[:P]b))

Alu\b

det(A)

det(Alu)

Atall = randn(3, 2)

Aqr = qrfact(Atall)

Aqr[:Q]

Aqr[:R]

Aqr[:Q]*ones(2)

Aqr[:Q]*ones(3)

Aqr[:Q]*ones(4)

Aqr[:Q][1]

Aqr\b

Atall\b

v = randn(3)
# Tomar el producto exterio de un vector consigo mismo nos da una matriz singular
singmatriz = v * v'

Aqrp = qrfact(singmatriz,Val{true})

Aqrp\b

Asym = A + A'
AsymEig = eigfact(Asym)

AsymEig[:values]

AsymEig[:vectors]

inv(AsymEig)*Asym

eig(Asym)

Asvd = svdfact(A[:,1:2])

Asvd\b

A

Diagonal(diag(A))

Diagonal(A)

LowerTriangular(tril(A))

LowerTriangular(A)

Symmetric(Asym)

SymTridiagonal(diag(Asym),diag(Asym,1))

n = 1000;
A = randn(n,n);

Asym = A + A';

Asym_noisy = copy(Asym); Asym_noisy[1,2] += 5eps();

println("Is Asym symmetric? ", issymmetric(Asym))
println("Is Asym_noisy symmetric? ", issymmetric(Asym_noisy))

@time eigvals(Asym);

@time eigvals(Asym_noisy);

@time eigvals(Symmetric(Asym_noisy));

n = 1_000_000;
A = SymTridiagonal(randn(n), randn(n-1));
@time eigmax(A)

rand(1:10,3,3)*rand(1:10,3,3)

1//2

Ar = convert(Matrix{Rational{BigInt}}, rand(1:10,3,3))/10

x = fill(1, (3))
b = Ar*x

Ar\b

lufact(Ar)

λ1,λ2,λ3 = 1//1,1//2,1//4
v1,v2,v3 = [1,0,0],[1,1,0],[1,1,1]
V,Λ = [v1 v2 v3], Diagonal([λ1,λ2,λ3])
A = V*Λ/V

max_stars_repo_path: es-es/intro-to-julia-ES/plutos/11.Factorizaciones y otras diversiones.jl

max_stars_repo_name: miguelraz/JuliaTutorials

max_issues_repo_name: miguelraz/JuliaTutorials

max_forks_repo_name: miguelraz/JuliaTutorials


text:
# Naive implementations of multiply_dimensionwise used to demonstrate the functionality
# without performance optimizations and for testing correctness of the optimized versions
# implemented below.
function multiply_dimensionwise_naive(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 2})
  size_out = size(matrix, 1)
  size_in  = size(matrix, 2)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out)

  for i in 1:size_out
    for ii in 1:size_in
      for v in 1:n_vars
        data_out[v, i] += matrix[i, ii] * data_in[v, ii]
      end
    end
  end

  return data_out
end

function multiply_dimensionwise_naive(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 3})
  size_out = size(matrix, 1)
  size_in  = size(matrix, 2)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out, size_out)

  for j in 1:size_out, i in 1:size_out
    for jj in 1:size_in, ii in 1:size_in
      for v in 1:n_vars
        data_out[v, i, j] += matrix[i, ii] * matrix[j, jj] * data_in[v, ii, jj]
      end
    end
  end

  return data_out
end

function multiply_dimensionwise_naive(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 4})
  size_out = size(matrix, 1)
  size_in  = size(matrix, 2)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out, size_out, size_out)

  for k in 1:size_out, j in 1:size_out, i in 1:size_out
    for kk in 1:size_in, jj in 1:size_in, ii in 1:size_in
      for v in 1:n_vars
        data_out[v, i, j, k] += matrix[i, ii] * matrix[j, jj] * matrix[k, kk] * data_in[v, ii, jj, kk]
      end
    end
  end

  return data_out
end

"""
    multiply_dimensionwise(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, NDIMS+1})

Multiply the array `data_in` by `matrix` in each coordinate direction, where `data_in`
is assumed to have the first coordinate for the number of variables and the remaining coordinates
are multiplied by `matrix`.
"""
function multiply_dimensionwise(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 2})
  # 1D
  # optimized version of multiply_dimensionwise_naive
  size_out = size(matrix, 1)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out)

  multiply_dimensionwise!(data_out, matrix, data_in)

  return data_out
end

function multiply_dimensionwise(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 3})
  # 2D
  # optimized version of multiply_dimensionwise_naive
  size_out = size(matrix, 1)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out, size_out)

  multiply_dimensionwise!(data_out, matrix, data_in)

  return data_out
end

function multiply_dimensionwise(matrix::AbstractMatrix, data_in::AbstractArray{<:Any, 4})
  # 3D
  # optimized version of multiply_dimensionwise_naive
  size_out = size(matrix, 1)
  n_vars   = size(data_in, 1)
  data_out = zeros(promote_type(eltype(data_in), eltype(matrix)), n_vars, size_out, size_out, size_out)

  multiply_dimensionwise!(data_out, matrix, data_in)

  return data_out
end

# In the following, there are several optimized in-place versions of multiply_dimensionwise.
# These may make use of advanced optimization features such as the macro `@tullio` from Tullio.jl,
# which basically uses an Einstein summation convention syntax.
# Another possibility is `@turbo` from LoopVectorization.jl. The runtime performance could be even
# optimized further by using `@turbo inline=true for` instead of `@turbo for`, but that comes at the
# cost of increased latency, at least on some systems...

# 1D version
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 2}, matrix::AbstractMatrix,
                                 data_in ::AbstractArray{<:Any, 2})
  # @tullio threads=false data_out[v, i] = matrix[i, ii] * data_in[v, ii]
  @turbo for i in axes(data_out, 2), v in axes(data_out, 1)
    res = zero(eltype(data_out))
    for ii in axes(matrix, 2)
      res += matrix[i, ii] * data_in[v, ii]
    end
    data_out[v, i] = res
  end

  return nothing
end

# 1D version for scalars
# Instead of having a leading dimension of size 1 in `data_out, data_in`, this leading dimension
# of size unity is dropped, resulting in one dimension less than in `multiply_dimensionwise!`.
function multiply_scalar_dimensionwise!(data_out::AbstractArray{<:Any, 1},
                                        matrix::AbstractMatrix,
                                        data_in ::AbstractArray{<:Any, 1})
  # @tullio threads=false data_out[i] = matrix[i, ii] * data_in[ii]
  @turbo for i in axes(data_out, 1)
    res = zero(eltype(data_out))
    for ii in axes(matrix, 2)
      res += matrix[i, ii] * data_in[ii]
    end
    data_out[i] = res
  end

  return nothing
end

# 1D version, apply matrixJ to data_inJ
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 2}, matrix1::AbstractMatrix,
                                 data_in1::AbstractArray{<:Any, 2}, matrix2::AbstractMatrix,
                                 data_in2::AbstractArray{<:Any, 2})
  # @tullio threads=false data_out[v, i] = matrix1[i, ii] * data_in1[v, ii] + matrix2[i, ii] * data_in2[v, ii]
  # TODO: LoopVectorization upgrade
  #   We would like to use `@turbo` for the outermost loop possibly fuse both inner
  #   loops, but that does currently not work because of limitations of
  #   LoopVectorizationjl. However, Chris Elrod is planning to address this in
  #   the future, cf. https://github.com/JuliaSIMD/LoopVectorization.jl/issues/230#issuecomment-810632972
  @turbo for i in axes(data_out, 2), v in axes(data_out, 1)
    res = zero(eltype(data_out))
    for ii in axes(matrix1, 2)
      res += matrix1[i, ii] * data_in1[v, ii]
    end
    data_out[v, i] = res
  end
  @turbo for i in axes(data_out, 2), v in axes(data_out, 1)
    res = zero(eltype(data_out))
    for ii in axes(matrix2, 2)
      res += matrix2[i, ii] * data_in2[v, ii]
    end
    data_out[v, i] += res
  end

  return nothing
end

# 2D version
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 3}, matrix::AbstractMatrix,
                                 data_in:: AbstractArray{<:Any, 3},
                                 tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix, 1), size(matrix, 2)))

  # Interpolate in x-direction
  # @tullio threads=false tmp1[v, i, j]     = matrix[i, ii] * data_in[v, ii, j]
  @turbo for j in axes(tmp1, 3), i in axes(tmp1, 2), v in axes(tmp1, 1)
    res = zero(eltype(tmp1))
    for ii in axes(matrix, 2)
      res += matrix[i, ii] * data_in[v, ii, j]
    end
    tmp1[v, i, j] = res
  end

  # Interpolate in y-direction
  # @tullio threads=false data_out[v, i, j] = matrix[j, jj] * tmp1[v, i, jj]
  @turbo for j in axes(data_out, 3), i in axes(data_out, 2), v in axes(data_out, 1)
    res = zero(eltype(data_out))
    for jj in axes(matrix, 2)
      res += matrix[j, jj] * tmp1[v, i, jj]
    end
    data_out[v, i, j] = res
  end

  return nothing
end

# 2D version for scalars
# Instead of having a leading dimension of size 1 in `data_out, data_in`, this leading dimension
# of size unity is dropped, resulting in one dimension less than in `multiply_dimensionwise!`.
function multiply_scalar_dimensionwise!(data_out::AbstractArray{<:Any, 2},
                                        matrix::AbstractMatrix,
                                        data_in:: AbstractArray{<:Any, 2},
                                        tmp1=zeros(eltype(data_out), size(matrix, 1), size(matrix, 2)))

  # Interpolate in x-direction
  # @tullio threads=false     tmp1[i, j] = matrix[i, ii] * data_in[ii, j]
  @turbo for j in axes(tmp1, 2), i in axes(tmp1, 1)
    res = zero(eltype(tmp1))
    for ii in axes(matrix, 2)
      res += matrix[i, ii] * data_in[ii, j]
    end
    tmp1[i, j] = res
  end

  # Interpolate in y-direction
  # @tullio threads=false data_out[i, j] = matrix[j, jj] * tmp1[i, jj]
  @turbo for j in axes(data_out, 2), i in axes(data_out, 1)
    res = zero(eltype(data_out))
    for jj in axes(matrix, 2)
      res += matrix[j, jj] * tmp1[i, jj]
    end
    data_out[i, j] = res
  end

  return nothing
end

# 2D version, apply matrixJ to dimension J of data_in
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 3},
                                 matrix1::AbstractMatrix, matrix2::AbstractMatrix,
                                 data_in:: AbstractArray{<:Any, 3},
                                 tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 2)))

  # Interpolate in x-direction
  # @tullio threads=false tmp1[v, i, j]     = matrix1[i, ii] * data_in[v, ii, j]
  @turbo for j in axes(tmp1, 3), i in axes(tmp1, 2), v in axes(tmp1, 1)
    res = zero(eltype(tmp1))
    for ii in axes(matrix1, 2)
      res += matrix1[i, ii] * data_in[v, ii, j]
    end
    tmp1[v, i, j] = res
  end

  # Interpolate in y-direction
  # @tullio threads=false data_out[v, i, j] = matrix2[j, jj] * tmp1[v, i, jj]
  @turbo for j in axes(data_out, 3), i in axes(data_out, 2), v in axes(data_out, 1)
    res = zero(eltype(data_out))
    for jj in axes(matrix2, 2)
      res += matrix2[j, jj] * tmp1[v, i, jj]
    end
    data_out[v, i, j] = res
  end

  return nothing
end

# 2D version, apply matrixJ to dimension J of data_in and add the result to data_out
function add_multiply_dimensionwise!(data_out::AbstractArray{<:Any, 3},
                                     matrix1::AbstractMatrix, matrix2::AbstractMatrix,
                                     data_in:: AbstractArray{<:Any, 3},
                                     tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 2)))

  # Interpolate in x-direction
  # @tullio threads=false tmp1[v, i, j]     = matrix1[i, ii] * data_in[v, ii, j]
  @turbo for j in axes(tmp1, 3), i in axes(tmp1, 2), v in axes(tmp1, 1)
    res = zero(eltype(tmp1))
    for ii in axes(matrix1, 2)
      res += matrix1[i, ii] * data_in[v, ii, j]
    end
    tmp1[v, i, j] = res
  end

  # Interpolate in y-direction
  # @tullio threads=false data_out[v, i, j] += matrix2[j, jj] * tmp1[v, i, jj]
  @turbo for j in axes(data_out, 3), i in axes(data_out, 2), v in axes(data_out, 1)
    res = zero(eltype(data_out))
    for jj in axes(matrix2, 2)
      res += matrix2[j, jj] * tmp1[v, i, jj]
    end
    data_out[v, i, j] += res
  end

  return nothing
end

# 3D version
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 4}, matrix::AbstractMatrix,
                                 data_in:: AbstractArray{<:Any, 4},
                                 tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix, 1), size(matrix, 2), size(matrix, 2)),
                                 tmp2=zeros(eltype(data_out), size(data_out, 1), size(matrix, 1), size(matrix, 1), size(matrix, 2)))

  # Interpolate in x-direction
  # @tullio threads=false tmp1[v, i, j, k]     = matrix[i, ii] * data_in[v, ii, j, k]
  @turbo for k in axes(tmp1, 4), j in axes(tmp1, 3), i in axes(tmp1, 2), v in axes(tmp1, 1)
    res = zero(eltype(tmp1))
    for ii in axes(matrix, 2)
      res += matrix[i, ii] * data_in[v, ii, j, k]
    end
    tmp1[v, i, j, k] = res
  end

  # Interpolate in y-direction
  # @tullio threads=false tmp2[v, i, j, k]     = matrix[j, jj] * tmp1[v, i, jj, k]
  @turbo for k in axes(tmp2, 4), j in axes(tmp2, 3), i in axes(tmp2, 2), v in axes(tmp2, 1)
    res = zero(eltype(tmp2))
    for jj in axes(matrix, 2)
      res += matrix[j, jj] * tmp1[v, i, jj, k]
    end
    tmp2[v, i, j, k] = res
  end

  # Interpolate in z-direction
  # @tullio threads=false data_out[v, i, j, k] = matrix[k, kk] * tmp2[v, i, j, kk]
  @turbo for k in axes(data_out, 4), j in axes(data_out, 3), i in axes(data_out, 2), v in axes(data_out, 1)
    res = zero(eltype(data_out))
    for kk in axes(matrix, 2)
      res += matrix[k, kk] * tmp2[v, i, j, kk]
    end
    data_out[v, i, j, k] = res
  end

  return nothing
end

# 3D version for scalars
# Instead of having a leading dimension of size 1 in `data_out, data_in`, this leading dimension
# of size unity is dropped, resulting in one dimension less than in `multiply_dimensionwise!`.
function multiply_scalar_dimensionwise!(data_out::AbstractArray{<:Any, 3}, matrix::AbstractMatrix,
                                        data_in:: AbstractArray{<:Any, 3},
                                        tmp1=zeros(eltype(data_out), size(matrix, 1), size(matrix, 2), size(matrix, 2)),
                                        tmp2=zeros(eltype(data_out), size(matrix, 1), size(matrix, 1), size(matrix, 2)))

  # Interpolate in x-direction
  # @tullio threads=false tmp1[i, j, k]     = matrix[i, ii] * data_in[ii, j, k]
  @turbo for k in axes(tmp1, 3), j in axes(tmp1, 2), i in axes(tmp1, 1)
    res = zero(eltype(tmp1))
    for ii in axes(matrix, 2)
      res += matrix[i, ii] * data_in[ii, j, k]
    end
    tmp1[i, j, k] = res
  end

  # Interpolate in y-direction
  # @tullio threads=false tmp2[i, j, k]     = matrix[j, jj] * tmp1[i, jj, k]
  @turbo for k in axes(tmp2, 3), j in axes(tmp2, 2), i in axes(tmp2, 1)
    res = zero(eltype(tmp2))
    for jj in axes(matrix, 2)
      res += matrix[j, jj] * tmp1[i, jj, k]
    end
    tmp2[i, j, k] = res
  end

  # Interpolate in z-direction
  # @tullio threads=false data_out[i, j, k] = matrix[k, kk] * tmp2[i, j, kk]
  @turbo for k in axes(data_out, 3), j in axes(data_out, 2), i in axes(data_out, 1)
    res = zero(eltype(data_out))
    for kk in axes(matrix, 2)
      res += matrix[k, kk] * tmp2[i, j, kk]
    end
    data_out[i, j, k] = res
  end

  return nothing
end

# 3D version, apply matrixJ to dimension J of data_in
function multiply_dimensionwise!(data_out::AbstractArray{<:Any, 4},
                                 matrix1::AbstractMatrix, matrix2::AbstractMatrix, matrix3::AbstractMatrix,
                                 data_in:: AbstractArray{<:Any, 4},
                                 tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 2), size(matrix1, 2)),
                                 tmp2=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 1), size(matrix1, 2)))

  # Interpolate in x-direction
  # @tullio threads=false tmp1[v, i, j, k]     = matrix1[i, ii] * data_in[v, ii, j, k]
  @turbo for k in axes(tmp1, 4), j in axes(tmp1, 3), i in axes(tmp1, 2), v in axes(tmp1, 1)
    res = zero(eltype(tmp1))
    for ii in axes(matrix1, 2)
      res += matrix1[i, ii] * data_in[v, ii, j, k]
    end
    tmp1[v, i, j, k] = res
  end

  # Interpolate in y-direction
  # @tullio threads=false tmp2[v, i, j, k]     = matrix2[j, jj] * tmp1[v, i, jj, k]
  @turbo for k in axes(tmp2, 4), j in axes(tmp2, 3), i in axes(tmp2, 2), v in axes(tmp2, 1)
    res = zero(eltype(tmp1))
    for jj in axes(matrix2, 2)
      res += matrix2[j, jj] * tmp1[v, i, jj, k]
    end
    tmp2[v, i, j, k] = res
  end

  # Interpolate in z-direction
  # @tullio threads=false data_out[v, i, j, k] = matrix3[k, kk] * tmp2[v, i, j, kk]
  @turbo for k in axes(data_out, 4), j in axes(data_out, 3), i in axes(data_out, 2), v in axes(data_out, 1)
    res = zero(eltype(data_out))
    for kk in axes(matrix3, 2)
      res += matrix3[k, kk] * tmp2[v, i, j, kk]
    end
    data_out[v, i, j, k] = res
  end

  return nothing
end

# 3D version, apply matrixJ to dimension J of data_in and add the result to data_out
function add_multiply_dimensionwise!(data_out::AbstractArray{<:Any, 4},
                                     matrix1::AbstractMatrix, matrix2::AbstractMatrix, matrix3::AbstractMatrix,
                                     data_in:: AbstractArray{<:Any, 4},
                                     tmp1=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 2), size(matrix1, 2)),
                                     tmp2=zeros(eltype(data_out), size(data_out, 1), size(matrix1, 1), size(matrix1, 1), size(matrix1, 2)))

  # Interpolate in x-direction
  # @tullio threads=false tmp1[v, i, j, k]     = matrix1[i, ii] * data_in[v, ii, j, k]
  @turbo for k in axes(tmp1, 4), j in axes(tmp1, 3), i in axes(tmp1, 2), v in axes(tmp1, 1)
    res = zero(eltype(tmp1))
    for ii in axes(matrix1, 2)
      res += matrix1[i, ii] * data_in[v, ii, j, k]
    end
    tmp1[v, i, j, k] = res
  end

  # Interpolate in y-direction
  # @tullio threads=false tmp2[v, i, j, k]     = matrix2[j, jj] * tmp1[v, i, jj, k]
  @turbo for k in axes(tmp2, 4), j in axes(tmp2, 3), i in axes(tmp2, 2), v in axes(tmp2, 1)
    res = zero(eltype(tmp1))
    for jj in axes(matrix2, 2)
      res += matrix2[j, jj] * tmp1[v, i, jj, k]
    end
    tmp2[v, i, j, k] = res
  end

  # Interpolate in z-direction
  # @tullio threads=false data_out[v, i, j, k] += matrix3[k, kk] * tmp2[v, i, j, kk]
  @turbo for k in axes(data_out, 4), j in axes(data_out, 3), i in axes(data_out, 2), v in axes(data_out, 1)
    res = zero(eltype(data_out))
    for kk in axes(matrix3, 2)
      res += matrix3[k, kk] * tmp2[v, i, j, kk]
    end
    data_out[v, i, j, k] += res
  end

  return nothing
end


text:
using MyProject_1
using Test

A = load_data("Chicago")
B = load_data("Chicago", false)

show_results(A)
show_results(B)
practice_plots(A)


text:
using Test
using Random
using LinearAlgebra
using ForwardDiff
using ReverseDiff
using Tracker
using DistributionsAD

using Bijectors
using Bijectors: Log, Exp, Shift, Scale, Logit, SimplexBijector, PDBijector, Permute, PlanarLayer, RadialLayer, Stacked, TruncatedBijector, ADBijector, RationalQuadraticSpline, LeakyReLU

Random.seed!(123)

struct MyADBijector{AD, N, B <: Bijector{N}} <: ADBijector{AD, N}
    b::B
end
MyADBijector(d::Distribution) = MyADBijector{Bijectors.ADBackend()}(d)
MyADBijector{AD}(d::Distribution) where {AD} = MyADBijector{AD}(bijector(d))
MyADBijector{AD}(b::B) where {AD, N, B <: Bijector{N}} = MyADBijector{AD, N, B}(b)
(b::MyADBijector)(x) = b.b(x)
(b::Inverse{<:MyADBijector})(x) = inv(b.orig.b)(x)

struct NonInvertibleBijector{AD} <: ADBijector{AD, 1} end

contains(predicate::Function, b::Bijector) = predicate(b)
contains(predicate::Function, b::Composed) = any(contains.(predicate, b.ts))
contains(predicate::Function, b::Stacked) = any(contains.(predicate, b.bs))

# Scalar tests
@testset "<: ADBijector{AD}" begin
    (b::NonInvertibleBijector)(x) = clamp.(x, 0, 1)

    b = NonInvertibleBijector{Bijectors.ADBackend()}()
    @test_throws Bijectors.SingularJacobianException logabsdetjac(b, [1.0, 10.0])
end

@testset "Univariate" begin
    # Tests with scalar-valued distributions.
    uni_dists = [
        Arcsine(2, 4),
        Beta(2,2),
        BetaPrime(),
        Biweight(),
        Cauchy(),
        Chi(3),
        Chisq(2),
        Cosine(),
        Epanechnikov(),
        Erlang(),
        Exponential(),
        FDist(1, 1),
        Frechet(),
        Gamma(),
        InverseGamma(),
        InverseGaussian(),
        # Kolmogorov(),
        Laplace(),
        Levy(),
        Logistic(),
        LogNormal(1.0, 2.5),
        Normal(0.1, 2.5),
        Pareto(),
        Rayleigh(1.0),
        TDist(2),
        truncated(Normal(0, 1), -Inf, 2),
        transformed(Beta(2,2)),
        transformed(Exponential()),
    ]

    for dist in uni_dists
        @testset "$dist: dist" begin
            td = @inferred transformed(dist)

            # single sample
            y = @inferred rand(td)
            x = @inferred inv(td.transform)(y)
            @test y ≈ @inferred td.transform(x)
            @test @inferred(logpdf(td, y)) ≈ @inferred(logpdf_with_trans(dist, x, true))

            # logpdf_with_jac
            lp, logjac = logpdf_with_jac(td, y)
            @test lp ≈ logpdf(td, y)
            @test logjac ≈ logabsdetjacinv(td.transform, y)

            # multi-sample
            y = @inferred rand(td, 10)
            x = inv(td.transform).(y)
            @test logpdf.(td, y) ≈ logpdf_with_trans.(dist, x, true)

            # logpdf corresponds to logpdf_with_trans
            d = dist
            b = @inferred bijector(d)
            x = rand(d)
            y = @inferred b(x)
            @test logpdf(d, inv(b)(y)) + logabsdetjacinv(b, y) ≈ logpdf_with_trans(d, x, true)
            @test logpdf(d, x) - logabsdetjac(b, x) ≈ logpdf_with_trans(d, x, true)

            # forward
            f = @inferred forward(td)
            @test f.x ≈ inv(td.transform)(f.y)
            @test f.y ≈ td.transform(f.x)
            @test f.logabsdetjac ≈ logabsdetjac(td.transform, f.x)
            @test f.logpdf ≈ logpdf_with_trans(td.dist, f.x, true)
            @test f.logpdf ≈ logpdf(td.dist, f.x) - f.logabsdetjac

            # verify against AD
            d = dist
            b = bijector(d)
            x = rand(d)
            y = b(x)
            # `ForwardDiff.derivative` can lead to some numerical inaccuracy,
            # so we use a slightly higher `atol` than default.
            @test log(abs(ForwardDiff.derivative(b, x))) ≈ logabsdetjac(b, x) atol=1e-6
            @test log(abs(ForwardDiff.derivative(inv(b), y))) ≈ logabsdetjac(inv(b), y) atol=1e-6
        end

        @testset "$dist: ForwardDiff AD" begin
            x = rand(dist)
            b = MyADBijector{Bijectors.ADBackend(:forwarddiff)}(dist)

            @test abs(det(Bijectors.jacobian(b, x))) > 0
            @test logabsdetjac(b, x) ≠ Inf

            y = b(x)
            b⁻¹ = inv(b)
            @test abs(det(Bijectors.jacobian(b⁻¹, y))) > 0
            @test logabsdetjac(b⁻¹, y) ≠ Inf
        end

        @testset "$dist: Tracker AD" begin
            x = rand(dist)
            b = MyADBijector{Bijectors.ADBackend(:reversediff)}(dist)

            @test abs(det(Bijectors.jacobian(b, x))) > 0
            @test logabsdetjac(b, x) ≠ Inf

            y = b(x)
            b⁻¹ = inv(b)
            @test abs(det(Bijectors.jacobian(b⁻¹, y))) > 0
            @test logabsdetjac(b⁻¹, y) ≠ Inf
        end
    end
end

@testset "Batch computation" begin
    bs_xs = [
        (Scale(2.0), randn(3)),
        (Scale([1.0, 2.0]), randn(2, 3)),
        (Shift(2.0), randn(3)),
        (Shift([1.0, 2.0]), randn(2, 3)),
        (Log{0}(), exp.(randn(3))),
        (Log{1}(), exp.(randn(2, 3))),
        (Exp{0}(), randn(3)),
        (Exp{1}(), randn(2, 3)),
        (Log{1}() ∘ Exp{1}(), randn(2, 3)),
        (inv(Logit(-1.0, 1.0)), randn(3)),
        (Identity{0}(), randn(3)),
        (Identity{1}(), randn(2, 3)),
        (PlanarLayer(2), randn(2, 3)),
        (RadialLayer(2), randn(2, 3)),
        (PlanarLayer(2) ∘ RadialLayer(2), randn(2, 3)),
        (Exp{1}() ∘ PlanarLayer(2) ∘ RadialLayer(2), randn(2, 3)),
        (SimplexBijector(), mapslices(z -> normalize(z, 1), rand(2, 3); dims = 1)),
        (stack(Exp{0}(), Scale(2.0)), randn(2, 3)),
        (Stacked((Exp{1}(), SimplexBijector()), (1:1, 2:3)),
         mapslices(z -> normalize(z, 1), rand(3, 2); dims = 1)),
        (RationalQuadraticSpline(randn(3), randn(3), randn(3 - 1), 2.), [-0.5, 0.5]),
        (LeakyReLU(0.1), randn(3)),
        (LeakyReLU(Float32(0.1)), randn(3)),
        (LeakyReLU(0.1; dim = Val(1)), randn(2, 3)),
    ]

    for (b, xs) in bs_xs
        @testset "$b" begin
            D = @inferred Bijectors.dimension(b)
            ib = @inferred inv(b)

            @test Bijectors.dimension(ib) == D

            x = D == 0 ? xs[1] : xs[:, 1]

            y = @inferred b(x)
            ys = @inferred b(xs)
            @inferred(b(param(xs)))

            x_ = @inferred ib(y)
            xs_ = @inferred ib(ys)
            @inferred(ib(param(ys)))

            result = @inferred forward(b, x)
            results = @inferred forward(b, xs)

            iresult = @inferred forward(ib, y)
            iresults = @inferred forward(ib, ys)

            # Sizes
            @test size(y) == size(x)
            @test size(ys) == size(xs)

            @test size(x_) == size(x)
            @test size(xs_) == size(xs)

            @test size(result.rv) == size(x)
            @test size(results.rv) == size(xs)

            @test size(iresult.rv) == size(y)
            @test size(iresults.rv) == size(ys)

            # Values
            @test ys ≈ hcat([b(xs[:, i]) for i = 1:size(xs, 2)]...)
            @test ys ≈ results.rv

            if D == 0
                # Sizes
                @test y == ys[1]

                @test length(logabsdetjac(b, xs)) == length(xs)
                @test length(logabsdetjac(ib, ys)) == length(xs)

                @test @inferred(logabsdetjac(b, param(xs))) isa Union{Array, TrackedArray}
                @test @inferred(logabsdetjac(ib, param(ys))) isa Union{Array, TrackedArray}

                @test size(results.logabsdetjac) == size(xs, )
                @test size(iresults.logabsdetjac) == size(ys, )

                # Values
                b_logjac_ad = [(log ∘ abs)(ForwardDiff.derivative(b, xs[i])) for i = 1:length(xs)]
                ib_logjac_ad = [(log ∘ abs)(ForwardDiff.derivative(ib, ys[i])) for i = 1:length(ys)]
                @test logabsdetjac.(b, xs) == @inferred(logabsdetjac(b, xs))
                @test @inferred(logabsdetjac(b, xs)) ≈ b_logjac_ad atol=1e-9
                @test logabsdetjac.(ib, ys) == @inferred(logabsdetjac(ib, ys))
                @test @inferred(logabsdetjac(ib, ys)) ≈ ib_logjac_ad atol=1e-9

                @test logabsdetjac.(b, param(xs)) == @inferred(logabsdetjac(b, param(xs)))
                @test logabsdetjac.(ib, param(ys)) == @inferred(logabsdetjac(ib, param(ys)))

                @test results.logabsdetjac ≈ vec(logabsdetjac.(b, xs))
                @test iresults.logabsdetjac ≈ vec(logabsdetjac.(ib, ys))
            elseif D == 1
                @test y == ys[:, 1]
                # Comparing sizes instead of lengths ensures we catch errors s.t.
                # length(x) == 3 when size(x) == (1, 3).
                # Sizes
                @test size(logabsdetjac(b, xs)) == (size(xs, 2), )
                @test size(logabsdetjac(ib, ys)) == (size(xs, 2), )

                @test @inferred(logabsdetjac(b, param(xs))) isa Union{Array, TrackedArray}
                @test @inferred(logabsdetjac(ib, param(ys))) isa Union{Array, TrackedArray}

                @test size(results.logabsdetjac) == (size(xs, 2), )
                @test size(iresults.logabsdetjac) == (size(ys, 2), )

                # Test all values
                @test @inferred(logabsdetjac(b, xs)) ≈ vec(mapslices(z -> logabsdetjac(b, z), xs; dims = 1))
                @test @inferred(logabsdetjac(ib, ys)) ≈ vec(mapslices(z -> logabsdetjac(ib, z), ys; dims = 1))

                @test results.logabsdetjac ≈ vec(mapslices(z -> logabsdetjac(b, z), xs; dims = 1))
                @test iresults.logabsdetjac ≈ vec(mapslices(z -> logabsdetjac(ib, z), ys; dims = 1))

                # FIXME: `SimplexBijector` results in ∞ gradient if not in the domain
                if !contains(t -> t isa SimplexBijector, b)
                    b_logjac_ad = [logabsdet(ForwardDiff.jacobian(b, xs[:, i]))[1] for i = 1:size(xs, 2)]
                    @test logabsdetjac(b, xs) ≈ b_logjac_ad atol=1e-9

                    ib_logjac_ad = [logabsdet(ForwardDiff.jacobian(ib, ys[:, i]))[1] for i = 1:size(ys, 2)]
                    @test logabsdetjac(ib, ys) ≈ ib_logjac_ad atol=1e-9
                end
            else
                error("tests not implemented yet")
            end
        end
    end

    @testset "Composition" begin
        @test_throws DimensionMismatch (Exp{1}() ∘ Log{0}())

        # Check that type-stable composition stays type-stable
        cb1 = Composed((Exp(), Log())) ∘ Exp()
        @test cb1 isa Composed{<:Tuple}
        cb2 = Exp() ∘ Composed((Exp(), Log()))
        @test cb2 isa Composed{<:Tuple}
        cb3 = cb1 ∘ cb2
        @test cb3 isa Composed{<:Tuple}

        @test logabsdetjac(cb1, 1.) isa Real
        @test logabsdetjac(cb1, 1.) == 1.

        @test inv(cb1) isa Composed{<:Tuple}
        @test inv(cb2) isa Composed{<:Tuple}
        @test inv(cb3) isa Composed{<:Tuple}

        # Check that type-unstable composition stays type-unstable
        cb1 = Composed([Exp(), Log()]) ∘ Exp()
        @test cb1 isa Composed{<:AbstractArray}
        cb2 = Exp() ∘ Composed([Exp(), Log()])
        @test cb2 isa Composed{<:AbstractArray}
        cb3 = cb1 ∘ cb2
        @test cb3 isa Composed{<:AbstractArray}

        @test logabsdetjac(cb1, 1.) isa Real
        @test logabsdetjac(cb1, 1.) == 1.

        @test inv(cb1) isa Composed{<:AbstractArray}
        @test inv(cb2) isa Composed{<:AbstractArray}
        @test inv(cb3) isa Composed{<:AbstractArray}

        # combining the two
        @test_throws ErrorException (Log() ∘ Exp()) ∘ cb1
        @test_throws ErrorException cb1 ∘ (Log() ∘ Exp())
    end

    @testset "Batch-computation with Tracker.jl" begin
        @testset "Scale" begin
            # 0-dim with `Real` parameter
            b = Scale(param(2.0))
            lj = logabsdetjac(b, 1.0)
            Tracker.back!(lj, 1.0)
            @test Tracker.extract_grad!(b.a) == 0.5

            # 0-dim with `Real` parameter for batch-computation
            lj = logabsdetjac(b, [1.0, 2.0, 3.0])
            Tracker.back!(lj, [1.0, 1.0, 1.0])
            @test Tracker.extract_grad!(b.a) == sum([0.5, 0.5, 0.5])

            # 1-dim with `Vector` parameter
            x = [3.0, 4.0, 5.0]
            xs = [3.0 4.0; 4.0 7.0; 5.0 8.0]
            a = [2.0, 3.0, 5.0]

            b = Scale(param(a))
            lj = logabsdetjac(b, x)
            Tracker.back!(lj)
            @test Tracker.extract_grad!(b.a) == ForwardDiff.gradient(a -> logabsdetjac(Scale(a), x), a)

            # batch
            lj = logabsdetjac(b, xs)
            Tracker.back!(mean(lj), 1.0)
            @test Tracker.extract_grad!(b.a) == ForwardDiff.gradient(a -> mean(logabsdetjac(Scale(a), xs)), a)

            # Forward when doing a composition
            y, logjac = logabsdetjac(b, xs)
            Tracker.back!(mean(logjac), 1.0)
            @test Tracker.extract_grad!(b.a) == ForwardDiff.gradient(a -> mean(logabsdetjac(Scale(a), xs)), a)
        end

        @testset "Shift" begin
            b = Shift(param(1.0))
            lj = logabsdetjac(b, 1.0)
            Tracker.back!(lj, 1.0)
            @test Tracker.extract_grad!(b.a) == 0.0

            # 0-dim with `Real` parameter for batch-computation
            lj = logabsdetjac(b, [1.0, 2.0, 3.0])
            @test lj isa TrackedArray
            Tracker.back!(lj, [1.0, 1.0, 1.0])
            @test Tracker.extract_grad!(b.a) == 0.0

            # 1-dim with `Vector` parameter
            b = Shift(param([2.0, 3.0, 5.0]))
            lj = logabsdetjac(b, [3.0, 4.0, 5.0])
            Tracker.back!(lj)
            @test Tracker.extract_grad!(b.a) == zeros(3)

            lj = logabsdetjac(b, [3.0 4.0 5.0; 6.0 7.0 8.0])
            @test lj isa TrackedArray
            Tracker.back!(lj, [1.0, 1.0, 1.0])
            @test Tracker.extract_grad!(b.a) == zeros(3)
        end
    end
end

@testset "Truncated" begin
    d = truncated(Normal(), -1, 1)
    b = bijector(d)
    x = rand(d)
    y = b(x)
    @test y ≈ link(d, x)
    @test inv(b)(y) ≈ x
    @test logabsdetjac(b, x) ≈ logpdf_with_trans(d, x, false) - logpdf_with_trans(d, x, true)

    d = truncated(Normal(), -Inf, 1)
    b = bijector(d)
    x = rand(d)
    y = b(x)
    @test y ≈ link(d, x)
    @test inv(b)(y) ≈ x
    @test logabsdetjac(b, x) ≈ logpdf_with_trans(d, x, false) - logpdf_with_trans(d, x, true)

    d = truncated(Normal(), 1, Inf)
    b = bijector(d)
    x = rand(d)
    y = b(x)
    @test y ≈ link(d, x)
    @test inv(b)(y) ≈ x
    @test logabsdetjac(b, x) ≈ logpdf_with_trans(d, x, false) - logpdf_with_trans(d, x, true)
end

@testset "Multivariate" begin
    vector_dists = [
        Dirichlet(2, 3),
        Dirichlet([1000 * one(Float64), eps(Float64)]),
        Dirichlet([eps(Float64), 1000 * one(Float64)]),
        MvNormal(randn(10), Diagonal(exp.(randn(10)))),
        MvLogNormal(MvNormal(randn(10), Diagonal(exp.(randn(10))))),
        Dirichlet([1000 * one(Float64), eps(Float64)]),
        Dirichlet([eps(Float64), 1000 * one(Float64)]),
        transformed(MvNormal(randn(10), Diagonal(exp.(randn(10))))),
        transformed(MvLogNormal(MvNormal(randn(10), Diagonal(exp.(randn(10))))))
    ]

    for dist in vector_dists
        @testset "$dist: dist" begin
            td = transformed(dist)

            # single sample
            y = rand(td)
            x = inv(td.transform)(y)
            @test inv(td.transform)(param(y)) isa TrackedArray
            @test y ≈ td.transform(x)
            @test td.transform(param(x)) isa TrackedArray
            @test logpdf(td, y) ≈ logpdf_with_trans(dist, x, true)

            # logpdf_with_jac
            lp, logjac = logpdf_with_jac(td, y)
            @test lp ≈ logpdf(td, y)
            @test logjac ≈ logabsdetjacinv(td.transform, y)

            # multi-sample
            y = rand(td, 10)
            x = inv(td.transform)(y)
            @test inv(td.transform)(param(y)) isa TrackedArray
            @test logpdf(td, y) ≈ logpdf_with_trans(dist, x, true)

            # forward
            f = forward(td)
            @test f.x ≈ inv(td.transform)(f.y)
            @test f.y ≈ td.transform(f.x)
            @test f.logabsdetjac ≈ logabsdetjac(td.transform, f.x)
            @test f.logpdf ≈ logpdf_with_trans(td.dist, f.x, true)

            # verify against AD
            # similar to what we do in test/transform.jl for Dirichlet
            if dist isa Dirichlet
                b = Bijectors.SimplexBijector{1, false}()
                x = rand(dist)
                y = b(x)
                @test b(param(x)) isa TrackedArray
                @test log(abs(det(ForwardDiff.jacobian(b, x)))) ≈ logabsdetjac(b, x)
                @test log(abs(det(ForwardDiff.jacobian(inv(b), y)))) ≈ logabsdetjac(inv(b), y)
            else
                b = bijector(dist)
                x = rand(dist)
                y = b(x)
                # `ForwardDiff.derivative` can lead to some numerical inaccuracy,
                # so we use a slightly higher `atol` than default.
                @test b(param(x)) isa TrackedArray
                @test log(abs(det(ForwardDiff.jacobian(b, x)))) ≈ logabsdetjac(b, x) atol=1e-6
                @test log(abs(det(ForwardDiff.jacobian(inv(b), y)))) ≈ logabsdetjac(inv(b), y) atol=1e-6
            end
        end
    end
end

@testset "Matrix variate" begin
    v = 7.0
    S = Matrix(1.0I, 2, 2)
    S[1, 2] = S[2, 1] = 0.5

    matrix_dists = [
        Wishart(v,S),
        InverseWishart(v,S),
        TuringWishart(v,S),
        TuringInverseWishart(v,S),
        LKJ(3, 1.)
    ]

    for dist in matrix_dists
        @testset "$dist: dist" begin
            td = transformed(dist)

            # single sample
            y = rand(td)
            x = inv(td.transform)(y)
            @test inv(td.transform)(param(y)) isa TrackedArray
            @test logpdf(td, y) ≈ logpdf_with_trans(dist, x, true)

            # TODO: implement `logabsdetjac` for these
            # logpdf_with_jac
            # lp, logjac = logpdf_with_jac(td, y)
            # @test lp ≈ logpdf(td, y)
            # @test logjac ≈ logabsdetjacinv(td.transform, y)

            # multi-sample
            y = rand(td, 10)
            x = inv(td.transform)(y)
            @test inv(td.transform)(param.(y)) isa Vector{<:TrackedArray}
            @test logpdf(td, y) ≈ logpdf_with_trans(dist, x, true)
        end
    end
end

@testset "Composition <: Bijector" begin
    d = Beta()
    td = transformed(d)

    x = rand(d)
    y = td.transform(x)

    b = @inferred Bijectors.composel(td.transform, Bijectors.Identity{0}())
    ib = @inferred inv(b)

    @test forward(b, x) == forward(td.transform, x)
    @test forward(ib, y) == forward(inv(td.transform), y)

    @test forward(b, x) == forward(Bijectors.composer(b.ts...), x)

    # inverse works fine for composition
    cb = @inferred b ∘ ib
    @test cb(x) ≈ x

    cb2 = @inferred cb ∘ cb
    @test cb(x) ≈ x

    # ensures that the `logabsdetjac` is correct
    x = rand(d)
    b = inv(bijector(d))
    @test logabsdetjac(b ∘ b, x) ≈ logabsdetjac(b, b(x)) + logabsdetjac(b, x)

    # order of composed evaluation
    b1 = MyADBijector(d)
    b2 = MyADBijector(Gamma())

    cb = inv(b1) ∘ b2
    @test cb(x) ≈ inv(b1)(b2(x))

    # contrived example
    b = bijector(d)
    cb = @inferred inv(b) ∘ b
    cb = @inferred cb ∘ cb
    @test @inferred(cb ∘ cb ∘ cb ∘ cb ∘ cb)(x) ≈ x

    # forward for tuple and array
    d = Beta()
    b = @inferred inv(bijector(d))
    b⁻¹ = @inferred inv(b)
    x = rand(d)

    cb_t = b⁻¹ ∘ b⁻¹
    f_t = forward(cb_t, x)

    cb_a = Composed([b⁻¹, b⁻¹])
    f_a = forward(cb_a, x)

    @test f_t == f_a

    # `composer` and `composel`
    cb_l = Bijectors.composel(b⁻¹, b⁻¹, b)
    cb_r = Bijectors.composer(reverse(cb_l.ts)...)
    y = cb_l(x)
    @test y == Bijectors.composel(cb_r.ts...)(x)

    k = length(cb_l.ts)
    @test all([cb_l.ts[i] == cb_r.ts[i] for i = 1:k])
end

@testset "Stacked <: Bijector" begin
    # `logabsdetjac` withOUT AD
    d = Beta()
    b = bijector(d)
    x = rand(d)
    y = b(x)

    sb1 = @inferred stack(b, b, inv(b), inv(b))             # <= Tuple
    res1 = forward(sb1, [x, x, y, y])
    @test sb1(param([x, x, y, y])) isa TrackedArray

    @test sb1([x, x, y, y]) ≈ res1.rv
    @test logabsdetjac(sb1, [x, x, y, y]) ≈ 0 atol=1e-6
    @test res1.logabsdetjac ≈ 0 atol=1e-6

    sb2 = Stacked([b, b, inv(b), inv(b)])        # <= Array
    res2 = forward(sb2, [x, x, y, y])
    @test sb2(param([x, x, y, y])) isa TrackedArray

    @test sb2([x, x, y, y]) ≈ res2.rv
    @test logabsdetjac(sb2, [x, x, y, y]) ≈ 0.0 atol=1e-12
    @test res2.logabsdetjac ≈ 0.0 atol=1e-12

    # `logabsdetjac` with AD
    b = MyADBijector(d)
    y = b(x)

    sb1 = stack(b, b, inv(b), inv(b))             # <= Tuple
    res1 = forward(sb1, [x, x, y, y])
    @test sb1(param([x, x, y, y])) isa TrackedArray

    @test sb1([x, x, y, y]) == res1.rv
    @test logabsdetjac(sb1, [x, x, y, y]) ≈ 0 atol=1e-12
    @test res1.logabsdetjac ≈ 0.0 atol=1e-12

    sb2 = Stacked([b, b, inv(b), inv(b)])        # <= Array
    res2 = forward(sb2, [x, x, y, y])
    @test sb2(param([x, x, y, y])) isa TrackedArray

    @test sb2([x, x, y, y]) == res2.rv
    @test logabsdetjac(sb2, [x, x, y, y]) ≈ 0.0 atol=1e-12
    @test res2.logabsdetjac ≈ 0.0 atol=1e-12

    # value-test
    x = ones(3)
    sb = @inferred stack(Bijectors.Exp(), Bijectors.Log(), Bijectors.Shift(5.0))
    res = forward(sb, x)
    @test sb(param(x)) isa TrackedArray
    @test sb(x) == [exp(x[1]), log(x[2]), x[3] + 5.0]
    @test res.rv == [exp(x[1]), log(x[2]), x[3] + 5.0]
    @test logabsdetjac(sb, x) == sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges[i]])) for i = 1:3])
    @test res.logabsdetjac == logabsdetjac(sb, x)

    # TODO: change when we have dimensionality in the type
    sb = @inferred Stacked((Bijectors.Exp(), Bijectors.SimplexBijector()), (1:1, 2:3))
    x = ones(3) ./ 3.0
    res = @inferred forward(sb, x)
    @test sb(param(x)) isa TrackedArray
    @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...]
    @test res.rv == [exp(x[1]), sb.bs[2](x[2:3])...]
    @test logabsdetjac(sb, x) == sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges[i]])) for i = 1:2])
    @test res.logabsdetjac == logabsdetjac(sb, x)

    x = ones(4) ./ 4.0
    @test_throws AssertionError sb(x)

    # Array-version
    sb = Stacked([Bijectors.Exp(), Bijectors.SimplexBijector()], [1:1, 2:3])
    x = ones(3) ./ 3.0
    res = forward(sb, x)
    @test sb(param(x)) isa TrackedArray
    @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...]
    @test res.rv == [exp(x[1]), sb.bs[2](x[2:3])...]
    @test logabsdetjac(sb, x) == sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges[i]])) for i = 1:2])
    @test res.logabsdetjac == logabsdetjac(sb, x)

    x = ones(4) ./ 4.0
    @test_throws AssertionError sb(x)

    # Mixed versions
    # Tuple, Array
    sb = Stacked([Bijectors.Exp(), Bijectors.SimplexBijector()], (1:1, 2:3))
    x = ones(3) ./ 3.0
    res = forward(sb, x)
    @test sb(param(x)) isa TrackedArray
    @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...]
    @test res.rv == [exp(x[1]), sb.bs[2](x[2:3])...]
    @test logabsdetjac(sb, x) == sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges[i]])) for i = 1:2])
    @test res.logabsdetjac == logabsdetjac(sb, x)

    x = ones(4) ./ 4.0
    @test_throws AssertionError sb(x)

    # Array, Tuple
    sb = Stacked((Bijectors.Exp(), Bijectors.SimplexBijector()), [1:1, 2:3])
    x = ones(3) ./ 3.0
    res = forward(sb, x)
    @test sb(param(x)) isa TrackedArray
    @test sb(x) == [exp(x[1]), sb.bs[2](x[2:3])...]
    @test res.rv == [exp(x[1]), sb.bs[2](x[2:3])...]
    @test logabsdetjac(sb, x) == sum([sum(logabsdetjac(sb.bs[i], x[sb.ranges[i]])) for i = 1:2])
    @test res.logabsdetjac == logabsdetjac(sb, x)

    x = ones(4) ./ 4.0
    @test_throws AssertionError sb(x)

    @testset "Stacked: ADVI with MvNormal" begin
        # MvNormal test
        dists = [
            Beta(),
            Beta(),
            Beta(),
            InverseGamma(),
            InverseGamma(),
            Gamma(),
            Gamma(),
            InverseGamma(),
            Cauchy(),
            Gamma(),
            MvNormal(zeros(2), I)
        ]

        ranges = []
        idx = 1
        for i = 1:length(dists)
            d = dists[i]
            push!(ranges, idx:idx + length(d) - 1)
            idx += length(d)
        end
        ranges = tuple(ranges...)

        num_params = ranges[end][end]
        d = MvNormal(zeros(num_params), I)

        # Stacked{<:Array}
        bs = bijector.(dists)     # constrained-to-unconstrained bijectors for dists
        ibs = inv.(bs)            # invert, so we get unconstrained-to-constrained
        sb = Stacked(ibs, ranges) # => Stacked <: Bijector
        x = rand(d)

        @test sb isa Stacked

        td = transformed(d, sb)  # => MultivariateTransformed <: Distribution{Multivariate, Continuous}
        @test td isa Distribution{Multivariate, Continuous}

        # check that wrong ranges fails
        @test_throws MethodError stack(ibs...)
        sb = Stacked(ibs)
        x = rand(d)
        @test_throws AssertionError sb(x)

        # Stacked{<:Tuple}
        bs = bijector.(tuple(dists...))
        ibs = inv.(bs)
        sb = @inferred Stacked(ibs, ranges)
        isb = @inferred inv(sb)
        @test sb isa Stacked{<:Tuple}

        # inverse
        td = @inferred transformed(d, sb)
        y = @inferred rand(td)
        x = @inferred isb(y)
        @test sb(x) ≈ y

        # verification of computation
        x = rand(d)
        y = sb(x)
        y_ = vcat([ibs[i](x[ranges[i]]) for i = 1:length(dists)]...)
        x_ = vcat([bs[i](y[ranges[i]]) for i = 1:length(dists)]...)
        @test x ≈ x_
        @test y ≈ y_

        # AD verification
        @test log(abs(det(ForwardDiff.jacobian(sb, x)))) ≈ logabsdetjac(sb, x)
        @test log(abs(det(ForwardDiff.jacobian(isb, y)))) ≈ logabsdetjac(isb, y)

        # Ensure `Stacked` works for a single bijector
        d = (MvNormal(zeros(2), I),)
        sb = Stacked(bijector.(d), (1:2, ))
        x = [.5, 1.]
        @test sb(x) == x
        @test logabsdetjac(sb, x) == 0
        @test forward(sb, x) == (rv = x, logabsdetjac = zero(eltype(x)))
    end
end

@testset "Example: ADVI single" begin
    # Usage in ADVI
    d = Beta()
    b = bijector(d)                # [0, 1] → ℝ
    ib = inv(b)                    # ℝ → [0, 1]
    td = transformed(Normal(), ib) # x ∼ 𝓝(0, 1) then f(x) ∈ [0, 1]
    x = rand(td)                   # ∈ [0, 1]
    @test 0 ≤ x ≤ 1
end

@testset "Jacobians of SimplexBijector" begin
    b = SimplexBijector()
    ib = inv(b)

    x = ib(randn(10))
    y = b(x)

    @test Bijectors.jacobian(b, x) ≈ ForwardDiff.jacobian(b, x)
    @test Bijectors.jacobian(ib, y) ≈ ForwardDiff.jacobian(ib, y)

    # Just some additional computation so we also ensure the pullbacks are the same
    weights = randn(10)

    # Tracker.jl
    x_tracked = Tracker.param(x)
    z = sum(weights .* b(x_tracked))
    Tracker.back!(z)
    Δ_tracker = Tracker.grad(x_tracked)

    # ForwardDiff.jl
    Δ_forwarddiff = ForwardDiff.gradient(z -> sum(weights .* b(z)), x)

    # Compare
    @test Δ_forwarddiff ≈ Δ_tracker

    # Tracker.jl
    y_tracked = Tracker.param(y)
    z = sum(weights .* ib(y_tracked))
    Tracker.back!(z)
    Δ_tracker = Tracker.grad(y_tracked)

    # ForwardDiff.jl
    Δ_forwarddiff = ForwardDiff.gradient(z -> sum(weights .* ib(z)), y)

    @test Δ_forwarddiff ≈ Δ_tracker
end

@testset "Equality" begin
    bs = [
        Identity{0}(),
        Identity{1}(),
        Identity{2}(),
        Exp{0}(),
        Exp{1}(),
        Exp{2}(),
        Log{0}(),
        Log{1}(),
        Log{2}(),
        Scale(2.0),
        Scale(3.0),
        Scale(rand(2,2)),
        Scale(rand(2,2)),
        Shift(2.0),
        Shift(3.0),
        Shift(rand(2)),
        Shift(rand(2)),
        Logit(1.0, 2.0),
        Logit(1.0, 3.0),
        Logit(2.0, 3.0),
        Logit(0.0, 2.0),
        InvertibleBatchNorm(2),
        InvertibleBatchNorm(3),
        PDBijector(),
        Permute([1.0, 2.0, 3.0]),
        Permute([2.0, 3.0, 4.0]),
        PlanarLayer(2),
        PlanarLayer(3),
        RadialLayer(2),
        RadialLayer(3),
        SimplexBijector(),
        Stacked((Exp{0}(), Log{0}())),
        Stacked((Log{0}(), Exp{0}())),
        Stacked([Exp{0}(), Log{0}()]),
        Stacked([Log{0}(), Exp{0}()]),
        Composed((Exp{0}(), Log{0}())),
        Composed((Log{0}(), Exp{0}())),
        # Composed([Exp{0}(), Log{0}()]),
        # Composed([Log{0}(), Exp{0}()]),
        TruncatedBijector(1.0, 2.0),
        TruncatedBijector(1.0, 3.0),
        TruncatedBijector(0.0, 2.0),
    ]
    for i in 1:length(bs), j in 1:length(bs)
        if i == j
            @test bs[i] == deepcopy(bs[j])
            @test inv(bs[i]) == inv(deepcopy(bs[j]))
        else
            @test bs[i] != bs[j]
        end
    end
end


text:
using AQFED.Math
using FastGaussQuadrature
export DeelstraBasketPricer, DeelstraLBBasketPricer, priceEuropean
using LinearAlgebra
using Roots

struct DeelstraBasketPricer
    δIndex::Int
    fIndex::Int
    x::Vector{Float64} #quadrature abscissae in -1, 1
    w::Vector{Float64} #quadrature weights
end

function DeelstraBasketPricer(δIndex::Int, fIndex::Int; N=33*4)
    xG, wG = gausslegendre(N)
    DeelstraBasketPricer(δIndex, fIndex, xG, wG)
end

function DeelstraLBBasketPricer(δIndex::Int, fIndex::Int)
    DeelstraBasketPricer(δIndex, fIndex, Vector{Float64}(undef, 0), Vector{Float64}(undef, 0))
end

DeelstraBasketPricer() = DeelstraBasketPricer(3,3)
DeelstraLBBasketPricer() = DeelstraLBBasketPricer(3,3)

function priceEuropean(
    p::DeelstraBasketPricer,
    isCall::Bool,
    strike::Number,
    discountFactor::Number, #discount factor to payment
    spot::AbstractArray{<:Number},
    forward::AbstractArray{TV}, #forward to option maturity
    totalVariance::AbstractArray{<:Number}, #vol^2 * τ
    weight::AbstractArray{<:Number},
    correlation::Matrix{TV},
)::Number where {TV <: Number}
    n = length(spot)
    δ = zeros(TV, n) #we store δ * S here.
    f = zero(eltype(forward))
    u1 = zero(eltype(forward))
    for (i, wi) in enumerate(weight)
        if p.δIndex == 1
            δ[i] = forward[i] * exp(-totalVariance[i] / 2)
        elseif p.δIndex == 2
            δ[i] = spot[i]
        else #p.δIndex == 3
            δ[i] = forward[i]
        end
        f += wi * δ[i]
        u1 += wi * forward[i]
    end
    sumwbd = zero(TV)
    wsdv = zeros(TV, n)

    for (i, wi) in enumerate(weight)
        wTildei = wi * δ[i] / f
        βi = log(forward[i]) - totalVariance[i] / 2
        sumwbd += wTildei * (βi - log(δ[i]))
        if totalVariance[i] > 0
            wsdv[i] = wi * δ[i] * sqrt(totalVariance[i])
        end
    end
    dGamma = f * (log(strike / f) - sumwbd)

    varGamma = zero(TV)
    r = zeros(TV, n)
    for (i, wsdi) in enumerate(wsdv)
        for (j, wsdj) in enumerate(wsdv)
            covar = wsdj * correlation[i, j]
            varGamma += covar * wsdi
            r[i] += covar
        end
    end
    varGamma = abs(varGamma)
    volGamma = sqrt(varGamma)
    r ./= volGamma
    priceCall = zero(Float64)
    if length(p.x) == 0
        function objective(λ::Number)::Tuple{<:Number,<:Number,<:Number}
            #fΔxΔΔx
            eS = zero(typeof(λ))
            DeS = zero(typeof(λ))
            D2eS = zero(typeof(λ))
            for (i, fi) in enumerate(forward)
                viT = totalVariance[i]
                eSi = weight[i] * fi
                if viT > 0
                    sqrti = sqrt(viT)
                    eSi *= exp(-r[i]^2 * viT / 2 + r[i] * sqrti  * (λ) / (volGamma))
                    DeS += eSi * (r[i] * sqrti)/volGamma
                    D2eS += eSi * ((r[i] * sqrti)/volGamma)^2
                end
                eS += eSi
            end
            # println(λ," ", eS-strike," ", DeS, " ",D2eS)
            return (eS - strike, (eS-strike)/DeS, DeS/D2eS)
        end
        lambdaMax = dGamma * 2
        if lambdaMax < zero(Float64)
            lambdaMax = dGamma / 2
        end
        #init guess dGamma
        lambdaOpt = find_zero(objective, dGamma, Roots.Halley(), atol=1e-8)
        dGammaStar = (lambdaOpt) / volGamma
        #println("dG ",dGammaStar)
        i1 = zero(Float64)
        for (i, fi) in enumerate(forward)
            if totalVariance[i] > 0
                i1 += weight[i] * fi * normcdf(r[i] * sqrt(totalVariance[i]) - dGammaStar)
            else
                i1 += weight[i] * fi * normcdf(- dGammaStar)
            end
        end
        i1 -= strike * normcdf(-dGammaStar)
        priceCall = i1
    else
        function integrand(λ)
            #depends on forward, weight, tvar, r, volGamma,correl, f, sumwbd, strike
            eS = zero(Float64)
            eS2 = zero(Float64)
            for (i, fi) in enumerate(forward)
                tempi = weight[i] * fi
                eSi = tempi
                viT = totalVariance[i]
                sqrti = zero(Float64)
                if viT > 0
                    sqrti = sqrt(viT)
                    eSi *= exp(-r[i]^2 * viT / 2 + r[i] * sqrti * (λ) / (volGamma))
                end
                eS += eSi
                for j = 1:i-1
                    vjT = totalVariance[j]
                    sqrtj = zero(Float64)
                    if vjT > 0
                        sqrtj = sqrt(vjT)
                    end
                    sigmaij = sqrt(viT + vjT + 2 * sqrti * sqrtj * correlation[i, j])
                    rij = (r[i] * sqrti + r[j] * sqrtj) / sigmaij
                    eFactor = -viT / 2 - vjT / 2 + (1 - rij^2) / 2 * sigmaij^2 + rij * sigmaij * (λ) / (volGamma)
                    eS2 += 2 * tempi * weight[j] * forward[j] * exp(eFactor)
                end
                eS2 += tempi^2 * exp(-viT + 2 * (1 - r[i]^2) * viT + r[i] * 2 * sqrti * (λ) / (volGamma))
            end

            fS = zero(Float64)
            if p.fIndex == 1
                #zero
            elseif p.fIndex == 2
                lnGF = λ / f + sumwbd
                fS = f * (1 + lnGF)
            else #p.fIndex >= 3
                #most accurate according to Deelstra paper Fig.2
                lnGF = λ / f + sumwbd
                gF = exp(lnGF)
                fS = f * gF
            end
            euphs2 = eS - fS
            if euphs2 < eps(Float64)
                fS = 0
                euphs2 = eS
            end
            e2ups2 = eS2 - 2 * fS * eS + fS^2

            if e2ups2 < eps(Float64)
                e2ups2 = eps(Float64)
            end
            uphs2 = log(euphs2)
            ups2 = log(e2ups2) / 2
            muS = (uphs2 - ups2 / 2) * 2
            sigmaS2 = ups2 - muS
            if sigmaS2 < eps(Float64)
                sigmaS2 = eps(Float64)
            end
            varLambda = varGamma
            pdf = exp(-λ^2 / (2 * varLambda)) / sqrt(Float64(pi) * 2 * varLambda)

            if strike <= fS
                return euphs2 * pdf
            end
            d1 = (ups2 - log(strike - fS)) / sqrt(sigmaS2)
            d2 = d1 - sqrt(sigmaS2)
            return (euphs2 * normcdf(d1) - (strike - fS) * normcdf(d2)) * pdf
        end
        λMin = -3 * sqrt(varGamma)
        i2 = zero(Float64)
        if dGamma > λMin
            a = λMin
            b = dGamma
            bma2 = (b - a) / 2
            bpa2 = (b + a) / 2
            i2 = bma2 * dot(p.w, integrand.(bma2 .* p.x .+ bpa2))
        end
        dGammaStar = (dGamma) / sqrt(varGamma)

        i1 = zero(Float64)
        for (i, fi) in enumerate(forward)
            if totalVariance[i] > 0
                i1 += weight[i] * fi * normcdf(r[i] * sqrt(totalVariance[i]) - dGammaStar)
            else
                i1 += weight[i] * fi * normcdf(- dGammaStar)
            end
        end
        i1 -= strike * normcdf(-dGammaStar)
        priceCall = (i1 + i2)
    end

    pricePut = strike - u1 + priceCall
    if isCall
        return priceCall * discountFactor
    else
        return pricePut * discountFactor
    end
end