mutable struct TimeStepWizard{FT, C, D}
                         cfl :: FT
               diffusive_cfl :: FT
                  max_change :: FT
                  min_change :: FT
                      max_Δt :: FT
                      min_Δt :: FT
                          Δt :: FT
    cell_advection_timescale :: C
    cell_diffusion_timescale :: D
end

infinite_diffusion_timescale(args...) = Inf # its not very limiting

"""
    TimeStepWizard(cfl=0.1, max_change=2.0, min_change=0.5, max_Δt=Inf, min_Δt=0.0, Δt=0.01)

A type for calculating adaptive time steps based on capping the CFL number at `cfl`.

On calling `update_Δt!(wizard, model)`, the `TimeStepWizard` computes a time-step such that
``cfl = max(u/Δx, v/Δy, w/Δz) Δt``, where ``max(u/Δx, v/Δy, w/Δz)`` is the maximum ratio
between model velocity and along-velocity grid spacing anywhere on the model grid. The new
`Δt` is constrained to change by a multiplicative factor no more than `max_change` or no
less than `min_change` from the previous `Δt`, and to be no greater in absolute magnitude
than `max_Δt` and no less than `min_Δt`.
"""
function TimeStepWizard(FT=Float64; cfl = 0.1,
                                    diffusive_cfl = Inf,
                                    max_change = 2.0,
                                    min_change = 0.5,
                                    max_Δt = Inf,
                                    min_Δt = 0.0,
                                    Δt = 0.01,
                                    cell_advection_timescale = cell_advection_timescale,
                                    cell_diffusion_timescale = infinite_diffusion_timescale)

    isfinite(diffusive_cfl) && # user wants to limit by diffusive CFL
    !(cell_diffusion_timescale === infinite_diffusion_timescale) && # user did not provide custom timescale
        (cell_diffusion_timescale = Oceananigans.TurbulenceClosures.cell_diffusion_timescale)

    C = typeof(cell_advection_timescale)
    D = typeof(cell_diffusion_timescale)

    return TimeStepWizard{FT, C, D}(cfl, diffusive_cfl, max_change, min_change, max_Δt, min_Δt, Δt,
                                    cell_advection_timescale, cell_diffusion_timescale)
end

using Oceananigans.Grids: topology

"""
    update_Δt!(wizard, model)

Compute `wizard.Δt` given the velocities and diffusivities of `model`, and the parameters
of `wizard`.
"""
function update_Δt!(wizard, model)

    Δt = min(
             wizard.cfl * wizard.cell_advection_timescale(model),          # advective
             wizard.diffusive_cfl * wizard.cell_diffusion_timescale(model) # diffusive
            )

    # Put the kibosh on if needed
    Δt = min(wizard.max_change * wizard.Δt, Δt)
    Δt = max(wizard.min_change * wizard.Δt, Δt)
    Δt = clamp(Δt, wizard.min_Δt, wizard.max_Δt)

    wizard.Δt = Δt

    return nothing
end

(c::CFL{<:TimeStepWizard})(model) = c.Δt.Δt / c.timescale(model)


import ..InterpolationAndGrids: interpolate_and_bin

"""
    resample(udata::AbstractUncertainIndexValueDataset, regularization_scheme::InterpolateAndBin{Linear})

Draw a single realisation of `udata` and interpolate-and-bin the data according to the
provided regularization scheme. Assumes points in `udata` are independent and sorts the draw
according to the index values before interpolating. See also [`InterpolateAndBin`](@ref).

## Example

```julia
npts = 50
y = rand(npts)

N = Normal(0, 1)

for t in 3:npts
    y[t,1] = 0.7*y[t-1,1] - 0.35*y[t-2,1] + rand(N)
end

# Assume data are unevenly spaced
time = sample(1.0:npts*5, npts, ordered = true, replace = false)

# Assign some uncertainties to both time indices and values and gather
# in an UncertainIndexValueDataset
utime = UncertainValue.(Normal.(time, 2))
uy = UncertainValue.(Normal.(y, 0.1))
udata = UncertainIndexValueDataset(utime, uy)

# Interpolation-and-binning scheme. First interpolate to a very fine grid,
# then gather the points falling in each of the coarser bins and summarise
# each bin using the mean of the points in each bin.
left_bin_edges = 0:10:npts*5
r = InterpolateAndBin(mean, left_bin_edges, Linear(), 0:0.1:1000, Flat(OnGrid()))

# The binned time axis:
time_binned = left_bin_edges[1:end-1] .+ step(left_bin_edges)/2

# Get a set corresponding resampled (interpolated+binned) values
y_binned = resample(udata, r)

# Plot some interpolated+binned draws
time_binned = left_bin_edges[1:end-1] .+ step(left_bin_edges)/2

p = plot(xlabel = "time", ylabel = "value")
for i = 1:100
    plot!(time_binned, resample(udata, r), lw = 0.3, α = 0.2, ms = 0.1, c = :red,
        marker = stroke(0.1), label = "")
end
plot!(time, y, c = :black, lw = 1, ms = 2, marker = stroke(2.0, :black), label = "")
plot!(udata, c = :black, lw = 1, ms = 2, marker = stroke(0.1, :black), [0.05, 0.95], [0.05, 0.95])
vline!(left_bin_edges, c = :black, α = 0.3, lw = 0.3, label = "")
```
"""
function resample(udata::AbstractUncertainIndexValueDataset, regularization_scheme::InterpolateAndBin{Linear})
    inds, vals = resample(udata)

    sortidxs = sortperm(inds)

    r = regularization_scheme
    vals_binned = interpolate_and_bin(r.f, r.left_bin_edges, inds[sortidxs], vals[sortidxs],
        r.intp, r.intp_grid, r.extrapolation_bc)

    return vals_binned
end

export resample


module LossTests

using Photon, Test

function test_distance(loss)
    p = randn(10,8)
    y = randn(10,8)
    r = loss(p,y)
    @test r > 0
    @test loss(p,y) == loss(y,p)

    r = loss(p,p)
    @test r == 0

    p = randn(10,1,8)
    y = randn(1,10,8)
    r = loss(p,y)
    @test r > 0

    p = KorA(randn(10,8))
    y = KorA(randn(10,8))
    r = loss(p,y)
    @test loss(p,y) == loss(y,p)
    @test r > 0

    r = loss(p,p)
    @test r == 0

end

function test_nondistance(loss)
    p = randn(10,8)
    y = randn(10,8)
    r = loss(p,y)
    @test r isa Number

    p = randn(10,1,8)
    y = randn(1,10,8)
    r = loss(p,y)
    @test r isa Number

    p = KorA(randn(10,8))
    y = KorA(randn(10,8))
    r = loss(p,y)
    @test r isa Number

    @test loss(p,p) <= loss(p,y)

end

function test_classification(loss)
    p = Float32.(rand(0:0.0001:1.0,10,16))
    y = Float32.(rand(0:1, 10,16))
    r = loss(p, y)
    @test r > 0

    r = loss(KorA(p), KorA((y)))
    @test r > 0

end

@testset "Loss" begin
    test_distance(L1Loss())
    test_distance(L1Loss(reduction=sum))
    test_distance(L2Loss())
    test_distance(LNLoss(3))

    test_nondistance(PseudoHuberLoss())

    test_classification(HingeLoss())
    test_classification(CrossEntropyLoss())
    test_classification(BCELoss())
    test_classification(FocalLoss())
end

end


using OMEinsum

N = 100
ca = randn(Float32, N, N)
ct = randn(Float32, N, N, N)
cq = randn(Float32, N, N, N, N)

function bfunc_star_cpu()
    ein"ji,kl,li->jkl"(ca,ca,ca)
end

function bfunc_t3_cpu()
    ein"ji,kli,li->jkl"(ca,ct,ca)
end

function bfunc_psum_cpu()
    ein"iikl->"(cq)
end

using BenchmarkTools
display(@benchmark bfunc_star_cpu() seconds = 1)
display(@benchmark bfunc_t3_cpu() seconds = 1)
display(@benchmark bfunc_psum_cpu() seconds = 1)


#broad_pdf.jl
using JSON

function create_broad_pdf(n_x::Int, n_y::Int, n_z::Int, dependence::Float64=0.) :: Dict{Tuple{Int,Int,Int},BigFloat}
    @assert n_x ≥ 2 "create_broad_pdf(): n_x >= 2 needed"
    @assert n_y ≥ 2 "create_broad_pdf(): n_y >= 2 needed"
    @assert n_z ≥ 2 "create_broad_pdf(): n_z >= 2 needed"
    @assert 0. ≤ dependence ≤ 1. "create_broad_pdf(): indep in [0,1] needed"
    pdf = Dict{Tuple{Int,Int,Int},BigFloat}()
    for y = 1:n_y
        for z = 1:n_z
            for x = 1:n_x
                pdf[ (x,y,z) ] = ( 1 + dependence*(rand()-.5) )/( n_x*n_y*n_z )
            end # for z
        end # for z
    end # for y
    return pdf
end # create_broad_pdf()

function broad_pdf_to_file(filename::String, n_x::Int, n_y::Int, n_z::Int, dependence::Float64=0.) :: Void
    open(filename,"w") do file
        pdf = create_broad_pdf(n_x,n_y,n_z,dependence)
        pdflist = []
        for k in keys(pdf)
            (x,y,z) = k
            e = [[x,y,z],pdf[(x,y,z)]]
            push!(pdflist, e )
        end
        JSON.print(file,pdflist)
    end
    return nothing
end # broad_pdf_to_file()

function loop_broad_pdfs_to_files() :: Void
    for n_x = 2:5
        for n_y = n_x : 5*n_x
            n_z = n_y
            filename = "broad_indep--"*string(n_x)*"-"*string(n_y)*"-"*string(n_z)*".dens"
            broad_pdf_to_file(filename,n_x,n_y,n_z)
            filename = "broad_dep--"*string(n_x)*"-"*string(n_y)*"-"*string(n_z)*".dens"
            broad_pdf_to_file(filename,n_x,n_y,n_z,1.)
        end
    end
    return nothing;
end

;
#EOF


# Define batched matrix multiply for GPU and Duals:
function batched_mul(A::CuArray{ForwardDiff.Dual{T,V,N},3}, B::CuArray{ForwardDiff.Dual{T,V,N},3}) where {T,V,N}
    # Extract values:
    Av = ForwardDiff.value.(A)
    Bv = ForwardDiff.value.(B)
    # Use strided batch for A*B (defined as gemm_strided_batched):
    Cv = Av ⊠ Bv
    # Compute derivatives ∂(AB)/∂x = A * ∂B/∂x + ∂A/∂x * B;
    dABdx = [Av ⊠ ForwardDiff.partials.(B,i) + ForwardDiff.partials.(A,i) ⊠ Bv for i=1:N];
    dABdx = ForwardDiff.Partials.(tuple.(dABdx...));
    return eltype(A).(Cv,dABdx);
end

# Overload of batch_inv! for Dual numbers
function batch_inv!(X::CuArray{ForwardDiff.Dual{T,V,N},3}, A::CuArray{ForwardDiff.Dual{T,V,N},3}) where {T,V,N}
    Atemp = ForwardDiff.value.(A)
    invA = similar(Atemp);
    # Set invA=A⁻¹
    batch_inv!(invA,Atemp)
    # Compute derivatives ∂A⁻¹/∂x = -A⁻¹ * ∂A/∂x * A⁻¹; using NNlib batched matrix multiply
    -invA ⊠ ForwardDiff.partials.(A,1) ⊠ invA
    dAdx = [-invA ⊠ ForwardDiff.partials.(A,i) ⊠ invA for i=1:N];
    # Pack into tuples again
    dAdx = ForwardDiff.Partials.(tuple.(dAdx...));
    X .= eltype(X).(invA,dAdx);

end


struct Variable{Name} <: AbstractVariable{Name}
end

copy(v::Variable) = v

macro polyvar(names...)
    exprs = [
        quote
            $(esc(name)) = Variable{$(esc(Expr(:quote, name)))}()
        end
        for name in names
    ]
    Expr(:block, exprs...)
end

struct Power{V}
    exponent::Int

    function Power{V}(e::Int) where {V}
        new{typeassert(V, Variable)}(e)
    end
end

variable(p::Power{V}) where {V} = V
exponent(p::Power) = p.exponent

struct Monomial{Vars, Exponents} <: AbstractMonomial{Vars}
    function Monomial{V, E}() where {V, E}
        @assert length(V) == length(E)
        new{typeassert(V, Tuple{Vararg{Variable}}),
            typeassert(E, Tuple{Vararg{Int}})}()
    end
end

exponents(m::Monomial{V, E}) where {V, E} = E

@generated function powers(::Type{Monomial{V, E}}) where {V, E}
    Expr(:tuple, [:($(Power{v}(e))) for (v, e) in zip(V, E)]...)
end

struct Term{T, M <: Monomial} <: AbstractTerm{T, M}
    coefficient::T
end

Term{M}(c::T) where {T, M <: Monomial} = Term{T, M}(c)
Term(c::T, m::M) where {T, M <: Monomial} = Term{T, M}(c)

coefficient(t::Term) = t.coefficient
monomial(::Type{Term{T, M}}) where {T, M} = M()
monomial(t::Term) = monomial(typeof(t))

struct Polynomial{Terms <: Tuple{Vararg{Term}}} <: AbstractPolynomial
    terms::Terms
end

terms(p::Polynomial) = p.terms

function literal_pow(^, v::V, ::Type{Val{x}}) where {V <: Variable, x}
    Monomial{(v,), (x,)}()
end

const MonomialLike = Union{<:Variable, <:Monomial}
const TermLike = Union{<:MonomialLike, <:Term}
const PolynomialLike = Union{<:TermLike, <:Polynomial}


"""
Returns an iterator over all combinations of theta-lambda quadrature nodes in the form of a CartesianIndices object
Returns empty weightsTensorGrid array for storing tensor quadrature weights
"""
function get_btg_iterator(nwθ::nodesWeights, nwλ::nodesWeights, quadType::Array{String,1})
    nt1 = nwθ.d   #number of dimensions of theta
    nt2 = nwθ.num #number of theta quadrature nodes
    nl1 = nwλ.d   #number of dimensions of lambda
    nl2 = nwλ.num #number of lambda quadrature nodes
    if endswith(quadType[1], "Carlo") && endswith(quadType[2], "Carlo")
        weightsTensorGrid = Array{Float64, 1}(undef, nt2)
        R = CartesianIndices(weightsTensorGrid)
        if quadType == ["SparseCarlo", "SparseCarlo"]
            weightsTensorGrid = nwθ.weights[:]
        end
    elseif quadType == ["Gaussian", "Gaussian"]
        weightsTensorGrid = Array{Float64, nt1+nl1}(undef, Tuple(vcat([nt2 for i = 1:nt1], [nl2 for i = 1:nl1]))) #initialize tensor grid
        R = CartesianIndices(weightsTensorGrid)
        for I in R #I is multi-index
            weightsTensorGrid[I] = getProd(nwθ.weights, nwλ.weights, I) #this step can be simplified because the tensor is symmetric (weights are the same along each dimension)
        end
    else
        weightsTensorGrid = Array{Float64, 2}(undef, nt2, nl2)
        R = CartesianIndices(weightsTensorGrid)
        weightsTensorGrid = repeat(nwλ.weights, nt2, 1) # add lambda weights
        if quadType[1] == "SparseGrid"
            weightsTensorGrid = Diagonal(reshape(nwθ.weights, nt2)) * weightsTensorGrid
        end
    end
    return R, weightsTensorGrid #R is a CartesianIndices iterator over weightsTensorGrid
end

"""
Takes in I and interprets it in the context of quadtype pair and (θ, λ)-nodeWeights objects stored in btg.
Returns (θ, λ) index slices (r1 and r2) and (θ, λ) nodes (t1 and t2)
"""
function get_index_slices(nwθ::nodesWeights, nwλ::nodesWeights, quadType::Array{String,1}, I)
    r1 = (endswith(quadType[1], "Carlo") && endswith(quadType[2], "Carlo")) ? I : Tuple(I)[1:end-1] #first n-1 coords or everything
    r2 = Tuple(I)[end] #last coord
    t1 = quadType[1] == "Gaussian" ? getNodeSequence(getNodes(nwθ), r1) : getNodes(nwθ)[:, r1[1]] #theta node combo
    t1 = length(t1) == 1 ? t1[1] : t1
    t2 = getNodeSequence(getNodes(nwλ), r2)
    #t2 = getNodes(nwλ)[:, r2] #lambda node
    if length(t1)==1
        @assert typeof(t1)<:Real
    end
    @assert typeof(t2)<:Real
    return (r1, r2, t1, t2) #The pair (t1, t2) is used a key for various dictionaries
end


"""
    BSpline(t, x, w, B, S)

Basis structure for the B-splines generated by the knot set `t`. `x`
and `w` are the associated quadrature roots and weights, respectively,
the columns of `B` correspond to the B-splines resolved on the
quadrature roots, and `S` is the (banded) B-spline overlap matrix.
"""
struct BSpline{T,R<:Real,
               KnotSet<:AbstractKnotSet{<:Any,<:Any,<:Any,R},
               XV<:AbstractVector{T},
               WV<:AbstractVector{R},
               BM<:AbstractMatrix{T},
               SM<:AbstractMatrix{T}} <: Basis{T}
    t::KnotSet
    x::XV
    w::WV
    B::BM
    S::SM
end

const RestrictedBSpline{T} = RestrictedQuasiArray{T,2,<:BSpline{T}}
const AdjointRestrictedBSpline{T} = AdjointRestrictedQuasiArray{T,2,<:BSpline{T}}

const BSplineOrRestricted{T} = BasisOrRestricted{<:BSpline{T}}
const AdjointBSplineOrRestricted{T} = AdjointBasisOrRestricted{<:BSpline{T}}

function basis_functions(t::AbstractKnotSet{k}, x::AbstractVector{T}, m=0) where {k,T}
    nf = numfunctions(t)
    B = spzeros(T, length(x), nf)

    nei = nonempty_intervals(t)
    # This is the amount of quadrature points per interval (assuming
    # the same amount per interval, which is currently the case), this
    # could conceivably be passed as an argument, instead.
    N = findlast(e -> e < t[nei[1]+1], x)

    ix = 1
    for i ∈ nei
        # These are the functions which are non-zero on the interval
        # t[i]..t[i+1].
        for j ∈ max(1,i-k+1):min(i,nf)
            eⱼ = UnitVector{T}(nf, j)
            for l ∈ ix:ix+N-1
                B[l,j] = deBoor(t, eⱼ, x[l], i, m)
            end
        end
        ix += N
    end

    B
end

basis_functions(B::BSpline, args...) = basis_functions(B.t, B.x, args...)
basis_functions(B::RestrictedBSpline, args...) =
    basis_functions(parent(B), args...)[:,indices(B,2)]

function overlap_matrix!(S::Union{BandedMatrix,Tridiagonal}, χ, ξ, w)
    m = min(size(S,1),size(χ,2))
    n = min(size(S,2),size(ξ,2))
    # It is assumed that the bandwidth is correct for the overlap
    # ⟨χ|ξ⟩.
    p = max(bandwidths(S)...)
    W = Diagonal(w)
    for i ∈ 1:m
        for j ∈ max(i-p,1):min(i+p,n)
            S[i,j] = χ[:,i]' * W * ξ[:,j]
        end
    end
    # χ == ξ && isreal(χ) && isreal(W) && isreal(ξ) ?
    #     Symmetric(S) : S
    S
end

function assert_compatible_bases(A::BSplineOrRestricted, B::BSplineOrRestricted)
    A = parent(A)
    B = parent(B)
    A.t.t == B.t.t &&
        A.x == B.x &&
        A.w == B.w ||
        throw(ArgumentError("Can only multiply B-spline bases with identical knot sets, resolved on the same Gauß–Legendre points"))
end

function overlap_matrix!(S::Union{BandedMatrix,Tridiagonal},
                         Ac::AdjointBSplineOrRestricted,
                         B::BSplineOrRestricted,
                         op=I)
    A = parent(Ac)
    assert_compatible_bases(A,B)
    χ = view(parent(A).B, :, indices(A,2))
    ξ = view(parent(B).B, :, indices(B,2))
    overlap_matrix!(S, χ, op*ξ, weights(parent(B)))
end

function BSpline(t::AbstractKnotSet{k}, x::AbstractVector{T}, w::AbstractVector) where {k,T}
    B = basis_functions(t, x)

    nf = numfunctions(t)
    S = BandedMatrix(Zeros{T}(nf, nf), (k-1,k-1))

    BSpline(t, x, w, B, overlap_matrix!(S, B, B, w))
end

"""
    BSpline(t[, N])

Create the B-spline basis corresponding to the knot set `t`. `N` is
the amount of Gauß–Legendre quadrature points per interval.
"""
BSpline(t::AbstractKnotSet, N) = BSpline(t, lgwt(t, N)...)

"""
    BSpline(t[; k′=3])
Create the B-spline basis corresponding to the knot set `t`. `k′` is
the highest polynomial order of operators for which it should be
possible to compute the matrix elements exactly (via Gauß–Legendre
quadrature). The default `k′=3` corresponds to operators O(x²).
"""
BSpline(t::AbstractKnotSet; k′=3) = BSpline(t, num_quadrature_points(order(t), k′))

# * Properties

axes(B::BSpline) = (Inclusion(first(B.t)..last(B.t)), Base.OneTo(numfunctions(B.t)))
size(B::BSpline) = (ℵ₁, numfunctions(B.t))
==(A::BSpline,B::BSpline) = A.t == B.t
Base.hash(B::BSpline, h::UInt) = hash(B.t, h)

distribution(B::BSpline) = distribution(B.t)

order(B::BSplineOrRestricted) = order(unrestricted_basis(B).t)

function show(io::IO, B::BSpline{T}) where T
    write(io, "BSpline{$(T)} basis with $(B.t)")
end

locs(B::BSpline) = B.x
weights(B::BSpline) = B.w

IntervalSets.leftendpoint(B::BSpline) = B.x[1]
IntervalSets.rightendpoint(B::BSpline) = B.x[end]

function centers(B::BSplineOrRestricted)
    x = B'*QuasiDiagonal(axes(B,1))*B
    S = B'B

    diag(x) ./ diag(S)
end

# # * Basis functions

"""
    deBoor(t, c, x[, i[, m=0]])

Evaluate the spline given by the knot set `t` and the set of control
points `c` at `x` using de Boor's algorithm. `i` is the index of the
knot interval containing `x`. If `m≠0`, calculate the `m`th derivative
at `x` instead.
"""
function deBoor(t::AbstractKnotSet, c::AbstractVector{T}, x,
                i=find_interval(t, x), m=0) where T
    isnothing(i) && return zero(eltype(t))
    k = order(t)
    nc = length(c)
    k == 1 && return (nc < i || m > 0) ? zero(T) : c[i]

    α = [r > 0 && r ≤ nc ? c[r] : zero(T)
         for r ∈ i-k+1:i]
    nt = length(t)
    nf = numfunctions(t)
    for j = 1:k-1
        for r = i:-1:max(i-k+j,1)
            jjj = r+k-j
            (jjj > nt || jjj < 1) && continue
            r′ = r - i + k
            r ≠ 1 && r′ == 1 && continue

            a = t[r+k-j]
            b = t[r]

            α[r′] = if j ≤ m
                (k-j)*if r == 1
                    -α[r′]
                elseif r == nf + j
                    α[r′-1]
                else
                    α[r′-1] - α[r′]
                end
            else
                if r == 1
                    (b-x)*α[r′]
                elseif r == nf + j
                    (x-a)*α[r′-1]
                else
                    ((x-a)*α[r′-1] + (b-x)*α[r′])
                end
            end
            α[r′] /= (b-a)
        end
    end

    α[end]
end

getindex(B::BSpline{T}, x::Real, j::Integer) where T =
    deBoor(B.t, UnitVector{T}(size(B,2), j),
           x, find_interval(B.t, x))

function basis_function!(χ, B::BSpline{T}, x::AbstractVector, j) where T
    eⱼ = UnitVector{T}(size(B,2), j)
    for (is,k) ∈ within_support(x, B.t, j)
        for i in is
            χ[i] = deBoor(B.t, eⱼ, x[i], k)
        end
    end
end

function getindex(B::BSpline{T}, x::AbstractVector, sel::AbstractVector) where T
    χ = spzeros(T, length(x), length(sel))
    o = sel[1] - 1
    for j in sel
        basis_function!(view(χ, :, j-o), B, x, j)
    end
    χ
end

function getindex(B::BSplineOrRestricted{T}, x::AbstractVector, j::Integer) where T
    χ = spzeros(T, length(x))
    basis_function!(χ, parent(B), x, j+first(indices(B,2))-1)
    χ
end

getindex(B::BSplineOrRestricted, x, ::Colon) =
    getindex(parent(B), x, indices(B,2))

getindex(B::BSplineOrRestricted, x::AbstractVector, sel::AbstractVector) =
    getindex(parent(B), x, indices(B,2)[sel])

# * Types

const SplineArray{T,N,B<:BSplineOrRestricted} = FuncArray{T,N,B}
const SplineVector{T,B<:BSplineOrRestricted} = FuncVector{T,B}
const SplineMatrix{T,B<:BSplineOrRestricted} = FuncMatrix{T,B}
const SplineVecOrMat{T,B<:BSplineOrRestricted} = FuncVecOrMat{T,B}

const AdjointSplineArray{T,N,B<:BSplineOrRestricted} = AdjointFuncArray{T,N,B}
const AdjointSplineVector{T,B<:BSplineOrRestricted} = AdjointFuncVector{T,B}
const AdjointSplineMatrix{T,B<:BSplineOrRestricted} = AdjointFuncMatrix{T,B}
const AdjointSplineVecOrMat{T,B<:BSplineOrRestricted} = AdjointFuncVecOrMat{T,B}

Base.show(io::IO, spline::SplineVector) =
    write(io, "Spline on $(spline.args[1])")

Base.show(io::IO, spline::SplineMatrix) =
    write(io, "$(size(spline, 2))d spline on $(spline.args[1])")

# * Matrix construction

function Matrix(::UndefInitializer, A::BSplineOrRestricted{T}, B::BSplineOrRestricted{T}, ::Type{U}=T) where {T,U}
    m,n = size(A,2), size(B,2)
    k = max(order(A),order(B))
    if k == 2 && m == n
        dl = Vector{U}(undef, m-1)
        d = Vector{U}(undef, m)
        du = Vector{U}(undef, m-1)
        Tridiagonal(dl, d, du)
    else
        ij = first(indices(B,2)) - first(indices(A,2))
        BandedMatrix{U}(undef, (m,n), (k-1+ij,k-1-ij))
    end
end

Matrix(::UndefInitializer, A::BSplineOrRestricted{T}, ::Type{U}=T) where {T,U} =
    Matrix(undef, A, A, U)

function Base.zeros(A::BSplineOrRestricted{T}, B::BSplineOrRestricted{T}=A, ::Type{U}=T) where {T,U}
    M = Matrix(undef, A, B, U)
    M .= zero(U)
    M
end

# * Mass matrix
@materialize function *(Ac::AdjointBSplineOrRestricted,
                        B::BSplineOrRestricted)
    T -> begin
        Matrix(undef, parent(Ac), B, T)
    end
    dest::AbstractMatrix{T} -> begin
        overlap_matrix!(dest, Ac, B)
    end
end

metric(B::BSpline) = B.S

# * Diagonal operators

@materialize function *(Ac::AdjointBSplineOrRestricted,
                        D::QuasiDiagonal,
                        B::BSplineOrRestricted)
    T -> begin
        Matrix(undef, parent(Ac), B, T)
    end
    dest::AbstractMatrix{T} -> begin
        # Evaluate the quasi-diagonal operator D on the quadrature roots
        # of B.
        op = Diagonal(getindex.(Ref(D.diag), locs(parent(B))))
        overlap_matrix!(dest, Ac, B, op)
    end
end

@materialize function *(Ac::AdjointBSplineOrRestricted,
                        D::QuasiDiagonal,
                        E::QuasiDiagonal,
                        B::BSplineOrRestricted)
    T -> begin
        Matrix(undef, parent(Ac), B, T)
    end
    dest::AbstractMatrix{T} -> begin
        # Evaluate the quasi-diagonal operators D & E on the quadrature roots
        # of B.
        op = Diagonal(getindex.(Ref(D.diag), locs(parent(B))) .*
                      getindex.(Ref(E.diag), locs(parent(B))))
        overlap_matrix!(dest, Ac, B, op)
    end
end

# * Function interpolation

function Base.:(\ )(B::BSpline, f::BroadcastQuasiArray)
    axes(f,1) == axes(B,1) ||
        throw(DimensionMismatch("Function on $(axes(f,1).domain) cannot be interpolated over basis on $(axes(B,1).domain)"))
    B.B \ getindex.(Ref(f), B.x)
end

function Base.:(\ )(B::RestrictedBSpline, f::BroadcastQuasiArray)
    axes(f,1) == axes(B,1) ||
        throw(DimensionMismatch("Function on $(axes(f,1).domain) cannot be interpolated over basis on $(axes(B,1).domain)"))
    # We need to evaluate the basis functions of the restricted
    # B-spline basis on /all/ quadrature points.
    x = locs(parent(B))
    V = B[x,:]
    V \ getindex.(Ref(f), x)
end

export BSpline


#break into kmers for integers
function collect_kmers(seq::Array{Int,1}, k::Int) #lookup_dic not used.
    return [seq[1+i:k+i] for i in 0:length(seq)-k]
end

#weights vector length is the number of unique kmers
#each kmer is counted once, contributes to one pos
function kmer_count!(kmer_counts::Array{UInt64,1}, kmer::Vector{Int}, n_chars::Int)
    kmer_counts[get_kmer_index(kmer,n_chars)] += unsigned(1)
end

#this function gets the index in lexicographic order...
#remember 1-based indexing in Julia...
function get_kmer_index(kmer::Vector{Int}, n_chars::Int)
    kmer_ix = 0
    for c in kmer
        kmer_ix = kmer_ix * n_chars + (c - 1)
    end
    return kmer_ix + 1
end

#with a `kmer_contribution()` function...
function kmer_embed(seq, k::Int, kmer_contribution::Function; lookup_dic = AA_DICT, missing_chars = Set{Char}())
    n = length(union(values(lookup_dic)))
    kmer_counts = zeros(UInt64, n^k)

    #break into kmers
    seqvec, missing_chars = string2encoding(seq, lookup_dic; missing_chars = missing_chars)
    kmers = collect_kmers(seqvec, k)

    #store kmer count in array
    for kmer in kmers
        if !(0 in kmer)
            kmer_contribution(kmer_counts, kmer, n)
        end
    end
    return kmer_counts, missing_chars
end


import Test.@testset
import Test.@test

using LightGraphs

push!(LOAD_PATH, pwd())
using Coils

const tol = 1e-10

@testset "Solving a system" begin
    g, vertex_positions = cuboid_system([1, 1, 1], [2, 2, 2])
    poi = [[0,0,0]]
    Bgoal(x) = [0, 0, 50]
    simpleloops, simpleloopscurrents = solve_system(g, vertex_positions, poi, Bgoal)

    field = field_loops(poi[1], simpleloops, simpleloopscurrents, vertex_positions)
    @test abs(field[3] - 50) < tol
end

if !(("TRAVIS" in keys(ENV)) & Sys.isapple())
    # do not execute this test on OSX in Travis - there is a problem with
    # matplotlib there
    @testset "Example code with plotting" begin
        include("example.jl")
    end
end


include("../../../../math/common/sparse/src/sparse_module.jl")
include("../src/mtx_module.jl")

module mtest
using ..sparse
using ..mtx
using PyPlot
PyPlot.pygui(true)

# url = raw"D:\programming\sparse_matrix_collection\add32\add32.mtx"
# @time mtx_content_add32 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\arc130\arc130.mtx"
# @time mtx_content_arc130 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\bp_1600\bp_1600.mtx"
# @time mtx_content_bp1600 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\cage3\cage3.mtx"
# @time mtx_content_cage3 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\epb1\epb1.mtx"
# @time mtx_content_epb1 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\g7jac100sc\g7jac100sc.mtx"
# @time mtx_content_g7jac100sc = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\G67\G67.mtx"
# @time mtx_content_G67 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\mark3jac140\mark3jac140.mtx"
# @time mtx_content_mark3jac140 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\nemeth01\nemeth01.mtx"
# @time mtx_content_nemeth01 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\nemeth06\nemeth06.mtx"
# @time mtx_content_nemeth06 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\nemeth19\nemeth19.mtx"
# @time mtx_content_nemeth19 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\nemeth20\nemeth20.mtx"
# @time mtx_content_nemeth20 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\stomach\stomach.mtx"
# @time mtx_content_stomach = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\cfd1\cfd1.mtx"
# @time mtx_content_cfd1 = mtx.read_mtx(url)

# url = raw"D:\programming\sparse_matrix_collection\worms20_10NN\worms20_10NN.mtx"
# @time mtx_content_worms20_10NN = mtx.read_mtx(url)

url = raw"D:\programming\sparse_matrix_collection\cz628\cz628.mtx"
@time mtx_content_cz628 = mtx.read_mtx(url)

# mtx_content_arc130
# mtx_content_cage3
# mtx_content_nemeth20

spmat = mtx.mtx_to_sparse_csc(mtx_content_cz628)

fmat = sparse.create_full(spmat)

PyPlot.figure()
PyPlot.spy(fmat)

end


function transfer_learn(model; opt=Descent(0.01), epochs=30_000,
                        train_batch_size=50, eval_batch_size=10, eval_interval=1000)

    weights = params(model)
    dist = Uniform(-5, 5)
    testx = Float32.(range(-5, stop=5, length=50))

    for i in 1:epochs
        task = SineWave()
        x = Float32.(rand(dist, train_batch_size))
        y = task(x)
        l = Flux.mse(model(x'), y')
        Flux.back!(l)
        Flux.Optimise._update_params!(opt, weights)

        if i % eval_interval == 0
            @printf("Iteration %d, evaluating model on random task...\n", i)
            evalx = Float32.(rand(dist, eval_batch_size))
            eval_model(model, evalx, testx, SineWave())
        end
    end
end


#=

Author: Shunsuke Hori

=#

using LS

using QuantEcon
using NLopt
using NLsolve
using Dierckx

Para=LS.Para
Planners_Allocation_Sequential = LS.Planners_Allocation_Sequential
"""
Bellman equation for the continuation of the Lucas-Stokey Problem
"""
mutable struct BellmanEquation{TP <: Para, TI <: Integer, TR <: Real}
    para::TP
    S::TI
    xbar::Array{TR}
    time_0::Bool
    z0::Array{Array}
    cFB::Vector{TR}
    nFB::Vector{TR}
    xFB::Vector{TR}
    zFB::Vector{Vector{TR}}
end

"""
Compute the planner's allocation by solving Bellman
equation.
"""
struct Planners_Allocation_Bellman{TP <: Para, TI <: Integer,
                                   TVg <: AbstractVector, TT <: Tuple}
    para::TP
    mc::MarkovChain
    S::TI
    T::BellmanEquation
    mugrid::TVg
    xgrid::TVg
    Vf::Array
    policies::TT
end

"""
Initializes the type from the calibration Para
"""
function Planners_Allocation_Bellman(para::Para, mugrid::AbstractArray)
    G = para.G
    S = size(para.Pi, 1) # number of states
    mc=MarkovChain(para.Pi)
    #now find the first best allocation
    Vf, policies, T, xgrid = solve_time1_bellman(para, mugrid)
    T.time_0 = true #Bellman equation now solves time 0 problem
    return Planners_Allocation_Bellman(para, mc, S, T, mugrid, xgrid, Vf, policies)
end

"""
Solve the time  1 Bellman equation for calibration Para and initial grid mugrid
"""
function solve_time1_bellman{TR <: Real}(para::Para{TR}, mugrid::AbstractArray)
    Pi = para.Pi
    S = size(para.Pi, 1)

    #First get initial fit from lucas stockey solution.
    #Need to change things to be ex_ante
    PP = LS.Planners_Allocation_Sequential(para)

    function incomplete_allocation(PP::LS.Planners_Allocation_Sequential,
                                    mu_::AbstractFloat, s_::Integer)
        c, n, x, V = LS.time1_value(PP, mu_)
        return c,n,dot(Pi[s_, :], x), dot(Pi[s_, :], V)
    end
    cf = Array{Function}(S, S)
    nf = Array{Function}(S, S)
    xprimef = Array{Function}(S, S)
    Vf = Vector{Function}(S)
    # cf = Array{LinInterp}(S, S)
    # nf = Array{LinInterp}(S, S)
    # xprimef = Array{LinInterp}(S, S)
    # Vf = Vector{LinInterp}(S)
    xgrid = Array{TR}(S, length(mugrid))
    for s_ in 1:S
        c = Array{TR}(length(mugrid), S)
        n = Array{TR}(length(mugrid), S)
        x = Array{TR}(length(mugrid))
        V = Array{TR}(length(mugrid))
        for (i_mu, mu) in enumerate(mugrid)
            c[i_mu, :], n[i_mu, :], x[i_mu], V[i_mu] =
                incomplete_allocation(PP, mu, s_)
        end
        xprimes = repmat(x, 1, S)
        xgrid[s_, :] = x
        for sprime = 1:S
            splc = Spline1D(x[end:-1:1], c[:, sprime][end:-1:1], k=3)
            spln = Spline1D(x[end:-1:1], n[:, sprime][end:-1:1], k=3)
            splx = Spline1D(x[end:-1:1], xprimes[:, sprime][end:-1:1], k=3)
            cf[s_, sprime] = y -> splc(y)
            nf[s_, sprime] = y -> spln(y)
            xprimef[s_, sprime] = y -> splx(y)
            # cf[s_, sprime] = LinInterp(x[end:-1:1], c[:, sprime][end:-1:1])
            # nf[s_, sprime] = LinInterp(x[end:-1:1], n[:, sprime][end:-1:1])
            # xprimef[s_, sprime] = LinInterp(x[end:-1:1], xprimes[:, sprime][end:-1:1])
        end
        splV = Spline1D(x[end:-1:1], V[end:-1:1], k=3)
        Vf[s_] = y -> splV(y)
        # Vf[s_] = LinInterp(x[end:-1:1], V[end:-1:1])
    end

    policies = [cf, nf, xprimef]

    #create xgrid
    xbar = [maximum(minimum(xgrid)), minimum(maximum(xgrid))]
    xgrid = linspace(xbar[1], xbar[2], length(mugrid))

    #Now iterate on Bellman equation
    T = BellmanEquation(para, xgrid, policies)
    diff = 1.0
    while diff > 1e-6
        PF = (i_x, x, s) -> get_policies_time1(T, i_x, x, s, Vf, xbar)
        Vfnew, policies = fit_policy_function(T, PF, xgrid)

        diff = 0.0
        for s=1:S
            diff = max(diff,maximum(abs,
                       (Vf[s].(xgrid)-Vfnew[s].(xgrid))./Vf[s].(xgrid)))
        end

        println("diff = $diff")
        Vf = copy(Vfnew)
    end

    return Vf, policies, T, xgrid
end

"""
Fits the policy functions
"""
function fit_policy_function{TF<:AbstractFloat}(
                    T::BellmanEquation,
                    PF::Function, xgrid::AbstractVector{TF})
    S = T.S
    # preallocation
    PFvec = Array{TF}(4S+1, length(xgrid))
    cf = Array{Function}(S, S)
    nf = Array{Function}(S, S)
    xprimef = Array{Function}(S, S)
    TTf = Array{Function}(S, S)
    Vf = Vector{Function}(S)
    # fit policy fuctions
    for s_ in 1:S
        for (i_x, x) in enumerate(xgrid)
            PFvec[:, i_x] = PF(i_x, x, s_)
        end
        splV = Spline1D(xgrid, PFvec[1,:], k=3)
        Vf[s_] = y -> splV(y)
        # Vf[s_] = LinInterp(xgrid, PFvec[1, :])
        for sprime=1:S
            splc = Spline1D(xgrid, PFvec[1+sprime,:], k=3)
            spln = Spline1D(xgrid, PFvec[1+S+sprime,:], k=3)
            splxprime = Spline1D(xgrid, PFvec[1+2S+sprime,:], k=3)
            splTT = Spline1D(xgrid, PFvec[1+3S+sprime,:], k=3)
            cf[s_,sprime] = y -> splc(y)
            nf[s_,sprime] = y -> spln(y)
            xprimef[s_,sprime] = y -> splxprime(y)
            TTf[s_,sprime] = y -> splTT(y)
        end
    end
    policies = (cf, nf, xprimef, TTf)
    return Vf, policies
end

"""
Computes Tau given c,n
"""
function Tau(pab::Planners_Allocation_Bellman,
            c::AbstractArray, n::AbstractArray)
    para = pab.para
    Uc, Un = para.Uc(c, n), para.Un(c, n)
    return 1+Un./(para.Theta .* Uc)
end

Tau(pab::Planners_Allocation_Bellman, c::Real, n::Real) =
    Tau(pab, [c], [n])

"""
Finds the optimal allocation given initial government debt B_ and state s_0
"""
function time0_allocation(pab::Planners_Allocation_Bellman,
                           B_::Real, s0::Integer)
    T, Vf = pab.T, pab.Vf
    xbar = T.xbar
    z0 = get_policies_time0(T, B_, s0, Vf, xbar)

    c0, n0, xprime0, T0 = z0[2], z0[3], z0[4], z0[5]
    return c0, n0, xprime0, T0
end

"""
Simulates planners policies for `T` periods
"""
function simulate{TF <: AbstractFloat}(pab::Planners_Allocation_Bellman,
                   B_::TF, s_0::Integer, T::Integer,
                   sHist::Vector=QuantEcon.simulate(pab.mc, T, init=s_0))
    para, mc, Vf, S = pab.para, pab.mc, pab.Vf, pab.S
    Pi, Uc = para.Pi, para.Uc
    cf, nf, xprimef, TTf = pab.policies

    cHist=Array{TF}(T)
    nHist=Array{TF}(T)
    Bhist=Array{TF}(T)
    xHist=Array{TF}(T)
    TauHist=Array{TF}(T)
    THist=Array{TF}(T)
    muHist=Array{TF}(T)

    #time0
    cHist[1], nHist[1], xHist[1], THist[1]  = time0_allocation(pab, B_, s_0)
    TauHist[1] = Tau(pab, cHist[1], nHist[1])[s_0]
    Bhist[1] = B_
    muHist[1] = Vf[s_0](xHist[1])

    #time 1 onward
    for t in 2:T
        s_, x, s = sHist[t-1], xHist[t-1], sHist[t]
        c = Array{TF}(S)
        n = Array{TF}(S)
        xprime = Array{TF}(S)
        TT = Array{TF}(S)
        for sprime=1:S
            c[sprime], n[sprime], xprime[sprime], TT[sprime] =
                cf[s_, sprime](x), nf[s_, sprime](x),
                xprimef[s_, sprime](x), TTf[s_, sprime](x)
        end

        Tau_val = Tau(pab, c, n)[s]
        u_c = Uc(c, n)
        Eu_c = dot(Pi[s_, :], u_c)

        muHist[t] = Vf[s](xprime[s])

        cHist[t], nHist[t], Bhist[t], TauHist[t] = c[s], n[s], x/Eu_c, Tau_val
        xHist[t], THist[t] = xprime[s], TT[s]
    end
    return cHist, nHist, Bhist, xHist, TauHist, THist, muHist, sHist
end

"""
Initializes the class from the calibration Para
"""
function BellmanEquation{TF <: AbstractFloat}(para::Para{TF},
                                              xgrid::AbstractVector{TF},
                                              policies0::Array)
    S = size(para.Pi, 1) # number of states
    xbar = [minimum(xgrid), maximum(xgrid)]
    time_0 = false
    z0 = Array{Array}(length(xgrid), S)
    cf, nf, xprimef = policies0[1], policies0[2], policies0[3]
    for s in 1:S
        for (i_x, x) in enumerate(xgrid)
            cs=Array{TF}(S)
            ns=Array{TF}(S)
            xprimes=Array{TF}(S)
            for j=1:S
                cs[j], ns[j], xprimes[j] = cf[s, j](x), nf[s, j](x), xprimef[s, j](x)
            end
            z0[i_x, s] = vcat(cs, ns, xprimes, zeros(S))
        end
    end
    cFB, nFB, IFB, xFB, zFB = LS.find_first_best(para, S, 2)
    return BellmanEquation(para, S, xbar, time_0, z0, cFB, nFB, xFB, zFB)
end

"""
Finds the optimal policies
"""
function get_policies_time1(T::BellmanEquation,
        i_x::Integer, x::Real, s_::Integer, Vf::AbstractArray{Function}, xbar::AbstractVector)
    para, S = T.para, T.S
    beta, Theta, G, Pi = para.beta, para.Theta, para.G, para.Pi
    U,Uc,Un = para.U, para.Uc, para.Un

    S_possible = sum(Pi[s_, :].>0)
    sprimei_possible = find(Pi[s_, :].>0)

    function objf(z, grad)
        c, xprime = z[1:S_possible], z[S_possible+1:2S_possible]
        n = (c+G[sprimei_possible])./Theta[sprimei_possible]
        Vprime = [Vf[sprimei_possible[si]](xprime[si]) for si in 1:S_possible]
        return -dot(Pi[s_, sprimei_possible], U.(c, n)+beta*Vprime)
    end

    function cons(out, z, grad)
        c, xprime, TT =
            z[1:S_possible], z[S_possible+1:2S_possible], z[2S_possible+1:3S_possible]
        n = (c+G[sprimei_possible])./Theta[sprimei_possible]
        u_c = Uc.(c, n)
        Eu_c = dot(Pi[s_, sprimei_possible], u_c)
        out .= x*u_c/Eu_c - u_c.*(c-TT)-Un(c, n).*n - beta*xprime
    end
    function cons_no_trans(out, z, grad)
        c, xprime = z[1:S_possible], z[S_possible+1:2S_possible]
        n = (c+G[sprimei_possible])./Theta[sprimei_possible]
        u_c = Uc.(c, n)
        Eu_c = dot(Pi[s_, sprimei_possible], u_c)
        out .= x*u_c/Eu_c - u_c.*c-Un(c, n).*n - beta*xprime
    end

    if para.transfers == true
        lb = vcat(zeros(S_possible), ones(S_possible)*xbar[1], zeros(S_possible))
        if para.n_less_than_one == true
            ub = vcat(100*ones(S_possible)-G[sprimei_possible],
                      ones(S_possible)*xbar[2], 100*ones(S_possible))
        else
            ub = vcat(100*ones(S_possible),
                      ones(S_possible)*xbar[2],
                      100*ones(S_possible))
        end
        init = vcat(T.z0[i_x, s_][sprimei_possible],
                    T.z0[i_x, s_][2S+sprimei_possible],
                    T.z0[i_x, s_][3S+sprimei_possible])
        opt = Opt(:LN_COBYLA, 3S_possible)
        equality_constraint!(opt, cons, zeros(S_possible))
    else
        lb = vcat(zeros(S_possible), ones(S_possible)*xbar[1])
        if para.n_less_than_one == true
            ub = vcat(ones(S_possible)-G[sprimei_possible], ones(S_possible)*xbar[2])
        else
            ub = vcat(ones(S_possible), ones(S_possible)*xbar[2])
        end
        init = vcat(T.z0[i_x, s_][sprimei_possible],
                    T.z0[i_x, s_][2S+sprimei_possible])
        opt = Opt(:LN_COBYLA, 2S_possible)
        equality_constraint!(opt, cons_no_trans, zeros(S_possible))
    end
    init[init .> ub] = ub[init .> ub]
    init[init .< lb] = lb[init .< lb]

    min_objective!(opt, objf)
    lower_bounds!(opt, lb)
    upper_bounds!(opt, ub)
    maxeval!(opt, 10000000)
    maxtime!(opt, 10)
    ftol_rel!(opt, 1e-8)
    ftol_abs!(opt, 1e-8)

    (minf, minx, ret) = optimize(opt, init)

    if ret != :SUCCESS && ret != :ROUNDOFF_LIMITED && ret != :MAXEVAL_REACHED &&
         ret != :FTOL_REACHED && ret != :MAXTIME_REACHED
        error("optimization failed: ret = $ret")
    end

    T.z0[i_x, s_][sprimei_possible] = minx[1:S_possible]
    T.z0[i_x, s_][S+sprimei_possible] = minx[1:S_possible]+G[sprimei_possible]
    T.z0[i_x, s_][2S+sprimei_possible] = minx[S_possible+1:2S_possible]
    if para.transfers == true
        T.z0[i_x, s_][3S+sprimei_possible] = minx[2S_possible+1:3S_possible]
    else
        T.z0[i_x, s_][3S+sprimei_possible] = zeros(S)
    end

    return vcat(-minf, T.z0[i_x, s_])
end

"""
Finds the optimal policies
"""
function get_policies_time0(T::BellmanEquation,
        B_::Real, s0::Integer, Vf::AbstractArray{Function}, xbar::AbstractVector)
    para = T.para
    beta, Theta, G = para.beta, para.Theta, para.G
    U, Uc, Un = para.U, para.Uc, para.Un

    function objf(z, grad)
        c, xprime = z[1], z[2]
        n=(c+G[s0])/Theta[s0]
        return -(U(c, n)+beta*Vf[s0](xprime))
    end

    function cons(z,grad)
        c, xprime, TT = z[1], z[2], z[3]
        n=(c+G[s0])/Theta[s0]
        return -Uc(c, n)*(c-B_-TT)-Un(c, n)*n - beta*xprime
    end
    cons_no_trans(z, grad) = cons(vcat(z, 0), grad)

    if para.transfers == true
        lb = [0.0, xbar[1], 0.0]
        if para.n_less_than_one == true
            ub = [1-G[s0], xbar[2], 100]
        else
            ub = [100.0, xbar[2], 100.0]
        end
        init = vcat(T.zFB[s0][1], T.zFB[s0][3], T.zFB[s0][4])
        init = [0.95124922, -1.15926816,  0.0]
        opt = Opt(:LN_COBYLA, 3)
        equality_constraint!(opt, cons)
    else
        lb = [0.0, xbar[1]]
        if para.n_less_than_one == true
            ub = [1-G[s0], xbar[2]]
        else
            ub = [100, xbar[2]]
        end
        init = vcat(T.zFB[s0][1], T.zFB[s0][3])
        init = [0.95124922, -1.15926816]
        opt = Opt(:LN_COBYLA, 2)
        equality_constraint!(opt, cons_no_trans)
    end
    init[init.> ub] = ub[init.> ub]
    init[init.< lb] = lb[init.< lb]

    min_objective!(opt, objf)
    lower_bounds!(opt, lb)
    upper_bounds!(opt, ub)
    maxeval!(opt, 100000000)
    maxtime!(opt, 30)

    (minf, minx, ret) = optimize(opt, init)

    if ret != :SUCCESS && ret != :ROUNDOFF_LIMITED && ret != :MAXEVAL_REACHED && ret != :FTOL_REACHED
        error("optimization failed: ret = $ret")
    end

    if para.transfers == true
        return -minf, minx[1], minx[1]+G[s0], minx[2], minx[3]
    else
        return -minf, minx[1], minx[1]+G[s0], minx[2], 0
    end
end


using Test
using CompScienceMeshes
using LinearAlgebra
using StaticArrays

for T in [Float64, Float32]
    rectangle = meshrectangle(T(1.0), T(1.0), T(1.0));

    @test vertices(rectangle,1) == point(T, 0.0, 0.0, 0.0)
    @test vertices(rectangle,2) == point(T, 0.0, 1.0, 0.0)
    @test vertices(rectangle,3) == point(T, 1.0, 0.0, 0.0)

    @test numvertices(rectangle) == 4
    @test numcells(rectangle) == 2
    @test dimension(rectangle) == 2

    faces = skeleton(rectangle, 2)
    edges = skeleton(rectangle, 1)
    verts = skeleton(rectangle, 0)

    @test numcells(verts) == 4
    @test numcells(edges) == 5
    @test numcells(faces) == 2

    Λ  = connectivity(verts, edges, sign)
    Σᵀ = connectivity(edges, faces, sign)

    @test size(Λ)  == (5,4)
    @test size(Σᵀ) == (2,5)

    @test norm(Σᵀ*Λ, Inf) ≈ 0

    bnd = boundary(rectangle)

    @test dimension(bnd) == 1
    @test numvertices(bnd) == 4
    @test numcells(bnd) == 4

    vtoe, num_vtoe = vertextocellmap(edges)
    vtof, num_vtof = vertextocellmap(faces)

    @test size(vtoe) == (4,3)
    @test size(vtof) == (4,2)

    @test size(num_vtoe) == (4,)
    @test size(num_vtof) == (4,)

    file = joinpath(dirname(@__FILE__), "sphere.in")
    sphere = readmesh(file)
    @test numvertices(sphere) == 335
    @test numcells(sphere) == 666

    m = meshrectangle(T(1.0),T(1.0),T(1.0))
    CompScienceMeshes.rotate!(m, [1,0,0]*T(π/2))

    @test CompScienceMeshes.isoriented(boundary(m))

    tol = sqrt(eps(T))
    @test norm(vertices(m,1) - point(T,0.0,0.0,0.0)) < tol
    @test norm(vertices(m,2) - point(T,0.0,0.0,1.0)) < tol
    @test norm(vertices(m,3) - point(T,1.0,0.0,0.0)) < tol
    @test norm(vertices(m,4) - point(T,1.0,0.0,1.0)) < tol

    ## test fliporientation
    V = [
        point(T,0,0,0),
        point(T,1,0,0),
        point(T,0,1,0),
    ]

    F = [
        index(1,2,3)
    ]

    m = Mesh(V,F)
    n = fliporientation(m)

    @test m.faces[1] == index(1,2,3)
    @test n.faces[1] == index(2,1,3)

    # test the construction of transposed connectivity matricees
    Σᵀ = connectivity(edges, faces, identity)
    Σ = connectivity(faces, edges, identity)
    @test norm(Σᵀ - Σ', Inf) == 0
end


@testset "Statistics" begin
  @testset "Solution" begin
    grid2D = RegularGrid{Float64}(3,3)
    realizations = Dict(:value => [i*ones(npoints(grid2D)) for i in 1:3])
    solution2D = SimulationSolution(grid2D, realizations)

    # mean
    mean2D = mean(solution2D)
    @test mean2D[1,:value] == 2.
    @test domain(mean2D) == solution2D.domain

    # variance
    variance2D = var(solution2D)
    @test variance2D[1,:value] == 1.
    @test domain(variance2D) == solution2D.domain

    # quantile (scalar)
    p = 0.5
    quantile2D = quantile(solution2D, p)
    @test quantile2D[1,:value] == 2.
    @test domain(quantile2D) == solution2D.domain

    # quantile (vector)
    ps = [0.0, 0.5, 1.0]
    quantiles2D = quantile(solution2D, ps)
    @test eltype(quantiles2D) <: SpatialStatistic

    # show methods
    @test sprint(show, mean2D) == "2D SpatialStatistic"
    @test sprint(show, MIME"text/plain"(), mean2D) == "2D SpatialStatistic\n  domain: 3×3 RegularGrid{Float64,2}\n  variables: value"
  end

  @testset "Data" begin
    # load data with bias towards large values (gold mine)
    sdata = readgeotable(joinpath(datadir,"clustered.csv"), coordnames=[:x,:y])

    # spatial mean
    μn = mean(sdata[:Au])
    μs = mean(sdata, :Au)
    @test abs(μn - 0.5) > abs(μs - 0.5)
    @test mean(sdata)[:Au] ≈ μs

    # spatial variance
    σn = var(sdata[:Au])
    σs = var(sdata, :Au)
    @test σn ≤ σs
    @test var(sdata)[:Au] ≈ σs

    # spatial quantile
    qn = quantile(sdata[:Au], 0.5)
    qs = quantile(sdata, :Au, 0.5)
    @test qn ≥ qs
    @test quantile(sdata, 0.5)[:Au] ≈ qs
  end
end


## Linear Solve

function linsolve(A::Operator,b;kwds...)
    if dimension(domain(A)) > 1
        pdesolve(A,b;kwds...)
    elseif isambiguous(domainspace(A))
        A=choosespaces(A,b)
        if isambiguous(domainspace(A))
            error("Cannot infer spaces")
        end
        linsolve(A,b;kwds...)
    else
        linsolve(qrfact(A),b;kwds...)
    end
end

function linsolve{OO<:Operator}(A::Array{OO},b;kwds...)
    if dimension(domain(A[end])) > 1
        pdesolve(A,b;kwds...)
    else
        linsolve(interlace(A),b;kwds...)
    end
end

for p in (1,2)
    @eval begin
        \{T<:Operator}(A::Array{T,$p},b::Union{Array,Number,Fun}) = linsolve(A,b)
        \{T<:Operator,F<:Fun}(A::Array{T,$p},b::Vector{F}) = linsolve(A,b)
        \{T<:Operator,F<:Number}(A::Array{T,$p},b::Vector{F}) = linsolve(A,b)
    end
end
\(A::Operator,b) = linsolve(A,b)


# mean, median, std etc.. functions don't work on Distribution.Truncated instances,
# so we'll define them using a resampling approach

import Distributions.Truncated
import StatsBase:
    mean,
    median,
    middle,
    std,
    mode,
    quantile

mean(d::Truncated; n::Int = 10000) = mean(rand(d, n))
median(d::Truncated; n::Int = 10000) = median(rand(d, n))
middle(d::Truncated; n::Int = 10000) = median(rand(d, n))
std(d::Truncated; n::Int = 10000) = std(rand(d, n))
mode(d::Truncated; n::Int = 10000) = mode(rand(d, n))
quantile(d::Truncated, q; n::Int = 10000) = quantile(rand(d, n), q)


"""
    Monin_Obukhov_length(Tair, pressure, ustar, H; constants)
    Monin_Obukhov_length!(df;constants=BigleafConstants())
    Monin_Obukhov_length(df;constants=BigleafConstants())

calculates the Monin-Obukhov length.

# Arguments
- `Tair`      : Air temperature (degC)
- `pressure`  : Atmospheric pressure (kPa)
- `ustar`     : Friction velocity (m s-1)
- `H`         : Sensible heat flux (W m-2)
- `df`        : DataFrame containing the above variables
optional
- `constants=`[`BigleafConstants`](@ref)`()`: Dictionary with entries
  - `Kelvin` - conversion degree Celsius to Kelvin
  - `cp` - specific heat of air for constant pressure (J K-1 1)
  - `k` - von Karman constant (-)
  - `g` - gravitational acceleration (m s-2)

# Details
The Monin-Obukhov length (L) is given by:

``L = - (\\rho * cp * ustar^3 * Tair) / (k * g * H)``

where ``\\rho`` is air density (kg m-3).

# Note
Note that L gets very small for very low ustar values with implications
for subsequent functions using L as input. It is recommended to filter
data and exclude low ustar values (ustar < ~0.2) beforehand.

# Value
Monin-Obukhov length L (m). The non-mutating DataFrame variant returns a vector,
the mutating variant add or modifies column `:MOL`.

# References
Foken, T, 2008: Micrometeorology. Springer, Berlin, Germany.

# See also
[`stability_parameter`](@ref)

```@example; output = false
Monin_Obukhov_length(
  Tair=25,pressure=100,
  ustar=seq(0.2,1,0.1),H=seq(40,200,20))
```
"""
function Monin_Obukhov_length(Tair::FT, pressure, ustar, H; constants=BigleafConstants()
  ) where FT
  rho  = air_density(Tair, pressure; constants)
  TairK = Tair + oftype(Tair, constants.Kelvin)
  #MOL  = (-rho*constants.cp*ustar^3*TairK) / (constants.k*constants.g*H)
  MOL  = (-rho*FT(constants.cp)*ustar^3*TairK) / (FT(constants.k)*FT(constants.g)*H)
end
function Monin_Obukhov_length!(df;constants=BigleafConstants())
  fr = (args...) -> Monin_Obukhov_length(args...; constants)
  transform!(df, SA[:Tair, :pressure, :ustar, :H] => ByRow(fr) => :MOL)
end
# function Monin_Obukhov_length(df::DFTable; kwargs...)
#     tmp = Monin_Obukhov_length.(df.Tair, df.pressure, df.ustar, df.H; kwargs...)
# end

"""
    stability_parameter(z,d,MOL)
    stability_parameter(z,d,Tair, pressure, ustar, H; constants)
    stability_parameter!(df::AbstractDataFrame; z,d, MOL=nothing, constants)

calculates stability parameter "zeta", a parameter characterizing stratification in the
lower atmosphere.

# Arguments
- `z`   : height (m)
- `d`   : Zero-plane displacement height (m)
- `MOL` : Monin-Obukhov-length L (m)
- `df`  : DataFrame containting the variables required by [`Monin_Obukhov_length`](@ref)
optional
- `constants=`[`BigleafConstants`](@ref)`()`

In the second variant and if `MOL=nothing` in the DataFrame variants, MOL is computed by
[`Monin_Obukhov_length`](@ref).

# Details
The stability parameter ``\\zeta`` is given by:

``\\zeta = (z - d) / L``

where L is the Monin-Obukhov length (m), calculated by
[`Monin_Obukhov_length`](@ref). The displacement height can
be estimated from the function [`roughness_parameters`](@ref).

# Value
``\\zeta``: stability parameter (-). The nonmutainting DataFrame variant returns a vector.
The mutating variant modifies or adds column [`:zeta`].

```jldoctest; output = false
using DataFrames
df = DataFrame(Tair=25.0, pressure=100.0, ustar=0.2:0.1:1.0, H=40:20.0:200)
z=40;d=15
zeta = stability_parameter.(z,d, df.Tair, df.pressure, df.ustar, df.H)
all(zeta .< 0)
# output
true
```
"""
function stability_parameter(z,d,MOL)
  zeta = (z - d) / MOL
end
function stability_parameter(z,d,Tair, pressure, ustar::FT, H;
  constants=BigleafConstants()) where FT
  MOL = Monin_Obukhov_length(Tair, pressure, ustar, H; constants);
  stability_parameter(FT(z),FT(d),MOL)
end
function stability_parameter!(df::AbstractDataFrame; z,d, MOL=nothing,
  constants=BigleafConstants())
  if isnothing(MOL)
    ft = (args...) -> stability_parameter.(z,d,args...; constants)
    transform!(df, SA[:Tair, :pressure, :ustar, :H] => ft => :zeta)
  else
    ft = () -> stability_parameter.(z,d,MOL)
    transform!(df, [] => ft => :zeta)
  end
end
# function stability_parameter(df::DFTable; z,d, MOL=nothing,
#   constants=BigleafConstants())
#   if isnothing(MOL) MOL = Monin_Obukhov_length(df; constants); end
#   stability_parameter.(z,d,MOL)
# end

abstract type StabilityCorrectionMethod end
struct Dyer1970 <: StabilityCorrectionMethod end
struct Businger1971 <: StabilityCorrectionMethod end
struct NoStabilityCorrection <: StabilityCorrectionMethod end

"""
    stability_correction(zeta;
      stab_formulation=Dyer1970())
    stability_correction(z,d, Tair,pressure,ustar,H; constants,
      stab_formulation=Dyer1970())
    stability_correction!(df; zeta, z, d;
      stab_formulation=Dyer1970(), constants =BigleafConstants())

Integrated Stability Correction Functions for Heat and Momentum

# Arguments
- `zeta`             : Stability parameter zeta (-)
- `Tair`,`pressure`,`ustar`,`H` : see [`Monin_Obukhov_length`](@ref)
- `z`,`d`            : see [`stability_parameter`](@ref)
- `df`  : DataFrame containting the variables required by [`Monin_Obukhov_length`](@ref)
- `stab_formulation` : Formulation for the stability function. Either
            `Dyer1970()`, or `Businger1971()` or `NoStabilityCorrection()`

In the second and third form computes `zeta` by [`stability_parameter`](@ref) and
[`Monin_Obukhov_length`](@ref) and requires respective arguments.

# Details
These dimensionless values are needed to correct deviations
from the exponential wind profile under non-neutral conditions.
The functions give the integrated form of the universal functions. They
depend on the value of the stability parameter ``\\zeta``,
a function of heigh `z`,
which can be calculated from the function [`stability_parameter`](@ref).
The integration of the universal functions is:

``\\psi = -x * \\zeta``

for stable atmospheric conditions (``\\zeta`` >= 0), and

``\\psi = 2 * log( (1 + y(\\zeta)) / 2)``

for unstable atmospheric conditions (``\\zeta`` < 0).

The different formulations differ in their value of x and y.

# Value
a NamedTuple with the following columns:
- `psi_h`: the value of the stability function for heat and water vapor (-)
- `psi_m`: the value of the stability function for momentum (-)

# References
- Dyer, A_J., 1974: A review of flux-profile relationships.
  Boundary-Layer Meteorology 7, 363-372.
- Dyer, A. J., Hicks, B_B., 1970: Flux-Gradient relationships in the
  constant flux layer. Quart. J. R. Meteorol. Soc. 96, 715-721.
- Businger, J_A., Wyngaard, J. C., Izumi, I., Bradley, E. F., 1971:
  Flux-Profile relationships in the atmospheric surface layer.
  J. Atmospheric Sci. 28, 181-189.
- Paulson, C_A., 1970: The mathematical representation of wind speed
  and temperature profiles in the unstable atmospheric surface layer.
  Journal of Applied Meteorology 9, 857-861.
  Foken, T, 2008: Micrometeorology. Springer, Berlin, Germany.

# Examples
```jldoctest; output = false
using DataFrames
zeta = -2:0.5:0.5
df2 = DataFrame(stability_correction.(zeta; stab_formulation=Businger1971()))
propertynames(df2) == [:psi_h, :psi_m]
# output
true
```
"""
function stability_correction(zeta::FT; stab_formulation=Dyer1970()) where FT
  # integration of universal functions (after Paulson1970 and Foken 2008)
  stab_formulation isa NoStabilityCorrection && return(
    (psi_h = zero(zeta), psi_m = zero(zeta)))
  ismissing(zeta) && return((psi_h = missing, psi_m = missing))
  is_stable = zeta >= 0
  if is_stable
    x_h::FT, x_m::FT = get_stability_coefs_stable(stab_formulation)
    psi_h = x_h * zeta
    psi_m = x_m * zeta
  else
    y_h::FT, y_m::FT = get_stability_coefs_unstable(stab_formulation, zeta)
    psi_h = 2 * log( (1 + y_h ) / 2)
    psi_m = 2 * log( (1 + y_m ) / 2) +
                      log( ( 1 + y_m^2 ) / 2)
                      -2 * atan(y_m) + pi/2
  end
  (;psi_h, psi_m)
end

get_stability_coefs_stable(::Businger1971) = (x_h = -7.8, x_m = -6)
function get_stability_coefs_unstable(::Businger1971, zeta)
  y_h = 0.95 * ( 1 - 11.6 * zeta)^0.5
  y_m = (1 - 19.3*zeta)^0.25
  (;y_h, y_m)
end

get_stability_coefs_stable(::Dyer1970) = (x_h = -5, x_m = -5)
function get_stability_coefs_unstable(::Dyer1970, zeta)
  y_h       = (1 - 16 * zeta)^0.5
  y_m       = (1 - 16 * zeta)^0.25
  (;y_h, y_m)
end

function stability_correction(z,d, Tair::Union{Missing,Number},pressure,ustar::FT,H;
  stab_formulation=Dyer1970(), constants=BigleafConstants()) where FT
  stab_formulation isa NoStabilityCorrection && return(
    (psi_h = zero(FT), psi_m = zero(FT)))
  FT <: Missing && return(missing)
  MOL = Monin_Obukhov_length(Tair,pressure,ustar,H; constants)
  zeta  = stability_parameter(FT(z),FT(d),MOL)
  psis = stability_correction(zeta; stab_formulation)
  psis
end
# function stability_correction(Tair,pressure,ustar,H, z,d; kwargs...)
#   Tables.columns(
#     stability_correction.(Tair,pressure,ustar,H, z,d; kwargs...))
# end
# function stability_correction(df::DFTable, z,d; kwargs...)
#   tmp1 = stability_correction.(df.Tair, df.pressure, df.ustar, df.H, z, d; kwargs...)
#   tmp2 = Tables.columns(tmp1)
# end
function stability_correction!(df; zeta=nothing, z=nothing, d=nothing,
  stab_formulation=Dyer1970(), constants =BigleafConstants())
  # cannot dispatch on keyword argument, hence need if-clause
  if stab_formulation isa NoStabilityCorrection
    zeroT = :ustar in propertynames(df) ? zero(eltype(skipmissing(df.ustar))) : zero(z)
    df[!,:psi_h] .= zeroT
    df[!,:psi_m] .= zeroT
    return(df)
  end
  if !isnothing(zeta)
    ft = () -> stability_correction.(zeta; stab_formulation)
    transform!(df, [] => ft => AsTable)
  else
    ft = (args...) -> stability_correction.(z, d, args...; stab_formulation, constants)
    transform!(df, SA[:Tair,:pressure,:ustar,:H] => ft => AsTable)
  end
end

function stability_correction(df::DFTable; z, d,
  stab_formulation=Dyer1970(), constants =BigleafConstants()
  )
  # do not provide zeta, because can simply invoke stability_correction.(zeta)
  if stab_formulation isa NoStabilityCorrection
    zeroT = zero(eltype(skipmissing(df.ustar)))
    rows = map(Tables.rows(df)) do row
      (psi_h = zeroT, psi_m = zeroT)
    end
    return(Tables.columns(rows))
  end
  Tables.columns(stability_correction.(
    z, d, df.Tair, df.pressure, df.ustar, df.H; stab_formulation, constants))
end


aux_field(::BernoulliLikelihood{<:LogisticLink}, Ω) = getproperty(Ω, :ω)

function init_aux_variables(rng::AbstractRNG, ::BernoulliLikelihood{<:LogisticLink}, n::Int)
    return TupleVector(; ω=rand(rng, PolyaGamma(1, 0.0), n))
end

function init_aux_posterior(T::DataType, ::BernoulliLikelihood{<:LogisticLink}, n::Int)
    return For(TupleVector(; c=zeros(T, n))) do φ
        NTDist(PolyaGamma(1, φ.c))
    end
end

function aux_full_conditional(::BernoulliLikelihood{<:LogisticLink}, ::Any, f::Real)
    return NTDist(PolyaGamma(1, abs(f)))
end

function aux_posterior!(
    qΩ,
    ::BernoulliLikelihood{<:LogisticLink},
    ::AbstractVector,
    qf::AbstractVector{<:Normal},
)
    map!(sqrt ∘ second_moment, only(qΩ.inds).c, qf)
    return qΩ
end

function auglik_potential(::BernoulliLikelihood{<:LogisticLink}, ::Any, y::AbstractVector)
    return (sign.(y .- 0.5) / 2,)
end

function auglik_precision(::BernoulliLikelihood{<:LogisticLink}, Ω, ::AbstractVector)
    return (Ω.ω,)
end

function expected_auglik_potential(
    lik::BernoulliLikelihood{<:LogisticLink}, qΩ, y::AbstractVector
)
    return auglik_potential(lik, qΩ, y)
end

function expected_auglik_precision(
    ::BernoulliLikelihood{<:LogisticLink}, qΩ, ::AbstractVector
)
    return (tvmean(qΩ).ω,)
end

function logtilt(::BernoulliLikelihood{<:LogisticLink}, ω::Real, y::Real, f::Real)
    return -log(2) + (sign(y - 0.5) * f - abs2(f) * ω) / 2
end

function aux_prior(lik::BernoulliLikelihood{<:LogisticLink}, y::AbstractVector)
    return For(length(y)) do _
        NTDist(aux_prior(lik))
    end
end

aux_prior(::BernoulliLikelihood{<:LogisticLink}) = PolyaGamma(1, 0.0)

function expected_logtilt(
    ::BernoulliLikelihood{<:LogisticLink}, qωᵢ::NTDist{<:PolyaGamma}, yᵢ::Real, qfᵢ::Normal
)
    m = mean(qfᵢ)
    θ = ntmean(qωᵢ)
    return -log(2) + (sign(yᵢ - 0.5) * m - (abs2(m) + var(qfᵢ)) * θ.ω) / 2
end


__precompile__(false)
module DiffDRO

using Conda
using PyCall

export DiffBasicWeightConvexSet,
    DiffBertsimas,
    DiffBenTal,
    DiffDelage,
    utility,
    expected_return,
    std_adjusted_expected_return,
    var_adjusted_expected_return

Conda.add("pip")

run(PyCall.python_cmd(`-m pip install numpy`))
run(PyCall.python_cmd(`-m pip install scs==2.1.3`))
run(PyCall.python_cmd(`-m pip install cvxpy`))
run(PyCall.python_cmd(`-m pip install torch`))
run(PyCall.python_cmd(`-m pip install cvxpylayers`))

py"""
import numpy as np
import cvxpy as cp
import torch
from cvxpylayers.torch import CvxpyLayer
"""

include("DiffBasicWeightConvexSet.jl")
include("DiffBertsimas.jl")
include("DiffBenTal.jl")
include("DiffDelage.jl")
include("DiffObj.jl")

include("test_utils.jl")

end


MatrixOrFunction = Union{AbstractMatrix,Function}
VectorOrFunction = Union{AbstractVector,Function}
ArrayOrFunction = Union{AbstractArray,Function}
RealOrFunction = Union{Real,Function}

#Delay behavoiur: stochastic or deterministic / Zero order holders or continuous delay

struct Delay{T}
    τ::T
end
(τ::Delay{<:Real})(t) = τ.τ
(τ::Delay{<:Function})(t) = τ.τ(t)
Delay(τ::T1, holdT::T1) where {T1 <: Real} = Delay(t->τ + t % holdT)
Base.convert(::Type{Delay}, τ::RealOrFunction) = Delay(τ)
Base.convert(::Type{Vector{<:Delay}}, τs::Vector{<:RealOrFunction}) = Delay.(τs)

abstract type Coefficients{d} end
abstract type CoefficientMatrix{d} <: Coefficients{d} end
abstract type AdditiveVector{d} <: Coefficients{d} end

Base.size(cs::T) where T<:CoefficientMatrix{d} where d = (d, d)
Base.size(cs::T) where T<:AdditiveVector{d} where d = (d,)
Base.size(cs::T, ::Integer) where T<:Coefficients{d} where d = d
Base.length(cs::T) where T<:AdditiveVector{d} where d = d
Base.length(cs::T) where T<:CoefficientMatrix{d} where d = d^2

struct ProportionalMX{d,T} <: CoefficientMatrix{d}
    MX::T # matrix
end
ProportionalMX(mx::mxT) where mxT <: Function = ProportionalMX{size(mx(0.), 1),mxT}(mx)
ProportionalMX(mx::mxT) where mxT <: SMatrix = ProportionalMX{size(mx, 1),mxT}(mx)
ProportionalMX(mx::mxT) where mxT <: AbstractMatrix{<:Real} = ProportionalMX(SMatrix{size(mx)...}(mx))
(PMX::ProportionalMX{d,<:Function})(t) where d = PMX.MX(t)
(PMX::ProportionalMX{d,<:AbstractMatrix{<:Real}})(t) where d = PMX.MX
Base.convert(::Type{ProportionalMX}, mx::mT) where mT <: AbstractMatrix{<:Real} = ProportionalMX(mx)

struct DelayMX{d,dT,bT} <: CoefficientMatrix{d}
    τ::Delay{dT} # delay
    MX::bT # matrix
end

DelayMX(τ::RealOrFunction, MX::MatrixOrFunction) = DelayMX(Delay(τ), MX)
DelayMX(τ::Delay{dT}, MX::mxT) where {dT,mxT<:Function} = DelayMX{size(MX(0.),1),dT,mxT}(τ, MX)
DelayMX(τ::Delay{dT}, MX::mxT) where {dT,mxT <: SMatrix} = DelayMX{size(MX, 1),dT,mxT}(τ, MX)

DelayMX(τ::Delay{dT}, MX::mxT) where {dT,mxT<:AbstractMatrix{<:Real}} = DelayMX(τ,SMatrix{size(MX)...}(MX))
(DMX::DelayMX{d,<:dT,<:AbstractMatrix{<:Real}})(t) where {d,dT} = DMX.MX
(DMX::DelayMX{d,<:dT,<:Function})(t) where {d,dT} = DMX.MX(t)

struct Additive{d,T} <: AdditiveVector{d}
    V::T
end
Additive(v::vT) where vT <: Function = Additive{length(v(0.)),vT}(v)
Additive(v::vT) where vT <: SVector{d} where d = Additive{d,vT}(v)
Additive(v::vT) where vT <: AbstractVector{<:Real} = Additive(SVector(v...))
Additive(d::Integer) = Additive{d,Vector{Nothing}}(Vector{Nothing}(undef, d))
(AV::Additive{d,<:AbstractVector{<:Real}})(t) where d = AV.V
(AV::Additive{d,<:Function})(t) where d = AV.V(t)
(AV::Additive{d,<:Array{<:Nothing}})(t) where d = zeros(d) # TODO check the effect on performance
Base.convert(::Type{Additive}, v::vT) where vT <: AbstractVector{<:Real} = Additive(v)

# Linear Delay Differential Equation Problem
abstract type AbstractLDDEProblem{d} end
struct LDDEProblem{d,AT <: ProportionalMX{d,<:MatrixOrFunction},BT <: DelayMX{d,<:Any,<:Any}, cT <: Additive{d,<:VectorOrFunction}} <: AbstractLDDEProblem{d}
    A::AT # Coefficient of the proportional term (present /non-discretised term)
    Bs::Vector{BT} # Coefficient of the delayed terms
    c::cT # Additive vector
    # d::Int64 # dimension of the state space
end
_problemsize(::LDDEProblem{d,AT,BT,cT}) where {d,AT,BT,cT} = d # size of the state space

LDDEProblem(A::AT, Bs::Vector{<:BT}, c::cT = Additive(size(A, 2))) where {AT <: ProportionalMX,BT <: DelayMX,cT <: Additive} = LDDEProblem{size(A, 2),AT,BT,cT}(A, Bs, c)
LDDEProblem(A::AT, B::BT, c::cT = Additive(size(A, 2))) where {AT <: ProportionalMX,BT <: DelayMX,cT <: Additive} = LDDEProblem{size(A, 2),AT,BT,cT}(A, [B], c)


## Stability derivative cases
function scale_inputs(freestream :: Freestream, b, c)
    # Building input vector
    x0 = [ freestream.alpha;
           freestream.beta;
           rate_coefficient(freestream.omega, freestream.V, b, c) ]

    # Unscaling non-dimensional rate coefficients
    scale = 2freestream.V .* [1/b, 1/c, 1/b]

    x0, scale
end

function print_case(nf, ff, derivs, comp)
    print_coefficients(nf, ff, comp)
    print_derivatives(derivs, comp)
end

function print_case(data, comp)
    nf, ff, derivs = data[comp]
    print_case(nf, ff, derivs, comp)
end

function solve_stability_case(wing :: Union{Wing, HalfWing}, freestream :: Freestream; rho_ref = 1.225, r_ref = zeros(3), area_ref = projected_area(wing), chord_ref = mean_aerodynamic_chord(wing), span_ref = span(wing), span_num :: Union{Integer, Vector{<: Integer}}, chord_num :: Integer, name = "Wing", viscous = false, x_tr = 0.3, print = false, spacing = ["sine"; fill("cosine", length(span_num) - 1)])
    # Reference values and scaling inputs
    S, b, c = area_ref, span_ref, chord_ref
    x, scale = scale_inputs(freestream, b, c)

    # Closure to generate results with input vector
    function stab(x)
        fs 	 = Freestream(freestream.V, rad2deg(x[1]), rad2deg(x[2]), x[3:end] .* scale)
        data = solve_case(wing, fs;
                          rho_ref   = rho_ref,
                          r_ref     = r_ref,
                          area_ref  = S,
                          span_ref  = b,
                          chord_ref = c,
                          chord_num = chord_num,
                          span_num  = span_num,
                          viscous 	= viscous,
                          x_tr 		= x_tr,
                          spacing   = spacing)

        [ data[1]; data[2] ] # Concatenate nearfield and farfield coefficients for DiffResults value
    end

    # Set up Jacobian system and evaluate
    y = ifelse(viscous, zeros(16), zeros(12))
    result = DiffResults.JacobianResult(y, x)
    result = ForwardDiff.jacobian!(result, stab, x)
    vars   = DiffResults.value(result)
    derivs = DiffResults.jacobian(result)

    # Gather results
    nf 	= ifelse(viscous, vars[1:11], vars[1:9]) 	  	# Nearfield coefficients
    ff 	= ifelse(viscous, vars[12:end], vars[10:end]) 	# Farfield coefficients
    dvs = ifelse(viscous, derivs[[1,4:8...],:], derivs[1:6,:])	# Nearfield stability derivatives, uses total drag for viscous cases just in case of generalisations

    # Print if needed
    if print; print_case(nf, ff, derivs, name) end

    nf, ff, dvs
end

function solve_stability_case(aircraft :: Dict{String, Tuple{Matrix{Panel3D{T}}, Matrix{SVector{3,T}}}}, freestream :: Freestream; rho_ref = 1.225, r_ref = zeros(3), area_ref = 1, chord_ref = 1, span_ref = 1, name = "Aircraft", print = false, print_components = false) where T <: Real
    # Reference values and scaling inputs
    S, b, c = area_ref, span_ref, chord_ref
    x, scale = scale_inputs(freestream, b , c)

    # Closure to generate results with input vector
    function stab(x)
        fs 	 = Freestream(freestream.V, rad2deg(x[1]), rad2deg(x[2]), x[3:end] .* scale)
        data = solve_case(aircraft, fs,
                          rho_ref   = rho_ref,
                          r_ref     = r_ref,
                          area_ref  = S,
                          span_ref  = b,
                          chord_ref = c,
                          name 		= name)

        # Creates array of nearfield and farfield coefficients for each component as a row vector.
        coeffs = reduce(hcat, let comps = data[name]; [ comps[1]; comps[2] ] end for name in names)
    end

    names 	= [ reduce(vcat, keys(aircraft)); name ]
    num_comps = length(names)

    y 		= zeros(12, num_comps)
    result 	= DiffResults.JacobianResult(y, x)
    result 	= ForwardDiff.jacobian!(result, stab, x)

    vars 	= DiffResults.value(result)
    derivs 	= DiffResults.jacobian(result)

    # Dictionary assembly
    ranges  = (1:num_comps) .* 12                     # Painful hacking
    bounds  = zip([ 1; ranges[1:end-1] .+ 1], ranges) # Painful hacking
    data 	= Dict(name => (vars[1:9, i], vars[10:end, i], derivs[first(inds):last(inds)-6,:]) for (i, (name, inds)) in (enumerate ∘ zip)(names, bounds))

    # Printing
    if print;            print_case(data, name)       end
    if print_components; print_case.(Ref(data), names)end

    data
end


using Printf
using Random: seed!
using Statistics: median

using BenchmarkTools
using Optim

using DZOptimization
using DZOptimization.ExampleFunctions:
    rosenbrock_objective, rosenbrock_gradient!,
    riesz_energy, riesz_gradient!

function run_dzopt_benchmark(Optimizer, name, f, g!,
                             initial_point, initial_step_size)

    membench = @benchmark begin
        while !opt.has_converged[]
            step!(opt)
        end
    end samples=1 evals=1 setup=(
        opt = $Optimizer($f, $g!, $initial_point, $initial_step_size))
    @assert membench.memory == 0
    @assert membench.allocs == 0

    result = @benchmark begin
        opt = $Optimizer($f, $g!, $initial_point, $initial_step_size)
        while !opt.has_converged[]
            step!(opt)
        end
    end

    final = DZOptimization.run_and_test!(Optimizer(
        f, g!, initial_point, initial_step_size))

    @printf("%s on %s: min %.16e in %5d iters; %8d ns (%8d bytes; %6d allocs)\n",
            rpad(string(Optimizer), 25), rpad(name, 10),
            final.current_objective_value[],
            final.iteration_count[],
            median(result.times), result.memory, result.allocs)
end

function run_optim_benchmark(Optimizer, name, f, g!, initial_point)

    result = @benchmark optimize($f, $g!, $initial_point,
                                 method=$Optimizer(),
                                 iterations=typemax(Int), g_tol=0.0)

    optimum = optimize(f, g!, initial_point,
                       method=Optimizer(),
                       iterations=typemax(Int), g_tol=0.0)

    @printf("%s on %s: min %.16e in %5d iters; %8d ns (%8d bytes; %6d allocs)\n",
            rpad(string(Optimizer), 25), rpad(name, 10),
            optimum.minimum,
            optimum.iterations,
            median(result.times), result.memory, result.allocs)

end

function run_comparative_benchmark(DZType, OptimType)

    run_dzopt_benchmark(DZType, "rosenbrock",
        rosenbrock_objective, rosenbrock_gradient!, [0.0, 0.0], 0.03)

    run_optim_benchmark(OptimType, "rosenbrock",
        rosenbrock_objective, rosenbrock_gradient!, [0.0, 0.0])

    seed!(0)
    points = 2.0 * rand(3, 10) .- 1.0

    run_dzopt_benchmark(DZType, "riesz",
        L2RegularizationWrapper(riesz_energy, 1.0),
        L2GradientWrapper(riesz_gradient!, 1.0),
        points, 0.01)

    run_optim_benchmark(OptimType, "riesz",
        L2RegularizationWrapper(riesz_energy, 1.0),
        L2GradientWrapper(riesz_gradient!, 1.0),
        points)

end

function main()
    run_comparative_benchmark(GradientDescentOptimizer, GradientDescent)
    run_comparative_benchmark(BFGSOptimizer, BFGS)
    run_comparative_benchmark(LBFGSOptimizer, LBFGS)
end

main()


sample = [3, 8, 9, 1, 2, 5, 4, 6, 7]
input = [1, 2, 3, 4, 8, 7, 5, 9, 6]

mutable struct Cup
    val::Int
    next::Union{Cup,Nothing}
end

function link!(cups)
    for i in 1:length(cups)
        next = i + 1 > length(cups) ? 1 : i + 1
        cups[i].next = cups[next]
    end
end

function play!(cups, rounds)
    addresses = Dict()
    for cup in cups
        addresses[cup.val] = cup
    end

    cmax = maximum(map(c -> c.val, cups))

    current = first(cups)

    for i in 1:rounds
        rem = rem!(current)
        d = dest(current, rem, cmax, addresses)
        ins!(d, rem)

        current = current.next
    end
end

function rem!(current)
    x = current.next
    y = x.next
    z = y.next
    current.next = z.next
    return [x, y, z]
end

function dest(current, rem, cmax, addresses)
    used = append!(map(c -> c.val, rem), [current.val])
    i = current.val
    while i in used
        i -= 1
        if i == 0
            i = cmax
        end
    end
    return addresses[i]
end

function ins!(dest, rem)
    tmp = dest.next
    dest.next = first(rem)
    last(rem).next = tmp
end

function val(cups)
    res = ""
    current = first(cups)

    while current.val != 1
        current = current.next
    end
    current = current.next
    while current.val != 1
        res = string(res, current.val)
        current = current.next
    end
    return parse(Int, res)
end

function stars(cups)
    current = first(cups)

    while current.val != 1
        current = current.next
    end
    return current.next.val * current.next.next.val
end

cups = map(n -> Cup(n, nothing), input)
link!(cups)

play!(cups, 100)
p1 = val(cups)

cups = map(n -> Cup(n, nothing), append!(input, collect((maximum(input) + 1):1000000)))
link!(cups)

play!(cups, 10000000)
p2 = stars(cups)

println("-----------------------------------------------------------------------")
println("crab cups -- part one :: $p1")
println("crab cups -- part two :: $p2")
println("-----------------------------------------------------------------------")

@assert(p1 == 47598263)
@assert(p2 == 248009574232)


using Revise
using ADCME
using PyCall
using Test
using MAT
using NNFEM
using PyPlot
using JLD2
using LinearAlgebra
using Distributions, Random
using ForwardDiff
using DelimitedFiles
np = pyimport("numpy")

"""
Force load on one edge of the plate,
The edge is [0,L], which is discretized to ne elements, with porder-th order polynomials
type: a string of load type, which can be "constant", or "Gaussian"
args: an array,
      for Constant load, it has p1 and p2, in tangential and normal direction
      for "Gaussian" load, it has p, x0, and σ, the force is in the normal direction

"""
function ComputeLoad(L, ne, porder, ngp, type,  args)
    @assert(ngp <= 4)
    dx = L/ne

    xx = np.linspace(0, L, ne*porder + 1)

    Ft, Fn = zeros(Float64, ne*porder + 1), zeros(Float64, ne*porder + 1)

    elem_coords = zeros(Float64, porder + 1, 2)

    # construct pressure load function
    if type == "Constant"
        pt, pn = args
        f_px, f_py = x->pt, x->pn
    elseif type == "Gaussian"
        pn, x0, σ = args
        f_px, f_py = x->0, x-> pn * 1.0/(sqrt(2*pi*σ^2)) * exp.(-0.5*(x .- x0).^2/σ^2)
    else
        error("Force load type is not recogonized ", type)
    end

    # assemble the force

    @show ne
    for e = 1:ne
        if porder == 1
            loc_id = [e, e+1]
        elseif porder == 2
            loc_id = [2*(e-1)+1, 2*(e-1)+3, 2*(e-1)+2]
        else
            error("porder error porder ==", porder)
        end

        loc_xx = xx[loc_id]         #x-coordinates of element nodes
        elem_coords[:,1] .= loc_xx

        #return list
        weights, hs = get1DElemShapeData( elem_coords, ngp)

        for igp = 1:ngp
            gp_xx = loc_xx' * hs[igp]   #x-coordinates of element Gaussian points
            Ft[loc_id] += f_px(gp_xx) * hs[igp] * weights[igp]
            Fn[loc_id] += f_py(gp_xx) * hs[igp] * weights[igp]

        end
    end

    return Ft, Fn
end

function BoundaryCondition(tid, nx, ny, porder=2, Lx = 1.0, Ly = 0.5; force_scale=5.0)
    nnodes, neles = (nx*porder + 1)*(ny*porder + 1), nx*ny

    x = np.linspace(0.0, Lx, nx*porder + 1)
    y = np.linspace(0.0, Ly, ny*porder + 1)

    X, Y = np.meshgrid(x, y)
    nodes = zeros(nnodes,2)
    nodes[:,1], nodes[:,2] = X'[:], Y'[:]
    ndofs = 2

    EBC, g = zeros(Int64, nnodes, ndofs), zeros(nnodes, ndofs)
    FBC, fext = zeros(Int64, nnodes, ndofs), zeros(nnodes, ndofs)
    ft = nothing
    gt = nothing

    #todo only change 203

    # setting EBC:
    # tid = 1XX   |==|-->
    # tid = 2XX
    if div(tid,100)==1
        EBC[collect(1:nx*porder+1), :] .= -1 # fix bottom
        FBC[collect((nx*porder+1)*ny*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -2 # force on the top
    elseif div(tid,100)==2
        EBC[collect(1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -1 # fix left
        FBC[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -2 # force on the right
    elseif div(tid,100)==3
        EBC[collect(1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -1 # fix left
        FBC[collect(2:nx*porder+1), :] .= -2 # force on the bottom
    elseif div(tid,100)==4
        EBC[collect(1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -1 # fix left
        EBC[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), :] .= -2 #pull right
    end

    # Uniform conservative pressure load
    # F1 = 2e2 /force_scale #gcm/ms^2 compress/pull
    # F2 = 2e1 /force_scale #gcm/ms^2 bend

    P1 = 8 /force_scale #gcm/ms^2 compress/pull
    P2 = 0.8 /force_scale #gcm/ms^2 bend
    P3 = 3/force_scale
    @show P1, P2
    ngp = 3
    #Bending or Pulling
    #
    if tid == 200
        F1, F2 = ComputeLoad(Ly, ny, porder, ngp, "Constant",  [P1, 0])
        fext[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), 1] .= F1
        fext[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), 2] .= F2

    elseif tid == 400
        g[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), 1] .= 0.00508e-3 *10#m/ms
        g[collect(nx*porder+1:nx*porder+1:(nx*porder+1)*(ny*porder+1)), 2] .= 0
    else
        error("tid = $tid is not understood")
    end

    dof_to_active = findall(FBC[:].==-2)
    if length(dof_to_active) != 0
        ft = t->fext[:][dof_to_active] #*sin(π*t/(T))
    end

    dof_to_active = findall(EBC[:].==-2)
    if length(dof_to_active) != 0
        gt = t->[g[:][dof_to_active]*t, g[:][dof_to_active], g[:][dof_to_active]*0.0]
    end

    npoints = (nx+1)*(ny+1)
    node_to_point = zeros(Int64, nnodes) .- 1
    pointid = 1
    for j = 1:2:(ny*porder + 1)
        for i = 1:2:(nx*porder + 1)
            node_to_point[(j-1)*(nx*porder + 1) + i] = pointid
            pointid += 1
        end
    end

    return nodes, EBC, g, gt, FBC, fext, ft, npoints, node_to_point
end


mutable struct Friction{T} <: Contact{T}
    Nx::Adjoint{T,SVector{3,T}}
    D::SMatrix{2,6,T,12}
    cf::T
    offset::SVector{6,T}
    b::SVector{2,T}

    function Friction(body::Body{T}, normal::AbstractVector, cf::Real; offset::AbstractVector = zeros(3)) where T
        @assert cf>0

        # Derived from plane equation a*v1 + b*v2 + distance*v3 = p - offset
        V1, V2, V3 = orthogonalcols(normal) # gives two plane vectors and the original normal axis
        A = [V1 V2 V3]
        Ainv = inv(A)
        ainv3 = Ainv[3,SA[1; 2; 3]]
        Nx = ainv3'
        D = [[V1 V2];szeros(3,2)]'
        offset = [offset;0.0;0.0;0.0]

        new{T}(Nx, D, cf, offset, zeros(2)), body.id
    end
end

function Base.show(io::IO, mime::MIME{Symbol("text/plain")}, constraint::Friction{T}) where {T}
    summary(io, constraint)
    println(io,"")
    println(io, " Nx:     "*string(constraint.Nx))
    println(io, " D:      "*string(constraint.D))
    println(io, " cf:     "*string(constraint.cf))
    println(io, " offset: "*string(constraint.offset))
end

@inline additionalforce(friction::Friction) = friction.D'*friction.b
@inline function calcFrictionForce!(mechanism, ineqc, friction::Friction, i, body::Body)
    calcFrictionForce!(mechanism, friction, body, ineqc.γsol[2][i])
    return
end


using ArnoldiMethod
import ContinuumArrays.QuasiArrays: AbstractQuasiMatrix
using LazyArrays

using ClassicalOrthogonalPolynomials

function orb_label((n,ℓ))
    ℓs = "spdfg"[ℓ+1]
    "$n$(ℓs)"
end

"""
    hydredwfn(n,ℓ[,Z=1])

Returns a function of `r` (in atomic units, i.e. Bohr radii) that
evaluates the _reduced_ hydrogenic radial orbital corresponding to the
quantum numbers `n` and `ℓ`, and for a nuclear charge of `Z`.
"""
function hydredwfn(n,ℓ,Z=1)
    n > 0 || throw(ArgumentError("Invalid n = $n"))
    0 ≤ ℓ && ℓ<n || throw(ArgumentError("Invalid ℓ = $ℓ"))

    q = 2Z/n
    N = √(q^3*factorial(n-ℓ-1)/(2n*factorial(n+ℓ)))

    r -> N*r*(q*r)^ℓ*laguerrel(n-ℓ-1,2ℓ+1,q*r)*exp(-q*r/2)
end

function get_orbitals(R::B, ℓ::Int, Z, nev::Int,
                      mode::Symbol) where {T,B<:AbstractQuasiMatrix{T}}
    if mode == :arnoldi
        println("Finding eigenstates for Z = $Z, ℓ = $ℓ")
        Tℓ = CoulombIntegrals.get_double_laplacian(R,ℓ,T)
        Tℓ ./= -2
        r = axes(R,1)
        H = Tℓ + R' * QuasiDiagonal(-Z ./ r) * R

        schur,history = partialschur(H, nev=nev, tol=sqrt(eps()), which=SR())
        println(history)
        println(diag(schur.R))

        [(n+ℓ,ℓ) => normalize!(applied(*, R, (schur.Q[:,n]*sign(schur.Q[1,n])))) for n = 1:nev]
    elseif mode == :symbolic
        r = axes(R,1)
        [(n,ℓ) => applied(*, R, R \ hydredwfn(n,ℓ,Z).(r))
         for n ∈ ℓ .+ (1:nev)]
    else
        throw(ArgumentError("Unknown mode $(mode)"))
    end
end


export structural, local_level, linear_trend, regression, tvm

"""
    structural(y::VecOrMat{Typ}, s::Int; X::VecOrMat{Typ} = Matrix{Typ}(undef, 0, 0)) where Typ <: Real

Build state-space system for a given structural model with observations `y`, seasonality `s`, and, optionally, exogenous variables `X`.

If `y` is provided as an `Array{Typ, 1}` it will be converted to an `Array{Typ, 2}` inside the `StateSpaceModel`. The same will happen to X,
if an `Array{Typ, 1}` it will be converted to an `Array{Typ, 2}` inside the `StateSpaceModel`.
"""
function structural(y::VecOrMat{Typ}, s::Int; X::VecOrMat{Typ} = Matrix{Typ}(undef, 0, 0)) where Typ <: Real

    # Number of observations and endogenous variables
    y = y[:, :]
    n, p = size(y)

    # Number of observations and exogenous variables
    X = X[:, :]
    n_exp, p_exp = size(X)

    if p_exp > 0 && n_exp < n
        error("Number of observations in X must be greater than or equal to the number of observations in y.")
    end

    m = (1 + s + p_exp)*p
    r = 3*p

    # Observation equation
    N = max(n, n_exp)
    if p_exp > 0 # exogenous variables: Z is time-variant
        Z = Array{Typ, 3}(undef, p, m, N)
        for t = 1:N
            Z[:, :, t] = kron(
                Matrix{Typ}(I, p, p),
                [
                    X[t, :]' 1 0 1 zeros(Typ, 1, s - 2)
                ]
                )
        end
    else # no exogenous variables: Z is time-invariant
        Z = kron(
            Matrix{Typ}(I, p, p),
            [
                1 0 1 zeros(Typ, 1, s - 2)
            ]
            )
    end

    # State equation
    if p_exp > 0
        T0 = [Matrix{Typ}(I, p_exp, p_exp) zeros(Typ, p_exp, 1 + s)]
        T = kron(
            Matrix{Typ}(I, p, p),
            [
                T0;
                zeros(Typ, 1, p_exp) 1 1 zeros(Typ, 1, s - 1);
                zeros(Typ, 1, p_exp) 0 1 zeros(Typ, 1, s - 1);
                zeros(Typ, 1, p_exp) 0 0 -ones(Typ, 1, s - 1);
                zeros(Typ, s - 2, p_exp) zeros(Typ, s - 2, 2) Matrix{Typ}(I, s - 2, s - 2) zeros(Typ, s - 2)
            ]
            )
    else
        T = kron(
            Matrix{Typ}(I, p, p),
            [
                1 1 zeros(Typ, 1, s - 1);
                0 1 zeros(Typ, 1, s - 1);
                0 0 -ones(Typ, 1, s - 1);
                zeros(Typ, s - 2, 2) Matrix{Typ}(I, s - 2, s - 2) zeros(Typ, s - 2)
            ]
            )
        end

    R = kron(
        Matrix{Typ}(I, p, p),
        [
            zeros(Typ, p_exp, 3);
            Matrix{Typ}(I, 3, 3);
            zeros(Typ, s - 2, 3)
        ]
        )

    return StateSpaceModel(y, Z, T, R)

end

"""
    local_level(y::VecOrMat{Typ}) where Typ <: Real

Build state-space system for a local level model with observations y.

If `y` is provided as an `Array{Typ, 1}` it will be converted to an `Array{Typ, 2}` inside the `StateSpaceModel`.
"""
function local_level(y::VecOrMat{Typ}) where Typ <: Real

    # Number of observations and endogenous variables
    y = y[:, :]
    n, p = size(y)

    m = p
    r = p

    # Observation equation
    Z = Matrix{Typ}(I, p, p)

    # State equation
    T = Matrix{Typ}(I, p, p)
    R = Matrix{Typ}(I, p, p)

    return StateSpaceModel(y, Z, T, R)
end

"""
    linear_trend(y::VecOrMat{Typ}) where Typ <: Real

Build state-space system for a linear trend model with observations y.

If `y` is provided as an `Array{Typ, 1}` it will be converted to an `Array{Typ, 2}` inside the `StateSpaceModel`.
"""
function linear_trend(y::VecOrMat{Typ}) where Typ <: Real

    # Number of observations and endogenous variables
    y = y[:, :]
    n, p = size(y)

    m = 2*p
    r = 2*p

    # Observation equation
    Z = kron(Matrix{Typ}(I, p, p),[1 0])

    # State equation
    T = kron(Matrix{Typ}(I, p, p),[1 1; 0 1])
    R = kron(Matrix{Typ}(I, p, p),[1 0; 0 1])

    return StateSpaceModel(y, Z, T, R)
end

"""
    regression(y::VecOrMat{Typ}, X::VecOrMat{Typ}) where Typ <: Real

Build state-space system for estimating a regression model ``y_t = X_t\\beta_t + \\varepsilon_t``. Once the model is estimated
the user can recover the parameter ``\\hat{\\beta}`` by the querying the smoothed states of the model.
"""
function regression(y::VecOrMat{Typ}, X::VecOrMat{Typ}) where Typ <: Real
    # Certify that they are matrices
    y = ensure_is_matrix(y)
    X = ensure_is_matrix(X)
    # Assert sizes are adequate
    n_y, p_y = size(y)
    n_X, p_X = size(X)
    @assert n_X == n_y

    if p_y > 1
        error("StateSpaceModels currently supports regression only for univariate cases.")
    end

    # Fill Z
    Z = Array{Typ, 3}(undef, p_y, p_X, n_X)
    for t in 1:n_X
        Z[:, :, t] = X[t, :]
    end
    T = Matrix{Typ}(I, p_X, p_X)
    R = zeros(Typ, p_X, 1)
    Q = zeros(Typ, 1, 1)
    d = zeros(Typ, n_y, p_y)
    c = zeros(Typ, n_X, p_X)
    # H is the only variable to estimate
    H = build_H(p_y, Typ)

    return StateSpaceModel(y, Z, T, R, d, c, H, Q)
end

"""
    tvm(y::VecOrMat{Typ}, X::VecOrMat{Typ}) where Typ <: Real

Build state-space system for estimating a time varying regression model ``y_t = X_t\\beta_t + \\varepsilon_t``. Once the model is estimated
the user can recover the parameter ``\\hat{\\beta}`` by the querying the smoothed states of the model.
"""

function tvm(y::VecOrMat{Typ}, X::VecOrMat{Typ}; X_var = zeros(Typ, size(X, 2))) where Typ <: Real

    #Time varying
    # Certify that they are matrices
    y = ensure_is_matrix(y)
    X = ensure_is_matrix(X)
    # Assert sizes are adequate
    n_y, p_y = size(y)
    n_X, p_X = size(X)

    #We are creating a m-dimensional state-space model
    m = p_X + p_y
    @assert n_X == n_y

    if p_y > 1
        error("StateSpaceModels currently supports regression only for univariate cases.")
    end

    # Fill Z
    Z = Array{Typ, 3}(undef, p_y, m, n_X)
    for t in 1:n_X
        Z[:, :, t] = vcat(Matrix{Typ}(I, p_y, p_y), X[t, :])
    end
    T = Matrix{Typ}(I, m, m)
    R = diagm(vcat([1.0], X_var))

    Q = diagm(vcat([1.0], X_var))
    Q[findall(isequal(1.0), Q)] .= NaN

    d = zeros(Typ, n_y, p_y)
    c = zeros(Typ, n_X, m)

    H = build_H(p_y, Typ)

    return StateSpaceModel(y, Z, T, R, d, c, H, Q)
end


"""
    KNN(k=5,
        metric=Euclidean,
        algorithm=:kdtree,
        leafsize=10,
        reorder=true,
        reduction=:maximum)

Calculate the anomaly score of an instance based on the distance to its k-nearest neighbors.

Parameters
----------
$_k_param

$_knn_shared

    reduction::Symbol
One of `(:maximum, :median, :mean)`. (`reduction=:maximum`) was proposed by [1]. Angiulli et al. [2] proposed sum to
reduce the distances, but mean has been implemented for numerical stability.

Examples
--------
$(_score_unsupervised("KNN"))

References
----------
[1] Ramaswamy, Sridhar; Rastogi, Rajeev; Shim, Kyuseok (2000): Efficient Algorithms for Mining Outliers from Large Data
Sets.

[2] Angiulli, Fabrizio; Pizzuti, Clara (2002): Fast Outlier Detection in High Dimensional Spaces.
"""
MMI.@mlj_model mutable struct KNN <: UnsupervisedDetector
    k::Integer = 5::(_ > 0)
    metric::DI.Metric = DI.Euclidean()
    algorithm::Symbol = :kdtree::(_ in (:kdtree, :balltree))
    leafsize::Integer = 10::(_ ≥ 0)
    reorder::Bool = true
    parallel::Bool = false
    reduction::Symbol = :maximum::(_ in (:maximum, :median, :mean))
end

struct KNNModel <: Model
    tree::NN.NNTree
end

function fit(detector::KNN, X::Data)::Fit
    # create the specified tree
    tree = buildTree(X, detector.metric, detector.algorithm, detector.leafsize, detector.reorder)

    # use tree to calculate distances
    _, dists = knn_others(tree, X, detector.k)

    # reduce distances to outlier score
    scores = _knn(dists, detector.reduction)
    Fit(KNNModel(tree), scores)
end

@score function score(detector::KNN, model::Fit, X::Data)::Result
    if detector.parallel
        idxs, dists = knn_parallel(model.tree, X, detector.k)
        return _knn(dists, detector.reduction)
    else
        idxs, dists = NN.knn(model.tree, X, detector.k)
        return _knn(dists, detector.reduction)
    end
end

@inline function _knn(distances::AbstractVector{<:AbstractVector}, reduction::Symbol)::Scores
    # Helper function to reduce `k` distances to a single distance.
    if reduction == :maximum
        return maximum.(distances)
    elseif reduction == :median
        return median.(distances)
    elseif reduction == :mean
        return mean.(distances)
    end
end


#!/usr/bin/env julia

using Luxor, Test

using Random
Random.seed!(42)

function findanddrawsometangentialcircles(c1center, c1radius, c2center, c2radius)
    for requiredradius in 1:150
        res = circletangent2circles(requiredradius, c1center, c1radius, c2center, c2radius)
        setopacity(0.5)
        if first(res) == 1
            circle(res[2], requiredradius, :stroke)
        elseif first(res) == 2
            circle(res[2], requiredradius, :stroke)
            circle(res[3], requiredradius, :stroke)
        elseif first(res) == 0
            sethue("red")
            circle(c1center, 5, :fill)
            circle(c2center, 5, :fill)
        end
    end
end

function test_circletangents(fname)
    pagewidth, pageheight = 1200, 1400
    Drawing(pagewidth, pageheight, fname)
    origin() # move 0/0 to center
    background("ivory")
    setline(0.5)
    tiles = Tiler(1200, 1400, 5, 5)
    for (pos, n) in tiles
        @layer begin
            translate(pos)
            circle1 = (O + Point(rand(-20:20), rand(-20:20)), rand(20:40))
            circle2 = (O + Point(rand(-20:20), rand(-20:20)), rand(20:40))
            sethue("black")
            setopacity(0.5)
            circle(circle1..., :fill)
            circle(circle2..., :fill)
            setopacity(1.0)
            circle(circle1..., :stroke)
            circle(circle2..., :stroke)
            sethue("purple")
            findanddrawsometangentialcircles(circle1[1], circle1[2], circle2[1], circle2[2])
        end
    end

    @test finish() == true
    println("...finished circle tangent test, saved in $(fname)")
end

fname = "circletangent-test.pdf"

test_circletangents(fname)