# -------------- #
# B-Spline Basis #
# -------------- #

struct Spline <: BasisFamily end

mutable struct SplineParams{T<:AbstractVector} <: BasisParams
    breaks::T
    evennum::Int
    k::Int

    # constructor to accept spline params arguments, do some pre-processing
    function SplineParams{T}(breaks::T, evennum::Int, k::Int) where T
        # error handling
        k < 0 && error("spline order must be positive")
        length(breaks) < 2 && error("Must have at least two breakpoints")
        any(diff(breaks) .< 0) && error("Breakpoints must be non-decreasing")

        if evennum == 0  # 43
            if length(breaks) == 2  # 44
                evennum = 2
            end
        else
            if length(breaks) == 2
                breaks = linspace(breaks[1], breaks[2], evennum)
            else
                error("Breakpoint squence must contain 2 values when evennum > 0")
            end
        end
        new{T}(breaks, evennum, k)
    end

end

SplineParams(breaks::T, evennum::Int, k::Int) where {T<:AbstractVector} =
    SplineParams{T}(breaks, evennum, k)

# constructor to take a, b, n and form linspace for breaks
SplineParams(n::Int, a::Real, b::Real, k::Int=3) =
    SplineParams(linspace(a, b, n), 0, k)

## BasisParams interface
# define these methods on the type, the instance version is defined over
# BasisParams
family(::Type{T}) where {T<:SplineParams} = Spline
family_name(::Type{T}) where {T<:SplineParams} = "Spline"
Base.issparse(::Type{T}) where {T<:SplineParams} = true
function Base.eltype(::Type{SplineParams{T}}) where T
    elT = eltype(T)
    elT <: Integer ? Float64 : elT
end

# methods that only make sense for instances
Base.min(p::SplineParams) = minimum(p.breaks)
Base.max(p::SplineParams) = maximum(p.breaks)
Base.length(p::SplineParams) = length(p.breaks) + p.k - 1

function Base.show(io::IO, p::SplineParams)
    m = string("$(p.k) order spline interpoland parameters from ",
               "$(p.breaks[1]), $(p.breaks[end])")
    print(io, m)
end

"""
Construct interpolation nodes, given SplineParams

Note that `p.k - 1` additional nodes will be inserted

##### Arguments

- `p::SplineParams`: `SplineParams` instance

##### Returns

- `x::Vector`: The Vector of 1d interpolation nodes.

"""
function nodes(p::SplineParams)
    breaks, evennum, k = p.breaks, p.evennum, p.k
    a = breaks[1]  # 20
    b = breaks[end]  # 21
    n = length(breaks) + k - 1  # 22
    x = cumsum(vcat(fill(a, k), breaks, fill(b, k)))  # 23
    x = (x[1+k:n+k] - x[1:n]) / k  # 24
    x[1] = a  # 25
    x[end] = b  # 26
    x
end

# TODO: define method derivative_op(::Type{SplineSparse}, p::SplineParams, order::Int)
function derivative_op(p::SplineParams, x::AbstractArray, order=1)
    breaks, evennum, k = p.breaks, p.evennum, p.k

    any(order .> k) && error("Order of differentiation can't be greater than k")

    # 38-40
    n = length(breaks) + k - 1
    kk = max(k - 1, k - order - 1)
    augbreaks = vcat(fill(breaks[1], kk), breaks, fill(breaks[end], kk))

    D = Array{SparseMatrixCSC{basis_eltype(p, x),Int64}}(abs(order), 1)

    if order > 0  # derivative
        temp = k ./ (augbreaks[k+1:n+k-1] - augbreaks[1:n-1])
        D[1] = spdiagm((-temp, temp), 0:1, n-1, n)

        for i in 2:order
            temp = (k+1-i) ./ (augbreaks[k+1:n+k-i] - augbreaks[i:n-1])
            D[i] = spdiagm((-temp, temp), 0:1, n-i, n+1-i)*D[i-1]
        end
    else
        error("not implemented")
    end

    D, SplineParams(breaks, evennum, k-order)
end

function _chk_evalbase(p::SplineParams, x::AbstractArray, order)
    breaks, evennum, k = p.breaks, p.evennum, p.k

    # error handling
    k < 0 && error("spline order must be positive")
    !(issorted(breaks)) && error("Breakpoints must be non-decreasing")
    any(order .>= k) && error("Order of differentiation must be less than k")
    size(x, 2) > 1 && error("x must be a column vector")

    m = size(x, 1)  # 54
    minorder = minimum(order)  # 55

    # Augment the breakpoint sequence 57-59
    n = length(breaks)+k-1
    a = breaks[1]
    b = breaks[end]
    augbreaks = vcat(fill(a, k-minorder), breaks, fill(b, k-minorder))

    ind = lookup(augbreaks, x, 3)  # 69

    n, m, minorder, augbreaks, ind
end

function evalbase(p::SplineParams, x::AbstractArray=nodes(p), order::Int=0)
    evalbase(p, x, [order])[1]
end

function evalbase(
        ::Type{SplineSparse}, p::SplineParams,
        x::AbstractArray=nodes(p), order::Int=0
    )
    evalbase(SplineSparse, p, x, [order])[1]
end

"""
Evaluate spline basis matrices for a certain order derivative at x

##### Arguments

- `p::SplineParams`: A `SplineParams` summarizing spline properties
- `x(nodes(p))` : the nodes at which to evaluate the basis matrices
- `order(0)` : The order(s) of derivative for which to evaluate the basis
matrices. `order=0` corresponds to the function itself, negative numbers
correspond to integrals.

##### Returns

- `B::SparseMatrixCSC` : Matrix containing the evaluation of basis functions
at each point in `x`. Each column represents a basis function.
- `x`: Points at which the functions were evaluated

"""
function evalbase(p::SplineParams, x::AbstractArray, order::AbstractVector{Int})
    n, m, minorder, augbreaks, ind = _chk_evalbase(p, x, order)

    max_repeat = p.k-minorder + 1
    T = basis_eltype(p, x)
    bas = zeros(T, m, max_repeat)  # 73
    bas[:, 1] = one(T)  # 74
    B = Array{SparseMatrixCSC{T,Int}}(length(order))  # 75

    # 76
    if maximum(order) > 0
        D = derivative_op(p, x, maximum(order))[1]
    end

    # 77
    if minorder < 0
        I = derivative_op(p, x, minorder)[1]
    end

    # We know what the rows and columns will be, we just compute them once and
    # then extract chunks as we need them.
    # note: I benchmarked and the version with collect was (2-5)x faster
    r = repmat(collect(1:m), max_repeat)
    c = (minorder-p.k:0)' .+ ind

    for j in 1:p.k-minorder  # 78
        for jj in j:-1:1  # 79
            for ix in eachindex(ind)
                b0 = augbreaks[ind[ix]+jj-j]  # 80
                b1 = augbreaks[ind[ix]+jj]  # 81
                temp = bas[ix, jj] / (b1 - b0)  # 82
                bas[ix, jj+1] = (x[ix] - b0) * temp + bas[ix, jj+1]  # 83
                bas[ix, jj] = (b1-x[ix]) * temp  # 84
            end
        end

        # bas now contains the order `j` spline basis
        ii = findfirst(order, p.k-j)  # 87
        if ii > 0
            # Put values in appropriate columns of a sparse matrix
            N = m * (p.k-order[ii]+1)
            B[ii] = sparse(view(r, 1:N), view(c, 1:N), view(bas, 1:N), m, n-order[ii])

            # 96-100
            if order[ii] > 0
                B[ii] = B[ii] * D[order[ii]]
            elseif order[ii] < 0
                B[ii] = B[ii] * I[-order[ii]]
            end
        end
    end

    B
end

function evalbase(
        ::Type{SplineSparse}, p::SplineParams, x::AbstractArray,
        order::AbstractVector{Int}
    )
    n, m, minorder, augbreaks, ind = _chk_evalbase(p, x, order)

    max_repeat = p.k-minorder + 1
    T = basis_eltype(p, x)
    bas = zeros(T, max_repeat, m)
    bas[1, :] = one(T)
    B = Array{SplineSparse{T,Int}}(length(order))  # 75

    if maximum(order) > 0
        D = derivative_op(p, x, maximum(order))[1]
    end

    if minorder < 0
        error("not supported yet")
        I = derivative_op(p, x, minorder)[1]
    end

    for j in 1:p.k-minorder  # 78
        for ix in eachindex(ind)
            for jj in j:-1:1  # 79
                b0 = augbreaks[ind[ix]+jj-j]  # 80
                b1 = augbreaks[ind[ix]+jj]  # 81
                temp = bas[jj, ix] / (b1 - b0)  # 82
                bas[jj+1, ix] = (x[ix] - b0) * temp + bas[jj+1, ix]  # 83
                bas[jj, ix] = (b1-x[ix]) * temp  # 84
            end
        end

        # bas now contains the order `j` spline basis
        ii = findfirst(order, p.k-j)  # 87
        if ii > 0
            ord = order[ii]
            cols = (ord - p.k) + ind
            if ord > 0
                D_ord = D[ord]
                new_vals = zeros(T, m * (p.k+1))
                ix = 0
                for row in 1:m  # row, col are rows/columns of the output
                    start_bas_col = cols[row]
                    end_bas_col = start_bas_col + p.k - ord
                    for col in 1:n
                        colptr_start = D_ord.colptr[col]
                        colptr_end = D_ord.colptr[col+1] - 1
                        start_D_row = D_ord.rowval[colptr_start]
                        end_D_row = D_ord.rowval[colptr_end]

                        # stop end_D_row is smaller than start_bas_col
                        end_D_row < start_bas_col && continue

                        # also stop if end_bas_col is smaller than start_D_row
                        end_bas_col < start_D_row && continue

                        # otherwise, we have some overlap
                        ix += 1

                        for D_ptr in colptr_start:colptr_end
                            D_row = D_ord.rowval[D_ptr]
                            if D_row < start_bas_col  # haven't made it to overlap
                                continue
                            elseif D_row > end_bas_col  # gone past overlap
                                break
                            else  # in the overlap
                                bas_ix = D_row - start_bas_col + 1
                                new_vals[ix] += D_ord.nzval[D_ptr] * bas[bas_ix, row]
                            end
                        end
                    end
                end

                B[ii] = SplineSparse{T,Int,1,p.k+1}(n, new_vals, cols)
            elseif ord == 0

                B[ii] = SplineSparse{T,Int,1,p.k+1}(n-ord, vec(bas), cols)
            else
                error("Shouldn't be here. Please file an issue")
            end
        end
    end
    return B
end


# ------------------------------------------------------------------------------------------
# # Optimizing DiffEq Code
#
# In this notebook we will walk through some of the main tools for optimizing your code in
# order to efficiently solve DifferentialEquations.jl. User-side optimizations are important
# because, for sufficiently difficult problems, most of the time will be spent inside of
# your `f` function, the function you are trying to solve. "Efficient" integrators are those
# that reduce the required number of `f` calls to hit the error tolerance. The main ideas
# for optimizing your DiffEq code, or any Julia function, are the following:
#
# - Make it non-allocating
# - Use StaticArrays for small arrays
# - Use broadcast fusion
# - Make it type-stable
# - Reduce redundant calculations
# - Make use of BLAS calls
# - Optimize algorithm choice
#
# We'll discuss these strategies in the context of small and large systems. Let's start with
# small systems.
#
# ## Optimizing Small Systems (<100 DEs)
#
# Let's take the classic Lorenz system from before. Let's start by naively writing the
# system in its out-of-place form:
# ------------------------------------------------------------------------------------------

function lorenz(u,p,t)
 dx = 10.0*(u[2]-u[1])
 dy = u[1]*(28.0-u[3]) - u[2]
 dz = u[1]*u[2] - (8/3)*u[3]
 [dx,dy,dz]
end

# ------------------------------------------------------------------------------------------
# Here, `lorenz` returns an object, `[dx,dy,dz]`, which is created within the body of
# `lorenz`.
#
# This is a common code pattern from high-level languages like MATLAB, SciPy, or R's
# deSolve. However, the issue with this form is that it allocates a vector, `[dx,dy,dz]`, at
# each step. Let's benchmark the solution process with this choice of function:
# ------------------------------------------------------------------------------------------

using DifferentialEquations, BenchmarkTools
u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz,u0,tspan)
@benchmark solve(prob,Tsit5())

# ------------------------------------------------------------------------------------------
# The BenchmarkTools package's `@benchmark` runs the code multiple times to get an accurate
# measurement. The minimum time is the time it takes when your OS and other background
# processes aren't getting in the way. Notice that in this case it takes about 5ms to solve
# and allocates around 11.11 MiB. However, if we were to use this inside of a real user code
# we'd see a lot of time spent doing garbage collection (GC) to clean up all of the arrays
# we made. Even if we turn off saving we have these allocations.
# ------------------------------------------------------------------------------------------

@benchmark solve(prob,Tsit5(),save_everystep=false)

# ------------------------------------------------------------------------------------------
# The problem of course is that arrays are created every time our derivative function is
# called. This function is called multiple times per step and is thus the main source of
# memory usage. To fix this, we can use the in-place form to ***make our code non-
# allocating***:
# ------------------------------------------------------------------------------------------

function lorenz!(du,u,p,t)
 du[1] = 10.0*(u[2]-u[1])
 du[2] = u[1]*(28.0-u[3]) - u[2]
 du[3] = u[1]*u[2] - (8/3)*u[3]
end

# ------------------------------------------------------------------------------------------
# Here, instead of creating an array each time, we utilized the cache array `du`. When the
# inplace form is used, DifferentialEquations.jl takes a different internal route that
# minimizes the internal allocations as well. When we benchmark this function, we will see
# quite a difference.
# ------------------------------------------------------------------------------------------

u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz!,u0,tspan)
@benchmark solve(prob,Tsit5())

@benchmark solve(prob,Tsit5(),save_everystep=false)

# ------------------------------------------------------------------------------------------
# There is a 4x time difference just from that change! Notice there are still some
# allocations and this is due to the construction of the integration cache. But this doesn't
# scale with the problem size:
# ------------------------------------------------------------------------------------------

tspan = (0.0,500.0) # 5x longer than before
prob = ODEProblem(lorenz!,u0,tspan)
@benchmark solve(prob,Tsit5(),save_everystep=false)

# ------------------------------------------------------------------------------------------
# since that's all just setup allocations.
#
# #### But if the system is small we can optimize even more.
#
# Allocations are only expensive if they are "heap allocations". For a more in-depth
# definition of heap allocations, [there are a lot of sources online](http://net-
# informations.com/faq/net/stack-heap.htm). But a good working definition is that heap
# allocations are variable-sized slabs of memory which have to be pointed to, and this
# pointer indirection costs time. Additionally, the heap has to be managed and the garbage
# controllers has to actively keep track of what's on the heap.
#
# However, there's an alternative to heap allocations, known as stack allocations. The stack
# is statically-sized (known at compile time) and thus its accesses are quick. Additionally,
# the exact block of memory is known in advance by the compiler, and thus re-using the
# memory is cheap. This means that allocating on the stack has essentially no cost!
#
# Arrays have to be heap allocated because their size (and thus the amount of memory they
# take up) is determined at runtime. But there are structures in Julia which are stack-
# allocated. `struct`s for example are stack-allocated "value-type"s. `Tuple`s are a stack-
# allocated collection. The most useful data structure for DiffEq though is the
# `StaticArray` from the package
# [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl). These arrays have their
# length determined at compile-time. They are created using macros attached to normal array
# expressions, for example:
# ------------------------------------------------------------------------------------------

using StaticArrays
A = @SVector [2.0,3.0,5.0]

# ------------------------------------------------------------------------------------------
# Notice that the `3` after `SVector` gives the size of the `SVector`. It cannot be changed.
# Additionally, `SVector`s are immutable, so we have to create a new `SVector` to change
# values. But remember, we don't have to worry about allocations because this data structure
# is stack-allocated. `SArray`s have a lot of extra optimizations as well: they have fast
# matrix multiplication, fast QR factorizations, etc. which directly make use of the
# information about the size of the array. Thus, when possible they should be used.
#
# Unfortunately static arrays can only be used for sufficiently small arrays. After a
# certain size, they are forced to heap allocate after some instructions and their compile
# time balloons. Thus static arrays shouldn't be used if your system has more than 100
# variables. Additionally, only the native Julia algorithms can fully utilize static arrays.
#
# Let's ***optimize `lorenz` using static arrays***. Note that in this case, we want to use
# the out-of-place allocating form, but this time we want to output a static array:
# ------------------------------------------------------------------------------------------

function lorenz_static(u,p,t)
 dx = 10.0*(u[2]-u[1])
 dy = u[1]*(28.0-u[3]) - u[2]
 dz = u[1]*u[2] - (8/3)*u[3]
 @SVector [dx,dy,dz]
end

# ------------------------------------------------------------------------------------------
# To make the solver internally use static arrays, we simply give it a static array as the
# initial condition:
# ------------------------------------------------------------------------------------------

u0 = @SVector [1.0,0.0,0.0]
tspan = (0.0,100.0)
prob = ODEProblem(lorenz_static,u0,tspan)
@benchmark solve(prob,Tsit5())

@benchmark solve(prob,Tsit5(),save_everystep=false)

# ------------------------------------------------------------------------------------------
# And that's pretty much all there is to it. With static arrays you don't have to worry
# about allocating, so use operations like `*` and don't worry about fusing operations
# (discussed in the next section). Do "the vectorized code" of R/MATLAB/Python and your code
# in this case will be fast, or directly use the numbers/values.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Exercise 1
#
# Implement the out-of-place array, in-place array, and out-of-place static array forms for
# the [Henon-Heiles System](https://en.wikipedia.org/wiki/H%C3%A9non%E2%80%93Heiles_system)
# and time the results.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Optimizing Large Systems
#
# ### Interlude: Managing Allocations with Broadcast Fusion
#
# When your system is sufficiently large, or you have to make use of a non-native Julia
# algorithm, you have to make use of `Array`s. In order to use arrays in the most efficient
# manner, you need to be careful about temporary allocations. Vectorized calculations
# naturally have plenty of temporary array allocations. This is because a vectorized
# calculation outputs a vector. Thus:
# ------------------------------------------------------------------------------------------

A = rand(1000,1000); B = rand(1000,1000); C = rand(1000,1000)
test(A,B,C) = A + B + C
@benchmark test(A,B,C)

# ------------------------------------------------------------------------------------------
# That expression `A + B + C` creates 2 arrays. It first creates one for the output of `A +
# B`, then uses that result array to `+ C` to get the final result. 2 arrays! We don't want
# that! The first thing to do to fix this is to use broadcast fusion. [Broadcast
# fusion](https://julialang.org/blog/2017/01/moredots) puts expressions together. For
# example, instead of doing the `+` operations separately, if we were to add them all at the
# same time, then we would only have a single array that's created. For example:
# ------------------------------------------------------------------------------------------

test2(A,B,C) = map((a,b,c)->a+b+c,A,B,C)
@benchmark test2(A,B,C)

# ------------------------------------------------------------------------------------------
# Puts the whole expression into a single function call, and thus only one array is required
# to store output. This is the same as writing the loop:
# ------------------------------------------------------------------------------------------

function test3(A,B,C)
    D = similar(A)
    @inbounds for i in eachindex(A)
        D[i] = A[i] + B[i] + C[i]
    end
    D
end
@benchmark test3(A,B,C)

# ------------------------------------------------------------------------------------------
# However, Julia's broadcast is syntactic sugar for this. If multiple expressions have a
# `.`, then it will put those vectorized operations together. Thus:
# ------------------------------------------------------------------------------------------

test4(A,B,C) = A .+ B .+ C
@benchmark test4(A,B,C)

# ------------------------------------------------------------------------------------------
# is a version with only 1 array created (the output). Note that `.`s can be used with
# function calls as well:
# ------------------------------------------------------------------------------------------

sin.(A) .+ sin.(B)

# ------------------------------------------------------------------------------------------
# Also, the `@.` macro applys a dot to every operator:
# ------------------------------------------------------------------------------------------

test5(A,B,C) = @. A + B + C #only one array allocated
@benchmark test5(A,B,C)

# ------------------------------------------------------------------------------------------
# Using these tools we can get rid of our intermediate array allocations for many vectorized
# function calls. But we are still allocating the output array. To get rid of that
# allocation, we can instead use mutation. Mutating broadcast is done via `.=`. For example,
# if we pre-allocate the output:
# ------------------------------------------------------------------------------------------

D = zeros(1000,1000);

# ------------------------------------------------------------------------------------------
# Then we can keep re-using this cache for subsequent calculations. The mutating
# broadcasting form is:
# ------------------------------------------------------------------------------------------

test6!(D,A,B,C) = D .= A .+ B .+ C #only one array allocated
@benchmark test6!(D,A,B,C)

# ------------------------------------------------------------------------------------------
# If we use `@.` before the `=`, then it will turn it into `.=`:
# ------------------------------------------------------------------------------------------

test7!(D,A,B,C) = @. D = A + B + C #only one array allocated
@benchmark test7!(D,A,B,C)

# ------------------------------------------------------------------------------------------
# Notice that in this case, there is no "output", and instead the values inside of `D` are
# what are changed (like with the DiffEq inplace function). Many Julia functions have a
# mutating form which is denoted with a `!`. For example, the mutating form of the `map` is
# `map!`:
# ------------------------------------------------------------------------------------------

test8!(D,A,B,C) = map!((a,b,c)->a+b+c,D,A,B,C)
@benchmark test8!(D,A,B,C)

# ------------------------------------------------------------------------------------------
# Some operations require using an alternate mutating form in order to be fast. For example,
# matrix multiplication via `*` allocates a temporary:
# ------------------------------------------------------------------------------------------

@benchmark A*B

# ------------------------------------------------------------------------------------------
# Instead, we can use the mutating form `A_mul_B!` into a cache array to avoid allocating
# the output:
# ------------------------------------------------------------------------------------------

@benchmark A_mul_B!(D,A,B) # same as D = A * B

# ------------------------------------------------------------------------------------------
# For repeated calculations this reduced allocation can stop GC cycles and thus lead to more
# efficient code. Additionally, ***we can fuse together higher level linear algebra
# operations using BLAS***. The package
# [SugarBLAS.jl](https://github.com/lopezm94/SugarBLAS.jl) makes it easy to write higher
# level operations like `alpha*B*A + beta*C` as mutating BLAS calls.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### Example Optimization: Gierer-Meinhardt Reaction-Diffusion PDE Discretization
#
# Let's optimize the solution of a Reaction-Diffusion PDE's discretization. In its
# discretized form, this is the ODE:
#
# $$ du = D_1 (A_y u + u A_x) + \frac{au^2}{v} + \bar{u} - \alpha u $$
# $$ dv = D_2 (A_y v + v A_x) + a u^2 + \beta v $$
#
# where $u$, $v$, and $A$ are matrices. Here, we will use the simplified version where $A$
# is the tridiagonal stencil $[1,-2,1]$, i.e. it's the 2D discretization of the LaPlacian.
# The native code would be something along the lines of:
# ------------------------------------------------------------------------------------------

# Generate the constants
p = (1.0,1.0,1.0,10.0,0.001,100.0) # a,α,ubar,β,D1,D2
N = 100
Ax = full(Tridiagonal([1.0 for i in 1:N-1],[-2.0 for i in 1:N],[1.0 for i in 1:N-1]))
Ay = copy(Ax)
Ax[2,1] = 2.0
Ax[end-1,end] = 2.0
Ay[1,2] = 2.0
Ay[end,end-1] = 2.0

function basic_version!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = r[:,:,1]
  v = r[:,:,2]
  Du = D1*(Ay*u + u*Ax)
  Dv = D2*(Ay*v + v*Ax)
  dr[:,:,1] = Du + a*u*u./v + ubar - α*u
  dr[:,:,2] = Dv + a*u*u - β*v
end

a,α,ubar,β,D1,D2 = p
uss = (ubar+β)/α
vss = (a/β)*uss^2
r0 = zeros(100,100,2)
r0[:,:,1] = uss+0.1*rand()
r0[:,:,2] = vss

prob = ODEProblem(basic_version!,r0,(0.0,0.1),p)

# ------------------------------------------------------------------------------------------
# In this version we have encoded our initial condition to be a 3-dimensional array, with
# `u[:,:,1]` being the `A` part and `u[:,:,2]` being the `B` part.
# ------------------------------------------------------------------------------------------

@benchmark solve(prob,Tsit5())

# ------------------------------------------------------------------------------------------
# While this version isn't very efficient,
#
# #### We recommend writing the "high-level" code first, and iteratively optimizing it!
#
# The first thing that we can do is get rid of the slicing allocations. The operation
# `r[:,:,1]` creates a temporary array instead of a "view", i.e. a pointer to the already
# existing memory. To make it a view, add `@view`. Note that we have to be careful with
# views because they point to the same memory, and thus changing a view changes the original
# values:
# ------------------------------------------------------------------------------------------

A = rand(4)
@show A
B = @view A[1:3]
B[2] = 2
@show A

# ------------------------------------------------------------------------------------------
# Notice that changing `B` changed `A`. This is something to be careful of, but at the same
# time we want to use this since we want to modify the output `dr`. Additionally, the last
# statement is a purely element-wise operation, and thus we can make use of broadcast fusion
# there. Let's rewrite `basic_version!` to ***avoid slicing allocations*** and to ***use
# broadcast fusion***:
# ------------------------------------------------------------------------------------------

function gm2!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = @view r[:,:,1]
  v = @view r[:,:,2]
  du = @view dr[:,:,1]
  dv = @view dr[:,:,2]
  Du = D1*(Ay*u + u*Ax)
  Dv = D2*(Ay*v + v*Ax)
  @. du = Du + a*u*u./v + ubar - α*u
  @. dv = Dv + a*u*u - β*v
end
prob = ODEProblem(gm2!,r0,(0.0,0.1),p)
@benchmark solve(prob,Tsit5())

# ------------------------------------------------------------------------------------------
# Now, most of the allocations are taking place in `Du = D1*(Ay*u + u*Ax)` since those
# operations are vectorized and not mutating. We should instead replace the matrix
# multiplications with `A_mul_B!`. When doing so, we will need to have cache variables to
# write into. This looks like:
# ------------------------------------------------------------------------------------------

Ayu = zeros(N,N)
uAx = zeros(N,N)
Du = zeros(N,N)
Ayv = zeros(N,N)
vAx = zeros(N,N)
Dv = zeros(N,N)
function gm3!(dr,r,p,t)
  a,α,ubar,β,D1,D2 = p
  u = @view r[:,:,1]
  v = @view r[:,:,2]
  du = @view dr[:,:,1]
  dv = @view dr[:,:,2]
  A_mul_B!(Ayu,Ay,u)
  A_mul_B!(uAx,u,Ax)
  A_mul_B!(Ayv,Ay,v)
  A_mul_B!(vAx,v,Ax)
  @. Du = D1*(Ayu + uAx)
  @. Dv = D2*(Ayv + vAx)
  @. du = Du + a*u*u./v + ubar - α*u
  @. dv = Dv + a*u*u - β*v
end
prob = ODEProblem(gm3!,r0,(0.0,0.1),p)
@benchmark solve(prob,Tsit5())

# ------------------------------------------------------------------------------------------
# But our temporary variables are global variables. We need to either declare the caches as
# `const` or localize them. We can localize them by adding them to the parameters, `p`. It's
# easier for the compiler to reason about local variables than global variables.
# ***Localizing variables helps to ensure type stability***.
# ------------------------------------------------------------------------------------------

p = (1.0,1.0,1.0,10.0,0.001,100.0,Ayu,uAx,Du,Ayv,vAx,Dv) # a,α,ubar,β,D1,D2
function gm4!(dr,r,p,t)
  a,α,ubar,β,D1,D2,Ayu,uAx,Du,Ayv,vAx,Dv = p
  u = @view r[:,:,1]
  v = @view r[:,:,2]
  du = @view dr[:,:,1]
  dv = @view dr[:,:,2]
  A_mul_B!(Ayu,Ay,u)
  A_mul_B!(uAx,u,Ax)
  A_mul_B!(Ayv,Ay,v)
  A_mul_B!(vAx,v,Ax)
  @. Du = D1*(Ayu + uAx)
  @. Dv = D2*(Ayv + vAx)
  @. du = Du + a*u*u./v + ubar - α*u
  @. dv = Dv + a*u*u - β*v
end
prob = ODEProblem(gm4!,r0,(0.0,0.1),p)
@benchmark solve(prob,Tsit5())

# ------------------------------------------------------------------------------------------
# We could then use the BLAS `gemmv` to optimize the matrix multiplications some more, but
# instead let's devectorize the stencil.
# ------------------------------------------------------------------------------------------

p = (1.0,1.0,1.0,10.0,0.001,100.0,N)
function fast_gm!(du,u,p,t)
  a,α,ubar,β,D1,D2,N = p

  @inbounds for j in 2:N-1, i in 2:N-1
    du[i,j,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end

  @inbounds for j in 2:N-1, i in 2:N-1
    du[i,j,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
            a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds for j in 2:N-1
    i = 1
    du[1,j,1] = D1*(2u[i+1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
            a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for j in 2:N-1
    i = 1
    du[1,j,2] = D2*(2u[i+1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
            a*u[i,j,1]^2 - β*u[i,j,2]
  end
  @inbounds for j in 2:N-1
    i = N
    du[end,j,1] = D1*(2u[i-1,j,1] + u[i,j+1,1] + u[i,j-1,1] - 4u[i,j,1]) +
           a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for j in 2:N-1
    i = N
    du[end,j,2] = D2*(2u[i-1,j,2] + u[i,j+1,2] + u[i,j-1,2] - 4u[i,j,2]) +
           a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds for i in 2:N-1
    j = 1
    du[i,1,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for i in 2:N-1
    j = 1
    du[i,1,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
              a*u[i,j,1]^2 - β*u[i,j,2]
  end
  @inbounds for i in 2:N-1
    j = N
    du[i,end,1] = D1*(u[i-1,j,1] + u[i+1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
  end
  @inbounds for i in 2:N-1
    j = N
    du[i,end,2] = D2*(u[i-1,j,2] + u[i+1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]
  end

  @inbounds begin
    i = 1; j = 1
    du[1,1,1] = D1*(2u[i+1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
              a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[1,1,2] = D2*(2u[i+1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
              a*u[i,j,1]^2 - β*u[i,j,2]

    i = 1; j = N
    du[1,N,1] = D1*(2u[i+1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[1,N,2] = D2*(2u[i+1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]

    i = N; j = 1
    du[N,1,1] = D1*(2u[i-1,j,1] + 2u[i,j+1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[N,1,2] = D2*(2u[i-1,j,2] + 2u[i,j+1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]

    i = N; j = N
    du[end,end,1] = D1*(2u[i-1,j,1] + 2u[i,j-1,1] - 4u[i,j,1]) +
             a*u[i,j,1]^2/u[i,j,2] + ubar - α*u[i,j,1]
    du[end,end,2] = D2*(2u[i-1,j,2] + 2u[i,j-1,2] - 4u[i,j,2]) +
             a*u[i,j,1]^2 - β*u[i,j,2]
   end
end
prob = ODEProblem(fast_gm!,r0,(0.0,0.1),p)
@benchmark solve(prob,Tsit5())

# ------------------------------------------------------------------------------------------
# Lastly, we can do other things like multithread the main loops, but these optimizations
# get the last 2x-3x out. The main optimizations which apply everywhere are the ones we just
# performed (though the last one only works if your matrix is a stencil. This is known as a
# matrix-free implementation of the PDE discretization).
#
# This gets us to about 20,000x faster than our original MATLAB/SciPy/R vectorized style
# code!
#
# The last thing to do is then ***optimize our algorithm choice***. We have been using
# `Tsit5()` as our test algorithm, but in reality this problem is a stiff PDE discretization
# and thus one recommendation is to use `CVODE_BDF()`. However, instead of using the default
# dense Jacobian, we should make use of the sparse Jacobian afforded by the problem. The
# Jacobian is the matrix $\frac{df_i}{dr_j}$, where $r$ is read by the linear index (i.e.
# down columns). But since the $u$ variables depend on the $v$, the band size here is large,
# and thus this will not do well with a Banded Jacobian solver. Instead, we utilize sparse
# Jacobian algorithms. `CVODE_BDF` allows us to use a sparse Newton-Krylov solver by setting
# `linear_solver = :GMRES` (see [the solver
# documentation](http://docs.juliadiffeq.org/latest/solvers/ode_solve.html#Sundials.jl-1),
# and thus we can solve this problem efficiently. Let's see how this scales as we increase
# the integration time.
# ------------------------------------------------------------------------------------------

prob = ODEProblem(fast_gm!,r0,(0.0,10.0),p)
@benchmark solve(prob,Tsit5())

@benchmark solve(prob,CVODE_BDF(linear_solver=:GMRES))

prob = ODEProblem(fast_gm!,r0,(0.0,100.0),p)
@benchmark solve(prob,Tsit5())

@benchmark solve(prob,CVODE_BDF(linear_solver=:GMRES))

# ------------------------------------------------------------------------------------------
# What's happening is that, because the problem is stiff, the number of steps required by
# the explicit Runge-Kutta method grows rapidly, whereas `CVODE_BDF` is taking large steps.
# Additionally, the `GMRES` linear solver form is quite an efficient way to solve the
# implicit system in this case. This is problem-dependent, and in many cases using a Krylov
# method effectively requires a preconditioner, so you need to play around with testing
# other algorithms and linear solvers to find out what works best with your problem.
#
# ## Conclusion
#
# Julia gives you the tools to optimize the solver "all the way", but you need to make use
# of it. The main thing to avoid is temporary allocations. For small systems, this is
# effectively done via static arrays. For large systems, this is done via in-place
# operations and cache arrays. Either way, the resulting solution can be immensely sped up
# over vectorized formulations by using these principles.
# ------------------------------------------------------------------------------------------


function extendSeries(X::Array{Float64}, method::String, length::String, n::Int)

  (N, ) = size(X)

  # determine final length 'n' of series after extension
  if (length == "arbitrary")
    if (n <= N)
     error("Invalid argument: 'n' must be greater than length of series when length='arbitrary'.")
    end
  elseif (length == "powerof2")
    k = N/(2^n)
    if (round(k) == k)
      error("Invalid argument: length of series should not be multiple of 2^j when length='powerof2'.")
    else
      n = ceil(k)*2^n
    end
  elseif (length == "double")
    error("The integer parameter is arbitrary, you probably made some error.")
  end

  if isa(X, Vector)
    # extend the series to length 'n'
    (method == "periodic") ? X = repeat(X, outer=[convert(Int,ceil(n/N))])[1:n] : ()
    (method == "reflection") ? X = repeat([X,X[end:-1:1]], outer=[convert(Int,ceil(n/N))])[1:n] : ()
    (method == "zeros") ? X = [X, zeros(N-n)] : ()
    (method == "mean") ? X = [X, mean(X).*ones(N-n)] : ()
    # if(method == "reflection.inverse") X = apply(X, 2, function(x,n,N){rep(c(x,2*x[N]-rev(x)),length=n)}, n=n, N=N)
  else
    # extend the series to length 'n'
    (method == "periodic") ? X = mapslices(x -> repeat(x, outer=[convert(Int,ceil(n/N))])[1:n], X, 1) : ()
    (method == "reflection") ? X = mapslices(x -> repeat(x, outer=[convert(Int,ceil(n/N))])[1:n], [X,X[end:-1:1]], 2) : ()
    (method == "zeros") ? X = mapslices(x -> [x, zeros(N-n)], X, 2) : ()
    (method == "mean") ? X = mapslices(x -> [x, mean(x).*ones(N-n)], X, 2) : ()
    # if(method == "reflection.inverse") X = apply(X, 2, function(x,n,N){rep(c(x,2*x[N]-rev(x)),length=n)}, n=n, N=N)
  end

  return X
end

function extendSeries(X::Array{Float64}, method::String, length::String)
  (N, ) = size(X)

  if length == "double"
    n = 2*N
  else
    error("The selected method does not correspond, you probably forgot to include another parameter.")
  end

  if isa(X, Vector)
    # extend the series to length 'n'
    (method == "periodic") ? X = repeat(X, outer=[convert(Int,ceil(n/N))])[1:n] : ()
    (method == "reflection") ? X = repeat([X,X[end:-1:1]], outer=[convert(Int,ceil(n/N))])[1:n] : ()
    (method == "zeros") ? X = [X, zeros(N-n)] : ()
    (method == "mean") ? X = [X, mean(X).*ones(N-n)] : ()
    # if(method == "reflection.inverse") X = apply(X, 2, function(x,n,N){rep(c(x,2*x[N]-rev(x)),length=n)}, n=n, N=N)
  else
    # extend the series to length 'n'
    (method == "periodic") ? X = mapslices(x -> repeat(x, outer=[convert(Int,ceil(n/N))])[1:n], X, 1) : ()
    (method == "reflection") ? X = mapslices(x -> repeat(x, outer=[convert(Int,ceil(n/N))])[1:n], [X,X[end:-1:1]], 2) : ()
    (method == "zeros") ? X = mapslices(x -> [x, zeros(N-n)], X, 2) : ()
    (method == "mean") ? X = mapslices(x -> [x, mean(x).*ones(N-n)], X, 2) : ()
    # if(method == "reflection.inverse") X = apply(X, 2, function(x,n,N){rep(c(x,2*x[N]-rev(x)),length=n)}, n=n, N=N)
  end

  return X
end


# This file is part of the IntervalArithmetic.jl package; MIT licensed

"""An `IntervalBox` is an `N`-dimensional rectangular box, given
by a Cartesian product of a vector of `N` `Interval`s.
"""
struct IntervalBox{N,T}
    v::SVector{N, Interval{T}}
end

# IntervalBox(x::Interval) = IntervalBox( SVector(x) )  # single interval treated as tuple with one element

IntervalBox(x::Interval...) = IntervalBox(SVector(x))
IntervalBox(x::SVector) = IntervalBox(interval.(x))
IntervalBox(x::Tuple) = IntervalBox(SVector(x))
IntervalBox(x::Real) = IntervalBox(interval.(x))
IntervalBox(x...) = IntervalBox(x)
IntervalBox(x) = IntervalBox(x...)
IntervalBox(X::IntervalBox, n) = foldl(×, Iterators.repeated(X, n))

Base.@propagate_inbounds Base.getindex(X::IntervalBox, i) = X.v[i]

setindex(X::IntervalBox, y, i) = IntervalBox( setindex(X.v, y, i) )

# iteration:

start(X::IntervalBox{N,T}) where {N,T} = 1

next(X::IntervalBox{N,T}, state) where {N,T} = (X[state], state+1)

done(X::IntervalBox{N,T}, state) where {N,T} = state > N

eltype(::Type{IntervalBox{N,T}}) where {N,T} = Interval{T} # Note that this is defined for the type

length(X::IntervalBox{N,T}) where {N,T} = N

# IntervalBox(xx) = IntervalBox(Interval.(xx))
# IntervalBox(xx::SVector) where {N,T} = IntervalBox(Interval.(xx))

## arithmetic operations
# Note that standard arithmetic operations are implemented automatically by FixedSizeArrays.jl
"""
    mid(X::IntervalBox, α=0.5)

Return a vector of the `mid` of each interval composing the `IntervalBox`.

See `mid(X::Interval, α=0.5)` for more informations.
"""
mid(X::IntervalBox) = mid.(X.v)
mid(X::IntervalBox, α) = mid.(X.v, α)

big(X::IntervalBox) = big.(X)

## set operations

# TODO: Update to use generator
⊆(X::IntervalBox{N,T}, Y::IntervalBox{N,T}) where {N,T} =
    all(X.v .⊆ Y.v)

∩(X::IntervalBox{N,T}, Y::IntervalBox{N,T}) where {N,T} =
    IntervalBox(X.v .∩ Y.v)
∪(X::IntervalBox{N,T}, Y::IntervalBox{N,T}) where {N,T} =
    IntervalBox(X.v .∪ Y.v)

#=
On Julia 0.6 can now write
∩{N,T}(X::IntervalBox{N,T}, Y::IntervalBox{N,T}) = IntervalBox(NTuple{N, Interval{Float64}}( (X[i] ∩ Y[i]) for i in 1:N))
=#

isempty(X::IntervalBox) = any(isempty, X.v)

diam(X::IntervalBox) = maximum(diam.(X.v))

emptyinterval(X::IntervalBox{N,T}) where {N,T} = IntervalBox(emptyinterval.(X.v))

isinf(X::IntervalBox) = any(isinf.(X))

isinterior(X::IntervalBox{N,T}, Y::IntervalBox{N,T}) where {N,T} = all(isinterior.(X, Y))

contains_zero(X::SVector) = all(contains_zero.(X))
contains_zero(X::IntervalBox) = all(contains_zero.(X))

×(a::Interval...) = IntervalBox(a...)
×(a::Interval, b::IntervalBox) = IntervalBox(a, b.v...)
×(a::IntervalBox, b::Interval) = IntervalBox(a.v..., b)
×(a::IntervalBox, b::IntervalBox) = IntervalBox(a.v..., b.v...)

IntervalBox(x::Interval, ::Type{Val{n}}) where {n} = IntervalBox(SVector(ntuple(i->x, Val{n})))

IntervalBox(x::Interval, n::Int) = IntervalBox(x, Val{n})

dot(x::IntervalBox, y::IntervalBox) = dot(x.v, y.v)


#=
using Revise
=#

include("HMC.jl")
using Distributions, Random
using Flux

using RCall
@rimport graphics as rgraphics
@rimport rcommon

struct State
  beta
  log_sig2
end

pretty(s::State) = "State(β => $(s.beta); σ => $(sqrt(exp(s.log_sig2[1]))))"

State(K::Int) = State(param(zeros(K)), param(ones(1)))

function sim_data(;N::Int=100, K::Int=10, sig=.3)
  X = randn(N, K)
  b = randn(K)
  y = X * b + randn(N) * sig
  return Dict(:X => X, :b => b, :y => y, :sig => sig)
end

function lpdf_std_normal(z)
  return -z ^ 2 / 2 - log(2 * pi) / 2
end

lpdf_normal(y, m, s) = lpdf_std_normal((y-m)/s) - log(s)

function loglike(X, y, state::State)
  # println(state)
  m = X * state.beta
  s = sum(sqrt.(exp.(state.log_sig2))) + 1e-6
  # s = 1.0
  ll = sum(lpdf_normal.(y, m, s))
  # println("ll: $(ll)")
  return ll
end

function logprior(state::State)
  lp_beta = sum(logpdf.(Normal(0, 10), state.beta))
  lp_sig2 = sum(logpdf.(InverseGamma(2, 1), exp.(state.log_sig2)))
  lp_sig2 += sum(state.log_sig2)  # add log abs jacobian
  lp = lp_beta + lp_sig2
  # println("lp: $(lp)")
  return lp
end

### MAIN ###
Random.seed!(1)

N, K = (1000, 20)
dat = sim_data(K=K)
log_prob(state::State) = loglike(dat[:X], dat[:y], state) + logprior(state)
state = State(K)
epsilon = 5 / N
num_leapfrog_steps = 100

# Compile
_ = HMC.hmc_update(state, log_prob, num_leapfrog_steps, epsilon)

function simulate(init; nburn=500, nsamps=1000)
  state = deepcopy(init)

  log_prob_hist = zeros(nsamps)
  samps = [state for i in 1:nsamps]
  for i in 1:(nburn + nsamps)
    print("\rProgress: $i / $(nburn + nsamps)")
    state, log_prob_curr = HMC.hmc_update(state, log_prob,
                                          num_leapfrog_steps, epsilon)
    if i > nburn
      samps[i - nburn] = state
      log_prob_hist[i - nburn] = log_prob_curr
    end
  end
  return samps, log_prob_hist
end

samps, log_prob_hist = simulate(state, nburn=500, nsamps=500)
println(pretty(samps[end]));
dat

acceptance_rate = length(unique(log_prob_hist)) / length(log_prob_hist)
println("Acceptance rate: $acceptance_rate")
rgraphics.plot(log_prob_hist, main="", xlab="iter", ylab="log prob", typ="l");

sig2_post = [exp(Tracker.data(s.log_sig2[1])) for s in samps]
b_first_2_post = cat([Tracker.data(s.beta[1:2])' for s in samps]..., dims=1)
rgraphics.plot(b_first_2_post, xlab="", ylab="", main="", typ="l");

rcommon.plotPosts(b_first_2_post);
rcommon.plotPost(sqrt.(sig2_post));
rgraphics.abline(v=dat[:sig], lwd=3, col="orange");


"""
	Implements all needed parts for constructing DeepSVDD model for detecting anomalies.
		(Depp One-Class Classification -> http://proceedings.mlr.press/v80/ruff18a/ruff18a.pdf)
	Code is inspired by original Pytorch implementation https://github.com/lukasruff/Deep-SVDD-PyTorch.

"""

mutable struct DeepSVDD
	encoder
	decoder
	objective
	c
	R
	nu
end

Flux.@functor DeepSVDD

function (svdd::DeepSVDD)(x)
	z = svdd.encoder(x)
	return z
end

function AE(svdd::DeepSVDD, x)
	z = svdd.encoder(x)
	x̂ = svdd.decoder(z)
	return x̂
end

mse_score(x̂,x;dims) = vec(Flux.sum((x̂ .- x).^2, dims=dims))

function anomaly_score(svdd::DeepSVDD, x)
	z = svdd(x)
	dist = mse_score(z, svdd.c, dims=1)
	#vec(Flux.sum((z .- svdd.c).^2, dims=1))
	score = (svdd.objective == "soft-boundary") ? dist .- svdd.R.^2 : dist
	return score
end

function init_c(svdd::DeepSVDD, data, eps=0.1; batchsize=50, to_testmode::Bool=true)
	(to_testmode == true) ? Flux.testmode!(svdd) : nothing
	c_ = deepcopy(svdd.c)
	for X in Flux.Data.DataLoader(data, batchsize=batchsize)
		X = getobs(X)|>gpu
		ŷ = svdd(X)
		c_ += Flux.sum(ŷ, dims=2)
	end
	c_ = c_ |> cpu
	c_ /= size(data, 4)
	c_[(abs.(c_) .< eps) .& (c_ .< 0)] .= -eps
	c_[(abs.(c_) .< eps) .& (c_ .> 0)] .= eps
	(to_testmode == true) ? Flux.testmode!(svdd, false) : nothing
	return c_
end

function get_radius(svdd::DeepSVDD, X)
	output = svdd(X)
	dist = vec(Flux.sum((output .- svdd.c).^2, dims=1))
	return Statistics.quantile(sqrt.(dist), 1-svdd.nu)
end

# LOSSES

# autoencoder loss function
ae_loss(ŷ, y; dims=(1,2,3)) = Flux.mean(Flux.sum((ŷ .- y).^2, dims=dims))

# objective / loss function for Soft-boundary "SVM"
function sb_loss(ŷ, c, R, nu)
	dist = Flux.sum((ŷ .- c).^2, dims=1)
	scores = dist .- R^2 #
	loss = R^2 + (1 / nu) * Flux.mean(max.(0, scores))
end

# objective / loss function for One-Class "SVM"
oc_loss(ŷ, c, R, nu) = Flux.mean(Flux.sum((ŷ .- c).^2, dims=1))

# anomaly score on GPU
function anomaly_score_gpu(svdd::DeepSVDD, real_input; batch_size=64, to_testmode::Bool=true)
	real_input = Flux.Data.DataLoader(real_input, batchsize=batch_size)
	(to_testmode == true) ? Flux.testmode!(svdd) : nothing
	svdd = svdd |> gpu
	scores = Array{Float32}([])
	for x in real_input
		x = x |> gpu
		score = anomaly_score(svdd, x)
		scores = cat(scores, score |> cpu, dims=1)
	end
	(to_testmode == true) ? Flux.testmode!(svdd, false) : nothing
	return scores
end

######################################################################################

function conv_svdd_constructor(
	;idim=(2,2,1),
	zdim::Int=1,
	activation="relu",
	kernelsizes=(1,1),
	channels=(1,1),
	scalings=(1,1),
	init_seed=nothing,
	objective="soft-boundary",
	batchnorm=false,
	nu::AbstractFloat = 0.1f0,
	R::AbstractFloat = 0f0,
	c = nothing,
	kwargs...
)
	# if seed is given, set it
	(init_seed !== nothing) ? Random.seed!(init_seed) : nothing
	# encoder
	encoder = Chain(
		conv_encoder(idim, zdim, kernelsizes, channels, scalings; activation=activation, batchnorm=batchnorm)...
	)
	# decoder
	decoder = Chain(
		conv_decoder(idim, zdim, reverse(kernelsizes), reverse(channels), reverse(scalings),
			activation=activation, vec_output=false, vec_output_dim=nothing, batchnorm=batchnorm)...,
			x->sigmoid.(x) # our dataset images are scaled from 0-1
	)
	# reset seed
	(init_seed !== nothing) ? Random.seed!() : nothing

	c_ = (c !== nothing) ? c : zeros(Float32, zdim)

	model = DeepSVDD(encoder, decoder, objective, c_, R, nu)
end

function svdd_constructor(
	;idim::Int=1,
	zdim::Int=1,
	hdim::Int=64,
	activation = "relu",
	nlayers::Int=3,
	init_seed=nothing,
	objective="soft-boundary",
	nu::AbstractFloat = 0.1f0,
	R::AbstractFloat = 0f0,
	c = nothing,
	kwargs...
	)
	(init_seed !== nothing) ? Random.seed!(init_seed) : nothing

	encoder = build_mlp(idim, hdim, zdim, nlayers, activation=activation, lastlayer="linear")

	decoder = build_mlp(zdim, hdim, idim, nlayers, activation=activation, lastlayer="linear")

	(init_seed !== nothing) ? Random.seed!() : nothing

	c_ = (c !== nothing) ? c : zeros(Float32, zdim)

	return DeepSVDD(encoder, decoder, objective, c_, R, nu)
end

##################################################################################

# fit Autoencoder
function fit_ae(svdd::DeepSVDD, optim, data, params)
	train_loader, val_loader = prepare_dataloaders(data, batch_size=params.batch_size, iters=params.ae_iters)
	history = MVHistory()
	total_iters = length(train_loader)
	progress = Progress(total_iters)
	# early stopping
	best_svdd = deepcopy(svdd)
	patience = params.patience
	val_batches = length(val_loader)
	best_val_loss = Inf

	if ndims(first(train_loader)) == 2
		dims = 1
	elseif ndims(first(train_loader)) == 4
		dims = (1,2,3)
	else
		error("unknown data type (no 2D or 4D tensor)")
	end

	svdd = svdd |> gpu
	ps_ae = Flux.params(svdd.encoder, svdd.decoder)
	for (iter, X) in enumerate(train_loader)
		X = getobs(X)|>gpu
		loss_ae, back = Flux.pullback(ps_ae) do
			ae_loss(AE(svdd, X), X, dims=dims)
		end
		grad_ae = back(1f0)
		Flux.Optimise.update!(optim, ps_ae, grad_ae)

		push!(history, :loss_ae, iter, loss_ae)

		next!(progress; showvalues=[
			(:iters, "$(iter)/$(total_iters)"),
			(:ae_loss, loss_ae)
			])

		if mod(iter, params.ae_check_every) == 0
			total_val_loss = 0
			Flux.testmode!(svdd)
			for X_val in val_loader
				X_val = X_val |> gpu
				total_val_loss += ae_loss(AE(svdd, X), X)
			end
			Flux.testmode!(svdd, false)
			push!(history, :validation_ae_loss, iter, total_val_loss/val_batches)
			if total_val_loss < best_val_loss
				best_val_loss = total_val_loss
				patience = params.patience
				best_svdd = deepcopy(svdd)
			else
				patience -= 1
				if patience == 0
					@info "Stopped training of Autoencoder after $(iter) iterations"
					break
				end
			end
		end
	end
	return best_svdd, history
end

function StatsBase.fit!(svdd::DeepSVDD, data, params)
	# prepare batches & loaders
	train_loader, val_loader = prepare_dataloaders(data, params)
	svdd_history = MVHistory()
	# prepare for early stopping
	best_val_loss = 1e10

	patience = params.patience
	val_batches = length(val_loader)

	#model to gpu
	svdd = svdd |> gpu

	# ADAMW(η = 0.001, β = (0.9, 0.999), decay = 0) = Optimiser(ADAM(η, β), WeightDecay(decay))
	opt_ae = haskey(params, :decay) ? ADAMW(params.lr_ae, (0.9, 0.999), params.decay) : ADAM(params.lr_ae)
	opt_svdd = haskey(params, :decay) ? ADAMW(params.lr_svdd, (0.9, 0.999), params.decay) : ADAM(params.lr_svdd)

	# loss_function for SVDD
	loss = (svdd.objective == "soft-boundary") ? sb_loss : oc_loss

	# Pretraining of AE
	svdd, ae_history = fit_ae(svdd, opt_ae, data, params)
	best_svdd = deepcopy(svdd)

	# init center of hypershpere
	svdd.c = init_c(svdd, data[1][1]) |> gpu

	# train "SVDD"
	progress = Progress(params.iters)
	ps_svdd = Flux.params(svdd.encoder)
	for (iter, X) in enumerate(train_loader)
		X = getobs(X)|>gpu
		loss_svdd, back = Flux.pullback(ps_svdd) do
			loss(svdd(X), svdd.c, svdd.R, svdd.nu)
		end
		grad_svdd = back(1f0)
		Flux.Optimise.update!(opt_svdd, ps_svdd, grad_svdd)

		# after warm up
		svdd.R = ((iter >= 0.05*params.iters) & (svdd.objective == "soft-boundary")) ? get_radius(svdd, X) : svdd.R

		push!(svdd_history, :loss_svdd, iter, loss_svdd)
		next!(progress; showvalues=[
			(:iters, "$(iter)/$(params.iters)"),
			(:svdd_loss, loss_svdd)
			])

		if mod(iter, params.check_every) == 0
			total_val_loss = 0
			Flux.testmode!(svdd)
			for X_val in val_loader
				X_val = X_val |> gpu
				total_val_loss += loss(svdd(X), svdd.c, svdd.R, svdd.nu)
			end
			Flux.testmode!(svdd, false)
			push!(svdd_history, :validation_svdd_loss, iter, total_val_loss/val_batches)
			if total_val_loss < best_val_loss
				best_val_loss = total_val_loss
				patience = params.patience
				best_svdd = deepcopy(svdd)
			else
				patience -= 1
				if patience == 0
					@info "Stopped training after $(iter) iterations"
					global iters = iter - params.check_every*params.patience
					break
				end
			end
		end
		if iter == params.iters
			global iters = params.iters
		end
	end

	return (svdd_history, ae_history), best_svdd, sum(map(p->length(p), Flux.params(svdd))), iters
end


function print_eigenvector(vec, precision)
	s = size(vec, 1)
	for i=1:s
		print(round(vec[i], precision), "\t")
	end
	println()
end

function eigenvalues_ham(hamiltonian, precision, print, Braak_form)
	size_ham = size(hamiltonian)
	λ = eigfact!(view(hamiltonian, 1:size_ham[1]-2, 1:size_ham[2]-2))
	states = Array{String}(length(λ[:values]))
	#print_matrix(hamiltonian, 2)
	if print
		println("\n Eigenvalues:")
		for i=1:size(λ[:values], 1)
			@printf("#%d |     ", i)
			@printf("%.2f \t current state |", round(λ[:values][i], precision))
			str = ""
			if Braak_form
				if i%2 == 0
					str = str * "+ "* join(convert(Int, floor((i-1)/2))) * ">"
				else
					str = str * "- " * join(convert(Int, floor(i/2))) * ">"
				end
			else
				if i%2 == 0
					str = str * "e, "* convert(Int, floor((i-1)/2))*">"
				else
					str = str * "g, " * convert(Int, floor(i/2)) *">"
				end
			end
			println(str)
			states[i] = str
		end
		#println("Extra λ |  ", round(hamiltonian[size_ham[1], size_ham[2]], precision))

		println("\n Eigenvectors:")
		for i=1:size(λ[:vectors], 2)
			@printf("#%d |     ", i)
			print_eigenvector(λ[:vectors][:, i], precision)
		end
	end
	return λ#(λ, states)
end

function sort_eigs(eigs, unsorted_eigs, ε)
	incr_ratio = eigs[:, 2] - eigs[:, 1]
	new_incr = unsorted_eigs - eigs[:, 2]
	for i=1:size(new_incr, 1)
		if abs(new_incr[i]-incr_ratio[i]) > ε ### out acceptable bounds...
			#info("qui")
			for j=i+1:size(new_incr, 1)
				if abs(unsorted_eigs[j] - eigs[i,2]-incr_ratio[i]) <= ε ##each other very close
					temp = unsorted_eigs[j]
					unsorted_eigs[j] = unsorted_eigs[i]
					unsorted_eigs[i] = temp
					new_incr = unsorted_eigs - eigs[:, 2]
				end
			end
		end
	end
	return unsorted_eigs
end

function get_corresponding_eigenvalue(hamiltonian, λ, str)
	size_ham = size(hamiltonian)
	level = eval(parse(str[search(str,',')+1:end-1]))
  # Grabbig qubit state
  index = -1
  if str[search(str,'|')+1] == 'g'
    index = level*2+1
  else
   index = level*2+2
  end
	test_hamiltonian = deepcopy(hamiltonian)
	test_hamiltonian[index, index] = 0.0
	λ_test = eigfact!(sub(test_hamiltonian, 1:size_ham[1]-1, 1:size_ham[2]-1))
	for x=1: size(λ, 1)
		found = false
		for y =1: size(λ_test[:values], 1)
			if round(λ[x], 5)==round(λ_test[:values][y, 1], 5)
				found = true
			end
		end
		found == false && return λ[x]
	end
	return NaN
end


"""
    gravy(sequence)

Calculate GRAVY index of `sequence`.
Source: https://www.ncbi.nlm.nih.gov/pubmed/7108955
"""
function gravy(sequence)
    aacomposition = BioSequences.composition(sequence)
    gravy = 0.0
    for (aa, n) in aacomposition
        gravy += n * get(hydropathy, aa, 0.0)
    end
    return gravy / length(sequence)
end


module ABBAj
include("Compress.jl")
include("Digitization.jl")
include("Reconstruction.jl")
using .Compress, .Digitization, .Reconstruction
using Statistics

export fit_transform, inverse_transform, standard_data, load_sample

"""
ABBA: adaptive Brownian bridge-based symbolic aggregation of time series
    A novel time series symbolic compressed technique with mean-based clustering.

Examples:
julia> time_series = load_samples(); # load time series samples

julia> symbols, model = fit_transform(time_series, 20, 0.1); # use 20 symbols with compressed tolerance of 0.1
julia> r_time_series = inverse_transform(symbols, model, time_series[1]); # inverse transform time series

julia> symbols, model = fit_transform(time_series, 20, 0.1, 10); # use 20 symbols with compressed tolerance of 0.1
                                                                 # and run in parallel kmeans with 10 threads
julia> r_time_series = inverse_transform(symbols, model, time_series[1]); # inverse transform time series

Reference
Elsworth, S., Güttel, S. ABBA: adaptive Brownian bridge-based symbolic aggregation of time series.
Data Min Knowl Disc 34, 1175–1200 (2020). https://doi.org/10.1007/s10618-020-00689-6

"""
global ABBA_PARAM_MU, ABBA_PARAM_SCL = Ref{Float64}(0.0), Ref{Float64}(1.0)

function fit_transform(series::AbstractVector, k::Int, tol=0.5::AbstractFloat,
                       nthreads=Threads.nthreads()::Int, max_len=typemax(Int)::Int
                       )
    @assert(k > 0 && typeof(k) == Int)
    @assert(tol > 0)
    @assert(max_len >=1)
    global ABBA_PARAM_MU, ABBA_PARAM_SCL;
    scl_series, ABBA_PARAM_MU, ABBA_PARAM_SCL= standard_data(series);
    if (ABBA_PARAM_SCL==0)
       ABBA_PARAM_SCL = 1;
    end
    # scl_series = copy(series);
    compressed = compress(scl_series, tol, max_len);
    symbols, model = digitization_k(compressed, k, nthreads);

    return symbols, model;
end

function inverse_transform(symbols::Array{Char,1},
                           model::Digitization.Model,
                           start=0.0::AbstractFloat
                           )
    pieces = inverse_digitize(symbols, model);
    pieces = quantize(pieces);
    reconstructed_ts = inverse_compress(pieces, start);
    reconstructed_ts = broadcast(*, reconstructed_ts, ABBA_PARAM_SCL);
    reconstructed_ts = broadcast(+, reconstructed_ts, ABBA_PARAM_MU);
    return reconstructed_ts;
end

function standard_data(series::AbstractVector)
    PARARM_SCL = Statistics.std(series);
    PARARM_MU =  mean(series);
    scl_series = broadcast(-, copy(series), PARARM_MU);
    scl_series = broadcast(/, scl_series, PARARM_SCL);
    return scl_series, PARARM_MU, PARARM_SCL;
end

function load_sample(len=1000::Int, freq=20.0::AbstractFloat)
    sample = zeros(len);
    j = 1;
    for i in range(0, len-1, step=1)*(1/freq)
        sample[j] = sin(i);
        j = j + 1;
    end
    return sample;
end

end


###############

function loss(model, S, c, u, λ, cTrue)
    return Flux.Losses.mse(S*vec(model(c)), u) +  λ*Flux.Losses.mse(vec(model(c)), vec(cTrue))
  end

  ###############

  function reconstructionDIP(S, N, u; model = make_model_unet(n),
                                      opt = ADAM(0.001, (0.7, 0.999)),
                                      cTrue = ones(Float32,N...),
                                      noiseData = vec(rand(Float32,N...).*Float32(0.7)),
                                      iterations = 100,
                                      λ = Float32(0),
                                      iters = nothing)

    parameters = params(model)
    data = [(model, S, noiseData, u, λ, device(cTrue))]; # λ, device(cTrue)

    @info "training"
    iter = 1:iterations
    #@showprogress 1.0 "Computing..."
    residual = Float64[]
    error = Float64[]
    ssims = Float64[]
    for i=iter
      Flux.train!(loss, parameters, data, opt)
      currentC = model(noiseData) |> cpu
      push!(residual, loss(data[1]...))
      push!(error, norm(vec(currentC)-vec(cTrue))/norm(vec(cTrue)) )
      push!(ssims, ssim(squeeze(currentC),squeeze(cTrue)) )
      if iters != nothing
        push!(iters, currentC)
      end
      if mod(i,20) == 1
        @info "res=$(residual[end]) err=$(error[end]) ssim=$(ssims[end]) time=$(i/iter[end])"
      end
    end
    cTrained = model(noiseData) |> cpu

    return cTrained, residual, error, ssims
  end

  function lossPre(model, c, u)
    return Flux.Losses.mae(vec(model(c)), vec(u))
  end

  ###############

  function reconstructionDIPPre(N; model = make_model_unet(n),
                                      opt = ADAM(0.001, (0.7, 0.999)),
                                      cTrue = ones(Float32,N...),
                                      noiseData = vec(rand(Float32,N...).*Float32(0.7)),
                                      iterations = 100)

    parameters = params(model)
    data = [(model, noiseData, cTrue)];

    @info "training"
    iter = 1:iterations
    #@showprogress 1.0 "Computing..."
    residual = Float64[]

    for i=iter
      Flux.train!(lossPre, parameters, data, opt)

      push!(residual, lossPre(data[1]...))

      if mod(i,20) == 1
        @info "res=$(residual[end])  time=$(i/iter[end])"
      end
    end
    cTrained = model(noiseData) |> cpu

    return cTrained, residual
  end


#Stationary iterative methods
#Templates, section 2.2
export jacobi, jacobi!, gauss_seidel, gauss_seidel!, sor, sor!, ssor, ssor!

jacobi(A::AbstractMatrix, b;
       tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2)^2) =
    jacobi!(zerox(A, b), A, b; tol=tol, maxiter=maxiter)

function jacobi!(x, A::AbstractMatrix, b;
        tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2)^2)
    history = ConvergenceHistory()
    history[:tol] = tol
    reserve!(history,:resnorm, maxiter)
    jacobi_method!(history, x, A, b; tol=tol, maxiter=maxiter)
    x, history
end

function jacobi_method!(log::ConvergenceHistory, x, A::AbstractMatrix, b;
        tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2)^2)
    iter=0
    n = size(A,2)
    xold = copy(x)
    z = zero(Amultype(A, x))
    tol = tol * norm(b)
	for iter=1:maxiter
        nextiter!(log,mvps=1)
		for i=1:n
			xi = z
			for j=[1:i-1;i+1:n]
				xi += A[i,j]*xold[j]
			end
			A[i,i]==0 && throw(SingularError())
			x[i]=(b[i]-xi)/A[i,i]
		end
		#check convergence
        resnorm = norm(A*x-b)
        push!(log,:resnorm,resnorm)
		resnorm < tol && (setconv(log, resnorm>=0); break)
		copy!(xold, x)
	end
    shrink!(log)
	x
end

gauss_seidel(A::AbstractMatrix, b;
        tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2)^2) =
    gauss_seidel!(zerox(A, b), A, b; tol=tol, maxiter=maxiter)

function gauss_seidel!(x, A::AbstractMatrix, b;
        tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2)^2)
    history = ConvergenceHistory()
    history[:tol] = tol
    reserve!(history,:resnorm, maxiter)
    gauss_seidel_method!(history, x, A, b; tol=tol, maxiter=maxiter)
    x, history
end

function gauss_seidel_method!(log::ConvergenceHistory, x, A::AbstractMatrix, b;
        tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2)^2)
    iter=0
    n = size(A,2)
    xold = copy(x)
    z = zero(Amultype(A, x))
    tol = tol * norm(b)
	for iter=1:maxiter
        nextiter!(log,mvps=1)
		for i=1:n
			σ=z
			for j=1:i-1
				σ+=A[i,j]*x[j]
			end
			for j=i+1:n
				σ+=A[i,j]*xold[j]
			end
			A[i,i]==0 && throw(SingularError())
			x[i]=(b[i]-σ)/A[i,i]
		end
		#check convergence
		resnorm = norm(A*x-b)
        push!(log,:resnorm,resnorm)
        resnorm < tol && (setconv(log, resnorm>=0); break)
		copy!(xold, x)
	end
    shrink!(log)
	x
end

#Successive overrelaxation
sor(A::AbstractMatrix, b, ω::Real;
    tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2)^2) =
    sor!(zerox(A, b), A, b, ω; tol=tol, maxiter=maxiter)

function sor!(x, A::AbstractMatrix, b, ω::Real;
        tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2)^2)
    history = ConvergenceHistory()
    history[:tol] = tol
    reserve!(history,:resnorm, maxiter)
    sor_method!(history, x, A, b, ω; tol=tol, maxiter=maxiter)
    x, history
end

function sor_method!(log::ConvergenceHistory, x, A::AbstractMatrix, b, ω::Real;
        tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2)^2)
	0 < ω < 2 || warn("ω = $ω lies outside the range 0<ω<2 which is required for convergence")

    iter=0
	n = size(A,2)
    xold = copy(x)
    z = zero(Amultype(A, x))
    tol = tol * norm(b)
	for iter=1:maxiter
        nextiter!(log,mvps=1)
		for i=1:n
			σ=z
			for j=1:i-1
				σ+=A[i,j]*x[j]
			end
			for j=i+1:n
				σ+=A[i,j]*xold[j]
			end
			A[i,i]==0 && throw(SingularError())
			σ=(b[i]-σ)/A[i,i]
			x[i]=xold[i]+ω*(σ-xold[i])
		end
		#check convergence
		resnorm = norm(A*x-b)
        push!(log,:resnorm,resnorm)
        resnorm < tol && (setconv(log, resnorm>=0); break)
		copy!(xold, x)
	end
    shrink!(log)
	x
end

#Symmetric successive overrelaxation
#A must be symmetric
ssor(A::AbstractMatrix, b, ω::Real;
     tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2)) =
    ssor!(zerox(A, b), A, b, ω; tol=tol, maxiter=maxiter)

function ssor!(x, A::AbstractMatrix, b, ω::Real;
        tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2))
    history = ConvergenceHistory()
    history[:tol] = tol
    reserve!(history,:resnorm, maxiter)
    ssor_method!(history, x, A, b, ω; tol=tol, maxiter=maxiter)
    x, history
end

function ssor_method!(log::ConvergenceHistory, x, A::AbstractMatrix, b, ω::Real;
        tol=size(A,2)^3*eps(typeof(real(b[1]))), maxiter=size(A,2))
	0 < ω < 2 || warn("ω = $ω lies outside the range 0<ω<2 which is required for convergence")

    iter=0
	n = size(A,2)
    xold = copy(x)
    z = zero(Amultype(A, x))
    tol = tol * norm(b)
	for iter=1:maxiter
        nextiter!(log,mvps=1)
		for i=1:n #Do a SOR sweep
			σ=z
			for j=1:i-1
				σ+=A[i,j]*x[j]
			end
			for j=i+1:n
				σ+=A[i,j]*xold[j]
			end
			A[i,i]==0 && throw(SingularError())
			σ=(b[i]-σ)/A[i,i]
			x[i]=xold[i]+ω*(σ-xold[i])
		end
		copy!(xold, x)
		for i=n:-1:1 #Do a backward SOR sweep
			σ=z
			for j=1:i-1
				σ+=A[i,j]*xold[j]
			end
			for j=i+1:n
				σ+=A[i,j]*x[j]
			end
			A[i,i]==0 && throw(SingularError())
			σ=(b[i]-σ)/A[i,i] #This line is missing in the Templates reference
			x[i]=xold[i]+ω*(σ-xold[i])
		end
		#check convergence
		resnorm = norm(A*x-b)
        push!(log,:resnorm,resnorm)
        resnorm < tol && (setconv(log, resnorm>=0); break)
		copy!(xold, x)
	end
    shrink!(log)
	x
end


# power_to_socp
#
# Functions that recursively simplify inequalities of the form
#
# x^n <= t^{n - m} s^m, t >= 0, s >= 0
#
# into a series of equations of the form x^2 <= st, which is
# second-order cone representable as
#
# norm([2x, t - s]) <= t + s, t >= 0, s >= 0.
#
# This reduction is documented in the pdf available at
# https://github.com/jump-dev/Convex.jl/raw/master/docs/supplementary/rational_to_socp.pdf

module psocp

mutable struct InequalityExpression
    # Represents an inequality of the form
    #
    # x^n <= t^{p_1} s^{p_2} u^{p_3},
    #
    # where we always have p_1 + p_2 + p_3 == n, and we must have p_3
    # belonging to {0, 1, 2}. This allows a particular form of recursive
    # enumeration of inequalities.

    # The powers p_1, p_2, p_3
    power1::Int
    power2::Int
    power3::Int

    # Index for variable represented by x in the above equation
    left_var_ind::Int
    # Index for variable represented by t
    t_var_ind::Int
    # Index for variable represented by s
    s_var_ind::Int
    # Index for variable represented by u (may be -1 if power3 == 0).
    u_var_ind::Int
end

struct SimpleInequalityExpression
    # Represents an inequality of the form
    #
    # x^2 <= t s,
    #
    # which is representable as the SOCP norm([2 x, t - s]) <= t + s. A
    # list of these is the first of the tuple returned by the method
    # ProductToSimpleInequalities.

    # Index for variable represented by x in the above equation
    left_var_ind::Int
    # Index for variable represented by t
    t_var_ind::Int
    # Index for variable represented by s
    s_var_ind::Int
end

# ProductToSimpleInequalities(first_power, second_power)
#
# Computes a sequence of inequalities of the form x^2 <= st, beginning
# from an initial inequality of the form
#
# x^n <= t^{p1} s^{p2},
#
# where n = p1 + p2 and p1 = first_power, p2 = second_power.
#
# Returns the tuple (inequalities, vars). The first of the tuple (the
# list inequalities) of which is an array of
# SimpleInequalityExpressions that consists of all the variable
# inequalities, the second of the tuple (the array vars) is an array
# of all the variable indices in the inequalities.
function ProductToSimpleInequalities(first_power::Int, second_power::Int)
    # Construct the first InequalityExpression, which is an inequality
    # of the form x^n <= t^p1 s^p2.
    @assert first_power > 0 && second_power > 0;
    n = first_power + second_power;
    var_list = [1, 2, 3];
    init_inequality = InequalityExpression(first_power, second_power, 0,
                                                                                1, 2, 3, -1);
    return ReducePowers(init_inequality, Array{SimpleInequalityExpression}(undef, 0),
                                            var_list);
end

# ReducePowers(curr_inequality::InequalityExpression,
#              ineq_array::Array{SimpleInequalityExpression, 1},
#              variable_array::Array{Int64, 1})
#
# Takes an initial inequality (curr_inequality) representing an
# expression of the form
#
# x^n <= t^{p1} s^{p2} u^{p3},
#
# where p3 is one of 0, 1, or 2, and n = p1 + p2 + p3, and recursively
# reduces it (see the file method-notes.tex) until the inequality
# array is fully populated by inequalities of the form x^2 <= t s,
# which are SOCP representable.
#
# The inequality array is populated (should be initially called with
# an empty array of SimpleInequalityExpressions), as is the
# variable_array, and returns the pair (ineq_array, variable_array)
# once both are fully populated.
function ReducePowers(curr_inequality::InequalityExpression,
                      ineq_array::Array{SimpleInequalityExpression, 1},
                      variable_array::Array{Int, 1})
    ReArrangePowers!(curr_inequality);
    p1 = curr_inequality.power1;
    p2 = curr_inequality.power2;
    p3 = curr_inequality.power3;
    n = (p1 + p2 + p3);
    @assert p1 >= 1;  # , "Must have at least 1 on first power");
    @assert p2 >= 1;  # , "Must have at least 1 on second power");
    # Evaluate cases for variables
    if (p3 == 0)
        # Double check if we have p1 == 1, p2 == 1
        if (p1 == 1 && p2 == 1)
            new_ineq = SimpleInequalityExpression(curr_inequality.left_var_ind,
                                                  curr_inequality.t_var_ind,
                                                  curr_inequality.s_var_ind);
            return (vcat(ineq_array, new_ineq), variable_array);
        end
        return ReduceThirdZero(curr_inequality,
                               ineq_array, variable_array);
    elseif (p3 == 1)
        return ReduceThirdOne(curr_inequality, ineq_array, variable_array);
    elseif (p3 == 2)
        return ReduceThirdTwo(curr_inequality, ineq_array, variable_array);
    else
        error("Third power in inequality (", curr_inequality.power3,
              ") must be 0, 1, or 2");
    end
end

# ReArrangePowers!(current_inequality::InequalityExpression)
#
# Function that takes in the current inequality, then re-orders
# variables so that the variable with smallest power is stored as
# u_var_ind. Also divides each power by their gcd.
function ReArrangePowers!(inequality::InequalityExpression)
    p1 = inequality.power1;
    p2 = inequality.power2;
    p3 = inequality.power3;
    divisor = gcd(p1, p2, p3);
    if (divisor > 1)
        inequality.power1 /= divisor;
        inequality.power2 /= divisor;
        inequality.power3 /= divisor;
        p1 = inequality.power1;
        p2 = inequality.power2;
        p3 = inequality.power3;
    end
    if (p1 <= min(p2, p3))
        # Swap t and u
        inequality.power3 = p1;
        inequality.power1 = p3;
        ind1 = inequality.t_var_ind;
        ind3 = inequality.u_var_ind;
        inequality.t_var_ind = ind3;
        inequality.u_var_ind = ind1;
    elseif (p2 <= min(p1, p3))
        # Swap s and u
        inequality.power3 = p2;
        inequality.power2 = p3;
        ind2 = inequality.s_var_ind;
        ind3 = inequality.u_var_ind;
        inequality.s_var_ind = ind3;
        inequality.u_var_ind = ind2;
    end  # In this case, p3 <= min(p1, p2), so no swap necessary
end

# ReduceThirdZero(curr_inequality::InequalityExpression,
#                 ineq_array::Array{SimpleInequalityExpression, 1},
#                 var_array::Array{Int64, 1})
#
# Recursively reduces powers when the third power (u) in
# curr_inequality is equal to 1.
function ReduceThirdZero(curr_inequality::InequalityExpression,
                         ineq_array::Array{SimpleInequalityExpression, 1},
                         var_array::Array{Int, 1})
    p1 = curr_inequality.power1;
    p2 = curr_inequality.power2;
    p3 = curr_inequality.power3;
    n = (p1 + p2);
    @assert p3 == 0;
    @assert n >= 2;
    if (mod(n, 2) == 0)
        # n is even, so check even-ness of power1, power2
        if (p1 == p2)
            # We are seriously done
            new_ineq = SimpleInequalityExpression(curr_inequality.left_var_ind,
                                                  curr_inequality.t_var_ind,
                                                  curr_inequality.s_var_ind);
            return (vcat(ineq_array, new_ineq), var_array);
        end
        if (mod(p1, 2) == 0)
            # Both are even, reduce and repeat
            curr_inequality.power1 /= 2;
            curr_inequality.power2 /= 2;
            return ReducePowers(curr_inequality, ineq_array, var_array);
        else
            # p1 is odd, p2 is odd, and they are unequal. Go to two
            # equations, which are
            #
            # x^(n / 2) <= t^((p1 - 1)/2) s^((p2 - 1)/2) u,
            # u^2 <= ts
            #
            # If n/2 == 2, then do not recurse, but simply catenate the
            # inequality on and return
            new_ind = var_array[end] + 1;
            new_ineq = SimpleInequalityExpression(new_ind,
                                                  curr_inequality.t_var_ind,
                                                  curr_inequality.s_var_ind);
            curr_inequality.power1 = (p1 - 1) / 2;
            curr_inequality.power2 = (p2 - 1) / 2;
            curr_inequality.power3 = 1;
            curr_inequality.u_var_ind = new_ind;
            # Recurse on the reduced inequality
            return ReducePowers(curr_inequality,
                                vcat(ineq_array, new_ineq),
                                vcat(var_array, new_ind));
        end  # if (mod(p1, 2) == 0)
    else    # n is odd, and we must recurse differently.
        if (mod(p1, 2) == 1)
            new_ind = var_array[end] + 1;
            new_ineq = SimpleInequalityExpression(new_ind,
                                                  curr_inequality.t_var_ind,
                                                  curr_inequality.left_var_ind);
            curr_inequality.power1 = (p1 - 1) / 2;
            curr_inequality.power2 = p2 / 2;
            curr_inequality.power3 = 1;
            curr_inequality.u_var_ind = new_ind;
            return ReducePowers(curr_inequality,
                                vcat(ineq_array, new_ineq),
                                vcat(var_array, new_ind));
        else    # mod(p2, 2) == 1
            new_ind = var_array[end] + 1;
            new_ineq = SimpleInequalityExpression(new_ind,
                                                  curr_inequality.s_var_ind,
                                                  curr_inequality.left_var_ind);
            curr_inequality.power1 = p1 / 2;
            curr_inequality.power2 = (p2 - 1) / 2;
            curr_inequality.power3 = 1;
            curr_inequality.u_var_ind = new_ind;
            return ReducePowers(curr_inequality,
                                vcat(ineq_array, new_ineq),
                                vcat(var_array, new_ind));
        end
    end  # if mod(n, 2) == 0
end

# ReduceThirdOne(curr_inequality::InequalityExpression,
#                ineq_array::Array{SimpleInequalityExpression, 1},
#                var_array::Array{Int64, 1})
#
# Recursively reduces powers when the third power (u) in
# curr_inequality is equal to 1.
function ReduceThirdOne(curr_inequality::InequalityExpression,
                        ineq_array::Array{SimpleInequalityExpression, 1},
                        var_array::Array{Int, 1})
    p1 = curr_inequality.power1;
    p2 = curr_inequality.power2;
    p3 = curr_inequality.power3;
    n = (p1 + p2 + p3);
    @assert p3 == 1; # , "Must have third power 1");
    @assert n >= 3; #  "Must have power n >= 3");
    if (mod(n, 2) == 0)
        # Exactly one of p1, p2 is odd, find it and reduce
        if (mod(p1, 2) == 1)
            new_ind = var_array[end] + 1;
            new_ineq = SimpleInequalityExpression(new_ind,
                                                  curr_inequality.t_var_ind,
                                                  curr_inequality.u_var_ind);
            curr_inequality.power1 = (p1 - 1)/2;
            curr_inequality.power2 = (p2 / 2);
            curr_inequality.power3 = 1;
            curr_inequality.u_var_ind = new_ind;
            return ReducePowers(curr_inequality,
                                vcat(ineq_array, new_ineq),
                                vcat(var_array, new_ind));
        else    # mod(p2, 2) == 1
            new_ind = var_array[end] + 1;
            new_ineq = SimpleInequalityExpression(new_ind,
                                                  curr_inequality.s_var_ind,
                                                  curr_inequality.u_var_ind);
            curr_inequality.power1 = p1 / 2;
            curr_inequality.power2 = (p2 - 1) / 2;
            curr_inequality.power3 = 1;
            curr_inequality.u_var_ind = new_ind;
            return ReducePowers(curr_inequality,
                                vcat(ineq_array, new_ineq),
                                vcat(var_array, new_ind));
        end
    else    # mod(n, 2) == 1
        # Either both of p1, p2 are odd or both are even
        if (mod(p1, 2) == 0)
            new_ind = var_array[end] + 1;
            new_ineq = SimpleInequalityExpression(new_ind,
                                                  curr_inequality.left_var_ind,
                                                  curr_inequality.u_var_ind);
            curr_inequality.power1 = p1 / 2;
            curr_inequality.power2 = p2 / 2;
            curr_inequality.power3 = 1;
            curr_inequality.u_var_ind = new_ind;
            return ReducePowers(curr_inequality,
                                vcat(ineq_array, new_ineq),
                                vcat(var_array, new_ind));
        else    # mod(p1, 2) == 1, so mod(p2, 2) == 1
            if (p1 == 1 && p2 == 1)
                # There is a simpler reduction to return: we replace
                #
                # x^3 <= tsu
                #
                # with the three inequalities
                #
                # x^4 <= w^2 v^2,  w^2 <= ts,  v^2 <= ux,
                #
                # which is equivalent to x^2 <= wv, w^2 <= ts, v^2 <= ux
                new_w_ind = var_array[end] + 1;
                new_v_ind = var_array[end] + 2;
                new_x_ineq = SimpleInequalityExpression(curr_inequality.left_var_ind,
                                                        new_w_ind, new_v_ind);
                new_w_ineq = SimpleInequalityExpression(new_w_ind,
                                                        curr_inequality.t_var_ind,
                                                        curr_inequality.s_var_ind);
                new_v_ineq = SimpleInequalityExpression(new_v_ind,
                                                        curr_inequality.u_var_ind,
                                                        curr_inequality.left_var_ind);
                return (vcat(ineq_array, [new_x_ineq, new_w_ineq, new_v_ineq]),
                        vcat(var_array, [new_w_ind, new_v_ind]));
            else
                # Do the full reduction
                new_w_ind = var_array[end] + 1;
                new_v_ind = var_array[end] + 2;
                new_y_ind = var_array[end] + 3;
                new_y_ineq = SimpleInequalityExpression(new_y_ind,
                                                        new_w_ind, new_v_ind);
                new_w_ineq = SimpleInequalityExpression(new_w_ind,
                                                        curr_inequality.t_var_ind,
                                                        curr_inequality.s_var_ind);
                new_v_ineq = SimpleInequalityExpression(new_v_ind,
                                                        curr_inequality.u_var_ind,
                                                        curr_inequality.left_var_ind);
                curr_inequality.power1 = (p1 - 1) / 2;
                curr_inequality.power2 = (p2 - 1) / 2;
                curr_inequality.power3 = 2;
                curr_inequality.u_var_ind = new_y_ind;
                return ReducePowers(curr_inequality,
                                    vcat(ineq_array, [new_y_ineq, new_w_ineq, new_v_ineq]),
                                    vcat(var_array, [new_w_ind, new_v_ind, new_y_ind]));
            end
        end
    end
end

# ReduceThirdTwo(curr_inequality::InequalityExpression,
#                ineq_array::Array{SimpleInequalityExpression, 1},
#                var_array::Array{Int64, 1})
#
# Recursively reduces powers when the third power (u) in
# curr_inequality is equal to 2.
function ReduceThirdTwo(curr_inequality::InequalityExpression,
                        ineq_array::Array{SimpleInequalityExpression, 1},
                        var_array::Array{Int, 1})
    p1 = curr_inequality.power1;
    p2 = curr_inequality.power2;
    p3 = curr_inequality.power3;
    n = (p1 + p2 + p3);
    @assert p3 == 2;  # , "Must have third power 2");
    @assert n >= 6;  # , "Must have power n >= 6");
    @assert p1 >= 2 && p2 >= 2;  # , "Did not rearrange powers properly");
    if (mod(n, 2) == 0)
        # Either both p1, p2 are even or both are odd
        if (mod(p1, 2) == 0)
            # Simply reduce powers on everything by a factor of two
            curr_inequality.power1 /= 2;
            curr_inequality.power2 /= 2;
            curr_inequality.power3 /= 2;
            return ReducePowers(curr_inequality, ineq_array, var_array);
        else
            # Both are odd, so introduce two new variables
            new_y_ind = var_array[end] + 1;
            new_w_ind = var_array[end] + 2;
            new_y_ineq = SimpleInequalityExpression(new_y_ind,
                                                    curr_inequality.u_var_ind,
                                                    new_w_ind);
            new_w_ineq = SimpleInequalityExpression(new_w_ind,
                                                    curr_inequality.s_var_ind,
                                                    curr_inequality.t_var_ind);
            curr_inequality.power1 = (p1 - 1) / 2;
            curr_inequality.power2 = (p2 - 1) / 2;
            curr_inequality.power3 = 2;
            curr_inequality.u_var_ind = new_y_ind;
            return ReducePowers(curr_inequality,
                                vcat(ineq_array, [new_y_ineq, new_w_ineq]),
                                vcat(var_array, [new_y_ind, new_w_ind]));
        end
    else    # mod(n, 2) == 1, so exactly one of p1 and p2 is odd
        if (mod(p1, 2) == 1)
            new_y_ind = var_array[end] + 1;
            new_w_ind = var_array[end] + 2;
            new_y_ineq = SimpleInequalityExpression(new_y_ind,
                                                    curr_inequality.u_var_ind,
                                                    new_w_ind);
            new_w_ineq = SimpleInequalityExpression(new_w_ind,
                                                    curr_inequality.left_var_ind,
                                                    curr_inequality.t_var_ind);
            curr_inequality.power1 = (p1 - 1) / 2;
            curr_inequality.power2 = p2 / 2;
            curr_inequality.power3 = 2;
            curr_inequality.u_var_ind = new_y_ind;
            return ReducePowers(curr_inequality,
                                vcat(ineq_array, [new_y_ineq, new_w_ineq]),
                                vcat(var_array, [new_y_ind, new_w_ind]));
        else    # mod(p2, 2) == 1
            new_y_ind = var_array[end] + 1;
            new_w_ind = var_array[end] + 2;
            new_y_ineq = SimpleInequalityExpression(new_y_ind,
                                                    curr_inequality.u_var_ind,
                                                    new_w_ind);
            new_w_ineq = SimpleInequalityExpression(new_w_ind,
                                                    curr_inequality.left_var_ind,
                                                    curr_inequality.s_var_ind);
            curr_inequality.power1 = p1 / 2;
            curr_inequality.power2 = (p2 - 1) / 2;
            curr_inequality.power3 = 2;
            curr_inequality.u_var_ind = new_y_ind;
            return ReducePowers(curr_inequality,
                                vcat(ineq_array, [new_y_ineq, new_w_ineq]),
                                vcat(var_array, [new_y_ind, new_w_ind]));
        end
    end
end

end  # module PSOCP


"""
    A = JopDifference(R::JetSpace[, dim=1])

`A*m` is similar to `diff(a,dims=dim)` where `a` is a vector in the space `R`.
In other words, it is the one-sided difference of `a` along the dimension `dim`.
"""
function JopDifference(spc::JetAbstractSpace, dim::Integer=1)
    JopLn(df! = JopDifference_df!, df′! = JopDifference_df′!, dom = spc, rng = spc, s = (dim=dim,))
end
export JopDifference

function JopDifference_df!(d, m; dim, kwargs...)
    d .= 0
    n,Rpre,Rpost = JopDifference_indices(m, dim)
    JopDifference_df_dim!(d, m, n, Rpre, Rpost)
end

function JopDifference_df′!(m, d; dim, kwargs...)
    m .= 0
    n,Rpre,Rpost = JopDifference_indices(m, dim)
    JopDifference_df′_dim!(m, d, n, Rpre, Rpost)
end

function JopDifference_indices(m, dim)
    Rpre = CartesianIndices(size(m)[1:dim-1])
    Rpost = CartesianIndices(size(m)[dim+1:end])
    n = size(m, dim)
    n,Rpre, Rpost
end

function JopDifference_df_dim!(d, m, n, Rpre, Rpost)
    for Ipost in Rpost
        for i = 1:n-1
            for Ipre in Rpre
                d[Ipre, i, Ipost] += m[Ipre, i+1, Ipost]
                d[Ipre, i, Ipost] -= m[Ipre, i, Ipost]
            end
        end
    end
    d
end

function JopDifference_df′_dim!(m, d, n, Rpre, Rpost)
    for Ipost in Rpost
        for i = 1:n-1
            for Ipre in Rpre
                m[Ipre, i+1, Ipost] += d[Ipre, i, Ipost]
                m[Ipre, i, Ipost] -= d[Ipre, i, Ipost]
            end
        end
    end
    m
end


using GLVisualize.ComposeBackend, Gadfly, Distributions
using Colors, GLVisualize

if !isdefined(:runtests)
    window = glscreen()
    composebackend = ComposeBackend.GLVisualizeBackend(window)
end

const static_example = true

sds = [1, 1/2, 1/4, 1/8, 1/16, 1/32]
n = 10
ys = [mean(rand(Distributions.Normal(0, sd), n)) for sd in sds]
ymins = ys .- (1.96 * sds / sqrt(n))
ymaxs = ys .+ (1.96 * sds / sqrt(n))

p = plot(x=1:length(sds), y=ys, ymin=ymins, ymax=ymaxs,
     Geom.point, Geom.errorbar)

draw(composebackend, p)

if !isdefined(:runtests)
	renderloop(window)
end


# Generate default parameters.
params = parameters_simple()

# Algebraic Tests
results_algebraic = stationary_algebraic_simple(params, settings_simple());
@test results_algebraic.g ≈ 0.020727753490265812; # from updating the eq in stationary
@test results_algebraic.ν ≈ 1.797254207263066;
@test results_algebraic.v(0) ≈ 34.58676260411018;
@test results_algebraic.v(2)*exp(2) ≈ 164.54095232228548; # correct for rescaling
@test results_algebraic.v(5)*exp(5) ≈ 3298.0717492785516;

# Grid Tests
grids = [
    unique([range(0.0, stop = 1.0, length = 500)' range(1.0, stop = 5.0, length = 201)']),
    unique([range(0.0, stop = 1.0, length = 1000)' range(1.0, stop = 2.0, length = 11)' range(2.0, stop = 5.0, length = 20)']),
    unique([range(0.0, stop = 1.0, length = 300)' range(1.0, stop = 2.0, length = 50)' range(2.0, stop = 7.0, length = 50)']),
    unique([range(0., stop = 0.1, length = 20)' range(0.1, stop = 1., length = 20)' range(1., stop = 5., length = 20)'])
]

gs = [
    0.019903560646302203,
    0.02008477004627781,
    0.020440266030220186,
    0.02080684556396779 # this is the most accurate, and uses the least gridpoints
]

@test var(gs) < 1e-6 # consistency

for i in 1:length(grids)
    results = stationary_numerical_simple(params, settings_simple(z_ex = grids[i]))
    @test results.g ≈ gs[i] # invariance
end


@testset "evaluating sphere" begin
    sphere1 = ImplicitSphere(v0, 1.0)
    @test sphere1.center === v0
    @test sphere1.radius === 1.0

	sp2 = ImplicitSphere([0,0,0], 15.0)
	@test sp2.center === SVector(0.0, 0, 0)
	@test sp2.radius === 15.0

	sp3 = ImplicitSphere([0,0,0], 15)
	@test sp3.center === SVector(0, 0, 0)
	@test sp3.radius === 15

	sp4 = ImplicitSphere([0,0,0.], 15)
	@test sp4.center === SVector(0., 0, 0)
	@test sp4.radius === 15.0

    @test isapprox(value(evaluate(sphere1, v0)), -1)

    @test isapprox(value(evaluate(sphere1, SVector(1, 0, 0))), 0)
    @test isapprox(value(evaluate(sphere1, SVector(1.0, 0.0, 0.0))), 0)
    @test isapprox(value(evaluate(sphere1, SVector(0.0, 1.0, 0.0))), 0)
    @test isapprox(value(evaluate(sphere1, SVector(0.0, 0.0, 1.0))), 0)

    @test value(evaluate(sphere1, SVector(0, 0, 1.0001))) > 0

    sphere2 = ImplicitSphere(SVector(-1.0, -1.0, -1.0), 1.0)
    @test isapprox(value(evaluate(sphere2, SVector(-1.0, -1.0, -1.0))), -1)

    @test isapprox(value(evaluate(sphere2, SVector(-1, -1, 0))), 0)
    @test isapprox(value(evaluate(sphere2, SVector(-1.0, -1.0, 0.0))), 0)
    @test isapprox(value(evaluate(sphere2, SVector(-1.0, 0.0, -1.0))), 0)
    @test isapprox(value(evaluate(sphere2, SVector(-1.0, -1.0, 0.0))), 0)

    @test value(evaluate(sphere2, SVector(0, 0, 0))) > 0
    @test value(evaluate(sphere2, SVector(-1.0001, -1.0, 0.0))) > 0

    sphere3 = ImplicitSphere(SVector(5, 5, 5), 5)

    @test isapprox(value(evaluate(sphere3, SVector(5, 5, 0))), 0)
    @test isapprox(value(evaluate(sphere3, SVector(5.0, 5.0, 0.0))), 0)
    @test isapprox(value(evaluate(sphere3, SVector(5.0, 0.0, 5.0))), 0)
    @test isapprox(value(evaluate(sphere3, SVector(0.0, 5.0, 5.0))), 0)

    @test value(evaluate(sphere3,SVector(0.0, 5.0001, 5.0))) > 0

    # Integer
    sphere4 = ImplicitSphere(SVector(0, 0, 0), 1)
    @test isapprox(value(evaluate(sphere1, v0)), -1)
end

@testset "evaluating plane" begin
    n1 = SVector(1,0,0.0);
    pl1 = ImplicitPlane(n1, n1)

    @test pl1.normal === n1
    @test pl1.point === n1

	pl2 = ImplicitPlane([0,0,0], [1,0,0])
	@test pl2.point === SVector(0,0,0)
	@test pl2.normal === SVector(1,0,0)

	pl3 = ImplicitPlane([0.,0,0], [1,0,0])
	@test pl3.point === SVector(0.,0,0)
	@test pl3.normal === SVector(1.,0,0)

    @test isapprox(value(evaluate(pl1, n1)), 0)
    @test isapprox(value(evaluate(pl1, SVector(0,0,0))), -1)
    @test isapprox(value(evaluate(pl1, SVector(2,0,0))), 1)

    @test value(evaluate(pl1, SVector(0.9999,0,0))) < 0
    @test value(evaluate(pl1, SVector(1.0001,0,0))) > 0

    @test isapprox(value(evaluate(pl1, SVector(1, 17.52, -37.12))), 0)
    @test isapprox(value(evaluate(pl1, SVector(1, 9446.644, -151677.09))), 0)

    @test value(evaluate(pl1, SVector(1.0001, 17.52, -37.12))) > 0
    @test value(evaluate(pl1, SVector(1.0001, 9446.644, -151677.09))) > 0

    @test value(evaluate(pl1, SVector(0.9999, 17.52, -37.12))) < 0
    @test value(evaluate(pl1, SVector(0.9999, 9446.644, -151677.09))) < 0
end

@testset "evaluating cylinder" begin
	n1 = SVector(0, 0, 1.0)
	cyl = ImplicitCylinder(n1, v0, 1.)

	@test cyl.center === v0
	@test cyl.axis === n1
	@test cyl.radius === 1.

	cyl2 = ImplicitCylinder([0,0,1], [0,0,0], 0.00001)
	@test cyl2.axis === SVector(0.0, 0, 1)
	@test cyl2.center === SVector(0.0, 0, 0)

	cyl3 = ImplicitCylinder([0,0,1.], [0,0,0], 15)
	@test cyl3.axis === SVector(0.0, 0, 1)
	@test cyl3.center === SVector(0.0, 0, 0)

	@test isapprox(value(evaluate(cyl, [1, 0, 0])), 0)
	@test isapprox(value(evaluate(cyl, [0, 1, 0])), 0)
	@test isapprox(value(evaluate(cyl, [-1, 0, 0])), 0)
	@test isapprox(value(evaluate(cyl, [0, -1, 0])), 0)

	@test isapprox(value(evaluate(cyl, [0, 0, 1])), -1)
	@test isapprox(value(evaluate(cyl, [0, 0, -1])), -1)

	@test value(evaluate(cyl, [1.0001, 0, 0])) > 0
	@test value(evaluate(cyl, [0, 1.0001, 0])) > 0
	@test value(evaluate(cyl, [-1.0001, 0, 0])) > 0
	@test value(evaluate(cyl, [0, -1.0001, 0])) > 0

	@test value(evaluate(cyl, [0.9999, 0, 0])) < 0
	@test value(evaluate(cyl, [0, 0.9999, 0])) < 0
	@test value(evaluate(cyl, [-0.9999, 0, 0])) < 0
	@test value(evaluate(cyl, [0, -0.9999, 0])) < 0
end

@testset "surface normals" begin
	@testset "plane" begin
		n1 = SVector(1,0,0.0);
		pl1 = ImplicitPlane(n1, n1)

		@test normal(pl1, v0) == n1
		@test normal(pl1, rand(SVector{3})) == n1
		@test isapprox(dot(normal(pl1, rand(SVector{3})), [1,0,0]), 1)
	end

	@testset "cylinder" begin
		n1 = SVector(0, 0, 1.0)
		cyl = ImplicitCylinder(n1, v0, 1.)

		cyl2 = ImplicitCylinder(n1, v0, 17)

		@test normal(cyl, [1,0,0]) == normal(cyl2, [1,0,0])
		@test normal(cyl, [1,0,0]) == [1, 0, 0]

		@test normal(cyl, [3.14, 0, 5]) == [1, 0, 0]
		@test normal(cyl, [-3.14, 0, -5]) == [-1, 0, 0]

		@test normal(cyl, [0, 3.14, 5]) == [0, 1, 0]
		@test normal(cyl, [0, -3.14, -5]) == [0, -1, 0]
	end

end


# This test file test detailed models in FractionalDiffEq.jl

using FractionalDiffEq
using Test

@testset "Test Bagley Torvik equation" begin
    @test isapprox(bagleytorvik(1, 1, 1, 1, 1, 0.5), [0; 0; 0.12773958089728293]; atol=1e-2)
end


using Optim
using ScenTreesMakie

#############################
# Regularization Ansatz

function myfunc(x)

    trr1.state = x # update states Vector
    #trr1.state[1] = trr2.state[1]
    dist = nestedWasserstein(trr1,trr2,2.0).^2.0

    return dist

end

###################
function myConstraint(x)

    # lower bound such that process is monotonically increasing
    lb = trr1.state[trr1.parent[2:end]]
    tmp = sum( max.([0;lb] .- x,0))

    #upper bound such that fitted tree does not exceed original tree (significantly)
    ub = maximum(trr2.state);
    tmp2 = sum(max.(x .- ub,0) )

    return( tmp + tmp2 )
end

function myRegularizedFunction(x)
    return( myfunc(x) + 1000*myConstraint(x) )

end

trr1 = Tree(404);
trr2 = Tree(404);
n = length(trr1.state);
trr1.state[2:end] = max.(trr1.state[trr1.parent[2:end]] .- trr1.state[2:end],0) + trr1.state[2:end];
x0 = trr1.state;

result = Optim.optimize(myRegularizedFunction, x0, ParticleSwarm(n_particles = 5), Optim.Options(iterations=2000, show_trace=false)  )

#result = Optim.optimize(myRegularizedFunction, x0, NelderMead(), Optim.Options(g_tol = 1e-12, show_trace=true)  )
Optim.minimizer(result)

plot(trr1)
plot(trr2)


using LinearAlgebra
using TensorOperations

#
# vhf tools
#
function get_coloumb(g, D)
    @tensor J[μ, ν] := g[μ, ν, ρ, σ] * D[ρ, σ]
    return J
end

function get_exchange(g, D)
    @tensor K[μ, ν] := g[μ, σ, ρ, ν] * D[ρ, σ]
    return K
end

function get_vhf(g, D)
    return get_coloumb(g, D) - 0.5 * get_exchange(g, D)
end

#
# RDM tools
#
function make_rdm1(mo_occ::Vector{Float64}, mo_coeff::Array{Float64})
    occ_coeff = mo_coeff[:, mo_occ.>0]
    return occ_coeff * diagm(0 => mo_occ[mo_occ.>0]) * occ_coeff'
end

#
# Fock tools
#
function get_fock(h, g, D)
    # This has to be Hermitian or eigenvalue solver with get messed up
    return Hermitian(h + get_vhf(g, D))
end

function solve_fock(h, g, D, s)
    res = eigen(get_fock(h, g, D), s)
    mo_energy = res.values
    mo_coeff = res.vectors
    return mo_energy, mo_coeff
end

function calc_energy_elec(h, g, dm)
    @tensor one_body = scalar(h[μ, ν] * dm[ν, μ])
    vhf = get_vhf(g, dm)
    @tensor two_body = 0.5 * scalar(vhf[μ, ν] * dm[ν, μ])
    return one_body + two_body, one_body, two_body
end


module IfPartialBase

include("./../../constraints/geometric/constants.jl")
using JuMP
using .Constants

export if_pb

function if_pb(m)

    #if (flbz_front_left_bot_z[i] == flbz_front_left_bot_z[j] + hob_height_of_box[j] &&
    #    flbz_front_left_bot_z[i] - hob_height_of_box[j] >= 0 &&
    #    flbx_front_left_bot_x[j] > flbx_front_left_bot_x[i] - box_i_x_axis_cover &&
    #    flbx_front_left_bot_x[j] < flbx_front_left_bot_x[i] + box_i_x_axis_cover &&
    #    flby_front_left_bot_y[j] > flby_front_left_bot_y[i] - box_i_y_axis_cover &&
    #    flby_front_left_bot_y[j] < flby_front_left_bot_y[i] + box_i_y_axis_cover)

    # IF 1
    @variable(m, if11_pbs[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @variable(m, if12_pbs[i=1:length(boxes), j=1:length(boxes); i!=j], Bin)
    @variable(m, if1_pbs[i=1:length(boxes), j=1:length(boxes); i!=j] >= 0, Bin)

    #Hz1 >= x-b+e
    #H(1-z1) >= b-x-e
    @constraint(m, c1_if11_pbs[i=1:length(boxes), j=1:length(boxes); i!=j],
        H * m[:if11_pbs][i, j] >=
        m[:flbz_front_left_bot_z][i] - (m[:flbz_front_left_bot_z][j] + m[:hob_height_of_box][j]) + 0.001
    )
    @constraint(m, c2_if11_pbs[i=1:length(boxes), j=1:length(boxes); i!=j],
        H * (1 - m[:if11_pbs][i, j]) >=
        (m[:flbz_front_left_bot_z][j] + m[:hob_height_of_box][j]) - m[:flbz_front_left_bot_z][i] - 0.001
    )

    #Hz2 >= b-x+e
    #H(1-z2) >= x-b-e
    @constraint(m, c1_if12_pbs[i=1:length(boxes), j=1:length(boxes); i!=j],
        H * m[:if12_pbs][i, j] >=
        (m[:flbz_front_left_bot_z][j] + m[:hob_height_of_box][j]) - m[:flbz_front_left_bot_z][i] + 0.001
    )
    @constraint(m, c2_if12_pbs[i=1:length(boxes), j=1:length(boxes); i!=j],
        H * (1 - m[:if12_pbs][i, j]) >=
        m[:flbz_front_left_bot_z][i] - (m[:flbz_front_left_bot_z][j] + m[:hob_height_of_box][j]) - 0.001
    )

    #If if11, if12 = 1, then box i and j have z-bot/top value within interval -e <= z <= +e, hence, expression = 1
    @constraint(m, c_if1_pbs[i=1:length(boxes), j=1:length(boxes); i!=j],
        m[:if11_pbs][i, j] + m[:if12_pbs][i, j] >=
        2 * m[:if1_pbs][i, j]
    )

    return m
end

end