#Ax_mul_Bxt

export Ax_mul_Bxt

"""
`Ax_mul_Bxt(A::AbstractOperator,B::AbstractOperator)`

Create an operator `P` such that:

`P == (Ax)*(Bx)'`

# Example: Matrix multiplication

```julia
julia> A,B = randn(4,4),randn(4,4);

julia> P = Ax_mul_Bxt(MatrixOp(A),MatrixOp(B))
▒*▒  ℝ^4 -> ℝ^(4, 4)

julia> x = randn(4);

julia> P*x == (A*x)*(B*x)'
true

```
"""
struct Ax_mul_Bxt{
                   L1 <: AbstractOperator,
                   L2 <: AbstractOperator,
                   C <: AbstractArray,
                   D <: AbstractArray,
                  } <: NonLinearOperator
  A::L1
  B::L2
  bufA::C
  bufB::C
  bufC::C
  bufD::D
  function Ax_mul_Bxt(A::L1, B::L2, bufA::C, bufB::C, bufC::C, bufD::D) where {L1,L2,C,D}
    if ndims(A,1) == 1
      if size(A) != size(B)
        throw(DimensionMismatch("Cannot compose operators"))
      end
    elseif ndims(A,1) == 2 && ndims(B,1) == 2 && size(A,2) == size(B,2)
      if size(A,1)[2] != size(B,1)[2]
        throw(DimensionMismatch("Cannot compose operators"))
      end
    else
      throw(DimensionMismatch("Cannot compose operators"))
    end
    new{L1,L2,C,D}(A,B,bufA,bufB,bufC,bufD)
  end
end

struct Ax_mul_BxtJac{
                      L1 <: AbstractOperator,
                      L2 <: AbstractOperator,
                      C <: AbstractArray,
                      D <: AbstractArray,
                     } <: LinearOperator
  A::L1
  B::L2
  bufA::C
  bufB::C
  bufC::C
  bufD::D
end

# Constructors
function Ax_mul_Bxt(A::AbstractOperator,B::AbstractOperator)
  s,t = size(A,1), codomainType(A)
  bufA = eltype(s) <: Int ? zeros(t,s) : ArrayPartition(zeros.(t,s)...)
  bufC = eltype(s) <: Int ? zeros(t,s) : ArrayPartition(zeros.(t,s)...)
  s,t = size(B,1), codomainType(B)
  bufB = eltype(s) <: Int ? zeros(t,s) : ArrayPartition(zeros.(t,s)...)
  s,t = size(A,2), domainType(A)
  bufD = eltype(s) <: Int ? zeros(t,s) : ArrayPartition(zeros.(t,s)...)
  Ax_mul_Bxt(A,B,bufA,bufB,bufC,bufD)
end

# Jacobian
function Jacobian(P::Ax_mul_Bxt{L1,L2,C,D}, x::AbstractArray) where {L1,L2,C,D}
  JA, JB = Jacobian(P.A, x), Jacobian(P.B, x)
  Ax_mul_BxtJac{typeof(JA),typeof(JB),C,D}(JA,JB,P.bufA,P.bufB,P.bufC,P.bufD)
end

# Mappings
function mul!(y, P::Ax_mul_Bxt{L1,L2,C,D}, b) where {L1,L2,C,D}
  mul!(P.bufA,P.A,b)
  mul!(P.bufB,P.B,b)
  mul!(y,P.bufA, P.bufB')
end

function mul!(y, J::AdjointOperator{Ax_mul_BxtJac{L1,L2,C,D}}, b) where {L1,L2,C,D}
  #y .= J.A.A'*(b*(J.A.bufB)) + J.A.B'*(b'*(J.A.bufA))
  mul!(J.A.bufC, b, J.A.bufB)
  mul!(y, J.A.A', J.A.bufC)
  mul!(J.A.bufB, b', J.A.bufA)
  mul!(J.A.bufD, J.A.B', J.A.bufB)
  y .+= J.A.bufD
  return y
end

size(P::Union{Ax_mul_Bxt,Ax_mul_BxtJac}) = ((size(P.A,1)[1],size(P.B,1)[1]),size(P.A,2))

fun_name(L::Union{Ax_mul_Bxt,Ax_mul_BxtJac}) = fun_name(L.A)*"*"*fun_name(L.B)

domainType(L::Union{Ax_mul_Bxt,Ax_mul_BxtJac})   = domainType(L.A)
codomainType(L::Union{Ax_mul_Bxt,Ax_mul_BxtJac}) = codomainType(L.A)

# utils
function permute(P::Ax_mul_Bxt{L1,L2,C,D},
                 p::AbstractVector{Int}) where {L1,L2,C,D <:ArrayPartition}
  Ax_mul_Bxt(permute(P.A,p),permute(P.B,p),P.bufA,P.bufB,P.bufC,ArrayPartition(P.bufD.x[p]) )
end

remove_displacement(P::Ax_mul_Bxt) =
Ax_mul_Bxt(remove_displacement(P.A), remove_displacement(P.B), P.bufA, P.bufB, P.bufC, P.bufD)


# Return flat index into arrray of (n, m) pairs.
# Assumes array is ordered as
#   [
#     (n, m)
#     for n in 0:n_max
#     for m in -n:n
# ]
@inline nm_index(n, m) = m + n * (n + 1) + 1

"""
    H_recursion_coefficients(ℓₘₐₓ, T)

Pre-compute constants used in Wigner H recursion.

"""
function H_recursion_coefficients(ℓₘₐₓ, ::Type{T}) where {T<:Real}
    aₙᵐ = T[√T((n+1+m)*(n+1-m)) for n in 0:ℓₘₐₓ+1 for m in 0:n]
    bₙᵐ = T[√T((n-m-1)*(n-m)) for n in 0:ℓₘₐₓ+1 for m in -n:n]
    dₙᵐ = T[(m<0 ? -1 : 1) * (√T((n-m)*(n+m+1))) for n in 0:ℓₘₐₓ+1 for m in -n:n]
    (aₙᵐ, bₙᵐ, dₙᵐ)
end

# function H(ℓₘₐₓ, expiβ::Complex{T}) where {T<:Real}
#     H(0, ℓₘₐₓ, ℓₘₐₓ, expiβ)
# end

"""
    H!(H, expiβ, ℓₘₐₓ, m′ₘₐₓ, H_rec_coeffs)
    H!(H, expiβ, ℓₘₐₓ, m′ₘₐₓ, H_rec_coeffs, Hindex)

Compute the ``H`` matrix defined by Gumerov and Duraiswami.

This computation forms the basis for computing Wigner's ``d`` and ``𝔇``
matrices via [`d!`](@ref) and [`D!`](@ref), the spin-weighted spherical
harmonics via [`Y!`](@ref), and for transforming from values of spin-weighted
spherical functions evaluated on a grid to the corresponding mode weights via
[`map2salm`](@ref).

Due to symmetries, we only need to compute ~1/4 of the elements of this matrix,
so only those elements with ``m ≥ |m′|`` are computed.  The relevant indices of
the `H` vector are computed based on the `Hindex` function — which defaults to
`WignerHindex`, but could reasonably be `WignerDindex` if the input `H` vector
contains all valid indices.  However, it is assumed that the storage scheme
used for `H` is such that the successive ``m`` values are located in successive
elements.

If ``m′ₘₐₓ < ℓₘₐₓ``, we don't even need 1/4 of the elements, and only values
with ``|m′| ≤ m′ₘₐₓ`` will be computed.  This is particularly useful for
computing spin-weighted spherical harmonics.

Note that the recursion coefficients `H_rec_coeffs` should be the quantity
returned by [`H_recursion_coefficients`](@ref).

"""
function H!(
    H::AbstractVector, expiβ::Complex{T}, ℓₘₐₓ, m′ₘₐₓ, (aₙᵐ,bₙᵐ,dₙᵐ), Hindex=WignerHindex
) where {T<:Real}
    m′ₘₐₓ = abs(m′ₘₐₓ)
    @assert m′ₘₐₓ ≤ ℓₘₐₓ
    @assert size(H, 1) ≥ Hindex(ℓₘₐₓ, m′ₘₐₓ, ℓₘₐₓ, m′ₘₐₓ)

    invsqrt2 = inv(√T(2))
    cosβ = expiβ.re
    sinβ = expiβ.im
    cosβ₊ = (1+cosβ)/2  # = cos²(β/2)
    cosβ₋ = (1-cosβ)/2  # = sin²(β/2)

    # Note: In step 3, we set the H^{1, m}_{ℓₘₐₓ} terms from H^{0, m+1}_{ℓₘₐₓ+1} data,
    # which of course should not appear in the final result, because it involves ℓ>ℓₘₐₓ.
    # This requires storing those extra terms somewhere temporarily during the execution
    # of this function.  Fortunately, this is only required if m′ₘₐₓ>0, which also implies
    # that there are at least ℓₘₐₓ slots in the input Hwedge array that are not needed until
    # step 5.  We use those slots for most of the temporary data.  However, there are still
    # two slots needed, which we provide as individual variables, representing the last and
    # second-to-last elements H^{0, ℓₘₐₓ+1}_{ℓₘₐₓ+1} and H^{0, ℓₘₐₓ}_{ℓₘₐₓ+1}:
    HΩ, HΨ = zero(eltype(H)), zero(eltype(H))

    @inbounds begin
        # Step 1: If n=0 set H_{0}^{0,0}=1
        H[1] = 1

        if ℓₘₐₓ > 0

            # Step 2: Compute H^{0,m}_{n}(β) for m=0,...,n and n=1,...,ℓₘₐₓ,ℓₘₐₓ+1?
            # Note that in this step only, we use notation derived from Xing et al., so that
            # the coefficients are denoted b̄ₙ, c̄ₙₘ, d̄ₙₘ, ēₙₘ.  Below, we will use notation
            # from Gumerov and Duraiswami, who denote their different coefficients aₙᵐ, etc.
            nₘₐₓstep2 = m′ₘₐₓ>0 ? ℓₘₐₓ+1 : ℓₘₐₓ
            # n = 1
            n0n_index = Hindex(1, 0, 0, m′ₘₐₓ)
            H[n0n_index] = cosβ
            H[n0n_index+1] = invsqrt2 * sinβ
            # n = 2, ..., ℓₘₐₓ, ℓₘₐₓ+1?
            for n in 2:nₘₐₓstep2
                superjacent = n > ℓₘₐₓ
                if superjacent
                    n00_index = Hindex(n-1, -1, 1, m′ₘₐₓ)
                else
                    n00_index = Hindex(n, 0, 0, m′ₘₐₓ)
                end
                nm100_index = Hindex(n-1, 0, 0, m′ₘₐₓ)
                inv2n = inv(T(2n))
                b̄ₙ = √(T(n-1)/n)

                # H^{0,0}_{n} = cosβ H^{0,0}_{n-1} - b̄ₙ sinβ H^{0,1}_{n-1}
                H[n00_index] = cosβ * H[nm100_index] - b̄ₙ  * sinβ * H[nm100_index+1]

                # H^{0,m}_{n} = c̄ₙₘ cosβ H^{0,m}_{n-1} - sinβ [ d̄ₙₘ H^{0,m+1}_{n-1} - ēₙₘ H^{0,m-1}_{n-1} ]
                for m in 1:n-2
                    c̄ₙₘ = 2inv2n * √T((n+m)*(n-m))
                    d̄ₙₘ = inv2n * √T((n-m)*(n-m-1))
                    ēₙₘ = inv2n * √T((n+m)*(n+m-1))
                    H[n00_index+m] = (
                        c̄ₙₘ * cosβ * H[nm100_index+m]
                        - sinβ * (
                            d̄ₙₘ * H[nm100_index+m+1]
                            - ēₙₘ * H[nm100_index+m-1]
                        )
                    )
                end
                let m = n-1
                    c̄ₙₘ = 2inv2n * √T(2n-1)
                    d̄ₙₘ = 0
                    ēₙₘ = inv2n * √T((2n-1)*(2n-2))
                    if superjacent
                        HΨ = (
                            c̄ₙₘ * cosβ * H[nm100_index+m]
                            + sinβ * ēₙₘ * H[nm100_index+m-1]
                        )
                    else
                        H[n00_index+m] = (
                            c̄ₙₘ * cosβ * H[nm100_index+m]
                            + sinβ * ēₙₘ * H[nm100_index+m-1]
                        )
                    end
                end
                let m = n
                    c̄ₙₘ = 0
                    d̄ₙₘ = 0
                    ēₙₘ = inv2n * √T(2n*(2n-1))
                    if superjacent
                        HΩ = sinβ * ēₙₘ * H[nm100_index+m-1]
                    else
                        H[n00_index+m] = sinβ * ēₙₘ * H[nm100_index+m-1]
                    end
                end

            end

            @fastmath if m′ₘₐₓ > 0

                # Step 3: Compute H^{1,m}_{n}(β) for m=1,...,n
                for n in 1:ℓₘₐₓ
                    # m = 1, ..., n
                    i1 = Hindex(n, 1, 1, m′ₘₐₓ)
                    if n+1 <= ℓₘₐₓ
                        i2 = Hindex(n+1, 0, 0, m′ₘₐₓ)
                        nₘₐₓstep3 = n-1
                    else  # n+1 == ℓₘₐₓ+1  &&  m′ₘₐₓ > 0
                        i2 = Hindex(n, -1, 1, m′ₘₐₓ)
                        nₘₐₓstep3 = n-3
                    end
                    i3 = nm_index(n+1, 0)
                    i4 = (n * (n + 1)) ÷ 2 + 2
                    inverse_b5 = inv(bₙᵐ[i3])
                    for i in 0:nₘₐₓstep3
                        b6 = bₙᵐ[-i+i3-2]
                        b7 = bₙᵐ[i+i3]
                        a8 = aₙᵐ[i+i4]
                        H[i+i1] = -inverse_b5 * (
                            b6 * cosβ₋ * H[i+i2+2]
                            + b7 * cosβ₊ * H[i+i2]
                            + a8 * sinβ * H[i+i2+1]
                        )
                    end
                    if n == ℓₘₐₓ  # &&  m′ₘₐₓ > 0
                        if n>1
                            let i = n-2
                                b6 = bₙᵐ[-i+i3-2]
                                b7 = bₙᵐ[i+i3]
                                a8 = aₙᵐ[i+i4]
                                H[i+i1] = -inverse_b5 * (
                                    b6 * cosβ₋ * HΨ
                                    + b7 * cosβ₊ * H[i+i2]
                                    + a8 * sinβ * H[i+i2+1]
                                )
                            end
                        end
                        let i = n-1
                            b6 = bₙᵐ[-i+i3-2]
                            b7 = bₙᵐ[i+i3]
                            a8 = aₙᵐ[i+i4]
                            H[i+i1] = -inverse_b5 * (
                                b6 * cosβ₋ * HΩ
                                + b7 * cosβ₊ * H[i+i2]
                                + a8 * sinβ * HΨ
                            )
                        end
                    end
                end  # Step 3

                # Step 4: Compute H^{m'+1, m}_{n}(β) for m'=1,...,n−1, m=m',...,n
                for n in 2:ℓₘₐₓ
                    for mp in 1:min(n, m′ₘₐₓ)-1
                        # m = m', ..., n-1
                        i1 = Hindex(n, mp+1, mp+1, m′ₘₐₓ) - 1
                        i2 = Hindex(n, mp-1, mp, m′ₘₐₓ)
                        i3 = Hindex(n, mp, mp, m′ₘₐₓ) - 1
                        i4 = Hindex(n, mp, mp+1, m′ₘₐₓ)
                        i5 = nm_index(n, mp)
                        i6 = nm_index(n, mp-1)
                        inverse_d5 = inv(dₙᵐ[i5])
                        d6 = dₙᵐ[i6]

                        d7 = d6
                        d8 = dₙᵐ[i5]
                        for i in 1:n-mp-1
                            d7 = d8
                            d8 = dₙᵐ[i+i5]
                            H[i+i1] = inverse_d5 * (
                                d6 * H[i+i2]
                                - d7 * H[i+i3]
                                + d8 * H[i+i4]
                            )
                        end

                        # m = n
                        let i=n-mp
                            d7 = d8
                            H[i+i1] = inverse_d5 * (
                                d6 * H[i+i2]
                                - d7 * H[i+i3]
                            )
                        end
                    end
                end  # Step 4

                # Step 5: Compute H^{m'−1, m}_{n}(β) for m'=0,...,−n+1, m=−m',...,n
                for n in 0:ℓₘₐₓ
                    for mp in 0:-1:1-min(n, m′ₘₐₓ)
                        # m = -m', ..., n-1
                        i1 = Hindex(n, mp-1, -mp+1, m′ₘₐₓ) - 1
                        i2 = Hindex(n, mp+1, -mp+1, m′ₘₐₓ) - 1
                        i3 = Hindex(n, mp, -mp, m′ₘₐₓ) - 1
                        i4 = Hindex(n, mp, -mp+1, m′ₘₐₓ)
                        i5 = nm_index(n, mp-1)
                        i6 = nm_index(n, mp)
                        i7 = nm_index(n, -mp-1)
                        i8 = nm_index(n, -mp)
                        inverse_d5 = inv(dₙᵐ[i5])
                        d6 = dₙᵐ[i6]
                        d7 = dₙᵐ[i7]
                        d8 = dₙᵐ[i8]
                        for i in 1:n+mp-1
                            d7 = d8
                            d8 = dₙᵐ[i+i8]
                            H[i+i1] = inverse_d5 * (
                                d6 * H[i+i2]
                                + d7 * H[i+i3]
                                - d8 * H[i+i4]
                            )
                        end
                        # m = n
                        let i=n+mp
                            d7 = d8
                            H[i+i1] = inverse_d5 * (
                                d6 * H[i+i2]
                                + d7 * H[i+i3]
                            )
                        end
                    end
                end  # Step 5

            end  # if m′ₘₐₓ > 0
        end  # if ℓₘₐₓ > 0
    end  # @inbounds
    H
end


struct AsymptoticBlockSKCETest{K<:Kernel,E,S,Z} <: HypothesisTests.ZTest
    """Kernel."""
    kernel::K
    """Number of observations per block."""
    blocksize::Int
    """Number of blocks of observations."""
    nblocks::Int
    """Calibration error estimate (average of evaluations of blocks of observations)."""
    estimate::E
    """Standard error of evaluations of blocks of observations."""
    stderr::S
    """z-statistic."""
    z::Z
end

function AsymptoticBlockSKCETest(
    kernel::Kernel, blocksize::Int, predictions::AbstractVector, targets::AbstractVector
)
    # obtain number of samples
    nsamples = length(predictions)
    nsamples ≥ blocksize || error("there must be at least ", blocksize, " samples")

    # compute number of blocks
    nblocks = nsamples ÷ blocksize
    nblocks ≥ 2 || error("there must be at least 2 blocks")

    # evaluate U-statistic of the first block
    istart = 1
    iend = blocksize
    x = UnbiasedSKCE(kernel)(view(predictions, istart:iend), view(targets, istart:iend))

    # initialize the estimate and the sum of squares
    estimate = x / 1
    S = zero(estimate)^2

    # for all other blocks
    for b in 2:nblocks
        # evaluate U-statistic
        istart += blocksize
        iend += blocksize
        x = UnbiasedSKCE(kernel)(view(predictions, istart:iend), view(targets, istart:iend))

        # update the estimate
        Δestimate = x - estimate
        estimate += Δestimate / b
        S += Δestimate * (x - estimate)
    end

    # compute standard error and z-statistic
    stderr = sqrt(S / (nblocks * (nblocks - 1)))
    z = estimate / stderr

    return AsymptoticBlockSKCETest(kernel, blocksize, nblocks, estimate, stderr, z)
end

# HypothesisTests interface

HypothesisTests.default_tail(::AsymptoticBlockSKCETest) = :right

## have to specify and check keyword arguments in `pvalue` and `confint` to
## force `tail = :right` due to the default implementation in HypothesisTests

function HypothesisTests.pvalue(test::AsymptoticBlockSKCETest; tail=:right)
    if tail === :right
        StatsFuns.normccdf(test.z)
    else
        throw(ArgumentError("tail=$(tail) is invalid"))
    end
end

# confidence interval by inversion
function HypothesisTests.confint(test::AsymptoticBlockSKCETest; level=0.95, tail=:right)
    HypothesisTests.check_level(level)

    if tail === :right
        q = StatsFuns.norminvcdf(level)
        lowerbound = test.estimate - q * test.stderr
        (max(zero(lowerbound), lowerbound), oftype(lowerbound, Inf))
    else
        throw(ArgumentError("tail = $(tail) is invalid"))
    end
end

HypothesisTests.testname(test::AsymptoticBlockSKCETest) = "Asymptotic block SKCE test"

# parameter of interest: name, value under H0, point estimate
function HypothesisTests.population_param_of_interest(test::AsymptoticBlockSKCETest)
    return "SKCE", zero(test.estimate), test.estimate
end

function HypothesisTests.show_params(io::IO, test::AsymptoticBlockSKCETest, ident="")
    println(io, ident, "number of observations per block: ", test.blocksize)
    println(io, ident, "number of blocks of observations: ", test.nblocks)
    println(io, ident, "z-statistic: ", test.z)
    return println(
        io, ident, "standard error of evaluations of pairs of observations: ", test.stderr
    )
end


using Zygote
using Tullio

export Tikhonov, GR, TV
export generate_spatial_grad_square, generate_GR, generate_TV

 # General hint
 # for the creation of the regularizers we are using meta programming
 # because the fastest way for automatic differentiation and Zygote
 # is Tullio.jl at the moment.
 # Our metaprogramming code was initially based on
 # https://github.com/mcabbott/Tullio.jl/issues/11
 #

"""
    generate_indices(num_dims, d, ind1, ind2)
Generates a list of symbols which can be used to generate Tullio expressions
via metaprogramming.
`num_dims` is the total number of dimensions.
`d` is the dimension where there is a offset in the index.
`ind1` and `ind2` are the offsets each.
 # Examples
 ```julia-repl
julia> a, b = generate_indices(5, 2, 1, 1)
(Any[:i1, :(i2 + 1), :i3, :i4, :i5], Any[:i1, :(i2 + 1), :i3, :i4, :i5])
```
"""
function generate_indices(num_dims, d, ind1, ind2)
    # create initial symbol
    ind = :i
    # create the array of symbols for each dimension
    inds1 = map(1:num_dims) do di
        # map over numbers and append the number of the position to symbol i
        i = Symbol(ind, di)
        # at the dimension where we want to do the step, add $ind1
        di == d ? :($i + $ind1) : i
    end
    inds2 = map(1:num_dims) do di
        i = Symbol(ind, di)
        # here we do the step but in the other dimension. ind2 should be
        # (-1) * ind1 or negative
        di == d ? :($i + $ind2) : i
    end
    return inds1, inds2
end

"""
    generate_laplace(num_dims, sum_dims_arr, weights)
Generate the Tullio statement for computing the abs2 of Laplacian.
`num_dims` is the dimension of the array.
`sum_dims_arr` is a array indicating over which dimensions we must sum over.
`weights` is a array of a weight for the different dimension.
"""
function generate_laplace(num_dims, sum_dims_arr, weights; debug=false)
    # create out list for the final expression
    # add accumulates the different add expressions
    out, add = [], []
    # loop over all dimensions which we want to sum. for each dimension must
    # be a weight provided
    for (d, w) in zip(sum_dims_arr, weights)
        # get the two lists of indices
        inds1, inds2 = generate_indices(num_dims, d, 1, -1)
        # critical part where we actually add the two expressions for
        # the steps in the dimension to the add array
        push!(add, :($w * arr[$(inds1...)] + $w * arr[$(inds2...)]))
    end
    # for laplace we need one final negative term at the position itself
    inds = map(1:num_dims) do di
        i = Symbol(:i, di)
    end
    # subtract this final term
    pre_factor = 2 ^ num_dims * sum(weights)
    push!(add, :(-$:($pre_factor * arr[$(inds...)])))
    # create final expressions by adding all elements of the add list
    if debug
        push!(out, :(res = abs2(+$(add...))))
    else
        push!(out, :(@tullio threads=false res = abs2(+$(add...))))
    end
    return out
end

"""
    create_Ndim_regularizer(expr, num_dims, sum_dims_arr, weights, ind1, ind2)
    A helper function to create a N-dimensional regularizer. In principle
    the same as `generate_laplace` but more general
    `expr` needs to be a function which takes `inds1`, `inds2` and a weight `w`-
    `num_dims` is the total amount of dimensions
    `sum_dims_arr` is a array indicating over which dimensions we must sum over.
    `weights` is a array of a weight for the different dimension.
    ``
"""
function create_Ndim_regularizer(expr, num_dims, sum_dims_arr, weights,
                                 ind1, ind2)
    out, add = [], []
    for (d, w) in zip(sum_dims_arr, weights)
        inds1, inds2 = generate_indices(num_dims, d, ind1, ind2)
        push!(add, expr(inds1, inds2, w))
    end
    push!(out, :(@tullio threads=false res = +($(add...))))
    return out
end

"""
    generate_spatial_grad_square(num_dims, sum_dims_arr, weights)
Generate the Tullio statement for calculating the squared spatial gradient
over n dimensions.
`num_dims` is the dimension of the array.
`sum_dims_arr` is a array indicating over which dimensions we must sum over.
`weights` is a array of a weight for the different dimension.
`ind1` and `ind2` are the offsets for the difference.
"""
function generate_spatial_grad_square(num_dims, sum_dims_arr, weights)
    expr(inds1, inds2, w) = :($w * abs2(arr[$(inds1...)] - arr[$(inds2...)]))
    @eval x = arr -> ($(create_Ndim_regularizer(expr, num_dims, sum_dims_arr,
                        weights, 1, -1)...))
    return x
end

"""
    Tikhonov(; <keyword arguments>)
This function returns a function to calculate the Tikhonov regularizer
of a n-dimensional array.
# Arguments
- `num_dims=2`:
- `sum_dims=[1, 2]`: A array containing the dimensions we want to sum over
- `weights=nothing`: A array containing weights to weight the contribution of
    different dimensions. If `weights=nothing` all dimensions are weighted equally.
- `step=1`: A integer indicating the step width for the array indexing
- `mode="laplace"`: Either `"laplace"`, `"spatial_grad_square"`, `"identity"` accounting for different
    modes of the Tikhonov regularizer. Default is `"laplace"`.
# Examples
To create a regularizer for a 3D dataset where the third dimension
has different contribution.
```julia-repl
julia> reg = Tikhonov(num_dims=2, sum_dims=[1, 2], weights=[1, 1], mode="identity");

julia> reg([1 2 3; 4 5 6; 7 8 9])
285
```
"""
function Tikhonov(;num_dims=2, sum_dims=[1, 2], weights=[1, 1], step=1, mode="laplace")
    if weights == nothing
        weights = ones(num_dims)
    end
    if mode == "laplace"
        Γ = @eval arr -> ($(generate_laplace(num_dims, sum_dims, weights)...))
    elseif mode == "spatial_grad_square"
        expr(inds1, inds2, w) = :($w * abs2(arr[$(inds1...)] - arr[$(inds2...)]))
        Γ = @eval arr -> ($(create_Ndim_regularizer(expr, num_dims, sum_dims,
                            weights, step, (-1) * step)...))
    elseif mode == "identity"
        Γ = arr -> sum(abs2.(arr))
    else
        throw(ArgumentError("The provided mode is not valid."))
    end

    return Γ
end

"""
    generate_GR(num_dims, sum_dims_arr, weights, ind1, ind2, ϵ)
Generate the Tullio statement for computing the Good's roughness.
`num_dims` is the dimension of the array. `sum_dims_arr` is a array
indicating over which dimensions we must sum over.
`weights` is a array of a weight for the different dimension.
`ind1` and `ind2` are the offsets for the difference.
"""
function generate_GR(num_dims, sum_dims_arr, weights, ind1, ind2; debug=false)
    out, add = [], []
    inds = map(1:num_dims) do di
        i = Symbol(:i, di)
    end

    for (d, w) in zip(sum_dims_arr, weights)
        inds1, inds2 = generate_indices(num_dims, d, ind1, ind2)
        push!(add, :($w * (arr[$(inds1...)] + arr[$(inds2...)])))
    end
    prefactor = - 4 / (abs(ind1) + abs(ind2))
    diff_factor = -sum(weights) * 2
    push!(add, :($diff_factor *arr[$(inds...)]))
    if debug
        push!(out, :(res = $prefactor * arr[$(inds...)] * +($(add...))))
    else
        push!(out, :(@tullio threads=false res = $prefactor * arr[$(inds...)] * +($(add...))))
    end
    return out
end

"""
    GR(; <keyword arguments>)
This function returns a function to calculate the Good's roughness regularizer
of a n-dimensional array.
# Arguments
- `num_dims=2`: Dimension of the array that should be regularized
- `sum_dims=[1, 2]`: A array containing the dimensions we want to sum over
- `weights=nothing`: A array containing weights to weight the contribution of
    different dimensions. If `weights=nothing` all dimensions are weighted equally.
- `step=1`: A integer indicating the step width for the array indexing
- `mode="forward"`: Either `"central"` or `"forward"` accounting for different
    modes of the spatial gradient. Default is "central".
# Examples
To create a regularizer for a 3D dataset where the third dimension
has different contribution. For the derivative we use forward mode.
```julia-repl
julia> reg = GR(num_dims=2, sum_dims=[1, 2], weights=[1, 1], mode="forward");

julia> reg([1 2 3; 4 5 6; 7 8 9])
-26.36561871738898
```
"""
function GR(; num_dims=2, sum_dims=[1, 2], weights=[1, 1], step=1,
              mode="central", ϵ=1f-8)
    if weights == nothing
        weights = ones(num_dims)
    end
    if mode == "central"
        GRf = @eval arr -> ($(generate_GR(num_dims, sum_dims, weights,
                                        step, (-1) * step)...))
    elseif mode == "forward"
        GRf = @eval arr -> ($(generate_GR(num_dims, sum_dims, weights,
                                        step, 0)...))
    else
        throw(ArgumentError("The provided mode is not valid."))
    end

    # we need to add a ϵ to prevent NaN in the derivative of it
    return arr -> GRf(sqrt.(arr .+ ϵ))
end

"""
    generate_TV(num_dims, sum_dims_arr, weights, ind1, ind2, ϵ)
Generate the Tullio statement for computing the Good's roughness.
`num_dims` is the dimension of the array.
`sum_dims_arr` is a array
indicating over which dimensions we must sum over.
`weights` is a array of a weight for the different dimension.
`ind1` and `ind2` are the offsets for the difference.
`ϵ` is a numerical constant to prevent division by zero.
    this is important for the gradient
"""
function generate_TV(num_dims, sum_dims_arr, weights, ind1, ind2, ϵ=1f-8; debug=false)
    out, add = [], []
    for (d, w) in zip(sum_dims_arr, weights)
        inds1, inds2 = generate_indices(num_dims, d, ind1, ind2)
        push!(add, :($w * abs2(arr[$(inds1...)] - arr[$(inds2...)])))
    end
    push!(add, ϵ)
    if debug
        push!(out, :(res = sqrt(+($(add...)))))
    else
        push!(out, :(@tullio threads=false res = sqrt(+($(add...)))))
    end
    return out
end

"""
    TV(; <keyword arguments>)
This function returns a function to calculate the Total Variation regularizer
of a n-dimensional array.
# Arguments
- `num_dims=2`:
- `sum_dims=[1, 2]`: A array containing the dimensions we want to sum over
- `weights=nothing`: A array containing weights to weight the contribution of
    different dimensions. If `weights=nothing` all dimensions are weighted equally.
- `step=1`: A integer indicating the step width for the array indexing
- `mode="central"`: Either `"central"` or `"forward"` accounting for different
    modes of the spatial gradient. Default is "central".
# Examples
To create a regularizer for a 3D dataset where the third dimension
has different contribution. For the derivative we use forward mode.
```julia-repl
julia> reg = TV(num_dims=2, sum_dims=[1, 2], weights=[1, 1], mode="forward");

julia> reg([1 2 3; 4 5 6; 7 8 9])
12.649111f0
```
"""
function TV(; num_dims=2, sum_dims=[1, 2], weights=nothing, step=1, mode="central")

    if weights == nothing
        weights = ones(num_dims)
    end

    if mode == "central"
        total_var = @eval arr -> ($(generate_TV(num_dims, sum_dims, weights,
                                        step, (-1) * step)...))
    elseif mode == "forward"
        total_var = @eval arr -> ($(generate_TV(num_dims, sum_dims, weights,
                                        step, 0)...))
    else
        throw(ArgumentError("The provided mode is not valid."))
    end
    return total_var
end


module LIPO

using LinearAlgebra

function lipo(f; k=0.0, n = 100, lower=0.0, upper=2.0, max_ndraws=1e4)
    Xdom = [0.0, upper]

    X = [rand()*(upper-lower)+lower]
    fX = [f(X[1])]

    U(Xtp1, X, fX, k) = minimum(fXt+k*norm(Xtp1 - Xt, 2) for (Xt, fXt) in zip(X, fX))
    t = 1
    ndraws = 0
    while t < n
       Xtp1 = rand()*(upper-lower)+lower
       ndraws += 1
       if U(Xtp1, X, fX, k) >= maximum(fX)
           X = push!(X, Xtp1)
           fX = push!(fX, f(X[end]))
           t += 1
       end
       if ndraws >= max_ndraws
           break
       end
    end
    return X, fX, ndraws
end

function adalipo(f; k_α=0.5, n = 100, lower=0.0, upper=2.0, max_ndraws=1e4, choose_max=true)
    Xdom = [lower, upper]

    # first, check that function isn't constant, run for max
    # safeguarding number of iterations (used to ensure a proper)

    X = [rand()*(upper-lower)+lower]
    fX = [f(X[1])]

    U(Xtp1, X, fX, k) = minimum(fXt+k*norm(Xtp1-Xt, 2) for (Xt, fXt) in zip(X, fX))
    t = 1
    ndraws = 0
    nexploit = 0
    nexplore = 0
    p = 0.4
    khat_i = 0.0
    while t < n
        if rand() < p
            Xtp1 = rand()*(upper-lower)+lower
            nexplore += 1
        else
            Xtp1 = exploit(fX, X, khat_i, choose_max, lower, upper)
            nexploit += 1
        end

       # a draw can come from
       ndraws += 1
       if U(Xtp1, X, fX, khat_i) >= maximum(fX)
           X = push!(X, Xtp1)
           fX = push!(fX, f(X[end]))
           t += 1
       end

       khat_i = khat(k_α, fX, X)

       if ndraws >= max_ndraws
           break
       end

    end
    return X, fX, ndraws
end

function exploit(fX, X, khat_i, choose_max, lower, upper)
    # pick either the largest or from the set of
    # tops
    val_at_tops, x_at_tops = find_tops(fX, X, khat_i, l)
    if chose_max
        i_m = argmax(val_at_tops)
    else
        i_m = rand(val_at_tops)
    end
    Xtp1 = x_at_tops[i_m]
end

function find_tops(fX, X, khat_i, lower, upper)
    xperm = sortperm(X)
    x = X[xperm]
    f = fX[perm]
    if f[1]+norm(x) < f[end]
        x_top = f[end]
        f_top = f[end]
    else
        x_top = x[1]
        f_top = f[1]
    end

    x_top = 0.0
    val_top = NaN
    for 1:length(x)-1
        if norm(x[i]-x[i+1])
    end
end

function khat(k_α, fX::AbstractVector{T}, X; max_depth = 100) where T
    #find estimates

    argmax_val = -T(1)
    for j = 1:length(X)
        for i = 1:length(X)
            if i == j
                continue
            end
            num = abs(fX[i] - fX[j])
            den = norm(X[i]- X[j], 2)
            candidate = num/den
            if candidate > argmax_val
                argmax_val = candidate
            end
        end
    end

    # Using the cloesd form solution in the paper
    k_i = ceil(log(argmax_val)/log(1+k_α))
    k_i
end

end # module


#!/usr/local/bin/julia
# last tested Julia version: 1.6.1
# Morris-Lecar model on a 2D lattice
# FvW 03/2018

using NPZ
using PyCall
using PyPlot
using Statistics
using VideoIO
@pyimport matplotlib.animation as anim

function ml2d(N, T, t0, dt, s, D, gL, VL, gCa, VCa, gK, VK, C, I0,
			  V0, V1, V2, V3, V4, phi, stim, blocks)
    C1 = 1.0/C
    # initialize Morris-Lecar system
    v = V0*ones(Float64,N,N)
    w = 0.015*ones(Float64,N,N)
    dv = zeros(Float64,N,N)
    dw = zeros(Float64,N,N)
    s_sqrt_dt = s*sqrt(dt)
    X = zeros(Float64,T,N,N)
    # stimulation protocol
    I = zeros(Float64,t0+T,N,N)
    for st in stim
        t_on, t_off = st[1]
        x0, x1 = st[2]
        y0, y1 = st[3]
        I[t0+t_on:t0+t_off, x0:x1, y0:y1] .= I0
    end

    # iterate
    for t in range(1, stop=t0+T, step=1)
        (t%100 == 0) && print("    t = ", t, "/", t0+T, "\r")
        # ML equations
        dv = C1.*(-gL.*(v.-VL) -gCa.*0.5.*(1.0 .+ tanh.((v.-V1)/V2)).*(v.-VCa)
                 -gK.*w.*(v.-VK) .+ I[t,:,:]) .+ D.*L(v)
        dw = phi*cosh.((v.-V3)/(2*V4)) .* (0.5*(1.0 .+ tanh.((v.-V3)/V4))-w)
        # Ito stochastic integration
        v += (dv*dt + s_sqrt_dt*randn(N,N))
        w += (dw*dt)
        # dead block(s):
        for bl in blocks
            v[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
            w[bl[1][1]:bl[1][2], bl[2][1]:bl[2][2]] .= 0.0
        end
        (t > t0) && (X[t-t0,:,:] = v)
    end
    println("\n")
    return X
end

function animate_pyplot(fname, data)
    """
    Animate 3D array as .mp4 using PyPlot, save as `fname`
    array dimensions:
        1: time
        2, 3: space
    """
    nt, nx, ny = size(data)
    vmin = minimum(data)
    vmax = maximum(data)
    # setup animation image
    fig = figure(figsize=(6,6))
    axis("off")
    t = imshow(data[1,:,:], origin="lower", cmap=ColorMap("gray"),
			   vmin=vmin, vmax=vmax)
    tight_layout()
    # frame generator
    println("[+] animate")
    function animate(i)
        (i%100 == 0) && print("    t = ", i, "/", nt, "\r")
        t.set_data(data[i+1,:,:])
    end
    # create animation
    ani = anim.FuncAnimation(fig, animate, frames=nt, interval=10)
    println("\n")
    # save animation
    ani[:save](fname, bitrate=-1,
               extra_args=["-vcodec", "libx264", "-pix_fmt", "yuv420p"])
    show()
end

function animate_video(fname, data)
    """
    Animate 3D array as .mp4 using VideoIO, save as `fname`
    array dimensions:
        1: time
        2, 3: space
    """
    # BW
    y = UInt8.(round.(255*(data .- minimum(data))/(maximum(data)-minimum(data))))
    # BW inverted
    #y = UInt8.(round.(255 .- 255*(data .- minimum(data))/(maximum(data)-minimum(data))))
    encoder_options = (color_range=2, crf=0, preset="medium")
    framerate=30
    T = size(data,1)
    open_video_out(fname, y[1,end:-1:1,:], framerate=framerate,
                   encoder_options=encoder_options) do writer
        for i in range(2,stop=T,step=1)
            write(writer, y[i,end:-1:1,:])
        end
    end
end

function L(x)
    # Laplace operator
    # periodic boundary conditions
    xU = circshift(x, [-1 0])
    xD = circshift(x, [1 0])
    xL = circshift(x, [0 -1])
    xR = circshift(x, [0 1])
    Lx = xU + xD + xL + xR - 4x
    # non-periodic boundary conditions
    Lx[1,:] .= 0.0
    Lx[end,:] .= 0.0
    Lx[:,1] .= 0.0
    Lx[:,end] .= 0.0
    return Lx
end

function main()
    println("Morris-Lecar (ML) lattice model\n")
    N = 128
    T = 15000
    t0 = 1000
	dt = 0.01
	s = 2.0
	D = 0.8
	# Morris-Lecar parameters
    gL = 2.0
    VL = -60.0
    gCa = 5.5 # 4.0
    VCa = 120.0
    gK = 8.0
    VK = -84.0
    C = 5.0 # 20.0
    I = 120.0
    V0 = -60.0
    V1 = -1.2
    V2 = 18.0
    V3 = 12.0 # 2.0
    V4 = 17.4 # 30.0
    phi = 1.0/15.0 # 0.04
    println("[+] Lattice size N: ", N)
    println("[+] Time steps T: ", T)
    println("[+] Warm-up steps t0: ", t0)
	println("[+] Integration time step dt: ", dt)
	println("[+] Noise intensity s: ", s)
	println("[+] Diffusion coefficient D: ", D)
    println("[+] ML parameter gL: ", gL)
    println("[+] ML parameter VL: ", VL)
    println("[+] ML parameter gCa: ", gCa)
    println("[+] ML parameter VCa: ", VCa)
    println("[+] ML parameter gK: ", gK)
    println("[+] ML parameter VK: ", VK)
    println("[+] ML parameter C: ", C)
    println("[+] ML parameter V0: ", V0)
    println("[+] ML parameter V1: ", V1)
    println("[+] ML parameter V2: ", V2)
    println("[+] ML parameter V3: ", V3)
    println("[+] ML parameter V4: ", V4)
    println("[+] ML parameter phi: ", phi)
	println("[+] Stimulation current I: ", I)

    # stim protocol, array of elements [[t0,t1], [x0,x1], [y0,y1]]
    # double spiral 1
    #stim = [ [[50,250], [1,5], [1,5]],
    #         [[5800,6200], [20,25], [1,35]] ]
    # double spiral 2
    #stim = [ [[50,250], [1,5], [1,5]],
    #         [[5800,6200], [18,23], [5,26]] ]
    # spiral wave 1
    #stim = [ [[50,250], [1,5], [1,5]],
    #         [[6300,6800], [30,35], [5,20]] ]
    # spiral wave 2
    stim = [ [[50,250], [1,5], [1,5]],
             [[6400,6900], [30,35], [5,20]] ]
    #stim = []

    # dead blocks, array of elementy [[x0,x1], [y0,y1]]
    blocks = [ [[1,15], [5,10]] ]
    #blocks = []

    # run simulation
    data = ml2d(N, T, t0, dt, s, D, gL, VL, gCa, VCa, gK, VK, C, I,
				V0, V1, V2, V3, V4, phi, stim, blocks)
    println("[+] Data dimensions: ", size(data))

    # plot mean voltage
    m = mean(reshape(data, (T,N*N)), dims=2)
    plot(m, "-k"); show()

    # save data
    I_str = rpad(I, 4, '0') # stim. current amplitude as 4-char string
    s_str = rpad(s, 4, '0') # noise as 4-char string
    D_str = rpad(D, 4, '0') # diffusion coefficient as 4-char string
    #fname1 = string("ml2d_I_", I_str, "_s_", s_str, "_D_", D_str, ".npy")
    #npzwrite(fname1, data)
    #println("[+] Data saved as: ", fname1)

    # video
    fname2 = string("ml2d_I_", I_str, "_s_", s_str, "_D_", D_str, ".mp4")
    #animate_pyplot(fname2, data) # slow
    animate_video(fname2, data) # fast
    println("[+] Data saved as: ", fname2)
end

main()


export PositiveSemidefinite

struct PositiveSemidefinite{T <: Real} <: ConvexSet
    r::T
end
PositiveSemidefinite() = PositiveSemidefinite(0.0)
```
This function projects the symmetric matrix Y onto the closest positive definite matrix.
```

function project!(s::PositiveSemidefinite{<:Real}, v::Matrix{T}, y::Matrix{T}) where {T <: Real}
   (D,V) = LinearAlgebra.eigen(y)
   X = zeros(size(y))
   for j = 1:size(y,1)
      if D[j] >= 0.0
         X = X .+ V[:,j] * D[j] * V[:,j]'
      end
   end
   v = X
end

function project(s::PositiveSemidefinite{<:Real}, y::Matrix{<:Real})
    project!(s, similar(y), y)
end


using Displaz

function example(N)
    t = linspace(0,1, N)'
    P = (10.+0.2*rand(size(t))) .* vcat(exp(5*t).*cos(100*pi*t),
                                        exp(5*t).*sin(100*pi*t),
                                        100*t);
    C = vcat(t, 1.-t, zeros(t));
    sz  = map(Float32, t)
    shape = mod(rand(UInt8, N), 6)
    plot3d!(P, color = C, label="Top series", markershape=shape)
    P[3,:] = -P[3,:]
    plot3d!(P, color = C, markersize = sz, markershape="-", label="Bottom series")
end

example(1_000_000)


export Cheb_dir_AD

function Cheb_dir_AD(nlp            :: AbstractNLPModel,
                     nlp_at_x       :: NLPAtX,
                     x              :: AbstractVector,
                     Nwtdirection   :: Function)

    dₙ = Nwtdirection(nlp_at_x.Hx, nlp_at_x.gx, verbose = false)

    ∇³fdₙdₙ = ∇f³xuv(nlp, x, dₙ, dₙ)

    dCp = Nwtdirection(nlp_at_x.Hx, ∇³fdₙdₙ, verbose = false)

    dC = dₙ + 0.5 * dCp

    return dC
end


# samplepath.jl

"""
    abstract type AbstractSamplePath <: Any

subtype: SamplePath.

Its subtype implements a continuous samplepath.
"""
abstract type AbstractSamplePath end

"""
    isincreasing(x::AbstractVector{T})::Bool where T <: Number
    isincreasing(x::Range{T})::Bool where T <: Number

Tests whether a vector of numbers is strictly increasing. Is internal to
NonparametricBayesForDiffusions.
"""
function isincreasing(x::S) where S<:AbstractArray{T} where T <: Number
  for i in 1:length(x)-1
    x[i+1] <= x[i] && return false
  end
  return true
end

isincreasing(x::S) where S<:Range{T} where T <: Number = step(x) > 0

"""
    SamplePath(timeinterval::S, samplevalues::T) where
        {S<:AbstractVector{Float64}, T<:AbstractVector{Float64}}

 Implements a continuous samplepath where the timeinterval is not necessarily
 equally spaced. Sample value samplevalues[k] is the value of the process at
 time timeinterval[k]. timeinterval is an increasing vector. timeinterval and
 samplevalues are of equal length.

 Performance hint: Use Range objects instead of vectors as much as possible, this
 will speed up computation time.

# Constructors:

```julia
    SamplePath(timeinterval, samplevalues)
    SamplePath(timeinterval, f::Function) =
        SamplePath(timeinterval, f.(timeinterval))
```

# Examples

```julia
t = 0.0:0.01:2.0
X = SamplePath(t, sinpi)
```

```julia
t = 0.0:0.01:1.0
v = map(x->floor(10*x), t)
X = SamplePath(t, v)
```
"""
struct SamplePath{S, T} <: AbstractSamplePath where {S<:AbstractVector{Float64},
        T<:AbstractVector{Float64}}
  timeinterval::S
  samplevalues::T

  function SamplePath(timeinterval::S, samplevalues::T) where {
      S<:AbstractVector{Float64}, T<:AbstractVector{Float64}}
    length(timeinterval) == length(samplevalues) ||
      throw(DimensionMismatch("Length of timeinterval should be equal
       to the length of the samplevalues vector"))
    isincreasing(timeinterval) ||
      throw(ArgumentError("Timeinterval should be increasing"))
    new{S,T}(timeinterval, samplevalues)
  end
end

step(X::SamplePath{S}) where S<:Range{T} where T<:Number = step(X.timeinterval)

#constructor
SamplePath(timeinterval, f::Function) =
    SamplePath(timeinterval, f.(timeinterval))

# Returns the length of timeinterval == samplevalues vectors, not the endtime!
# TO DO: Answers to examples
"""
    length(X::SamplePath)

Returns the length of the vector timeinterval == length vector samplevalues.

# Examples
```julia
X = SamplePath([0.,1.,2.], [3.,5., -1.])
length(X)
```

```julia
X = SamplePath(0.0:0.1:2π, sin)
length(X) == length(0.0:0.1:2π)
```
"""
length(X::SamplePath)=length(X.timeinterval)


function DLDR(ALPHA,BETA)
      A=[0.0050    0.0070    0.0130    0.0180    0.0150    0.0210    0.0230
         0.0170    0.0160    0.0130    0.0150    0.0140    0.0110    0.0100
         0.0140    0.0140    0.0110    0.0150    0.0130    0.0100    0.0110
         0.0100    0.0140    0.0120    0.0140    0.0130    0.0110    0.0110
        -0.0050    0.0130    0.0110    0.0140    0.0120    0.0100    0.0110
         0.0090    0.0090    0.0090    0.0140    0.0110    0.0090    0.0100
         0.0190    0.0120    0.0080    0.0140    0.0110    0.0080    0.0080
         0.0050    0.0050    0.0050    0.0150    0.0100    0.0100    0.0100
              0         0   -0.0020    0.0130    0.0080    0.0060    0.0060
        -0.0050    0.0040    0.0050    0.0110    0.0080    0.0050    0.0140
        -0.0110    0.0090    0.0030    0.0060    0.0070         0    0.0200
         0.0080    0.0070    0.0050    0.0010    0.0030    0.0010         0]
      #
      S= 0.2 * ALPHA
      K= floor(S)+3
      if(K <= 1)
            K= 2
      end
      if(K >=  12)
            K=  11
      end
      DA= S - K +3
      L = K + fix( 1.1*sign(DA) )
      #
      S= 0.1 * BETA
      M= floor(S)+4
      if(M <= 1)
            M= 2
      end
      if(M >=  7)
            M=  6
      end
      DB= S - M + 4
      N= M + fix( 1.1*sign(DB) )
      #
      K=convert(Int64, K)
      L=convert(Int64, L)
      M=convert(Int64, M)
      N=convert(Int64, N)
      T= A[K,M]
      U= A[K,N]
      V= T + abs(DA) * (A[L,M] - T)
      W= U + abs(DA) * (A[L,N] - U)
      return dldrv = V + (W-V)  * abs(DB)
end


mutable struct TripleBandLinearOp{FM <: FdmMesher}
  direction::Int
  mesher::FM
  i0::Vector{Int}
  i2::Vector{Int}
  reverseIndex::Vector{Int}
  lower::Vector{Float64}
  _diag::Vector{Float64}
  upper::Vector{Float64}
end

function TripleBandLinearOp(direction::Int, mesher::FdmMesher)
  sz = mesher.layout.size
  i0 = Vector{Int}(undef, sz)
  i2 = Vector{Int}(undef, sz)
  reverseIndex = Vector{Int}(undef, sz)
  lower = Vector{Float64}(undef, sz)
  _diag = Vector{Float64}(undef, sz)
  upper = Vector{Float64}(undef, sz)

  newDim = copy(mesher.layout.dim)
  newDim[1], newDim[direction] = newDim[direction], newDim[1] # swap

  newSpacing = FdmLinearOpLayout(newDim).spacing
  newSpacing[1], newSpacing[direction] = newSpacing[direction], newSpacing[1] # swap

  coords = ones(Int, length(mesher.layout.dim))
  @inbounds @simd for i = 1:sz
    ## TripleBandLinearOp part
    i0[i] = neighborhood(mesher.layout, i, coords, direction, -1)
    i2[i] = neighborhood(mesher.layout, i, coords, direction, 1)

    newIndex = dot(coords .- 1, newSpacing) + 1

    reverseIndex[newIndex] = i

    iter_coords!(coords, mesher.layout.dim)
  end

  return TripleBandLinearOp(direction, mesher, i0, i2, reverseIndex, lower, _diag, upper)
end

function FirstDerivativeOp(direction::Int, mesher::FdmMesher)
  sz = mesher.layout.size
  i0 = Vector{Int}(undef, sz)
  i2 = Vector{Int}(undef, sz)
  reverseIndex = Vector{Int}(undef, sz)
  lower = Vector{Float64}(undef, sz)
  _diag = Vector{Float64}(undef, sz)
  upper = Vector{Float64}(undef, sz)

  newDim = copy(mesher.layout.dim)
  newDim[1], newDim[direction] = newDim[direction], newDim[1] # swap

  newSpacing = FdmLinearOpLayout(newDim).spacing
  newSpacing[1], newSpacing[direction] = newSpacing[direction], newSpacing[1] # swap

  coords = ones(Int, length(mesher.layout.dim))
  @inbounds @simd for i = 1:sz
    ## TripleBandLinearOp part
    i0[i] = neighborhood(mesher.layout, i, coords, direction, -1)
    i2[i] = neighborhood(mesher.layout, i, coords, direction, 1)

    newIndex = dot(coords .- 1, newSpacing) + 1

    reverseIndex[newIndex] = i

    ## FirstDerivativeOp part
    hm = get_dminus(mesher, coords, direction)
    hp = get_dplus(mesher, coords, direction)

    zetam1 = hm * (hm + hp)
    zeta0 = hm * hp
    zetap1 = hp * (hm + hp)

    if coords[direction] == 1
      # upwinding scheme
      lower[i] = 0.0
      upper[i] = 1.0 / hp
      _diag[i] = -(upper[i])
    elseif coords[direction] == mesher.layout.dim[direction]
      # downwinding scheme
      _diag[i] = 1.0 / hm
      lower[i] = -_diag[i]
      upper[i] = 0.0
    else
      lower[i] = -hp / zetam1
      _diag[i] = (hp - hm) / zeta0
      upper[i] = hm / zetap1
    end
    iter_coords!(coords, mesher.layout.dim)
  end

  return TripleBandLinearOp(direction, mesher, i0, i2, reverseIndex, lower, _diag, upper)
end

function SecondDerivativeOp(direction::Int, mesher::FdmMesher)
  sz = mesher.layout.size
  i0 = Vector{Int}(undef, sz)
  i2 = Vector{Int}(undef, sz)
  reverseIndex = Vector{Int}(undef, sz)
  lower = Vector{Float64}(undef, sz)
  _diag = Vector{Float64}(undef, sz)
  upper = Vector{Float64}(undef, sz)

  newDim = copy(mesher.layout.dim)
  newDim[1], newDim[direction] = newDim[direction], newDim[1] # swap

  newSpacing = FdmLinearOpLayout(newDim).spacing
  newSpacing[1], newSpacing[direction] = newSpacing[direction], newSpacing[1] # swap

  coords = ones(Int, length(mesher.layout.dim))
  @inbounds @simd for i = 1:sz
    ## TripleBandLinearOp part
    i0[i] = neighborhood(mesher.layout, i, coords, direction, -1)
    i2[i] = neighborhood(mesher.layout, i, coords, direction, 1)

    newIndex = dot(coords .- 1, newSpacing) + 1

    reverseIndex[newIndex] = i

    ## SecondDerivativeOp part
    hm = get_dminus(mesher, coords, direction)
    hp = get_dplus(mesher, coords, direction)

    zetam1 = hm * (hm + hp)
    zeta0 = hm * hp
    zetap1 = hp * (hm + hp)

    co = coords[direction]

    if co == 1 || co == mesher.layout.dim[direction]
      lower[i] = upper[i] = _diag[i] = 0.0
    else
      lower[i] = 2.0 / zetam1
      _diag[i] = -2.0 / zeta0
      upper[i] = 2.0 / zetap1
    end
    iter_coords!(coords, mesher.layout.dim)
  end

  return TripleBandLinearOp(direction, mesher, i0, i2, reverseIndex, lower, _diag, upper)
end

function mult!(trpBandLinOp::TripleBandLinearOp, u::Vector{T}) where {T <: Number}
  # probably should protect for array dims here
  trpBandLinOp.lower = trpBandLinOp.lower .* u
  trpBandLinOp._diag = trpBandLinOp._diag .* u
  trpBandLinOp.upper = trpBandLinOp.upper .* u

  return trpBandLinOp
end

function add!(trpBandLinOp::TripleBandLinearOp, m::TripleBandLinearOp)
  trpBandLinOp.lower += m.lower
  trpBandLinOp._diag += m._diag
  trpBandLinOp.upper += m.upper

  return trpBandLinOp
end

function axpyb!(trpBandLinOp::TripleBandLinearOp, a::Vector{Float64}, x::TripleBandLinearOp, y::TripleBandLinearOp, b::Vector{Float64})
  sz = trpBandLinOp.mesher.layout.size

  if length(a) == 0
    if length(b) == 0
      trpBandLinOp._diag[1:sz] = y._diag
      trpBandLinOp.lower[1:sz] = y.lower
      trpBandLinOp.upper[1:sz] = y.upper
    else
      addB = length(b) > 1 ? b[1:sz] : b[1]
      trpBandLinOp._diag[1:sz] = y._diag + addB
      trpBandLinOp.lower[1:sz] = y.lower
      trpBandLinOp.upper[1:sz] = y.upper
    end
  elseif length(b) == 0
    # this might be improved
    trpBandLinOp._diag[1:sz] = y._diag + (length(a) > 0 ? a .* x._diag : a[1] * x._diag)
    trpBandLinOp.lower[1:sz] = y.lower + (length(a) > 0 ? a .* x.lower : a[1] * x.lower)
    trpBandLinOp.upper[1:sz] = y.upper + (length(a) > 0 ? a .* x.upper : a[1] * x.upper)
  else
    addB = length(b) > 1 ? b[1:sz] : b[1]
    trpBandLinOp._diag[1:sz] = y._diag + (length(a) > 0 ? a .* x._diag : a[1] * x._diag) + addB
    trpBandLinOp.lower[1:sz] = y.lower + (length(a) > 0 ? a .* x.lower : a[1] * x.lower)
    trpBandLinOp.upper[1:sz] = y.upper + (length(a) > 0 ? a .* x.upper : a[1] * x.upper)
  end

  return trpBandLinOp
end

function apply(trpBandLinOp::TripleBandLinearOp, r::Vector{Float64})
  idx = trpBandLinOp.mesher.layout
  length(r) == idx.size || error("inconsistent length of r")

  retArray = Vector{Float64}(undef, length(r))

  for i = 1:idx.size
    retArray[i] = r[trpBandLinOp.i0[i]] * trpBandLinOp.lower[i] + r[i] * trpBandLinOp._diag[i] + r[trpBandLinOp.i2[i]] * trpBandLinOp.upper[i]
  end

  return retArray
end

function solve_splitting(trpBandLinOp::TripleBandLinearOp, r::Vector{Float64}, a::Float64, b::Float64)
  layout = trpBandLinOp.mesher.layout
  length(r) == layout.size || error("inconsistent length of r")

  retArray = Vector{Float64}(undef, length(r))
  tmp = Vector{Float64}(undef, length(r))

  # solving a tridiagonal system (we could use Julia built in f'ns here)
  rim1 = trpBandLinOp.reverseIndex[1]
  bet = 1.0 / (a * trpBandLinOp._diag[rim1] + b)
  bet != 0.0 || error("division by zero")

  retArray[rim1] = r[rim1] * bet

  @inbounds @simd for j = 2:layout.size
    ri = trpBandLinOp.reverseIndex[j]
    tmp[j] = a * trpBandLinOp.upper[rim1] * bet

    bet = b + a * (trpBandLinOp._diag[ri] - tmp[j] * trpBandLinOp.lower[ri])
    bet != 0.0 || error("division by zero")

    bet = 1.0 / bet

    retArray[ri] = (r[ri] - a * trpBandLinOp.lower[ri] * retArray[rim1]) * bet
    rim1 = ri
  end

  @inbounds @simd for j = layout.size - 1:-1:2
    retArray[trpBandLinOp.reverseIndex[j]] -= tmp[j + 1] * retArray[trpBandLinOp.reverseIndex[j + 1]]
  end

  retArray[trpBandLinOp.reverseIndex[1]] -= tmp[2] * retArray[trpBandLinOp.reverseIndex[2]]

  return retArray
end

mutable struct SecondOrderMixedDerivativeOp{FD <: FdmMesher} <: NinePointLinearOp
  d1::Int
  d2::Int
  i00::Vector{Int}
  i10::Vector{Int}
  i20::Vector{Int}
  i01::Vector{Int}
  i21::Vector{Int}
  i02::Vector{Int}
  i12::Vector{Int}
  i22::Vector{Int}
  a00::Vector{Float64}
  a10::Vector{Float64}
  a20::Vector{Float64}
  a01::Vector{Float64}
  a11::Vector{Float64}
  a21::Vector{Float64}
  a02::Vector{Float64}
  a12::Vector{Float64}
  a22::Vector{Float64}
  mesher::FD
end

function SecondOrderMixedDerivativeOp(d1::Int, d2::Int, mesher::FD) where {FD <: FdmMesher}
  sz = mesher.layout.size
  i00 = Vector{Int}(undef, sz)
  i10 = Vector{Int}(undef, sz)
  i20 = Vector{Int}(undef, sz)
  i01 = Vector{Int}(undef, sz)
  i21 = Vector{Int}(undef, sz)
  i02 = Vector{Int}(undef, sz)
  i12 = Vector{Int}(undef, sz)
  i22 = Vector{Int}(undef, sz)
  a00 = Vector{Float64}(undef, sz)
  a10 = Vector{Float64}(undef, sz)
  a20 = Vector{Float64}(undef, sz)
  a01 = Vector{Float64}(undef, sz)
  a11 = Vector{Float64}(undef, sz)
  a21 = Vector{Float64}(undef, sz)
  a02 = Vector{Float64}(undef, sz)
  a12 = Vector{Float64}(undef, sz)
  a22 = Vector{Float64}(undef, sz)

  coords = ones(Int, length(mesher.layout.dim))
  @inbounds @simd for i = 1:sz
    # NinePointLinearOp part
    i10[i] = neighborhood(mesher.layout, i, coords, d2, -1)
    i01[i] = neighborhood(mesher.layout, i, coords, d1, -1)
    i21[i] = neighborhood(mesher.layout, i, coords, d1,  1)
    i12[i] = neighborhood(mesher.layout, i, coords, d2,  1)
    i00[i] = neighborhood(mesher.layout, i, coords, d1, -1, d2, -1)
    i20[i] = neighborhood(mesher.layout, i, coords, d1,  1, d2, -1)
    i02[i] = neighborhood(mesher.layout, i, coords, d1, -1, d2,  1)
    i22[i] = neighborhood(mesher.layout, i, coords, d1,  1, d2,  1)

    # SecondOrderMixedDerivativeOp part
    hm_d1 = get_dminus(mesher, coords, d1)
    hp_d1 = get_dplus(mesher, coords, d1)
    hm_d2 = get_dminus(mesher, coords, d2)
    hp_d2 = get_dplus(mesher, coords, d2)

    zetam1 = hm_d1 * (hm_d1 + hp_d1)
    zeta0 = hm_d1 * hp_d1
    zetap1 = hp_d1 * (hm_d1 + hp_d1)

    phim1 = hm_d2 * (hm_d2 + hp_d2)
    phi0 = hm_d2 * hp_d2
    phip1 = hp_d2 * (hm_d2 + hp_d2)

    c1 = coords[d1]
    c2 = coords[d2]

    if c1 == 1 && c2 == 1
      # lower left corner
      a00[i] = a01[i] = a02[i] = a10[i] = a20[i] = 0.0
      a11[i] = a22[i] = 1.0 / (hp_d1 * hp_d2)
      a21[i] = a12[i] = -a11[i]
    elseif c1 == mesher.layout.dim[d1] && c2 == 1
      # upper left corner
      a22[i] = a21[i] = a20[i] = a10[i] = a00[i] = 0.0
      a01[i] = a12[i] = 1.0 / (hm_d1 * hp_d2)
      a11[i] = a02[i] = -a01[i]
    elseif c1 == 1 && c2 == mesher.layout.dim[d2]
      # lower right corner
      a00[i] = a01[i] = a02[i] = a12[i] = a22[i] = 0.0
      a10[i] = a21[i] = 1.0 / (hp_d1 * hm_d2)
      a20[i] = a11[i] = -a10[i]
    elseif c1 == mesher.layout.dim[d1] && c2 == mesher.layout.dim[d2]
      # upper right corner
      a20[i] = a21[i] = a22[i] = a12[i] = a02[i] = 0.0
      a00[i] = a11[i] = 1.0 / (hm_d1 * hm_d2)
      a10[i] = a01[i] = -a00[i]
    elseif c1 == 1
      # lower side
      a00[i] = a01[i] = a02[i] = 0.0
      a10[i] = hp_d2 / (hp_d1 * phim1)
      a20[i] = -a10[i]
      a21[i] = (hp_d2 - hm_d2) / (hp_d1 * phi0)
      a11[i] = -a21[i]
      a22[i] = hm_d2 / (hp_d1 * phip1)
      a12[i] = -a22[i]
    elseif c1 == mesher.layout.dim[d1]
      # upper side
      a20[i] = a21[i] = a22[i] = 0.0
      a00[i] = hp_d2 / (hm_d1 * phim1)
      a10[i] = -a00[i]
      a11[i] = (hp_d2 - hm_d2) / (hm_d1 * phi0)
      a01[i] = -a11[i]
      a12[i] = hm_d2 / (hm_d1 * phip1)
      a02[i] = -a12[i]
    elseif c2 == 1
      # left side
      a00[i] = a10[i] = a20[i] = 0.0
      a01[i] = hp_d1 / (zetam1 * hp_d2)
      a02[i] = -a01[i]
      a12[i] = (hp_d1 - hm_d1) / (zeta0 * hp_d2)
      a11[i] = -a12[i]
      a22[i] = hm_d1 / (zetap1 * hp_d2)
      a21[i] = -a22[i]
    elseif c2 == mesher.layout.dim[d2]
      # right side
      a22[i] = a12[i] = a02[i] = 0.0
      a00[i] = hp_d1 / (zetam1 * hm_d2)
      a01[i] = -a00[i]
      a11[i] = (hp_d1 - hm_d1) / (zeta0 * hm_d2)
      a10[i] = -a11[i]
      a21[i] = hm_d1 / (zetap1 * hm_d2)
      a20[i] = -a21[i]
    else
      a00[i] = hp_d1 * hp_d2 / (zetam1 * phim1)
      a10[i] = -(hp_d1 - hm_d1) * hp_d2 / (zeta0 * phim1)
      a20[i] = -hm_d1 * hp_d2 / (zetap1 * phim1)
      a01[i] = -hp_d1 * (hp_d2 - hm_d2) / (zetam1 * phi0)
      a11[i] = (hp_d1 - hm_d1) * (hp_d2 - hm_d2) / (zeta0 * phi0)
      a21[i] = hm_d1 * (hp_d2 - hm_d2) / (zetap1 * phi0)
      a02[i] = -hp_d1 * hm_d2 / (zetam1 * phip1)
      a12[i] = hm_d2 * (hp_d1 - hm_d1) / (zeta0 * phip1)
      a22[i] = hm_d1 * hm_d2 / (zetap1 * phip1)
    end
    iter_coords!(coords, mesher.layout.dim)
  end

  return SecondOrderMixedDerivativeOp{FD}(d1, d2, i00, i10, i20, i01, i21, i02, i12, i22, a00, a10, a20, a01, a11, a21, a02, a12, a22, mesher)
end

function mult!(ninePointLin::NinePointLinearOp, u::Vector{T}) where {T <: Number}
  ninePointLin.a11 = ninePointLin.a11 .* u
  ninePointLin.a01 = ninePointLin.a01 .* u
  ninePointLin.a10 = ninePointLin.a10 .* u
  ninePointLin.a21 = ninePointLin.a21 .* u
  ninePointLin.a22 = ninePointLin.a22 .* u
  ninePointLin.a00 = ninePointLin.a00 .* u
  ninePointLin.a02 = ninePointLin.a02 .* u
  ninePointLin.a20 = ninePointLin.a20 .* u
  ninePointLin.a12 = ninePointLin.a12 .* u

  return ninePointLin
end

function apply(ninePointLin::NinePointLinearOp, u::Vector{Float64})
  length(u) == ninePointLin.mesher.layout.size || error("inconsistent length of r")

  retVal = ninePointLin.a00 .* u[ninePointLin.i00] + ninePointLin.a01 .* u[ninePointLin.i01] + ninePointLin.a02 .* u[ninePointLin.i02] + ninePointLin.a10 .* u[ninePointLin.i10] +
          ninePointLin.a11 .* u + ninePointLin.a12 .* u[ninePointLin.i12] + ninePointLin.a20 .* u[ninePointLin.i20] + ninePointLin.a21 .* u[ninePointLin.i21] +
          ninePointLin.a22 .* u[ninePointLin.i22]

  return retVal
end

mutable struct FdmG2Op{FD <: FdmMesher} <: FdmLinearOpComposite
  direction1::Int
  direction2::Int
  x::Vector{Float64}
  y::Vector{Float64}
  dxMap::TripleBandLinearOp{FD}
  dyMap::TripleBandLinearOp{FD}
  corrMap::SecondOrderMixedDerivativeOp{FD}
  mapX::TripleBandLinearOp{FD}
  mapY::TripleBandLinearOp{FD}
  model::G2
end

function FdmG2Op(mesher::FD, model::G2, direction1::Int, direction2::Int) where {FD <: FdmMesher}
  x = get_locations(mesher, direction1)
  y = get_locations(mesher, direction2)

  dxMap = add!(mult!(FirstDerivativeOp(direction1, mesher), (-x * get_a(model))),
              mult!(SecondDerivativeOp(direction1, mesher), (0.5 * get_sigma(model) * get_sigma(model) * ones(mesher.layout.size))))

  dyMap = add!(mult!(FirstDerivativeOp(direction2, mesher), (-y * get_b(model))),
              mult!(SecondDerivativeOp(direction2, mesher), (0.5 * get_eta(model) * get_eta(model) * ones(mesher.layout.size))))

  corrMap = mult!(SecondOrderMixedDerivativeOp(direction1, direction2, mesher), fill(get_rho(model) * get_sigma(model) * get_eta(model), mesher.layout.size))

  mapX = TripleBandLinearOp(direction1, mesher)
  mapY = TripleBandLinearOp(direction2, mesher)

  return FdmG2Op{FD}(direction1, direction2, x, y, dxMap, dyMap, corrMap, mapX, mapY, model)
end

mutable struct FdmHullWhiteOp{FD <: FdmMesher} <: FdmLinearOpComposite
  direction::Int
  x::Vector{Float64}
  dzMap::TripleBandLinearOp{FD}
  mapT::TripleBandLinearOp{FD}
  model::HullWhite
end

function FdmHullWhiteOp(mesher::FD, model::HullWhite, direction::Int) where {FD <: FdmMesher}
  x = get_locations(mesher, direction)

  dzMap = add!(mult!(FirstDerivativeOp(direction, mesher), (-x * get_a(model))),
              mult!(SecondDerivativeOp(direction, mesher), (0.5 * get_sigma(model) * get_sigma(model)) * ones(mesher.layout.size)))

  mapT = TripleBandLinearOp(direction, mesher)

  return FdmHullWhiteOp{FD}(direction, x, dzMap, mapT, model)
end

# OP methods #
function set_time!(op::FdmG2Op, t1::Float64, t2::Float64)
  dynamics = get_dynamics(op.model)

  phi = 0.5 * (short_rate(dynamics, t1, 0.0, 0.0) + short_rate(dynamics, t2, 0.0, 0.0))

  hr = -0.5 * (op.x + op.y .+ phi)
  axpyb!(op.mapX, Vector{Float64}(), op.dxMap, op.dxMap, hr)
  axpyb!(op.mapY, Vector{Float64}(), op.dyMap, op.dyMap, hr)

  return op
end

function set_time!(op::FdmHullWhiteOp, t1::Float64, t2::Float64)
  dynamics = get_dynamics(op.model)

  phi = 0.5 * (short_rate(dynamics, t1, 0.0) + short_rate(dynamics, t2, 0.0))

  axpyb!(op.mapT, Vector{Float64}(), op.dzMap, op.dzMap, -(op.x .+ phi))

  return op
end

function apply_direction(op::FdmG2Op, direction::Int, r::Vector{Float64})
  if direction == op.direction1
    return apply(op.mapX, r)
  elseif direction == op.direction2
    return apply(op.mapY, r)
  else
    return zeros(length(r))
  end
end

function apply_direction(op::FdmHullWhiteOp, direction::Int, r::Vector{Float64})
  if direction == op.direction
    return apply(op.mapT, r)
  else
    return zeros(length(r))
  end
end

function solve_splitting(op::FdmG2Op, direction::Int, r::Vector{Float64}, a::Float64)
  if direction == op.direction1
    return solve_splitting(op.mapX, r, a, 1.0)
  elseif direction == op.direction2
    return solve_splitting(op.mapY, r, a, 1.0)
  else
    return zeros(length(r))
  end
end

function solve_splitting(op::FdmHullWhiteOp, direction::Int, r::Vector{Float64}, a::Float64)
  if direction == op.direction
    return solve_splitting(op.mapT, r, a, 1.0)
  else
    return zeros(length(r))
  end
end

apply_mixed(op::FdmG2Op, r::Vector{Float64}) = apply(op.corrMap, r)
apply_mixed(op::FdmHullWhiteOp, r::Vector{Float64}) = zeros(length(r))

apply(op::FdmG2Op, r::Vector{Float64}) = apply(op.mapX, r) + apply(op.mapY, r) + apply_mixed(op, r)
apply(op::FdmHullWhiteOp, r::Vector{Float64}) = apply(op.mapT, r)

get_size(op::FdmG2Op) = 2
get_size(op::FdmHullWhiteOp) = 1


using TrajectoryOptimization
using Plots
# Pendulum
n = 2
m = 1
model, = Dynamics.pendulum # same as UDP pendulum dynamics

Qf = 30.0*Matrix(I,n,n)
Q = 0.3*Matrix(I,n,n)
R = 0.3*Matrix(I,m,m)
dt = 0.1
tf = 5.0
x0 = [0.0; 0.0]
xf = [pi; 0.0]

obj = LQRObjective(Q,R,Qf,tf,x0,xf)

opts = SolverOptions()
opts.cost_tolerance = 1e-5
opts.gradient_norm_tolerance = 1e-5
opts.verbose = true
opts.cache = true
opts.c2 = 100.0

solver = TrajectoryOptimization.Solver(model,obj,dt=dt,opts=opts)
U = zeros(solver.model.m, solver.N)
results, stats = TrajectoryOptimization.solve(solver,U)

# Cartpole
n = 4
m = 1
model = Model(Dynamics.cartpole_dynamic_udp!,n,m)

dt = .1; # time step
# quadratic costs
Qf = 1000*Matrix(I,n,n);
Q = .1*Matrix(I,n,n)
R = .01*Matrix(I,m,m);

tf = 5.0
x0 = [0.0; 0.0; 0.0; 0.0]
xf = [0.0; pi; 0.0; 0.0]

obj = LQRObjective(Q,R,Qf,tf,x0,xf)

opts = SolverOptions()
opts.cost_tolerance = 1e-5
opts.gradient_norm_tolerance = 1e-5
opts.verbose = true
opts.cache = true

solver = TrajectoryOptimization.Solver(model,obj,dt=dt,opts=opts)
U = zeros(solver.model.m, solver.N)
results, stats = TrajectoryOptimization.solve(solver,U)
println(results.X[:,end])
# plot(results.cost[1:results.termination_index])


function centre_dual!(point::Class_point, comp_feas::Float64)
  y_c = point.mu ./ point.s

  point.y = min( y_c / comp_feas, max(point.y, comp_feas * y_c))
end

function dual_scale(iter::Class_iterate, pars::Class_parameters)
    if pars.term.dual_scale_mode == :ipopt
      return pars.term.dual_scale_threshold / max(mean(get_y(iter)), pars.term.dual_scale_threshold)
    elseif pars.term.dual_scale_mode == :max_dual
      return pars.term.dual_scale_threshold / max(norm(get_y(iter),Inf), pars.term.dual_scale_threshold)
    elseif pars.term.dual_scale_mode == :sqrt
      return pars.term.dual_scale_threshold / max(sqrt(norm(get_y(iter), Inf)), pars.term.dual_scale_threshold)
    elseif pars.term.dual_scale_mode == :exact
      return 1.0
    elseif pars.term.dual_scale_mode == :primal_dual
      return pars.term.dual_scale_threshold / max(sqrt(norm(get_y(iter), Inf) * norm(get_s(iter), Inf)), + pars.term.dual_scale_threshold)
    else
      throw("dual scale type does not exist")
    end
end

function scaled_dual_feas(iter::Class_iterate, pars::Class_parameters)
    return scaled_dual_feas(iter, iter.point.mu, pars)
end

function scaled_dual_feas(iter::Class_iterate, mu::Float64, pars::Class_parameters)
    return norm(eval_grad_lag(iter, mu),Inf) * dual_scale(iter, pars)
end

function check_for_nan(point::Class_point)
  if isbad(point.s)
    error("NaN in s")
  end

  if isbad(point.y)
    #@show point.y, point.s
    error("NaN in y")
  end

  if isbad(point.x)
    error("NaN in x")
  end

  if isbad(point.mu)
    error("NaN in mu")
  end
end

function is_feasible(it::Class_iterate, comp_feas::Float64)
    check_for_nan(it.point)

    mu = get_mu(it);
    s =  get_s(it)
    y =  get_y(it)

    Sy = s .* y
    if length(Sy) > 0
      return all(s .> 0.0) && all(y .> 0.0) && maximum(Sy) / mu  <= 1.0 / comp_feas && minimum(Sy) / mu >= comp_feas
    else
      return true
    end
end


## Pendulum
T = Float64

# model
model = TrajectoryOptimization.Dynamics.car
n = model.n; m = model.m
model_d = TO.discretize_model(model,:rk4)

# cost
Q = Array(1e-3*Diagonal(I,n))
R = Array(1e-3*Diagonal(I,m))
Qf = Array(Diagonal(I,n)*100.0)
x0 = zeros(n)
xf = [0.0;1.0;0.0]

# problem
N = 31
U = [ones(m) for k = 1:N-1]
obj = TrajectoryOptimization.LQRObjective(Q,R,Qf,xf,N)
dt = 0.15
prob = Problem(model_d,obj,U,dt=dt,x0=x0)
rollout!(prob)

## unconstrained
#bp
opts_ilqr = iLQRSolverOptions{T}()
ilqr_solver = iLQRSolver(prob,opts_ilqr)
jacobian!(prob,ilqr_solver)
TO.cost_expansion!(prob,ilqr_solver)
ΔV = TO._backwardpass!(prob,ilqr_solver)

#sqrt bp
opts_ilqr_sqrt = iLQRSolverOptions{T}(square_root=true)
ilqr_solver_sqrt = iLQRSolver(prob,opts_ilqr_sqrt)
jacobian!(prob,ilqr_solver_sqrt)
TO.cost_expansion!(prob,ilqr_solver_sqrt)
ΔV_sqrt = TO._backwardpass_sqrt!(prob,ilqr_solver_sqrt)
@test isapprox(ΔV_sqrt,ΔV)
@test all(isapprox(ilqr_solver.K,ilqr_solver_sqrt.K))
@test all(isapprox(ilqr_solver.d,ilqr_solver_sqrt.d))
for k = 1:N
    @test isapprox(ilqr_solver.S[k].xx,ilqr_solver_sqrt.S[k].xx'*ilqr_solver_sqrt.S[k].xx)
    @test isapprox(ilqr_solver.S[k].x,ilqr_solver_sqrt.S[k].x)
end

## constrained
opts_al = AugmentedLagrangianSolverOptions{T}(opts_uncon=opts_ilqr)
opts_al_sqrt = AugmentedLagrangianSolverOptions{T}(opts_uncon=opts_ilqr_sqrt)

# constraints
u_bnd = 5.
bnd = BoundConstraint(n,m,u_min=-u_bnd,u_max=u_bnd,trim=true)
bnd
goal_con = goal_constraint(xf)
con = [bnd, goal_con]
prob = update_problem(prob,constraints=Constraints(con,N))
rollout!(prob)

#bp
solver_ilqr = iLQRSolver(prob,opts_ilqr)
solver_al = AugmentedLagrangianSolver(prob,opts_al)
prob_al = TO.AugmentedLagrangianProblem(prob,solver_al)
TO.update_constraints!(prob_al.obj.C,prob_al.obj.constraints, prob.X, prob.U)
TO.update_active_set!(prob_al.obj)
jacobian!(prob_al,solver_ilqr)
TO.cost_expansion!(prob_al,solver_ilqr)
ΔV = TO._backwardpass!(prob_al,solver_ilqr)

#bp sqrt
solver_al_sqrt = AugmentedLagrangianSolver(prob,opts_al_sqrt)
solver_ilqr_sqrt = solver_al_sqrt.solver_uncon
prob_al_sqrt = TO.AugmentedLagrangianProblem(prob,solver_al_sqrt)
TO.update_constraints!(prob_al_sqrt.obj.C,prob_al_sqrt.obj.constraints, prob_al_sqrt.X, prob_al_sqrt.U)
TO.update_active_set!(prob_al_sqrt.obj)
jacobian!(prob_al_sqrt,solver_ilqr_sqrt)
TO.cost_expansion!(prob_al_sqrt,solver_ilqr_sqrt)
ΔV_sqrt = TO._backwardpass_sqrt!(prob_al_sqrt,solver_ilqr_sqrt)

@test isapprox(ΔV_sqrt,ΔV)
@test all(isapprox(solver_ilqr.K,solver_ilqr_sqrt.K))
@test all(isapprox(solver_ilqr.d,solver_ilqr_sqrt.d))
for k = 1:N
    @test isapprox(solver_ilqr.S[k].xx,solver_ilqr_sqrt.S[k].xx'*solver_ilqr_sqrt.S[k].xx)
    @test isapprox(solver_ilqr.S[k].x,solver_ilqr_sqrt.S[k].x)
end


"""
    BFGS

Quasi-Newton descent method.
"""
mutable struct BFGS{T<:AbstractFloat,
                    V<:AbstractVector{T},
                    M<:AbstractMatrix{T}} <: OptBuffer
    invH::M
    x::V
    g::V
    xpre::V
    gpre::V
    d::V
    xdiff::V
    gdiff::V
    y::T
    ypre::T
end

function BFGS(x::AbstractVector{T}) where {T}
    F = float(T)
    bfgs = BFGS(
        sqmatr(x, F),
        similar(x, F),
        similar(x, F),
        similar(x, F),
        similar(x, F),
        similar(x, F),
        similar(x, F),
        similar(x, F),
        F(NaN),
        F(NaN),
    )
    reset!(bfgs)
    return bfgs
end

function init!(
    optfn!, M::BFGS{T}, x0;
    reset, constrain_step = infstep
) where {T}
    optfn!(x0, zero(T), x0)
    if reset
        M.xpre, M.x = M.x, M.xpre
        M.gpre, M.g = M.g, M.gpre
        M.ypre = M.y
        map!(-, M.d, M.gpre)
        αmax = constrain_step(M.xpre, M.d)
        α = strong_backtracking!(
            optfn!, M.xpre, M.d, M.ypre, M.gpre;
            αmax=αmax, β=one(T)/100, σ=one(T)/10
        )
        map!(-, M.xdiff, M.x, M.xpre)
        map!(-, M.gdiff, M.g, M.gpre)

        invH = M.invH
        nr, nc = size(invH)
        scale = dot(M.xdiff, M.gdiff) / dot(M.gdiff, M.gdiff)
        fill!(invH, 0)
        for i in 1:nc
            elt = M.xdiff[i] / M.gdiff[i]
            invH[i, i] = 1e-5 < elt / scale < 1e5 ? elt : scale
        end
    end
    return M
end

@inline function reset!(M::BFGS)
    invH = M.invH
    nr, nc = size(invH)
    for j in 1:nc, i in 1:nr
        invH[i, j] = i == j
    end
    return M
end

function reset!(M::BFGS, x0, scale::Real=1)
    copy!(M.x, x0)
    invH = M.invH
    nr, nc = size(invH)
    for j in 1:nc, i in 1:nr
        invH[i, j] = (i == j) * scale
    end
    return M
end

@inline function callfn!(fdf, M::BFGS, x, α, d)
    __update_arg!(M, x, α, d)
    y, g = fdf(M.x, M.g)
    __update_grad!(M, g)
    M.y = y
    return y, g
end

function __descent_dir!(M::BFGS)
    mul!(M.d, M.invH, M.gpre, -1, 0)
    return M.d
end

function step!(optfn!::F, M::BFGS; constrain_step::S=infstep) where {F,S}
    #=
    argument and gradient from the end of the last
    iteration are stored into `xpre` and `gpre`
    =#
    M.gpre, M.g = M.g, M.gpre
    M.xpre, M.x = M.x, M.xpre
    M.ypre = M.y

    x, xpre, g, gpre, invH = M.x, M.xpre, M.g, M.gpre, M.invH
    d = __descent_dir!(M)
    maxstep = constrain_step(xpre, d)
    α = strong_backtracking!(optfn!, xpre, d, M.ypre, gpre, αmax = maxstep, β = 0.01, σ = 0.9)
    if α > 0
        #=
        BFGS update:
                δγ'B + Bγδ'   ⌈    γ'Bγ ⌉ δδ'
        B <- B - ---------- + |1 + -----| ---
                    δ'γ       ⌊     δ'γ ⌋ δ'γ
        =#
        δ, γ = M.xdiff, M.gdiff
        map!(-, γ, g, gpre)
        map!(-, δ, x, xpre)
        denom = dot(δ, γ)
        δscale = 1 + dot(γ, invH, γ) / denom
        # d <- B * γ
        mul!(d, invH, γ, 1, 0)
        invH .= invH .- (δ .* d' .+ d .* δ') ./ denom .+ δscale .* δ .* δ' ./ denom
    else
        fill!(M.xdiff, 0)
    end
    return α
end

@inline function __update_arg!(M::BFGS, x, α, d)
    map!(M.x, d, x) do a, b
        muladd(α, a, b)
    end
    return M.x
end

@inline function __update_arg!(M::BFGS, x)
    if x !== M.x
        copy!(M.x, x)
    end
    return M.x
end

@inline function __update_grad!(M::BFGS, g)
    if M.g !== g
        copy!(M.g, g)
    end
    return M.g
end


using LightGraphs, SimpleWeightedGraphs, GraphPlot

g = SimpleWeightedGraph(3)

add_edge!(g, 1, 2, 0.5)
add_edge!(g, 1, 3, 2.5)
add_edge!(g, 2, 3, 0.8)

gplot(g, nodelabel=vertices(g), edgelabel=weight.(collect(edges(g))))


using Hungarian
using Test
using LinearAlgebra
using DelimitedFiles

@testset "simple examples" begin
    A = [ 0.891171  0.0320582   0.564188  0.8999    0.620615;
          0.166402  0.861136    0.201398  0.911772  0.0796335;
          0.77272   0.782759    0.905982  0.800239  0.297333;
          0.561423  0.170607    0.615941  0.960503  0.981906;
          0.748248  0.00799335  0.554215  0.745299  0.42637]

    assign, cost = @inferred hungarian(A)
    @test assign == [2, 3, 5, 1, 4]

    B = [ 24     1     8;
           5     7    14;
           6    13    20;
          12    19    21;
          18    25     2]

    assign, cost = hungarian(B)
    @test assign == [2, 1, 0, 0, 3]
    @test cost == 8

    assign, cost = hungarian(B')
    @test assign == [2, 1, 5]
    @test cost == 8

    assign, cost = hungarian(ones(5,5) - I)
    @test assign == [1, 2, 3, 4, 5]
    @test cost == 0
end

@testset "forbidden edges" begin
    # result checked against Python package munkres: https://github.com/bmc/munkres/blob/master/munkres.py
    # Python code:
    #   m = Munkres()
    #   matrix = [[DISALLOWED, 1, 1], [1, 0, 1], [1, 1, 0]]
    #   m.compute(matrix)
    # Result: [(0, 1), (1, 0), (2, 2)]
    A = Union{Int, Missing}[missing 1 1; 1 0 1; 1 1 0]
    assign, cost = hungarian(A)
    @test assign == [2, 1, 3]
    @test cost == 2
end

@testset "issue #6" begin
    # load the cost matrix produced from the `solve_LSAP` function contained in the
    # [clue package](https://cran.r-project.org/web/packages/clue/index.html) in R.
    costMatrix = readdlm(joinpath(@__DIR__, "R-clue-cost.txt"))
    assignmentH, costH = hungarian(copy(transpose(costMatrix)))
    assignmentR = vec(readdlm(joinpath(@__DIR__, "R-clue-solution.txt"), Int64))
    costR = sum(costMatrix[ii, assignmentR[ii]] for ii in 1:size(costMatrix,1))
    @test costH ≤ costR
end


module IntervalStatistic

include("Mean.jl")
include("Variance.jl")
include("Check.jl")

function mean(values :: Vector, method :: Mean.BaseEstimator)
    Mean.estimate(values, method)
end

function var(values :: Vector, method :: Variance.BaseEstimator)
    Variance.estimate(values, method)
end

function isDistribution(values :: Vector, method :: Check.BaseCheck)
    Check.isDistribution(values, method)
end

end


function test_cashflow_irr()
    daycount = ActualActualISDA()
    evaluationdate = Date(2014, 6, 4)
    payments = [
        Payment(evaluationdate, -996.90)
        Payment(Date(2014, 12, 29), 531.5)
        Payment(Date(2015, 12, 29), 15.75)
        Payment(Date(2016, 12, 29), 15.75)
        Payment(Date(2017, 12, 29), 15.75)
        Payment(Date(2017, 12, 29), 500.0)
    ]
    cf = CashFlow(payments, daycount, evaluationdate)
    @test irr(cf) ≈ 0.040142375624945

    daycount = ActualActualISDA()
    evaluationdate = Date(2014, 12, 5)
    payments = [
        Payment(evaluationdate, -1007.572602739726)
        Payment(Date(2015, 12, 3), 31.5)
        Payment(Date(2016, 12, 3), 31.5)
        Payment(Date(2017, 12, 4), 31.5)
        Payment(Date(2018, 12, 3), 31.5)
        Payment(Date(2018, 12, 3), 1000.0)
    ]
    cf = CashFlow(payments, daycount, evaluationdate)
    @test irr(cf) ≈ 0.029507633486583
end

function test_cashflow_all()
    print(rpad("Finance.CashFlow...", 50, ' '))

    test_cashflow_irr()

    println("PASS")
end


using Plots, Printf, LinearAlgebra

const do_visu = parse(Bool, ENV["DO_VIZ"])
const do_save = parse(Bool, ENV["DO_SAVE"])
const nx = parse(Int, ENV["NX"])
const ny = parse(Int, ENV["NY"])

# enable plotting by default
# if !@isdefined do_visu; do_visu = false end

# macros to avoid array allocation
macro qHx(ix,iy)  esc(:( -(0.5*(H[$ix,$iy+1]+H[$ix+1,$iy+1]))*(0.5*(H[$ix,$iy+1]+H[$ix+1,$iy+1]))*(0.5*(H[$ix,$iy+1]+H[$ix+1,$iy+1])) * (H[$ix+1,$iy+1]-H[$ix,$iy+1])*_dx )) end
macro qHy(ix,iy)  esc(:( -(0.5*(H[$ix+1,$iy]+H[$ix+1,$iy+1]))*(0.5*(H[$ix+1,$iy]+H[$ix+1,$iy+1]))*(0.5*(H[$ix+1,$iy]+H[$ix+1,$iy+1])) * (H[$ix+1,$iy+1]-H[$ix+1,$iy])*_dy )) end
macro dtau(ix,iy) esc(:(  (1.0/(min_dxy2 / (H[$ix+1,$iy+1]*H[$ix+1,$iy+1]*H[$ix+1,$iy+1]) / 4.1) + _dt)^-1  )) end

function compute_update!(H2, dHdtau, H, Hold, _dt, damp, min_dxy2, _dx, _dy)
    Threads.@threads for iy=1:size(dHdtau,2)
    # for iy=1:size(dHdtau,2)
        for ix=1:size(dHdtau,1)
            dHdtau[ix,iy] = -(H[ix+1, iy+1] - Hold[ix+1, iy+1])*_dt +
                             (-(@qHx(ix+1,iy)-@qHx(ix,iy))*_dx -(@qHy(ix,iy+1)-@qHy(ix,iy))*_dy) +
                             damp*dHdtau[ix,iy]                        # damped rate of change
            H2[ix+1,iy+1] = H[ix+1,iy+1] + @dtau(ix,iy)*dHdtau[ix,iy]  # update rule, sets the BC as H[1]=H[end]=0
        end
    end
    return
end

function compute_residual!(ResH, H, Hold, _dt, _dx, _dy)
    Threads.@threads for iy=1:size(ResH,2)
    # for iy=1:size(ResH,2)
        for ix=1:size(ResH,1)
            ResH[ix,iy] = -(H[ix+1, iy+1] - Hold[ix+1, iy+1])*_dt +
                           (-(@qHx(ix+1,iy)-@qHx(ix,iy))*_dx -(@qHy(ix,iy+1)-@qHy(ix,iy))*_dy)
        end
    end
    return
end

function assign!(Hold, H)
    Threads.@threads for iy=1:size(H,2)
    # for iy=1:size(H,2)
        for ix=1:size(H,1)
            Hold[ix,iy] = H[ix,iy]
        end
    end
    return
end

@views function diffusion_2D_damp_perf_loop_fun(; do_visu=true, save_fig=false)
    # Physics
    lx, ly = 10.0, 10.0   # domain size
    ttot   = 1.0          # total simulation time
    dt     = 0.2          # physical time step
    # Numerics
    # nx, ny = 512, 512     # numerical grid resolution
    nout   = 100          # check error every nout
    tol    = 1e-6         # tolerance
    itMax  = 1e5          # max number of iterations
    damp   = 1-35/nx      # damping (this is a tuning parameter, dependent on e.g. grid resolution)
    # Derived numerics
    dx, dy = lx/nx, ly/ny # grid size
    xc, yc = LinRange(dx/2, lx-dx/2, nx), LinRange(dy/2, ly-dy/2, ny)
    # Array allocation
    ResH   = zeros(nx-2, ny-2) # normal grid, without boundary points
    dHdtau = zeros(nx-2, ny-2) # normal grid, without boundary points
    # Initial condition
    H      = exp.(.-(xc.-lx/2).^2 .-(yc.-ly/2)'.^2)
    Hold   = copy(H)
    H2     = copy(H)
    _dx, _dy, _dt = 1.0/dx, 1.0/dy, 1.0/dt
    min_dxy2 = min(dx,dy)^2
    t = 0.0; it = 0; ittot = 0; t_tic = 0.0; niter = 0
    # Physical time loop
    while t<ttot
        iter = 0; err = 2*tol
        # Picard-type iteration
        while err>tol && iter<itMax
            if (it==1 && iter==0) t_tic = Base.time(); niter = 0 end
            compute_update!(H2, dHdtau, H, Hold, _dt, damp, min_dxy2, _dx, _dy)
            H, H2 = H2, H  # pointer swap
            if iter % nout == 0
                compute_residual!(ResH, H, Hold, _dt, _dx, _dy)
                err = norm(ResH)/length(ResH)
            end
            iter += 1; niter += 1
        end
        ittot += iter; it += 1; t += dt
        assign!(Hold, H)
    end
    t_toc = Base.time() - t_tic
    A_eff = (2*2+1)/1e9*nx*ny*sizeof(Float64)  # Effective main memory access per iteration [GB]
    t_it  = t_toc/niter                        # Execution time per iteration [s]
    T_eff = A_eff/t_it                         # Effective memory throughput [GB/s]
    @printf("Time = %1.3f sec, T_eff = %1.2f GB/s (niter = %d)\n", t_toc, round(T_eff, sigdigits=2), niter)
    # Visualize
    if do_visu
        fontsize = 12
        opts = (aspect_ratio=1, yaxis=font(fontsize, "Courier"), xaxis=font(fontsize, "Courier"),
                ticks=nothing, framestyle=:box, titlefontsize=fontsize, titlefont="Courier", colorbar_title="",
                xlabel="Lx", ylabel="Ly", xlims=(xc[1],xc[end]), ylims=(yc[1],yc[end]), clims=(0.,1.))
        display(heatmap(xc, yc, H'; c=:davos, title="damped diffusion (nt=$it, iters=$ittot)", opts...))
        if save_fig savefig("diff2D_damp.png") end
    end
    if do_save open("./output/out_diffusion_2D_damp_perf_loop_fun.txt","a") do io; println(io, "$(nx) $(ny) $(t_toc) $(T_eff)") end end
    return
end

diffusion_2D_damp_perf_loop_fun(; do_visu=do_visu)


module mstraight_duct_H8_1
using FinEtools
using Compat.Test
function test()
    t0  =  time()

    rho = 1.21*phun("kg/m^3");# mass density
    c  = 343.0*phun("m/s");# sound speed
    bulk =  c^2*rho;
    omega =  54.5901*phun("rev/s")
    vn0 =  -1.0*phun("m/s")
    Lx = 10.0*phun("m");# length of the box, millimeters
    Ly = 1.0*phun("m"); # length of the box, millimeters
    n = 20;#number of elements along the length

    # println("""
    #
    # Straight duct with anechoic termination.
    # Example from Boundary element acoustics: Fundamentals and computer codes, TW Wu, page 44.
    # Both real and imaginary components of the pressure should have amplitude of
    # rho*c = $(rho*c).
    #
    # Hexahedral mesh.
    # """)

    fens,fes  =  H8block(Lx,Ly,Ly,n,2,2); # Mesh
    bfes  =  meshboundary(fes)
    L0 = selectelem(fens,bfes, facing = true, direction = [-1.0 0.0 0.0])
    L10 = selectelem(fens,bfes, facing = true, direction = [+1.0 0.0 0.0])
    nLx = selectnode(fens,box = [0.0 Lx  0.0 0.0 0.0 0.0], inflate = Lx/1.0e5)

    geom  =  NodalField(fens.xyz)
    P  =  NodalField(zeros(Complex128,size(fens.xyz,1),1))

    numberdofs!(P)

    material = MatAcoustFluid(bulk,rho)
    femm  =  FEMMAcoust(IntegData(fes, GaussRule(3, 2)), material)

    S  =  acousticstiffness(femm, geom, P);
    C  =  acousticmass(femm, geom, P);

    E10femm  =  FEMMAcoustSurf(IntegData(subset(bfes,L10),GaussRule(2, 2)), material)
    D  =  acousticABC(E10femm, geom, P);

    E0femm  =  FEMMBase(IntegData(subset(bfes,L0), GaussRule(2,  2)))
    fi  =  ForceIntensity(-1.0im*omega*rho*vn0);
    F  =  distribloads(E0femm, geom, P, fi, 2);

    p = (-omega^2*S +omega*1.0im*D + C)\F
    scattersysvec!(P, p[:])

    # println("Pressure amplitude bounds: ")
    # println("  real $(minimum(real(P.values)))/$(maximum(real(P.values)))")
    # println("  imag $(minimum(imag(P.values)))/$(maximum(imag(P.values)))")
    #
    # println("Total time elapsed  =  ",time() - t0,"s")

    File  =   "straight_duct.vtk"
    scalars = real(P.values);
    vtkexportmesh(File, connasarray(fes), geom.values, FinEtools.MeshExportModule.H8;
    scalars = [("Pressure", scalars)])
    try rm(File) catch end
    # @async run(`"paraview.exe" $File`)

    ref=rho*c
    @test abs(minimum(real(P.values))-(-ref))/ref < 0.02
    @test abs(minimum(imag(P.values))-(-ref))/ref < 0.02
    @test abs(maximum(real(P.values))-(+ref))/ref < 0.02
    @test abs(maximum(imag(P.values))-(+ref))/ref < 0.02

    # plotly()
    # ix = sortperm(geom.values[nLx,1])
    # plot(geom.values[nLx,1][ix], real(P.values)[nLx][ix], color = :blue, label = "real")
    # plot!(geom.values[nLx,1][ix], imag(P.values)[nLx][ix], color = :red, label  =  "imag")
    # plot!(title = "Straight duct with anechoic termination",
    # xlabel = "x", ylabel = "Pressure")
    # gui()

    true

end
end
using .mstraight_duct_H8_1
mstraight_duct_H8_1.test()


function sum_2(x::A) where {T, N, A<:AbstractArray{T,N}}
    n = length(x)
    if n < 2
        iszero(n) && return zero(T)
        isone(n)  && return x[1]
    end
    return summation_(n, x)
end

@inline sum_(x::A) where {T, N, A<:AbstractArray{T,N}} = sum_2(x)

function summation_(n::Int, x::A) where {T, N, A<:AbstractArray{T,N}}
   hi = x[1]
   lo = zero(T)
   for i in 2:n
       hi, low = add_(hi, x[i])
       lo += low
   end
   hi, lo = add_hilo_(hi, lo)
   return hi, lo
end


######################################################################
# pto_linear.jl
#
# like pto, but for linear systems. Should be faster
######################################################################

export pto_linear

function pto_linear(em::ErgodicManager, tm::TrajectoryManager; verbose::Bool=true, logging::Bool=false, max_iters::Int=100, es_crit::Float64=0.0, dd_crit::Float64=1e-6)
	xd0, ud0 = initialize(tm.initializer, em, tm)
	pto_linear(em, tm, xd0, ud0; verbose=verbose, logging=logging, max_iters=max_iters, es_crit=es_crit, dd_crit = dd_crit)
end

function pto_linear(em::ErgodicManager, tm::TrajectoryManager, xd0::VVF, ud0::VVF; verbose::Bool=true, logging::Bool=false, max_iters::Int=100, es_crit::Float64=0.0, dd_crit::Float64=1e-6)

	# let's not overwrite the initial trajectories
	xd = deepcopy(xd0)
	ud = deepcopy(ud0)
	N = tm.N

	# matrices for gradients
	ad = zeros(tm.dynamics.n, N+1)
	bd = zeros(tm.dynamics.m, N)

	# prepare for logging if need be
	if logging
		outfile = open("temp.csv", "w")
		save(outfile, xd)
	end

	if verbose; print_header(); end
	i = 1
	not_finished = true
	es = 0.; cs = 0.; ts = 0.; dd = 0.; step_size = 0.; es_diff=0.0
	es_prev = 0.0
	es_count = 1
	A, B = linearize(tm.dynamics, xd, ud, tm.h)
	K, Ginv = LQ2(A, B, tm.Qn, tm.Rn, tm.N)
	while not_finished

		# determine gradients used in optimization
		gradients!(ad, bd, em, tm, xd, ud)

		# find gains K and descent direction (zd, vd) using LQ
		#A, B = linearize(tm.dynamics, xd, ud, tm.h)
		C = LQ(K, Ginv, A, B, ad, bd, tm.Qn, tm.Rn, tm.N)
		zd, vd = apply_LQ_gains(A, B, K, C)

		# determine step size and descend
		step_size = get_step_size2(tm.descender, em, tm, xd, ud, zd, vd, ad, bd, K, i)

		# descend and project
		#xd, ud = project(em, tm, K, xd, ud, zd, vd, step_size)
		xd, ud = project2(em, tm, K, xd, ud, zd, vd, step_size)

		# compute statistics and report
		es, cs, ts = all_scores(em, tm, xd, ud)
		dd = directional_derivative(ad, bd, zd, vd)
		if verbose; step_report(i, es, cs, ts, dd, step_size); end
		if logging; save(outfile, xd); end

		# check convergence
		i += 1
		es_count = abs(es - es_prev) < 1e-7 ? es_count + 1 : 0
		es_prev = es
		not_finished = check_convergence(es,es_crit,i,max_iters,dd,dd_crit,verbose, es_count)
	end

	# now that we are done, print a special finished report
	if verbose
		print_header()
		if verbose; step_report(i-1, es, cs, ts, dd, step_size); end
	end

	if logging; close(outfile); end

	return xd, ud
end


function macWithCarry(a::ArbDigit, b::ArbDigit, c::ArbDigit, carry::DoubleArbDigit)
    @assert carry < (one(DoubleArbDigit) << BITS)
    carry += DoubleArbDigit(a) + DoubleArbDigit(b) * DoubleArbDigit(c)
    return carry % ArbDigit, carry >> BITS
end

function mulWithCarry(a::ArbDigit, b::ArbDigit, cin::DoubleArbDigit)
    cout = DoubleArbDigit(a) * DoubleArbDigit(b) + cin
    return cout % ArbDigit, cout >> BITS
end

function macDigit!(acc::AbstractVector{ArbDigit}, b::AbstractVector{ArbDigit}, c::ArbDigit)
    iszero(c) && return
    # we need to make sure that we have enough space
    # meaning length(b) +1 for carry
    @assert length(acc) > length(b) "accumulator needs extra space for carry!"

    carry = zero(DoubleArbDigit)
    a_lo = @view acc[begin:length(b)]
    a_hi = @view acc[length(b)+1:end]

    # we want to add `b*c+carry` into `acc`,
    # so we add elementwise into the lower
    # length(b) digits of `acc`.
    # whatever is left over is our true carry.
    for (a,(i,b)) in zip(a_lo, enumerate(b))
        a_lo[i], carry = macWithCarry(a, b, c, carry)
    end

    carry_hi, carry_lo = fromDoubleArbDigit(carry)

    # if we have carry from multiplying, we add it here
    final_carry = if iszero(carry_hi)
        _add2!(a_hi, carry_lo)
    else
        _add2!(a_hi, [carry_hi, carry_lo])
    end

    # there should never be anything left over
    # i.e. if this assert fires, the one above
    # should also have fired already
    @assert iszero(final_carry) "carry overflow during multiplication!"
end

function long_mul!(acc::AbstractVector{ArbDigit}, b::AbstractVector{ArbDigit}, c::AbstractVector{ArbDigit})
    # naive long multiplication
    # multiply each digit of `c` with `b` and
    # write to `acc`
    for (i,ci) in enumerate(c)
        macDigit!(@view(acc[i:end]), b, ci)
    end
end

function karatsuba!(acc::AbstractVector{ArbDigit}, a::AbstractVector{ArbDigit}, b::AbstractVector{ArbDigit})
    #= karatsuba multiplication
      exponents are just for numbering, not power function
      power is denoted by ^

      write x*y in two parts:
      a = a¹ * sh + a²
      b = b¹ * sh + b²

      now this becomes
      a*b = (a¹ * sh + a²) * (b¹ * sh + b²)
          =  a¹ * b¹            * sh^2
          + (a² * b¹ + a¹ * b²) * sh
          +  a² * b²

      now set p¹ = a¹ * b¹ and p³ = a² * b²:
      a*b =  p¹                 * sh^2
          + (a² * b¹ + a¹ * b²) * sh
          +  p³

      here comes karatsuba's trick:
              a¹ * b² + a² * b¹
          =   a¹ * b² + a² * b¹ - p¹ + p¹ - p³ + p³
          =   a¹ * b² + a² * b¹ - a¹ * b¹ - a² * b²  + p¹ + p³
          = -(a¹ * b¹ - a¹ * b² - a² * b¹ + a² * b²) + p¹ + p³
          = -((a² - a¹) * (b² - b¹))                 + p¹ + p³

      so we now have p² = (a² - a¹) * (b² - b¹) and thus
      a*b = p¹             * sh^2
          + (p¹ + p³ - p²) * sh
          + p³

      so our intermediate products are
      p¹  = a¹ * b¹
      p²  = (a² - a¹) * (b² - b¹)
      p³  = a² * b²

      which we can use to rearrange the above like so
      a*b = p¹ * sh^2
          + p¹ * sh
          + p³ * sh
          - p² * sh
          + p³
          = p³
          + p³ * sh
          + p¹ * sh
          + p¹ * sh^2
          - p² * sh
      with evaluation order from top to bottom. This means,
      by offsetting our index instead of really multiplying
      by `sh`, we can only calculate p¹, p² and p³ once each
      and reuse the result when adding. Subtracting p²*sh at
      the end is done to ensure we don't go negative. An
      informal argument for that is that we're multiplying
      positive integers - that can't ever give a negative
      result and since all transformations above are valid,
      the result has to be positive as well. However, to ensure
      we don't go negative during an intermediate step, we
      make sure to subtract as the last step.
    =#

    short, long = length(a) < length(b) ? (a,b) : (b,a)
    @assert length(short) <= length(long)
    @assert length(acc) >= length(short)+length(long)

    short_half_len = div(length(short), 2)
    short² = @view short[begin:short_half_len]
    short¹ = @view short[short_half_len+1:end]
    long²  = @view  long[begin:short_half_len]
    long¹  = @view  long[short_half_len+1:end]

    # uppers are longer, so we need at most that space
    p_len = length(short¹) + length(long¹)
    p = Vector{ArbDigit}(undef, p_len)
    p_num = ArbUInt(p)

    resize!(p, p_len)
    fill!(p, zero(ArbDigit))

    # p³ = a² * b²
    mac3!(p, short², long²)
    normalize!(p_num) # remove excess zeros
    # acc += p³
    add2!(acc, p)

    # acc += p³*sh
    add2!(@view(acc[short_half_len+1:end]), p)

    # zero buffer
    resize!(p, p_len)
    fill!(p, zero(ArbDigit))

    # p¹ = a¹ * b¹
    mac3!(p, short¹, long¹)
    normalize!(p_num) # remove excess zeros
    # acc += p¹*sh
    add2!(@view(acc[short_half_len+1:end]), p)

    # acc += p¹*sh^2
    add2!(@view(acc[2*short_half_len+1:end]), p)

    # now for p²
    # p²  = (a² - a¹) * (b² - b¹)
    # tmp¹  = a² - a¹
    # tmp²  = b² - b¹
    # this could be negative, so we have to take care
    sign¹, tmp¹ = sub_sign(short², short¹)
    sign², tmp² = sub_sign(long², long¹)

    if sign¹ * sign² > 0
        # both subtractions are positive, so we
        # have to _subtract_ from acc
        resize!(p, p_len)
        fill!(p, zero(ArbDigit))

        mac3!(p, tmp¹, tmp²)
        normalize!(p_num) # remove excess zeros

        # acc -= p²*sh
        sub2!(@view(acc[short_half_len+1:end]), p)

    elseif sign¹ * sign² < 0
        # one subtraction is negative, so we
        # have to _add_ to acc
        mac3!(@view(acc[short_half_len+1:end]), tmp¹, tmp²)

    else
        # one of the subtractions results in 0
        # we don't have to do anything!
    end

    # we're a mutating function - so return nothing
    return nothing
end

function mac3!(acc::AbstractVector{ArbDigit}, b::AbstractVector{ArbDigit}, c::AbstractVector{ArbDigit})
    if !iszero(length(b)) && iszero(first(b))
        nz = findfirst(!iszero, b)
        nz === nothing && return # only zeros!
        b = @view b[nz:end]
        acc = @view acc[nz:end]
    end
    if !iszero(length(c)) && iszero(first(c))
        nz = findfirst(!iszero, c)
        nz === nothing && return # only zeros!
        c = @view c[nz:end]
        acc = @view acc[nz:end]
    end

    short, long = length(b) < length(c) ? (b,c) : (c,b)

    if length(short) <= 32
        long_mul!(acc, long, short)
    else #if length(short) <= 256
        karatsuba!(acc, long, short)
    #else
        # TODO: implement Toom-3  multiplication for numbers larger than 256 digits
        #throw(ArgumentError("Multiplication for numbers with more than 256 digits has not been implemented yet."))
    end
end

function sub_sign(a::AbstractVector{ArbDigit}, b::AbstractVector{ArbDigit})
    if !isempty(a) && iszero(last(a))
        idx = something(findlast(!iszero, a), 0)
        a = @view a[begin:idx]
    end
    if !isempty(b) && iszero(last(b))
        idx = something(findlast(!iszero, b), 0)
        b = @view b[begin:idx]
    end

    ret = if a == b
        (0, ArbDigit[])
    elseif _less(<, a, b)
        retArr = copy(b)
        sub2!(retArr, a)
        (-1, retArr)
    else
        retArr = copy(a)
        sub2!(retArr, b)
        (1, retArr)
    end
    return ret
end

function mul3(x::AbstractVector{ArbDigit}, y::AbstractVector{ArbDigit})
    len = length(x) + length(y)
    prod = zeros(ArbDigit, len)

    mac3!(prod, x, y)
    ArbUInt(prod)
end

function scalar_mul!(a::ArbUInt, b::ArbDigit)
    if iszero(b)
        set_zero!(a)
    elseif isone(b)
        return
    else
        if ispow2(b)
            # this will break with too large numbers
            shl!(a, trailing_zeros(b))
        else
            carry = zero(DoubleArbDigit)
            for i in eachindex(a.data)
                a[i], carry = mulWithCarry(a[i], b, carry)
            end
            !iszero(carry) && push!(a.data, carry % ArbDigit)
        end
    end
end

Base.:(*)(a::ArbUInt, b::ArbUInt) = mul3(a.data, b.data)
Base.:(*)(a::ArbUInt, b::Unsigned) = scalar_mul!(deepcopy(a), ArbDigit(b))
Base.:(*)(b::Unsigned, a::ArbUInt) = scalar_mul!(deepcopy(a), ArbDigit(b))


#=================================
Methods using Taylor Models
=================================#
using .TaylorModels

@inline zeroBox(N) = IntervalBox(0..0, N)
@inline symBox(N) = IntervalBox(-1..1, N)

function _enclose(tme::TaylorModelsEnclosure, f::Function, dom::Interval_or_IntervalBox;
                  kwargs...)

    if tme.normalize
        R = _enclose_TaylorModels_norm(f, dom, tme.order)
    else
        R = _enclose_TaylorModels(f, dom, tme.order)
    end
    return R
end

# univariate
function _enclose_TaylorModels(f::Function, dom::Interval, order::Int)
    x0 = Interval(mid(dom))
    x = TaylorModel1(order, x0, dom)
    return evaluate(f(x-x0), dom)
end

# normalized univariate
function _enclose_TaylorModels_norm(f::Function, dom::Interval, order::Int)
    x0 = Interval(mid(dom))
    x = TaylorModel1(order, x0, dom)
    xnorm = normalize_taylor(x.pol, dom - x0, true)
    xnormTM = TaylorModel1(xnorm, 0..0, 0..0, -1..1)
    return evaluate(f(xnormTM), -1..1)
end

# multivariate
function _enclose_TaylorModels(f::Function, dom::IntervalBox{N}, order::Int) where {N}
    x0 = mid(dom)
    set_variables(Float64, "x", order=2order, numvars=N)
    x = [TaylorModelN(i, order, IntervalBox(x0), dom) for i=1:N]
    return evaluate(f(x...), dom - x0)
end

# normalized multivariate
function _enclose_TaylorModels_norm(f::Function, dom::IntervalBox{N}, order::Int) where {N}
    x0 = mid(dom)
    set_variables(Float64, "x", order=2order, numvars=N)

    zBoxN = zeroBox(N)
    sBoxN = symBox(N)
    x = [TaylorModelN(i, order, IntervalBox(x0), dom) for i=1:N]
    xnorm = [normalize_taylor(xi.pol, dom - x0, true) for xi in x]
    xnormTM = [TaylorModelN(xi_norm, 0..0, zBoxN, sBoxN) for xi_norm in xnorm]
    return evaluate(f(xnormTM...), sBoxN)
end