text:
__precompile__()

module MixedModelsBLB

using MathProgBase
using Reexport
using Distributions
using Random
using StatsModels
using StatsBase
using DataFrames
using InteractiveUtils
using Permutations
using LinearAlgebra
using Tables
using TableOperations
using Distributed
using Query
using DBInterface
using WiSER

using LinearAlgebra: BlasReal, copytri!
import LinearAlgebra: BlasFloat, checksquare

@reexport using Ipopt
@reexport using NLopt

export blblmmObs, blblmmModel
export update_w!, init_ls!, fit!, loglikelihood! # lmm.jl
export SubsetEstimates, blbEstimates, save_bootstrap_result!, blb_one_subset, blb_full_data, blb_db # blb.jl
export confint, fixef, vc, coeftable # blb.jl
export NonparametricBootSimulator, ParametricBootSimulator, simulate! # simulate.jl
export ◺ # multivariate_calculus.jl

"""
    blblmmObs

BLB linear mixed model observation type. Contains data from a single cluster, working arrays and so on.
"""
struct blblmmObs{T <: LinearAlgebra.BlasReal}
    # data
    y::Vector{T}
    X::Matrix{T} # X should include a column of 1's
    Z::Matrix{T}
    n::Int
    # grad and hess
    ∇β::Vector{T}   # gradient wrt β
    ∇σ²::Vector{T}   # gradient wrt σ²
    ∇L::Matrix{T}   # gradient wrt L
    Hββ::Matrix{T}   # Hessian wrt β
    Hσ²σ²::Vector{T}   # Hessian wrt σ²
    Hσ²L::Vector{T}   # Hessian cross term
    HLL::Matrix{T}   # Hessian wrt L
    # pre-compute
    yty::Vector{T}
    xty::Vector{T}
    zty::Vector{T}
    xtx::Matrix{T}  # Xi'Xi (p-by-p)
    ztx::Matrix{T}  # Zi'Xi (q-by-p)
    ztz::Matrix{T}  # Zi'Zi (q-by-q)
    # working arrays
    # obj::Vector{T}
    xtr::Vector{T}
    ztr::Vector{T}
    storage_p::Vector{T}
    storage_q_1::Vector{T}
    storage_q_2::Vector{T}
    storage_qq_1::Matrix{T}
    storage_qq_2::Matrix{T}
    storage_qq_3::Matrix{T}
    storage_qp::Matrix{T}
end

"""
    blblmmObs(y, X, Z)

The constructor for the blblmmObs type.

# Positional arguments
- `y`: response vector
- `X`: design matrix for fixed effects
- `Z`: design matrix for random effects.
"""
function blblmmObs(
    y::Vector{T},
    X::Matrix{T},
    Z::Matrix{T}
    ) where T <: BlasReal
    n, p = size(X)
    q = size(Z, 2)
    q◺ = ◺(q)
    @assert length(y) == n "length(y) should be equal to size(X, 1)"
    ∇β     = Vector{T}(undef, p)
    ∇σ²    = Vector{T}(undef, 1)
    ∇L     = Matrix{T}(undef, q, q)
    Hββ    = Matrix{T}(undef, p, p)
    Hσ²σ²  = Vector{T}(undef, 1)
    Hσ²L   = Vector{T}(undef, q◺)
    HLL    = Matrix{T}(undef, q◺, q◺)
    # obj = Vector{T}(undef, 1)
    yty = [dot(y, y)]
    xty = transpose(X) * y
    zty = transpose(Z) * y
    xtx = transpose(X) * X
    ztx = transpose(Z) * X
    ztz = transpose(Z) * Z
    xtr = Vector{T}(undef, p)
    ztr = Vector{T}(undef, q)
    storage_p = Vector{T}(undef, p)
    storage_q_1 = Vector{T}(undef, q)
    storage_q_2 = Vector{T}(undef, q)
    storage_qq_1 = Matrix{T}(undef, q, q)
    storage_qq_2 = Matrix{T}(undef, q, q)
    storage_qq_3 = Matrix{T}(undef, q, q)
    storage_qp = Matrix{T}(undef, q, p)
    blblmmObs{T}(
        y, X, Z,
        n,
        ∇β, ∇σ², ∇L,
        Hββ, Hσ²σ², Hσ²L, HLL,
        yty, xty, zty, xtx, ztx, ztz,
        # obj,
        xtr, ztr, storage_p, storage_q_1, storage_q_2,
        storage_qq_1, storage_qq_2, storage_qq_3,
        storage_qp
    )
end

"""
    blblmmModel

BLB linear mixed model type, which contains a vector of `blblmmObs` as data, model parameters, and working arrays.
"""
struct blblmmModel{T <: BlasReal} <: MathProgBase.AbstractNLPEvaluator
    # data
    data::Vector{blblmmObs{T}}
    fenames::Vector{String} # a vector of the fixed effect variable names
    renames::Vector{String} # a vector of the random effect variable names
    #N::Int    # total number of unique IDs (individuals) in the full data set
    #b::Int    # total number of unique IDs (individuals) in the subset
    p::Int            # number of fixed effect parameters
    q::Int            # number of random effect parameters
    q◺::Int           # number of parameters in the cholesky factor of Σ
    w::Vector{Int}      # a vector of weights from bootstraping the subset
    # model parameters
    β::Vector{T}     # fixed effects
    σ²::Vector{T}    # error variance
    Σ::Matrix{T}     # covariance of random effects
    ΣL::Matrix{T}    # lower cholesky factor of Σ
    # grad and hess
    ∇β::Vector{T}
    ∇σ²::Vector{T}
    ∇L::Matrix{T}
    Hββ::Matrix{T}
    Hσ²σ²::Vector{T}
    Hσ²L::Vector{T}
    HLL::Matrix{T}
    # storage
    xtx::Matrix{T}
    xty::Vector{T}
    ztz2::Matrix{T}
    ztr2::Vector{T}
    # the diag indices of L
    diagidx::Vector{Int64}
    # whether to use multi-threading in evaluating the loglikelihood
    # use_threads::Bool
end

"""
    blblmmModel(obsvec)

The constructor for  the blblmmModel type.

# Positional arguments
- `obsvec`: a vector of type blblmmObs
"""
function blblmmModel(
    obsvec::Vector{blblmmObs{T}},
    fenames::Vector{String},
    renames::Vector{String},
    N::Int64
    # use_threads::Bool
    ) where T <: BlasReal
    # T = eltype(obsvec[1].X)
    b, p, q = length(obsvec), size(obsvec[1].X, 2), size(obsvec[1].Z, 2)
    q◺ = ◺(q)
    npar = p + 1 + (q * (q + 1)) >> 1
    # since X includes a column of 1, p is the number of mean parameters
    # the cholesky factor for the qxq random effect mx has (q * (q + 1))/2 values
    # the arithmetic shift right operation has the effect of division by 2^n, here n = 1
    # then there is the error variance
    w      = ones(Int64, b) # initialize weights to be 1
    β      = Vector{T}(undef, p)
    σ²     = Vector{T}(undef, 1)
    Σ      = Matrix{T}(undef, q, q)
    ΣL     = Matrix{T}(undef, q, q)
    ∇β     = Vector{T}(undef, p)
    ∇σ²    = Vector{T}(undef, 1)
    ∇L     = Matrix{T}(undef, q, q)
    Hββ    = Matrix{T}(undef, p, p)
    Hσ²σ²  = Vector{T}(undef, 1)
    Hσ²L   = Vector{T}(undef, q◺)
    HLL    = Matrix{T}(undef, q◺, q◺)
    xtx    = Matrix{T}(undef, p, p)
    xty    = Vector{T}(undef, p)
    ztz2    = Matrix{T}(undef, abs2(q), abs2(q))
    ztr2    = Vector{T}(undef, abs2(q))
    diagidx = diag_idx(q)
    # ntotal = 0
    blblmmModel{T}(
        obsvec, fenames, renames,
        # N, b, p, q, q◺, w,
        p, q, q◺, w,
        β, σ², Σ, ΣL,
        ∇β, ∇σ², ∇L, Hββ, Hσ²σ², Hσ²L, HLL,
        xtx, xty, ztz2, ztr2, diagidx
        # use_threads
    )
end

include("lmm.jl")
include("blb.jl")
include("simulate.jl")
include("multivariate_calculus.jl")

end # module


text:
abstract type CurveInterpolationMethod end

"""
    ProjectionCurveInterpolation(embedding_interpolation_method::InterpolationType)

Interpolation in the embedding using embedding_interpolation_method and projection
onto the manifold.
"""
struct ProjectionCurveInterpolation{TIM<:Interpolations.InterpolationType} <:
       CurveInterpolationMethod
    embedding_interpolation_method::TIM
end

function ProjectionCurveInterpolation()
    im = Gridded(Linear())
    return ProjectionCurveInterpolation{typeof(im)}(im)
end

"""
    DiscretizedCurves(M::AbstractManifold, grid::AbstractVector)

Space of curves on manifold `M` discretized on the given `grid`.
"""
struct DiscretizedCurves{
    𝔽,
    TM<:AbstractManifold{𝔽},
    TG<:AbstractVector,
    TIM<:CurveInterpolationMethod,
} <: Manifolds.AbstractPowerManifold{𝔽,TM,Manifolds.ArrayPowerRepresentation}
    manifold::TM
    grid::TG
    interpolation_method::TIM
end

function DiscretizedCurves(M::AbstractManifold{𝔽}, grid::AbstractVector) where {𝔽}
    itpm = ProjectionCurveInterpolation()
    return DiscretizedCurves{𝔽,typeof(M),typeof(grid),typeof(itpm)}(M, grid, itpm)
end

struct ProjectionInterpolant{TM<:AbstractManifold,TEITP}
    M::TM
    embedding_itp::TEITP
end

function (pitp::ProjectionInterpolant)(t::Number)
    embedding_val = pitp.embedding_itp(t)
    return project(pitp.M, embedding_val)
end

function make_interpolant(
    M::DiscretizedCurves{<:Any,<:Any,<:Any,<:ProjectionCurveInterpolation},
    p,
)
    rep_size = representation_size(M.manifold)
    embedded_p = [embed(M.manifold, _read(M, rep_size, p, i)) for i in get_iterator(M)]
    embedding_itp = extrapolate(
        interpolate(
            (M.grid,),
            embedded_p,
            M.interpolation_method.embedding_interpolation_method,
        ),
        Flat(),
    )
    return ProjectionInterpolant(M.manifold, embedding_itp)
end

"""
    UniformDiscretizedCurves
"""
const UniformDiscretizedCurves{TM} =
    DiscretizedCurves{𝔽,TM,<:AbstractRange} where {𝔽,TM<:AbstractManifold{𝔽}}

embed!(M::DiscretizedCurves, q, p) = copyto!(q, p)
embed!(M::DiscretizedCurves, Y, p, X) = copyto!(Y, X)

get_iterator(M::DiscretizedCurves) = axes(M.grid, 1)

function manifold_dimension(M::DiscretizedCurves)
    return manifold_dimension(M.manifold) * length(M.grid)
end

function representation_size(M::DiscretizedCurves)
    return (representation_size(M.manifold)..., length(M.grid))
end

function reverse_srvf!(M::UniformDiscretizedCurves, c_out, c_X, initial_point)
    dt = step(M.grid)
    tmp = Ref(allocate(initial_point))
    last_point = Ref(allocate(initial_point))
    copyto!(last_point[], initial_point)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        copyto!(_write(M, rep_size, c_out, i), last_point[])
        vector_transport_to!(
            M.manifold,
            tmp[],
            last_point[],
            _read(M, rep_size, c_X, i),
            last_point[],
        )
        scale = dt * norm(tmp[])
        exp!(M.manifold, last_point[], last_point[], scale * tmp[])
    end
    return c_out
end

function srvf(M::UniformDiscretizedCurves, c, backend::ProjectedDifferenceBackend)
    vel = velocity_curve(M, c, backend)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        copyto!(
            _write(M, rep_size, vel, i),
            q_function(M.manifold, _read(M, rep_size, c, i), _read(M, rep_size, vel, i)),
        )
    end
    return vel
end

function transport_srvf!(M::UniformDiscretizedCurves, c_out, c_p, c_X, p)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        vector_transport_to!(
            M.manifold,
            _write(M, rep_size, c_out, i),
            _read(M, rep_size, c_p, i),
            _read(M, rep_size, c_X, i),
            p,
        )
    end
    return c_out
end

function velocity_curve(M::UniformDiscretizedCurves, c, backend::ProjectedDifferenceBackend)
    N = length(M.grid)
    factor = N - 1
    c_out = allocate(c)
    rep_size = representation_size(M.manifold)
    for i in get_iterator(M)
        wri = _write(M, rep_size, c_out, i)
        if i < N
            log!(M.manifold, wri, _read(M, rep_size, c, i), _read(M, rep_size, c, i + 1))
        else
            log!(M.manifold, wri, _read(M, rep_size, c, N - 1), _read(M, rep_size, c, N))
        end
        wri .*= factor
    end
    return c_out
end

ext: jl


text:
using DataFrames
using Test

m = page_model()

for gas in [:CO2, :CH4, :N2O, :Lin]
    abatementcostparameters = MimiPAGE2009.addabatementcostparameters(m, gas)
    abatementcosts = MimiPAGE2009.addabatementcosts(m, gas)

    abatementcostparameters[:yagg] = readpagedata(m,"test/validationdata/yagg_periodspan.csv")
    abatementcostparameters[:cbe_absoluteemissionreductions] = abatementcosts[:cbe_absoluteemissionreductions]

    abatementcosts[:zc_zerocostemissions] = abatementcostparameters[:zc_zerocostemissions]
    abatementcosts[:q0_absolutecutbacksatnegativecost] = abatementcostparameters[:q0_absolutecutbacksatnegativecost]
    abatementcosts[:blo] = abatementcostparameters[:blo]
    abatementcosts[:alo] = abatementcostparameters[:alo]
    abatementcosts[:bhi] = abatementcostparameters[:bhi]
    abatementcosts[:ahi] = abatementcostparameters[:ahi]
end

p = load_parameters(m)
p["y_year_0"] = 2008.
p["y_year"] = Mimi.dim_keys(m.md, :time)
set_leftover_params!(m, p)

run(m)

@test !isnan(m[:AbatementCostsCO2, :tc_totalcost][10, 5])
@test !isnan(m[:AbatementCostsCH4, :tc_totalcost][10, 5])
@test !isnan(m[:AbatementCostsN2O, :tc_totalcost][10, 5])
@test !isnan(m[:AbatementCostsLin, :tc_totalcost][10, 5])

#compare output to validation data
tc_compare_co2=readpagedata(m, "test/validationdata/tc_totalcosts_co2.csv")
tc_compare_ch4=readpagedata(m, "test/validationdata/tc_totalcosts_ch4.csv")
tc_compare_n2o=readpagedata(m, "test/validationdata/tc_totalcosts_n2o.csv")
tc_compare_lin=readpagedata(m, "test/validationdata/tc_totalcosts_linear.csv")

zc_compare_co2=readpagedata(m, "test/validationdata/zc_zerocostemissionsCO2.csv")
zc_compare_ch4=readpagedata(m, "test/validationdata/zc_zerocostemissionsCH4.csv")
zc_compare_n2o=readpagedata(m, "test/validationdata/zc_zerocostemissionsN2O.csv")
zc_compare_lin=readpagedata(m, "test/validationdata/zc_zerocostemissionsLG.csv")

@test m[:AbatementCostsCO2, :tc_totalcost] ≈ tc_compare_co2 rtol=1e-2
@test m[:AbatementCostsCH4, :tc_totalcost] ≈ tc_compare_ch4 rtol=1e-2
@test m[:AbatementCostsN2O, :tc_totalcost] ≈ tc_compare_n2o rtol=1e-2
@test m[:AbatementCostsLin, :tc_totalcost] ≈ tc_compare_lin rtol=1e-2

@test m[:AbatementCostParametersCO2, :zc_zerocostemissions] ≈ zc_compare_co2 rtol=1e-2
@test m[:AbatementCostParametersCH4, :zc_zerocostemissions] ≈ zc_compare_ch4 rtol=1e-3
@test m[:AbatementCostParametersN2O, :zc_zerocostemissions] ≈ zc_compare_n2o rtol=1e-3
@test m[:AbatementCostParametersLin, :zc_zerocostemissions] ≈ zc_compare_lin rtol=1e-3


text:
module Jets

using CRC32c, LinearAlgebra, Random

abstract type JetAbstractSpace{T,N} end

"""
    eltype(R)

Return the element type of the space `R::JetAbstractSpace`.
"""
Base.eltype(R::JetAbstractSpace{T}) where {T} = T
Base.eltype(R::Type{JetAbstractSpace{T,N}}) where {T,N} = T
Base.eltype(R::Type{JetAbstractSpace{T}}) where {T} = T
Base.ndims(R::JetAbstractSpace{T,N}) where {T,N} = N

"""
    length(R)

Return the dimension the space `R::JetAbstractSpace`
"""
Base.length(R::JetAbstractSpace) = prod(size(R))

"""
    size(R[,i])

Return the shape of the array associated to the Jet space R::JetAbstractSpace.
If `i` is specifid, then returns the length along the ith array dimension.
"""
Base.size(R::JetAbstractSpace, i) = size(R)[i]

"""
    reshape(x, R)

Returns an array that is consistent with the shape of the space `R::JetAbstractSpace`, and shares
memory with `x`.
"""
Base.reshape(x::AbstractArray, R::JetAbstractSpace) = reshape(x, size(R))

struct JetSpace{T,N} <: JetAbstractSpace{T,N}
    n::NTuple{N,Int}
end

"""
    JetSpace(T, n)

Construct and return a JetSpace of type `T` and size `n`

# Examples
Create a 100 dimension space with array elelment type Float64 and array size (100,)
```
R1 = JetSpace(Float64, 100)
```

Create a 100 dimension space with array element type Float32 and array size (5, 20)
```
R2 = JetSpace(Float32, 5, 20)
```
"""
JetSpace(_T::Type{T}, n::Vararg{Int,N}) where {T,N} = JetSpace{T,N}(n)
JetSpace(_T::Type{T}, n::NTuple{N,Int}) where {T,N} = JetSpace{T,N}(n)

Base.size(R::JetSpace) = R.n
Base.eltype(R::Type{JetSpace{T,N}}) where {T,N} = T
Base.eltype(R::Type{JetSpace{T}}) where {T} = T
Base.vec(R::JetSpace) = JetSpace(eltype(R), length(R))

@doc """
    Array(R)

Construct an uninitialized array of the type and size defined by `R::JetsAbstractSpace`.
"""
Array

@doc """
    ones(R)

Construct an array of the type and size defined by `R::JetAbstractSpace{T}` and filled with `one(T)`.
"""
ones

@doc """
    rand(R)

Construct an array of the type and size defined by the `R::JetAbstractSpace`, and filled with random values.
"""
rand

@doc """
    zeros(R)

Construct an array of the type and size defined by `R::JetAbstractSpace{T}` and filled with `zero(T)`.
"""
zeros

for f in (:ones, :rand, :zeros)
    @eval (Base.$f)(R::JetSpace{T,N}) where {T,N} = ($f)(T,size(R))::Array{T,N}
end
Base.Array(R::JetSpace{T,N}) where {T,N} = Array{T,N}(undef, size(R))

"""
    randperm(R)

Construct a list of random linear indices over the dimensions of `R::JetAbstractSpace`.  The list
is useful for selecting a random subset of a multi-dimensional image.

# Example
```julia
using Jets
R = JetSpace(Float64, 10, 2)
x = rand(R)
y = x[randperm(R)[1:10]] # get 10 elements at random from x
```
"""
Random.randperm(R::JetAbstractSpace, k::Int) = sort(randperm(length(R))[1:k])

space(x::AbstractArray{T,N}) where {T,N} = JetSpace{T,N}(size(x))

jet_missing(m) = error("not implemented")

mutable struct Jet{D<:JetAbstractSpace,R<:JetAbstractSpace,F<:Function,DF<:Function,DF′<:Function,U<:Function,M<:AbstractArray,S<:NamedTuple}
    dom::D
    rng::R
    f!::F
    df!::DF
    df′!::DF′
    upstate!::U
    mₒ::M
    s::S
end

"""
    Jet(;dom, rng, f!, df!, df′!, upstate!, s)

Return a `Jet` with domain `dom::JetAbstractSpace`, range `rng::JetSAbstractpace`, with
forward mapping  `f!::Function`, linearized forward mapping `df!::Function`, linearized adjoint
mapping `df′!::Function`, Jacobian state modification function `upstate!::Function`, and state
`s::NamedTuple`.

A jet describes a function `f!` and its linearization (forward `df!, and adjoint `df′!``) about
a point.

If one of `f!` or `df!` is specified, and `df′!` is not, then we assume that `f!=df!=df′!`. This means
that the operator is linear and self-adjoint.

If `f!` and `df!` are sepecified, bug `df′!` is not, then we assume that `df′!=df!`.  This means
that the operator is nonlinear and self-adjoint.

# Example
Consider a nonlinear mapping with a self-adjoint linearization ``f(x)=x^2``
```julia
using Jets
g!(m) = m.^2
dg!(δm; mₒ) = @. 2*mₒ*δm
jet = Jet(;dom=JetSpace(Float64,2), rng=JetSpace(Float64,2), f! = g!, df! = dg!)
```
"""
function Jet(;
        dom,
        rng,
        f! = jet_missing,
        df! = jet_missing,
        df′! = jet_missing,
        upstate! = (m,s) -> nothing,
        s = NamedTuple())
    if isa(f!, typeof(jet_missing)) && isa(df!, typeof(jet_missing))
        error("must set at-least one of f! and df!")
    end
    if isa(f!, typeof(jet_missing))
        f! = df!
    end
    if isa(df′!, typeof(jet_missing))
        df′! = df!
    end
    Jet(dom, rng, f!, df!, df′!, upstate!, Array{eltype(dom)}(undef, ntuple(i->0,ndims(dom))), s)
end

f!(d, jet::Jet, m; kwargs...) = jet.f!(d, m; kwargs...)
df!(d, jet::Jet, m; kwargs...) = jet.df!(d, m; kwargs...)
df′!(m, jet::Jet, d; kwargs...) = jet.df′!(m, d; kwargs...)

abstract type Jop{T<:Jet} end

struct JopNl{T<:Jet} <: Jop{T}
    jet::T
end

"""
    JopNl(; kwargs ...)

Construct a `JopNl` with `Jet` constructed from keyword arguments `kwargs`.
This is equivalent to `JopNl(Jet(;kwargs...))`.  Please see `Jet` for more
information.
"""
JopNl(;kwargs...) = JopNl(Jet(;kwargs...))

struct JopLn{T<:Jet} <: Jop{T}
    jet::T
end
JopLn(jet::Jet, mₒ::AbstractArray) = JopLn(point!(jet, mₒ))

"""
    JopLn(; kwargs ...)

Construct a `JopLn` with `Jet` constructed from keyword arguments `kwargs`.
This is equivalent to `JopLn(Jet(;kwargs...))`.  Please see `Jet` for more
information.
"""
JopLn(;kwargs...) = JopLn(Jet(;kwargs...))

JopLn(A::JopLn) = A
JopLn(F::JopNl) = JopLn(jet(F))

struct JopAdjoint{J<:Jet,T<:Jop{J}} <: Jop{J}
    op::T
end

Base.copy(jet::Jet, copymₒ=true) = Jet(jet.dom, jet.rng, jet.f!, jet.df!, jet.df′!, jet.upstate!, copymₒ ? copy(jet.mₒ) : jet.mₒ, deepcopy(jet.s))
Base.copy(A::JopLn, copymₒ=true) = JopLn(copy(jet(A), copymₒ))
Base.copy(A::JopAdjoint, copymₒ=true) = JopAdjoint(copy(A.op, copymₒ))
Base.copy(F::JopNl, copymₒ=true) = JopNl(copy(jet(F), copymₒ))

JopLn(A::JopAdjoint) = A

"""
    R = domain(A)

Return `R::JetAbstractSpace`, which is the domain of `A::Union{Jet, Jop, AbstractMatrix}`.
"""
domain(jet::Jet) = jet.dom

"""
    R = range(A)

Return `R::JetAbstractSpace`, which is the range of `A::Union{Jet, Jop, AbstractMatrix}`.
"""
Base.range(jet::Jet) = jet.rng

"""
    eltype(A::Union{Jet,Jop,JopAdjoint})

Return the element type of `A`.
"""
Base.eltype(jet::Jet) = promote_type(eltype(domain(jet)), eltype(range(jet)))

"""
    state(A::Union{Jet,Jop,JopAdjoint}[, key])

If `key::Symbol` is specified, then return the state corresponding to `key`.
Otherwise, return the state of A as a `NamedTuple`.
"""
state(jet::Jet) = jet.s
state(jet::Jet{D,R,F}, key) where {D,R,F<:Function} = jet.s[key]

"""
    state!(A::Union{Jet,Jop,JopAdjoint}, s)

Updates and merges the state of the `A` with `s`.
"""
state!(jet, s) = begin jet.s = merge(jet.s, s); jet end

"""
    perfstat(A)

Return a `Dictionary` with performance information for A::Union{Jet,Jop,JopAdjoint}. the
`perfstat(jet::Jet)` method that can be implemented by the author of an operator to track
performance metrics.
"""
perfstat(jet::T) where {D,R,F<:Function,T<:Jet{D,R,F}} = nothing

"""
    point(F)

Return the linearization point (model vector) `mₒ` associated with `F::Union{Jet, JopLn, JopAdjoint}`.
"""
point(jet::Jet) = jet.mₒ

Base.close(j::Jet{D,R,F}) where {D,R,F<:Function} = false

"""
    point!(F, mₒ)

Update the linearization point (model vector) for `F::Union{Jet, JopLn, JopAdjoint}` to model vector `mₒ`.
"""
function point!(jet::Jet, mₒ::AbstractArray)
    jet.mₒ = mₒ
    jet.upstate!(mₒ, state(jet))
    jet
end

"""
    jet(A)

Return the `Jet` associated with `A::Union{Jop, JopAdjoint}`.
"""
jet(A::Jop) = A.jet
jet(A::JopAdjoint) = jet(A.op)
Base.eltype(A::Jop) = eltype(jet(A))
point(A::JopLn) = point(jet(A))
point(A::JopAdjoint) = point(jet(A.op))
state(A::Jop) = state(jet(A))
state(A::Jop, key) = state(jet(A), key)
state!(A::Jop, s) = state!(jet(A), s)
perfstat(A::Jop) = perfstat(jet(A))
Base.close(A::Jop) = close(jet(A))

domain(A::Jop) = domain(jet(A))
Base.range(A::Jop) = range(jet(A))

domain(A::JopAdjoint) = range(A.op)
Base.range(A::JopAdjoint) = domain(A.op)

domain(A::AbstractMatrix{T}) where {T} = JetSpace(T, size(A,2))
Base.range(A::AbstractMatrix{T}) where {T} = JetSpace(T, size(A,1))

function shape(A::Union{Jet,Jop}, i)
    if i == 1
        return size(range(A))
    end
    size(domain(A))
end

"""
    shape(A[, i])

Return the shape of the range and domain of `A::Union{Jet, Jop, AbstractMatrix}`.
With no arguments, return `(size(range(A)), size(domain(A)))`.  With `i` specified,
return `size(range(A))` for `i = 1` and return `size(domain(A))` otherwise.
"""
shape(A::Union{Jet,Jop}) = (shape(A, 1), shape(A, 2))

shape(A::AbstractMatrix,i) = (size(A, i),)
shape(A::AbstractMatrix) = ((size(A, 1),), (size(A, 2),))

"""
    size(A[, i])

Return the size of the range and domain of `A::Union{Jet,Jop}`.  With no arguments, return
`(length(range(A)), length(domain(A)))`.  With `i` specified, return `length(range(A))` for `i = 1`
and return `length(domain(A))` otherwise.
"""
Base.size(A::Union{Jet,Jop}, i) = prod(shape(A, i))
Base.size(A::Union{Jet,Jop}) = (size(A, 1), size(A, 2))

"""
    jacobian!(F, m₀)

Return the jacobian of `F::Union{Jet, Jop, AbstractMatrix}` at the
linearization point `m₀`. The jacobian shares the underlying `Jet` with `F`.
This means that if the jacobian may mutate `F`.
"""
jacobian!(jet::Jet, mₒ::AbstractArray) = JopLn(jet, mₒ)
jacobian!(F::JopNl, mₒ::AbstractArray) = jacobian!(jet(F), mₒ)
jacobian!(A::Union{JopLn,AbstractMatrix}, mₒ::AbstractArray) = A

"""
    jacobian(F, m₀)

Return the jacobian of `F::Union{Jet, Jop, AbstractMatrix}` at the point `m₀`.
The linearization constructs a new underlying `Jet`.
"""
jacobian(F::Union{Jet,Jop}, mₒ::AbstractArray) = jacobian!(copy(F, false), copy(mₒ))
jacobian(A::AbstractMatrix, mₒ::AbstractArray) = copy(A)

"""
    adjoint(A::Union{JopLn, JopAdjoint})

Return the adjoint of A.
"""
Base.adjoint(A::JopLn) = JopAdjoint(A)
Base.adjoint(A::JopAdjoint) = A.op

"""
    mul!(d, F, m)

In place version of `d=F*m` where F is a Jets linear (e.g. `JopLn`) or nonlinear (`JopNl`) operator.
"""
LinearAlgebra.mul!(d::AbstractArray, F::JopNl, m::AbstractArray) = f!(d, jet(F), m; state(F)...)
LinearAlgebra.mul!(d::AbstractArray, A::JopLn, m::AbstractArray) = df!(d, jet(A), m; mₒ=point(A), state(A)...)
LinearAlgebra.mul!(m::AbstractArray, A::JopAdjoint{J,T}, d::AbstractArray) where {J<:Jet,T<:JopLn} = df′!(m, jet(A), d; mₒ=point(A), state(A)...)

"""
    :*(F, m)

Constructs `F*m` where F is a Jets linear (e.g. `JopLn`) or nonlinear (`JopNl`) operator.
"""
Base.:*(A::Jop, m::AbstractArray) = mul!(zeros(range(A)), A, m)

Base.show(io::IO, A::JopLn) = show(io, "Jet linear operator, $(size(domain(A))) → $(size(range(A)))")
Base.show(io::IO, A::JopAdjoint) = show(io, "Jet adjoint operator, $(size(domain(A))) → $(size(range(A)))")
Base.show(io::IO, F::JopNl) = show(io, "Jet nonlinear operator, $(size(domain(F))) → $(size(range(F)))")

#
# Symmetric spaces / arrays
#
struct JetSSpace{T,N,F<:Function} <: JetAbstractSpace{T,N}
    n::NTuple{N,Int}
    M::NTuple{N,Int}
    map::F
end

"""
    JetSSpace(_T, n, M, map::F)

Construct and return a symmetric space `JetSSpace`.

# parameters
* `_T` is the type, usually `Complex{Float32}` or `Complex{Float64}`.
* `n` is a tuple that defines the dimensionality of the space.
* `M` is a tuple that defines which dimensions are symmetric. Note that currently only a single symmetric dimension is supported by the API.
* `F` is a function that maps indices for the symmetric dimension, described below.

An example requiring a `JetSSpace` is the Fourier transform: the Fourier transform of a real vector is in
a complex space with Hermittian symmetry. Only the positive frequencies are needed, and the spectrum at
negative frequencies is the Hermittian conjugate of the spectrum at the corresponding positive frequencies:
`S(-f) = conj(S(f)`. For this example the map `F` is a function that returns the multi-dimensional index
of `f` when given the multi-dimensional index of `-f`.

See also: `JopFft` in the `JetPackTransforms` package.
"""
JetSSpace(_T::Type{T}, n::NTuple{N,Int}, M::NTuple{N,Int}, map::F) where {T,N,F} = JetSSpace{T,N,F}(n, M, map)

Base.size(R::JetSSpace) = R.n
Base.eltype(R::Type{JetSSpace{T,N,F}}) where {T,N,F} = T
Base.eltype(R::Type{JetSSpace{T,N}}) where {T,N} = T
Base.eltype(R::Type{JetSSpace{T}}) where {T} = T
symspace() = nothing

struct SymmetricArray{T,N,F<:Function} <: AbstractArray{T,N}
    A::Array{T,N}
    n::NTuple{N,Int}
    map::F
end

space(x::SymmetricArray{T,N,F}) where {T,N,F} = JetSSpace{T,N,F}(x.n, size(x.A), x.map)

Base.parent(A::SymmetricArray) = A.A

# SymmetricArray array interface implementation <--
Base.IndexStyle(::Type{T}) where {T<:SymmetricArray} = IndexCartesian()
Base.size(A::SymmetricArray) = A.n

function Base.getindex(A::SymmetricArray{T,N}, I::Vararg{Int,N}) where {T,N}
    for idim = 1:ndims(A)
        if I[idim] > size(A.A, idim)
            return conj(A.A[A.map(I)])
        end
    end
    A.A[CartesianIndex(I)]
end

function Base.getindex(A::SymmetricArray, i::Int)
    I = CartesianIndices(size(A))[i]
    getindex(A, I.I...)
end

function Base.setindex!(A::SymmetricArray{T,N}, v, I::Vararg{Int,N}) where {T,N}
    for idim = 1:ndims(A)
        if I[idim] > size(A.A, idim)
            A.A[A.map(I)] = conj(v)
            return A.A[A.map(I)]
        end
    end
    A.A[CartesianIndex(I)] = v
end

function Base.setindex!(A::SymmetricArray, v, i::Int)
    I = CartesianIndices(size(A))[i]
    setindex!(A, v, I.I...)
end

Base.similar(A::SymmetricArray, ::Type{T}) where {T<:Complex} = SymmetricArray(similar(A.A), A.n, A.map)
Base.similar(A::SymmetricArray, ::Type{T}) where {T<:Real} = Array{T}(undef, size(A))
Base.similar(A::SymmetricArray{T}) where {T} = similar(A, T)
# -->

# SymmetricArray broadcasting interface implementation --<
Base.BroadcastStyle(::Type{<:SymmetricArray}) = Broadcast.ArrayStyle{SymmetricArray}()

Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{SymmetricArray}}, ::Type{T}) where {T} = similar(find_symmetricarray(bc), T)

find_symmetricarray(bc::Broadcast.Broadcasted) = find_symmetricarray(bc.args)
find_symmetricarray(args::Tuple) = find_symmetricarray(find_symmetricarray(args[1]), Base.tail(args))
find_symmetricarray(x) = x
find_symmetricarray(a::SymmetricArray, rest) = a
find_symmetricarray(::Any, rest) = find_symmetricarray(rest)

get_symmetricarray_parent(bc::Broadcast.Broadcasted, ::Type{S}) where {S} = Broadcast.Broadcasted{S}(bc.f, map(arg->get_symmetricarray_parent(arg, S), bc.args))
get_symmetricarray_parent(A::SymmetricArray, ::Type{<:Any}) = parent(A)
get_symmetricarray_parent(A, ::Type{<:Any}) = A

function Base.copyto!(dest::SymmetricArray{T,N}, bc::Broadcast.Broadcasted{Nothing}) where {T,N}
    S = Broadcast.DefaultArrayStyle{N}
    copyto!(parent(dest), get_symmetricarray_parent(bc, S))
    dest
end
# -->

for f in (:ones, :rand, :zeros)
    @eval (Base.$f)(R::JetSSpace{T,N,F}) where {T,N,F} = SymmetricArray(($f)(T,R.M), R.n, R.map)::SymmetricArray{T,N,F}
end
Base.Array(R::JetSSpace{T,N,F}) where {T,N,F} = SymmetricArray{T,N,F}(Array{T,N}(undef, R.M), R.n, R.map)

#
# composition, f ∘ g
#

JetComposite(ops) = Jet(f! = JetComposite_f!, df! = JetComposite_df!, df′! = JetComposite_df′!, dom = domain(ops[end]), rng = range(ops[1]), s = (ops=ops,))

function JetComposite_f!(d::T, m; ops, kwargs...) where {T<:AbstractArray}
    g = mapreduce(i->(_m->mul!(zeros(range(ops[i])), ops[i], _m)), ∘, 1:length(ops))
    d .= g(m)
    d::T
end

function JetComposite_df!(d::T, m; ops, kwargs...) where {T<:AbstractArray}
    dg = mapreduce(i->(_m->mul!(zeros(range(JopLn(ops[i]))), JopLn(ops[i]), _m)), ∘, 1:length(ops))
    d .= dg(m)
    d::T
end

function JetComposite_df′!(m::T, d; ops, kwargs...) where {T<:AbstractArray}
    dg′ = mapreduce(i->(_d->mul!(zeros(domain(JopLn(ops[i]))), (JopLn(ops[i]))', _d)), ∘, length(ops):-1:1)
    m .= dg′(d)
    m::T
end

jops_comp(op::Jop) = (op,)
jops_comp(op::JopLn{J}) where {D,R,J<:Jet{D,R,typeof(JetComposite_f!)}} = state(jet(op)).ops
jops_comp(op::JopNl{J}) where {D,R,J<:Jet{D,R,typeof(JetComposite_f!)}} = state(jet(op)).ops

function jops_comp(op::JopAdjoint{J,T}) where {D,R,J<:Jet{D,R,typeof(JetComposite_f!)},T<:JopLn{J}}
    ops = state(op.op).ops
    n = length(ops)
    ntuple(i->JopAdjoint(ops[n-i+1]), n)
end

"""
    :∘(A₂, A₁)

Construct the composition of the two `Jets` operators.  Note that when applying the composition
operator, operators are applied in order from right to left: first `A₁` and then `A₂`.

# Example
```julia
using Jets
dg!(d,m;mₒ) = @. d = 2*m
A₁ = JopLn(Jet(;dom=JetSpace(Float64,2), rng=JetSpace(Float64,2), df! = dg!))
A₂ = JopLn(Jet(;dom=JetSpace(Float64,2), rng=JetSpace(Float64,2), df! = dg!))
C = A₂ ∘ A₁
m = rand(domain(C))
d = C * m
```
"""
Base.:∘(A₂::Union{JopAdjoint,JopLn}, A₁::Union{JopAdjoint,JopLn}) = JopLn(JetComposite((jops_comp(A₂)..., jops_comp(A₁)...)))
Base.:∘(A₂::Jop, A₁::Jop) = JopNl(JetComposite((jops_comp(A₂)..., jops_comp(A₁)...)))
Base.:∘(A₂::AbstractMatrix, A₁::AbstractMatrix) = A₂*A₁

_matmul_df!(d, m; A, kwargs...) = mul!(d, A, m)
_matmul_df′!(m, d; A, kwargs...) = mul!(m, adjoint(A), d)
Base.:∘(A₂::Jop, A₁::AbstractMatrix) = A₂ ∘ JopLn(;dom = domain(A₁), rng = range(A₁), df! = _matmul_df!, df′! = _matmul_df′!, s=(A=A₁,))
Base.:∘(A₂::AbstractMatrix, A₁::Jop) = JopLn(;dom = domain(A₂), rng = range(A₂), df! = _matmul_df!, df′! = _matmul_df′!, s=(A=A₂,)) ∘ A₁

function point!(j::Jet{D,R,typeof(JetComposite_f!)}, mₒ::AbstractArray) where {D<:JetAbstractSpace,R<:JetAbstractSpace}
    j.mₒ = mₒ
    ops = state(j).ops
    _m = copy(mₒ)
    for i = length(ops):-1:1
        point!(jet(ops[i]), _m)
        if i > 1
            _m = mul!(zeros(range(ops[i])), ops[i], _m)
        end
    end
    j
end

function Base.close(j::Jet{D,R,typeof(JetComposite_f!)}) where {D,R}
    ops = state(j).ops
    close.(ops)
    nothing
end

function perfstat(j::Jet{D,R,typeof(JetComposite_f!)}) where {D,R}
    ops = state(j).ops
    s = nothing
    for op in ops
        s = perfstat(op)
        s == nothing || break
    end
    s
end

function state(j::Jet{D,R,typeof(JetComposite_f!)}, key) where {D,R}
    key ∈ keys(state(j)) && (return state(j)[key])

    haskey = 0
    local _op
    for op in state(j).ops
        if key ∈ keys(state(op))
            haskey += 1
            _op = op
        end
    end

    haskey == 0 && error("key $key does not exist in the state of the composite operator")
    haskey > 1 && error("ambiguous: key $key exists in more than one operator in the composition")

    state(_op, key)
end

#
# Composition, f ± g
#
JetSum(ops, sgns) = Jet(f! = JetSum_f!, df! = JetSum_df!, df′! = JetSum_df′!, dom = domain(ops[1]), rng = range(ops[1]), s = (ops=ops, sgns=sgns))

function JetSum_f!(d, m; ops, sgns, kwargs...)
    d .= 0
    _d = zeros(range(ops[1]))
    for i = 1:length(ops)
        broadcast!(sgns[i], d, d, mul!(_d, ops[i], m))
    end
    d
end

function JetSum_df!(d, m; ops, sgns, kwargs...)
    d .= 0
    _d = zeros(range(ops[1]))
    for i = 1:length(ops)
        broadcast!(sgns[i], d, d, mul!(_d, JopLn(ops[i]), m))
    end
    d
end

function JetSum_df′!(m, d; ops, sgns, kwargs...)
    m .= 0
    _m = zeros(domain(ops[1]))
    for i = 1:length(ops)
        broadcast!(sgns[i], m, m, mul!(_m, JopLn(ops[i])', d))
    end
    m
end

jops_sum(op::Jop) = (op,)
jops_sum(op::JopLn{J}) where {D,R,J<:Jet{D,R,typeof(JetSum_f!)}} = state(jet(op)).ops
jops_sum(op::JopNl{J}) where {D,R,J<:Jet{D,R,typeof(JetSum_f!)}} = state(jet(op)).ops

function jops_sum(op::JopAdjoint{J,T}) where {D,R,J<:Jet{D,R,typeof(JetSum_f!)},T<:JopLn{J}}
    ops = state(op.op).ops
    n = length(ops)
    ntuple(i->JopAdjoint(ops[i]), n)
end

function flipsgn(sgn, sgnnew)
    typeof(sgnnew) == typeof(+) && (return sgn)
    typeof(sgnnew) == typeof(-) && typeof(sgn) == typeof(-) && (return +)
    typeof(sgnnew) == typeof(-) && typeof(sgn) == typeof(+) && (return -)
end

sgns(op::Jop, r) = (r,)
sgns(op::JopLn{J}, r) where {D,R,J<:Jet{D,R,typeof(JetSum_f!)}} = ntuple(i->flipsgn(state(op).sgns[i],r), length(state(op).sgns))
sgns(op::JopNl{J}, r) where {D,R,J<:Jet{D,R,typeof(JetSum_f!)}} = ntuple(i->flipsgn(state(op).sgns[i],r), length(state(op).sgns))
sgns(op::JopAdjoint{J,T}, r) where {D,R,J<:Jet{D,R,typeof(JetSum_f!)},T<:JopLn{J}} = ntuple(i->flipsgn(state(op).sgns[i],r), length(state(op.op).sgns))

"""
    :+(A₂, A₁)

Construct and return the linear combination of the two `Jets` operators `A₁::Jop` and `A₂::Jop`.
Note that `A₁` and `A₂` must have consistent (same size and type) domains and ranges.

# Example
```
A = 1.0*A₁ - 2.0*A₂ + 3.0*A₃
```
"""
Base.:+(A₂::Union{JopAdjoint,JopLn}, A₁::Union{JopAdjoint,JopLn}) = JopLn(JetSum((jops_sum(A₂)..., jops_sum(A₁)...), (sgns(A₂,+)..., sgns(A₁,+)...)))
Base.:+(A₂::Jop, A₁::Jop) = JopNl(JetSum((jops_sum(A₂)..., jops_sum(A₁)...), (sgns(A₂,+)..., sgns(A₁,+)...)))
Base.:+(A₂::Jop, A₁::AbstractMatrix) = A₂ + JopLn(;dom = domain(A₁), rng = range(A₁), df! = _matmul_df!, df′! = _matmul_df′!, s=(A=A₁,))
Base.:+(A₂::AbstractMatrix, A₁::Jop) = JopLn(;dom = domain(A₂), rng = range(A₂), df! = _matmul_df!, df′! = _matmul_df′!, s=(A=A₂,)) + A₁

"""
    :-(A₂, A₁)

Construct and return the linear combination of the two `Jets` operators `A₁::Jop` and `A₂::Jop`.
Note that `A₁` and `A₂` must have consistent (same size and type) domains and ranges.

# Example
```
A = 1.0*A₁ - 2.0*A₂ + 3.0*A₃
```
"""
Base.:-(A₂::Union{JopAdjoint,JopLn}, A₁::Union{JopAdjoint,JopLn}) = JopLn(JetSum((jops_sum(A₂)..., jops_sum(A₁)...), (sgns(A₂,+)..., sgns(A₁,-)...)))
Base.:-(A₂::Jop, A₁::Jop) = JopNl(JetSum((jops_sum(A₂)..., jops_sum(A₁)...), (sgns(A₂,+)..., sgns(A₁,-)...)))
Base.:-(A₂::Jop, A₁::AbstractMatrix) = A₂ - JopLn(;dom = domain(A₁), rng = range(A₁), df! = _matmul_df!, df′! = _matmul_df′!, s=(A=A₁,))
Base.:-(A₂::AbstractMatrix, A₁::Jop) = JopLn(;dom = domain(A₂), rng = range(A₂), df! = _matmul_df!, df′! = _matmul_df′!, s=(A=A₂,)) - A₁

function point!(j::Jet{D,R,typeof(JetSum_f!)}, mₒ::AbstractArray) where {D<:JetAbstractSpace,R<:JetAbstractSpace}
    for op in state(j).ops
        point!(jet(op), mₒ)
    end
    j
end

function Base.close(j::Jet{D,R,typeof(JetSum_f!)}) where {D,R}
    ops = state(j).ops
    close.(ops)
    nothing
end

function perfstat(j::Jet{D,R,typeof(JetSum_f!)}) where {D,R}
    ops = state(j).ops
    s = nothing
    for op in ops
        s = perfstat(op)
        s == nothing || break
    end
    s
end

#
# Block operator
#
struct JetBSpace{T,S<:JetAbstractSpace} <: JetAbstractSpace{T,1}
    spaces::Vector{S}
    indices::Vector{UnitRange{Int}}
    function JetBSpace(spaces)
        S = mapreduce(typeof, promote_type, spaces)
        T = eltype(S)
        indices = [0:0 for i=1:length(spaces)]
        stop = 0
        for i = 1:length(indices)
            start = stop + 1
            stop = start + length(spaces[i]) - 1
            indices[i] = start:stop
        end
        new{T,S}(spaces, indices)
    end
end

Base.:(==)(x::JetBSpace, y::JetBSpace) = x.spaces == y.spaces && x.indices == y.indices

Base.size(R::JetBSpace) = (R.indices[end][end],)
Base.eltype(R::Type{JetBSpace{T,S}}) where {T,S} = T
Base.eltype(R::Type{JetBSpace{T}}) where {T} = T
Base.vec(R::JetBSpace) = R

"""
    indices(R, iblock)

Return the linear indices associated with block `iblock` in the `Jets` block space
`R::JetBSpace`.

# Example
Consider a block operator with 2 row-blocks and 3 column-blocks.  We can use `indices`
to determine the elements of the vector that are associatd with the first block of its
domain:
```
using Pkg
pkg"add Jets JetPack"
using Jets, JetPack
A = @blockop [JopDiagonal(rand(10)) for irow=1:2, icol=1:3]
indices(domain(A), 1) # returns indices 1:10
```
"""
indices(R::JetBSpace, iblock::Integer) = R.indices[iblock]

"""
    space(R, iblock)

Return the `Jets` sub-space associated with block `iblock` in the `Jets` block space
`R::JetBSpace`.

# Example
Consider a block operator with 2 row-blocks and 3 column-blocks.  We can use `space` to
determine the sub-space associated with the first block of its domain:
```
using Pkg
pkg"add Jets JetPack"
using Jets, JetPack
A = @blockop [JopDiagonal(rand(10)) for irow=1:2, icol=1:3]
space(domain(A), 1) # JetSpace(Float64,10)
```
"""
space(R::JetBSpace, iblock::Integer) = R.spaces[iblock]

"""
    nblocks(R)

Return the number of blocks in the `Jets` block space `R::JetBSpace`.
"""
nblocks(R::JetBSpace) = length(R.spaces)
nblocks(R::JetAbstractSpace) = 1

struct BlockArray{T,A<:AbstractArray{T}} <: AbstractArray{T,1}
    arrays::Vector{A}
    indices::Vector{UnitRange{Int}}
end

space(x::BlockArray) = JetBSpace([space(_x) for _x in x.arrays])

# BlockArray array interface implementation <--
Base.IndexStyle(::Type{T}) where {T<:BlockArray} = IndexLinear()
Base.size(A::BlockArray) = (A.indices[end][end],)

function Base.getindex(A::BlockArray, i::Int)
    j = findfirst(rng->i∈rng, A.indices)::Int
    A.arrays[j][i-A.indices[j][1]+1]
end
function Base.setindex!(A::BlockArray, v, i::Int)
    j = findfirst(rng->i∈rng, A.indices)::Int
    A.arrays[j][i-A.indices[j][1]+1] = v
end

Base.similar(A::BlockArray, ::Type{T}) where {T} = BlockArray([similar(A.arrays[i], T) for i=1:length(A.arrays)], A.indices)
Base.similar(A::BlockArray{T}) where {T} = similar(A, T)
Base.similar(x::BlockArray, ::Type{T}, n::Integer) where {T} = length(x) == n ? similar(x, T) : Array{T}(undef, n)
Base.similar(x::BlockArray, ::Type{T}, dims::Tuple{Int}) where {T} = similar(x, T, dims[1])

function LinearAlgebra.norm(x::BlockArray{T}, p::Real=2) where {T}
    if p == Inf
        mapreduce(_x->norm(_x,p), max, x.arrays)
    elseif p == -Inf
        mapreduce(_x->norm(_x,p), min, x.arrays)
    elseif p == 1
        mapreduce(_x->norm(_x,p), +, x.arrays)
    elseif p == 0
        mapreduce(_x->norm(_x,p), +, x.arrays)
    else
        _T = float(real(T))
        _p = _T(p)
        mapreduce(_x->norm(_x,p)^_p, +, x.arrays)^(one(_T)/_p)
    end
end

function LinearAlgebra.dot(x::BlockArray{T}, y::BlockArray{T}) where {T}
    a = zero(T)
    for i = 1:length(x.arrays)
        a += dot(x.arrays[i], y.arrays[i])
    end
    a
end

indices(x::BlockArray, i) = x.indices[i]

nblocks(x::BlockArray) = length(x.indices)

function Base.convert(::Type{Array}, x::BlockArray{T}) where {T}
    _x = Vector{T}(undef, length(x))
    for i = 1:length(x.indices)
        _x[indices(x,i)] .= vec(x.arrays[i])
    end
    _x
end

function Base.extrema(x::BlockArray{T}) where {T}
    mn,mx = extrema(x.arrays[1])
    for i = 2:length(x.arrays)
        _mn, _mx = extrema(x.arrays[i])
        _mn < mn && (mn = _mn)
        _mx > mx && (mx = _mx)
    end
    mn,mx
end

function Base.fill!(x::BlockArray, a)
    for i = 1:length(x.arrays)
        fill!(x.arrays[i], a)
    end
    x
end
# -->

# BlockArray broadcasting implementation --<
struct BlockArrayStyle <: Broadcast.AbstractArrayStyle{1} end
Base.BroadcastStyle(::Type{<:BlockArray}) = BlockArrayStyle()
BlockArrayStyle(::Val{1}) = BlockArrayStyle()

Base.similar(bc::Broadcast.Broadcasted{BlockArrayStyle}, ::Type{T}) where {T} = similar(find_blockarray(bc), T)
find_blockarray(bc::Broadcast.Broadcasted) = find_blockarray(bc.args)
find_blockarray(args::Tuple) = find_blockarray(find_blockarray(args[1]), Base.tail(args))
find_blockarray(x) = x
find_blockarray(a::BlockArray, rest) = a
find_blockarray(::Any, rest) = find_blockarray(rest)

getblock(bc::Broadcast.Broadcasted, ::Type{S}, iblock, indices) where {S} = Broadcast.Broadcasted{S}(bc.f, map(arg->getblock(arg, S, iblock, indices), bc.args))
getblock(A::BlockArray, ::Type{<:Any}, iblock, indices) = getblock(A, iblock)
getblock(A::AbstractArray, ::Type{<:Any}, iblock, indices) = A[indices]
getblock(A, ::Type{<:Any}, iblock, indices) = A

function Base.copyto!(dest::BlockArray{T,<:AbstractArray{T,N}}, bc::Broadcast.Broadcasted{BlockArrayStyle}) where {T,N}
    S = Broadcast.DefaultArrayStyle{N}
    for iblock = 1:nblocks(dest)
        copyto!(getblock(dest, iblock), getblock(bc, S, iblock, dest.indices[iblock]))
    end
    dest
end
# -->

getblock(x::BlockArray, iblock) = x.arrays[iblock]
getblock!(x::BlockArray, iblock, xblock::AbstractArray) = xblock .= x.arrays[iblock]
setblock!(x::BlockArray, iblock, xblock) = x.arrays[iblock] .= xblock

getblock(x::AbstractArray, iblock) = x
getblock!(x::AbstractArray, iblock, xblock::AbstractArray) = xblock .= x
setblock!(x::AbstractArray, iblock, xblock) = x .= xblock

for f in (:Array, :ones, :rand, :zeros)
    @eval (Base.$f)(R::JetBSpace{T,S}) where {T,S<:JetAbstractSpace} = BlockArray([($f)(space(R, i)) for i=1:length(R.spaces)], R.indices)
end

function JetBlock(ops::AbstractMatrix{T}; dadom=false, kwargs...) where {T<:Jop}
    dom = (size(ops,2) == 1 && !dadom) ? domain(ops[1,1]) : JetBSpace([domain(ops[1,i]) for i=1:size(ops,2)])
    rng = JetBSpace([range(ops[i,1]) for i=1:size(ops,1)])
    Jet(f! = JetBlock_f!, df! = JetBlock_df!, df′! = JetBlock_df′!, dom = dom, rng = rng, s = (ops=ops,dom=dom,rng=rng))
end
JopBlock(ops::AbstractMatrix{T}; kwargs...) where {T<:Union{JopLn,JopAdjoint}} = JopLn(JetBlock(ops; kwargs...))
JopBlock(ops::AbstractMatrix{T}; kwargs...) where {T<:Jop} = JopNl(JetBlock(ops; kwargs...))
JopBlock(ops::AbstractVector{T}; kwargs...) where {T<:Jop} = JopBlock(reshape(ops, length(ops), 1); kwargs...)

"""
    JopZeroBlock(dom, rng)

Construct a Jets operator that is equivalent to a matrix of zeros, and that maps from `dom::JetAbstractSpace`
to `rng::JetAbstractSpace`.  This can be useful when forming block operators that contain zero blocks.
"""
JopZeroBlock(dom::JetAbstractSpace, rng::JetAbstractSpace) = JopLn(df! = JopZeroBlock_df!, dom = dom, rng = rng)
JopZeroBlock_df!(d, m; kwargs...) = d .= 0

"""
    iszero(A::Union{Jet, Jop})

Return true if `A` was constructed via `JopZeroBlock`.
"""
Base.iszero(jet::Jet{D,R,typeof(JopZeroBlock_df!)}) where {D<:JetAbstractSpace,R<:JetAbstractSpace} = true
Base.iszero(jet::Jet) = false
Base.iszero(A::Jop) = iszero(jet(A))

macro blockop(ex)
    :(JopBlock($(esc(ex))))
end

"""
    @blockop(ex, kwargs)

Construct a `Jets` block operator, a combination of Jet operators
analogous to a block matrix.

# Examples
example with 1 row-block and 3 column-blocks:
```julia
using Pkg
pkg"add Jets JetPack"
using Jets, JetPack
A = @blockop [JopDiagonal(rand(10)) for icol=1:3]
```

example with 2 row-blocks and 3 column-blocks.
```
using Pkg
pkg"add Jets JetPack"
using Jets, JetPack
A = @blockop [JopDiagonal(rand(10)) for irow=1:2, icol=1:3]
```
"""
macro blockop(ex, kw)
    :(JopBlock($(esc(ex)); $(esc(kw))))
end

macro blockop(ex, kw1, kw2)
    :(JopBlock($(esc(ex)); $(esc(kw1)), $(esc(kw2))))
end

function JetBlock_f!(d::AbstractArray, m::AbstractArray; ops, dom, rng, kwargs...)
    local dtmp
    if size(ops, 2) > 1
        dtmp = zeros(range(ops[1,1]))
    end
    for iblkrow = 1:size(ops, 1)
        _d = getblock(d, iblkrow)
        if size(ops, 2) > 1
            dtmp = size(dtmp) == size(range(ops[iblkrow,1])) ? dtmp : zeros(range(ops[iblkrow,1]))
        end
        for iblkcol = 1:size(ops, 2)
            _m = getblock(m, iblkcol)
            if size(ops, 2) > 1
                _d .+= mul!(dtmp, ops[iblkrow,iblkcol], _m)
            else
                mul!(_d, ops[iblkrow,iblkcol], m)
            end
        end
    end
    d
end

function JetBlock_df!(d::AbstractArray, m::AbstractArray; ops, dom, rng, kwargs...)
    local dtmp
    if size(ops, 2) > 1
        dtmp = zeros(range(ops[1,1]))
    end
    for iblkrow = 1:size(ops, 1)
        _d = getblock(d, iblkrow)
        if size(ops, 2) > 1
            dtmp = size(dtmp) == size(range(ops[iblkrow,1])) ? dtmp : zeros(range(ops[iblkrow,1]))
        end
        for iblkcol = 1:size(ops, 2)
            _m = getblock(m, iblkcol)
            if !iszero(ops[iblkrow,iblkcol])
                if size(ops, 2) > 1
                    _d .+= mul!(dtmp, JopLn(ops[iblkrow,iblkcol]), _m)
                else
                    mul!(_d, JopLn(ops[iblkrow,iblkcol]), _m)
                end
            end
        end
    end
    d
end

function JetBlock_df′!(m::AbstractArray, d::AbstractArray; ops, dom, rng, kwargs...)
    local mtmp
    if size(ops, 1) > 1
        mtmp = zeros(domain(ops[1,1]))
    end
    for iblkcol = 1:size(ops, 2)
        _m = getblock(m, iblkcol)
        if size(ops, 1) > 1
            _m .= 0
            mtmp = size(mtmp) == size(domain(ops[1,iblkcol])) ? mtmp : zeros(domain(ops[1,iblkcol]))
        end
        for iblkrow = 1:size(ops, 1)
            _d = getblock(d, iblkrow)
            if !iszero(ops[iblkrow,iblkcol])
                if size(ops, 1) > 1
                    _m .+= mul!(mtmp, (JopLn(ops[iblkrow,iblkcol]))', _d)
                else
                    mul!(_m, (JopLn(ops[iblkrow,iblkcol]))', _d)
                end
            end
        end
    end
    m
end

function point!(j::Jet{D,R,typeof(JetBlock_f!)}, mₒ::AbstractArray) where {D<:JetAbstractSpace,R<:JetAbstractSpace}
    ops = state(j).ops
    dom = domain(j)
    for icol = 1:size(ops, 2), irow = 1:size(ops, 1)
        point!(jet(ops[irow,icol]), getblock(mₒ, icol))
    end
    j
end

"""
    nblocks(A[, i])

Return the number of blocks in the range and domain of the `Jets` block operator `A::Union{Jet, Jop}`.
With `i` specified, return `nblocks(range(jet))` for `i = 1` and return `nblocks(domain(jet))` otherwise.
"""
nblocks(jet::Jet) = (nblocks(range(jet)), nblocks(domain(jet)))
nblocks(jet::Jet, i) = i == 1 ? nblocks(range(jet)) : nblocks(domain(jet))
nblocks(A::Jop) = (nblocks(range(A)), nblocks(domain(A)))
nblocks(A::Jop, i) = i == 1 ? nblocks(range(A)) : nblocks(domain(A))

"""
    getblock(A, i, j)

Return the block of the Jets block operator `A` that corresponds to row block `i`
and column block `j`.
"""
getblock(jet::Jet{D,R,typeof(JetBlock_f!)}, i, j) where {D,R} = state(jet).ops[i,j]
getblock(A::JopLn{T}, i, j) where {T<:Jet} = JopLn(getblock(jet(A), i, j))
getblock(A::JopNl{T}, i, j) where {T<:Jet} = getblock(jet(A), i, j)
getblock(A::T, i, j) where {J<:Jet,T<:JopAdjoint{J}} = getblock(A.op, j, i)'
getblock(::Type{JopNl}, A::Jop{T}, i, j) where {T<:Jet} = getblock(jet(A), i, j)::JopNl
getblock(::Type{JopLn}, A::Jop{T}, i, j) where {T<:Jet} = JopLn(getblock(jet(A), i, j))

"""
    isblockop(A)

Return true if `A` is a Jets block operator.
"""
isblockop(A::Jop{<:Jet{<:JetAbstractSpace,<:JetAbstractSpace,typeof(JetBlock_f!)}}) = true
isblockop(A::Jop) = false

function getblock(jet::Jet{D,R,typeof(JetComposite_f!)}, i, j) where {D,R}
    ops = []
    for op in state(jet).ops
        if isblockop(op)
            push!(ops, getblock(op, i, j))
        else
            push!(ops, op)
        end
    end
    mapreduce(identity, ∘, ops)
end

Base.reshape(x::AbstractArray, R::JetBSpace) = BlockArray([reshape(view(x, R.indices[i]), R.spaces[i]) for i=1:length(R.indices)], R.indices)

# BlockArrays are vectors
function Base.reshape(x::BlockArray, R::JetBSpace)
    length(x) == length(R) || error("dimension mismatch, unable to reshape block array")
    x
end

function Base.close(j::Jet{D,R,typeof(JetBlock_f!)}) where {D,R}
    ops = state(j).ops
    close.(ops)
    nothing
end

#
# Vectorized operator
#
JetVec(op::Jop) = Jet(f! = JetVec_f!, df! = JetVec_df!, df′! = JetVec_df′!, dom = vec(domain(op)), rng = vec(range(op)), s=(op=op,))
JetVec(op::Jop{<:J}) where {T,U,D<:JetAbstractSpace{T,1},R<:JetAbstractSpace{T,1},J<:Jet{D,R}} = op
JopVec(op::Union{JopLn,JopAdjoint}) = JopLn(JetVec(op))
JopVec(op::Jop) = JopNl(JetVec(op))

JetVec_f!(d::T, m::AbstractVector; op, kwargs...) where {T<:AbstractVector} = vec(mul!(reshape(d, range(op)), op, reshape(m, domain(op))))::T
JetVec_df!(d::T, m::AbstractVector; op, kwargs...) where {T<:AbstractVector} = vec(mul!(reshape(d, range(op)), JopLn(op), reshape(m, domain(op))))::T
JetVec_df′!(m::T, d::AbstractVector; op, kwargs...) where {T<:AbstractVector} = vec(mul!(reshape(m, domain(op)), JopLn(op)', reshape(d,range(op))))::T

"""
    B = vec(A)

B is equivelent to A except that its domain and range are "vectorized".  This is
useful when calling algorithms that expect vectors in the domain and range of the
operator.  One example of this is the `lsqr` method in the IterativeSolvers package.

# Example
```julia
using Jets, JetPack, IterativeSolvers

A = JopDiagonal(rand(10,11))
d = rand(range(A))
m = lsqr(vec(A), vec(d))
```
"""
Base.vec(op::Jop) = JopVec(op)

#
# multiply operator by a scalar
#
_constdiag_df!(d, m; a, kwargs...) = d .= a * m
_constdiag_df′!(m, d; a, kwargs...) = m .= conj(a) * d
function Base.:*(a::Number, A::Jop)
    _a = JopLn(dom = domain(A), rng = domain(A), df! = _constdiag_df!, df′! = _constdiag_df′!, s=(a=a,))
    _a ∘ A
end

#
# utilities
#
"""
    convert(Array, A::JopLn)

Convert a linear Jets operator into its equivalent matrix.
"""
function Base.convert(::Type{T}, A::Jop) where {T<:Array}
    m = zeros(domain(A))
    d = zeros(range(A))
    B = zeros(eltype(A), size(A))
    for icol = 1:size(A, 2)
        m .= 0
        d .= 0
        m[icol] = 1
        B[:,icol] .= mul!(d, A, m)[:]
    end
    B
end

"""
    lhs,rhs = dot_product_test(A, m, d; mmask, dmask)

Compute and return the left and right hand sides of the *dot product test*:

`<d,Am> ≈ <Aᴴd,m>`

Here `Aᴴ` is the conjugate transpose or adjoint of `A`, and `<x, y>` denotes the inner
product of vectors `x` and `y`. The left and right hand sides of the dot product test are
expected to be equivalent close to machine precision for operator `A`. If the equality does not
hold this can indicate a problem with the implementation of the operator `A`.

This function provides the optional named arguments `mmask` and `dmask` which are vectors in the
domain and range of `A` that are applied via elementwise multiplication to mask the vectors
`m` and `d` before applying of the operator, as shown below. Here we use `∘` to represent
the Hadamard product (elementwise multiplication) of two vectors.

`<dmask ∘ d, A (mmask ∘ m)> ≈ <Aᵀ (dmask ∘ d), mmask ∘ m>`

You can test the relative accuracy of the operator with this relation for the left hand side `lhs` and
right hand side `rhs` returned by this function:

`|lhs - rhs| / |lhs + rhs| < ϵ`
"""
function dot_product_test(op::JopLn, m::AbstractArray, d::AbstractArray; mmask=[], dmask=[])
    mmask = length(mmask) == 0 ? ones(domain(op)) : mmask
    dmask = length(dmask) == 0 ? ones(range(op)) : dmask

    ds = op * (mmask .* m)
    ms = op' * (dmask .* d)

    lhs = dot(mmask.*m, ms)
    rhs = dot(ds, dmask.*d)

    if eltype(lhs) <: Complex && eltype(rhs) <: Complex
        return lhs, rhs
    else
        return real(lhs), real(rhs)
    end
end

"""
    μobs, μexp = linearization_test(F, mₒ; μ)

Thest that the jacobian, `J`, of `F` satisfies the Taylor expansion:

`F(m) = F(m_o) + F'(m_o)δm + O(δm^2)`
"""
function linearization_test(F::JopNl, mₒ::AbstractArray;
        μ=[1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125], δm = [], mmask=[], dmask = [], seed=Inf)
    mmask = length(mmask) == 0 ? ones(domain(F)) : mmask
    dmask = length(dmask) == 0 ? ones(range(F)) : dmask

    isfinite(seed) && Random.seed!(seed)
    if length(δm) == 0
        δm = mmask .* (-1 .+ 2 .* rand(domain(F)))
    else
        δm .*= mmask
    end

    Fₒ = F*mₒ
    Jₒ = jacobian!(F, mₒ)
    Jₒδm = Jₒ*δm

    μ = convert(Array{eltype(mₒ)}, sort(μ, rev=true))

    μobs = zeros(length(μ)-1)
    μexp = zeros(length(μ)-1)
    ϕ = zeros(length(μ))
    for i = 1:length(μ)
        d_lin  = Fₒ .+ μ[i] .* Jₒδm
        d_non  = F*(mₒ .+ μ[i] .* δm)
        ϕ[i] = norm(dmask .* (d_non .- d_lin))
        if i>1
            μobs[i-1] = ϕ[i-1]/ϕ[i]
            μexp[i-1] = (μ[i-1]/μ[i]).^2
        end
    end
    μobs, μexp
end

"""
    lhs,rhs = linearity_test(A::Jop)

test the the linear Jet operator `A` satisfies the following test
for linearity:

`A(m_1+m_2)=Am_1 + A_m2`
"""
function linearity_test(A::Union{JopLn,JopAdjoint}, m1 = [], m2 = [])
    m1 = length(m1) == 0 ? -1 .* 2 .* rand(domain(A)) : m1
    m2 = length(m2) == 0 ? -1 .* 2 .* rand(domain(A)) : m2
    lhs = A*(m1 + m2)
    rhs = A*m1 + A*m2
    lhs, rhs
end

# for hashing models <--
CRC32c.crc32c(m::Array{<:Union{UInt32,Float32,Float64,Complex{Float32},Complex{Float64}}}) = CRC32c.crc32c(unsafe_wrap(Array, convert(Ptr{UInt8}, pointer(m)), length(m)*sizeof(eltype(m)), own=false))
#-->

export Jet, JetAbstractSpace, JetBSpace, JetSpace, JetSSpace, Jop, JopAdjoint, JopLn, JopNl,
    JopZeroBlock, @blockop, domain, getblock, getblock!, dot_product_test, getblock,
    getblock!, indices, jacobian, jacobian!, jet, linearity_test, linearization_test,
    nblocks, perfstat, point, setblock!, shape, space, state, state!, symspace

end


text:
using StaticArrays

function topolar(pos)
    theta = atan(pos[2]/pos[1])
    r = sqrt(sum(abs2, pos))
    return Polar(r,theta)
end

function tocart(pos::Polar)
    x = pos[1]*cos(pos[2])
    y = pos[1]*sin(pos[2])
    return SA[x,y]
end

function reflect(pos,r)
    pos1 = topolar(pos)
    dist = pos1[1] - r
    pos2 = SA[r - dist,pos1[2]]
    posf = tocart(pos2)
    return posf
end

#=
function slope(pos1,pos2)               #finds the slope between two points
    return (pos1[2]-pos2[2])/(pos1[1]-pos2[1])
end

function midpoint(pos1,pos2)            #finds the midpoint between two points
    x = (pos1[1]+pos2[1])/2
    y = (pos1[2]+pos2[2])/2
    return SA[x,y]
 end

function distance(pos1,pos2)            #finds the distance between two points
    return ((pos2[1]-pos1[1])^2+(pos2[2]-pos1[2])^2)^(1/2)
end

function circleintersect(pos1,pos2,r)   #estimates the intersection between the line formed by two points and a circle of radius r
    if(!(((pos1[1]^2)+(pos1[2]^2))^(1/2) < r) || !(((pos2[1]^2)+(pos2[2]^2))^(1/2) > r))
throw(DomainError("pos1 must be inside the circle and pos2 must be outside the circle"))
    end
    mp = midpoint(pos1,pos2)
    if(isapprox(((pos1[1]^2)+(pos1[2]^2))^(1/2),r))
        return pos1
    elseif(((mp[1]^2)+(mp[2]^2))^(1/2) < r)
        circleintersect(mp, pos2, r)
    else
        circleintersect(pos1, mp, r)
    end
end

function slope(pos1,pos2)               #finds the slope between two points
    return (pos1[2]-pos2[2])/(pos1[1]-pos2[1])
end

function midpoint(pos1,pos2)            #finds the midpoint between two points
    x = (pos1[1]+pos2[1])/2
    y = (pos1[2]+pos2[2])/2
    return SA[x,y]
 end

function circleintersect(pos1,pos2,r)   #estimates the intersection between the line formed by two points and a circle of radius r
    mp = midpoint(pos1,pos2)
    if(isapprox(((pos1[1]^2)+(pos1[2]^2))^(1/2),r))
        return pos1
    elseif(((mp[1]^2)+(mp[2]^2))^(1/2) < r)
        circleintersect(mp, pos2, r)
    else
        circleintersect(pos1, mp, r)
    end
end

function rslope(posC)
    return posC[2]/posC[1]
end

function perpslope(m)
    return -1/m
end

function solve(m1,m2,x1,y1) #slope of radius, slope of perpendicular, point on perpendicular line (x and y)
    A = [-m1 1;-m2 1]
    b = [0, y1 - m2*x1]
    temp = A\b
    return SA[temp[1],temp[2]]
end

function posFinal(pos1,posM)
    deltax = posM[1] - pos1[1]
    deltay = posM[2] - pos1[2]
    return SA[posM[1] + deltax, posM[2] + deltay]
end

#=
function circleslope(pos1)              #finds the slope of a circle at a point
    return -pos1[1]/pos1[2]
end

function angleofincidence(pos1,pos2,r)  #finds the angle of incidence of the line between pos1 and pos 2 as it reflects of a circle with radius r
    posC = circleintersect(pos1,pos2,r)
    m1 = slope(pos1,pos2)
    m2 = circleslope(posC)
    return atan((m1-m2)/(1+m1*m2))
end

function returnslope(pos1,pos2,r)       #finds the slope of the line between pos1 and pos 2 after reflection on a circle with radius r
    theta = angleofincidence(pos1,pos2,r)
    posC = circleintersect(pos1, pos2, r)
    m2 = circleslope(posC)
    return (tan(-theta)+m2)/(1-tan(-theta)*m2)
end

function finallocation(pos1,pos2,r)     #currently broken-tries to find where a particle moving from pos1 to pos2 would reflect to after hitting a circle with radius r
    m = returnslope(pos1,pos2,r)
    posC = circleintersect(pos1,pos2,r)
    dist = distance(posC,pos2)
    posTemp1 = SA[posC[1] + 100000, posC[2] + 100000*m]
    posFinalA = circleintersect(posC, posTemp1, dist)       #the broken part is here-circle intersect uses a circle centered on 0, I need a circle centered on posC
    posTemp2 = SA[posC[1] - 100000, posC[2] - 100000*m]
    posFinalB = circleintersect(posC, posTemp2, dist)
    if(((posFinalA[1]^2)+(posFinalA[2]^2))^(1/2) < r)
        return posFinalA
    else
        return posFinalB
    end
end
=#
=#
using Test

@testset "Diffusion.jl" begin
#   @test slope(SA[1,2],SA[3,4]) == 1
#   @test midpoint(SA[1,1],SA[3,3]) == SA[2,2]
#   @test distance(SA[1,1],SA[2,1]) == 1
#   @test isapprox(circleintersect(SA[1,0], SA[7,0], 2),SA[2,0])
#   @test circleslope(SA[1,2]) == -0.5
#   @test angleofincidence(SA[1,1], SA[3,3], 2) == 90π/180
#   @test isapprox(returnslope(SA[1,1], SA[3,3], 2), 1)
#   I don't know how to test final location
#   @test rslope(SA[2,3]) == 3/2
#   @test perpslope(3) == -1/3
#   @test solve(1,-2,0,3) == SA[1,1]
#   @test posFinal(SA[1,2],SA[1.5,1.5]) == SA[2,1]
#   @test topolar(SA[3,4]) == SA[5, 0.93]
#   @test topolar(SA[-12,-5]) == SA[13, 3.54]
#   @test topolar(SA[-3,10]) == SA[10.4, 1.86]
#   @test topolar(SA[5,-8]) == SA[9.4, 5.27]
#   @test tocart(SA[5, 0.93]) == SA[3,4]
#   @test tocart(SA[13, 3.54]) == SA[-12,-5]
#   @test tocart(SA[10.4, 1.86]) == SA[-3,10]
#   @test tocart(SA[9.4, 5.27]) == SA[5,-8]
    @test reflect(SA[4,4], sqrt(18)) == SA[2,2]
end


text:
"""
    nll,nllc = rd_dico()
Retourne nll la liste des bons mots, et nllc la liste des mots rejetés
"""
function rd_dico()
    ll = Array{String,1}()
    # https://infolingu.univ-mlv.fr/DonneesLinguistiques/Dictionnaires/telechargement.html
    f = open("dela-fr-public.dic",enc"UTF-16LE","r")
    for l in eachline(f)
        push!(ll, l)
    end
    close(f)
    nll = Array{String,1}()
    nllc = Array{String,1}()
    function mot(l)
        i = findfirst(",",l)[1]
        nl = l[1:prevind(l, i, 1)]
    end
    for l in ll
    	nl = mot(l)
   	if isnothing(findfirst(" ",nl)) && isnothing(findfirst("-",nl))
            push!(nll, nl)
        else
            push!(nllc, nl)
        end
    end
    (nll,nllc)
end

nll,nllc = rd_dico()

"""
    setoflettres(nll)
Retourne l'ensemble des lettres utilisées dans les mots de la liste nll
"""
function setoflettres(nll)
    lettres = Set{Char}()
    for l in nll
        for c in l
            push!(lettres, c)
        end
    end
    lettres
end

"""
    dico,sol = dicoflower(sl)
Retourne le dictionnaire dico lettre -> lettre minuscule
et l'ensemble des lettres minuscules
"""
function dicoflower(sl)
    dico = Dict{Char,Char}()
    sol = Set{Char}()
    for c in sl
        cl = lowercase(c)
        dico[c] = cl
        push!(sol, cl)
    end
    dico,sol
end

sl = setoflettres(nll)
sll,sol = dicoflower(sl)

"""
Dictionnaire des lettres accentuées
"""
dicos = Dict{Char,Union{Char,Missing}}()
dicos['\''] = '\'';
dicos['-'] = missing;
dicos['0'] = missing;
dicos['1'] = missing;
dicos['\\'] = missing;
dicos['a'] = 'a';
dicos['b'] = 'b';
dicos['c'] = 'c';
dicos['d'] = 'd';
dicos['e'] = 'e';
dicos['f'] = 'f';
dicos['g'] = 'g';
dicos['h'] = 'h';
dicos['i'] = 'i';
dicos['j'] = 'j';
dicos['k'] = 'k';
dicos['l'] = 'l';
dicos['m'] = 'm';
dicos['n'] = 'n';
dicos['o'] = 'o';
dicos['p'] = 'p';
dicos['q'] = 'q';
dicos['r'] = 'r';
dicos['s'] = 's';
dicos['t'] = 't';
dicos['u'] = 'u';
dicos['v'] = 'v';
dicos['w'] = 'w';
dicos['x'] = 'x';
dicos['y'] = 'y';
dicos['z'] = 'z';
dicos['à'] = 'a';
dicos['â'] = 'a';
dicos['ã'] = 'a';
dicos['ä'] = 'a';
dicos['ç'] = 'c';
dicos['è'] = 'e';
dicos['é'] = 'e';
dicos['ê'] = 'e';
dicos['ë'] = 'e';
dicos['î'] = 'i';
dicos['ï'] = 'i';
dicos['ñ'] = 'n';
dicos['ó'] = 'o';
dicos['ô'] = 'o';
dicos['õ'] = 'o';
dicos['ö'] = 'o';
dicos['ù'] = 'u';
dicos['û'] = 'u';
dicos['ü'] = 'u';
dicos[''] = missing;

"""
Dictionnaire Letre quelconque -> minuscule non accentuée
"""
DICASCII = Dict{Char,Union{Char,Missing}}()
for k in sll
    DICASCII[k[1]] = dicos[k[2]]
end

arrg = Array{Array{UInt8,1},1}()
setg = Set{String}()

for m in nll
    nm = Array{UInt8,1}(undef,length(m))
    nsm = ""
    i = 0
    ok = true
    for c in m
        i += 1
        ca = DICASCII[c]
        if ismissing(ca)
            println(m)
            ok = false
            continue
        else
            nm[i] = ca
            nsm *= ca
        end
    end
    if ok
        if !(nsm in setg)
            push!(arrg, nm)
        end
        push!(setg, nsm)
    end
end

"""
a = Array{UInt8,1}(['a','b','c'])
length(a)     # 3
s = String(a) # "abc"
length(a)     # 0 Bug or feature ? feature Cf. doc.
"""

myString(a) = String(copy(a))

# voir le pb. de prudh'homme

myString(arrg[312157])

struct Momo
    c::UInt8
    t::Array{Momo,1}
end

function pushtomomo!(momo::Momo, m::Array{UInt8,1})
    if length(m) == 0
        push!(momo.t, Momo(0x00, []))
        return
    end
    c = m[1]
    found = false
    for mt in momo.t
        if c == mt.c
            pushtomomo!(mt, m[2:end])
            found = true
            break
        end
    end
    if !found
        nmomo = Momo(c, [])
        pushtomomo!(nmomo, m[2:end])
        push!(momo.t, nmomo)
    end
end

function buildmomo()
    alpha = Momo(0x00, [])
    for m in arrg
        pushtomomo!(alpha, m)
    end
    alpha
end

function findc(momo::Momo, c::UInt8)
    for nm in momo.t
        if nm.c == c
            return nm
        end
    end
    nothing
end

function findm(momo::Momo, m::Array{UInt8,1})
    if length(m) == 0
        return momo
    end
    nm = findc(momo, m[1])
    if nm == 0x00
        return nothing
    end
    return findm(nm, m[2:end])
end

function explore(out::String, momo::Momo, c::UInt8)
    nm = findc(momo, c)
    if !isnothing(nm)
        return out*Char(nm.c),nm
    else
        return out,nm
    end
end

# Facilités
findc(momo::Momo, c::Char) = findc(momo, UInt8(DICASCII[c]))


text:
# This file was generated, do not modify it. # hide
coerce!(df0, :TotalCharges => Continuous);

max_stars_repo_name: alan-turing-institute/MLJTutorials

max_issues_repo_name: alan-turing-institute/MLJTutorials

max_forks_repo_name: alan-turing-institute/MLJTutorials


text:
#= Install the primes package: (if unavailable)
import Pkg; Pkg.add("Primes") =#
using Primes

# Public key: (e,n)
struct PublicKey
    n::Integer
    e::Integer
end

# Private key: (d,n)
struct PrivateKey
    n::Integer
    d::Integer
end

function keyGenerator(p::Integer, q::Integer, e::Integer)
    # Basic corner cases:
    isprime(p) || error("p must be a prime number")
    isprime(q) || error("q must be a prime number")
    e > 0 || error("e must be positive")
    # Computing Euler's totient:
    phi = (p - 1) * (q - 1)
    # Computing greatest common divisor and Bézout coefficients (i.e. u and v, such that u*e + u*phi = gcd(e,phi)):
    (g, u, v) = gcdx(e,phi)
    # If gcd(e,phi)!=1, they aren't coprime:
    g == 1 || error("phi(n) and e must be coprime")
    # Computing private exponent d based on (+ve)/(-ve) value of u through a ternary operator:
    u < 0 ? d = (u % phi) + phi : d = u % phi
    # Computing n as the product of the two entered primes and returning the computed public and private key components:
    n = p * q
    return PublicKey(n, e), PrivateKey(n, d)
end

# Encryption function:
function encrypt(m::Integer, k::PublicKey)
    0 <= m || error("m must be non-negative")
    m < k.n || error("m is too large")
    return powermod(m, k.e, k.n)
end

# Decryption function:
function decrypt(c::Integer, k::PrivateKey)
    0 <= c || error("c must be non-negative")
    c < k.n || error("c is too large")
    return powermod(c, k.d, k.n)
end

print("Enter a prime number: ")
p = parse(Int, readline())
print("Enter another prime number: ")
q = parse(Int, readline())
print("Enter public key exponent: ")
e = parse(Int, readline())
(publicKey, privateKey) = keyGenerator(p, q, e)
println("Public key (n,e) = ($(publicKey.n),$(publicKey.e)); (to be shared)")
println("Private key (n,d) = ($(privateKey.n),$(privateKey.d)); (to be kept with oneself)")
println("Enter an integer to encrypt:")
m = parse(Int, readline())
println("Original integer: $m")
c = encrypt(m, publicKey)
println("Encrypted integer: $c")
println("Decrypted integer: $(decrypt(c,privateKey))")


text:
# --------------------------------------------------------------------------
# ACE.jl: Julia implementation of the Atomic Cluster Expansion
# Copyright (c) 2019 Christoph Ortner <christophortner0@gmail.com>
# All rights reserved.
# --------------------------------------------------------------------------

module Bonds

include("cylindrical_coords.jl")
include("fourier.jl")

include("bondalist.jl")

include("bond_basis.jl")

# include("bond_eval.jl")

include("utils.jl")

end


text:
# Multinomial distribution

export Multinomial

@parameterized Multinomial(n, p)

basemeasure(d::Multinomial{(:n, :p)}) = CountingMeasure(ℤ) ^ length(d.p)

@kwstruct Multinomial(n,p)

function logdensity(d::Multinomial{(:n, :p)}, x)
    p = d.p
    s = 0.0
    for j in eachindex(x)
        s += x[j] * log(p[j])
    end
    return s
end

Base.rand(rng::AbstractRNG, T::Type, μ::Multinomial) = rand(rng, Dists.Multinomial(μ.n, μ.p))

distproxy(d::Multinomial{(:p,)}) = Dists.Multinomial(d.n,d.p)

# Based on
# https://github.com/JuliaMath/Combinatorics.jl/blob/c2114a71ccfc93052efb9a9379e62b81b9388ef8/src/factorials.jl#L99
function logmultinomial(k)
    s = 0
    result = 1
    @inbounds for i in k
        s += i
        (Δresult, _) = logabsbinomial(s, i)
        result += Δresult
    end
    result
end

function testvalue(d::Multinomial{(:n,:p)})
    n = d.n
    l = length(d.p)
    q,r = divrem(n, l)
    x = fill(q, l)
    x[1] += r
    return x
end


text:
# This file includes baseband modulator

export Modulator, scheme, alphabet, modulate, constellation, basis, iswaveform

"""
    $TYPEDEF

Baseband digital modulator.

!!! note `tsample` determines the type of modulation. If `tsample` is `NaN`, the modulator transmits the vectors whose
    elements are the coefficients of the expansion of the symbol waveform over the orthogonal basis. Otherwise, the
    modulator trasmits the symbol waveform sampled by `tsample`.

# Fields

    $TYPEDFIELDS
"""
struct Modulator{T<:AbstractScheme, S<:AbstractPulse}
    "Modulation scheme"
    scheme::T
    "Modulation pulse energy"
    pulse::S
    "Sampling time."
    tsample::Float64
end
Modulator(scheme) = Modulator(scheme, RectangularPulse(), NaN)

show(io::IO, modulator::Modulator) = print(io,
    "Modulator(scheme:$(modulator.scheme), pulse:$(modulator.pulse))")

"""
    $SIGNATURES

Returns true if modulator a waveform modulator
"""
iswaveform(modulator::Modulator) = modulator.tsample !== NaN

# Modulation
(modulator::Modulator)(ms) = iswaveform(modulator) ? wavemodulate(modulator, ms) : vecmodulate(modulator, ms)

function vecmodulate(modulator::Modulator{T,S}, stream) where {T, S}
    # Note: In this case, `modulator` maps each element of `stream` to a vector whose elements(possibly complex) are the
    # coefficients of the expansion of symbol waveform over the basis of the modulator.
    scheme = modulator.scheme
    pulse = modulator.pulse
    Eg = energy(pulse)
    T <: FSK ? alphabet(scheme)[stream] * sqrt(2 * Eg) : alphabet(scheme)[stream] * sqrt(Eg)
end

function wavemodulate(modulator::Modulator{T, S}, symbols) where {T, S}
    # Note: In this case, `modulator` maps each element of `stream` to a waveform that is spanned by the basis of the
    # modulator.
    pulse = modulator.pulse
    ts = modulator.tsample
    tp = pulse.duration
    t = 0 : ts : tp - ts
    basisvals = map(base -> base.(t), basis(modulator)) # Precomputation of the values of the basis functions
    si = map(symbol -> sum(alphabet(modulator.scheme)[symbol] .* basisvals), symbols)
    Eg = energy(pulse)
    T <: FSK ? si * sqrt(2 * Eg)  : si * sqrt(Eg)
end

"""
    $SIGNATURES

Returns basis `modulator`
"""
function basis(modulator::Modulator{T, S}) where {T, S}
    iswaveform(modulator) || error("Expected waveform modulator, got a vector modulator")
    pulse = modulator.pulse
    Eg = energy(pulse)
    M = constelsize(modulator.scheme)
    if T <: FSK
        # Multidimensional signalling
        Δf = 1 / pulse.duration  # Minimum frequency sepeartion for orthogonality of the basis.
        map(m -> (t -> pulse(t) * exp(1im * 2π * m * Δf * t)), 1 : M)
    else
        # Two dimensional signalling
        [t -> pulse(t) / sqrt(Eg)]
    end
end

"""
    $SIGNATURES

Returns scheme of the modulation `modulator`.
"""
scheme(modulator::Modulator) = modulator.scheme

"""
    $SIGNATURES

Returns the constellation size of `modulator`
"""
constelsize(modulator::Modulator) = constelsize(modulator.scheme)


text:
# NOTE: This file is mainly for records-keeping purposes, and the interface may have evolved.
# User settings.
  # Time and space grid.
  z_min = 0.0
  z_max = 5.0
  P = 1000
  T = 20.0
# Experiment settings.
  d_0 = 5.0
  d_T = 2.3701
  Δ_E = 1e-06
# Overall parameters.
  params = (ρ = 0.02, σ = 4.2508, N = 10, θ = 5.1269, γ = 1.00, κ = 0.013, ζ = 1, η = 0, Theta = 1, χ = 1/(2.1868), υ = 0.0593, μ = 0, δ = 0.053)
# Solver settings.
  tstops = nothing

# Construct intermediate objects.
z_ex = range(z_min, stop = z_max, length = P+2)
z = z_ex[2:end-1]
settings = (z = z, z_ex = z_ex, tstops = tstops, Δ_E = Δ_E, T_U_bar = 1.0)
params_0 = merge(params, (d = d_0,))
params_T = merge(params, (d = d_T,))

# Get (numerical) stationary solution.
stationary_0 = stationary_numerical(params_0, z_ex)
stationary_T = stationary_numerical(params_T, z_ex)

# Process those
Ω_0 = stationary_0.Ω
Ω_T = stationary_T.Ω
v_T = stationary_T.v_tilde
g_T = stationary_T.g
z_hat_T = stationary_T.z_hat

# Define more interim quantities.
Ω(t) = Ω_T # This is constant.
E(t) = (log(Ω(t + Δ_E)) - (log(Ω(t - Δ_E))))/(2*Δ_E) + params.δ # central differences.

@testset "Constant Omega" begin
    sol = solve_dynamics(params_0, stationary_0, settings, T, Ω, E)
  # Spot-checks.
    @test sol.sol.t[5] == 19.852
    @test sol.results[:λ_ii][1] ≈ 0.993204827349257
    @test sol.sol.u[4][3] ≈ 1.1338849871884573
    @test sol.sol.prob.u0[1] ≈ 1.1675290237815987
  # Detailed checks.
    @test sol.results[:g][1] ≈ 0.01739899362603311 # g check.
    @test sol.results[:z_hat][1] ≈  2.827474674128036
    @test sol.results[:z_hat][(end-9)] ≈ 2.8259905922985666
    # Sub-pieces of L_tilde
    @test sol.results[:L_tilde_a] + sol.results[:L_tilde_x] + sol.results[:L_tilde_E] ≈ sol.results[:L_tilde]
end


text:
module DestructuringKoans

function add_separately(x, y)
end

function add_coordinates(___, ___)
end

end


text:
"""
    groupby_dbscan_temporallimit()

Group nearby localizations separated by time into `Molecule`s

Identifies local maxima of `Localization`s temporally distant from other points by `t_off` frames. The effect is to
consider flashes correlated in time as originating from the same molecule. Only localizations within `radius` (same
units as position) of the cluster center are joined into the molecule. The algorithm is repeated on the remaining
unjoined localizations.
"""
function groupby_localmax_temporallimit(localizations::Vector{Localization}, radius, t_off)
    molecules = Molecule[]
    locs = copy(localizations)
    while length(locs) > 0
        neighborsdict = findtemporalneighbors(locs, radius, t_off)
        localmaxima, localmaxima_map = findlocalmaxima(locs, neighborsdict)
        setdiff!(locs, localmaxima)
        remaining_localmaxima_map = IdDict(Pair(k, localmaxima_map[k]) for k ∈ setdiff(keys(localmaxima_map), localmaxima))
        newmolecules = map(l -> Molecule(l), localmaxima)
        buildmolecules!(newmolecules, locs, remaining_localmaxima_map, radius)
        append!(molecules, newmolecules)
    end

    molecules
end

"""
    merge_close_molecules()

Merge closely clustered molecules (which may be physically associated) into single molecules.
"""
function merge_close_molecules(molecules::Vector{Molecule}, radius)
    mergedmolecules = Molecule[]

    while length(molecules) > 0
        neighborsdict = findneighbors(molecules, radius)
        localmaxima, localmaxima_map = findlocalmaxima(molecules, neighborsdict)
        setdiff!(molecules, localmaxima)
        remaining_localmaxima_map = IdDict([Pair(k, localmaxima_map[k]) for k ∈ setdiff(keys(localmaxima_map), localmaxima)])
        buildmolecules!(localmaxima, molecules, remaining_localmaxima_map, radius)
        append!(mergedmolecules, localmaxima)
    end

    reindex!(mergedmolecules)
    mergedmolecules
end

"""
    findtemporalneighbors()

Find `Localiation`s that occur near each other within `t_off` frames of each other.
"""
function findtemporalneighbors(localizations::Vector{Localization}, radius, t_off)
        neighborsdict = findneighbors(localizations, radius)

        function istemporallyclose(localization1, localization2)
            if localization1.frame < localization2.frame
                loc1end = localization1.frame + localization2.length - 1
                return localization2.frame - loc1end ≤ t_off
            else
                loc2end = localization2.frame + localization2.length - 1
                return localization1.frame - loc2end ≤ t_off
            end
        end

        foreach(l -> filter!(n -> istemporallyclose(n, l), neighborsdict[l]), localizations)
        neighborsdict
end

"""
    findneighbors()

Find `Localiation`s or `Molecule`s within `radius` of each other.
"""
function findneighbors(points::Vector{<:DataEntity}, radius)
    coordinates = extractcoordinates(points)
    neighbortree = BallTree(coordinates)
    ineighbors = inrange(neighbortree, coordinates, radius, true)
    neighbors = map(i -> points[i], ineighbors)
    IdDict(zip(points, neighbors))
end

findnearestneighbor(neighbors::Vector{<:DataEntity}, molecule::Molecule) =
    findnearestneighbor(neighbors, molecule.group)

"""
    findnearestneighbor()

Find the closest `Localiation` or `Molecule` to another.
"""
function findnearestneighbor(neighbors::Vector{<:DataEntity}, point::DataEntity)
    coordinates = extractcoordinates(point)
    neighborcoordinates = extractcoordinates(neighbors)
    neighbortree = BallTree(neighborcoordinates)
    inearestneighbor, nearestneighbor_distance = nn(neighbortree, coordinates)
    neighbors[inearestneighbor], nearestneighbor_distance
end

"""
    findnearestunvisitedmaxneighbor()

Find the nearest neighbor that hadn't been previously visited.
"""
function findnearestunvisitedmaxneighbor(entity, neighborsdict, visited)
    neighbor_counts = length.(values(neighborsdict))
    maxneighbor_indexes = argmaxall(neighbor_counts)
    maxneighbors = collect(keys(neighborsdict))[maxneighbor_indexes]

    # if there are multiple neighbors with max count, choose closest
    if length(maxneighbors) > 1
        setdiff!(maxneighbors, visited)
        return length(maxneighbors) > 0 ?
            first(findnearestneighbor(maxneighbors, entity)) :
            entity
    end

    return first(maxneighbors)
end

"""
    create_neighborsofneighborsdict()

Create a dictionary containing the 2nd-level neighbors of each `Localization` or `Molecule`.
"""
function create_neighborsofneighborsdict(entities::Vector{T}, neighborsdict::IdDict{T, Vector{T}}) where T <: DataEntity
    neighborsofneighborsdict = IdDict{T, IdDict{T, Vector{T}}}()
    foreach(e -> neighborsofneighborsdict[e] = IdDict(Pair(k, neighborsdict[k]) for k ∈ neighborsdict[e]), entities)
    neighborsofneighborsdict
end

"""
    findlocalmaxima()

Find the `Localization`s or `Molecule`s corresponding to local density maxima.
"""
function findlocalmaxima(entities::Vector{T}, neighborsdict::IdDict{T, Vector{T}}) where T <: DataEntity
    localmaxima = T[]

    # map of each localization to the index of its associated local maximum
    localmaxima_map = IdDict{T, Int}()

    updatemap(chain, i) =
        foreach(x -> localmaxima_map[x] = i, chain)

    isvisited(entity) =
        haskey(localmaxima_map, entity)

    neighborsofneighborsdict = create_neighborsofneighborsdict(entities, neighborsdict)

    for entity ∈ entities
        isvisited(entity) && continue

        visited_chain = T[]

        while true
            push!(visited_chain, entity)

            localneighborsdict = neighborsofneighborsdict[entity]

            !isempty(localneighborsdict) || break

            maxneighbor = findnearestunvisitedmaxneighbor(entity, localneighborsdict, visited_chain)

            if (reached_maximum = maxneighbor === entity)
                push!(localmaxima, entity)
                updatemap(visited_chain, entity.index)
                break
            else
                if (joined_chain = isvisited(maxneighbor))
                    localmaxindex = localmaxima_map[maxneighbor]
                    push!(localmaxima_map, maxneighbor => localmaxindex)
                    updatemap(visited_chain, localmaxindex)
                    break
                end

                entity = maxneighbor
            end
        end
    end
    (localmaxima, localmaxima_map)
end

"""
    buildmolecules!()

Build grouped molecules from the molecule seeds and the other entities (`Localization`s or `Molecule`s) to join.
"""
function buildmolecules!(seedmolecules::Vector{Molecule}, entities::Vector{T},
    localmaxima_map::IdDict{T, Int}, radius) where T <: DataEntity

    moleculeneighbors_map = invertdictionary(localmaxima_map)
    entityindexes = IdDict(zip(entities, eachindex(entities)))

    todelete = Int[]

    for molecule ∈ seedmolecules
        haskey(moleculeneighbors_map, molecule.index) || continue

        clusterneighbors = moleculeneighbors_map[molecule.index]
        clusterneighborcoordinates = extractcoordinates(clusterneighbors)
        clusterneighborsused = falses(length(clusterneighbors))
        clusterneighbortree = BallTree(clusterneighborcoordinates)
        coordinates = extractcoordinates(molecule)

        for i ∈ eachindex(clusterneighbors)
            inearestneighbor, nearestneighbor_distance = nn(clusterneighbortree, coordinates, j -> clusterneighborsused[j])
            nearestneighbor = clusterneighbors[inearestneighbor]

            isinrange = nearestneighbor_distance < radius + nearestneighbor.accuracy + molecule.accuracy

            isinrange || break

            push!(molecule, nearestneighbor)
            clusterneighborsused[inearestneighbor] = true
            push!(todelete, entityindexes[nearestneighbor])
            # try batching this one more time; filter not efficient
        end

    end

    sort!(todelete)
    deleteat!(entities, todelete)

    seedmolecules
end


text:
export elu

doc"""
    elu(x::Var)

Exponential Linear Unit.

# References
* Clevert et al., ["Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)"](https://arxiv.org/abs/1511.07289), arXiv 2015.

```math
f(x) =
\begin{cases}
x & x > 0 \\
\alpha (e^{x}-1) & x\leq0
\end{cases}
```
where ``\alpha=1``.
"""
elu(x::Var) = Var(elu.(x.data), (elu,x))
elu(x::T) where T = x > zero(T) ? x : exp(x)-1
elu(x::Node) = Node(elu, x)

function addgrad!(y::Var, ::typeof(elu), x::Var)
    isvoid(x.grad) && return
    ∇elu!(y.data, y.grad, x.data, x.grad)
end

function ∇elu!(y::Array{T}, gy::Array{T}, x::Array{T}, gx::Array{T}) where T
    alpha = one(T)
    @inbounds for i = 1:length(x)
        gx[i] += x[i] > zero(T) ? gy[i] : gy[i]*(y[i]+alpha)
    end
end


text:
module OOPMacro
export @class, @super
include("OOPMacroImpl.jl")
end #module


text:
using JLD2
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation
#using SparseArrays, LinearAlgebra
using ViewerGL
GL = ViewerGL

# created by surface_extraction_parallel_ircad01.jl
@load "test/liver01.jld2" V FV
FW = [[v1,v2,v4,v3] for (v1,v2,v3,v4) in FV]
GL.VIEW([ GL.GLGrid(V,FW,GL.COLORS[2]) ]);

# created by surface_extraction_parallel_ircad01.jl
@load "test/liver02.jld2" V FV
TV = cat([[[v1,v2,v3],[v1,v4,v3]] for (v1,v2,v3,v4) in FV])
TV = convert(Lar.Cells,TV)
GL.VIEW([ GL.GLGrid(V,TV,GL.COLORS[1]) ]);


text:
# Test functionality of PDDL axioms / derived predicates
path = joinpath(dirname(pathof(PDDL)), "..", "test", "axioms")

domain = load_domain(joinpath(path, "domain.pddl"))
@test domain.name == Symbol("blocksworld-axioms")
@test domain.predicates[:above] == pddl"(above ?x ?y)"
@test @julog(above(X, Y) <<= or(on(X, Y), and(on(X, Z), above(Z, Y)))) in domain.axioms

problem = load_problem(joinpath(path, "problem.pddl"))
@test problem.name == Symbol("blocksworld-problem")
@test problem.objects == @pddl("a", "b", "c")

state = init_state(problem)
state = execute(pddl"(pickup b)", state, domain)
@test state[domain, pddl"(holding b)"] == true
state = execute(pddl"(stack b c)", state, domain)
@test state[domain, pddl"(on b c)"] == true
state = execute(pddl"(pickup a)", state, domain)
@test state[domain, pddl"(holding a)"] == true
state = execute(pddl"(stack a b)", state, domain)
@test state[domain, pddl"(above a c)"] == true

@test satisfy(problem.goal, state, domain)[1] == true


text:
using PowerModels: parse_file

@testset "PowerModels integration" begin
    matpower_path = joinpath(@__DIR__, "data", "nesta_case118_ieee.m")
    network_data = parse_file(matpower_path)

    line_lengths = estimate_length(network_data)

    @test maximum(collect(values(line_lengths))) ≈ 348276.672 atol=atol
    @test line_lengths["106"] ≈ 220329.17271676302 atol=atol

    acsr_specs = acsr_interpolation(network_data)
    labels = [a.label for a in values(acsr_specs)]
    @test length(unique(labels)) == 19
    @test acsr_specs["41"].label == "Oriole"
end


text:
function roots_and_ends(F, x1, x2)
    r = [x1, x2]
    for f in F
        _roots = find_zeros(f, x1, x2)
        append!(r, _roots)
    end
    return unique!(sort!(r))
end

function extend(x0::T, k, x) where {T<:Number}
    if k == 1
        return [x, x0]
    elseif k == 2
        return [x0, x]
    else
        throw(ArgumentError("Expected k ∈ {1,2}, got $k"))
    end
end

function extend(x0::V, k, x::T) where {V<:AbstractVector,T<:Number}
    dim = length(x0)
    if dim == 1
        return extend(x0[1], k, x)
    else
        throw(ArgumentError("Current support for dim = 1, got dim = $dim"))
    end
end

function extend(x0::V, k, x::M) where {V<:AbstractVector,M<:AbstractMatrix}
    old_dim = length(x0)
    dim, npoints = size(x)
    if dim != 1
        msg = "Current support for dim = 1, got dim = $dim"
        throw(ArgumentError(msg))
    end

    old_row = repeat(x0, inner = (1, npoints))
    if old_dim == 1 && k == 1
        return vcat(x, old_row)
    elseif old_dim == 1 && k == 2
        return vcat(old_row, x)
    else
        throw(
            ArgumentError(
                "Current support for dim = 1 and k ∈ {1,2}, got dim = $dim, k = $k",
            ),
        )
    end
end
function height_direction(grad, xc)
    grad = vec(grad(xc))
    if norm(grad) ≈ 0.0
        return 0
    else
        return argmax(abs.(grad))
    end
end

function height_direction(grad, box::IntervalBox{N,T}) where {N,T}
    return height_direction(grad, mid(box))
end

function curvature_measure(grad, k, C)
    return sum(grad .^ 2) - C^2 * grad[k]^2
end

function is_suitable(
    height_dir,
    grad,
    box;
    tol = 1e-3,
    C = 4,
    perturbation = 0.0,
)

    if height_dir == 0
        return false, 0
    else
        gradk(x) = grad(x)[height_dir]
        s = sign(gradk, box, tol = tol, perturbation = perturbation)
        curvatureflag = true
        if s != 0
            curve(x) = curvature_measure(grad(x), height_dir, C)
            t = sign(curve, box, tol = tol)
            curvatureflag = t == -1 ? false : true
        end
        flag = (s == 0 || curvatureflag) ? false : true
        return flag, s
    end
end

function Base.sign(m::Z, s::Z, S::Bool, sigma::Z) where {Z<:Integer}

    if S || m == sigma * s
        return sigma * m
    else
        return 0
    end
end

function combine_quadratures(splitquads)
    points = hcat([sq.points for sq in splitquads]...)
    weights = vcat([sq.weights for sq in splitquads]...)
    return TemporaryQuadrature(points, weights)
end

function update_interpolating_gradient!(interpgrad, poly)
    points = interpolation_points(basis(interpgrad))
    interpvals = mapslices(x -> vec(gradient(poly, x)), points, dims = 1)
    PolynomialBasis.update!(interpgrad, interpvals)
end

max_stars_repo_name: ArjunNarayanan/ImplicitDomainQuadrature

max_issues_repo_name: ArjunNarayanan/ImplicitDomainQuadrature

max_forks_repo_name: ArjunNarayanan/ImplicitDomainQuadrature


text:
# Miscellaneous utilities.

export issomething,
        order_asc,
        isdisjoint,
        pushrand!,
        get_bit,
        init_array,
        subseteq_fast,
        similar!,
        Var,
        Lit,
        var2lit,
        lit2var,
        variables,
        num_variables,
        always,
        never,
        uniform,
        logsumexp,
        noop

using CUDA: CuArray, CuVector, CuMatrix, CUDA

"Is the argument not `nothing`?"
@inline issomething(x) = !isnothing(x)

"Order the arguments in a tuple in ascending order"
@inline order_asc(x, y) = x > y ? (y, x) : (x , y)

"Are the given sets disjoint (no shared elements)?"
function Base.isdisjoint(set1::AbstractSet, set2::AbstractSet, set3::AbstractSet, othersets::AbstractSet...)::Bool
    seen = set1 # can be sped up by copying seen first then reusing it with union!?
    for set in [set2, set3, othersets...]
        if !isempty(intersect(seen,set))
            return false
        else
            seen = union(seen, set)
        end
    end
    return true
end

"""
Push `element` into random position in vector`v`
"""
function pushrand!(v::AbstractVector{<:Any}, element)
    len = length(v)
    i = rand(1:len + 1)
    if i == len + 1
        push!(v, element)
    else
        splice!(v, i:i, [element, v[i]])
    end
    v
end

"Retrieve the jth bit from a `BitVector` chunk"
get_bit(chunk, j) = chunk & (UInt64(1) << Base._mod64(j-1)) != 0

"An array of undetermined values (fast) for the given element type"
@inline init_array(::Type{T}, dims::Int...) where T<:Number = Array{T}(undef, dims...)
@inline init_array(::Type{T}, dims::Int...) where T<:Bool= BitArray(undef, dims...)

"Replacement for BitSet.⊆ that does not allocate a new BitSet"
function subseteq_fast(s::BitSet,t::BitSet) # see https://discourse.julialang.org/t/issubset-bitset-bitset-is-slow-and-allocates-unnecessarily/43828
    # TODO see if https://github.com/JuliaLang/julia/pull/36882/files has fixed this issue in 1.6
    s === t && return true
    s.offset < t.offset && return false
    length(s.bits) > length(t.bits) && return false
    length(s) > length(t) && return false
    return all(e -> e ∈ t, s)
end

"Reuse a given array if it has the right type and size, otherwise make a new one"
function similar!(reuse, ::Type{A}, desired_size...) where A<:AbstractArray
    if reuse isa A && size(reuse) == (desired_size...,)
        reuse
    elseif reuse isa A && ndims(reuse) == 1 && length(desired_size) == 1
        # Use resize! to save some effort when the desired array is 1-dimentional
        resize!(reuse, desired_size...)
    else
        A(undef, desired_size...)
    end
end

#####################
# General logic
#####################

"""
Variables are represented as 32-bit unsigned integers
"""
const Var = UInt32 # variable ids

"""
Literals are represented as 32-bit signed integers.
Positive literals are positive integers identical to their variable. Negative literals are their negations. Integer 0 should not be used to represent literals.
"""
const Lit = Int32 # variable with a positive or negative sign

"Convert a variable to the corresponding positive literal"
@inline var2lit(v::Var)::Lit = convert(Lit,v)

"Convert a literal its variable, removing the sign of the literal"
@inline lit2var(l::Lit)::Var = convert(Var,abs(l))

"Get the `BitSet` of variables in the data structure"
function variables end

"Number of variables in the data structure"
@inline num_variables(x)::Int = length(variables(x))

#####################
# probability semantics
#####################

"An array of 100% probabilities for the given element type"
@inline always(::Type{T}, dims::Int...) where T<:Number = ones(T, dims...)
@inline always(::Type{T}, dims::Int...) where T<:Bool = trues(dims...)
@inline always(::Type{<:CuArray{T}}, dims::Int...) where T = CUDA.ones(T, dims...)
@inline always(::Type{T}, dims::Int...) where T<:AbstractArray = always(eltype(T), dims...)::T

"An array of 0% probabilities for the given element type"
@inline never(::Type{T}, dims::Int...) where T<:Number = zeros(T, dims...)
@inline never(::Type{T}, dims::Int...) where T<:Bool = falses(dims...)
@inline never(::Type{<:CuArray{T}}, dims::Int...) where T = CUDA.zeros(T, dims...)
@inline never(::Type{T}, dims::Int...) where T<:AbstractArray = never(eltype(T), dims...)::T

"An array of uniform probabilities"
@inline uniform(dims::Int...) = uniform(Float64, dims...)
@inline uniform(::Type{F}, dims::Int...) where F<:AbstractFloat = always(F, dims...) ./ prod(dims)

import StatsFuns: logsumexp #extend

"Marginalize out dimensions `dims` from log-probability tensor"
function logsumexp(a::AbstractArray, dims)
    return dropdims(mapslices(logsumexp, a, dims=dims), dims=dims)
end

#####################
# functional programming
#####################

"Function that does nothing"
@inline noop(_...) = nothing

max_issues_repo_name: Juice-jl/LogicCircuits

max_forks_repo_name: Juice-jl/LogicCircuits


text:
immutable Atm
    altitude::@quantity(Float64,Meter)
    delta_temperature::@quantity(Float64,Kelvin)
    temperature::@quantity(Float64,Kelvin)
    temperature_ratio::Float64
    pressure::@quantity(Float64,N/m^2)
    pressure_ratio::Float64
    density::@quantity(Float64,kg/m^3)
    density_ratio::Float64
    speed_of_sound::@quantity(Float64,m/s)
end

SeaLevelTemperature = 288.15K
SeaLevelDensity = 1.225kg/m^3
SeaLevelPressure = 101327.0N/m^2
SeaLevelSpeedOfSound = 340.2979m/s
LowerStratosphereTemperature = 216.649K
LowerStratosphereTemperatureRatio = LowerStratosphereTemperature/SeaLevelTemperature
LowerStratospherePressure = 22633.0N/m^2
LowerStratospherePressureRatio = LowerStratospherePressure/SeaLevelPressure
HydrostaticConstant = 34.163195K/m

function Atm(altitude::@quantity(Float64,Meter),delta_temperature::@quantity(Float64,Kelvin) = 0.0K)
    # http://www.pdas.com/programs/atmos.f90
    if (altitude < 11000m)
        temperature_ratio = 1 - altitude/11000m*(1-LowerStratosphereTemperatureRatio);
        pressure_ratio = temperature_ratio^(HydroStaticConstant/(6.5K/m));
    elseif altitude < 25000m
        temperature_ratio = LowerStratosphereTemperatureRatio;
        pressure_ratio = LowerStratospherePressureRatio*exp(-HydrostaticConstant*(altitude-11000m)/LowerStratosphereTemperature);
    else
        throw(ErrorException("Atmosphere above 25000m not supported"))
    end
    density_ratio = pressure_ratio/temperature_ratio;
    delta_temperature_ratio = 1+delta_temperature/(temperature_ratio*SeaLevelTemperature);
    temperature_ratio *= delta_temperature_ratio;
    density_ratio /= delta_temperature_ratio;

    temperature = temperature_ratio * SeaLevelTemperature;
    density = density_ratio * SeaLevelDensity;
    pressure = pressure_ratio * SeaLevelPressure;

    speed_of_sound = SeaLevelSpeedOfSound*sqrt(pressure_ratio);
    Atm(altitude,
        delta_temperature,
        temperature,
        temperature_ratio,
        pressure,
        pressure_ratio,
        density,
        density_ratio,
        speed_of_sound)
end

function show(io::IO,atm::Atm)
    print(io,"Atm(")
    print(io,"Altitude=",round(atm.altitude,1),", ")
    print(io,"Temperature=",round(atm.temperature,1),", ")
    print(io,"Pressure=",round(atm.pressure,1),", ")
    print(io,"Density=",round(atm.density,3),")")
end

# function Atm{N1<:Real,N2<:Real,U1<:LengthUnit,U2<:TemperatureUnit}(altitude::Quantity{N1,U1},delta_temperature::Quantity{N2,U2} = 0.0K)
#     Atm(convert(SIQuantity{Float64},altitude),convert(SIQuantity{Float64},delta_temperature))
# end