import Base: *, copyto!
import LinearAlgebra: mul!
using SparseArrays

abstract type AbstractDerivOperator{T,N} end

abstract type AbstractFiniteDiff{T,N} <: AbstractDerivOperator{T,N} end

struct PeriodicFD{T<:AbstractFloat,N,S} <: AbstractFiniteDiff{T,N}
    derivative_order        :: Int
    approximation_order     :: Int
    dx                      :: T
    len                     :: Int
    stencil_length          :: Int
    stencil_offset          :: Int
    stencil_coefs           :: S
end

struct FiniteDiffDeriv{T<:AbstractFloat,N,S1,S2,S3} <: AbstractFiniteDiff{T,N}
    derivative_order        :: Int
    approximation_order     :: Int
    dx                      :: T
    len                     :: Int
    stencil_length          :: Int
    stencil_offset          :: Int
    stencil_coefs           :: S1
    boundary_stencil_length :: Int
    # these are assumed to be ordered from the boundary inwards
    low_boundary_offsets    :: Vector{Int}
    low_boundary_coefs      :: S2
    high_boundary_offsets   :: Vector{Int}
    high_boundary_coefs     :: S3
end

struct SpectralDeriv{T<:Real,N,S} <: AbstractDerivOperator{T,N}
    derivative_order        :: Int
    len                     :: Int
    D                       :: S
end

struct CenteredDiff{N} end

function CenteredDiff{N}(derivative_order::Int,
                         approximation_order::Int, dx::T,
                         len::Int) where {T<:AbstractFloat,N}
    @assert approximation_order > 1 "approximation_order must be greater than 1."

    stencil_length = derivative_order + approximation_order - 1 +
        (derivative_order+approximation_order)%2

    stencil_offset = div(stencil_length-1,2)
    weights = calculate_weights(derivative_order, stencil_length-1, stencil_offset)

    stencil_coefs = (1/dx^derivative_order) .* weights

    PeriodicFD{T,N,typeof(stencil_coefs)}(derivative_order, approximation_order,
                                          dx, len, stencil_length,
                                          stencil_offset, stencil_coefs)
end

CenteredDiff(args...) = CenteredDiff{1}(args...)

struct EqualSizeStencilFD{N} end

"""
    EqualSizeStencilFD{N}(derivative_order, approximation_order, dx, len)

Finite difference operator where interior and boundary stencils have the same
length and same approximation order. Low (high) boundary operators will use as
many points to the left (right) as possible. For approximation_order > 2 there
is more than one boundary point, and each boundary point uses a different
stencil. Note that these are *not* truncation matched.
"""
function EqualSizeStencilFD{N}(derivative_order::Int,
                               approximation_order::Int, dx::T,
                               len::Int) where {T<:AbstractFloat,N}
    @assert approximation_order > 1 "approximation_order must be greater than 1."

    stencil_length = derivative_order + approximation_order - 1 +
        (derivative_order+approximation_order)%2
    stencil_offset = div(stencil_length-1,2)
    weights = calculate_weights(derivative_order, stencil_length-1, stencil_offset)

    stencil_coefs = (1/dx^derivative_order) .* weights

    boundary_stencil_length = stencil_length

    low_boundary_size  = div(approximation_order,2)
    high_boundary_size = div(approximation_order,2)

    # remember to order from the boundary inwards, which means the offset array
    # is increasing in the low boundary and decreasing in the high boundary
    low_boundary_offsets  = collect(0:low_boundary_size-1)
    high_boundary_offsets = collect(boundary_stencil_length-1:-1:high_boundary_size+1)

    low_weights = [calculate_weights(derivative_order, stencil_length-1, offset)
                   for offset in low_boundary_offsets]
    low_boundary_coefs = (1/dx^derivative_order) .* low_weights

    high_weights = [calculate_weights(derivative_order, stencil_length-1, offset)
                    for offset in high_boundary_offsets]
    high_boundary_coefs = (1/dx^derivative_order) .* high_weights

    FiniteDiffDeriv{T,N,typeof(stencil_coefs),typeof(low_boundary_coefs),
                    typeof(high_boundary_coefs)}(
                        derivative_order, approximation_order,
                        dx, len, stencil_length,
                        stencil_offset, stencil_coefs,
                        boundary_stencil_length, low_boundary_offsets,
                        low_boundary_coefs, high_boundary_offsets, high_boundary_coefs
                    )
end

EqualSizeStencilFD(args...) = EqualSizeStencilFD{1}(args...)

struct ChebDeriv{N} end

function ChebDeriv{N}(derivative_order::Int, xmin::T, xmax::T, len::Int) where {T<:Real,N}
    x, Dx, Dxx = cheb(xmin, xmax, len)

    if derivative_order == 1
        D = Dx
    elseif derivative_order == 2
        D = Dxx
    else
        error("derivative_order not implemented")
    end

    SpectralDeriv{T,N,typeof(D)}(derivative_order, len, D)
end

ChebDeriv(args...) = ChebDeriv{1}(args...)

struct FourierDeriv{N} end

function FourierDeriv{N}(derivative_order::Int, xsize::T, len::Int) where {T<:Real,N}
    Dx, Dxx = fourier(xsize, len)

    if derivative_order == 1
        D = Dx
    elseif derivative_order == 2
        D = Dxx
    else
        error("derivative_order not implemented")
    end

    SpectralDeriv{T,N,typeof(D)}(derivative_order, len, D)
end

FourierDeriv(args...) = FourierDeriv{1}(args...)

function Base.getindex(A::FiniteDiffDeriv{T}, i::Int, j::Int) where {T}
    s = A.stencil_offset + 1

    # interior points
    if s <= i <= (A.len-s+1)
        idx  = j - i + s
        coeffs = A.stencil_coefs
        ss     = s
        stencil_length = A.stencil_length
    elseif i < s
        coeffs = A.low_boundary_coefs[i]
        ss     = A.low_boundary_offsets[i] + 1
        idx    = j - i + ss
        stencil_length = A.boundary_stencil_length
    else # i > (A.len-s+1)
        ii     = A.len-i+1
        coeffs = A.high_boundary_coefs[ii]
        ss     = A.high_boundary_offsets[ii] + 1
        idx    = j - i + ss
        stencil_length = A.boundary_stencil_length
    end

    if idx < 1 || idx > stencil_length
        return zero(T)
    else
        return coeffs[idx]
    end
end

function Base.getindex(A::PeriodicFD{T}, i::Int, j::Int) where {T}
    N      = A.len
    coeffs = A.stencil_coefs
    s      = A.stencil_offset + 1

    if 1 <= j - i + s <= N
        idx  = j - i + s
    else
        # imposing periodicity
        idx  = mod1(j - i + s, N)
    end

    if idx < 1 || idx > A.stencil_length
        return zero(T)
    else
        return coeffs[idx]
    end
end

@inline Base.getindex(A::AbstractFiniteDiff, i::Int, ::Colon) =
    [A[i,j] for j in 1:A.len]

@inline Base.getindex(A::SpectralDeriv, i, j) = A.D[i,j]

# make it a callable struct, to compute derivatives at a given point
function (A::AbstractFiniteDiff{T,N})(f::AbstractArray{T,M},
                                      idx::Vararg{Int,M}) where {T<:Real,N,M}
    # make sure axis of differentiation is contained in the dimensions of f
    @assert N <= M

    s = A.stencil_offset + 1
    i = idx[N] # point where derivative will be taken (with respect to the N-axis)

    if s <= i <= (A.len-s+1)
        return _D_interior(A, f, idx)
    elseif i < s
        return _D_low_bdr(A, f, idx)
    else # i > (A.len-s+1)
        return _D_high_bdr(A, f, idx)
    end
end

# interior points
function _D_interior(A::AbstractFiniteDiff{T,N}, f::AbstractArray, idx) where {T<:Real,N}
    coeffs = A.stencil_coefs
    s      = A.stencil_offset + 1
    i      = idx[N] # point where derivative will be taken (with respect to the N-axis)

    sum_i  = zero(T)
    @fastmath @inbounds for aa in 1:A.stencil_length
        i_circ = i - (s - aa)
        I      = Base.setindex(idx, i_circ, N)
        sum_i += coeffs[aa] * f[I...]
    end

    sum_i
end

# boundary points
function _D_low_bdr(A::PeriodicFD{T,N}, f::AbstractArray, idx) where {T<:Real,N}
    coeffs = A.stencil_coefs
    s      = A.stencil_offset + 1
    i      = idx[N]

    sum_i  = zero(T)
    @fastmath @inbounds for aa in 1:A.stencil_length
        # imposing periodicity
        i_circ = mod1(i - (s-aa), A.len)
        I      = Base.setindex(idx, i_circ, N)
        sum_i += coeffs[aa] * f[I...]
    end

    sum_i
end
function _D_high_bdr(A::PeriodicFD{T,N}, f::AbstractArray, idx) where {T<:Real,N}
    coeffs = A.stencil_coefs
    s      = A.stencil_offset + 1
    i      = idx[N]

    sum_i  = zero(T)
    @fastmath @inbounds for aa in 1:A.stencil_length
        # imposing periodicity
        i_circ = mod1(i - (s-aa), A.len)
        I      = Base.setindex(idx, i_circ, N)
        sum_i += coeffs[aa] * f[I...]
    end

    sum_i
end

# low boundary points
function _D_low_bdr(A::FiniteDiffDeriv{T,N}, f::AbstractArray, idx) where {T<:Real,N}
    i      = idx[N]

    sum_i  = zero(T)
    @fastmath @inbounds for aa in 1:A.boundary_stencil_length
        coeffs = A.low_boundary_coefs[i]
        s      = A.low_boundary_offsets[i] + 1
        i_circ = i - (s - aa)
        I      = Base.setindex(idx, i_circ, N)
        sum_i += coeffs[aa] * f[I...]
    end

    sum_i
end

# high boundary points
function _D_high_bdr(A::FiniteDiffDeriv{T,N}, f::AbstractArray, idx) where {T<:Real,N}
    i      = idx[N]
    ii     = A.len-i+1

    sum_i  = zero(T)
    @fastmath @inbounds for aa in 1:A.boundary_stencil_length
        coeffs = A.high_boundary_coefs[ii]
        s      = A.high_boundary_offsets[ii] + 1
        i_circ = i - (s - aa)
        I      = Base.setindex(idx, i_circ, N)
        sum_i += coeffs[aa] * f[I...]
    end

    sum_i
end

function LinearAlgebra.mul!(df::AbstractVector, A::AbstractFiniteDiff, f::AbstractVector)
    @fastmath @inbounds for idx in eachindex(f)
        df[idx] = A(f,idx)
    end
    nothing
end

function LinearAlgebra.mul!(df::AbstractArray, A::AbstractFiniteDiff, f::AbstractArray)
    @fastmath @inbounds for idx in CartesianIndices(f)
        df[idx] = A(f,idx.I...)
    end
    nothing
end

# mul! done by standard matrix multiplication for Chebyshev differentiation
# matrices. This can also be done through FFT, but after some testing the FFT
# route actually seemed to be performing slower, so let's stick with this
function LinearAlgebra.mul!(df::AbstractVector, A::SpectralDeriv, f::AbstractVector)
    mul!(df, A.D, f)
end

# and now for Arrays
function LinearAlgebra.mul!(df::AbstractArray{T}, A::SpectralDeriv{T,N,S},
                            f::AbstractArray{T}) where {T,N,S}
    Rpre  = CartesianIndices(axes(f)[1:N-1])
    Rpost = CartesianIndices(axes(f)[N+1:end])

    _mul_loop!(df, A, f, Rpre, size(f,N), Rpost)
end

# this works as follows. suppose we are taking the derivative of a 4-dimensional
# array g, with coordinates w,x,y,z, and size
#
#   size(g) = (nw, nx, ny, nz)
#
# let's further suppose that we want the derivative along the x-direction (here,
# the second entry). this fixes
#
#   N = 2
#
# and then
#
#   Rpre  = CartesianIndices((nw,))
#   Rpost = CartesianIndices((ny,nz))
#
# now, the entry [Ipre,:,Ipost]
#
# slices *only* along the x-direction, for fixed Ipre and Ipost. Ipre and Ipost
# loop along Rpre and Rpost (respectively) without touching the x-direction. we
# further take care to loop in memory-order, since julia's arrays are faster in
# the first dimension.
#
# adapted from http://julialang.org/blog/2016/02/iteration
#
@noinline function _mul_loop!(df::AbstractArray{T}, A, f, Rpre, n, Rpost) where {T}
    @fastmath @inbounds for Ipost in Rpost
        @inbounds for i in 1:n
            @inbounds for Ipre in Rpre
                sum_i = zero(T)
                @inbounds @simd for j in 1:n
                    sum_i += A.D[i,j] * f[Ipre,j,Ipost]
                end
                df[Ipre,i,Ipost] = sum_i
            end
        end
    end
    nothing
end

function *(A::AbstractDerivOperator, x::AbstractArray)
    y = similar(x)
    mul!(y, A, x)
    y
end

# make SpectralDeriv a callable struct, to compute derivatives only at a given point.
function (A::SpectralDeriv{T,N,S})(f::AbstractArray{T,M},
                                   idx::Vararg{Int,M}) where {T<:Real,N,M,S}
    # make sure axis of differentiation is contained in the dimensions of f
    @assert N <= M

    i  = idx[N] # point where derivative will be taken (with respect to the N-axis)

    sum_i = zero(T)
    @fastmath @inbounds for ii in 1:A.len
        I = Base.setindex(idx, ii, N)
        sum_i += A.D[i,ii] * f[I...]
    end
    sum_i
end

# now for cross-derivatives. we assume that A acts on the first and B on the
# second axis of the x Matrix.
function (A::AbstractFiniteDiff{T,N1})(B::AbstractFiniteDiff{T,N2},
                                       f::AbstractArray{T,M},
                                       idx::Vararg{Int,M}) where {T<:Real,N1,N2,M}
    NA   = A.len
    NB   = B.len
    sA   = A.stencil_offset + 1
    sB   = B.stencil_offset + 1

    # make sure axes of differentiation are contained in the dimensions of f
    @assert N1 < N2 <= M

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    if sA <= i <= NA-sA+1 && sB <= j <= NB-sB+1
        return _D_interior(A, B, f, idx)

    elseif sA <= i <= NA-sA+1 && j < sB
        return _D_intA_lowB(A, B, f, idx)

    elseif sA <= i <= NA-sA+1 && j > NB-sB+1
        return _D_intA_highB(A, B, f, idx)

    elseif i < sA && sB <= j <= NB-sB+1
        return _D_lowA_intB(A, B, f, idx)

    elseif i > NA-sA+1 && sB <= j <= NB-sB+1
        return _D_highA_intB(A, B, f, idx)

    elseif i < sA && j < sB
        return _D_lowA_lowB(A, B, f, idx)

    elseif i < sA && j > NB-sB+1
        return _D_lowA_highB(A, B, f, idx)

    elseif i > NA-sA+1 && j < sB
        return _D_highA_lowB(A, B, f, idx)

    elseif i > NA-sA+1 && j > NB-sB+1
        return _D_highA_highB(A, B, f, idx)

    else
        error("I shouldn't be here...")
    end

end

# interior points
function _D_interior(A::AbstractFiniteDiff{T,N1}, B::AbstractFiniteDiff{T,N2},
                     f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    NB   = B.len
    qA   = A.stencil_coefs
    qB   = B.stencil_coefs
    sA   = A.stencil_offset + 1
    sB   = B.stencil_offset + 1

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

# boundary points
function _D_intA_lowB(A::AbstractFiniteDiff{T,N1}, B::PeriodicFD{T,N2},
                      f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    NB   = B.len
    qA   = A.stencil_coefs
    qB   = B.stencil_coefs
    sA   = A.stencil_offset + 1
    sB   = B.stencil_offset + 1

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        # imposing periodicity
        j_circ = mod1(j - (sB-jj), NB)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_intA_highB(A::AbstractFiniteDiff{T,N1}, B::PeriodicFD{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    NB   = B.len
    qA   = A.stencil_coefs
    qB   = B.stencil_coefs
    sA   = A.stencil_offset + 1
    sB   = B.stencil_offset + 1

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        # imposing periodicity
        j_circ = mod1(j - (sB-jj), NB)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_lowA_intB(A::PeriodicFD{T,N1}, B::AbstractFiniteDiff{T,N2},
                      f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    NB   = B.len
    qA   = A.stencil_coefs
    qB   = B.stencil_coefs
    sA   = A.stencil_offset + 1
    sB   = B.stencil_offset + 1

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            # imposing periodicity
            i_circ = mod1(i - (sA-ii), NA)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_highA_intB(A::PeriodicFD{T,N1}, B::AbstractFiniteDiff{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    NB   = B.len
    qA   = A.stencil_coefs
    qB   = B.stencil_coefs
    sA   = A.stencil_offset + 1
    sB   = B.stencil_offset + 1

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            # imposing periodicity
            i_circ = mod1(i - (sA-ii), NA)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_lowA_lowB(A::PeriodicFD{T,N1}, B::PeriodicFD{T,N2},
                      f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    NB   = B.len
    qA   = A.stencil_coefs
    qB   = B.stencil_coefs
    sA   = A.stencil_offset + 1
    sB   = B.stencil_offset + 1

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        # imposing periodicity
        j_circ = mod1(j - (sB-jj), NB)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            # imposing periodicity
            i_circ = mod1(i - (sA-ii), NA)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_lowA_highB(A::PeriodicFD{T,N1}, B::PeriodicFD{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    NB   = B.len
    qA   = A.stencil_coefs
    qB   = B.stencil_coefs
    sA   = A.stencil_offset + 1
    sB   = B.stencil_offset + 1

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        # imposing periodicity
        j_circ = mod1(j - (sB-jj), NB)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            # imposing periodicity
            i_circ = mod1(i - (sA-ii), NA)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_highA_lowB(A::PeriodicFD{T,N1}, B::PeriodicFD{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    NB   = B.len
    qA   = A.stencil_coefs
    qB   = B.stencil_coefs
    sA   = A.stencil_offset + 1
    sB   = B.stencil_offset + 1

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        # imposing periodicity
        j_circ = mod1(j - (sB-jj), NB)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            # imposing periodicity
            i_circ = mod1(i - (sA-ii), NA)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_highA_highB(A::PeriodicFD{T,N1}, B::PeriodicFD{T,N2},
                        f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    NB   = B.len
    qA   = A.stencil_coefs
    qB   = B.stencil_coefs
    sA   = A.stencil_offset + 1
    sB   = B.stencil_offset + 1

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        # imposing periodicity
        j_circ = mod1(j - (sB-jj), NB)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            # imposing periodicity
            i_circ = mod1(i - (sA-ii), NA)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

# now for FiniteDiffDeriv structs

function _D_intA_lowB(A::AbstractFiniteDiff{T,N1}, B::FiniteDiffDeriv{T,N2},
                      f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    NB   = B.len
    qA   = A.stencil_coefs
    sA   = A.stencil_offset + 1

    # points where derivative will be taken (along their respective axes)
    i  = idx[N1]
    j  = idx[N2]

    qB  = B.low_boundary_coefs[j]
    sB  = B.low_boundary_offsets[j] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.boundary_stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

function _D_intA_highB(A::AbstractFiniteDiff{T,N1}, B::FiniteDiffDeriv{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NB   = B.len
    qA   = A.stencil_coefs
    sA   = A.stencil_offset + 1

    i     = idx[N1]
    j     = idx[N2]
    jhigh = NB-j+1

    qB    = B.high_boundary_coefs[jhigh]
    sB    = B.high_boundary_offsets[jhigh] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.boundary_stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

function _D_lowA_intB(A::FiniteDiffDeriv{T,N1}, B::AbstractFiniteDiff{T,N2},
                      f::AbstractArray, idx) where {T<:Real,N1,N2}
    qB = B.stencil_coefs
    sB = B.stencil_offset + 1

    i  = idx[N1]
    j  = idx[N2]

    qA = A.low_boundary_coefs[i]
    sA = A.low_boundary_offsets[i] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.boundary_stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

function _D_highA_intB(A::FiniteDiffDeriv{T,N1}, B::AbstractFiniteDiff{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA   = A.len
    qB   = B.stencil_coefs
    sB   = B.stencil_offset + 1

    i     = idx[N1]
    j     = idx[N2]
    ihigh = NA-i+1

    qA    = A.high_boundary_coefs[ihigh]
    sA    = A.high_boundary_offsets[ihigh] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.boundary_stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

function _D_lowA_lowB(A::FiniteDiffDeriv{T,N1}, B::FiniteDiffDeriv{T,N2},
                      f::AbstractArray, idx) where {T<:Real,N1,N2}
    i  = idx[N1]
    j  = idx[N2]

    qA = A.low_boundary_coefs[i]
    sA = A.low_boundary_offsets[i] + 1
    qB = B.low_boundary_coefs[j]
    sB = B.low_boundary_offsets[j] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.boundary_stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.boundary_stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

function _D_lowA_highB(A::FiniteDiffDeriv{T,N1}, B::FiniteDiffDeriv{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NB = B.len

    i     = idx[N1]
    j     = idx[N2]
    jhigh = NB-j+1

    qA    = A.low_boundary_coefs[i]
    sA    = A.low_boundary_offsets[i] + 1
    qB    = B.high_boundary_coefs[jhigh]
    sB    = B.high_boundary_offsets[jhigh] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.boundary_stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.boundary_stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

function _D_highA_lowB(A::FiniteDiffDeriv{T,N1}, B::FiniteDiffDeriv{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA = A.len

    i     = idx[N1]
    j     = idx[N2]
    ihigh = NA-i+1

    qA    = A.high_boundary_coefs[ihigh]
    sA    = A.high_boundary_offsets[ihigh] + 1
    qB    = B.low_boundary_coefs[j]
    sB    = B.low_boundary_offsets[j] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.boundary_stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.boundary_stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

function _D_highA_highB(A::FiniteDiffDeriv{T,N1}, B::FiniteDiffDeriv{T,N2},
                        f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA = A.len
    NB = B.len

    i     = idx[N1]
    j     = idx[N2]
    ihigh = NA-i+1
    jhigh = NB-j+1

    qA    = A.high_boundary_coefs[ihigh]
    sA    = A.high_boundary_offsets[ihigh] + 1
    qB    = B.high_boundary_coefs[jhigh]
    sB    = B.high_boundary_offsets[jhigh] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.boundary_stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.boundary_stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

# cross derivatives with FiniteDiffDeriv and PeriodicFD

function _D_lowA_lowB(A::PeriodicFD{T,N1}, B::FiniteDiffDeriv{T,N2},
                      f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA = A.len
    qA = A.stencil_coefs
    sA = A.stencil_offset + 1

    i  = idx[N1]
    j  = idx[N2]

    qB = B.low_boundary_coefs[j]
    sB = B.low_boundary_offsets[j] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.boundary_stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            # imposing periodicity
            i_circ = mod1(i - (sA-ii), NA)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_lowA_lowB(A::FiniteDiffDeriv{T,N1}, B::PeriodicFD{T,N2},
                      f::AbstractArray, idx) where {T<:Real,N1,N2}
    NB = B.len
    qB = B.stencil_coefs
    sB = B.stencil_offset + 1

    i  = idx[N1]
    j  = idx[N2]

    qA = A.low_boundary_coefs[i]
    sA = A.low_boundary_offsets[i] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        # imposing periodicity
        j_circ = mod1(j - (sB-jj), NB)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.boundary_stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

function _D_lowA_highB(A::PeriodicFD{T,N1}, B::FiniteDiffDeriv{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA    = A.len
    qA    = A.stencil_coefs
    sA    = A.stencil_offset + 1

    i     = idx[N1]
    j     = idx[N2]
    jhigh = NB-j+1

    qB    = B.high_boundary_coefs[jhigh]
    sB    = B.high_boundary_offsets[jhigh] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.boundary_stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            # imposing periodicity
            i_circ = mod1(i - (sA-ii), NA)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_lowA_highB(A::FiniteDiffDeriv{T,N1}, B::PeriodicFD{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NB    = B.len
    qB    = B.stencil_coefs
    sB    = B.stencil_offset + 1

    i     = idx[N1]
    j     = idx[N2]
    jhigh = NB-j+1

    qA    = A.low_boundary_coefs[i]
    sA    = A.low_boundary_offsets[i] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        # imposing periodicity
        j_circ = mod1(j - (sB-jj), NB)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.boundary_stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

function _D_highA_lowB(A::PeriodicFD{T,N1}, B::FiniteDiffDeriv{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA    = A.len
    qA    = A.stencil_coefs
    sA    = A.stencil_offset + 1

    i     = idx[N1]
    j     = idx[N2]
    ihigh = NA-i+1

    qB    = B.low_boundary_coefs[j]
    sB    = B.low_boundary_offsets[j] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.boundary_stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            # imposing periodicity
            i_circ = mod1(i - (sA-ii), NA)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_highA_lowB(A::FiniteDiffDeriv{T,N1}, B::PeriodicFD{T,N2},
                       f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA    = A.len
    NB    = B.len
    qB    = B.stencil_coefs
    sB    = B.stencil_offset + 1

    i     = idx[N1]
    j     = idx[N2]
    ihigh = NA-i+1

    qA    = A.high_boundary_coefs[ihigh]
    sA    = A.high_boundary_offsets[ihigh] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        # imposing periodicity
        j_circ = mod1(j - (sB-jj), NB)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.boundary_stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

function _D_highA_highB(A::PeriodicFD{T,N1}, B::FiniteDiffDeriv{T,N2},
                        f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA    = A.len
    NB    = B.len

    qA    = A.stencil_coefs
    sA    = A.stencil_offset + 1

    i     = idx[N1]
    j     = idx[N2]
    ihigh = NA-i+1
    jhigh = NB-j+1

    qB    = B.high_boundary_coefs[jhigh]
    sB    = B.high_boundary_offsets[jhigh] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.boundary_stencil_length
        j_circ = j - (sB-jj)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.stencil_length
            # imposing periodicity
            i_circ = mod1(i - (sA-ii), NA)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end
function _D_highA_highB(A::FiniteDiffDeriv{T,N1}, B::PeriodicFD{T,N2},
                        f::AbstractArray, idx) where {T<:Real,N1,N2}
    NA    = A.len
    NB    = B.len
    qB    = B.stencil_coefs
    sB    = B.stencil_offset + 1

    i     = idx[N1]
    j     = idx[N2]
    ihigh = NA-i+1
    jhigh = NB-j+1

    qA    = A.high_boundary_coefs[ihigh]
    sA    = A.high_boundary_offsets[ihigh] + 1

    sum_ij = zero(T)
    @fastmath @inbounds for jj in 1:B.stencil_length
        # imposing periodicity
        j_circ = mod1(j - (sB-jj), NB)
        Itmp   = Base.setindex(idx, j_circ, N2)
        sum_i  = zero(T)
        @inbounds for ii in 1:A.boundary_stencil_length
            i_circ = i - (sA-ii)
            I      = Base.setindex(Itmp, i_circ, N1)
            sum_i += qA[ii] * qB[jj] * f[I...]
        end
        sum_ij += sum_i
    end

    sum_ij
end

# Casting to matrix types

copyto!(M::AbstractMatrix{T}, A::SpectralDeriv) where {T<:Real} =
    copyto!(M, A.D)

function copyto!(M::AbstractMatrix{T}, A::AbstractFiniteDiff) where {T<:Real}
    for idx in CartesianIndices(M)
        M[idx] = A[idx.I...]
    end
    M
end

LinearAlgebra.Array(A::AbstractDerivOperator{T}) where {T} =
    copyto!(zeros(T, A.len, A.len), A)

SparseArrays.SparseMatrixCSC(A::AbstractFiniteDiff{T}) where {T} =
    copyto!(spzeros(T, A.len, A.len), A)


myfun(x) = x^2 + x


@testset "radial basis function" begin
    xc = rand(10)
    yc = rand(10)
    e = rand(10)
    c = rand(10)
    d = rand(3)
    r = RBF2D(xc, yc; c=c, eps=e, d=d, kind = 0)

    x = rand(5)
    y = rand(5)
    o = r(x, y)

    O = zeros(5)
    for i = 1:5
        for j = 1:10
            dist = sqrt((x[i]-xc[j])^2 + (y[i]-yc[j])^2)
            O[i] += c[j] * exp(-(e[j]*dist)^2)
        end
        O[i] += d[1] + d[2] * x[i] + d[3] * y[i]
    end
    init(sess)
    @test norm(run(sess, o)-O)<1e-5
end

@testset "interp1" begin
    Random.seed!(233)
    x = sort(rand(10))
    y = @. x^2 + 1.0
    z = [x[1]; x[2]; rand(5) * (x[end]-x[1]) .+ x[1]; x[end]]
    u = interp1(x,y,z)
    @test norm(run(sess, u)-[1.026422850882909
                    1.044414684090653
                    1.312604319732756
                    1.810845361128137
                    1.280789421523103
                    1.600084940795178
                    1.930560200260898
                    1.972130181835701]) < 1e-10
end

@testset "rbf3d" begin

    function f0(r, e)
        return exp(-(e*r)^2)
    end

    function f1(r, e)
        return sqrt((1+(e*r)^2))
    end

    function f2(r, e)
        return 1/(1+(e*r)^2)
    end

    function f3(r, e)
        return 1/sqrt(1+(e*r)^2)
    end

    fs = [f0, f1, f2, f3]
    nc = 100
    n = 10
    xc = rand(nc)
    yc = rand(nc)
    zc = rand(nc)
    x = rand(n)
    y = rand(n)
    z = rand(n)
    c = rand(nc)
    d = rand(4)
    e = rand(nc)

    for kind = 1:4
        out = zeros(n)
        for i = 1:n
            r = @. sqrt((x[i] - xc)^2 + (y[i] - yc)^2 + (z[i] - zc)^2)
            out[i] = sum(c .* fs[kind].(r, e)) + d[1] + d[2] * x[i] + d[3] * y[i] + d[4] * z[i]
        end

        rbf = RBF3D(xc, yc, zc, c = c, eps = e, d = d, kind = kind - 1)
        o = rbf(x,y,z)

        o_ = run(sess, o)

        @test maximum(abs.(out .- o_)) < 1e-8
    end

    for kind = 1:4
        out = zeros(n)
        for i = 1:n
            r = @. sqrt((x[i] - xc)^2 + (y[i] - yc)^2 + (z[i] - zc)^2)
            out[i] = sum(c .* fs[kind].(r, e)) + d[1]
        end

        rbf = RBF3D(xc, yc, zc, c = c, eps = e, d = d[1:1], kind = kind - 1)
        o = rbf(x,y,z)

        o_ = run(sess, o)

        @test maximum(abs.(out .- o_)) < 1e-8
    end

    for kind = 1:4
        out = zeros(n)
        for i = 1:n
            r = @. sqrt((x[i] - xc)^2 + (y[i] - yc)^2 + (z[i] - zc)^2)
            out[i] = sum(c .* fs[kind].(r, e))
        end

        rbf = RBF3D(xc, yc, zc, c = c, eps = e, kind = kind - 1)
        o = rbf(x,y,z)

        o_ = run(sess, o)

        @test maximum(abs.(out .- o_)) < 1e-8
    end
end


using ConstraintSolver
using JSON

CS = ConstraintSolver
include("../../test/sudoku_fcts.jl")

function parseJSON(json_sums)
    sums = []
    for s in json_sums
        indices = Tuple[]
        for ind in s["indices"]
            push!(indices, tuple(ind...))
        end

        push!(sums, (result=s["result"], indices=indices))
    end
    return sums
end

function solve_all(filenames; benchmark=false, single_times=true)
    ct = time()
    for (i,filename) in enumerate(filenames)
        sums = parseJSON(JSON.parsefile("benchmark/killer_sudoku/data/$(filename)"))

        com = CS.init()
        grid = zeros(Int, (9,9))

        com_grid = create_sudoku_grid!(com, grid)

        for s in sums
            add_constraint!(com, CS.eq_sum, [com_grid[CartesianIndex(ind)] for ind in s.indices]; rhs=s.result)
        end

        add_sudoku_constr!(com, com_grid)

        if single_times
            GC.enable(false)
            t = time()
            status = solve!(com);
            t = time()-t
            GC.enable(true)
            println(i-1,", ", t)
        else
            GC.enable(false)
            status = solve!(com);
            GC.enable(true)
        end
        if !benchmark
            @show com.info
            println("Status: ", status)
            @assert fulfills_sudoku_constr(com_grid)
        end
    end
    println("")
    tt = time()-ct
    println("total time: ", tt)
    println("avg: ", tt/length(filenames))
end

function main(; benchmark=false, single_times=true)
    solve_all(["en_wikipedia" for i=1:100]; benchmark=benchmark, single_times=single_times)
    # solve_all(from_file("hardest.txt"), "hardest")
end


using Zygote, LinearAlgebra
using GenericLinearAlgebra: _eigvals!

function test_eigmin(A::AbstractMatrix)
    vals = _eigvals!(copy(A))
    ev = minimum(real.(vals))
    return ev
end

A = rand(4, 4); A = A + A';

@show eigmin(A)
@show test_eigmin(A)
@show Zygote.gradient(test_eigmin, A)


export TikhonovReg, getTikhonovReg, regularizer

"""
Tikhnov Regularizer R(x) = 0.5*alpha*(H*(x-xref), x-xref), with H spd
"""
mutable struct TikhonovReg{T}
    H
    alpha::T
    xref::Array{T}
end

getTikhonovReg(TYPE;H=I,alpha=1.0,xref=zeros(0))=TikhonovReg{TYPE}(H,convert(TYPE,alpha),convert.(TYPE,vec(xref)))

function regularizer(this::TikhonovReg,x)
    u = vec(copy(x))
    if !isempty(this.xref)
        u.-=this.xref
    end
    d2S = this.alpha*this.H
    dS = d2S*u
    S  = 0.5*dot(dS,u)
    return S,dS,d2S
end


"""
$(TYPEDSIGNATURES)

Generates the Liouville transformed set of ODEs, which is the original
ODE system with a new variable `trJ` appended, corresponding to the
-tr(Jacobian). This variable is used for properties like uncertainty
propagation from a given initial distribution density.

For example, if ``u'=p*u`` and `p` follows a probability distribution
``f(p)``, then the probability density of a future value with a given
choice of ``p`` is computed by setting the inital `trJ = f(p)`, and
the final value of `trJ` is the probability of ``u(t)``.

Example:

```julia
using ModelingToolkit, OrdinaryDiffEq, Test

@parameters t α β γ δ
@variables x(t) y(t)
D = Differential(t)

eqs = [D(x) ~ α*x - β*x*y,
       D(y) ~ -δ*y + γ*x*y]

sys = ODESystem(eqs)
sys2 = liouville_transform(sys)
@variables trJ

u0 = [x => 1.0,
      y => 1.0,
      trJ => 1.0]

prob = ODEProblem(sys2,u0,tspan,p)
sol = solve(prob,Tsit5())
```

Where `sol[3,:]` is the evolution of `trJ` over time.

Sources:

Probabilistic Robustness Analysis of F-16 Controller Performance: An
Optimal Transport Approach

Abhishek Halder, Kooktae Lee, and Raktim Bhattacharya
https://abhishekhalder.bitbucket.io/F16ACC2013Final.pdf
"""
function liouville_transform(sys::AbstractODESystem)
      t = get_iv(sys)
      @variables trJ
      D = ModelingToolkit.Differential(t)
      neweq = D(trJ) ~ trJ*-tr(calculate_jacobian(sys))
      neweqs = [equations(sys);neweq]
      vars = [states(sys);trJ]
      ODESystem(neweqs,t,vars,parameters(sys))
end


using ConservationLawsParticles
using SpecialFunctions, RecursiveArrayTools, DifferentialEquations, Plots
using Random

# reference solution
boxsol(t, x; a, b) = (erf((x-a)/2sqrt(t)) - erf((x-b)/2sqrt(t))) / (2*(b-a))
# refsol(t, x) = sum(boxsol(t, x; a=x0[i-1], b=x0[i]) for i in eachindex(x0)[2:end]) / (length(x0)-1)
# refsol(t, x) = 1 / (pi * (1 + x^2))

# model
V(t, x) = -2x / (1 + x^2) # 2 + 0.5sin(2x) + 0.3sin(2sqrt(2)*x) + 0.2cos(2sqrt(7)*x)
W(t, x) = 0.0 # x^2/2 - abs(x)
W′(t, x) = 0.0 # x - sign(x)
mob(ρ) = 1.0 # max((ρ-1)^2 * (2ρ+1), 0)
A(ρ) = ρ^2
simple_model = DiffusiveSampledModel((V,), ((W′,),), (mob,), (SimpleDiffusion(A),))
mean_model = DiffusiveSampledModel((V,), ((W′,),), (mob,), (MeanDiffusion(A),))
min_model = DiffusiveSampledModel((V,), ((W′,),), (mob,), (MinDiffusion(A),))

# initial condition
Random.seed!(42)
n = 60
x0 = ArrayPartition(sort(vcat(randn(n) .- 1, .5randn(n) .+ 1)))

# time span
tspan = (0., 15.)

# ODE system for the particles
simple_prob = ODEProblem(velocities_diff!, x0, tspan, simple_model)
mean_prob = ODEProblem(velocities_diff!, x0, tspan, mean_model)
min_prob = ODEProblem(velocities_diff!, x0, tspan, min_model)

# solve it
abstol = reltol = 2e-7
@time simple_sol = solve(simple_prob, BS5(); abstol=abstol, reltol=reltol);
@time mean_sol = solve(mean_prob, BS5(); abstol=abstol, reltol=reltol);
@time min_sol = solve(min_prob, BS5(); abstol=abstol, reltol=reltol);

# plot the particle trajectories
plot(title="Trajectories", xlabel="time", ylabel="position", legend=false, ylims=(-5,5))
plot!(simple_sol, vars=eachindex(x0)[1:3:end], color=1)
plot!(mean_sol, vars=eachindex(x0)[1:3:end], color=2)
plot!(min_sol, vars=eachindex(x0)[1:3:end], color=3)
savefig("diffusion9.png")

# plot an animation of the density
anim = @animate for t in range(tspan...; length=100)
    plot(title="Density", xlabel="position", ylabel="density",
        xlims=(-2,2), ylims=(0,.6), legend=:topleft, size=(1200, 800))
    # plot!(x -> refsol(t, x); color=:black, label="analytic")
    densityplot!(simple_sol(t); color=1, label="simple")
    densityplot!(mean_sol(t); color=2, label="mean")
    densityplot!(min_sol(t); color=3, label="min")
end
gif(anim, "diffusion9.gif")


function eftAdd{T<:StdFloat}(a::T, b::T)
  hi = a + b
  t = hi - a
  lo = (a - (hi - t)) + (b - t)
  hi,lo
end

function eftAddGTE{T<:StdFloat}(a::T, b::T)
  hi = a + b
  lo = b - (hi - a)
  hi,lo
end

function eftAddLTE{T<:StdFloat}(a::T, b::T)
  hi = a + b
  lo = a - (hi - b)
  hi,lo
end

function eftSub{T<:StdFloat}(a::T, b::T)
  hi = a - b
  t = hi - a
  lo = (a - (hi - t)) - (b + t)
  hi,lo
end

function eftSubGTE{T<:StdFloat}(a::T, b::T)
  hi = a - b
  lo = (a - hi) - b
  hi,lo
end

function eftMul{T<:StdFloat}(a::T, b::T)
    hi = a * b
    lo = fma(a, b, -hi)
    hi,lo
end

function accDiv{T<:StdFloat}(a::T, b::T)
    hi = a/b
    remainder = fma(-b,hi,a)
    lo = remainder/b
    remainder = fma(-b,lo,remainder)
    lo += remainder/b
    hi,lo
end

# for use with directed rounding, when only the sign of the lo part is needed
function accDivForSign{T<:StdFloat}(a::T, b::T)
    hi = a/b
    loApprox = -fma(b,hi,-a) # sign is correct, value is actual*b
    hi,loApprox
end

for fn in (:eftAdd, :eftAddGTE, :eftSub, :eftSubGTE, :eftMul, :accDiv, :accDivForSign)
  @eval begin
    ($fn){T1<:Real,T2<:StdFloat}(a::T1, b::T2) = ($fn)(convert(Float64,a), b)
    ($fn){T1<:Real,T2<:StdFloat}(a::T2, b::T1) = ($fn)(a, convert(Float64,b))
    ($fn){T1<:Real,T2<:Real}(a::T1, b::T2) = ($fn)(convert(Float64,a), convert(Float64,b))
  end
end


function _integrate_spectrum(h::Float64, γ, λ, f) #interpreted as point mass spectrum
    denom = (λ / h + 1 / (1 + γ * f))
    1 / denom
end

#interpreted as
function _integrate_spectrum(h::Distribution, γ, λ, f)
    expectation(u->_integrate_spectrum(u, γ, λ, f), h)
end

function fixed_point_function(hs, γs, λs)
    γ = sum(γs)
    fixed_point_f = f -> f - sum(γs ./ γ .* _integrate_spectrum.(hs, γ, λs, f))
    find_zero(fixed_point_f, (0.0, 100.0))
end

function risk_formula(hs, γs, αs, λs)
    λs = min.(λs, 10_000) #hack for now until properly dealing with Infinity
    γ = sum(γs)
    fixed_pt = λs_tilde -> fixed_point_function(hs, γs, λs_tilde)
    f = fixed_pt(λs)
    ∇f = grad(central_fdm(5, 1), fixed_pt, λs)[1]
    #return ∇f
    #return γ ./ γs .* (γs .* λs - αs.^2 .* λs.^2) .* ∇f
    1 + γ * f + sum(γ ./ γs .* (γs .* λs - αs .^ 2 .* λs .^ 2) .* ∇f)
end

function r_squared(hs, γs, αs, λs)
    response_var = 1 + sum(abs2, αs)
    risk = risk_formula(hs, γs, αs, λs)
    1 - risk / response_var
end

function optimal_r_squared(αs, γs, hs)
    λs_opt = γs ./ αs .^ 2
    r_squared(hs, γs, αs, λs_opt)
end

function optimal_λs(γs, αs)
    γs ./ αs .^ 2
end

function optimal_risk(hs, γs, αs)
    λs_opt = optimal_λs(γs, αs)
    risk_formula(hs, γs, αs, λs_opt)
end

function optimal_single_λ(γs, αs)
    λ_opt = sum(γs) / sum(abs2, αs)
    λs_opt = fill(λ_opt, 2)
end

function optimal_single_λ_risk(hs, γs, αs)
    λs_opt = optimal_single_λ(γs, αs)
    risk_formula(hs, γs, αs, λs_opt)
end

function optimal_ignore_second_group_λs(γs, αs)
    λ1_opt = γs[1] * (1 + αs[2]^2) / αs[1]^2
    λs_opt = [λ1_opt; Inf]
end

function optimal_ignore_second_group_risk(hs, γs, αs)
    λs_opt = optimal_ignore_second_group_λs(γs, αs)
    risk_formula(hs, γs, αs, λs_opt)
end


"""
    logit10(x::Real)
Function transforming interval [0,1] to [-Inf, Inf] using logit transformation.
"""
function logit10(x::Real)
    log10(x / (1.0 - x))
end

"""
    logistic10(x::Real)
Function transforming interval [-Inf, Inf] to [0,1] using logistic transformation.
Inverse function for [`logit10`](@ref).
"""
function logistic10(x::Real)
    exp10(x) / (exp10(x) + 1.0)
end

"""
    scaling(x::Real, scale::Symbol = :direct)
Transforms values from specific scale to range [-Inf, Inf] based on option.

## Return
Transformed value.

## Arguments
* `x`: input value.
* `scale`: transformation type: `:direct, :log, :logit`.
"""
function scaling(x::Real, scale::Symbol = :direct)
    if scale == :direct
        return x
    elseif scale == :log
        return log10(x)
    elseif scale == :logit
        return logit10(x)
    else
        throw(DomainError(scale, "scale type is not supported"))
    end
end

"""
    unscaling(x::Real, scale::Symbol = :direct)
Transforms values from [-Inf, Inf] to specific scale based on option. Inverse function
for [`scaling`](@ref).

## Return
Transformed value.

## Arguments
* `x`: input value.
* `scale`: transformation type: `:direct, :log, :logit`.
"""
function unscaling(x::Real, scale::Symbol = :direct)
    if scale == :direct
        return x
    elseif scale == :log
        return exp10(x)
    elseif scale == :logit
        return logistic10(x)
    else
        throw(DomainError(scale, "scale type is not supported"))
    end
end

scaling(x_tup::Tuple{Float64,Float64}, scale::Symbol = :direct) = scaling.(x_tup, scale)

unscaling(x_tup::Tuple{Float64,Float64}, scale::Symbol = :direct) = unscaling.(x_tup, scale)

scaling(::Nothing, ::Symbol) = nothing

unscaling(::Nothing, ::Symbol) = nothing

"""
    function get_endpoint(
        theta_init::Vector{Float64},
        theta_num::Int,
        loss_func::Function,
        method::Symbol,
        direction::Symbol = :right;

        loss_crit::Float64 = 0.0,
        scale::Vector{Symbol} = fill(:direct, length(theta_init)),
        theta_bounds::Vector{Tuple{Float64,Float64}} = unscaling.(
            fill((-Inf, Inf), length(theta_init)),
            scale
            ),
        scan_bound::Float64 = unscaling(
            (direction==:left) ? -9.0 : 9.0,
            scale[theta_num]
            ),
        scan_tol::Float64 = 1e-3,
        loss_tol::Float64 = 1e-3,
        local_alg::Symbol = :LN_NELDERMEAD,
        kwargs...
        )

Calculates confidence interval's right or left endpoints for a given parameter `theta_num`.

## Return
[`EndPoint`](@ref) object storing confidence interval's endpoint and intermediate profile points.

## Arguments
- `theta_init`: starting values of parameter vector ``\\theta``. The starting values should not necessary be the optimum values of `loss_func` but `loss_func(theta_init)` should be lower than `loss_crit`.
- `theta_num`: index of vector component for identification: `theta_init(theta_num)`.
- `loss_func`: loss function ``\\Lambda\\left(\\theta\\right)`` for profile likelihood-based (PL) identification. Usually we use log-likelihood for PL analysis: ``\\Lambda( \\theta ) = - 2 ln\\left( L(\\theta) \\right)``.
- `method`: computational method to estimate confidence interval's endpoint. Currently the following methods are implemented: `:CICO_ONE_PASS`, `:LIN_EXTRAPOL`, `:QUADR_EXTRAPOL`.
- `direction`: `:right` or `:left` endpoint to estimate.

## Keyword arguments
see [`get_interval`](@ref)

"""
function get_endpoint(
    theta_init::Vector{Float64},
    theta_num::Int,
    loss_func::Function,
    method::Symbol,
    direction::Symbol = :right;

    loss_crit::Float64 = 0.0,
    # :direct, :log, :logit
    scale::Vector{Symbol} = fill(:direct, length(theta_init)),
    theta_bounds::Vector{Tuple{Float64,Float64}} = unscaling.(
        fill((-Inf, Inf), length(theta_init)),
        scale
        ),
    scan_bound::Float64 = unscaling(
        (direction==:left) ? -9.0 : 9.0,
        scale[theta_num]
        ),
    scan_tol::Float64 = 1e-3,
    loss_tol::Float64 = 1e-3,
    local_alg::Symbol = :LN_NELDERMEAD,
    kwargs... # other options for get_right_endpoint
    )
    isLeft = direction == :left

    # checking arguments
    # theta_bound[1] < theta_init < theta_bound[2]
    theta_init_outside_theta_bounds = .! [theta_bounds[i][1] < theta_init[i] < theta_bounds[i][2] for i in 1:length(theta_init)]
    if any(theta_init_outside_theta_bounds)
        throw(ArgumentError("theta_init is outside theta_bound: $(findall(theta_init_outside_theta_bounds))"))
    end
    # scan_bound should be within theta_bounds
    !(theta_bounds[theta_num][1] < scan_bound < theta_bounds[theta_num][2]) &&
        throw(ArgumentError("scan_bound are outside of the theta_bounds $(theta_bounds[theta_num])"))
    # theta_init should be within scan_bound
    if (theta_init[theta_num] >= scan_bound && !isLeft) || (theta_init[theta_num] <= scan_bound && isLeft)
        throw(ArgumentError("init values are outside of the scan_bound $scan_bound"))
    end
    # 0 <= theta_bound[1] for :log
    less_than_zero_theta_bounds = (scale .== :log) .& [theta_bounds[i][1] < 0 for i in 1:length(theta_init)]
    if any(less_than_zero_theta_bounds)
        throw(ArgumentError(":log scaled theta_bound min is negative: $(findall(less_than_zero_theta_bounds))"))
    end
    # 0 <= theta_bounds <= 1 for :logit
    less_than_zero_theta_bounds = (scale .== :logit) .& [theta_bounds[i][1] < 0 || theta_bounds[i][2] > 1 for i in 1:length(theta_init)]
    if any(less_than_zero_theta_bounds)
        throw(ArgumentError(":logit scaled theta_bound min is outside range [0,1]: $(findall(less_than_zero_theta_bounds))"))
    end
    # loss_func(theta_init) < loss_crit
    !(loss_func(theta_init) < loss_crit) &&
        throw(ArgumentError("Check theta_init and loss_crit: loss_func(theta_init) should be < loss_crit"))

    # set counter in the scope
    counter::Int = 0
    # set supreme, maximal or minimal value of scanned parameter inside critical
    supreme_gd = nothing

    # transforming
    theta_init_gd = scaling.(theta_init, scale)
    if isLeft theta_init_gd[theta_num] *= -1 end # change direction
    function loss_func_gd(theta_gd)
        theta_g = copy(theta_gd)
        if isLeft theta_g[theta_num] *= -1 end # change direction
        theta = unscaling.(theta_g, scale)
        # calculate function
        loss_norm = loss_func(theta) - loss_crit

        # update counter
        counter += 1
        # update supreme ?
        update_supreme = (loss_norm < 0.) &&
            (typeof(supreme_gd)==Nothing || (theta_gd[theta_num] > supreme_gd))
        if update_supreme
            supreme_gd = theta_gd[theta_num]
        end

        return loss_norm
    end
    theta_bounds_gd = scaling.(theta_bounds, scale)
    if isLeft theta_bounds_gd[theta_num] = (-1*theta_bounds_gd[theta_num][2], -1*theta_bounds_gd[theta_num][1]) end # change direction
    scan_bound_gd = scaling(scan_bound, scale[theta_num])
    if isLeft scan_bound_gd *= -1 end # change direction

    # calculate endpoint using base method
    (optf_gd, pp_gd, status) = get_right_endpoint(
        theta_init_gd,
        theta_num,
        loss_func_gd,
        Val(method);
        theta_bounds = theta_bounds_gd,
        scan_bound = scan_bound_gd,
        scan_tol = scan_tol,
        loss_tol = loss_tol,
        local_alg = local_alg,
        kwargs... # options for local fitter
    )

    # transforming back
    if (isLeft && typeof(optf_gd)!==Nothing) optf_gd *= -1 end # change direction
    optf = unscaling(optf_gd, scale[theta_num])
    temp_fun = (pp::ProfilePoint) -> begin
        if isLeft pp.params[theta_num] *= -1 end # change direction
        ProfilePoint(
            unscaling(pp.params[theta_num], scale[theta_num]),
            pp.loss + loss_crit,
            unscaling.(pp.params, scale),
            pp.ret,
            pp.counter
        )
    end
    pps = [ temp_fun(pp_gd[i]) for i in 1:length(pp_gd) ]
    # transforming supreme back
    if (isLeft && typeof(supreme_gd)!==Nothing) supreme_gd *= -1 end # change direction
    supreme = unscaling(supreme_gd, scale[theta_num])

    EndPoint(optf, pps, status, direction, counter, supreme)
end

"""
    function get_endpoint(
        theta_init::Vector{Float64},
        scan_func::Function,
        loss_func::Function,
        method::Symbol,
        direction::Symbol = :right;

        loss_crit::Float64 = 0.0,
        scale::Vector{Symbol} = fill(:direct, length(theta_init)),
        theta_bounds::Vector{Tuple{Float64,Float64}} = unscaling.(
            fill((-Inf, Inf), length(theta_init)),
            scale
            ),
        scan_bound::Float64 = unscaling(
            (direction==:left) ? -9.0 : 9.0,
            scale[theta_num]
            ),
        scan_tol::Float64 = 1e-3,
        loss_tol::Float64 = 1e-3,
        local_alg::Symbol = :LN_NELDERMEAD,
        kwargs...
        )

Calculates confidence interval's right or left endpoints for a function of parameters `scan_func`.

## Return
[`EndPoint`](@ref) object storing confidence interval's endpoint and intermediate profile points.

## Arguments
- `theta_init`: starting values of parameter vector ``\\theta``. The starting values should not necessary be the optimum values of `loss_func` but `loss_func(theta_init)` should be lower than `loss_crit`.
- `scan_func`: function of parameters.
- `loss_func`: loss function ``\\Lambda\\left(\\theta\\right)`` for profile likelihood-based (PL) identification. Usually we use log-likelihood for PL analysis: ``\\Lambda( \\theta ) = - 2 ln\\left( L(\\theta) \\right)``.
- `method`: computational method to estimate confidence interval's endpoint. Currently the only supported method is: `:CICO_ONE_PASS`.
- `direction`: `:right` or `:left` endpoint to estimate.

## Keyword arguments
see [`get_interval`](@ref)

"""
function get_endpoint(
    theta_init::Vector{Float64},
    scan_func::Function,
    loss_func::Function,
    method::Symbol,
    direction::Symbol = :right;

    loss_crit::Float64 = 0.0,
    # :direct, :log, :logit
    scale::Vector{Symbol} = fill(:direct, length(theta_init)),
    theta_bounds::Vector{Tuple{Float64,Float64}} = unscaling.(
        fill((-Inf, Inf), length(theta_init)),
        scale
        ),
    scan_bound::Float64 = (direction==:left) ? -1e9 : 1e9, # log scan bound is not implemented
    scan_tol::Float64 = 1e-3,
    loss_tol::Float64 = 1e-3,
    local_alg::Symbol = :LN_NELDERMEAD,
    kwargs... # other options for get_right_endpoint
    )
    isLeft = direction == :left

    # checking arguments
    # theta_bound[1] < theta_init < theta_bound[2]
    theta_init_outside_theta_bounds = .! [theta_bounds[i][1] < theta_init[i] < theta_bounds[i][2] for i in 1:length(theta_init)]
    if any(theta_init_outside_theta_bounds)
        throw(ArgumentError("theta_init is outside theta_bound: $(findall(theta_init_outside_theta_bounds))"))
    end

    # 0 <= theta_bound[1] for :log
    less_than_zero_theta_bounds = (scale .== :log) .& [theta_bounds[i][1] < 0 for i in 1:length(theta_init)]
    if any(less_than_zero_theta_bounds)
        throw(ArgumentError(":log scaled theta_bound min is negative: $(findall(less_than_zero_theta_bounds))"))
    end
    # 0 <= theta_bounds <= 1 for :logit
    less_than_zero_theta_bounds = (scale .== :logit) .& [theta_bounds[i][1] < 0 || theta_bounds[i][2] > 1 for i in 1:length(theta_init)]
    if any(less_than_zero_theta_bounds)
        throw(ArgumentError(":logit scaled theta_bound min is outside range [0,1]: $(findall(less_than_zero_theta_bounds))"))
    end
    # loss_func(theta_init) < loss_crit
    !(loss_func(theta_init) < loss_crit) &&
        throw(ArgumentError("Check theta_init and loss_crit: loss_func(theta_init) should be < loss_crit"))

    # set counter in the scope
    counter::Int = 0
    # set supreme, maximal or minimal value of scanned parameter inside critical
    supreme_gd = nothing

    # transforming
    theta_init_gd = scaling.(theta_init, scale)

    function scan_func_gd(theta_gd)
        theta_g = copy(theta_gd)
        theta = unscaling.(theta_g, scale)
        scan_val = scan_func(theta)
        scan_val_gd = isLeft ? (-1)*scan_val : scan_val

        return scan_val_gd
    end

    function loss_func_gd(theta_gd)
        theta_g = copy(theta_gd)
        theta = unscaling.(theta_g, scale)
        # calculate function
        loss_norm = loss_func(theta) - loss_crit

        # update counter
        counter += 1
        # update supreme ?
        #scan_val = scan_func(theta)
        #scan_val_gd = isLeft ? (-1)*scan_val : scan_val
        #update_supreme = (loss_norm < 0.) &&
        #    (typeof(supreme_gd)==Nothing || (scan_val_gd > supreme_gd))
        #if update_supreme
        #    supreme_gd = scan_val_gd
        #end

        return loss_norm
    end
    theta_bounds_gd = scaling.(theta_bounds, scale)

    # TODO: transformed by scan_scale: scaling(scan_bound, scan_scale)
    scan_bound_gd = scan_bound
    if isLeft scan_bound_gd *= -1 end # change direction

    # calculate endpoint using base method
    (optf_gd, pp_gd, status) = get_right_endpoint(
        theta_init_gd,
        scan_func_gd,
        loss_func_gd,
        Val(method);
        theta_bounds = theta_bounds_gd,
        scan_bound = scan_bound_gd,
        scan_tol = scan_tol,
        loss_tol = loss_tol,
        local_alg = local_alg,
        kwargs... # options for local fitter
    )

    # transforming back
    temp_fun = (pp::ProfilePoint) -> begin
        value = isLeft ? (-1)*pp.value : pp.value # TODO: scan_scale
        ProfilePoint(
            value,
            pp.loss + loss_crit,
            unscaling.(pp.params, scale),
            pp.ret,
            pp.counter
        )
    end
    pps = [ temp_fun(pp_gd[i]) for i in 1:length(pp_gd) ]
    # transforming supreme back
    if (isLeft && typeof(supreme_gd)!==Nothing) supreme_gd *= -1 end # change direction
    supreme = unscaling(supreme_gd, :direct) # scan_scale ??

    if (isLeft && typeof(optf_gd)!==Nothing) optf_gd *= -1 end # change direction
    # optf = unscaling(optf_gd, scan_scale)
    optf = optf_gd

    EndPoint(optf, pps, status, direction, counter, supreme)
end


#== MFDFA1.jl  by Tehran Davis, 2020

This is an Julia-based translation of the Matlab script based on EAF Ihlen:
Front. Physiol., 04 June 2012 | https://doi.org/10.3389/fphys.2012.00141

==#

using StatsBase, Polynomials, DataFrames, GLM, Polynomials.PolyCompat

function MFDFA(signal, scale, q, m)

#== Multifractal detrended fluctuation analysis (MFDFA)
[Hq,tq,hq,Dq,Fq]=MFDFA(signal,scale,q,m);
INPUT PARAMETERS---------------------------------------------------------
signal:       input signal
scale:        vector of scales
q:            q-order that weights the local variations
m:            polynomial order for the detrending

OUTPUT VARIABLES---------------------------------------------------------
Hq:           q-order Hurst exponent
tq:           q-order mass exponent
hq:           q-order singularity exponent
Dq:           q-order dimension
Fq:           q-order scaling function
--------------------------------------------------------------------==#

    X=cumsum(signal.-mean(signal), dims=1);

    RMS_scale = [Float64[] for i=1:length(scale)]
    qRMS = [Float64[] for i=1:length(q),j=1:length(scale)]
    Fq = Array{Float64}(undef, length(q), length(scale))
    Hq = Float64[]
    Hq_r2 = Float64[]
    qRegLine = []

    # if minimum(size(X))!=1||minimum(size(scale))!=1||minimum(size(q))!=1;
    #     error("Input arguments signal, scale and q must be a vector");
    # end
    # if size(X,2)==1;
    #    X=transpose(X);
    # end
    # if min(scale)<m+1
    #    error('The minimum scale must be larger than trend order m+1')
    # end

    for ns=1:length(scale)
        # segments[ns]=floor(length(X)/scale[ns]);
        segments = floor(length(X)/scale[ns])
        for v=1:segments
            Index=Int.([(((v-1)*scale[ns])+1):v*scale[ns];])
            C=Polynomials.PolyCompat.polyfit(Index,X[Index],m);
            pfit = Polynomials.PolyCompat.polyval(C,Index)
            push!(RMS_scale[ns],sqrt(mean((X[Index]-pfit).^2)))
        end
        for nq=1:length(q)
            qRMS[nq,ns] = RMS_scale[ns].^q[nq]
            Fq[nq,ns]=mean(qRMS[nq,ns]) .^(1/q[nq])
        end
        Fq[findall(x->x==0,q),ns].=exp(0.5*mean(log.(RMS_scale[ns].^2)));
    end
    for nq=1:length(q)

        # using Polynomials - old method
        # C = polyfit(log2.(scale),log2.(Fq[nq,:]),1)
        # push!(Hq,C[1])
        # push!(qRegLine,polyval(C,log2.(scale)))

        # using GLM: benefit we get R-squared
        C = GLM.lm(@formula(y~x),DataFrame(y=log2.(Fq[nq,:]),x=log2.(scale)))
        push!(Hq,GLM.coef(C)[2])
        push!(qRegLine,GLM.predict(C))
        push!(Hq_r2, GLM.r²(C))

    end

    tq = (Hq.*q).-1
    hq = diff(tq)./(q[2]-q[1])
    Dq = q[1:(length(q)-1)].*hq-tq[1:(length(tq)-1)]

    mfw = maximum(hq) - minimum(hq);

    output = Dict("Hq" => Hq,
                    "Hq_r2" => Hq_r2,
                    "tq" => tq,
                    "hq" => hq,
                    "Dq" => Dq,
                    "Fq" => Fq,
                    "q" => q,
                    "mfw" => mfw
                    )

    return output

end


## --- Dealing with different number representations

    """
    ```julia
    nearest(T, x)
    ```
    Convert `x` to the nearest representable value in type T, rounding if inexact
    """
    function nearest(::Type{T}, x) where T <: Number
        T(max(min(x, typemax(T)), typemin(T)))
    end
    function nearest(::Type{T}, x) where T <: Integer
        round(T, max(min(x, typemax(T)), typemin(T)))
    end
    export nearest

## --- Fast inverse square-root

    """
    ```julia
    fast_inv_sqrt(x)
    ```
    The infamous fast inverse square root of `x`, in 32 and 64 bit versions.
    Can be up to 10x faster than base `1/sqrt(x)`, though with nontrivial loss
    of precision. The implementations here are good to about 4 ppm.
    """
    function fast_inv_sqrt(x::Float64)
        x_2 = 0.5 * x
        result = Base.sub_int(9.603007803048109e153, Base.lshr_int(x,1)) # Floating point magic
        result *= ( 1.5 - (x_2 * result * result )) # Newton's method
        result *= ( 1.5 - (x_2 * result * result )) # Newton's method (again)
        return result
    end
    function fast_inv_sqrt(x::Float32)
        x_2 = 0.5f0 * x
        result = Base.sub_int(1.321202f19, Base.lshr_int(x,1)) # Floating point magic
        result *= ( 1.5f0 - (x_2 * result * result) ) # Newton's method
        result *= ( 1.5f0 - (x_2 * result * result) ) # Newton's method (again)
        return result
    end
    export fast_inv_sqrt

## --- Fast vectorized mean
    function _vmean(A::AbstractArray{T}, ::Colon) where T
        m = zero(T)
        @turbo for i ∈ eachindex(A)
            m += A[i]
        end
        return m / length(A)
    end

## --- Base-10 version of log1p
#
#     log10f(x::Number,from::Number=-1) = log10(abs(x-from)+1)*sign(x-from)
#

## --- Some mathematical constants

    const SQRT2 = sqrt(2)
    const SQRT2PI = sqrt(2*pi)
    const INVSQRT2 = 1/sqrt(2)
    const AN = Union{AbstractArray{<:Number},Number}

## --- Gaussian distribution functions

    """
    ```julia
    normpdf(mu,sigma,x)
    ```
    Probability density function of the Normal (Gaussian) distribution

    ``ℯ^{-(x-μ)^2 / (2σ^2)} / σ√2π``

    with mean `mu` and standard deviation `sigma`, evaluated at `x`
    """
    @inline normpdf(mu,sigma,x) = exp(-(x-mu)*(x-mu) / (2*sigma*sigma)) / (SQRT2PI*sigma)
    @inline normpdf(mu::Number,sigma::Number,x::Number) = exp(-(x-mu)*(x-mu) / (2*sigma*sigma)) / (SQRT2PI*sigma)
    normpdf(mu::AN,sigma::AN,x::AN) = @turbo @. exp(-(x-mu)*(x-mu) / (2*sigma*sigma)) / (SQRT2PI*sigma)
    export normpdf

    """
    ```julia
    normpdf_ll(mu, sigma, x)
    ```
    Fast log likelihood corresponding to the natural logarithm of the probability
    density function of a Normal (Gaussian) distribution with mean `mu` and
    standard deviation `sigma`, evaluated at `x`.

    If `x`, [`mu`, and `sigma`] are given as arrays, the sum of the log likelihood
    over all `x` will be returned.

    See also `normpdf`
    """
    @inline normpdf_ll(mu,sigma,x) = -(x-mu)*(x-mu) / (2*sigma*sigma)
    @inline normpdf_ll(mu::Number,sigma::Number,x::Number) = -(x-mu)*(x-mu) / (2*sigma*sigma)
    function normpdf_ll(mu::Number,sigma::Number,x::AbstractArray)
        inv_s2 = 1/(2*sigma*sigma)
        ll = zero(typeof(inv_s2))
        @turbo for i ∈ eachindex(x)
            ll -= (x[i]-mu)*(x[i]-mu) * inv_s2
        end
        return ll
    end
    function normpdf_ll(mu::AbstractArray,sigma::Number,x::AbstractArray)
        inv_s2 = 1/(2*sigma*sigma)
        ll = zero(typeof(inv_s2))
        @turbo for i ∈ eachindex(x)
            ll -= (x[i]-mu[i])*(x[i]-mu[i]) * inv_s2
        end
        return ll
    end
    function normpdf_ll(mu::Number,sigma::AbstractArray,x::AbstractArray)
        ll = zero(float(eltype(sigma)))
        @turbo for i ∈ eachindex(x)
            ll -= (x[i]-mu)*(x[i]-mu) / (2*sigma[i]*sigma[i])
        end
        return ll
    end
    function normpdf_ll(mu::AbstractArray,sigma::AbstractArray,x::AbstractArray)
        ll = zero(float(eltype(sigma)))
        @turbo for i ∈ eachindex(x)
            ll -= (x[i]-mu[i])*(x[i]-mu[i]) / (2*sigma[i]*sigma[i])
        end
        return ll
    end
    export normpdf_ll

    """
    ```julia
    normcdf(mu,sigma,x)
    ```
    Cumulative distribution function of the Normal (Gaussian) distribution

    ``1/2 + erf(\frac{x-μ}{σ√2})/2``

    with mean `mu` and standard deviation `sigma`, evaluated at `x`.
    """
    @inline normcdf(mu,sigma,x) = 0.5 + 0.5 * erf((x-mu) / (sigma*SQRT2))
    @inline normcdf(mu::Number,sigma::Number,x::Number) = 0.5 + 0.5 * erf((x-mu) / (sigma*SQRT2))
    normcdf(mu::AN,sigma::AN,x::AN) = @turbo @. 0.5 + 0.5 * erf((x-mu) / (sigma*SQRT2))
    export normcdf

    """
    ```julia
    normcdf_ll(mu, sigma, x)
    ```
    Fast log likelihood corresponding to the natural logarithm of the cumulative
    distribution function of a Normal (Gaussian) distribution with mean `mu` and
    standard deviation `sigma`, evaluated at `x`.

    If `x`, [`mu`, and `sigma`] are given as arrays, the sum of the log likelihood
    over all `x` will be returned.

    See also `normcdf`
    """
    @inline function normcdf_ll(xₛ::Number)
        if xₛ < -1.0
            return log(0.5*erfcx(-xₛ * INVSQRT2)) - 0.5*abs2(xₛ)
        else
            return log1p(-0.5*erfc(xₛ * INVSQRT2))
        end
    end
    function normcdf_ll(xₛ::AbstractArray)
        ll = zero(float(eltype(xₛ)))
        @inbounds for i ∈ eachindex(xₛ)
            ll += normcdf_ll(xₛ[i])
        end
        return ll
    end
    @inline function normcdf_ll(mu::Number, sigma::Number, x::Number)
        xₛ = (x - mu) / sigma
        return normcdf_ll(xₛ)
    end
    function normcdf_ll(mu::Number,sigma::Number,x::AbstractArray)
        inv_sigma = 1/sigma
        ll = zero(typeof(inv_sigma))
        @inbounds for i ∈ eachindex(x)
            xₛ = (x[i] - mu) * inv_sigma
            ll += normcdf_ll(xₛ)
        end
        return ll
    end
    function normcdf_ll(mu::AbstractArray,sigma::Number,x::AbstractArray)
        inv_sigma = 1/sigma
        ll = zero(typeof(inv_sigma))
        @inbounds for i ∈ eachindex(x)
            xₛ = (x[i] - mu[i]) * inv_sigma
            ll += normcdf_ll(xₛ)
        end
        return ll
    end
    function normcdf_ll(mu::Number,sigma::AbstractArray,x::AbstractArray)
        ll = zero(float(eltype(sigma)))
        @inbounds for i ∈ eachindex(x)
            xₛ = (x[i] - mu) / sigma[i]
            ll += normcdf_ll(xₛ)
        end
        return ll
    end
    function normcdf_ll(mu::AbstractArray,sigma::AbstractArray,x::AbstractArray)
        ll = zero(float(eltype(sigma)))
        @inbounds for i ∈ eachindex(x)
            xₛ = (x[i] - mu[i]) / sigma[i]
            ll += normcdf_ll(xₛ)
        end
        return ll
    end
    export normcdf_ll

    """
    ```julia
    normcdf_ll!(mu, sigma, x)
    ```
    Fast log likelihood corresponding to the natural logarithm of the cumulative
    distribution function of a Normal (Gaussian) distribution with mean `mu` and
    standard deviation `sigma`, evaluated at `x`.

    As `normcdf_ll`, but in-place (using `x` as a buffer).
    """
    function normcdf_ll!(xₛ::AbstractArray)
        @inbounds for i ∈ eachindex(xₛ)
            xₛ[i] = normcdf_ll(xₛ[i])
        end
        ll = zero(float(eltype(xₛ)))
        @turbo for i ∈ eachindex(xₛ)
            ll += xₛ[i]
        end
        return ll
    end
    function normcdf_ll!(mu::Number,sigma::Number,x::AbstractArray)
        inv_sigma = 1/sigma
        @inbounds for i ∈ eachindex(x)
            xₛ = (x[i] - mu) * inv_sigma
            x[i] = normcdf_ll(xₛ)
        end
        ll = zero(typeof(inv_sigma))
        @turbo for i ∈ eachindex(x)
            ll += x[i]
        end
        return ll
    end
    function normcdf_ll!(mu::AbstractArray,sigma::Number,x::AbstractArray)
        inv_sigma = 1/sigma
        @inbounds for i ∈ eachindex(x)
            xₛ = (x[i] - mu[i]) * inv_sigma
            x[i] = normcdf_ll(xₛ)
        end
        ll = zero(typeof(inv_sigma))
        @turbo for i ∈ eachindex(x)
            ll += x[i]
        end
        return ll
    end
    function normcdf_ll!(mu::Number,sigma::AbstractArray,x::AbstractArray)
        @inbounds for i ∈ eachindex(x)
            xₛ = (x[i] - mu) / sigma[i]
            x[i] = normcdf_ll(xₛ)
        end
        ll = zero(float(eltype(sigma)))
        @turbo for i ∈ eachindex(x)
            ll += x[i]
        end
        return ll
    end
    function normcdf_ll!(mu::AbstractArray,sigma::AbstractArray,x::AbstractArray)
        @inbounds for i ∈ eachindex(x)
            xₛ = (x[i] - mu[i]) / sigma[i]
            x[i] = normcdf_ll(xₛ)
        end
        ll = zero(float(eltype(sigma)))
        @turbo for i ∈ eachindex(x)
            ll += x[i]
        end
        return ll
    end
    export normcdf_ll!

    """
    ```julia
    normcdf!(result,mu,sigma,x)
    ```
    In-place version of `normcdf`
    """
    function normcdf!(result::Array, mu::Number, sigma::Number, x::AbstractArray)
        T = eltype(result)
        inv_sigma_sqrt2 = one(T)/(sigma*T(SQRT2))
        @turbo for i ∈ 1:length(x)
            result[i] = T(0.5) + T(0.5) * erf((x[i]-mu) * inv_sigma_sqrt2)
        end
        return result
    end
    export normcdf!

    """
    ```julia
    norm_quantile(F::Number)
    ```
    How far away from the mean (in units of sigma) should we expect proportion
    F of the samples to fall in a standard Gaussian (Normal[0,1]) distribution
    """
    @inline norm_quantile(F) = SQRT2*erfinv(2*F-1)
    export norm_quantile

    """
    ```julia
    norm_width(N::Number)
    ```
    How dispersed (in units of sigma) should we expect a sample of N numbers
    drawn from a standard Gaussian (Normal[0,1]) distribution to be?
    """
    @inline norm_width(N) = 2*norm_quantile(1 - 1/(2N))
    export norm_width

    """
    ```julia
    normproduct(μ1, σ1, μ2, σ2)
    ```
    The integral of the product of two normal distributions N[μ1,σ1] * N[μ2,σ2].
    This is itself just another Normal distribution! Specifically, one with
    variance σ1^2 + σ2^2, evaluated at distance |μ1-μ2| from the mean
    """
    normproduct(μ1, σ1, μ2, σ2) = normpdf(μ1, sqrt.(σ1.*σ1 + σ2.*σ2), μ2)
    export normproduct

    """
    ```julia
    normproduct_ll(μ1, σ1, μ2, σ2)
    ```
    Log likelihood corresponding to the integral of N[μ1,σ1] * N[μ2,σ2]
    As `normproduct`, but using the fast log likelihood of a Normal distribution
    """
    normproduct_ll(μ1, σ1, μ2, σ2) = normpdf_ll(μ1, sqrt.(σ1.*σ1 + σ2.*σ2), μ2)
    export normproduct_ll

## --- Geometry

    """
    ```julia
    inpolygon(x,y,point)
    ```
    Check if a 2D polygon defined by the arrays `x`, `y` contains a given `point`.
    Returns boolean (true or false)
    """
    function inpolygon(x,y,point)
        # Check that we have the right kind of input data
        if length(x) != length(y)
            error("polygon must have equal number of x and y points\n")
        end
        if length(x) < 3
            error("polygon must have at least 3 points\n")
        end
        if length(point) != 2
            error("point must be an ordered pair (x,y)\n")
        end

        # Extract x and y data of point
        point_x = point[1]
        point_y = point[2]
        # For first point, previous point is last
        x_here = x[end]
        y_here = y[end]
        # Ensure we are not sitting parallel to a vertex by infinitessimally moving the point
        if y_here == point_y
            point_y = nextfloat(float(point_y))
        end
        if x_here == point_x
            point_x = nextfloat(float(point_x))
        end

        # Check how many times a line projected right along x-axis from point intersects the polygon
        intersections = 0
        @inbounds for i ∈ eachindex(x)
            # Recycle our vertex
            x_last = copy(x_here)
            y_last = copy(y_here)

            # Get a new vertex
            x_here = x[i]
            y_here = y[i]

            # Ensure we are not sitting parallel to a vertex by infinitessimally moving the point
            if y_here == point_y
                point_y = nextfloat(float(point_y))
            end
            if x_here == point_x
                point_x = nextfloat(float(point_x))
            end

            if y_last > point_y && y_here > point_y
                # If both ys above point, no intersection
                continue
            elseif y_last < point_y && y_here < point_y
                # If both ys below point, no intersection
                continue
            elseif x_last < point_x && x_here < point_x
                # If both x's left of point, no intersection
                continue
            elseif x_last > point_x && x_here > point_x
                # By elimination
                # We have one y above and y below our point
                # If both y's are right of line, then definite intersection
                intersections += 1
                continue
            else
                # By elimination
                # One y above and one y below
                # One x to the right and one x to the left
                # We must project
                dy = y_here - y_last
                if abs(dy) > 0
                    dx = x_here - x_last
                    inv_slope = dx / dy
                    x_proj = x_last + (point_y - y_last) * inv_slope
                    if x_proj > point_x
                        intersections += 1
                    end
                end
            end
        end

        # If number of intersections is odd, point is in the polygon
        return Bool(mod(intersections,2))
    end
    export inpolygon

    """
    ```julia
    (rows, columns) = find_grid_inpolygon(grid_x, grid_y, poly_x, poly_y)
    ```
    Find the indexes of grid points that fall within a polygon for a grid with
    cell centers given by grid_x (j-columns of grid) and grid_y (i-rows of grid).
    Returns a list of rows and columns in the polygon
    """
    function find_grid_inpolygon(grid_x, grid_y, poly_x, poly_y)
        # Check that we have the right kind of input data
        if length(poly_x) != length(poly_y)
            error("polygon must have equal number of x and y points\n")
        end
        if length(poly_x) < 3
            error("polygon must have at least 3 points\n")
        end

        # Find maximum x and y range of polygon
        (xmin, xmax) = extrema(poly_x)
        (ymin, ymax) = extrema(poly_y)

        # Find the matrix indices within the range of the polygon (if any)
        column_inrange = findall((grid_x .>= xmin) .& (grid_x .<= xmax))
        row_inrange = findall((grid_y .>= ymin) .& (grid_y .<= ymax))

        # Keep a list of matrix indexes in the polygon
        row = Array{Int}(undef,length(column_inrange) * length(row_inrange))
        column = Array{Int}(undef,length(column_inrange) * length(row_inrange))
        n = 0
        for j = 1:length(column_inrange)
            for i = 1:length(row_inrange)
                point = [grid_x[column_inrange[j]], grid_y[row_inrange[i]]]
                if inpolygon(poly_x, poly_y, point)
                    n += 1
                    row[n] = row_inrange[i]
                    column[n] = column_inrange[j]
                end
            end
        end

        return (row[1:n], column[1:n])
    end
    export find_grid_inpolygon

    """
    ```julia
    arcdistance(latᵢ,lonᵢ,lat,lon)
    ```
    Calculate the distance on a sphere between the point (`latᵢ`,`lonᵢ`) and any
    number of points in (`lat`,`lon`).
    Latitude and Longitude should be specified in decimal degrees
    """
    function arcdistance(latᵢ,lonᵢ,lat,lon)
        # Argument for acos()
        arg = sin.(latᵢ .* pi/180) .* sin.(lat .* pi/180) .+ cos.(latᵢ*pi/180) .* cos.(lat .* pi/180).*cos.((lonᵢ .- lon) .* pi/180)

        # Avoid domain errors from imprecise sine and cosine math
        arg[arg .> 1.0] .= 1.0
        arg[arg .< -1.0] .= -1.0

        # Calculate angular distance
        theta = 180/pi .* acos.(arg)
        return theta
    end
    export arcdistance

## --- Linear regression

    """
    ```julia
    (a,b) = linreg(x::AbstractVector, y::AbstractVector)
    ```

    Returns the coefficients for a simple linear least-squares regression of
    the form `y = a + bx`
    """
    linreg(x::AbstractVector, y::AbstractVector) = hcat(fill!(similar(x), 1), x) \ y
    export linreg

## --- End of File


function _one_iterate(cutoff_multiplier::Real, target::Real, draws::Array{T,2}, covar_matrix, tuning_parameter::Real) where T<:Real
    cutoffs = cutoff_multiplier .* sqrt.(diag(covar_matrix))
    number_of_draws = size(draws)[1]
    in_confidence_area = 0
    for i in 1:number_of_draws
        in_confidence_area += all(abs.(draws[i,:]) .< cutoffs)
    end
    mass_in_area = in_confidence_area/number_of_draws
    confidence_gap = target - mass_in_area
    return cutoff_multiplier + confidence_gap * tuning_parameter
end

function _randoms(chol, num::Integer, Seed::Integer)
    twist = MersenneTwister(Seed)
    dims = size(chol)[1]
    float_type = typeof(chol[1,1])
    array = Array{float_type,2}(undef, num, dims)
    for i in 1:num
        sobs = rand(twist, dims)
        normal_draw = quantile.(Ref(Normal()), sobs)
        scaled_draw = chol * normal_draw
        array[i,:] = scaled_draw
    end
    return array
end

"""
    get_confidence_hypercube(covar::ForwardCovariance, confidence_level::Real, data::Array{T,2}; tuning_parameter::Real = 1.0)
This returns the endpoints of a hypercube that contains confidence_level (%) of the dataset.
"""
function get_confidence_hypercube(covar::ForwardCovariance, confidence_level::Real, data::Array{T,2}; tuning_parameter::Real = 1.0, ConvergenceMetricThreshold::Real = 1e-10) where T<:Real
    # Using a univariate guess as we can get these pretty cheaply.
    guess = quantile(Normal(), 0.5*(1+confidence_level))
    # This runs once so that any error is explictly thrown and is traceable rather than being obscured by FixedPoint's try-catch
    #_ = _one_iterate.(guess, Ref(confidence_level), Ref(data), Ref(covar.covariance_), Ref(tuning_parameter))
    FP = fixed_point(x -> _one_iterate.(x, Ref(confidence_level), Ref(data), Ref(covar.covariance_), Ref(tuning_parameter)), [guess]; ConvergenceMetricThreshold = ConvergenceMetricThreshold, MaxIter = 10000)
    cutoff_multiplier = FP.FixedPoint_[1]
    cutoffs = vcat(zip(-cutoff_multiplier .* sqrt.(diag(covar.covariance_)) , cutoff_multiplier .* sqrt.(diag(covar.covariance_)))...)
    return Dict{Symbol,Tuple{T,T}}(covar.covariance_labels_ .=> cutoffs)
end

function get_confidence_hypercube(covar::ForwardCovariance, confidence_level::Real, num::Integer; tuning_parameter::Real = 1.0,  ConvergenceMetricThreshold::Real = 1e-10)
    dims = length(covar.covariance_labels_)
    data = _randoms(covar.chol_, num, 1)
    return get_confidence_hypercube(covar, confidence_level, data; tuning_parameter = tuning_parameter, ConvergenceMetricThreshold = ConvergenceMetricThreshold)
end


if basename(pwd()) == "aoc"
    cd("2015/17")
end
using Combinatorics

function part1(filename::AbstractString, target::Integer)
    count(==(target) ∘ sum, combinations(parse.(Int, eachline(filename))))
end

@assert part1("example.txt", 25) == 4
part1("input.txt", 150)

function part2(filename::AbstractString, target::Integer)
    lengths = map(length,
        Iterators.filter(
            ==(target) ∘ sum,
            combinations(parse.(Int, eachline(filename)))))
    count(==(first(lengths)), lengths)
end

@assert part2("example.txt", 25) == 3
part2("input.txt", 150)


using DiffEqSensitivity, Flux, Zygote, OrdinaryDiffEq, Test # , Plots

function lotka_volterra(du,u,p,t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = (α - β*y)x
  du[2] = dy = (δ*x - γ)y
end
p = [2.2, 1.0, 2.0, 0.4]
u0 = [1.0,1.0]
prob = ODEProblem(lotka_volterra,u0,(0.0,10.0),p)

# Reverse-mode

function predict_rd(p)
  Array(solve(prob,Tsit5(),p=p,saveat=0.1,reltol=1e-4,sensealg=TrackerAdjoint()))
end
loss_rd(p) = sum(abs2,x-1 for x in predict_rd(p))
loss_rd() = sum(abs2,x-1 for x in predict_rd(p))
loss_rd()

grads = Zygote.gradient(loss_rd, p)
@test !iszero(grads[1])

opt = ADAM(0.1)
cb = function ()
  display(loss_rd())
  # display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
end

# Display the ODE with the current parameter values.
loss1 = loss_rd()
Flux.train!(loss_rd, Flux.params(p), Iterators.repeated((), 100), opt, cb = cb)
loss2 = loss_rd()
@test 10loss2 < loss1

# Forward-mode, R^n -> R^m layer

p = [2.2, 1.0, 2.0, 0.4]
function predict_fd()
  vec(Array(solve(prob,Tsit5(),p=p,saveat=0.0:0.1:1.0,reltol=1e-4,sensealg=ForwardDiffSensitivity())))
end
loss_fd() = sum(abs2,x-1 for x in predict_fd())
loss_fd()

ps = Flux.params(p)
grads = Zygote.gradient(loss_fd, ps)
@test !iszero(grads[p])

data = Iterators.repeated((), 100)
opt = ADAM(0.1)
cb = function ()
  display(loss_fd())
  # display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
end

# Display the ODE with the current parameter values.
loss1 = loss_fd()
Flux.train!(loss_fd, ps, data, opt, cb = cb)
loss2 = loss_fd()
@test 10loss2 < loss1

# Adjoint sensitivity
p = [2.2, 1.0, 2.0, 0.4]
ps = Flux.params(p)
function predict_adjoint()
    solve(remake(prob,p=p),Tsit5(),saveat=0.1,reltol=1e-4)
end
loss_reduction(sol) = sum(abs2,x-1 for x in vec(sol))
loss_adjoint() = loss_reduction(predict_adjoint())
loss_adjoint()

grads = Zygote.gradient(loss_adjoint, ps)
@test !iszero(grads[p])

data = Iterators.repeated((), 100)
opt = ADAM(0.1)
cb = function ()
  display(loss_adjoint())
  # display(plot(solve(remake(prob,p=p),Tsit5(),saveat=0.1),ylim=(0,6)))
end

# Display the ODE with the current parameter values.
loss1 = loss_adjoint()
Flux.train!(loss_adjoint, ps, data, opt, cb = cb)
loss2 = loss_adjoint()
@test 10loss2 < loss1


# # Wiring diagrams in Graphviz
#
#md # [![](https://img.shields.io/badge/show-nbviewer-579ACA.svg)](@__NBVIEWER_ROOT_URL__/generated/graphics/graphviz_wiring_diagrams.ipynb)
#
# Catlab can draw wiring diagrams using the `dot` program in
# [Graphviz](https://www.graphviz.org/). This feature requires that Graphviz be
# installed, but does not require any additional Julia packages.

using Catlab.WiringDiagrams, Catlab.Graphics

# ## Examples

# ### Symmetric monoidal category

using Catlab.Doctrines

A, B = Ob(FreeSymmetricMonoidalCategory, :A, :B)
f = Hom(:f, A, B)
g = Hom(:g, B, A)
h = Hom(:h, otimes(A,B), otimes(A,B));

# To start, here are a few very simple examples.

to_graphviz(f)
#-
to_graphviz(compose(f,g))
#-
to_graphviz(otimes(f,g))

# In the next example, notice how Graphviz automatically "untwists" the double
# braiding to minimize edge crossings.

to_graphviz(compose(braid(A,A), otimes(f,f), braid(B,B)))

# Here is a larger composite morphism.

composite = compose(otimes(g,f), h, otimes(f,g))
to_graphviz(composite)

# By default, the wiring diagram is laid out from top to bottom. Other layout
# orientations can be requested, such as left-to-right or bottom-to-top:

to_graphviz(composite, orientation=LeftToRight)
#-
to_graphviz(composite, orientation=BottomToTop)

# When working with very large diagrams (larger than the ones shown here), it is
# sometimes convenient to omit the ports of the outer box and any wires attached
# to them.

to_graphviz(composite, outer_ports=false)

# ### Biproduct category

A, B = Ob(FreeBiproductCategory, :A, :B)
f = Hom(:f, A, B)
g = Hom(:g, B, A);

# By default, copies and merges are drawn the way they are represented
# internally, as multiple wires.

f1 = compose(mcopy(A), otimes(f,f))
to_graphviz(f1)
#-
f2 = compose(mcopy(A), otimes(f,f), mmerge(B))
to_graphviz(f2)

# To draw nodes for copies and merges, we need to add junctions to the wiring
# diagram.

to_graphviz(add_junctions!(to_wiring_diagram(f1)))
#-
to_graphviz(add_junctions!(to_wiring_diagram(f2)))

# ### Traced monoidal category

A, B, X, Y = Ob(FreeTracedMonoidalCategory, :A, :B, :X, :Y)
f = Hom(:f, otimes(X,A), otimes(X,B))

to_graphviz(trace(X, A, B, f))
#-
to_graphviz(trace(X, A, B, f), orientation=LeftToRight)
#-
g, h = Hom(:g, A, A), Hom(:h, B, B)

trace_naturality = trace(X, A, B, compose(otimes(id(X),g), f, otimes(id(X),h)))
to_graphviz(trace_naturality, orientation=LeftToRight)

# ## Custom styles

# The visual appearance of wiring diagrams can be customized by setting Graphviz
# [attributes](https://www.graphviz.org/doc/info/attrs.html) at the graph, node,
# edge, and cell levels. Graph, node, and edge attributes are described in the
# Graphviz documentation. Cell attributes are passed to the primary cell of the
# [HTML-like label](https://www.graphviz.org/doc/info/shapes.html#html) used for
# the boxes.

A, B, C = Ob(FreeSymmetricMonoidalCategory, :A, :B, :C)
f, g = Hom(:f, A, B), Hom(:g, B, C)

to_graphviz(compose(f,g),
  labels = true, label_attr=:headlabel,
  node_attrs = Dict(
    :fontname => "Courier",
  ),
  edge_attrs = Dict(
    :fontname => "Courier",
    :labelangle => "25",
    :labeldistance => "2",
  ),
  cell_attrs = Dict(
    :bgcolor => "lavender",
  )
)

# ## Output formats

# The function `to_graphviz` returns an object of a type `Graphviz.Graph`,
# representing a Graphviz graph as an abstract syntax tree. When displayed
# interactively, this object is automatically run through Graphviz and rendered
# as an SVG image. Sometimes it is convenient to perform this process manually,
# to change the output format or further customize the generated dot file.

# To generate a dot file, use the builtin pretty-printer. This feature does not
# require Graphviz to be installed.

using Catlab.Graphics: Graphviz

graph = to_graphviz(compose(f,g))
Graphviz.pprint(graph)

# Catlab provides a simple wrapper around the Graphviz command-line programs.
# For example, here the JSON output for the graph.

import JSON

JSON.parse(Graphviz.run_graphviz(graph, format="json0"))


# ---
# title: 1006. Clumsy Factorial
# id: problem1006
# author: Tian Jun
# date: 2020-10-31
# difficulty: Medium
# categories: Math
# link: <https://leetcode.com/problems/clumsy-factorial/description/>
# hidden: true
# ---
#
# Normally, the factorial of a positive integer `n` is the product of all
# positive integers less than or equal to `n`.  For example, `factorial(10) = 10
# * 9 * 8 * 7 * 6 * 5 * 4 * 3 * 2 * 1`.
#
# We instead make a _clumsy factorial:_  using the integers in decreasing order,
# we swap out the multiply operations for a fixed rotation of operations:
# multiply (*), divide (/), add (+) and subtract (-) in this order.
#
# For example, `clumsy(10) = 10 * 9 / 8 + 7 - 6 * 5 / 4 + 3 - 2 * 1`.  However,
# these operations are still applied using the usual order of operations of
# arithmetic: we do all multiplication and division steps before any addition or
# subtraction steps, and multiplication and division steps are processed left to
# right.
#
# Additionally, the division that we use is _floor division_  such that `10 * 9
# / 8` equals `11`.  This guarantees the result is an integer.
#
# `Implement the clumsy` function as defined above: given an integer `N`, it
# returns the clumsy factorial of `N`.
#
#
#
# **Example 1:**
#
#
#
#     Input: 4
#     Output:  7
#     Explanation: 7 = 4 * 3 / 2 + 1
#
#
# **Example 2:**
#
#
#
#     Input: 10
#     Output: 12
#     Explanation: 12 = 10 * 9 / 8 + 7 - 6 * 5 / 4 + 3 - 2 * 1
#
#
#
#
# **Note:**
#
#   1. `1 <= N <= 10000`
#   2. `-2^31 <= answer <= 2^31 - 1`  (The answer is guaranteed to fit within a 32-bit integer.)
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


using PETSc2

import MPI
MPI.Init()

PetscInitialize(["-ksp_monitor","-malloc","-malloc_debug","-malloc_dump"]);
#
#   Create a vector and put values in it

comm = MPI.COMM_WORLD
comm_size = MPI.Comm_size(MPI.COMM_WORLD)
comm_rank = MPI.Comm_rank(MPI.COMM_WORLD)

sys_size_local = 3
println("my comm rank = ", comm_rank)

b = PetscVec(comm);
PetscVecSetType(b,"mpi");
PetscVecSetSizes(b,sys_size_local, comm_size*sys_size_local);

x = PetscVec(comm);
PetscVecSetType(x,"mpi");
PetscVecSetSizes(x,sys_size_local, comm_size*sys_size_local);

rhs = zeros(sys_size_local)
for i=1:sys_size_local
  idxm = [(comm_rank)*10 + i]  # index
  PetscVecSetValues(b,idxm,[ convert(Float64,i)],INSERT_VALUES);
  rhs[i] = i
end

PetscVecAssemblyBegin(b);
println("began vector assembly")
PetscVecAssemblyEnd(b);
println("finished vector assembly")
PetscView(b);

println("finished assembling vector")
#
#    Create a matrix

A_julia = [1.0 2.0 3; 4 5 7; 7 8 9]
println("rhs = ", rhs)
println("A_julia = ", A_julia)
x_julia = A_julia\rhs
println("x_julia = ", x_julia)

A = PetscMat(comm);
PetscMatSetType(A,"mpiaij");
PetscMatSetSizes(A,sys_size_local,sys_size_local,comm_size*sys_size_local,comm_size*sys_size_local);
PetscSetUp(A);
for i=1:sys_size_local
  for j = 1:sys_size_local
    idxm = [(comm_rank)*sys_size_local + i]  # row index
    idxn = [(comm_rank)*sys_size_local + j]  # column index
    PetscMatSetValues(A,idxm, idxn, [A_julia[i,j]],INSERT_VALUES);
  end
end

println("finished setting matrix values")
PetscMatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
PetscMatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
println("finished assembling matrix")
PetscView(A);

println("finished viewing matrix")

println("setting up linear solve")
# setup linear solve
ksp = KSP(comm)
KSPSetOperators(ksp, A, A)
KSPSetFromOptions(ksp)
println("performing linear solve")
KSPSolve(ksp, b, x)
println("finished performing linear solve")
PetscView(x)

# copy solution back to Julia
x_copy = zeros(sys_size_local)
idx = collect(0:2)

PetscVecGetValues(x, sys_size_local, idx, x_copy)

println("x_copy = ", x_copy)

err_norm = norm(x_copy - x_julia)
println("err norm = ", err_norm)

PetscDestroy(b)
PetscDestroy(A)
PetscDestroy(ksp)
# remove references to all Petsc variables
# perform solution
#A_julia\rhs
#println("x = ", x)

#PetscFinalize()