#
# gqtpar.jl --
#
# Julia implementation of GQTPAR algorithm for "Computing A Trust Region Step"
# by Moré, J.J. & Sorensen, D.C. (SIAM J. Sci. Stat. Comp., 1983, vol. 4,
# pp. 553-572).
#
#-----------------------------------------------------------------------------
#
# Original FORTRAN code dgqt.f and destsv.f from MINPACK-2 project.
#
# Copyright (C) 1993, Brett M. Averick and Jorge J. Moré (MINPACK-2 Project,
# Argonne National Laboratory and University of Minnesota).
#
# Copyright (C) 1994, Brett M. Averick, Richard Carter and Jorge J. Moré
# (MINPACK-2 Project, Argonne National Laboratory and University of
# Minnesota).
#
# Julia version of GQTPAR is part of OptimPackNextGen package licensed under
# the MIT "Expat" License.
#
# Copyright (C) 2008-2018, Éric Thiébaut
# <https://github.com/emmt/OptimPackNextGen.jl>.
#

module MoreSorensen

export
    gqtpar,
    gqtpar!

import LinearAlgebra.BLAS: libblas, @blasfunc, trsv!, nrm2,
    BlasInt, BlasReal, BlasFloat, BlasComplex

# We need to apply BLAS TRSV to a submatrix.  BEWARE that no bound check is
# performed.
for (fname, elty) in ((:dtrsv_,:Float64),
                      (:strsv_,:Float32),
                      (:ztrsv_,:ComplexF64),
                      (:ctrsv_,:ComplexF32))
    @eval begin
                #       SUBROUTINE DTRSV(UPLO,TRANS,DIAG,N,A,LDA,X,INCX)
                #       .. Scalar Arguments ..
                #       INTEGER INCX,LDA,N
                #       CHARACTER DIAG,TRANS,UPLO
                #       .. Array Arguments ..
                #       DOUBLE PRECISION A(LDA,*),X(*)
        function _trsv!(uplo::Char, trans::Char, diag::Char, n::Integer,
                        A::StridedMatrix{$elty}, lda::Integer,
                        x::StridedVector{$elty}, incx::Integer = stride(x,1))
            ccall((@blasfunc($fname), libblas), Cvoid,
                (Ref{UInt8}, Ref{UInt8}, Ref{UInt8}, Ref{BlasInt},
                 Ptr{$elty}, Ref{BlasInt}, Ptr{$elty}, Ref{BlasInt}),
                 uplo, trans, diag, n, A, lda, x, incx)
        end
    end
end

function _trsv!(uplo::Char, trans::Char, diag::Char, n::Integer,
                A::StridedMatrix{T}, x::StridedVector{T},
                incx::Integer = stride(x,1)) where {T<:BlasFloat}
    _trsv!(uplo, trans, diag, n, A, max(1,stride(A,2)), x, incx)
end

"""
# Computing a Trust Region Step

**GQTPAR** is Moré & Sorensen method for computing a trust region step.

The possible calling sequences are:

```julia
gqtpar(A, uplo, b, delta, rtol, atol, itmax, par) -> info, par, x, fx, iter
```

or:

```julia
gqtpar!(A, uplo, b, delta,
        rtol, atol, itmax, par, x [, z, d, w]) -> info, par, x, fx, iter
```

Given an `n×n` symmetric matrix `A`, an `n`-vector `b`, and a positive number
`delta`, this subroutine determines a vector `x` which approximately minimizes
the quadratic function:

    f(x) = (1/2)*x'*A*x + b'*x

subject to the Euclidean norm constraint

    norm(x) ≤ delta.

This subroutine computes an approximation `x` and a Lagrange multiplier `par`
such that either `par` is zero and

    norm(x) ≤ (1+rtol)*delta,

or `par` is positive and

    abs(norm(x) - delta) ≤ rtol*delta.

If `xsol` is the solution to the problem, the approximation `x` satisfies

    f(x) ≤ ((1 - rtol)**2)*f(xsol)

The arguments are:

* `A` is a 2D floating-point array of size `(n,n)`.  On entry the full upper or
   lower triangle of `A` must contain the full upper triangle of the symmetric
   matrix `A`.  On exit the array contains the matrix `A`.

* `uplo` specifies which part of the symmetric matrix `A` is stored in array
  `A` on entry: 'U' for the full upper triangle of `A`, 'L' for the full
  lower triangle of `A`.

* `b` is a 1D floating-point array of length `n` which specifies the linear
  term in the quadratic fucntion.  The contents of `b` is left unchanged.

* `delta` is a bound on the Euclidean norm of `x`.

* `rtol` is the relative accuracy desired in the solution. Convergence occurs
  if `f(x) ≤ ((1 - rtol)^2)*f(xsol)`.

* `atol` is the absolute accuracy desired in the solution. Convergence occurs
  when `norm(x) ≤ (1 + rtol)*delta` or `max(-f(x),-f(xsol)) ≤ atol`.

* `itmax` specifies the maximum number of iterations.

* `par` is an initial estimate of the Lagrange multiplier for the constraint
  `norm(x) ≤ delta`.

For `gqtpar!`:

* `x` is a 1D floating-point array of length `n` to store the approximate
  solution.

* `z`, `d` and `w` are optional work arrays of length `n`.

The returned values are:

* `info` is an integer code:

  - `info = 1` if the function value `f(x)` has the relative accuracy specified
    by `rtol`.

  - `info = 2` if the function value `f(x)` has the absolute accuracy specified
    by `atol`.

  - `info = 3` if rounding errors prevent further progress.  On exit `x` is the
    best available approximation.

  - `info = 4` indicates a failure to converge after `itmax` iterations.  On
    exit `x` is the best available approximation.

* `par` is the final estimate of the Lagrange multiplier.

* `x` is the final estimate of the solution.

* `fx` is `f(x)` at the output `x`.

* `iter` is the number of necessary iterations.

## References

* J. J. Moré & D. C. Sorensen, "*Computing a Trust Region Step*", SIAM
  J. Sci. Stat. Comp., vol. **4**, pp. 553-572 (1983).

* MINPACK-2 Project. July 1994.
  Argonne National Laboratory and University of Minnesota.
  Brett M. Averick, Richard Carter, and Jorge J. Moré.

"""
function gqtpar(A::StridedMatrix{T},
                uplo::Char,
                b::StridedVector{T},
                delta::Real, rtol::Real, atol::Real,
                itmax::Integer, par::Real) where {T<:BlasReal}
    return gqtpar!(A, uplo, b, delta, rtol, atol, itmax, par,
                   Array{T}(undef, length(b)))
end

function gqtpar!(A::StridedMatrix{T},
                 uplo::Char,
                 b::StridedVector{T},
                 delta::Real, rtol::Real, atol::Real,
                 itmax::Integer, par::Real,
                 x::StridedVector{T},
                 z::StridedVector{T} = Array{T}(undef, length(b)),
                 d::StridedVector{T} = Array{T}(undef, length(b)),
                 w::StridedVector{T} = Array{T}(undef, length(b))) where {T<:BlasReal}
    return gqtpar!(A, uplo, b, T(delta), T(rtol), T(atol),
                   Int(itmax), T(par), x, z, d, w)
end

@doc @doc(gqtpar) gqtpar!

function gqtpar!(A::StridedMatrix{T},
                 uplo::Char,
                 b::StridedVector{T},
                 delta::T,
                 rtol::T,
                 atol::T,
                 itmax::Int,
                 par::T,
                 x::StridedVector{T},
                 z::StridedVector{T},
                 d::StridedVector{T},
                 w::StridedVector{T}) where {T<:BlasReal}

    @inbounds begin

        n = length(b)
        @assert size(A) == (n, n)
        @assert length(x) == n
        @assert length(z) == n
        @assert length(d) == n
        @assert length(w) == n
        lda = max(1, stride(A, 2))

        # Local constants and variables.
        ZERO  = convert(T, 0.0)
        ONE   = convert(T, 1.0)
        HALF  = convert(T, 0.5)
        SMALL = convert(T, 0.001)
        local alpha::T
        local anorm::T
        local bnorm::T
        local prod::T
        local rxnorm::T
        local rznorm::T
        local xnorm::T
        local parc::T
        local parf::T
        local parl::T
        local pars::T
        local paru::T
        local temp::T
        local temp1::T
        local temp2::T
        local temp3::T

        # Initialization.
        if ! isfinite(par) || par < ZERO
            par = ZERO
        end
        parf = ZERO
        xnorm = ZERO
        rxnorm = ZERO
        alpha = ZERO
        rznorm = ZERO  # avoids compiler warnings
        vzero!(x)
        vzero!(z)
        if uplo == 'U'
            # Copy upper triangular part of A in its lower triangular part and
            # save its diagonal.
            for j in 1:n
                d[j] = A[j,j]
                @simd for i in j+1:n
                    A[i,j] = A[j,i]
                end
            end
        elseif uplo == 'L'
            # Copy lower triangular part of A in its upper triangular part and
            # save its diagonal.
            for j in 1:n
                d[j] = A[j,j]
                @simd for i in j+1:n
                    A[j,i] = A[i,j]
                end
            end
        else
            error("invalid value for uplo='$uplo' (expecting 'U' or 'L')")
        end

        # Calculate the l1-norm of A, the Gershgorin row sums,
        # and the l2-norm of b.
        anorm = ZERO
        for j in 1:n
            temp = ZERO
            @simd for i in 1:n
                temp += abs(A[i,j])
            end
            anorm = max(anorm, temp)
            w[j] = temp - abs(d[j])
        end
        bnorm = nrm2(b)

        # Calculate a lower bound, PARS, for the domain of the problem.
        # Also calculate an upper bound, PARU, and a lower bound, PARL,
        # for the Lagrange multiplier.
        pars = parl = paru = -anorm
        for j in 1:n
            pars = max(pars, -d[j])
            parl = max(parl, w[j] + d[j])
            paru = max(paru, w[j] - d[j])
        end
        temp = bnorm/delta
        parl = max(ZERO, temp - parl, pars)
        paru = max(ZERO, temp + paru)

        # If the input PAR lies outside of the interval (PARL,PARU),
        # set PAR to the closer endpoint.
        par = min(max(par, parl), paru)

        # Special case: parl = paru.
        paru = max(paru, (ONE + rtol)*parl)

        # Beginning of an iteration.
        info = 0
        iter = 0
        while true

            # Safeguard PAR.
            if par ≤ pars && paru > ZERO
                par = max(sqrt(parl/paru), SMALL)*paru
            end

            # Copy the lower triangle of A into its upper triangle and
            # compute A + par*I.  FIXME: not the first time
            for j in 1:n
                A[j,j] = d[j] + par
                @simd for i in j+1:n
                    A[j,i] = A[i,j]
                end
            end

            # Attempt the Cholesky factorization of A without referencing
            # the lower triangular part.
            indef = Int(potrf!('U', A)[2])
            rednc = false

            if indef == 0

                ###############################################
                ##  Case 1: A + par⋅I is positive definite.  ##
                ###############################################

                # Compute an approximate solution x and save the
                # last value of par with A + par*I positive definite.
                parf = par
                copyto!(w, b)
                trsv!('U', 'T', 'N', A, w)
                rxnorm = nrm2(w)
                trsv!('U', 'N', 'N', A, w)
                @simd for j in 1:n
                    x[j] = -w[j]
                end
                xnorm = nrm2(x)

                # Test for convergence.
                if (abs(xnorm - delta) ≤ rtol*delta
                    || (par == ZERO && xnorm ≤ (ONE + rtol)*delta))
                    # RTOL test satisfied.
                    info = 1
                end

                # Compute a direction of negative curvature and use this
                # information to improve pars.
                rznorm = estsv!(A, z)
                pars = max(pars, par - rznorm*rznorm)

                # Compute a negative curvature solution of the form
                # x + alpha*z where norm(x + alpha*z) = delta.
                if xnorm < delta

                    # Compute ALPHA.
                    prod = DOT(z, x)/delta
                    temp = ((delta + xnorm)/delta)*(delta - xnorm)
                    alpha = temp/(abs(prod) + sqrt(prod*prod + temp/delta))
                    alpha = copysign(alpha, prod)

                    # Test to decide if the negative curvature step
                    # produces a larger reduction than with z = 0.
                    rznorm *= abs(alpha)
                    temp1 = (rznorm/delta)^2
                    rednc = (temp1 + par*(xnorm/delta)^2 ≤ par)

                    # Test for convergence.
                    temp2 = par + (rxnorm/delta)^2
                    if HALF*temp1 ≤ rtol*(ONE - HALF*rtol)*temp2
                        info = 1
                    elseif info == 0 && HALF*temp2 ≤ (atol/delta)/delta
                        info = 2
                    elseif xnorm == ZERO
                        info = 1
                    end
                end

                # Compute the Newton correction PARC to PAR.
                if xnorm == ZERO
                    parc = -par
                else
                    temp = ONE/xnorm
                    vscale!(w, temp, x)
                    trsv!('U', 'T', 'N', A, w)
                    temp = nrm2(w)
                    parc = (((xnorm - delta)/delta)/temp)/temp
                end

                # Update PARL or PARU.
                if xnorm > delta && parl < par
                    parl = par
                end
                if xnorm < delta && paru > par
                    paru = par
                end

            else

                ###################################################
                ##  Case 2: A + par⋅I is not positive definite.  ##
                ###################################################

                # Use the rank information from the Cholesky
                # decomposition to update PAR.
                if indef > 1
                    # Restore column indef to A + par*I.
                    @simd for j in 1:indef-1
                        A[j,indef] = A[indef,j]
                    end
                    A[indef,indef] = d[indef] + par

                    # Compute PARC.
                    @simd for j in 1:indef-1
                        w[j] = A[j,indef]
                    end
                    _trsv!('U', 'T', 'N', indef-1, A, w)
                    A[indef,indef] -= nrm2(indef-1, w, stride(w,1))^2
                    _trsv!('U', 'N', 'N', indef-1, A, w)
                end
                w[indef] = -ONE
                temp = nrm2(indef, w, stride(w,1))
                parc = -(A[indef,indef]/temp)/temp
                pars = max(pars, par, par + parc)

                # If necessary, increase PARU slightly.  This is needed
                # because in some exceptional situations PARU is the optimal
                # value of PAR.
                paru = max(paru, (ONE + rtol)*pars)
            end

            # Use PARS to update PARL.
            parl = max(parl, pars)

            # Test for termination.
            iter += 1
            if info == 0
                if paru == ZERO
                    info = 2
                elseif paru ≤ (ONE + HALF*rtol)*pars
                    info = 3
                elseif iter ≥ itmax
                    info = 4
                end
            end

            # If exiting, store the best approximation and restore
            # the upper triangle of A.
            if info != 0

                # Compute the best current estimates for x and f.
                par = parf
                if rednc
                    fx = -HALF*((rxnorm^2 + par*delta^2) - rznorm^2)
                    if alpha != ZERO
                        @simd for j in 1:n
                            x[j] += alpha*z[j]
                        end
                    end
                else
                    fx = -HALF*(rxnorm^2 + par*xnorm^2)
                end

                # Restore the upper triangle of A.
                for j in 1:n
                    A[j,j] = d[j]
                    @simd for i in j+1:n
                        A[j,i] = A[i,j]
                    end
                end

                return info, par, x, fx, iter
            end

            # Compute an improved estimate for PAR.
            par = max(par + parc, parl)

        end # End of an iteration.
    end
end

"""

```julia
estsv!(R, z) -> svmin
```

Given an `n×n` upper triangular matrix `R`, this subroutine estimates the
smallest singular value and the associated singular vector of `R`.

In the algorithm a vector `e` is selected so that the solution `y` to the
system `R'*y = e` is large. The choice of sign for the components of `e` cause
maximal local growth in the components of `y` as the forward substitution
proceeds.  The vector `z` is the solution of the system `R*z = y`, and the
estimate `svmin` is `norm(y)/norm(z)` in the Euclidean norm.

On exit `svmin` is returned and `z` contains a singular vector associated with
the estimate `svmin` such that `norm(R*z) = svmin` and `norm(z) = 1` in the
Euclidean norm.

## References

* MINPACK-2 Project. October 1993.
  Argonne National Laboratory.
  Brett M. Averick and Jorge J. Moré.

"""
function estsv!(R::StridedMatrix{T},
                z::StridedVector{T}) where {T<:BlasReal}

    # Local constants and variables.
    n = length(z)
    @assert size(R) = (n,n)
    ZERO  = convert(T, 0.0)
    ONE   = convert(T, 1.0)
    SMALL = convert(T, 0.01)
    local e::T, s::T, sm::T, temp::T, w::T, wm::T, ynorm::T, znorm::T

    @inbounds begin

        # Clear vector z.
        @simd for i in 1:n
            z[i] = ZERO
        end

        # This choice of e makes the algorithm scale invariant.
        e = abs(R[1,1])
        if e == ZERO
            z[1] = ONE
            return ZERO
        end

        # Solve R'*y = e.
        for i in 1:n

            # Scale y. The factor of 0.01 reduces the number of scalings.
            e = copysign(e, -z[i])
            if abs(e - z[i]) > abs(R[i,i])
                temp = min(abs(R[i,i]/(e - z[i])), SMALL)
                vscale!(z, temp)
                e *= temp
            end

            # Determine the two possible choices of y[i].
            if R[i,i] == ZERO
                w = ONE
                wm = ONE
            else
                w  =  (e - z[i])/R[i,i]
                wm = -(e + z[i])/R[i,i]
            end

            # Choose y[i] based on the predicted value of y(j) for j > i.
            s  = abs(e - z[i])
            sm = abs(e + z[i])
            for j in i+1:n
                sm += abs(z[j] + wm*R[i,j])
            end
            @simd for j in i+1:n
                temp = z[j] + w*R[i,j]
                z[j] = temp
                s += abs(temp)
            end
            if s < sm
                temp = wm - w
                w = wm
                if temp != ZERO
                    @simd for j in i+1:n
                        z[j] += temp*R[i,j]
                    end
                end
            end
            z[i] = w
        end
        ynorm = nrm2(z)

        # Solve R**z = y.
        for j in n:-1:1
            # Scale z.
            if abs(z[j]) > abs(R[j,j])
                temp = min(abs(R[j,j]/z[j]), SMALL)
                vscale!(z, temp)
                ynorm *= temp
            end
            if R[j,j] == ZERO
                z[j] = ONE
            else
                z[j] /= R[j,j]
            end
            if (temp = z[j]) != ZERO
                @simd for i in 1:j-1
                    z[i] -= temp*R[i,j]
                end
            end
        end

        # Normalize z and return svmin.
        znorm = ONE/nrm2(z)
        vscale!(z, znorm)
        return ynorm*znorm
    end
end

end # module


using SweepContractor

n=8; d=2

TN=TensorNetwork(undef,n)
for i=1:n
    TN[i]=Tensor(
        setdiff(1:n,i),
        randn(d*ones(Int,n-1)...),
        randn(),
        randn()
    )
end

println("example_rand")
for χ∈d.^(4:9)
    sweep=sweep_contract(TN,χ,2*χ;connected=false)
    println("χ=$χ:\t",ldexp(sweep...))
end


include("./exportData.jl")
using .exportData, DifferentialEquations, DiffEqParamEstim, PyPlot, Optim, DelimitedFiles

const N = Int(round(3/0.015 * sqrt(2)))
const D1= 0.0010/(0.015*0.015)
const D2= 0.0002/(0.015*0.015)

# define fenkar system
function H(x;k=100.0)
	return 0.5*(1.0 + tanh(k*x))
end
function fenkar!(dx, x, p, t)

	# parameters
	@views tsi,tv1m,tv2m,tvp,twm,twp,td,to,tr,xk,uc,uv,ucsi,θ,t0,TI,IA = p[1:17]

	@inbounds for n=1:N

		# fenkar dynamics
		dx[n,1] =-(x[n,1]*H(uc-x[n,1])/to + H(x[n,1]-uc)/tr - x[n,2]*H(x[n,1]-uc)*(1.0-x[n,1])*(x[n,1]-uc)/td - x[n,3]*H(x[n,1]-ucsi;k=xk)/tsi)
		dx[n,2] = H(uc-x[n,1])*(1.0-x[n,2])/(tv1m*H(uv-x[n,1]) + tv2m*H(x[n,1]-uv)) - H(x[n,1]-uc)*x[n,2]/tvp
		dx[n,3] = H(uc-x[n,1])*(1.0-x[n,3])/twm - H(x[n,1]-uc)*x[n,3]/twp

		# no-flux with angular dependence on two diffusions
		if n==1
			dx[n,1] = dx[n,1] + (D1*cos(θ)^2 + D2*sin(θ)^2)*(2*x[2,1] - 2*x[n,1])
		elseif n==N
			dx[n,1] = dx[n,1] + (D1*cos(θ)^2 + D2*sin(θ)^2)*(2*x[N-1,1] - 2*x[n,1])
		else
			dx[n,1] = dx[n,1] + (D1*cos(θ)^2 + D2*sin(θ)^2)*(x[n-1,1] + x[n+1,1] - 2*x[n,1])
		end
	end

	@inbounds for n=1:10
		dx[n,1] = dx[n,1] + H(t-t0;k=1.0)*IA*sin(pi*(t-t0)/TI)^500
	end

	return nothing
end

# fenkar system parameters in p[1:13] (BR; from https://doi.org/10.1063/1.166311)
# and then p[14:14] = [fiber_angle]
# and then p[15:17] = [offset_time, stimulus_period, stimulus_amplitude]
   p = [29.0, 19.6, 1250.0, 3.33, 41.0, 870.0, 0.25, 12.5, 33.3, 10.0, 0.13, 0.04, 0.85, 0.0,  1.0,   125.0, 0.1]

tspan = (0.0, 3000.0)
u0 = zeros(Float64,N,3)

#solve and take data
prob = ODEProblem(fenkar!, u0, tspan, p)
sol = solve(prob, Tsit5(); saveat=tspan[begin]:5.0:tspan[end])

x = collect(0:(N-1)).*0.015;
plt.pcolormesh(sol.t, x, sol[:,1,:], rasterized=true, snap=true, vmin=0.0, vmax=1.0)
plt.colorbar()
plt.savefig("./pde/optim/baseline.pdf")
plt.close()

# getExpData
ind=78
t, data, p[15] = getExpData(ind; tInds=1:1000)

# remake problem
prob = remake(prob, tspan=(0.0, t[end]), p=p)

# use the save_idxs to select the mid-point of the u-variable
sol = solve(prob, Tsit5(); saveat=t, save_idxs=Int(round(N//2)):Int(round(N//2)))

fig = plt.figure(figsize=(4,2))
plt.plot(t, data[:], "-", label="Exp. Sample")
plt.plot(sol.t, sol[1,:], "--", label="Initial Parameters")
plt.xlabel("\$ t \$")

# define the cost function
cost_function = build_loss_objective(prob, Tsit5(), L2Loss(t,data), maxiters=100000, verbose_opt=true, verbose_steps=5, saveat=t, save_idxs=Int(round(N//2)):Int(round(N//2)));

# define limits for parameters
#          tsi,  tv1m,   tv2m,   tvp,   twm,    twp,  td,    to,    tr,   xk,  uc,    uv, ucsi,     θ,   t0,    TI, IA
#   p = [ 33.4,  25.7, 1162.6,   4.00, 26.7,  662.8, 0.26, 18.7,  29.0, 10.0, 0.26, 0.13,  0.41,  0.0,  1.0,   BCL, 0.1]
lower = [  1.0,  10.0,  400.0,   2.33, 10.0,  100.0, 0.10, 10.0,  10.0,  9.0, 0.10, 0.01,  0.10, -0.3,  0.0, 100.0, 0.0]
upper = [ 50.0,  50.0, 2000.0,   4.33, 50.0, 1000.0, 0.50, 50.0,  50.0, 11.0, 0.90, 0.90,  0.90,  0.3, t[1], 150.0, 1.0]

@assert all(p .>= lower) && all(p .<= upper)

# fit
result = optimize(cost_function, lower, upper, p, Fminbox(NelderMead()))
@info "Initial p = $(p) \n"
@info "Minimizer = $(result.minimizer) \n"
prob = remake(prob, p=result.minimizer)
sol = solve(prob, Tsit5(); saveat=t, save_idxs=Int(round(N//2)):Int(round(N//2)))
plt.plot(sol.t, sol[1,:], "--", label="Nelder-Mead Parameters")

result = optimize(cost_function, lower, upper, result.minimizer, Fminbox(BFGS()))
@info "Initial p = $(p) \n"
@info "Minimizer = $(result.minimizer) \n"
prob = remake(prob, p=result.minimizer)
sol = solve(prob, Tsit5(); saveat=t, save_idxs=Int(round(N//2)):Int(round(N//2)))
plt.plot(sol.t, sol[1,:], "--", label="BFGS Parameters")

plt.savefig("./pde/optim/$(ind)_fittings.pdf")
plt.close()

open("./pde/optim/$(ind)_params.txt", "w") do io
	writedlm(io, ["tsi" "tv1m" "tv2m" "tvp" "twm" "twp" "td" "to" "tr" "xk" "uc" "uv" "ucsi" "θ" "t0" "TI" "IA"])
	writedlm(io, "\n")
	writedlm(io, transpose(p))
	writedlm(io, "\n")
	writedlm(io, transpose(result.minimizer))
end


using SymbolicUtils, Test
using SymbolicUtils: Term, Sym, istree, operation, arguments, symtype

SymbolicUtils.istree(ex::Expr) = ex.head == :call
SymbolicUtils.operation(ex::Expr) = ex.args[1]
SymbolicUtils.arguments(ex::Expr) = ex.args[2:end]

for f ∈ [:+, :-, :*, :/, :^]
    @eval begin
        Base.$f(x::Union{Expr, Symbol}, y::Number) = Expr(:call, $f, x, y)
        Base.$f(x::Number, y::Union{Expr, Symbol}) = Expr(:call, $f, x, y)
        Base.$f(x::Union{Expr, Symbol}, y::Union{Expr, Symbol}) = (Expr(:call, $f, x, y))
    end
end

ex = 1 + (:x - 2)

@test simplify(ex) == ex

SymbolicUtils.symtype(::Expr) = Real
SymbolicUtils.symtype(::Symbol) = Real
@test simplify(ex) == -1 + :x
@test simplify(:a * (:b + -1 * :c) + -1 * (:b * :a + -1 * :c * :a), expand=true) == 0


#=
The tests here merely check that the interface functions work.
Testing of the implementation is done in ldlt.jl.
=#

using Test
using Random
using SparseArrays
using LinearAlgebra
using SparseFactorizations

@testset "ldlt API" begin
    @testset "ldlt_structure" begin
        Random.seed!(1)
        A = sprand(10,10,0.1)
        A = I + A + A'
        Fref = lu(Matrix(A),Val(false))
        P = ldlt_structure(A)
        @test P == (Fref.L .!= 0)
    end

    @testset "level_of_fill" begin
        Random.seed!(1)
        A = sprand(10,10,0.1)
        A = I + A + A'
        Fref = lu(Matrix(A),Val(false))
        P = level_of_fill(A)
        nonzeros(P) .= 1
        @test P == (Fref.L .!= 0)
    end

    @testset "simplicial_ldlt" begin
        @testset for (conj,adj) in ((conj,adjoint),(identity,transpose))
            Random.seed!(1)
            A = sprand(ComplexF64,10,10,0.1)
            A = I + A + adj(A)
            Fref = lu(Matrix(A),Val(false))
            F = simplicial_ldlt(A; conj=conj)

            @test F.L ≈ Fref.L
            @test F.d ≈ diag(Fref.U)
            @test F.Lt == adj(F.L)
            @test F.D == Diagonal(F.d)
        end
    end

    @testset "ildlt" begin
        @testset for (conj,adj) in ((conj,adjoint),(identity,transpose))
            Random.seed!(1)
            A = sprand(ComplexF64,10,10,0.1)
            A = I + A + adj(A)
            Fref = lu(Matrix(A),Val(false))
            P = ldlt_structure(A)
            F = ildlt(A,P; conj=conj)

            @test F.L ≈ Fref.L
            @test F.d ≈ diag(Fref.U)
            @test F.Lt == adj(F.L)
            @test F.D == Diagonal(F.d)
        end
    end

    @testset "ildlt_lof" begin
        @testset for (conj,adj) in ((conj,adjoint),(identity,transpose))
            Random.seed!(1)
            A = sprand(ComplexF64,10,10,0.1)
            A = I + A + adj(A)
            Fref = lu(Matrix(A),Val(false))
            F = ildlt_lof(A,size(A,2); conj=conj)

            @test F.L ≈ Fref.L
            @test F.d ≈ diag(Fref.U)
            @test F.Lt == adj(F.L)
            @test F.D == Diagonal(F.d)
        end
    end

    @testset "ildlt_tol" begin
        @testset for (conj,adj) in ((conj,adjoint),(identity,transpose))
            Random.seed!(1)
            A = sprand(ComplexF64,10,10,0.1)
            A = I + A + adj(A)
            Fref = lu(Matrix(A),Val(false))
            F = ildlt_tol(A,0.0; conj=conj)

            @test F.L ≈ Fref.L
            @test F.d ≈ diag(Fref.U)
            @test F.Lt == adj(F.L)
            @test F.D == Diagonal(F.d)
        end
    end
end


#=
day1:
- Julia version: 1.5.2
- Author: Paul.Mealor
- Date: 2020-12-01
=#

using AoC,Test

powerConsumption(lines) = γ(lines) * ϵ(lines)

γ(lines::AbstractArray{<:AbstractString}) = γ([[c=='1' for c in line] for line in lines])
function γ(lines::AbstractArray{<:AbstractArray{Bool}})
    total = zeros(Int, length(lines[1]))
    for line in  lines
        total += line
    end
    threshold = length(lines)/2
    return [b >= threshold for b in total] |> bb -> foldl(((g,b)->g*2+b), bb; init = 0)
end

ϵ(lines::Array{<:AbstractString}) = ϵ([[c=='1' for c in line] for line in lines])
function ϵ(lines::AbstractArray{<:AbstractArray{Bool}})
    total = zeros(Int, length(lines[1]))
    for line in  lines
        total += line
    end
    threshold = length(lines)/2
    return [b < threshold for b in total] |>( bb -> foldl(((g,b)->g*2+b), bb) )
end

lifeSupportRating(lines) = oxygenGeneratorRating(lines) * co2ScrubberRating(lines)

oxygenGeneratorRating(lines::AbstractArray{<:AbstractString}) = oxygenGeneratorRating([[c=='1' for c in line] for line in lines])
function oxygenGeneratorRating(lines::AbstractArray{<:AbstractArray{Bool}})
    for bit in 1:length(lines[1])
        bitValue = γ(lines) & (1<<(length(lines[1])-bit)) != 0
        lines = filter(l->l[bit] == bitValue, lines)
        if length(lines) == 1
            break
        end
    end
    return reduce((total,next)->total*2+next, lines[1])
end

@test oxygenGeneratorRating(exampleLines(3,1)) == 23

co2ScrubberRating(lines::AbstractArray{<:AbstractString}) = co2ScrubberRating([[c=='1' for c in line] for line in lines])
function co2ScrubberRating(lines)
    for bit in 1:length(lines[1])
        bitValue = ϵ(lines) & (1<<(length(lines[1])-bit)) != 0
        lines = filter(l->l[bit] == bitValue, lines)
        if (length(lines) == 1)
            break
        end
    end
    return reduce((total,next)->total*2+next, lines[1])
end

@test co2ScrubberRating(exampleLines(3,1)) == 10

@test lifeSupportRating(exampleLines(3,1)) == 230

show(exampleLines(3,1) |> ll -> @time powerConsumption(ll))
show(lines(3) |> ll -> @time powerConsumption(ll))
show(exampleLines(3,1) |> ll -> @time lifeSupportRating(ll))
show(lines(3) |> ll -> @time lifeSupportRating(ll))


#
#   Auxiliarfunctions.jl file:
#       ErrorEst:
#       MyNorm
#       Rdigits

function ErrorEst(U,F,dt,beta,abstol,reltol)

    uiType=eltype(U[1])

    (s,)=size(F)
	D=length(U[1])

	est=zero(uiType)

    @inbounds begin
	for k in eachindex(U[1])

		sum=zero(uiType)
		maxU=zero(uiType)
		for is in 1:s
        	sum+=beta[is]*F[is][k]
			maxU=max(maxU,abs(U[is][k]))
        end

		est+=(abs(dt*sum))^2/(abstol+maxU^2*reltol)

	end
    end

    return(est/D)

end

function MyNorm(u,abstol,reltol)

    uiType=eltype(u)
	norm=zero(uiType)

    @inbounds begin
	for k in eachindex(u)
	    aux=u[k]/(abstol+abs(u[k])*reltol)
		norm+=aux*aux
	end
    end

	norm=sqrt(norm/(length(u)))

    return(norm)

end

function Rdigits(x::Real,r::Integer)

	aux=copy(x)
	mx=r*aux
	mxx=mx+aux
	return (mxx-mx)

end


# ------------------------------------------------------------------
# Licensed under the MIT License. See LICENSE in the project root.
# ------------------------------------------------------------------

# helper function to extract raw data
# from uniform scaling objects
raw(a::UniformScaling) = a.λ
raw(a) = a

"""
    NestedVariogram(cs, γs)

A nested variogram model `γ = c₁γ₁ + c₂γ₂ + ⋯ + cₙγₙ` with
coefficients `cs = (c₁, c₂, ..., cₙ)` and variogram models
`γs = (γ₁, γ₂, ..., γₙ)`.
"""
struct NestedVariogram{CS,GS} <: Variogram{Number,Metric}
  cs::CS
  γs::GS

  function NestedVariogram{CS,GS}(cs, γs) where {CS,GS}
    @assert all(issymmetric.(cs)) "coefficients must be symmetric"
    new(cs, γs)
  end
end

NestedVariogram(cs, γs) = NestedVariogram{typeof(cs),typeof(γs)}(cs, γs)

# variogram interface
(g::NestedVariogram)(h)                  = raw(sum(c*γ(h) for (c, γ) in zip(g.cs, g.γs)))
(g::NestedVariogram)(x::Point, y::Point) = raw(sum(c*γ(x, y) for (c, γ) in zip(g.cs, g.γs)))
sill(g::NestedVariogram)                 = raw(sum(c*sill(γ) for (c, γ) in zip(g.cs, g.γs)))
nugget(g::NestedVariogram)               = raw(sum(c*nugget(γ) for (c, γ) in zip(g.cs, g.γs)))
Base.range(g::NestedVariogram)           = maximum(range(γ) for γ in g.γs)
isstationary(g::NestedVariogram)         = all(isstationary(γ) for γ in g.γs)

# algebraic structure
*(c, γ::Variogram)                          = NestedVariogram((c,), (γ,))
*(c, γ::NestedVariogram)                    = NestedVariogram(map(x->c.*x, γ.cs), γ.γs)
+(γ₁::Variogram, γ₂::Variogram)             = NestedVariogram((1, 1), (γ₁, γ₂))
+(γ₁::NestedVariogram, γ₂::Variogram)       = NestedVariogram((γ₁.cs..., 1), (γ₁.γs..., γ₂))
+(γ₁::Variogram, γ₂::NestedVariogram)       = NestedVariogram((1, γ₂.cs...), (γ₁, γ₂.γs...))
+(γ₁::NestedVariogram, γ₂::NestedVariogram) = NestedVariogram((γ₁.cs..., γ₂.cs...), (γ₁.γs..., γ₂.γs...))

"""
    structures(γ)

Return the individual structures of a (possibly nested)
variogram as a tuple. The structures are the total nugget
`cₒ`, and the coefficients (or contributions) `cs` for the
remaining non-trivial structures `γs` after normalization
(i.e. sill=1, nugget=0).
"""
function structures(γ::Variogram)
  cₒ = nugget(γ)
  c  = sill(γ) - nugget(γ)
  T  = typeof(c)
  γ  = @set γ.sill = one(T)
  γ  = @set γ.nugget = zero(T)
  cₒ, (c,), (γ,)
end

function structures(γ::NestedVariogram)
  ks, gs = γ.cs, γ.γs

  # total nugget and contributions
  cₒ = raw(sum(@. ks * nugget(gs)))
  cs = @. raw(ks * (sill(gs) - nugget(gs)))

  # discard nugget effect terms
  inds = findall(g->!(g isa NuggetEffect), gs)
  cs, γs = cs[inds], gs[inds]

  # adjust sill and nugget
  γs = map(γs) do γ
    T = typeof(sill(γ))
    γ = @set γ.sill = one(T)
    γ = @set γ.nugget = zero(T)
  end

  cₒ, cs, γs
end

# -----------
# IO METHODS
# -----------

function Base.show(io::IO, g::NestedVariogram)
  N = length(g.cs)
  print(io, "NestedVariogram{$N}")
end

function Base.show(io::IO, ::MIME"text/plain", g::NestedVariogram)
  coeffs = 1 .* raw.(g.cs)
  models = [nameof(typeof(γ)) for γ in g.γs]
  params = ["range=$(range(γ)), sill=$(sill(γ)), nugget=$(nugget(γ))" for γ in g.γs]
  println(io, g)
  println(io, "  structures")
  lines  = ["    └─$γ($p)" for (γ, p) in zip(models, params)]
  println(io, join(lines, "\n"))
  println(io, "  coefficients")
  lines  = ["    └─$c" for c in coeffs]
  print(io, join(lines, "\n"))
end


# This file is a part of JuliaFEM.
# License is MIT: see https://github.com/JuliaFEM/Mortar2D.jl/blob/master/LICENSE

using Test
using Mortar2D: calculate_mortar_assembly

# calculate mortar assembly
ns = 10 # number of slaves in interface
nm = 9 # number of masters in interface

Xs = Dict(i => [x, 1/2] for (i, x) in enumerate(range(0.0, stop=1.0, length=ns)))
Es = Dict(i => [i, i+1] for i=1:(ns-1))
Xm = Dict(ns+i => [x, 1/2] for (i, x) in enumerate(range(0.0, stop=1.0, length=nm)))
Em = Dict(i => [i+1, i+2] for i=ns:ns+nm-2)
element_types = Dict(i => :Seg2 for i=1:(ns+nm))

slave_element_ids = collect(keys(Es))
master_element_ids = collect(keys(Em))
elements = merge(Es, Em)
coords = merge(Xs, Xm)

s, m, D, M = calculate_mortar_assembly(elements,
                                       element_types,
                                       coords,
                                       slave_element_ids,
                                       master_element_ids)

# test case 1, constant pressure load
um = ones(nm)
us = D[s,s] \ (M[s,m]*um)
us_expected = ones(ns)
@test isapprox(us, us_expected)

# test case 2, linear load
um = range(0.0, stop=1.0, length=nm)
us = D[s,s] \ (M[s,m]*um)
us_expected = range(0.0, stop=1.0, length=ns)
@test isapprox(us, us_expected)


include("tests_ebayes_samples.jl")
include("tests_example_priors.jl")
include("tests_sure.jl")
include("tests_ebcf.jl")


QPSolver = Union{MPB.AbstractMathProgSolver, Function}
get_solver(subsolver::MPB.AbstractMathProgSolver) = subsolver
get_solver(generator::Function)::MPB.AbstractMathProgSolver = generator()

mutable struct LQSolver{M <: MPB.AbstractLinearQuadraticModel, S <: MPB.AbstractMathProgSolver}
    lqmodel::M
    optimsolver::S

    function LQSolver(model::JuMP.Model, optimsolver::MPB.AbstractMathProgSolver)
        lqmodel = MPB.LinearQuadraticModel(optimsolver)
        solver = new{typeof(lqmodel), typeof(optimsolver)}(lqmodel, optimsolver)
        MPB.loadproblem!(solver.lqmodel, loadLP(model)...)
        MPB.setvartype!(lqmodel, model.colCat)
        return solver
    end
end

(solver::LQSolver)() = solver(rand(MPB.numvar(solver.lqmodel)))
function (solver::LQSolver)(x₀::AbstractVector)
    if applicable(MPB.setwarmstart!, solver.lqmodel, x₀)
        try
            MPB.setwarmstart!(solver.lqmodel, x₀)
        catch
            @warn "Could not set warm start for some reason"
        end
    end
    MPB.optimize!(solver.lqmodel)
    optimstatus = MPB.status(solver.lqmodel)
    !(optimstatus ∈ [:Optimal,:Infeasible,:Unbounded]) && error("LP could not be solved, returned status: ", optimstatus)
    return nothing
end

function getsolution(solver::LQSolver)
    n = MPB.numvar(solver.lqmodel)
    optimstatus = MPB.status(solver.lqmodel)
    if optimstatus == :Optimal
        return MPB.getsolution(solver.lqmodel)
    else
        @warn "LP was not solved to optimality, returned status: $optimstatus"
        return fill(NaN, n)
    end
end

function getobjval(solver::LQSolver)
    optimstatus = MPB.status(solver.lqmodel)
    if optimstatus == :Optimal
        return MPB.getobjval(solver.lqmodel)
    elseif optimstatus == :Infeasible
        return Inf
    elseif optimstatus == :Unbounded
        return -Inf
    else
        @warn "LP could not be solved, returned status: $optimstatus"
        return NaN
    end
end

function getredcosts(solver::LQSolver)
    cols = MPB.numvar(solver.lqmodel)
    optimstatus = MPB.status(solver.lqmodel)
    if optimstatus == :Optimal
        return try
            MPB.getreducedcosts(solver.lqmodel)[1:cols]
        catch
            fill(NaN, cols)
        end
    else
        @warn "LP was not solved to optimality. Return status: $optimstatus"
        return fill(NaN, cols)
    end
end

function getduals(solver::LQSolver)
    rows = MPB.numconstr(solver.lqmodel)
    optimstatus = MPB.status(solver.lqmodel)
    if optimstatus == :Optimal
        return try
            MPB.getconstrduals(solver.lqmodel)[1:rows]
        catch
            fill(NaN, rows)
        end
    elseif optimstatus ==:Infeasible
        return try
            MPB.getinfeasibilityray(solver.lqmodel)[1:rows]
        catch
            fill(NaN, rows)
        end
    else
        @warn "LP was not solved to optimality, and the model was not infeasible. Return status: $optimstatus"
        return fill(NaN, rows)
    end
end

status(solver::LQSolver) = MPB.status(solver.lqmodel)

function feasibility_problem!(solver::LQSolver)
    MPB.setobj!(solver.lqmodel, zeros(MPB.numvar(solver.lqmodel)))
    for i = 1:MPB.numconstr(solver.lqmodel)
        MPB.addvar!(solver.lqmodel, [i], [1.0], 0.0, Inf, 1.0)
        MPB.addvar!(solver.lqmodel, [i], [-1.0], 0.0, Inf, 1.0)
    end
end

function loadLP(m::JuMP.Model)
    l = m.colLower
    u = m.colUpper
    # Build objective
    # ==============================
    c = JuMP.prepAffObjective(m)
    c *= m.objSense == :Min ? 1 : -1
    # Build constraints
    # ==============================
    # Non-zero row indices
    I = Vector{Int}()
    # Non-zero column indices
    J = Vector{Int}()
    # Non-zero values
    V = Vector{Float64}()
    # Lower constraint bound
    lb = zeros(Float64,length(m.linconstr))
    # Upper constraint bound
    ub = zeros(Float64,length(m.linconstr))
    # Linear constraints
    for (i, constr) in enumerate(m.linconstr)
        coeffs = constr.terms.coeffs
        vars = constr.terms.vars
        @inbounds for (j,var) = enumerate(vars)
            if var.m == m
                push!(I, i)
                push!(J, var.col)
                push!(V, coeffs[j])
            end
        end
        lb[i] = constr.lb
        ub[i] = constr.ub
    end
    A = sparse(I, J, V, length(m.linconstr), m.numCols)
    return A,l,u,c,lb,ub,:Min
end


# Plot recipes

# Recipe for Connectivity Histogram (Cardinality Histogram)
@recipe function f(W::SpatialWeights)
    xguide --> "Number of Neighbors"
    legend --> false
    seriestype --> :histogram
    cardinalities(W)
end

# Recipe for Moran plot
@recipe function f(x::Vector{Float64}, W::SpatialWeights)

    Wx = slag(W, x)

    zstandardize  = get(plotattributes, :standardize, true)
    if zstandardize
        z = standardize(ZScoreTransform, x)
        Wz = standardize(ZScoreTransform, Wx)
    else
        z = x
        Wz = Wx
    end

    title --> "Moran Scatterplot"
    xguide --> "Attribute"
    yguide --> "Spatial Lag of " * plotattributes[:xguide]
    legend --> false
    grid --> false

    # Variable and spatial lag of variable
    @series begin
        seriestype := :scatter
        smooth := true
        z, Wz
    end

    # Vertical line at variable mean
    zmean = mean(z)
    @series begin
        seriestype := :vline
        linestyle := :dash
        seriescolor --> :red
        [zmean]
    end

    # Horizontal line at W * variable mean
    Wzmean = mean(Wz)
    @series begin
        seriestype := :hline
        linestyle := :dash
        seriescolor --> :red
        [Wzmean]
    end

end

# Recipe for Global Spatial Autocorrelation reference distribution
@recipe function plot(x::AbstractGlobalSpatialAutocorrelation)

    S = score(x)
    Sperms = scoreperms(x)

    legend --> false

    # Histogram of permutated values
    @series begin
        seriestype := :histogram
        Sperms
    end

    # Vertical line at Moran's i
    @series begin
        seriestype := :vline
        seriescolor --> :red
        [S]
    end

end


using Distributions
using Statistics
using JLD

println("imported")
# %%
mutable struct Agent
    R::Float64
    V::Float64
end

init(n, σ) = Agent.(1., rand(Normal(0., σ), n))
Statistics.mean(arr::Array{Agent}) = mean(map(arr) do a; cis(a.R); end)
function update!(agent::Agent, r::Float64, Ψ::Float64, κ::Float64, dt::Float64)
    agent.R += (agent.V + κ * r * sin(Ψ - agent.R)) * dt
end
function evolve!(agents::Array{Agent}, κ::Float64, dt::Float64)
    meanvec = mean(agents)
    r, Ψ = abs(meanvec), angle(meanvec)
    update!.(agents, r, Ψ, κ, dt)
    meanvec
end
function simulate(steps::Int, n::Int, σ::Float64, κ::Float64, dt::Float64)
    agents = init(n, σ)
    means = Array{Complex,1}(undef, Int(steps))
    for step in 1:steps
        means[step] = evolve!(agents, κ, dt)
    end
    means
end

function graph(κs::Array{Float64}, σs::Array{Float64})
    arr = zeros(Float64, length(κs), length(σs))
    for (i, κ) in enumerate(κs), (j, σ) in enumerate(σs)
        arr[i,j] = mean(
                abs,
                @view simulate(Int(1e5), 1000, σ, κ, 1e-4)[end - 1000:end]
	)
    end
    arr
end
# %%
κs, σs = collect(0.0:0.1:2.5), collect(0.0:0.1:2.5)
zs = graph(κs, σs)
# %%
save("save.jld", "zs", zs, "ks", κs, "ss", σs)


using Random
using LinearAlgebra, Krylov
using Plots
using RandomizedLasso
const RL = RandomizedLasso

## Example with random data
n, r = 1000, 500
μ = 1e-3
A = randn(n, r)
A = A*A' + μ*I
λA = sort(eigvals(A), by=x->-x)
plt_eigvals = plot(
    λA,
    dpi=300,
    lw=3,
    title="Eigenvalues of A",
    legend=false
)
savefig(plt_eigvals, joinpath(@__DIR__, "figs/eigvals.pdf"))
deff_A = RL.deff(A, μ)

rs = 10:10:1000
ks = Int.(round.(0.9 .* rs .- 1))
tail = zeros(length(ks))
for (ind, k) in enumerate(ks)
    tail[ind] = sum(λA[k+1:end])
end

bound = max.(1e-8, λA[10:10:end] + ks ./ (rs .- ks .- 1) .* tail)

errors = Float64[]
sizehint!(errors, length(rs))
for (r, k) in zip(rs, ks)
    Anys = RL.NystromApprox(A, k, r)
    E = maximum(svdvals(A - Anys.U*Anys.Λ*Anys.U'))
    push!(errors, E)
    r % 100 == 0 && @info "Finished with r = $r"
end

plt_sketch_error = plot(ks[1:70],
    [errors[1:70], bound[1:70]],
    yaxis=:log,
    dpi=300,
    lw=3,
    label=["Error" "Expectation Bound"],
    ylabel="||A - Anys||",
    xlabel="rank k",
    title="Error vs. r with k = 0.9r - 1"
)
savefig(plt_sketch_error, joinpath(@__DIR__, "figs/sketch_error.pdf"))


import CLBLAS
import OpenCL.cl

const clblas = CLBLAS

function randomMatrixSize(num::Int)
   # Format is (m,k,n) so matrix mul is (m,n)
   matSizes = Array(Tuple, num)
   for i in 1:num
      matSizes[i] = (rand(20:100) * 32, rand(20:100) * 32, rand(20:100) * 32)
   end
   return matSizes
end

function cpu_gemm(A, B, C, alpha, beta)
    startT = time()
    C = alpha*A*B + beta*C
    runtime = time() - startT
    runtime
end

function gpu_gemm(A, B, C, alpha, beta)
     const alphaF = cl.cl_float(alpha)
     const betaF  = cl.cl_float(beta)
     const u_zero = uint32(0)

     startT = time()
     future = clblas.clblasSgemm(u_zero, u_zero, alphaF, A, B, betaF, C)
     cl.api.clWaitForEvents(1, pointer(future.event))
     runtime = time() - startT
     runtime
end

function gpu_clblas_gemm(ctx, queue, AData, BData, CData, alpha, beta)
     A = cl.Buffer(Float32, ctx, (:rw, :copy), hostbuf=AData)
     B = cl.Buffer(Float32, ctx, (:rw, :copy), hostbuf=BData)
     C = cl.Buffer(Float32, ctx, (:rw, :copy), hostbuf=CData)

     m, k = size(AData)
     k, n = size(BData)
     m = convert(Csize_t, m)
     k = convert(Csize_t, k)
     n = convert(Csize_t, n)

     lda = m
     ldb = k
     ldc = m

     event = cl.UserEvent(ctx)
     ptr_event = [pointer(event)]

     startT = time()
     clblas.clblasSgemm(1, 0, 0, m, n, k,
                             alpha, pointer(A), 0, lda,
                             pointer(B), 0, ldb, beta,
                             pointer(C), 0, ldc,
                             1, [pointer(queue)],
                             0, C_NULL, pointer(ptr_event))

     cl.api.clWaitForEvents(1, pointer(ptr_event))
     runtime = time() - startT
     runtime
end

function gpu_gemm_block(ctx, queue, kernel, AData, BData, CData, alpha, beta)
     A = cl.Buffer(Float32, ctx, (:r, :copy), hostbuf=AData)
     B = cl.Buffer(Float32, ctx, (:r, :copy), hostbuf=BData)
     C = cl.Buffer(Float32, ctx, (:rw, :copy), hostbuf=CData)

     m, k = size(AData)
     k, n = size(BData)

     blocksize = 16
     localmem1 = cl.LocalMem(Float32, blocksize^2)
     localmem2 = cl.LocalMem(Float32, blocksize^2)

     evt = kernel[queue, (m, n), (blocksize, blocksize)](convert(Csize_t, m), convert(Csize_t, k), alpha, beta, C, A, B, localmem1, localmem2)
     runtime = evt[:profile_duration] / 1e9
     runtime
end

function many_gc(num)
   for id in 1:num
      gc()
   end
end

function run(device, context, queue, kernel, num)
   cpu_times     =  zeros(Float32, num)
   gpu_times     =  zeros(Float32, num)
   gpu_times_fut = zeros(Float32, num)
   gpu_times_blk = zeros(Float32, num)

   matrix_sizes = randomMatrixSize(num)
   alpha = 5.0
   beta = 2.0

   #gc_disable()
   for id in 1:num
      m, k, n = matrix_sizes[1]
      A = rand(Float32, m, k)
      B = rand(Float32, k, n)
      C = rand(Float32, m, n)

      gpu_times[id]     = gpu_clblas_gemm(context, queue, A, B, C, alpha, beta)
      gpu_times_fut[id] = gpu_gemm(A, B, C, alpha, beta)
      gpu_times_blk[id] = gpu_gemm_block(context, queue, kernel, A, B, C, alpha, beta)
      cpu_times[id]     = cpu_gemm(A, B, C, alpha, beta)
   end
   #gc_enable()

   (cpu_times, gpu_times, gpu_times_fut, gpu_times_blk)
end

function simulate(num, verbose = false)
    clblas.setup(true)
    if(verbose)
        for platform in cl.platforms()
            if platform[:name] == "Portable Computing Language"
                warn("Portable Computing Language platform not yet supported")
                continue
            end

            for device in cl.devices(platform)
                @printf("====================================================\n")
                @printf("Platform name:    %s\n",  platform[:name])
                @printf("Platform profile: %s\n",  platform[:profile])
                @printf("Platform vendor:  %s\n",  platform[:vendor])
                @printf("Platform version: %s\n",  platform[:version])
                @printf("----------------------------------------------------\n")
                @printf("Device name: %s\n", device[:name])
                @printf("Device type: %s\n", device[:device_type])
                @printf("Device mem: %i MB\n",           device[:global_mem_size] / 1024^2)
                @printf("Device max mem alloc: %i MB\n", device[:max_mem_alloc_size] / 1024^2)
                @printf("Device max clock freq: %i MHZ\n",  device[:max_clock_frequency])
                @printf("Device max compute units: %i\n",   device[:max_compute_units])
                @printf("Device max work group size: %i\n", device[:max_work_group_size])
                @printf("Device max work item size: %s\n",  device[:max_work_item_size])
           end
        end
    end

    device, context, queue = clblas.get_next_compute_context()
    kernel_source = open(readall, "GEMM-NVIDIA.cl")
    prg  = cl.Program(context, source=kernel_source) |> cl.build!
    mmul = cl.Kernel(prg, "mmul")

    run(device, context, queue, mmul, 2)

    cpu_times, gpu_times, gpu_times_fut, gpu_times_blk = run(device, context, queue, mmul, num)
    cpu_avg     = sum(cpu_times)/num
    gpu_avg     = sum(gpu_times)/num
    gpu_fut_avg = sum(gpu_times_fut)/num
    gpu_blk_avg = sum(gpu_times_blk)/num
    println("Average CPU    :", cpu_avg)
    println("Average GPU    :", gpu_avg)
    println("Average GPU FUT:", gpu_fut_avg)
    println("Average GPU BLK:", gpu_blk_avg)
    println("GPU BLAS xer CPU    :", cpu_avg / gpu_avg)
    println("GPU FUT  xer CPU    :", cpu_avg / gpu_fut_avg)
    println("GPU BLK  xer CPU    :", cpu_avg / gpu_blk_avg)

    clblas.teardown()
end

num = int(ARGS[1])
simulate(num)

#@pyimport matplotlib.pyplot as plt
#plt.plot([1:num], gpuMmulTimes2, color="green", linewidth=2.0, linestyle="--", label="GPU-Mul2")
#plt.plot([1:num], cpuTimes, color="red", linewidth=2.0, linestyle="--", label="CPU")
#plt.plot([1:num], gpuTimes, color="blue", linewidth=2.0, linestyle="--", label="GPU")
##plt.plot([1:num], gpuMmulTimes, color="yellow", linewidth=2.0, linestyle="--", label="GPU-Mul")

#plt.ylabel("time ns")
#plt.title("CPU vs GPU times")
#plt.show()


#=- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
 Unit Conversion
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -=#

export @unit_convert!

macro unit_convert!(src, from_unit::Symbol, to_unit::Symbol)
    return Expr(:call, Symbol(from_unit, "2", to_unit, "!"), esc(src))
end

#=- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
 Converters
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -=#

thz2nj!(f::Real) = 6.626070150e-13 * f
nj2thz!(e::Real) = e / 6.626070150e-13

thz2meV!(f::Real) = 4.135667696 * f
meV2thz!(e::Real) = e / 4.135667696

nm2nj!(λ::Real) = 1.98644586e-7 / λ
nj2nm!(e::Real) = 1.98644586e-7 / e

nm2eV!(λ::Real) = 1239.84198 / λ
eV2nm!(e::Real) = 1239.84198 / e

cm⁻¹2nm!(λ::Real) = 1e7 / λ
nm2cm⁻¹!(λ::Real) = 1e7 / λ

ps2μm!(t::Real) = round(149.896225 * t)
μm2ps!(d::Real) = d / 149.896225

#=- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
 Vectorization
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -=#

for fname! in (
        :thz2nj!, :nj2thz!, :thz2meV!, :meV2thz!,
        :nm2nj!, :nj2nm!, :nm2eV!, :eV2nm!,
        :cm⁻¹2nm!, :nm2cm⁻¹!, :ps2μm!, :μm2ps!
    )
    @eval begin
        function $(fname!)(x::AbstractArray{T}) where T<:Real
            @inbounds @simd for i in eachindex(x)
                x[i] = $(fname!)(x[i])
            end
        end
    end
end


using Calculus
using Symbolics
using SparseArrays
using ForwardDiff
using ReverseDiff
using LinearAlgebra
using LightGraphs
using fastmaxent
using BenchmarkTools
using Latexify

# supporting functions - define the metric in terms of the adjacency matrix
"""using the adjacency matrix"""
function degree(i::Int,A::AbstractArray{T,2}) where T<:Real
    return sum(@view A[i,:]) - A[i,i]
end

"""using the symbolic representation"""
function degree(i::Int, A::Symbolics.Arr{Num, 2})
    return sum(@view A[i,:]) - A[i,i]
end

"""using the computed parameters"""
function degree(i, x::AbstractVector)
    res = - a(i,i,x)
    for j in eachindex(x)
        res += a(i,j,x)
    end
    return res
end

"""using the adjacency matrix"""
function ANND(i::Int, A::AbstractArray{T,2}) where T<:Real
    return sum(A[i,j] * degree(j,A) for j=1:size(A,1) if j≠i) / degree(i,A)
end

"""using the symbolic representation"""
function ANND(i::Int, A::Symbolics.Arr{Num, 2})
    return sum(A[i,j] * degree(j,A) for j=1:size(A,1) if j≠i) / degree(i,A)
end

"""using the computed parameters"""
function ANND(i::Int, x::Symbolics.Arr{Num, 1})
    @warn "ANND(i::Int, x::Symbolics.Arr{Num, 1}) is used!"
    res = zero(Num)
    for j in eachindex(x)
        @info "j: $(j), $(typeof(j)) - i: $(i), $(typeof(i))"
        if j[1] ≠ i
            res += a(i,j,x) * degree(j, x)
        end
    end
    return res / degree(i, x)
end

"""using the computed parameters"""
function ANND(i::Int, x::AbstractVector{T}) where T
    @warn "ANND(i::Int, x::AbstractVector{T}) where T is used!"
    res = zero(T)
    for j in eachindex(x)
        @info "j: $(j), $(typeof(j)) - i: $(i), $(typeof(i))"
        if j ≠ i
            res += a(i,j,x) * degree(j, x)
        end
    end
    return res / degree(i, x)
end

# second method: make use of Symbolics.jl (symbolic computations => sparsity pattern)
n = 7
@variables A[1:n, 1:n]
@variables x[1:n]
A_num = Array{Float64}(Symmetric(rand(Bool, n,n)))

i = 2

X = ANND(i, A) # X is the symbolic expression of our topological property
X_num = ANND(i,A_num) # X_num is the numerical version of our topological property
X_sym_to_num = eval(Symbolics.build_function(X, A))(A_num) # X_sym_to_num is the numerical evaluation of the symbolic expression

# compute each contribution ∂X/∂a_{t,s} and store it in an array
∂X∂a = zeros(Num,n,n)
for t = 1:n
    for s = 1:n
        ∂X∂a[t,s] = Symbolics.derivative(X, A[t,s])
    end
end

# you can eavluate the numerical value of each entry using
ΔxΔA = Symbolics.build_function(∂X∂a, A) # leads to two function definitions (with or without output arg)
ΔxΔA_exact = eval(ΔxΔA[1]) # this is the function
ΔxΔA_exact(A_num) # this is the results based on the numerical values

# third method: make use of ForwardDiff.jl (automated differentiation => numerical result)

# appropriate function that takes only one input argument
ANND_w(A::AbstractArray{R,2}) where {R<:Real} = ANND(i, A)
ANND_w(x::AbstractVector{R}) where {R<:Real} = ANND(i, x)
ForwardDiff.gradient(ANND_w, A_num)

# Challenges:
# 1. translate this to a closure for a single contribution, i.e. one ∂X/∂a_{ij}(a) at a time instead of ∇f(a) at once
# 2. inlude vector of parameters that allows to avoid computing matrix a as a whole

#@btime ΔxΔA_exact($A_num)
#@btime ForwardDiff.gradient($ANND_w, $A_num)
#@btime ReverseDiff.gradient($ANND_w, $A_num)

a_slow(x,i,j) = x[i]*x[j] / (1 + x[i]*x[j])
a_fast(x,i,j) = @inbounds f(x[i]*x[j])
f(x::T) where {T} = x / (one(T) + x)

# Include the entire method for an actual network

# supporting functions for computations
f(x::T) where {T} = x / (one(T) + x) # helper function UBCM]–
a(i,j,x::AbstractVector) = @inbounds f(x[i]*x[j]) # UBCM

# generate random graph
N = n
G = barabasi_albert(N, 2, 2, seed=-5)

model = UBCM(G)
model_c = UBCMCompact(G)
solve!(model)
solve!(model_c)
@info model.x
A_star = adjacency_matrix(G)                # observed adjacency matrix
d_star = LightGraphs.degree(G)              # observed degree sequence
ANND_star = map(i -> ANND(i, A_star), 1:N)  # observed ANND

A_exp = reshape([a(i,j,model.x) for j = 1:N for i in 1:N],N,:)  # expected adjacency_matrix
d_exp = map(x->degree(x, A_exp), 1:N)                           # expected degree sequence
ANND_exp = map(i -> ANND(i, A_exp), 1:N)                        # expected ANND

# variance for specific entry e.g. ANND (node number 2)
ΔA⁺Δa_exact = ΔxΔA_exact(A_exp)                                 # ∂ANND_i∂a_ij using analytical expression
ΔA⁺Δa_num = ForwardDiff.gradient(ANND_w, A_exp)                 # ∂ANND_i∂a_ij using autodiff (entire matrix A required)

# Trying vector based method
ΔMΔx = ForwardDiff.gradient(ANND_w, model.x)                    # gradient of M wrt computed parameters x
∂xi∂aij(i::Int, j::Int, x::Vector{T}) where T<: Real = (1 + x[i]*x[j])^2 / x[j]        # supposed form of ∂xi∂aij
∂M∂aij(ΔM::Vector, x::Vector, i::Int,j::Int) = ΔM[i] * ∂xi∂aij(i,j,x) + ΔM[j] * ∂xi∂aij(j,i,x) # unique contribution
# first problem stems from symmetrical case i = j (∂xi∂aij changes with factor 2, not main cause)
#
ΔA⁺Δa_piecewise = map(t -> t[1]==t[2] ? 0. : ∂M∂aij(ΔMΔx, model.x, t[1], t[2]), Base.product(1:N, 1:N));

## Run and compare the numerical method with the expected analytical result

# 1. Write M in function of aij & xi: OK
M_A = ANND(i, A) # M_A is the symbolic expression of our topological property using the matrix A
M_x = ANND(i, x) # M_x is the symbolic expression of our topological property using the vector x

#latexify(M_A)
# this is wrong, because there are other values that exist as well, it is better to solve the problem in the
# compact space in order to obtain the desired result. This reduces the matrix from N² to Ŋ² elements that
# can be reused :-)
#Symbolics.derivative(M_x, x[1]) / Symbolics.derivative(a(1,2,x), x[1]) + Symbolics.derivative(M_x, x[2]) / Symbolics.derivative(a(1,2,x), x[2])

##### Compact Vs extended model

model = UBCM(G)
model_c = UBCMCompact(G)
solve!(model)
solve!(model_c)
N_c = length(model_c.x0)

A_exp = reshape([i ≠ j ? a(i,j,model.x) : 0.0 for j = 1:N for i in 1:N],N,:)
A_exp_c = reshape([ a(i,j,model_c.x)  for j = 1:N_c for i in 1:N_c],N_c,:)

# Lets expand the matrix
nodemap = similar(model.x, Int)
for (key, val) in model_c.revmap
    for node in val
        nodemap[node] = key
    end
end

exp_ind(i,nodemap) = nodemap[i]
exp_ind_f(i,nodemap) = @inbounds nodemap[i]

A_exp_CL = reshape([A_exp_c[exp_ind(i, nodemap), exp_ind(j, nodemap)] for j in 1:N for i in 1:N], N,:)
ΔA⁺Δa_num_c = ForwardDiff.gradient(ANND_w, A_exp_c)

σ²ₐ(i,j,x) = (x[i]*x[j]) / (1 + x[i]*x[j])^2

σ²ₓ(A,x) = sum( σ²ₐ(i,j,x) * A[i,j]^2 for i in eachindex(x) for j in eachindex(x) if i≠j)

σ²ₓ(ΔA⁺Δa_num, model.x)

σ²ₓ_mat = zeros(length(model.x),length(model.x))
for i in eachindex(model.x)
    for j in eachindex(model.x)
        if i ≠ j
            σ²ₓ_mat[i,j] = σ²ₐ(i,j,model.x) * ΔA⁺Δa_num[i,j]^2
        end
    end
end
σ²ₓ_mat_c = zeros(length(model_c.x),length(model_c.x))
for i in eachindex(model_c.x)
    for j in eachindex(model_c.x)
        if i ≠ j
            σ²ₓ_mat_c[i,j] = σ²ₐ(i,j,model_c.x) * ΔA⁺Δa_num_c[i,j]^2
        end
    end
end

σ²ₓ(ΔA⁺Δa_num_c, model_c.x)

# first method: make use of Calculus.jl (expression) combined with symbolic differentiation

# - not trivial
# - text/symbol-based, so should work rather fast => hard to generate expressions

#=

function degree(i::Int,A::Array{String,2})
    res = IOBuffer()
    for j in 1:size(A,1)
        if j≠i
            write(res, " + $(A[i,j])")
        end
    end

    return String(take!(res))
end

function degree!(i::Int,A::Array{String,2}, res::IOBuffer)
    for j in 1:size(A,1)
        if j≠i
            write(res, " + $(A[i,j])")
        end
    end

    return String(take!(res))
end

n = 20
bin = IOBuffer()
A = rand(Bool,n,n)
As = reshape(["a_$(i)_$(j)" for j in 1:n for i in 1:n],n,:)
using BenchmarkTools

@btime degree(1,$As)
@btime degree!(1,$As,$bin)

d_i = degree(1, As)

@btime differentiate($d_i, :a_1_7)
=#

#=
struct CompactArray{T,N,A<:AbstractArray{T,N},J<:Integer,I<:AbstractArray{J,N} } <: AbstractArray{T,N}
    values::A
    index::I
end
Base.IndexStyle(::Type{CompactArray{T,N,A,J,I}}) where {T,N,A,J,I} = IndexStyle(A)
Base.show(io::IO,A::CompactArray{T,N}) where {T,N} = print(io, "CompactArray{$(T)}\n$(Base.print_matrix(io,[1 2;4 5]))")
Base.display(A::CompactArray{T,N}) where {T,N} = show(A)
Base.size(A::CompactArray{T,N}) where {T,N} = length(A.index) .* size(A.values)
Base.getindex(A::CompactArray, i::Int) = getindex(A.values,i)
Base.getindex(A::CompactArray{T,N}, I::Vararg{Int,N}) where {T,N} = getindex(A.values, I)
Base.setindex(A::CompactArray{T,N}, v::T, i::Int) where {T,N} = setindex(A.values, v, i)
Base.setindex!(A::CompactArray{T,N}, v, I::Vararg{Int,N}) where {T,N}  = setindex(A.values, v, I)
val = [1 2;3 4]
inds = [1;2;3;4;5]
CompactArray(val, inds)

Base.size(A::CompactArray{T,N}) where {T,N} = size(A.index)
Base.getindex(A::CompactArray, i::Int) = getindex(A.values,i)
Base.getindex(A::CompactArray{T,N}, I::Vararg{Int,N}) where {T,N} = getindex(A.values, I)
Base.setindex(A::CompactArray{T,N}, v::T, i::Int) where {T,N} = setindex(A.values, v, i)
Base.setindex!(A::CompactArray{T,N}, v, I::Vararg{Int,N}) where {T,N}  = setindex(A.values, v, I)
#=

struct CompactArray{T,N,I, A<:AbstractArray{T,N}, V<:AbstractVector{I}} <: AbstractArray{T,N}
    values::A
    index::V
    function CompactArray{T,N,I,A,V}(values::AbstractArray{T,N}, index::AbstractVector{I}) where {T,N,I,A,V}
        I <: Integer || error("index must have integer eltype")
        new{T,N,I,A,V}(values, index)
    end
end

const CompactVector{T,N,I,A<:AbstractVector,V} = IndirectArray{T,N,I,A,V}

#Base.size(A::CompactArray) = length(A.index) .* size(A.values)
Base.size(A::CompactArray{T,N,I,A,V}) where {T,N,I,A,V} = length(A.index) .* size(A.values)
Base.size(A::IndirectArray) = size(A.index)
Base.getindex(A::CompactArray, inds...) =
x = CompactArray(model_c.x, nodemap)

=#
=#

# Large model


# First run some mcmc sampling, then test the util functions..

f_ec = (x::Vector{Float64}) -> ( sin(x[1])^2 + sin(x[2])^2 ) / ( 0.1 + x[1]*x[1] + x[2]*x[2] )^(1/2.5)
f_double_ec = (x::Vector{Float64}) -> max( f_ec( x+[pi/4;pi/4] ) , f_ec( x-[pi/4;pi/4] ) )

xlim = [-20.,20.]
ylim = [-20.,20.]
bounding_box_G =[eye(2);-eye(2)]
bounding_box_h =[-xlim[1];xlim[2];-ylim[1];ylim[2]]

# and create likelihood from "inverted egg crate"
f_likelihood = (x::Vector{Float64}) -> f_ec(x)
f_likelihood_2p = (x,y) -> f_ec([x;y])
#f_likelihood = (x::Vector{Float64}) -> f_double_ec(x)
#f_likelihood_2p = (x,y) -> f_double_ec([x;y])

# starting point
x0 = [5.;5.]
# pure mcmc sampling using mh:
conf_only_mh = SamplingHS3.DecorrelationConfig( 20000 , 1.0*eye(2) , 0.5 , 0.5 )
f_density    = f_likelihood
# this is quite slow.. most time is spent for the generation of the multivar. normal random numbers..
tic()
mcmc_result = SamplingHS3.run_mcmc( x0 , f_density(x0) , f_density , (x) -> 0 , 1.0 , bounding_box_G  , bounding_box_h , conf_only_mh )
toc()

# Test utilts:
z_sampled      = SamplingHS3.bin2d( mcmc_result.dd.xx , 40 , 40 )
z_sampled_mean = SamplingHS3.bin2d( mcmc_result.dd.xx , mcmc_result.dd.xx_fdens , 40 , 40 , mean)

# And test interp1_lin:
xx_01  = collect(linspace(0.0,10.0,2000))
yy_01  = sin.(xx_01)
x_test = [1.0;2.0;5.0;5.23]
y_test = SamplingHS3.interp1_lin(xx_01,yy_01, x_test )
@test maximum( abs.( y_test - sin.(x_test) ) ) < 0.001


"""
    This is a port of the IDL script 'TSC.pro' by Joop Schaye from Feb 1999.
    Original: https://idlastro.gsfc.nasa.gov/ftp/pro/math/tsc.pro
"""

module TriangularShapedCloudInterpolation

    export TSCInterpolation,
           get_tsc_positions

    """
        get_tsc_positions(pos::Array{<:Real}, res_elements::Array{<:Integer})

    Returns an array with interpolation positions for a given `pos` array.
    Number of interpolation positions in each dimension is given by the `res_elements` array.
    """
    function get_tsc_positions(pos::Array{<:Real}, res_elements::Array{<:Integer})

        dim   = size(pos,1)
        Npart = size(pos,2)

        pos_tsc = Array{eltype(pos[1]),2}(undef, Npart, dim)

        @inbounds for i = 1:dim
            minx = minimum(pos[i,:])
            maxx = maximum(pos[i,:]) .* (1.0+1.e-6)
            dx   = -(minx - maxx) / res_elements[i]
            for j = 1:Npart
                pos_tsc[j,i] = ( pos[i,j] .- minx ) ./ dx
            end
        end

        return pos_tsc
    end

    """
        get_tsc_positions(pos::Array{<:Real}, res_elements::Integer)

    Helper function if resolution elements in all dimensions are the same.
    """
    get_tsc_positions(pos::Array{<:Real}, res_elements::Integer) = get_tsc_positions(pos, [res_elements, res_elements, res_elements])

    @inline function get_distances_weights(pos::Array{<:Real}, n_steps::Integer;
                                           wraparound::Bool, isolated::Bool)

        # Coordinates of nearest grid point (ngp)
        if wraparound
            ng = floor.(pos  .+ 0.5)
        else
            ng = floor.(pos) .+ 0.5
        end

        # Distance from sample to ngp.
        dng = ng .- pos

        if wraparound
            k2 = ng
        else
            k2 = floor.(Int64, ng .- 0.5)
        end

        # Weight of ngp.
        w2 = 0.75 .- dng .* dng

        # Point before ngp.
        k1 = k2 .- 1  # Index.
        dx  = 1.0 .- dng  # Distance to sample.
        w1 = 0.5 .* (1.5 .- dx).^2  # TSC-weight.

        # Point after ngp.
        k3 = k2 .+ 1  # Index.
        dx  = 1.0 .+ dng # Distance to sample.
        w3 = 0.5 .* (1.5 .- dx).^2  # TSC-weight.

        # Periodic boundary conditions.
        bad = findall( k2 .== 0 )
        if length(bad) > 0
            k1[bad] .= n_steps - 1
            if isolated
                w1[bad] .= 0.0
            end
        end

        bad = findall( k2 .== (n_steps - 1) )
        if length(bad) > 0
            k3[bad] .= 0
            if isolated
                w3[bad] .= 0.0
            end
        end

        if wraparound
            bad = findall( k2 .== n_steps )
            if length(bad) > 0
                k2[bad] .= 0
                k3[bad] .= 1
            end
        end

        return k1, k2, k3, w1, w2, w3

    end

    """
        calculate_weights!(field::Array{<:Real}, tottscweight::Array{<:Real},
                                        index::Array{<:Integer}, tscweight::Array{<:Real},
                                        value::Array{<:Real}, Nsamples::Integer;
                                        average::Bool=false)

    Helper function to update the weights arrays.
    """
    @inline function calculate_weights!(field::Array{<:Real}, tottscweight::Array{<:Real},
                                        index::Array{<:Integer}, tscweight::Array{<:Real},
                                        value::Array{<:Real}, Nsamples::Integer;
                                        average::Bool=false)

        for j = 1:Nsamples
            field[index[j]]  += tscweight[j] * value[j]
        end

        if average
            for j = 1:Nsamples
                tottscweight[index[j]] += tscweight[j]
            end
        end

        return field, tottscweight
    end

    """
        find_dim_bounds(dim::Integer)

    Helper function to get loop limits for dimension loops.
    """
    @inline function find_dim_bounds(dim::Integer)

        if dim == 1
            dimx = 3
            dimy = 1
            dimz = 1
        elseif dim == 2
            dimx = 3
            dimy = 3
            dimz = 1
        elseif dim == 3
            dimx = 3
            dimy = 3
            dimz = 3
        end

        return dimx, dimy, dimz
    end

    """
        TSCInterpolation( value::Array{<:Real},
                          pos::Array{<:Real},
                          res_elements::Array{<:Integer},
                          average::Bool=true,
                          wraparound::Bool=false,
                          isolated::Bool=false    )

    Runs a TSC interpolation on the `value` array based on the provided positions.
    Returns a 3D array with interpolated values.
    """
    function TSCInterpolation(  value::Array{<:Real},
                                pos::Array{<:Real},
                                res_elements::Array{<:Integer};
                                average::Bool=true,
                                wraparound::Bool=false,
                                isolated::Bool=false    )

        Nsamples = length(value)

        # allocate arrays for indices and weights
        kx = zeros(Int64, Nsamples,3)
        ky = zeros(Int64, Nsamples,3)
        kz = zeros(Int64, Nsamples,3)

        wx = ones(Nsamples,3)
        wy = ones(Nsamples,3)
        wz = ones(Nsamples,3)

        dim = size(pos,2)

        nx = res_elements[1]

        # x direction
        kx[:,1], kx[:,2], kx[:,3], wx[:,1], wx[:,2], wx[:,3] = get_distances_weights(pos[:,1], res_elements[1],
                                                                                wraparound = wraparound,
                                                                                isolated   = isolated)

        # y direction
        if dim > 1
            ny = res_elements[2]

            ky[:,1], ky[:,2], ky[:,3], wy[:,1], wy[:,2], wy[:,3] = get_distances_weights(pos[:,2], res_elements[2],
                                                                    wraparound = wraparound,
                                                                    isolated   = isolated)
        else
            ny = 1
        end

        # z direction
        if dim > 2
            nz = res_elements[3]

            kz[:,1], kz[:,2], kz[:,3], wz[:,1], wz[:,2], wz[:,3] = get_distances_weights(pos[:,3], res_elements[3],
                                                                    wraparound = wraparound,
                                                                    isolated   = isolated)
        else
            nz = 1
        end

        nxny = nx*ny

        # find maximum indices for dimension loop
        dimx, dimy, dimz = find_dim_bounds(dim)

        # allocate arrays for interpolated data and weighting
        field = zeros( nx * ny * nz )
        tottscweight = zeros( nx * ny * nz )

        # run interpolation
        @inbounds @fastmath for i = 1:dimx, j = 1:dimy, k = 1:dimz

            index     = @. kx[:,i] + ky[:,j] * nx + kz[:,k] * nxny + 1
            tscweight = @. wx[:,i] * wy[:,j] * wz[:,k]

            calculate_weights!( field, tottscweight,
                                index, tscweight, value, Nsamples,
                                average=average )

        end

        if average
            good = findall(tottscweight .> 0.0)
            field[good] ./= tottscweight[good]
        end

        return reshape(field, (nx, ny, nz))
    end

    """
        TSCInterpolation( value::Array{<:Real},
                          pos::Array{<:Real},
                          res_elements::Integer;
                          average::Bool=true,
                          wraparound::Bool=false,
                          isolated::Bool=false  )

    Helper function to run TSC interpolation with the same number of resolution elements in all dimensions.
    """
    function TSCInterpolation(  value::Array{<:Real},
                                pos::Array{<:Real},
                                res_elements::Integer;
                                average::Bool=true,
                                wraparound::Bool=false,
                                isolated::Bool=false    )

        dim = length(pos[1,:])

        res = res_elements .* ones(Int64, dim)

        return TSCInterpolation(value, pos, res,
                                average=average,
                                wraparound=wraparound,
                                isolated=isolated )
    end

end # module


using Test, DynamicNLPModels, MadNLP, Random, JuMP, LinearAlgebra, SparseArrays, CUDA
include("sparse_lq_test.jl")

N  = 3 # number of time steps
ns = 2 # number of states
nu = 1 # number of inputs

# generate random Q, R, A, and B matrices
Random.seed!(10)
Q_rand = Random.rand(ns, ns)
Q = Q_rand * Q_rand' + I
R_rand   = Random.rand(nu,nu)
R    = R_rand * R_rand' + I

A_rand = rand(ns, ns)
A = A_rand * A_rand' + I
B = rand(ns, nu)

# generate upper and lower bounds
sl = rand(ns)
ul = fill(-15.0, nu)
su = sl .+ 4
uu = ul .+ 10
s0 = sl .+ 2

su_with_inf = copy(su)
sl_with_inf = copy(sl)

su_with_inf[1] = Inf
sl_with_inf[1] = -Inf

Qf_rand = Random.rand(ns,ns)
Qf = Qf_rand * Qf_rand' + I

E  = rand(3, ns)
F  = rand(3, nu)
gl = fill(-5.0, 3)
gu = fill(15.0, 3)

S = rand(ns, nu)

K = rand(nu, ns)

# Test with no bounds
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0)
dnlp      = LQDynamicData(s0, A, B, Q, R, N)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-7

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-6
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-6

# Test with lower bounds
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, sl = sl, ul = ul)
dnlp      = LQDynamicData(s0, A, B, Q, R, N;  sl = sl, ul=ul)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-7

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-6
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-6

# Test with upper bounds
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, su = su, uu = uu)
dnlp      = LQDynamicData(s0, A, B, Q, R, N; su = su, uu = uu)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; su = su, uu = uu)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; su = su, uu = uu)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-7

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-6
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-6

# Test with upper and lower bounds
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, sl = sl, ul = ul, su = su, uu = uu)
dnlp      = LQDynamicData(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; sl=sl, ul=ul, su = su, uu = uu)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; sl=sl, ul=ul, su = su, uu = uu)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-6
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-6

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-6
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-6

# Test with Qf matrix
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, sl = sl, ul = ul, su = su, uu = uu, Qf=Qf)
dnlp      = LQDynamicData(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, Qf = Qf)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, Qf = Qf)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, Qf = Qf)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-6
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-6

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-6
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-6

# Test get_u and get_s functions with no K matrix
s_values = value.(all_variables(model)[1:(ns * (N + 1))])
u_values = value.(all_variables(model)[(1 + ns * (N + 1)):(ns * (N + 1) + nu * N)])

@test s_values ≈ get_s(solution_ref_sparse, lq_sparse) atol = 1e-7
@test u_values ≈ get_u(solution_ref_sparse, lq_sparse) atol = 1e-7
@test s_values ≈ get_s(solution_ref_dense, lq_dense) atol = 1e-6
@test u_values ≈ get_u(solution_ref_dense, lq_dense) atol = 1e-7

# Test with E and F matrix bounds
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu)
dnlp      = LQDynamicData(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-5
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-5

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-6
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-6

# Test edge case where one state is unbounded, other(s) is bounded
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, sl = sl_with_inf, ul = ul, su = su_with_inf, uu = uu, E = E, F = F, gl = gl, gu = gu)
dnlp      = LQDynamicData(s0, A, B, Q, R, N; sl = sl_with_inf, ul = ul, su = su_with_inf, uu = uu, E = E, F = F, gl = gl, gu = gu)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; sl = sl_with_inf, ul = ul, su = su_with_inf, uu = uu, E = E, F = F, gl = gl, gu = gu)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; sl = sl_with_inf, ul = ul, su = su_with_inf, uu = uu, E = E, F = F, gl = gl, gu = gu)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-5
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-5

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-6
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-6

@test size(lq_dense.data.A, 1) == size(E, 1) * 3 + sum(su_with_inf .!= Inf .|| sl_with_inf .!= -Inf) * N

# Test S matrix case
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, S = S)
dnlp      = LQDynamicData(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, S = S)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, S = S)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, S = S)

optimize!(model)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-5
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-5

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-6
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-6

# Test K matrix case without S
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, K = K)
dnlp      = LQDynamicData(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, K = K)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, K = K)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, K = K)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-5
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-5

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-6
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-6

# Test K matrix case with S
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, K = K, S = S)
dnlp      = LQDynamicData(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, K = K, S = S)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, K = K, S = S)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, K = K, S = S)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-5
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-5

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-6
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-6

# Test K matrix case with S and with partial bounds on u

nu = 2 # number of inputs

# generate random Q, R, A, and B matrices
Random.seed!(3)
R_rand   = Random.rand(nu,nu)
R    = R_rand * transpose(R_rand) + I

B = rand(ns, nu)

# generate upper and lower bounds
ul = fill(-20.0, nu)
uu = ul .+ 30

ul_with_inf = copy(ul)
uu_with_inf = copy(uu)

uu_with_inf[1] = Inf
ul_with_inf[1] = -Inf

F  = rand(3, nu)

S = rand(ns, nu)

K = rand(nu, ns)

model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, sl = sl, ul = ul_with_inf, su = su, uu = uu_with_inf, E = E, F = F, gl = gl, gu = gu, K = K, S = S)
dnlp      = LQDynamicData(s0, A, B, Q, R, N; sl = sl, ul = ul_with_inf, su = su, uu = uu_with_inf, E = E, F = F, gl = gl, gu = gu, K = K, S = S)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul_with_inf, su = su, uu = uu_with_inf, E = E, F = F, gl = gl, gu = gu, K = K, S = S)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; sl = sl, ul = ul_with_inf, su = su, uu = uu_with_inf, E = E, F = F, gl = gl, gu = gu, K = K, S = S)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-5
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-5

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-5
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-5

# Test K with no bounds
model     = build_QP_JuMP_model(Q,R,A,B, N;s0=s0, E = E, F = F, gl = gl, gu = gu, K = K)
dnlp      = LQDynamicData(s0, A, B, Q, R, N; E = E, F = F, gl = gl, gu = gu, K = K)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

lq_sparse_from_data = SparseLQDynamicModel(s0, A, B, Q, R, N; E = E, F = F, gl = gl, gu = gu, K = K)
lq_dense_from_data  = DenseLQDynamicModel(s0, A, B, Q, R, N; E = E, F = F, gl = gl, gu = gu, K = K)

optimize!(model)
solution_ref_sparse           = madnlp(lq_sparse, max_iter=100)
solution_ref_dense            = madnlp(lq_dense, max_iter=100)
solution_ref_sparse_from_data = madnlp(lq_sparse_from_data, max_iter=100)
solution_ref_dense_from_data  = madnlp(lq_dense_from_data, max_iter=100)

@test objective_value(model) ≈ solution_ref_sparse.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense.objective atol = 1e-5
@test objective_value(model) ≈ solution_ref_sparse_from_data.objective atol = 1e-7
@test objective_value(model) ≈ solution_ref_dense_from_data.objective atol = 1e-5

@test solution_ref_sparse.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense.solution atol =  1e-5
@test solution_ref_sparse_from_data.solution[(ns * (N + 1) + 1):(ns * (N + 1) + nu*N)] ≈ solution_ref_dense_from_data.solution atol =  1e-5

# Test get_u and get_s functions with K matrix
s_values = value.(all_variables(model)[1:(ns * (N + 1))])
u_values = value.(all_variables(model)[(1 + ns * (N + 1)):(ns * (N + 1) + nu * N)])

@test s_values ≈ get_s(solution_ref_sparse, lq_sparse) atol = 1e-7
@test u_values ≈ get_u(solution_ref_sparse, lq_sparse) atol = 1e-7
@test s_values ≈ get_s(solution_ref_dense, lq_dense) atol = 1e-7
@test u_values ≈ get_u(solution_ref_dense, lq_dense) atol = 1e-7

# Test get_* and set_* functions

dnlp      = LQDynamicData(s0, A, B, Q, R, N; sl = sl, ul = ul, su = su, uu = uu, E = E, F = F, gl = gl, gu = gu, K = K, S = S)
lq_sparse = SparseLQDynamicModel(dnlp)
lq_dense  = DenseLQDynamicModel(dnlp)

@test get_A(dnlp) == A
@test get_A(lq_sparse) == A
@test get_A(lq_dense) == A

rand_val = rand()
Qtest = copy(Q)

Qtest[1, 2] = rand_val
set_Q!(dnlp, 1,2, rand_val)
@test get_Q(dnlp) == Qtest

Qtest[1, 1] = rand_val
set_Q!(lq_sparse, 1, 1, rand_val)
@test get_Q(lq_sparse) == Qtest

Qtest[2, 1] = rand_val
set_Q!(lq_dense, 2, 1, rand_val)
@test get_Q(lq_dense) == Qtest

rand_val = rand()
gltest = copy(gl)

gltest[1] = rand_val
set_gl!(dnlp, 1, rand_val)
@test get_gl(dnlp) == gltest

gltest[2] = rand_val
set_gl!(lq_sparse, 2, rand_val)
@test get_gl(lq_sparse) == gltest

gltest[3] = rand_val
set_gl!(lq_dense, 3, rand_val)
@test get_gl(lq_dense) == gltest

# Test non-default vector/matrix on GenericArrays
s0 = randn(Float32,2)
A = randn(Float32,2,2)
B = randn(Float32,2,2)
Q = randn(Float32,2,2)
R = randn(Float32,2,2)
S = randn(Float32,2,2)
K = randn(Float32,2,2)
E = randn(Float32,2,2)
F = randn(Float32,2,2)
gl = randn(Float32,2)
gu = gl .+ 2
sl = s0 .- 1
su = s0 .+ 1
ul = randn(Float32,2)
uu = ul .+ 2

s0 = Test.GenericArray(s0)
A = Test.GenericArray(A)
B = Test.GenericArray(B)
Q = Test.GenericArray(Q)
R = Test.GenericArray(R)
S = Test.GenericArray(S)
K = Test.GenericArray(K)
E = Test.GenericArray(E)
F = Test.GenericArray(F)
gl = Test.GenericArray(gl)
gu = Test.GenericArray(gu)
sl = Test.GenericArray(sl)
su = Test.GenericArray(su)
ul = Test.GenericArray(ul)
uu = Test.GenericArray(uu)

@test (DenseLQDynamicModel(s0, A, B, Q, R, 10; S = S, E = E, F = F, gl = gl, gu = gu, ul = ul, uu = uu, sl = sl, su = su) isa
    DenseLQDynamicModel{Float32, GenericArray{Float32, 1}, GenericArray{Float32, 2}, GenericArray{Float32, 2}, GenericArray{Float32, 2}, SparseMatrixCSC{Float32, Int64}, Nothing})
@test (DenseLQDynamicModel(s0, A, B, Q, R, 10; K = K, S = S, E = E, F = F, gl = gl, gu = gu, ul = ul, uu = uu, sl = sl, su = su) isa
    DenseLQDynamicModel{Float32, GenericArray{Float32, 1}, GenericArray{Float32, 2}, GenericArray{Float32, 2}, GenericArray{Float32, 2}, SparseMatrixCSC{Float32, Int64}, GenericArray{Float32, 2}})
@test (SparseLQDynamicModel(s0, A, B, Q, R, 10; S = S, E = E, F = F, gl = gl, gu = gu, ul = ul, uu = uu, sl = sl, su = su) isa
    SparseLQDynamicModel{Float32, GenericArray{Float32, 1}, SparseMatrixCSC{Float32, Int64}, SparseMatrixCSC{Float32, Int64}, GenericArray{Float32, 2}, Nothing})
@test (SparseLQDynamicModel(s0, A, B, Q, R, 10; K = K, S = S, E = E, F = F, gl = gl, gu = gu, ul = ul, uu = uu, sl = sl, su = su) isa
    SparseLQDynamicModel{Float32, GenericArray{Float32, 1}, SparseMatrixCSC{Float32, Int64}, SparseMatrixCSC{Float32, Int64}, GenericArray{Float32, 2}, GenericArray{Float32, 2}})

# Test with CUDA Array type
if CUDA.has_cuda_gpu()
    A = randn(2,2)
    B = randn(2,2)
    Q = randn(2,2)
    R = randn(2,2)
    K = randn(2,2)
    E = randn(2,2)
    F = randn(2,2)
    S = randn(2,2)
    gu = randn(2)
    gl = gu .- 2
    uu = randn(2)
    ul = uu .- 2
    su = randn(2)
    sl = su .- 2
    s0 = su .- 1

    s0c = CuArray{Float64}(undef, size(s0))
    Ac  = CuArray{Float64}(undef, size(A))
    Bc  = CuArray{Float64}(undef, size(B))
    Qc  = CuArray{Float64}(undef, size(Q))
    Rc  = CuArray{Float64}(undef, size(R))
    Kc  = CuArray{Float64}(undef, size(K))
    Sc   = CuArray{Float64}(undef, size(S))
    Ec  = CuArray{Float64}(undef, size(E))
    Fc  = CuArray{Float64}(undef, size(F))
    glc = CuArray{Float64}(undef, size(gl))
    guc = CuArray{Float64}(undef, size(gu))
    ulc = CuArray{Float64}(undef, size(ul))
    uuc = CuArray{Float64}(undef, size(uu))
    slc = CuArray{Float64}(undef, size(sl))
    suc = CuArray{Float64}(undef, size(su))

    LinearAlgebra.copyto!(Ac, A)
    LinearAlgebra.copyto!(Bc, B)
    LinearAlgebra.copyto!(Qc, Q)
    LinearAlgebra.copyto!(Rc, R)
    LinearAlgebra.copyto!(s0c, s0)
    LinearAlgebra.copyto!(Kc, K)
    LinearAlgebra.copyto!(Sc, S)
    LinearAlgebra.copyto!(Ec, E)
    LinearAlgebra.copyto!(Fc, F)
    LinearAlgebra.copyto!(glc, gl)
    LinearAlgebra.copyto!(guc, gu)
    LinearAlgebra.copyto!(ulc, ul)
    LinearAlgebra.copyto!(uuc, uu)
    LinearAlgebra.copyto!(slc, sl)
    LinearAlgebra.copyto!(suc, su)

    @test (SparseLQDynamicModel(s0c,Ac,Bc,Qc,Rc,10) isa SparseLQDynamicModel{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},
        SparseMatrixCSC{Float64, Int64}, SparseMatrixCSC{Float64, Int64}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Nothing})

    @test (SparseLQDynamicModel(s0c,Ac,Bc,Qc,Rc,10; K = Kc, S = Sc, E = Ec, F = Fc, gl = glc, gu = guc, ul = ulc, uu = uuc, sl = slc, su = suc) isa
        SparseLQDynamicModel{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, SparseMatrixCSC{Float64, Int64}, SparseMatrixCSC{Float64, Int64},
        CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}})

    @test (DenseLQDynamicModel(s0c,Ac,Bc,Qc,Rc,10) isa DenseLQDynamicModel{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer},
        CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, SparseMatrixCSC{Float64, Int64}, Nothing})

    @test (DenseLQDynamicModel(s0c,Ac,Bc,Qc,Rc,10; K = Kc, S = Sc, E = Ec, F = Fc, gl = glc, gu = guc, ul = ulc, uu = uuc, sl = slc, su = suc) isa
        DenseLQDynamicModel{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer},
        CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, SparseMatrixCSC{Float64, Int64}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}})
end


@userplot DimensionalityEstimate
@recipe function f(D::DimensionalityEstimate)
    if D.args[1] isa Inference
        I = D.args[1]
        M = I.model
        F = Array(I.F̂)
    else
        F = Array(D.args[1])
        if length(D.args) > 2
            model = D.args[2]
        else
            model = principalcomponents
        end
        M = model(F)
    end
    σ² = residualVariance(M, F)
    ξ² = explainedVariance(M)

    legend --> false
    #link := :x
    bottom_margin --> 20Plots.mm
    left_margin --> 20Plots.mm
    grid --> false
    size --> (1000, 400)
    framestyle := :box
    layout --> @layout [ev rv]

    @series begin
        seriestype := :path
        markersize --> 5
        subplot := 1
        marker --> :circle
        label --> nothing
        seriescolor --> cornflowerblue
        xguide --> "Principal Components"
        yguide --> "Residual Variance"
        (x, y) = (0:length(σ²), [1, σ²...])
    end

    @series begin
        seriestype := :path
        markersize --> 5
        subplot := 2
        marker --> :circle
        label --> nothing
        seriescolor --> crimson
        xguide --> "Principal Components"
        yguide --> "Explained Variance"
        (x, y) = (0:length(ξ²), [0, ξ²...])
    end
end

@userplot OrthonormalisedBaseline
@recipe function f(D::OrthonormalisedBaseline)
    F = D.args[1]
    Fₕ = D.args[2]
    if D.args[3] isa AbstractArray
        Fₗ = D.args[3]
        interval = (x, y) -> NonstationaryProcesses.rampInterval(0, 1, x, y)
        𝑏 = orthonormalHiloBaseline(Fₗ, Fₕ; interval)
    else
        𝑏 = D.args[3]
    end

    legend --> false
    framestyle := :box

    ξ²ₕ = sort(explainedVariance(𝑏(Fₕ)), rev=true)
    ξ² = sort(explainedVariance(𝑏(F)), rev=true)

    @series begin
        seriestype := :path
        linestyle := :dash
        linewidth := 2.5
        label --> nothing
        seriescolor --> :gray
        (x, y) = (0:length(ξ²), LinRange(0, 1, length(ξ²)+1))
    end

    @series begin
        seriestype := :path
        markersize --> 5
        marker --> :circle
        label --> "Fₕ"
        seriescolor --> :black
        (x, y) = (0:length(ξ²ₕ), cumsum([0.0, ξ²ₕ...]))
    end

    @series begin
        seriestype := :path
        markersize --> 5
        marker --> :circle
        label --> "F"
        seriescolor --> :crimson
        xguide --> "Principal Components"
        yguide --> "ξ²"
        (x, y) = (0:length(ξ²), cumsum([0.0, ξ²...]))
    end
end


#=
This file uses the Kaplan Yorke dimension to provide high accuracy estimates for
the fractal dimension for the Lorenz-96 system.
=#
using DrWatson
@quickactivate :FractalDimension # uses DynamicalSystems, PyPlot
using DynamicalSystems, Statistics, OrdinaryDiffEq

Dmax = 13
ics = 10 # how many random initial conditions to obtain
diffeq = (alg = Vern9(), reltol = 1e-9, abstol = 1e-9, maxiters = typemax(Int))
Ds = 5:Dmax
K = [Float64[] for _ ∈ 1:length(Ds)]
Λ = [Vector{Float64}[] for _ ∈ 1:length(Ds)]
T = 100000.0
F = 8.0

for (i, D) ∈ enumerate(Ds)
    @show D
    for j ∈ 1:ics
        ds = Systems.lorenz96(D; F = F)
        @time λs = lyapunovs(ds, T; Ttr = 100.0, diffeq...)
        push!(Λ[i], λs)
        push!(K[i], kaplanyorke_dim(λs))
    end
end

name = savename("lorenz96", @dict(Dmax, ics, T, F), "jld2")
result = @strdict Ds ics K Λ
@tagsave(datadir(name), result)