import LinearAlgebra
using LinearAlgebra: norm, dot, I, UniformScaling, Symmetric, eigvals, eigen
using Printf

export dgmres, dlanczos, darnoldi

"""
`dirder(x,w,f,f0)`

Finite difference directional derivative, approximate f'(x) w

### Inputs
* x, w = point and direction
* f = function
* f0 = f(x), in nonlinear iterations f(x) has usually been computed
                before the call to dirder

### Output
* `z`: approximation to `f'(x) w`

based on a code by C. T. Kelley, April 1, 2003
"""
function dirder(x, w, f, f0, h=1e-7)
   # initialise difference increment
   epsnew = h
   n = length(x)
   # scale the step
   if norm(w) == 0
      return zeros(n)
   end
   # Now scale the difference increment.
   xs = dot(x, w) / norm(w)
   if xs != 0.0
      epsnew *= max(abs(xs), 1.0) * sign(xs)
   end
   epsnew /= norm(w)
   f1 = f(x + epsnew * w)
   return (f1 - f0) / epsnew
end

"""
`dirderinf(f, f0, xc, v, h)`:
directional derivative operator to mimic the ∂f(xc) * v action
via (f(x+h'*v) - f(x)) / h' where h' is rescaled against ||v||_∞
"""
function dirderinf(f, f0, xc, v, h=1e-7)
   h = h / norm(v, Inf)
   return (f(xc + h * v) - f0) / h
end

"""
`function givapp(c, s, vin, k) -> vrot`

Apply a sequence of k Givens rotations, used within gmres codes.
"""
function givapp(c, s, vin, k)
   vrot = copy(vin)                   # TODO: can we do this in-place?
   for i = 1:k
      w1 = c[i] * vrot[i] - s[i] * vrot[i+1]
      # Here's a modest change that makes the code work in complex
      # arithmetic. Thanks to Howard Elman for this.
      # w2 = s(i)*vrot(i)+c(i)*vrot(i+1);
      w2 = s[i] * vrot[i] + c[i]' * vrot[i+1]
      vrot[i:i+1] = [w1, w2]
   end
   return vrot
end

"""
`function dgmres(f0, f, xc, errtol, kmax, [reorth, xinit]) -> x, error, total_iters`

GMRES linear equation solver for use in Newton-GMRES solver

### Input Parameters

* f0 : function at current point, f(xc)
* f : nonlinear function the format for f is  function fx = f(x)
              Note that for Newton-GMRES we incorporate any
              preconditioning into the function routine
* xc : current point
* errtol : relative residual reduction factor
* kmax : max number of iterations
* reorth : reorthogonalization method
      - 1 -- Brown/Hindmarsh condition (default)
      - 2 -- Never reorthogonalize (not recommended)
      - 3 -- Always reorthogonalize (not cheap!)
* xinit : initial iterate. xinit = zeros(f0) is the default. This
              is a reasonable choice unless restarted GMRES
              will be used as the linear solver.

### Output Parameters

* x : solution
* error : vector of residual norms for the history of the iteration
* total_iters : number of iterations
"""
function dgmres(f0, f, xc, errtol, kmax, reorth = 1, x = zeros(length(f0)); hfd = 1e-7)
   # The right side of the linear equation for the step is -f0.
   b = - f0
   n = length(b)
   r = - dirder(xc, x, f, f0, hfd) - f0

   h = zeros(kmax+1, kmax+1)
   v = zeros(n, kmax+1)
   c = zeros(kmax+1)
   s = zeros(kmax+1)
   rho = norm(r)
   g = rho * Matrix(1.0*I, kmax+1, 1)
   errtol = errtol * norm(b)
   error = Float64[]

   # Test for termination on entry.
   push!(error, rho)
   total_iters = 0
   if (rho < errtol)
      # early termination
      return x, error, total_iters
   end

   v[:,1] = r / rho
   beta = rho
   k = 0

   # GMRES iteration
   while rho > errtol && k < kmax
      k = k+1;

      # Call directional derivative function.
      v[:, k+1] = dirder(xc, v[:,k], f, f0, hfd)
      normav = norm(v[:,k+1])

      # Modified Gram-Schmidt
      for j = 1:k
         h[j, k] = dot(v[:,j], v[:,k+1])
         v[:, k+1] = v[:,k+1] - h[j,k] * v[:,j]
      end
      h[k+1,k] = norm(v[:,k+1])
      normav2 = h[k+1,k]

      # Reorthogonalize?
      if  ((reorth == 1) && (normav + .001 * normav2 == normav)) || (reorth ==  3)
         for j = 1:k
            hr = dot(v[:,j], v[:,k+1])
            h[j,k] = h[j,k] + hr
            v[:,k+1] = v[:,k+1] - hr * v[:,j]
         end
         h[k+1, k] = norm(v[:,k+1])
      end

      # Watch out for happy breakdown.
      if h[k+1,k] != 0
         v[:,k+1] = v[:,k+1] / h[k+1,k]
      end

      # Form and store the information for the new Givens rotation.
      if k > 1
         h[1:k,k] = givapp(c[1:k-1], s[1:k-1], h[1:k,k], k-1)
      end

      # Don't divide by zero if solution has  been found.
      nu = norm(h[k:k+1,k])
      if nu != 0
         # c[k] = h[k,k]/nu
         c[k] = conj(h[k,k] / nu)
         s[k] = -h[k+1,k] / nu;
         h[k,k] = c[k] * h[k,k] - s[k] * h[k+1,k]
         h[k+1,k] = 0
         g[k:k+1] = givapp(c[k], s[k], g[k:k+1], 1)
      end

      # Update the residual norm.
      rho = abs(g[k+1])
      push!(error, rho)
   end # end of the main while loop

   # At this point either k > kmax or rho < errtol.
   # It's time to compute x and leave.
   y = h[1:k,1:k] \ g[1:k]
   total_iters = k
   x = x + v[1:n, 1:k] * y

   return x, error, total_iters, total_iters
end

"""
`pushcol(V::Matrix, v::Vector) -> V`:
append the vector `v` as a column to `V`
"""
function pushcol(V::Matrix, v::Vector)
   rows, cols = size(V)
   V = V[:]
   append!(V, v)
   return reshape(V, rows, cols+1)
end

"""
`sorted_eig(A::Union{SymTridiagonal, Symmetric})`:
diagonalise A, and return sorted eigenvalues and analogously sorted eigenvectors
"""
function sorted_eig(A)
   D, Q = eigen(A)
   I = sortperm(real(D))
   return D[I], Q[:, I]
end

LinearAlgebra.dot(x, A::UniformScaling, y) = A.λ * dot(x,y)
LinearAlgebra.dot(x, A::AbstractMatrix, y) = dot(x, A*y)
LinearAlgebra.norm(P, x) = sqrt(abs(dot(x, P*x)))
dualnorm(P, f) = sqrt(abs(dot(f, P \ f)))
# nkdualnorm(P, f) = dualnorm(P, f)  # TODO: revisit this!!!
nkdualnorm(P, f) = norm(f)
# There seems to be no consistent benefit to either!

"""
`orthogonalise(w, V::Matrix, P=I) -> w'`:
Assuming V is P-orthogonal, P-orthogonalise w to V and return; if
w is already in the span of V, then a random vector is orthogonalised
and returned instead. The returned vector is normalised.
"""
function orthogonalise(w, V::Matrix, P=I; maxiter = 10, ORTHTOL=1e-10)
   @assert size(V,2) < size(V,1)
   for n = 1:maxiter
      w /= norm(P, w)
      # make w orthogonal to all columns of V
      for i = 1:size(V,2)
         w -= dot(w, P, V[:,i]) * V[:,i]
      end
      # test that the resulting w is truly orthogonal:
      nrmw = norm(P, w)
      if nrmw < 1e-15 && n == maxiter; break; end
      if nrmw < 1e-15
         w = P \ (rand(length(w)) - 0.5)
         w /= norm(P, w)
      else
         w /= norm(P, w)
         if norm( (P*w)' * V, Inf ) < ORTHTOL
            return w
         end
      end
   end
   warn("""`orthogonalise` did not terminate succesfully, most likely
      the returned vector is not orthogonal to V""")
   return w
end

"""
`appendkrylov(V, AxV, Y, v, Hmul, P)`:
appends [V v], [AxV A*v], [Y P \\ A*v]
"""
function appendkrylov(V, AxV, Y, v, Hmul, P)
   V   = pushcol(V,   v)
   AxV = pushcol(AxV, Hmul(v))
   Y   = pushcol(Y,   P \ AxV[:, end])
   return V, AxV, Y
end

"""
`immutable LanczosMatrix`: (todo: write documentation)
"""
struct LanczosMatrix
   V::Matrix{Float64}
   D::Vector{Float64}
   P
end

Base.length(K::LanczosMatrix) = length(K.D)
Base.getindex(K::LanczosMatrix, i) = K.D[i], K.V[:,i]
import Base: *, \
import LinearAlgebra: mul!
*(K::LanczosMatrix, v) = K.P * (K.V * (K.D .* (K.V' * (K.P * v))))
mul!(out::Array{Float64,1}, L::LanczosMatrix, x::Array{Float64,1}) = copy!(out, L * x)
\(K::LanczosMatrix, f) = K.V * ( K.D.^(-1) .* (K.V' * f) )
# rank
# isnewton

"spectrum transformation for `dlanczos` for minimisation"
minim(D) = abs.(D)
"spectrum transformation for `dlanczos` for index-1 saddles"
index1(D) = [-abs(D[1]); abs.(D[2:end])]
"spectrum transformation for `dlanczos` for index-2 saddles"
index2(D) = [-abs.(D[1:2]); abs.(D[3:end])]

"spectrum transformation for general order saddle"
indexp(D, p::Integer) = [-abs.(D[1:p]); abs.(D[p+1:end])]

# TODO:
#   * look into re-orthogonalising
#   * look into getting rid of AxV
#   * this code does not make proper use of the lanczos recursion
#     at all; either switch to an arnoldi method, or clean it up
#   * `isnewton` could be incorporated into LanczosMatrix

"""
dlanczos(f0, f, xc, errtol, kmax, transform; kwargs...)
      -> x, G, numf, success, isnewton

Given E : ℝᴺ → ℝ with hessian H = ∇²E(x) and b ∈ ℝᴺ  `dlanczos` computes an
approximate solution `u` to the system
   g(H) u = b
where g(H) ∈ ℝᴺˣᴺ is a transformation of J = Q D Q'  of the form
   g(H) = Q g(D) Q'
Here, D = diag(λ₁, λ₂, …) is the diagonal matrix of ordered eigenvalues
(λ₁ being the smallest). More on the transformation g see below.

`dlanczos` first uses a (possibly preconditioned) (block-)lanczos method to compute
a reasonable approximation to H in the form H = P V T V' P. Then it diagonalises
T = Q D Q', replaces D with g(D) as above. This gives an approximation G to g(H).
The method terminates if `pinv(G) b` yields a solution of sufficiently high accuracy
(residual).

In the lanzcos iterations the matrix vector product H * u
is replaced with the finite-difference operation (∇E(xc + h u) - ∇E(xc))/h;
see also `dirder` and `dirderinf`.

## Required Parameters

* `f0` : ∇E(xc)
* `f` : function to evaluate ∇E
* `xc` : current point
* `b` : right-hand side in the equation
* `errtol` : max-norm residual tolerance for termination
* `kmax` : maximum number of lanzcos iterations
* `transform` : the g transformation, default: id

## KW parameters

* `P` : preconditioner (default: I;  minimally needs to define `*` and `\\` )
* `V0` : initial subspace (default: P \\ b)
* `eigatol`, `eigrtol`: tolerance on the first eigenvalue
* `debug`: show debug information (true/false)
* `hfd` : finite-difference parameter
* `dirder` : function to compute the directional derivative; see
         `dirderinf` for format
* `ORTHTOL` : orthogonality tolerance parameter
* `Hmul` : the finite-difference operation can in principle be replaced by
         and arbitrary operator specifying the matrix-vector product

## Returns

* `x` : approximate solution
* `G` : of type `LanczosMatrix`, to extract information about the computed operator
* `numf` : number of f evaluations
* `success` : true if termination criteria are satisfied, false if kmax is reached
* `isnewton` : true if the step if the hessian spectrum was left unmodified

## Further Notes

* Acknowledgement: This code is inspired by C T Kelley's Matlab Newton-Krylov solver `nsoli.m`;
see also `nsoli` and `dgmres` which are directly ported with permission from `nsoli.m`.
* Termination criterion: If g ≠ id then G * u - b is not in fact the residual
of the equation we are trying to solve but only an approximation to that residual.
Therefore the termination criterion is only approximately satisfied.
"""
function dlanczos(f0, f, xc, b, errtol, kmax, transform = identity;
                  P = I, V0 = P \ b,
                  eigatol = 1e-1, eigrtol = 1e-1,
                  debug = false, hfd = 1e-7,
                  dirder = dirderinf,
                  Hmul = z -> dirder(f, f0, xc, z, hfd),
                  ORTHTOL = 1e-12, nevals=1 )

   # make sure V0 is given in the correct format (Matrix)
   if isa(V0, Vector)
      V0 = reshape(V0, length(V0), 1)
   end

   # initialise some variables
   p = size(V0, 2)     # block size
   d = length(f0)      # problem dimension
   @assert kmax <= d
   numf = 0            # count f evaluations
   isnewton = false    # remember whether the output is a newton direction
   success = false     # flag for termination

   # initialise Krylov subspace and more
   V = zeros(d,0)      # store the Krylov basis
   AxV = zeros(d, 0)   # store A vⱼ
   Y = zeros(d, 0)     # store P \ A vⱼ
   Q = Matrix{Float64}(undef, 0,0)
   E = Vector{Float64}(undef, 0)

   # initialise Krylov subspace; TODO: detect if a vector is in the span of previous ones
   for j = 1:size(V0, 2)
      vj = orthogonalise(V0[:,j], V, P)
      V, AxV, Y = appendkrylov(V, AxV, Y, vj, Hmul, P)
      numf += 1
   end

   # prepare for the Block-Lanczos loop
   j = 1     # the index of the krylov vector which will next be multiplied by J = A
   x = zeros(d)
   λ = sort(eigvals(Symmetric(V'*AxV)))[1:nevals]
   λ_old = Vector{Float64}[]

   if debug
      @printf("      numf     λ    err_λ     |Ax-b|/|b| \n")
   end

   # start the block-lanczos loop; when we have kmax v-vectors we stop
   while size(V, 2) <= kmax

      # At this point we have V, AxV and Y = P \ AxV  available, so we can
      # now solve the projected linear system and eigenvalue problem
      #
      # there is probably an elegant way to assemble the projected
      # linear system, but for now we just do it brute-force:
      #    A vj = P V T V' P vj
      #    yj = P \ A qj = V T  V' P vj
      #    V' P Y = V' P V T V' P V = T
      n = size(V, 2)
      T = Symmetric(V' * AxV)       # T = V' * P * Y = V' * AV
      #  TODO: could replace above with with [dot(V[:,a], P, Y[:,b] for a=1:n, b=1:n]
      #        to avoid storing AxV altogether

      # make the specified transformation
      D, Q = sorted_eig(T)
      E = transform(D)
      # residual estimate for the old x    >>> TODO: should we switch to P^{-1}-norm?
      res = nkdualnorm(P,  P * (V * (Q * (E .* (Q' * (V' * (P * x)))))) - b )
      # new x and λ (remember the old)
      g = Q * (E .\ (Q' * (V' * b)))
      x, x_old = V * g, x
      push!(λ_old, λ)
      λ = D[1:nevals]
      # if E == D then g(H) = H hence we can estimate the *actual* and *current* residual
      if (isnewton = (norm(E - D, Inf) < 1e-7) )
         res = nkdualnorm(P, AxV * g - b )    # TODO: should we switch to P^{-1}-norm?
      end
      # check for termination
      err_λ = maximum(norm(λ - λ_old[i], Inf)
                      for i = max(1,length(λ)-p+1):length(λ))
      if debug
         @printf("      %d   %.2f   %.2e    %.2e \n", numf, λ[1], err_λ, res/norm(b))
      end

      # CHECK FOR SUCCESFUL TERMINATION
      #  * if size(V,2) == d then we have assembled the entire matrix
      if (size(V,2) == d) ||  (
         (res < errtol) && ((λ == E[1:nevals]) || (err_λ < eigatol + eigrtol * abs(λ[1])))  )
         success = true
         break
      end

      # CHECK FOR UNSUCCESFUL TERMINATION
      if size(V,2) == kmax
         success = false
         break
      end

      # add the next Krylov vector
      if j <= size(V,2)   # we have Krylov vectors left to cycle through (99.9% of the time)
         w = orthogonalise(Y[:, j], V, P)
         V, AxV, Y = appendkrylov(V, AxV, Y, w, Hmul, P)
         numf += 1
         j += 1
      else
         error("how did we get here?")
      end
      # else    # no vectors left, try to add a random vector (how can this happen???)
      #    warn("no vectors left; had to add a random vector")
      #    w = orthogonalise(P \ rand(d), V, P)
      #    V, AxV, Y = appendkrylov(V, AxV, Y, w, Hmul, P)
      #    numf += 1
      # end
   end
   # if we are here it means that kmax is reached, i.e. we terminate with
   # warning or error >>> TODO: return to how to best handle this?!
   # warn("`dcg_index1` did not converge within kmax = $(kmax) iterations")
   return x, LanczosMatrix(V * Q, E, P), numf, success, isnewton
end

# """
# `immutable KrylovMatrix`: (todo: write documentation)
#
# We represent a matrix projected to a Krylov subspace as
#
# P * V * X * diagm(D) * X⁻¹ * V' * P
#
# and its inverse as
#
# V * X * diagm(D.⁻¹) * X⁻¹ * V'
# """
# immutable KrylovMatrix{T}
#    V::Matrix{T}   # orthogonal component
#    X::Matrix{T}   # non-orthogonal component
#    D::Vector{T}   # diagonal component
#    P              # preconditioner (we don't know anything about this object)
# end
#
# Base.length(K::KrylovMatrix) = length(K.D)
# Base.rank(K::KrylovMatrix) = length(K)
# Base.getindex(K::KrylovMatrix, i) = K.D[i], K.V[:,i]
# import Base: *, \, A_mul_B!
# *(K::LanczosMatrix, v) = K.P * (K.V * (K. X * (K.D .* (K.X \ (K.V' * (K.P * v))))))
# A_mul_B!(out::Array{Float64,1}, L::LanczosMatrix, x::Array{Float64,1}) = copy!(out, L * x)
# \(K::LanczosMatrix, f) = K.V * ( K.D.^(-1) .* (K.V' * f) )

"""
`stabilise(D::Vector) = D .* sign.(real.(D))`

a Lipschitz-continuous spectrum transformation that stabilises the origin
"""
stabilise(D::Vector) = D .* sign.(real.(D))

"""
darnoldi(f0, f, xc, errtol, kmax, transform; kwargs...)
      -> u, G, numf, success, isnewton

Given f : ℝᴺ → ℝᴺ with jacobian J = ∂f(x) and b ∈ ℝᴺ  `darnoldi` computes an
approximate solution `u` to the system
   g(J) u = b
where g(J) ∈ ℝᴺˣᴺ is a transformation of J = X D X⁻¹  of the form
   g(J) = X g(D) X⁻¹
Here, D = diag(λ₁, λ₂, …) is the diagonal matrix of ordered eigenvalues
(λ₁ being the smallest). More on the transformation g see below. Also
see below what happens if J is not diagonalisable.

`darnoldi` first uses a (possibly preconditioned) (block-)arnoldi iteration to compute
an approximation to  J in the form J = P V H V'. Then it diagonalises
H = X D V⁻¹, replaces D with g(D) as above. This gives an approximation G to g(J).
The method terminates if `pinv(G) b` yields a solution of sufficiently high accuracy
(residual).

In the Arnoldi iterations the matrix vector product J * u
is replaced with the finite-difference operation (f(xc + h u) - f(xc))/h;
see also `dirder` and `dirderinf`.

## Required Parameters

* `f0` : f(xc)
* `f` : function to evaluate ∇E
* `xc` : current point
* `b` : right-hand side in the equation
* `errtol` : max-norm residual tolerance for termination
* `kmax` : maximum number of lanzcos iterations
* `transform` : the g transformation, default: id

## KW parameters

* `P` : preconditioner (default: I;  minimally needs to define `*` and `\\` )
* `V0` : initial subspace (default: P \\ b)
* `debug`: show debug information (true/false)
* `hfd` : finite-difference parameter
* `dirder` : function to compute the directional derivative; see
         `dirderinf` for format
* `ORTHTOL` : orthogonality tolerance parameter
* `Hmul` : the finite-difference operation can in principle be replaced by
         an arbitrary operator specifying the matrix-vector product

## Returns

* `x` : approximate solution
/** `G` : of type `LanczosMatrix`, to extract information about the computed operator*/
* `numf` : number of f evaluations
* `success` : true if termination criteria are satisfied, false if kmax is reached
* `isnewton` : true if the hessian spectrum was left unmodified

For Further Notes see `?dlanczos`
"""
function darnoldi( f0, f, xc, b, errtol, kmax, transform = identity;
                   P = I, V0 = P \ b,
                   debug = false, hfd = 1e-7,
                   dirder = dirderinf,
                   Hmul = z -> dirder(f, f0, xc, z, hfd),
                   ORTHTOL = 1e-10 )

   # make sure V0 is given in the correct format (Matrix)
   if isa(V0, Vector)
      V0 = reshape(V0, length(V0), 1)
   end

   # initialise some variables
   p = size(V0, 2)     # block size
   d = length(f0)      # problem dimension
   @assert kmax <= d
   numf = 0            # count f evaluations
   isnewton = false    # remember whether the output is a newton direction
   success = false     # flag for termination

   # initialise Krylov subspace and more
   V = zeros(d,0)      # store the Krylov basis
   AxV = zeros(d, 0)   # store A vⱼ
   Y = zeros(d, 0)     # store P \ A vⱼ
   X = Matrix{ComplexF64}()     # H = X * D * inv(X) becomes
   Dmod = Vector{ComplexF64}()  #    Hmod = X * Dmod * inv(X)

   # initialise Krylov subspace;
   for j = 1:size(V0, 2)
      vj = orthogonalise(V0[:,j], V)
      if dot(vj, V0[:,j]) < ORTHTOL
         warn("the initial subspace matrix V0 does not have full rank")
      end
      V, AxV, Y = appendkrylov(V, AxV, Y, vj, Hmul, P)
      numf += 1
   end

   # prepare for the Arnoldi loop
   j = 1
   x = zeros(d)
   bP = P \ b

   if debug
      @printf("      numf     |Ax-b|/|b| \n")
   end

   # start the block-arnoldi process; when we have kmax v-vectors we stop
   while size(V, 2) <= kmax

      # At this point we have V, AxV and Y = P \ AxV  available, so we can
      # now solve the projected linear system and eigenvalue problem
      n = size(V, 2)
      H = V' * Y
      # make the specified spectrum transformation
      D, X = eigen(H)
      Dmod = transform(D)
      # residual estimate for the old x
      res = norm(  P * (V * (X * (Dmod .* (X \ (V'*x))))) - b )
      # new x (remember the old)
      #     A ≈ P V V' * Y V' = P V H V'     =>     A⁻¹ ≈ V H⁻¹ V' P⁻¹
      xV = X * (Dmod .\ (X \ (V' * bP)))  # coefficients of x in V
      x, x_old = V * xV, x
      # if Dmod == D then this was a newton step
      isnewton = (norm(Dmod - D, Inf) < 1e-7)
      if isnewton
         res = norm(AxV * xV - b)
      end
      if debug
         @printf("      %d    %.2e   %.2e \n", numf, res/norm(b), norm(x-x_old, Inf))
      end

      # CHECK FOR SUCCESFUL TERMINATION
      #  * if size(V,2) == d then we have assembled the entire matrix
      if (size(V,2) == d) || (res < errtol)
         success = true
         break
      end

      # CHECK FOR UNSUCCESFUL TERMINATION
      if size(V,2) == kmax
         success = false
         break
      end

      # add the next Krylov vector
      if j <= size(V,2)   # we have Krylov vectors left to cycle through
         w = orthogonalise(Y[:, j], V)
         V, AxV, Y = appendkrylov(V, AxV, Y, w, Hmul, P)
         numf += 1 # (the call is in appendkrylov)
         j += 1
      else
         error("how did we get here?")
      end
   end
   return x, numf, success, isnewton
end


using Pkg
Pkg.activate(joinpath(Pkg.devdir(), "MLCourse"))
using Plots, MLCourse, CSV, MLJ, DataFrames, MLJLinearModels, Random, Flux, MLJFlux, MLJXGBoostInterface, MLJDecisionTreeInterface
include("./Data_Processing.jl")

#Load the data
drop_std = generate(option = "drop", std = "true", valid = "false", test = "true");
med_std = generate(option = "med", std = "true", valid = "false", test = "true");
drop = generate(option = "drop", std = "false", valid = "false", test = "true");
med = generate(option = "med", std = "false", valid = "false", test = "true");

#Random Forest Classifier
m_forest_drop = machine(RandomForestClassifier(n_trees = 500), select(drop.train[:,:], Not(:precipitation_nextday)), drop.train.precipitation_nextday) |> fit!
m_forest_med = machine(RandomForestClassifier(n_trees = 1200), select(med.train[:,:], Not(:precipitation_nextday)), med.train.precipitation_nextday) |> fit!
m_forest_drop_std = machine(RandomForestClassifier(n_trees = 500), select(drop_std.train[:,:], Not(:precipitation_nextday)), drop_std.train.precipitation_nextday) |> fit!
m_forest_med_std = machine(RandomForestClassifier(n_trees = 1200), select(med_std.train[:,:], Not(:precipitation_nextday)), med_std.train.precipitation_nextday) |> fit!

auc_forest_drop = MLJ.auc(predict(m_forest_drop, select(drop.test[:,:], Not(:precipitation_nextday))), drop.test.precipitation_nextday) # 0.909
auc_forest_med = MLJ.auc(predict(m_forest_med, select(med.test[:,:], Not(:precipitation_nextday))), med.test.precipitation_nextday) # 0.927
auc_forest_drop_std = MLJ.auc(predict(m_forest_drop_std, select(drop_std.test[:,:], Not(:precipitation_nextday))), drop_std.test.precipitation_nextday) # 0.908
auc_forest_med_std = MLJ.auc(predict(m_forest_med_std, select(med_std.test[:,:], Not(:precipitation_nextday))), med_std.test.precipitation_nextday) # 0.926

#The med data seems to have the best AUC (not a big difference between standardized or not)

function TunedModel_forest(data)
    Random.seed!(2711)
    model = RandomForestClassifier()
    self_tuning_model = TunedModel(model = model,
                                resampling = CV(nfolds = 5),
                                tuning = Grid(),
                                range = range(model, :n_trees, values=100:50:1200),
                                measure = auc)
    self_tuning_mach = machine(self_tuning_model,
                            select(data, Not(:precipitation_nextday)),
                            data.precipitation_nextday) |> fit!
    self_tuning_mach
end

#Test the tuned model on the drop data set and med data set with both standardized and not-standardized data:
rep1 = report(TunedModel_forest(drop_std.train));
rep1.best_history_entry.measurement
rep1.best_model.n_trees
scatter(reshape(rep1.plotting.parameter_values, :), rep1.plotting.measurements, xlabel = "N", ylabel = "AUC")

rep2 = report(TunedModel_forest(med_std.train));
rep2.best_history_entry.measurement
rep2.best_model.n_trees
scatter(reshape(rep2.plotting.parameter_values, :), rep2.plotting.measurements, xlabel = "N", ylabel = "AUC")

rep3 = report(TunedModel_forest(drop.train));
rep3.best_history_entry.measurement
rep3.best_model.n_trees
scatter(reshape(rep3.plotting.parameter_values, :), rep3.plotting.measurements, xlabel = "N", ylabel = "AUC")

rep4 = report(TunedModel_forest(med.train));
rep4.best_history_entry.measurement
rep4.best_model.n_trees
scatter(reshape(rep4.plotting.parameter_values, :), rep4.plotting.measurements, xlabel = "N", ylabel = "AUC")

#Test the best model over the test set with the best hyper-parameters to pick the best one:
best_mach1 = machine(RandomForestClassifier(n_trees = rep1.best_model.n_trees), select(drop_std.train[:,:], Not(:precipitation_nextday)), drop_std.train.precipitation_nextday)|> fit!
best_mach2 = machine(RandomForestClassifier(n_trees = rep2.best_model.n_trees), select(med_std.train[:,:], Not(:precipitation_nextday)), med_std.train.precipitation_nextday)|> fit!
best_mach3 = machine(RandomForestClassifier(n_trees = rep3.best_model.n_trees), select(drop.train[:,:], Not(:precipitation_nextday)), drop.train.precipitation_nextday)|> fit!
best_mach4 = machine(RandomForestClassifier(n_trees = rep4.best_model.n_trees), select(med.train[:,:], Not(:precipitation_nextday)), med.train.precipitation_nextday)|> fit!
losses(best_mach1, select(drop_std.test, Not(:precipitation_nextday)), drop_std.test.precipitation_nextday)
losses(best_mach2, select(med_std.test, Not(:precipitation_nextday)), med_std.test.precipitation_nextday)
losses(best_mach3, select(drop.test, Not(:precipitation_nextday)), drop.test.precipitation_nextday)
losses(best_mach4, select(med.test, Not(:precipitation_nextday)), med.test.precipitation_nextday)

#Write in the submission file with the best machine trained on all data
train, test = generate(option = "med", std = "true", valid = "false", test = "false");
best_mach = machine(RandomForestClassifier(n_trees = rep2.best_model.n_trees), select(train[:,:], Not(:precipitation_nextday)), train.precipitation_nextday) |> fit!

sample = CSV.read(joinpath(@__DIR__, "data", "sample_submission.csv"), DataFrame);
pred = pdf.(predict(best_mach, test), true);
sample.precipitation_nextday = pred;
CSV.write(joinpath(@__DIR__, "results", "RandomForest_submission.csv"), sample);


function charge(sequence, pKvalues = pKtable["solomon"], pH = 7.0)
    aa = BioSequences.composition(sequence)

    charge = 0.0

    charge += (1 / (1 + 10^(1 * (pH - pKvalues[1]))))
    charge += (-1 / (1 + 10^(-1 * (pH - pKvalues[2]))))

    charge += aa[AA_C] * (-1 / (1 + 10^(-1 * (pH - pKvalues[3]))))
    charge += aa[AA_D] * (-1 / (1 + 10^(-1 * (pH - pKvalues[4]))))
    charge += aa[AA_E] * (-1 / (1 + 10^(-1 * (pH - pKvalues[5]))))
    charge += aa[AA_Y] * (-1 / (1 + 10^(-1 * (pH - pKvalues[9]))))
    charge += aa[AA_H] * (1 / (1 + 10^(1 * (pH - pKvalues[6]))))
    charge += aa[AA_K] * (1 / (1 + 10^(1 * (pH - pKvalues[7]))))
    charge += aa[AA_R] * (1 / (1 + 10^(1 * (pH - pKvalues[8]))))

    return charge
end

"""
    pI(sequence, [pKvalues; gamma])

Estimate isoelectric point of `sequence`. `pKvalues` can be a vector containing pK values or a string specifying one of the pK-value sources in `pKtable`. See the docstring for `pKtable` for details. Default is "IPC protein".
"""
pI(sequence::LongAminoAcidSeq, pKvalues::AbstractString; gamma = 0.001) = pI(sequence, pKtable[pKvalues]; gamma = gamma)
function pI(sequence::LongAminoAcidSeq, pKvalues = pKtable["IPC protein"]; gamma = 0.001)
    pH = 7.0

    if charge(sequence, pKvalues, pH) < 0
        pHs = 0:gamma:7
        charges = charge.(Ref(sequence), Ref(pKvalues), pHs)
        pH = pHs[findmin(abs.(charges))[2]]
        return pH
    end

    if charge(sequence, pKvalues, pH) > 0
        pHs = 7:gamma:14
        charges = charge.(Ref(sequence), Ref(pKvalues), pHs)
        pH = pHs[findmin(abs.(charges))[2]]
        return pH
    end

    return pH
end


using Arblib: Acb, AcbRef, AcbVector, AcbRefVector, AcbMatrix, AcbRefMatrix

export certify,
    SolutionCertificate,
    CertificationResult,
    CertificationCache,
    is_certified,
    is_real,
    is_complex,
    is_positive,
    solution_candidate,
    certified_solution_interval,
    certified_solution_interval_after_krawczyk,
    certificate_index,
    solution_approximation,
    certificates,
    distinct_certificates,
    ncertified,
    nreal_certified,
    ncomplex_certified,
    ndistinct_certified,
    ndistinct_real_certified,
    ndistinct_complex_certified,
    show_straight_line_program,
    save,
    # deprecated
    initial_solution,
    certified_solution

"""
    SolutionCertificate

Result of [`certify`](@ref) for a single solution. Contains the initial solutions
and if the certification was successfull a vector of complex intervals where the true
solution is contained in.
"""
Base.@kwdef struct SolutionCertificate
    solution_candidate::AbstractVector
    certified::Bool
    real::Bool = false
    complex::Bool = false
    index::Union{Nothing,Int} = nothing
    prec::Int = 53
    # I, I′ ∈ 𝕀ℂⁿ and a certified solution has I′ ⊊ I
    I::Union{Nothing,AcbMatrix} = nothing
    I′::Union{Nothing,AcbMatrix} = nothing
    x̃::Union{Nothing,AcbMatrix} = nothing
    Y::Union{Nothing,AcbMatrix} = nothing

    # We also store a double precision representation of the midpoint of x₁
    # as the best available double precision estimate of the solution
    solution::Union{Nothing,Vector{ComplexF64}} = nothing
end

"""
    solution_candidate(certificate::SolutionCertificate)

Returns the given provided solution candidate.
"""
solution_candidate(C::SolutionCertificate) = C.solution_candidate
@deprecate initial_solution(cert::SolutionCertificate) solution_candidate(cert)

"""
    is_certified(certificate::SolutionCertificate)

Returns `true` if `certificate` is a certificate that `certified_solution_interval(certificate)`
contains a unique zero.
"""
is_certified(C::SolutionCertificate) = C.certified

"""
    is_real(certificate::SolutionCertificate)

Returns `true` if `certificate` certifies that the certified solution interval
contains a true real zero of the system.
If `false` is returned then this does not necessarily mean that the true solution is not real.
"""
is_real(C::SolutionCertificate) = C.real

"""
    is_complex(certificate::SolutionCertificate)

Returns `true` if `certificate` certifies that the certified solution interval
contains a true complex zero of the system.
"""
is_complex(C::SolutionCertificate) = C.complex

"""
    is_positive(certificate::SolutionCertificate)

Returns `true` if `is_certified(certificate)` is `true` and the unique zero contained
in `certified_solution_interval(certificate)` is real and positive.
"""
function is_positive(C::SolutionCertificate)
    if !is_certified(C) || !is_real(C)
        return false
    else
        all(i -> Arblib.is_positive(real(Arblib.ref(C.I′, i, 1))), 1:length(C.I′))
    end
end

"""
    is_positive(certificate::SolutionCertificate, i)

Returns `true` if `is_certified(certificate)` is `true` and `i`-th coordinate of the
unique zero contained in `certified_solution_interval(certificate)` is real and positive.
"""
function is_positive(C::SolutionCertificate, i::Integer)
    if !is_certified(C) || !is_real(C)
        return false
    else
        Arblib.is_positive(real(Arblib.ref(C.I′, i, 1)))
    end
end

"""
    certified_solution_interval(certificate::SolutionCertificate)

Returns an `Arblib.AcbMatrix` representing a vector of complex intervals where a unique
zero of the system is contained in.
Returns `nothing` if `is_certified(certificate)` is `false`.
"""
certified_solution_interval(certificate::SolutionCertificate) = certificate.I
@deprecate certified_solution(certificate) certified_solution_interval(certificate)

"""
    certified_solution_interval_after_krawczyk(certificate::SolutionCertificate)

Returns an `Arblib.AcbMatrix` representing a vector of complex intervals where a unique
zero of the system is contained in.
This is the result of applying the Krawczyk operator to `certified_solution_interval(certificate)`.
Returns `nothing` if `is_certified(certificate)` is `false`.
"""
certified_solution_interval_after_krawczyk(certificate::SolutionCertificate) =
    certificate.I′

"""
    certificate_index(certificate::SolutionCertificate)

Return the index of the solution certificate. Here the index refers to the index of the
provided solution candidates.
"""
certificate_index(certificate::SolutionCertificate) = C.index

"""
    precision(certificate::SolutionCertificate)

Returns the maximal precision used to produce the given `certificate`.
"""
Base.precision(certificate::SolutionCertificate) = certificate.prec

"""
    solution_approximation(certificate::SolutionCertificate)

If `is_certified(certificate)` is `true` this returns the midpoint of the
[`certified_solution_interval`](@ref) of the given `certificate` as a `Vector{ComplexF64}`.
Returns `nothing` if `is_certified(certificate)` is `false`.
"""
solution_approximation(certificate::SolutionCertificate) = certificate.solution

"""
    krawczyk_operator_parameters(cert::SolutionCertificate)

Returns a `NamedTuple` `(x, Y)` with the parameters of the Krawczyk operator following [BRT20].
"""
krawczyk_operator_parameters(cert::SolutionCertificate) = (x = cert.x̃, Y = cert.Y)

function Base.show(f::IO, cert::SolutionCertificate)
    println(f, "SolutionCertificate:")
    println(f, "solution_candidate = [")
    for z in solution_candidate(cert)
        print(f, "  ")
        print(f, z)
        println(f, ",")
    end
    println(f, "]")
    print(f, "is_certified = ", is_certified(cert))
    if !isnothing(certified_solution_interval(cert))
        println(f)
        println(f, "certified_solution_interval = [")
        for z in certified_solution_interval(cert)
            print(f, "  ")
            print(f, z)
            println(f, ",")
        end
        println(f, "]")
        println(f, "precision = ", cert.prec)
        print(f, "is_real = ", is_real(cert))
    end
    if !isnothing(cert.index)
        println(f)
        print(f, "index = ", cert.index)
    end
end

"""
    CertificationResult

The result of [`certify`](@ref) for multiple solutions.
Contains a vector of [`SolutionCertificate`](@ref) as well as a list of certificates
which correspond to the same true solution.
"""
struct CertificationResult
    certificates::Vector{SolutionCertificate}
    duplicates::Vector{Vector{Int}}
    slp::ModelKit.Interpreter
    slp_jacobian::ModelKit.Interpreter
end

"""
    certificates(R::CertificationResult)

Obtain the stored [`SolutionCertificate`](@ref)s.
"""
certificates(R::CertificationResult) = R.certificates

"""
    distinct_certificates(R::CertificationResult)

Obtain the certificates corresponding to the determined distinct solution intervals.
"""
function distinct_certificates(C::CertificationResult)
    cs = certificates(C)
    isempty(C.duplicates) && return cs
    indices = trues(length(cs))
    for d in C.duplicates, k = 2:length(d)
        indices[d[k]] = false
    end
    cs[indices]
end

"""
    ncertified(R::CertificationResult)

Returns the number of certified solutions.
"""
ncertified(R::CertificationResult) = count(is_certified, R.certificates)

"""
    nreal_certified(R::CertificationResult)

Returns the number of certified real solutions.
"""
nreal_certified(R::CertificationResult) =
    count(r -> is_certified(r) && is_real(r), R.certificates)

"""
    ncomplex_certified(R::CertificationResult)

Returns the number of certified complex solutions.
"""
ncomplex_certified(R::CertificationResult) =
    count(r -> is_certified(r) && is_complex(r), R.certificates)

"""
    ndistinct_certified(R::CertificationResult)

Returns the number of distinct certified solutions.
"""
function ndistinct_certified(R::CertificationResult)
    ncert = ncertified(R)
    if isempty(R.duplicates)
        return ncert
    else
        return ncert - sum(length, R.duplicates) + length(R.duplicates)
    end
end

"""
    ndistinct_real_certified(R::CertificationResult)

Returns the number of distinct certified real solutions.
"""
function ndistinct_real_certified(R::CertificationResult)
    ncert = nreal_certified(R)
    if isempty(R.duplicates)
        return ncert
    else
        ncert - sum(R.duplicates) do dup
            is_real(R.certificates[dup[1]]) ? length(dup) - 1 : 0
        end
    end
end

"""
    ndistinct_complex_certified(R::CertificationResult)

Returns the number of distinct certified complex solutions.
"""
function ndistinct_complex_certified(R::CertificationResult)
    ncert = ncomplex_certified(R)
    if isempty(R.duplicates)
        return ncert
    else
        ncert - sum(R.duplicates) do dup
            is_complex(R.certificates[dup[1]]) ? length(dup) - 1 : 0
        end
    end
end

"""
    show_straight_line_program(R::CertificationResult)
    show_straight_line_program(io::IO, R::CertificationResult)

Print a representation of the used straight line program.
"""
show_straight_line_program(R::CertificationResult) = show_straight_line_program(stdout, R)
show_straight_line_program(io::IO, R::CertificationResult) =
    ModelKit.show_instructions(io, R.slp)

function Base.show(io::IO, R::CertificationResult)
    println(io, "CertificationResult")
    println(io, "===================")
    println(io, "• $(length(R.certificates)) solution candidates given")
    ncert = ncertified(R)
    print(io, "• $ncert certified solution intervals")
    nreal = nreal_certified(R)
    ncomplex = ncomplex_certified(R)
    print(io, " ($nreal real, $ncomplex complex")
    if nreal + ncomplex < ncert
        println(io, ", $(ncert - (nreal + ncomplex)) undecided)")
    else
        println(io, ")")
    end
    ndist_cert = ndistinct_certified(R)
    print(io, "• $ndist_cert distinct certified solution intervals")
    ndist_real = ndistinct_real_certified(R)
    ndist_complex = ndistinct_complex_certified(R)
    print(io, " ($ndist_real real, $ndist_complex complex")
    if ndist_real + ndist_complex < ndist_cert
        print(io, ", $(ndist_cert - (ndist_real + ndist_complex)) undecided)")
    else
        print(io, ")")
    end
end

"""
    save(filename, C::CertificationResult)

Store a text representation of the certification result `C` on disk.
"""
function save(filename, R::CertificationResult)
    open(filename, "w") do f
        println(f, "## Summary")
        show(f, R)
        println(f, "\n\n## Certificates")
        for cert in certificates(R)
            show(f, cert)
            println(f)
        end
    end
    filename
end

"""
    CertificationCache(F::AbstractSystem)

Contains the necessary data structures for [`certify`](@ref).
"""
Base.@kwdef mutable struct CertificationCache{T₁,T₂}
    eval_interpreter::ModelKit.Interpreter{T₁}
    jac_interpreter::ModelKit.Interpreter{T₂}
    newton_cache::NewtonCache{MatrixWorkspace{Matrix{ComplexF64}}}
    # data for krawczyc_step
    C::Matrix{ComplexF64}
    r₀::Vector{IComplexF64}
    Δx₀::Vector{IComplexF64}
    ix̃₀::Vector{IComplexF64}
    Jx₀::Matrix{IComplexF64}
    M::Matrix{IComplexF64}
    δx::Vector{IComplexF64}
    interpreter_cache::ModelKit.InterpreterCache{Vector{IComplexF64},Nothing}

    arb_prec::Int
    arb_C::AcbRefMatrix
    arb_r₀::AcbRefMatrix # m × 1
    arb_Δx₀::AcbRefMatrix # m × 1
    arb_x̃₀::AcbRefMatrix # m × 1
    arb_x₀::AcbRefMatrix # m × 1
    arb_x₁::AcbRefMatrix # m × 1
    arb_J_x₀::AcbRefMatrix
    arb_M::AcbRefMatrix
    arb_δx::AcbRefMatrix # m × 1
    arb_mag::Arblib.Mag
    arb_interpreter_cache::ModelKit.InterpreterCache{AcbRefVector,Acb}
end
function CertificationCache(F::AbstractSystem)
    m, n = size(F)
    m == n || error("We can only certify square systems.")
    f = System(F)
    interpreter = ModelKit.interpreter(f)
    jac_interpreter = ModelKit.jacobian_interpreter(f)
    tape_length = max(
        ModelKit.cache_min_length(interpreter),
        ModelKit.cache_min_length(jac_interpreter),
    )
    arb_prec = 128
    CertificationCache(;
        eval_interpreter = interpreter,
        jac_interpreter = jac_interpreter,
        newton_cache = NewtonCache(F; optimize_data_structure = false),
        C = zeros(ComplexF64, m, m),
        r₀ = zeros(IComplexF64, m),
        Δx₀ = zeros(IComplexF64, m),
        ix̃₀ = zeros(IComplexF64, m),
        Jx₀ = zeros(IComplexF64, m, m),
        M = zeros(IComplexF64, m, m),
        δx = zeros(IComplexF64, m),
        interpreter_cache = ModelKit.InterpreterCache(IComplexF64, tape_length),
        arb_prec = arb_prec,
        arb_C = AcbRefMatrix(m, m; prec = arb_prec),
        arb_r₀ = AcbRefMatrix(m, 1; prec = arb_prec),
        arb_Δx₀ = AcbRefMatrix(m, 1; prec = arb_prec),
        arb_x̃₀ = AcbRefMatrix(m, 1; prec = arb_prec),
        arb_x₀ = AcbRefMatrix(m, 1; prec = arb_prec),
        arb_x₁ = AcbRefMatrix(m, 1; prec = arb_prec),
        arb_J_x₀ = AcbRefMatrix(m, m; prec = arb_prec),
        arb_M = AcbRefMatrix(m, m; prec = arb_prec),
        arb_δx = AcbRefMatrix(m, 1; prec = arb_prec),
        arb_mag = Arblib.Mag(),
        arb_interpreter_cache = ModelKit.InterpreterCache(
            Acb,
            tape_length;
            prec = arb_prec,
        ),
    )
end

Base.setprecision(M::AcbRefMatrix, p::Int) = AcbRefMatrix(M.acb_mat, p)
function set_arb_precision!(cache::CertificationCache, p::Int)
    cache.arb_prec == p && return cache
    cache.arb_prec = p
    cache.arb_r₀ = setprecision(cache.arb_r₀, p)
    cache.arb_Δx₀ = setprecision(cache.arb_Δx₀, p)
    cache.arb_x̃₀ = setprecision(cache.arb_x̃₀, p)
    cache.arb_x₀ = setprecision(cache.arb_x₀, p)
    cache.arb_x₁ = setprecision(cache.arb_x₁, p)
    cache.arb_J_x₀ = setprecision(cache.arb_J_x₀, p)
    cache.arb_M = setprecision(cache.arb_M, p)
    cache.arb_δx = setprecision(cache.arb_δx, p)
    ModelKit.set_arb_precision!(cache.arb_interpreter_cache, p)

    cache
end

"""
    certify(F, solutions, [p, certify_cache]; options...)
    certify(F, result, [p, certify_cache]; options...)

Attempt to certify that the given approximate `solutions` correspond to true solutions
of the polynomial system ``F(x;p)``. The system ``F`` has to be an (affine) square
polynomial system. Also attemps to certify for each solutions whether it
approximates a real solution. The certification is done using interval arithmetic and the
Krawczyk method[^Moo77]. Returns a [`CertificationResult`](@ref) which additionall returns
the number of distinct solutions. For more details of the implementation see [^BRT20].

## Options

* `show_progress = true`: If `true` shows a progress bar of the certification process.
* `max_precision = 256`: The maximal accuracy (in bits) that is used in the certification process.
* `compile = false`: See the [`solve`](@ref) documentation.

## Example
We take the [first example](https://www.juliahomotopycontinuation.org/guides/introduction/#a-first-example) from our
introduction guide.
```julia
@var x y
# define the polynomials
f₁ = (x^4 + y^4 - 1) * (x^2 + y^2 - 2) + x^5 * y
f₂ = x^2+2x*y^2 - 2y^2 - 1/2
F = System([f₁, f₂], variables = [x,y])
result = solve(F)
```
```
Result with 18 solutions
========================
• 18 paths tracked
• 18 non-singular solutions (4 real)
• random seed: 0xcaa483cd
• start_system: :polyhedral
```
We see that we obtain 18 solutions and it seems that 4 solutions are real. However,
this is based on heuristics. To be absolute certain we can certify the result

```julia
certify(F, result)
```
```
CertificationResult
===================
• 18 solution candidates given
• 18 certified solution intervals (4 real, 14 complex)
• 18 distinct certified solution intervals (4 real, 14 complex)
```

and see that there are indeed 18 solutions and that they are all distinct.

[^Moo77]: Moore, Ramon E. "A test for existence of solutions to nonlinear systems." SIAM Journal on Numerical Analysis 14.4 (1977): 611-615.
[^BRT20]: Breiding, P., Rose, K. and Timme, S. "Certifying zeros of polynomial systems using interval arithmetic." arXiv:2011.05000.
"""
function certify end

struct CertificationParameters
    params::Vector{ComplexF64}
    interval_params::Vector{IComplexF64}
    arb_interval_params::AcbRefVector
end

function CertificationParameters(p::AbstractVector; prec::Int = 256)
    arb_ip = AcbRefVector(length(p); prec = prec)
    for (i, p_i) in enumerate(p)
        x = arb_ip[i]
        x[] = p_i
    end
    CertificationParameters(
        convert(Vector{ComplexF64}, p),
        convert(Vector{IComplexF64}, p),
        arb_ip,
    )
end

certification_parameters(p::AbstractVector; prec::Int = 256) = CertificationParameters(p)
certification_parameters(::Nothing; prec::Int = 256) = nothing

complexF64_params(C::CertificationParameters) = C.params
complexF64_interval_params(C::CertificationParameters) = C.interval_params
arb_interval_params(C::CertificationParameters) = C.arb_interval_params
complexF64_params(::Nothing) = nothing
complexF64_interval_params(::Nothing) = nothing
arb_interval_params(::Nothing) = nothing

function _certify(
    F::AbstractSystem,
    solution_candidates::AbstractArray{<:AbstractArray{<:Number}},
    p::Union{Nothing,CertificationParameters},
    cache::CertificationCache;
    show_progress::Bool = true,
    max_precision::Int = 256,
    check_distinct::Bool = true,
)
    certificates = SolutionCertificate[]

    m, n = size(F)
    m == n || throw(ArgumentError("We can only certify solutions to square systems."))

    if isnothing(p) && nparameters(System(F)) > 0
        throw(ArgumentError("The given system expects parameters but none are given."))
    end

    if !show_progress
        for (i, s) in enumerate(solution_candidates)
            push!(certificates, certify_solution(F, s, p, cache, i))
        end
    else
        # Create progress meter
        n = length(solution_candidates)
        desc = "Certifying $n solutions... "
        barlen = min(ProgressMeter.tty_width(desc, stdout, false), 40)
        progress = ProgressMeter.Progress(
            n;
            dt = 0.2,
            desc = desc,
            barlen = barlen,
            color = :green,
            output = stdout,
        )
        progress.tlast += progress.dt

        # Go over all solution
        ncertified = 0
        nreal_certified = 0
        for (i, s) in enumerate(solution_candidates)
            r = certify_solution(F, s, p, cache, i)
            push!(certificates, r)
            ncertified += is_certified(r)
            nreal_certified += is_certified(r) && is_real(r)
            update_certify_progress!(progress, i, ncertified, nreal_certified)
        end
    end
    if check_distinct
        duplicates = find_duplicates(certificates)
    else
        duplicates = Vector{Vector{Int}}()
    end

    CertificationResult(
        certificates,
        duplicates,
        cache.eval_interpreter,
        cache.jac_interpreter,
    )
end

function update_certify_progress!(progress, k, ncertified, nreal_certified)
    t = time()
    if k == progress.n || t > progress.tlast + progress.dt
        showvalues = make_certify_showvalues(k, ncertified, nreal_certified)
        ProgressMeter.update!(progress, k; showvalues = showvalues)
    end
    nothing
end
@noinline function make_certify_showvalues(k, ncertified, nreal_certified)
    (
        ("# solutions candidates considered", k),
        ("# certified solution intervals (real)", "$(ncertified) ($nreal_certified)"),
    )
end

function certify_solution(
    F::AbstractSystem,
    solution_candidate::AbstractVector,
    cert_params::Union{Nothing,CertificationParameters},
    cert_cache::CertificationCache,
    index::Int,
    max_precision::Int = 256,
    refine_solution::Bool = true,
)
    @unpack C, arb_C, arb_x̃₀ = cert_cache
    n = size(C, 1)

    # refine solution to machine precicision
    res = newton(
        F,
        solution_candidate,
        complexF64_params(cert_params),
        InfNorm(),
        cert_cache.newton_cache;
        rtol = 8 * eps(),
        atol = 0.0,
        # this should already be an approximate zero
        extended_precision = true,
        max_iters = 8,
    )

    x̃₀ = solution(res)
    LA.inv!(C, cert_cache.newton_cache.J)

    certified, x₁, x₀, is_real = ε_inflation_krawczyk(x̃₀, cert_params, C, cert_cache)

    if certified
        return SolutionCertificate(
            solution_candidate = solution_candidate,
            certified = true,
            real = is_real,
            complex = any(xi -> !(0.0 in imag(xi)), x₁),
            index = index,
            prec = 53,
            I = AcbMatrix(x₀; prec = 53),
            I′ = AcbMatrix(x₁; prec = 53),
            x̃ = AcbMatrix(x̃₀; prec = 53),
            Y = AcbMatrix(C; prec = 53),
            solution = mid.(x₁),
        )
    end

    # If not certified we do the computation in higher precision using Arb
    prec = 128
    set_arb_precision!(cert_cache, prec)

    # We keep the same C matrix for now.
    for j = 1:n, i = 1:n
        arb_C[i, j][] = C[i, j]
    end
    for (i, x̃₀_i) in enumerate(x̃₀)
        arb_x̃₀[i][] = x̃₀_i
    end
    while (prec <= max_precision)
        certified, arb_x₁, arb_x₀, is_real =
            arb_ε_inflation_krawczyk(arb_x̃₀, cert_params, arb_C, cert_cache; prec = prec)

        if certified
            I = AcbMatrix(n, 1; prec = prec)
            I′ = AcbMatrix(n, 1; prec = prec)
            x̃ = AcbMatrix(n, 1; prec = prec)
            Y = AcbMatrix(n, n; prec = prec)
            Arblib.set!(I, arb_x₀)
            Arblib.set!(I′, arb_x₁)
            Arblib.set!(x̃, arb_x̃₀)
            Arblib.set!(Y, arb_C)
            return SolutionCertificate(
                solution_candidate = solution_candidate,
                certified = true,
                real = is_real,
                complex = any(xi -> !Arblib.contains_zero(Arblib.imagref(xi)), arb_x₁),
                index = index,
                prec = prec,
                I = I,
                I′ = I′,
                x̃ = x̃,
                Y = Y,
                solution = [ComplexF64(arb_x₁[i]) for i = 1:n],
            )
        end

        prec += 64
        set_arb_precision!(cert_cache, prec)
    end

    # certification failed
    return SolutionCertificate(
        solution_candidate = solution_candidate,
        certified = false,
        index = index,
    )
end

function Base.Vector{ComplexF64}(A::Arblib.AcbMatrixLike)
    @assert size(A, 2) == 1
    [ComplexF64(Arblib.ref(A, i, 1).acb_ptr) for i = 1:size(A, 1)]
end

function ε_inflation_krawczyk(x̃₀, p::Union{Nothing,CertificationParameters}, C, cert_cache)
    @unpack C, r₀, Δx₀, ix̃₀, Jx₀, M, δx, interpreter_cache = cert_cache
    J_x₀ = Jx₀

    ix̃₀ .= IComplexF64.(x̃₀)
    # r₀ = F(ix̃₀)
    ModelKit.execute!(
        r₀,
        cert_cache.eval_interpreter,
        ix̃₀,
        complexF64_interval_params(p),
        interpreter_cache,
    )
    # iΔx = C * r₀
    sqr_mul!(Δx₀, C, r₀)

    # Perform ε-inflation. We choose a different εᵢ per coordinate. This matches our strategy
    # to use a weighted norm.
    x₀ = map(x̃₀, Δx₀) do x̃₀_i, Δx₀_i
        εᵢ = 512 * max(mag(Δx₀_i), eps())
        complex(
            Interval(real(x̃₀_i) - εᵢ, real(x̃₀_i) + εᵢ),
            Interval(imag(x̃₀_i) - εᵢ, imag(x̃₀_i) + εᵢ),
        )
    end

    ModelKit.execute!(
        nothing,
        J_x₀,
        cert_cache.jac_interpreter,
        x₀,
        complexF64_interval_params(p),
        interpreter_cache,
    )

    x₁ = similar(x₀)
    # x₁ = (x̃₀ - C * F([x̃₀])) + (I - C * J(x₀)) * (x₀ - x̃₀)
    #    = (x̃₀ - C * F([x̃₀])) - (C * J(x₀) - I) * (x₀ - x̃₀)
    #    = (x̃₀ - Δx₀) - (C * J(x₀) - I) * (x₀ - x̃₀)

    # Define M = C * J(x₀) - I
    sqr_mul!(M, C, J_x₀)
    for i = 1:size(M, 1)
        M[i, i] -= 1
    end
    # Necessary condition is ||M|| < 1 / √2
    # We lower bound 1 / √2 by 0.7071
    if IntervalArithmetic.inf_norm_bound(M) < 0.7071
        δx .= x₀ .- x̃₀
        sqr_mul!(x₁, M, δx)
        x₁ .= (x̃₀ .- Δx₀) .- x₁
        certified = all2(isinterior, x₁, x₀)
        is_real =
            certified ? all2((x₁_i, x₀_i) -> isinterior(conj(x₁_i), x₀_i), x₁, x₀) : false
    else
        certified = false
        is_real = false
    end
    certified, x₁, x₀, is_real
end

function arb_ε_inflation_krawczyk(
    x̃₀::AcbRefMatrix,
    p::Union{Nothing,CertificationParameters},
    C::Arblib.AcbMatrixLike,
    cert_cache;
    prec::Int,
)
    @unpack arb_r₀, arb_Δx₀, arb_x̃₀, arb_J_x₀, arb_M, arb_δx, arb_mag, arb_x₀, arb_x₁ =
        cert_cache
    r₀, Δx₀, x̃₀, x₀, x₁, J_x₀, M, δx, m =
        arb_r₀, arb_Δx₀, arb_x̃₀, arb_x₀, arb_x₁, arb_J_x₀, arb_M, arb_δx, arb_mag
    @unpack arb_interpreter_cache = cert_cache

    m = arb_mag
    x̃₀′ = x̄₁ = Δx₀
    # Perform a simple newton refinement using arb_C until we cannot improve the
    # accuracy anymore
    acc = Inf
    max_iters = 10
    for i = 1:max_iters
        ModelKit.execute!(
            r₀,
            cert_cache.eval_interpreter,
            x̃₀,
            arb_interval_params(p),
            arb_interpreter_cache,
        )
        Arblib.mul!(Δx₀, C, r₀)

        acc_new = Arblib.get(Arblib.bound_inf_norm!(m, Δx₀))
        if acc_new > acc || i == max_iters
            break
        end
        acc = acc_new
        Arblib.sub!(x̃₀, x̃₀, Δx₀)
        Arblib.get_mid!(x̃₀, x̃₀)
    end

    # Perform ε-inflation
    n = length(x̃₀)
    Arblib.get_mid!(x₀, x̃₀)
    # We increase the radius by 2^(prec/4) to avoid hitting the precision limit.
    # We choose a dynamic increase to account for bad situations where any fixed choice
    # would be insufficient.
    incr_factor = exp2(div(prec, 4))
    mach_eps = exp2(-prec)
    for i = 1:n
        m[] = max(magF64(Δx₀[i], m), mach_eps) * incr_factor
        Arblib.add_error!(x₀[i], m)
    end

    ModelKit.execute!(
        nothing,
        J_x₀,
        cert_cache.jac_interpreter,
        x₀,
        arb_interval_params(p),
        arb_interpreter_cache,
    )

    # x₁ = (x̃₀ - C * F([x̃₀])) + (I - C * J(x₀)) * (x₀ - x̃₀)
    #    = (x̃₀ - C * F([x̃₀])) - (C * J(x₀) - I) * (x₀ - x̃₀)
    #    = (x̃₀ - Δx₀) - (C * J(x₀) - I) * (x₀ - x̃₀)

    # Define M = C * J(x₀) - I
    Arblib.mul!(M, C, J_x₀)
    for i = 1:n
        Arblib.sub!(M[i, i], M[i, i], 1)
    end

    # Necessary condition is ||M|| < 1 / √2
    # We lower bound 1 / √2 by 0.7071
    if Arblib.get(Arblib.bound_inf_norm!(m, M)) < 0.7071
        # x₀ - x̃₀
        Arblib.sub!(δx, x₀, x̃₀)
        # (C * J(x₀) - I) * (x₀ - x̃₀)
        Arblib.mul!(x₁, M, δx)
        # (x̃₀ - Δx₀)
        Arblib.sub!(x̃₀′, x̃₀, Δx₀)
        # (x̃₀ - Δx₀) - (C * J(x₀) - I) * (x₀ - x̃₀)
        Arblib.sub!(x₁, x̃₀′, x₁)

        certified = Bool(Arblib.contains(x₀, x₁))

        # check realness
        Arblib.conjugate!(x̄₁, x₁)
        is_real = Bool(Arblib.contains(x₀, x̄₁))
    else
        certified = false
        is_real = false
    end

    if !certified
        # Update the approximation of the inverse
        Arblib.get_mid!(J_x₀, J_x₀)
        Arblib.inv!(C, J_x₀)
    end

    certified, AcbMatrix(x₁), AcbMatrix(x₀), is_real
end

magF64(x, mag) = Arblib.get(Arblib.get!(mag, x))

"Simple matrix multiplication using the muladd method which is much faster four our case."
function sqr_mul!(C, A, B)
    n = size(A, 1)
    C .= zero(eltype(C))
    @inbounds for j = 1:size(B, 2)
        for k = 1:n
            bkj = B[k, j]
            for i = 1:n
                C[i, j] = muladd(A[i, k], bkj, C[i, j])
            end
        end
    end
    C
end

"""
    find_duplicates(certificates)

Find all duplicate solutions. A solution is considered to be duplicate if the unique
region overlap.
"""
function find_duplicates(certificates::AbstractVector{SolutionCertificate})
    # The strategy to indentify duplication candidates is as following.
    # We first compute the squared euclidean distance to a random point
    # in interval arithmetic. We then compute all overlapping intervals and return this
    # list of tuples.
    # The complexity of this is O(n log(n)) (we need to perform an array sort)

    # 1 Compute squared euclidean distance to a random point
    original_indices = Int[]

    n = length(solution_candidate(first(certificates)))
    pt = randn(ComplexF64, n)

    sq_eucl_distances = Interval{Float64}[]

    a, b = Arblib.Arf(prec = 53), Arblib.Arf(prec = 53)
    for (i, cert) in enumerate(certificates)
        if is_certified(cert)
            push!(original_indices, i)

            d = zero(Interval{Float64})
            for i = 1:n
                yᵢ = IComplexF64(cert.I′[i], a, b)
                d +=
                    IntervalArithmetic.sqr(real(yᵢ) - real(pt[i])) +
                    IntervalArithmetic.sqr(imag(yᵢ) - imag(pt[i]))
            end
            push!(sq_eucl_distances, d)
        end
    end

    # 2) Find all overlapping intervals

    # strategy:
    # Look at close and open of an interval separately and merge all these in one large vector
    # (keeping track if open or close of an interval),
    # sort this (first by value, then open, then close).
    # Now iterate through sorted array and keep track of open intervals

    # value, index, and false if opening and true if closing
    interval_bounds = Tuple{Float64,Int,Bool}[]
    for (i, d) in enumerate(sq_eucl_distances)
        push!(interval_bounds, (d.lo, i, false), (d.hi, i, true))
    end

    sort!(interval_bounds; lt = (a, b) -> a[1] != b[1] ? a[1] < b[1] : !a[3])

    current_open = Int[]
    # dict to help group duplicates:
    # Assume (1,3) and (3,5) overlap then we want to have entries 1=>1, 3=>1 and 5=>1
    duplicate_grouping = Dict{Int,Int}()
    # final duplicates. From our example: 1 => [1,3,5]
    duplicates_dict = Dict{Int,Vector{Int}}()

    for (_v, index, is_closing) in interval_bounds
        if is_closing
            if index == current_open[1]
                popfirst!(current_open)
            else
                for (k, idx) in enumerate(current_open)
                    if idx == index
                        deleteat!(current_open, k)
                        break
                    end
                end
            end
        else
            if !isempty(current_open)
                for idx in current_open
                    # 3) Interval (idx, index) is overlapping. Check for duplicate
                    check_duplicate_candidate!(
                        duplicate_grouping,
                        duplicates_dict,
                        original_indices[idx],
                        original_indices[index],
                        certificates,
                    )
                end
            end
            push!(current_open, index)
        end
    end

    isempty(duplicates_dict) ? Vector{Int}[] : collect(values(duplicates_dict))
end

function check_duplicate_candidate!(duplicate_grouping, duplicates_dict, i, j, certificates)
    i_is_dupl = haskey(duplicate_grouping, i)
    j_is_dupl = haskey(duplicate_grouping, j)
    #If i and j are already in a duplicate cluster then there is no need to
    # check again overlapping by transitivity of our clustering
    i_is_dupl && j_is_dupl && return true
    if Bool(Arblib.overlaps(certificates[i].I′, certificates[j].I′))
        if i_is_dupl
            duplicate_grouping[j] = duplicate_grouping[i]
            push!(duplicates_dict[duplicate_grouping[i]], j)
        elseif j_is_dupl
            duplicate_grouping[i] = duplicate_grouping[j]
            push!(duplicates_dict[duplicate_grouping[j]], i)
        else
            duplicate_grouping[i] = i
            duplicate_grouping[j] = i
            duplicates_dict[i] = [i, j]
        end
        true
    else
        false
    end
end

function certify(
    F::AbstractVector{Expression},
    X,
    p = nothing;
    parameters = Variable[],
    variables = setdiff(variables(F), parameters),
    variable_ordering = variables,
    variable_groups = nothing,
    kwargs...,
)
    sys = System(
        F,
        variables = variable_ordering,
        parameters = parameters,
        variable_groups = variable_groups,
    )
    certify(sys, X, p; kwargs...)
end
function certify(
    F::AbstractVector{<:MP.AbstractPolynomial},
    X,
    p = nothing;
    parameters = similar(MP.variables(F), 0),
    variables = setdiff(MP.variables(F), parameters),
    variable_ordering = variables,
    variable_groups = nothing,
    target_parameters = nothing,
    kwargs...,
)
    if !isnothing(p)
        target_parameters = p
    end
    # as in solver_startsolutions():
    # handle special case that we have no parameters
    # to shift the coefficients of the polynomials to the parameters
    # this was the behaviour of HC.jl v1
    if isnothing(target_parameters) && isempty(parameters)
        sys, target_parameters = ModelKit.system_with_coefficents_as_params(
            F,
            variables = variable_ordering,
            variable_groups = variable_groups,
        )
    else
        sys = System(
            F,
            variables = variable_ordering,
            parameters = parameters,
            variable_groups = variable_groups,
        )
    end
    certify(sys, X, p; target_parameters = target_parameters, kwargs...)
end

function certify(F::System, args...; compile::Union{Bool,Symbol} = false, kwargs...)
    certify(fixed(F; compile = compile), args...; kwargs...)
end

function certify(
    F::AbstractSystem,
    X::MonodromyResult,
    cache::CertificationCache = CertificationCache(F);
    max_precision::Int = 256,
    kwargs...,
)
    _certify(
        F,
        solutions(X),
        certification_parameters(parameters(X)),
        cache;
        max_precision = max_precision,
        kwargs...,
    )
end

function certify(
    F::AbstractSystem,
    r::PathResult,
    p::Union{Nothing,AbstractArray} = nothing,
    cache::CertificationCache = CertificationCache(F);
    target_parameters = nothing,
    max_precision::Int = 256,
    kwargs...,
)
    cert_params =
        certification_parameters(isnothing(p) ? target_parameters : p; prec = max_precision)
    _certify(F, [solution(r)], cert_params, cache; max_precision = max_precision, kwargs...)
end

function certify(
    F::AbstractSystem,
    r::AbstractVector{PathResult},
    p::Union{Nothing,AbstractArray} = nothing,
    cache::CertificationCache = CertificationCache(F);
    target_parameters = nothing,
    max_precision::Int = 256,
    kwargs...,
)
    cert_params =
        certification_parameters(isnothing(p) ? target_parameters : p; prec = max_precision)
    _certify(F, solution.(r), cert_params, cache; max_precision = max_precision, kwargs...)
end

function certify(
    F::AbstractSystem,
    x::AbstractVector{<:Number},
    p::Union{Nothing,AbstractArray} = nothing,
    cache::CertificationCache = CertificationCache(F);
    target_parameters = nothing,
    max_precision::Int = 256,
    kwargs...,
)
    cert_params =
        certification_parameters(isnothing(p) ? target_parameters : p; prec = max_precision)
    _certify(F, [x], cert_params, cache; max_precision = max_precision, kwargs...)
end

function certify(
    F::AbstractSystem,
    X::Result,
    p::Union{Nothing,AbstractArray} = nothing,
    cache::CertificationCache = CertificationCache(F);
    target_parameters = nothing,
    max_precision::Int = 256,
    kwargs...,
)
    params = isnothing(p) ? target_parameters : p
    cert_params = certification_parameters(params; prec = max_precision)
    _certify(
        F,
        solutions(X; only_nonsingular = true),
        cert_params,
        cache;
        max_precision = max_precision,
        kwargs...,
    )
end

function certify(
    F::AbstractSystem,
    X,
    p::Union{Nothing,AbstractArray} = nothing,
    cache::CertificationCache = CertificationCache(F);
    target_parameters = nothing,
    max_precision::Int = 256,
    kwargs...,
)
    cert_params =
        certification_parameters(isnothing(p) ? target_parameters : p; prec = max_precision)
    _certify(F, X, cert_params, cache; max_precision = max_precision, kwargs...)
end


function  Bistochastic_Normalization!(X::Array{Float64,2},iterations::Int)

    n,m = size(X)

    for _ in 1:iterations
        for i = 1:n
            X[i,:] ./= norm(X[i,:],1)
        end

        for j = 1:m
            X[:,j] ./= norm(X[:,j],1)
        end
    end

end

function Bistochastic_Normalization!(x::Array{Float64,1},iterations::Int,n::Int,m::Int)

    X = reshape(x,n,m)

    for _ in 1:iterations
        for i = 1:n
            X[i,:] ./= norm(X[i,:],1)
        end

        for j = 1:m
            X[:,j] ./= norm(X[:,j],1)
        end
    end

end

function build_assignment(X::Array{Float64,2};use_colsort::Bool=false,
                          return_matrix=false)
    """-------------------------------------------------------------------------
        For a HOFASM_iteration matrix result X, computing a greedy matching on
      the maximum elements of the columns sequentially. if using the option
      use_colsort, will find the matching in the columns ordered by the largest
      element of each column. If return_matrix is true, a matrix is returned of
      the same dimensions as X, but X(i,j) =  1 iff i is matched to j.

      Notes:
      ------
      * Currently implemented for square matrices
    -------------------------------------------------------------------------"""
    n, m = size(X)
    selected_vertices = Set{Int}()
    match_size = minimum((n,m))
    matching = Array{Tuple{Int,Int},1}(undef,match_size)
    index = 1

    if use_colsort
        col_order = sortperm([maximum(X[:,i]) for i in 1:m],rev=true)
    else
        col_order = 1:match_size
    end

    for j in col_order #search over columns
        max_val = -Inf
        arg_max = -1
        for i = 1:n
            if (X[i,j] > max_val) && !(i in selected_vertices)
                max_val = X[i,j]
                arg_max = i
            end
        end
        matching[index] = (j,arg_max)
        index += 1
        push!(selected_vertices,arg_max)
        if index > match_size
            break
        end
    end
    if return_matrix
        Z = zeros(n,m)
        for (i,j) in matching
            Z[i,j] = 1
        end
        return Z
    else
        return matching
    end
end

function HOFASM_iterations(Bn_indices_list::Array{Array{Int,2},1},Bn_vals_list::Array{Array{Float64,1},1},
                           Hn_indices_list::Array{Array{Int,2},1},n::Int,m::Int,max_iterations::Int=12)

    x_k::Array{Float64,1} = rand(Float64,n*m)
    x_k ./= norm(x_k)
    beta = 30 # constant from the paper

    i = 1
    x_k_1::Array{Float64,1} = zeros(n*m)

    while true

        implicit_kronecker_model_marginalization_contraction!(Bn_indices_list,Bn_vals_list,
                                                              Hn_indices_list,x_k,x_k_1,n,m)

        max_elem = maximum(x_k_1)
        inflation_normalization = x -> exp(x*beta/max_elem)
        broadcast!(inflation_normalization,x_k_1,x_k_1)

        while true

            x_normalized = copy(x_k_1)
            Bistochastic_Normalization!(x_normalized,15,n,m)
  #          print(f"{norm(X_normalized - X_k_1,'fro')}")
            if norm(x_normalized - x_k_1) < 1e-4
                x_k_1 = x_normalized
                break
            else
                x_k_1 = copy(x_normalized)
            end
        end
#        return X_k_1

        x_k_1 ./= norm(x_k_1)
        res = norm(x_k_1 - x_k) / norm(x_k)

        #print("Finished iteration $i with difference $res")
        i += 1
        if i > max_iterations
            return x_k_1
        end

        if res < 1e-4
            return x_k_1
        else
            x_k = copy(x_k_1)
        end
    end
end

function HOFASM_iterations(matrices::Array{SparseMatrixCSC{Float64,Int64},1},n::Int,m::Int,max_iterations::Int=12)

    #    X_k = np.ones((n,m))
    x_k::Array{Float64,1} = rand(Float64,n*m)
    x_k ./= norm(x_k)
    beta = 30 # constant from the paper

    i = 1

    while true

        x_k_1::Array{Float64,1} = zeros(n*m)
        for Tn in matrices
            x_k_1 .+= Tn*x_k
        end

        max_elem = maximum(x_k_1)
        inflation_normalization = x -> exp(x*beta/max_elem)
        broadcast!(inflation_normalization,x_k_1,x_k_1)

        while true

            x_normalized = copy(x_k_1)
            Bistochastic_Normalization!(x_normalized,15,n,m)
  #          print(f"{norm(X_normalized - X_k_1,'fro')}")
            if norm(x_normalized - x_k_1) < 1e-4
                x_k_1 = x_normalized
                break
            else
                x_k_1 = copy(x_normalized)
            end
        end
#        return X_k_1

        x_k_1 ./= norm(x_k_1)
        res = norm(x_k_1 - x_k) / norm(x_k)

        #print("Finished iteration $i with difference $res")
        i += 1
        if i > max_iterations
            return x_k_1
        end

        if res < 1e-4
            println("tolerance met")
            return x_k_1
        else
            x_k = copy(x_k_1)
        end
    end
end

function HOFASM_iterations(tensor_pairs::Array{NTuple{2,SparseMatrixCSC{Float64,Int64}},1},n::Int,
                           m::Int,max_iterations::Int=12,tol=1e-4)

    #    X_k = np.ones((n,m))
    X_k::Array{Float64,2} = rand(Float64,n,m)
    X_k ./= norm(X_k)
    beta = 30 # constant from the paper

    X_k_1 = Array{Float64,2}(undef,n,m)
    i = 1

    preprocessed_pairs = [(findnz(H)...,B) for (H,B) in tensor_pairs]

    while true

#        HOFASM_contraction!_test(tensor_pairs,X_k,X_k_1)
        HOFASM_contraction!(preprocessed_pairs,X_k,X_k_1)

 #       SIMD_sparse_matrix_multiplication(Hn_matrices,Bn_matrices,X_k,X_k_1)
#        packed_X_k = pack(X_k')
 #       for (Hn,Bn) in tensor_pairs

#            X_k_1 .+= Hn*(Bn*X_k')'
#            X_k_1 .+= (Bn*(Hn*X_k)' )'

#            temp = unpack(Bn*packed_X_k)'
#            X_k_1 .+= unpack(Hn*pack(temp))
           #println(typeof(Hn*(Bn*X_k')'))
           # X_k_1 .+= Hn*(X_k*Bn')
  #      end

 #       println("done")

        max_elem = maximum(X_k_1)
        inflation_normalization = x -> exp(x*beta/max_elem)
        broadcast!(inflation_normalization,X_k_1,X_k_1)

        while true

            X_normalized = copy(X_k_1)
            Bistochastic_Normalization!(X_normalized,15)
  #          println("BN")
   #         print("{norm(X_normalized - X_k_1,'fro')}")
            if norm(X_normalized - X_k_1) < 1e-4
                X_k_1 = X_normalized
                break
            else
                X_k_1 = copy(X_normalized)
            end
        end
#        return X_k_1

        X_k_1 ./= norm(X_k_1)
        res = norm(X_k_1 - X_k) / norm(X_k)

  #      println("Finished iteration $i with difference $res")
        i += 1
        if i > max_iterations
            return X_k_1
        end

        if res < tol
            return X_k_1
        else
            X_k = copy(X_k_1)
        end
    end
end

#TODO: could generalize to take a contraction function to unify all graduated assignment routines
function HOM_graduated_assignment(marg_tensor::SparseMatrixCSC{Float64,Int64},m::Int=-1,max_iterations::Int=12,tol::Float64=1e-4)
    println(m)
    println(max_iterations)
    #    X_k = np.ones((n,m))
    n = size(marg_tensor,1)
    if m == -1
        m = Int(sqrt(n))
    end
    #m = Int(sqrt(n))
    x_k::Array{Float64,1} = rand(Float64,n)
    x_k ./= norm(x_k)
    beta = 30 # constant from the paper
    i = 1

    while true

        x_k_1 = marg_tensor*x_k

        max_elem = maximum(x_k_1)
        inflation_normalization = x -> exp(x*beta/max_elem)
        broadcast!(inflation_normalization,x_k_1,x_k_1)

        while true

            x_normalized = copy(x_k_1)
            Bistochastic_Normalization!(x_normalized,10,Int(n/m),m)
  #          print(f"{norm(X_normalized - X_k_1,'fro')}")
            if norm(x_normalized - x_k_1) < 1e-4
                x_k_1 = x_normalized
                break
            else
                x_k_1 = copy(x_normalized)
            end
        end
#        return X_k_1

        x_k_1 ./= norm(x_k_1)
        res = norm(x_k_1 - x_k) / norm(x_k)

        #print("Finished iteration $i with difference $res")
        i += 1
        if i > max_iterations
            return x_k_1
        end

        if res < 1e-4
            println("tolerance met")
            return x_k_1
        else
            x_k = copy(x_k_1)
        end
    end

end

#Sequential Second-order expansion for higher order matching
function SSoEfHOM(tensor_pairs::Array{NTuple{2,SparseMatrixCSC{Float64,Int64}},1},n::Int,
                  m::Int, max_iterations::Int=12,tol::Float64=1e-4)

    X_k::Array{Float64,2} = rand(Float64,n,m)
    beta = 30 # constant from the paper

    X_k_1 = Array{Float64,2}(undef,n,m)
    B_k_1 = Array{Float64,2}(undef,n,m)
    temp2 = Array{Float64,2}(undef,n,m)
    i = 1

    preprocessed_pairs = [(findnz(H)...,B) for (H,B) in tensor_pairs]

    HOFASM_contraction!(preprocessed_pairs,X_k,X_k_1)
    S_opt = dot(X_k,X_k_1)
    x_opt = build_assignment(rand(n,m),return_matrix=true)

    while true

        #Graduated assignment routine
        max_elem = maximum(X_k_1)
        inflation_normalization = x -> exp(x*beta/max_elem)
        broadcast!(inflation_normalization,B_k_1,X_k_1)

        while true

            X_normalized = copy(B_k_1)
            Bistochastic_Normalization!(X_normalized,15)

            if norm(X_normalized - B_k_1) < 1e-6
                B_k_1 = X_normalized
                break
            else
                B_k_1 = copy(X_normalized)
            end
        end

        B_k_1 = build_assignment(B_k_1,return_matrix=true)

        diff = B_k_1-X_k_1
        println(norm(diff))
        C = dot(X_k_1,diff)
    #    println(C)

        HOFASM_contraction!(preprocessed_pairs,diff,temp2)
        D = dot(diff,temp2)
#        println(D)

        if D >= 0
            X_k = copy(B_k_1)
        else
            r = minimum((-C/D,1))
            X_k .+= r * diff
        end

        temp3 = copy(B_k_1)
        HOFASM_contraction!(preprocessed_pairs,B_k_1,temp3)
        S = dot(temp3,B_k_1)
        println(S)

        if S >= S_opt
            x_opt = copy(B_k_1)
            S_opt = S
        end

  #      println(X_k)

        i += 1
        if i >= max_iterations
            break
        else
            #tensor vector contraction
            HOFASM_contraction!(preprocessed_pairs,X_k,X_k_1)
        end
    end

    return x_opt
end

#=
#TODO: could generalize to take a contraction function to unify all graduated assignment routines
function HOM_graduated_assignment(A::ssten.COOTen,max_iterations::Int=12,tol::Float64=1e-4)

    #    X_k = np.ones((n,m))
    n = A.cubical_dimension
    m = Int(sqrt(n))
    x_k::Array{Float64,1} = rand(Float64,n)
    x_k ./= norm(x_k)
    beta = 30 # constant from the paper
    i = 1

    while true

        x_k_1 = ssten.contract_k_1(A,x_k)

        max_elem = maximum(x_k_1)
        inflation_normalization = x -> exp(x*beta/max_elem)
        broadcast!(inflation_normalization,x_k_1,x_k_1)

        while true

            x_normalized = copy(x_k_1)
            Bistochastic_Normalization!(x_normalized,10,m,m)
  #          print(f"{norm(X_normalized - X_k_1,'fro')}")
            println(norm(x_normalized - x_k_1))
            if norm(x_normalized - x_k_1) < 1e-4
                x_k_1 = x_normalized
                break
            else
                x_k_1 = copy(x_normalized)
            end
        end
#        return X_k_1

        x_k_1 ./= norm(x_k_1)
        res = norm(x_k_1 - x_k) / norm(x_k)

        #print("Finished iteration $i with difference $res")
        i += 1
        if i > max_iterations
            return x_k_1
        end

        if res < 1e-4
            println("tolerance met")
            return x_k_1
        else
            x_k = copy(x_k_1)
        end
    end

end
=#


using Distributions
using OmegaCore
using Test

function testsolution()
  μ = 1 ~ Normal(0, 1)
  y = 2 ~ ((id, ω) -> Normal(μ(ω), 1)(id, ω))
  μc = μ |ᶜ (y ==ₚ 5.0)
  @test_throws OmegaCore.ConditionException μc(defω())
  sol = solution(μc)
end

function testlogpdfsol()
  μ = 1 ~ Normal(0, 1)
  y = 2 ~ ((id, ω) -> Normal(μ(ω), 1)(id, ω))
  μc = μ |ᶜ (y ==ₚ 5.0)
  ω = solution(μc)
  μ_ = μ(ω)
  @test logpdf(ω) == logpdf(Normal(0, 1), μ_) + logpdf(Normal(μ_, 1), 5.0)
end

function test_solution_intervene()
  μ = 1 ~ Normal(0, 1)
  y = 2 ~ ((id, ω) -> Normal(μ(ω), 1)(id, ω))
  μc = μ |ᶜ (y ==ₚ 5.0)
  yi = y |ᵈ (μ => (ω -> 100.0))
  yi2 = y |ᵈ (μ => (ω -> 100))
  ω = solution(μc)
  a, b = (yi(ω), y(ω))
  @test a != b
end

@testset "Solution" begin
  testsolution()
  testlogpdfsol()
  test_solution_intervene()
end


"Log( data::SingleCellExpr )

Log the expression data using log(x+1), not log2 or log10
This will not touch the 0 or -1 values

"

function Log( data::SingleCellExpr )
  function logthis(x::SparseVector)
    ok = find( x .> 0 );
    x[ok] .= log(x[ok]+1);
    return x
  end
  rws = 1:size(data.data)[1]
  for  i =  1:size(data.data)[2]
    data.data[rws,i] = logthis( data.data[rws,i] )
  end
  return data
end


mutable struct FullCovMatrix{T<:AbstractFloat, K <: MatTypes} <: CovMatrix
    C::K
    c_1::T
    c_μ::T
end

function FullCovMatrix(::Type{K}, n, μ_w) where K <: MatTypes
    c_1 = 2/((n+1.3)^2 + μ_w)
    c_μ = min(1-c_1,2*(μ_w-2+1/μ_w)/((n+2)^2 +μ_w))
    C = init_matrix(K,n)

    FullCovMatrix(C,c_1,c_μ)
end

cov_matrix(cov_mat::FullCovMatrix) = cov_mat.C
eigvals(cov_mat::FullCovMatrix) = eigen(cov_mat.C).values #FIXME maybe a bit costly

mv_rand(n,C) = rand(MultivariateNormal(zeros(n),C))
mv_rand(n,C::Diagonal) = rand(MultivariateNormal(zeros(n),diag(C)))

function update_p_σ(::Type{FullCovMatrix{T,K}},c_σ, p_σ, μ_w, C, Δm, σ) where {T,K}
    (1-c_σ)*p_σ  + √(c_σ*(2-c_σ)*μ_w) * C^(-1/2) * Δm/σ
end

function update_p_c(::Type{FullCovMatrix{T,K}},c_c, p_c, H_σ, c_σ, μ_w, Δm, σ) where {T,K}

#    @show (H_σ*√(c_σ*(2 - c_σ)*μ_w), Δm/σ)
#    (1-c_c/3)*p_c + H_σ*√(c_σ*(2 - c_σ)*μ_w) * 2.5Δm/σ
    (1-c_c)*p_c + H_σ*√(c_σ*(2 - c_σ)*μ_w) * Δm/σ
end

function update_C(::Type{FullCovMatrix{T,K}}, C::K, c_1, c_μ, H_σ, c_c, p_c, w, σ, x, m, μ) where {T,K}

#    c_1 = 0.75*c_1
#    c_µ = 2*c_µ

    C = (1-c_1-c_µ + (1-H_σ)*c_1*c_c*(2-c_c)) * C +
        c_1 * p_c * p_c' +
        c_μ * ∑( w[i]/σ^2*( (x[i]-m)*(x[i]-m)') for i=1:μ)

#    @show (1-c_1-c_µ + (1-H_σ)*c_1*c_c*(2-c_c)) * C
#    @show c_1 * p_c * p_c'
#    @show c_μ * ∑( w[i]/σ^2*( (x[i]-m)*(x[i]-m)') for i=1:μ)
#    @show p_c

    C = (C+C')/2 #keep symmetric part
    C
end

#this is probably wrong
function update_C(::Type{FullCovMatrix{T,K}}, C::K, c_1, c_μ, H_σ, c_c, p_c, w, σ, x, m, μ) where {T,K <: Diagonal}

    C = (1-c_1-c_µ + (1-H_σ)*c_1*c_c*(2-c_c)) * C +
        c_1 * Diagonal(p_c.^2) +
        c_μ * ∑( w[i]/σ^2*( Diagonal((x[i]-m).^2) ) for i=1:μ)
    C
end

function update_state!(d, state::CMAState{T,CM}, method::CMA) where {T,CM <: FullCovMatrix}

    @unpack w,n,μ_w,c_σ,d_σ,c_c,p_σ,p_c,cov_mat,m,chi_D,σ,t,xs = state
    @unpack C,c_1,c_μ = cov_mat

    state.f_x_previous = state.f_x
    copy!(state.x_previous, state.x)

    λ = method.λ
    μ = floor(Int,λ/2)

    x = [m + σ*mv_rand(n,C) for i=1:λ]

    fx = eval_objfun(T, λ, d, x)

    idx = sortperm(fx)
    x, fx = x[idx], fx[idx]

    m_t = m
    m = ∑( w[i] * x[i] for i=1:μ)

    Δm = m-m_t

#    @show C

    p_σ = update_p_σ(CM, c_σ, p_σ, μ_w, C, Δm, σ)
    H_σ = update_H_σ(CM, p_σ, c_σ, t, n)
    p_c = update_p_c(CM, c_c, p_c, H_σ, c_σ, μ_w, Δm, σ)
    C = update_C(CM, C, c_1, c_μ, H_σ, c_c, p_c, w, σ, x, m, μ)

    σ = update_σ(CM, σ, c_σ, d_σ, p_σ, chi_D)

    D = eigvals(cov_mat)

#    @printf("σ: %2.2f H_σ: %2.2f axis-ratio: %2.2e \n",
#            σ, H_σ, maximum(D) / minimum(D) )

    t += 1
    state.x = x[1]
    state.f_x = fx[1]
    xs = x

    @pack! cov_mat = C,c_1,c_μ
    @pack! state = w,n,μ_w,c_σ,d_σ,c_c,p_σ,p_c,cov_mat,m,chi_D,σ,t,xs

    false # should the procedure force quit?
end


# (coeff, featurenum, lag, windowlen) for each factor and possibly an intercept.
"""
A moving average linear regression model for time series prediction. This is
similar to but more general than a lagged factor model since the factors can
be moving averages over a range of lags rather than simply a single lag.
In theory, this could "smoothen out" noise in the dependence from one feature
to another while still allowing non-moving average (lagged factors) by setting
the length of the moving window to 1. The `minlag` and `maxlag` limits
how far back in time we can look.
"""
mutable struct MovingAverageLITRM <: AbstractLinearITRM
    nfactors::Int
    nfeatures::Int
    minlag::Int
    maxlag::Int
    maxcoef::Float64
    mincoef::Float64
    hasintercept::Bool
    floorcoefs::Bool
    ndigits::Int # When printing
    params::Union{Nothing,Vector{Float64}}
    featurenames::Vector{String}
    MovingAverageLITRM(nfactors::Int, nfeatures::Int;
        minlag::Int = 3, maxlag::Int = 10,
        maxcoef = 1e3, mincoef = -maxcoef,
        hasintercept = true, floorcoefs = false, ndigits = 2,
        params = nothing,
        featurenames = String[]) = begin
        new(nfactors, nfeatures, minlag, maxlag, maxcoef, mincoef,
            hasintercept, floorcoefs, ndigits, params, featurenames)
    end
end

function parambounds(m::MovingAverageLITRM)
    bounds = Tuple{Float64, Float64}[]
    for _ in 1:nfactors(m)
        # Order of params: coef, feature, startlag, windowlen
        push!(bounds, (m.mincoef, m.maxcoef))
        # Since we will floor(Int, v) the following we add almost 1.0 to:
        push!(bounds, (1.0, float(m.nfeatures) + AlmostRoundUpTo1))
        push!(bounds, (float(m.minlag), float(m.maxlag) + AlmostRoundUpTo1))
        # Window length 1.0 means it is just a lagged factor with no
        # moving average window.
        push!(bounds, (1.0, float(m.maxlag) - float(m.minlag) + 1 + AlmostRoundUpTo1))
    end
    hasintercept(m) && push!(bounds, (m.mincoef, m.maxcoef))
    bounds
end

function isvalidparams(m::MovingAverageLITRM, ps::Vector{Float64})
    spans = windowspans(m, ps)
    (minlag(m) <= minimum(map(first, spans))) &&
    (maximum(map(last, spans)) <= maxlag(m)) &&
    all(t -> 1 <= t[2] <= (maxlag(m) - minlag(m) + 1),  spans)
end

windowspans(m::MovingAverageLITRM, ps::Vector{Float64} = params(m)) =
    Tuple{Int, Int, Int}[getfactor(m, fi, ps)[[3,4,5]] for fi in 1:nfactors(m)]

startofwindows(m::MovingAverageLITRM, ps::Vector{Float64} = params(m)) =
    map(first, windowspans(m, ps))

windowlengths(m::MovingAverageLITRM, ps::Vector{Float64} = params(m)) =
    map(t -> t[2], windowspans(m, ps))

endofwindows(m::MovingAverageLITRM, ps::Vector{Float64} = params(m)) =
    map(last, windowspans(m, ps))

# If we have params the smallest (since they are positive in the params vector)
# lag value is our horizon. If we don't have params we use the
actualminlag(m::MovingAverageLITRM, ps::Vector{Float64} = params(m)) =
    hasparams(m) ? minimum(startofwindows(m, ps)) : minlag(m)

actualmaxlag(m::MovingAverageLITRM, ps::Vector{Float64} = params(m)) =
        hasparams(m) ? maximum(endofwindows(m, ps)) : maxlag(m)

horizon(m::MovingAverageLITRM) = actualminlag(m)

function predictforrow(m::MovingAverageLITRM, X::AbstractMatrix{N}, row::Int,
    ps::Vector{Float64} = params(m)) where {N<:Number}
    prediction = hasintercept(m) ? ps[end] : 0.0
    for fi in 1:nfactors(m)
        coef, featureidx, winstart, winlen, wend = getfactor(m, fi, ps)
        #@show (row, winstart, winlen, wend, actualmaxlag(m), maxlag(m), ps)
        if winlen > 1
            windowavg = mean(X[(row-wend):(row-winstart), featureidx])
        else
            windowavg = X[row-winstart, featureidx]
        end
        prediction += (coef * windowavg)
    end
    return prediction
end

function windowfromparams(m::MovingAverageLITRM, ws::N1, wl::N2) where {N1<:Number,N2<:Number}
    wstart = floor(Int, ws)
    wend = min(maxlag(m), wstart + floor(Int, wl) - 1)
    wlen = wend - wstart + 1
    return wstart, wlen, wend
end

function getfactor(m::MovingAverageLITRM, fi::Int, ps::Vector{Float64} = params(m))
    i = 1 + (fi-1)*4
    c, f, ws, wl = ps[i:(i+3)]
    coef = m.floorcoefs ? floor(Int, c) : c
    return coef, floor(Int, f), windowfromparams(m, ws, wl)...
end

modelname(m::MovingAverageLITRM) =
    "MovingAverage factor LITRM (Linear Interpretable Timeseries Regression Model)"

function writefactor(io::IO, m::MovingAverageLITRM, fi::Int)
    c, f, ws, wlen, wend = getfactor(m, fi)
    coef = r(m, c)
    if coef != 0
        op = (coef > 0.0) ? "+" : "-"
        write(io, " $op $(abs(coef)) * ")
        f = "$(featurename(m, f))[T-$(ws)]"
        s = (wlen > 1) ? "MAvg($(f), $(wlen))" : f
        write(io, s)
    end
end


#===============
    STRUCTS
===============#

mutable struct sphere <: RayTracer.RTObject
    id::UUID
    radius::Float64
    position::rt_tuple
    transform::Array{Float64,2}
    material::RayTracer.RTMaterial
    function sphere()
        return new(uuid4(),1.0,RayTracer.point(0.,0.,0.),Matrix(I,4,4),RayTracer.material())
    end
end

#===============
    FUNCTIONS
===============#

function intersect(s::sphere,r::ray)::Array{intersection,1}
    r2 = transform(r,inv(s.transform))
    x0 =  r2.origin - s.position
    a = RayTracer.dot(r2.direction,r2.direction)
    b = 2*RayTracer.dot(r2.direction,x0)
    c = RayTracer.dot(x0,x0) - s.radius^2

    discriminant = b^2 - 4*a*c
    if discriminant<0
        return []
    end
    t1 = -(b+sqrt(discriminant))/(2*a)
    t2 = -(b-sqrt(discriminant))/(2*a)
    return [intersection(t1,s),intersection(t2,s)]
end

function normal_at(s::sphere,p::RayTracer.rt_tuple)::RayTracer.rt_tuple
    object_normal = inv(s.transform)*p - s.position
    return_transform = zeros(4,4)
    return_transform[1:3,1:3] = inv(RayTracer.submatrix(s.transform,4,4))'
    return_transform[4,4] = 1
    return RayTracer.normalize(return_transform*object_normal)
end


"Distance between two values"
function dist end

# zt: these choices seem a little adhoc.  Should I specialise by type

@inline dist(x::Real, y::Real) = (xy = (x - y); xy * xy)
# @inline dist(x::Vector{<:Real}, y::Vector{<:Real}) = norm(x - y)
@inline dist(x::AbstractVector{<:Real}, y::AbstractVector{<:Real}) = sum(dist.(x, y))
@inline dist(x::NTuple{N,<: Real}, y::NTuple{N,<:Real}) where N = sum(dist.(x, y))
@inline dist(x::AbstractArray{<:Real}, y::AbstractArray{<:Real}) = norm(x[:] - y[:])

"Distance from x to [a, b]"
function bound_loss(x, a, b)
  # @pre b >= a
  if x < a
    a - x
  elseif x > b
    x - b
  else
    zero(x)
  end
end


using SimpleBirthDeathProcess;
using DelimitedFiles;
using Plots;

N = 200;
F = Float64;
I = 200;
J = 100;
T = 1;
L = F(10);
M = F(10);

# values obtained with Maple 2018.2
correct_values = readdlm("log_trans_prob_values_bivariate.csv", ',', BigFloat);
correct_values = reshape(correct_values, N, :);

λset = [L * F(n) / N for n = 1:N];
μset = [M * F(n) / N for n = 1:N];
logp = map((λ, μ) -> trans_prob(I, J, T, [λ; μ]),
           repeat(reshape(λset, 1, :), N, 1),
           repeat(μset, 1, N));

relative_error = abs.(1 .- logp ./ correct_values);

colgr = ColorGradient([colorant"#ffffff", colorant"#525252"], [0.0, 1.0]);
pl = plot(λset, μset, log10.(relative_error), seriestype=:contour,
          levels=[-20; -16; -14; -13; -12], fill=true, contour_labels=false,
          legend=true, linecolour=:black, aspect_ratio=:equal, fillcolour=colgr,
          xlim=(0, 10), ylim=(0, 10), xlab="\\lambda", ylab="\\mu");

savefig(pl, "Figure_4.svg")


import Distances
include("structs.jl")

function stepped_dynamics(obj::Contraption, settings::PhysicsSettings, Δt::Float64)
    # Name binding for brevity
    bounds = settings.bounds
    n = length(obj.mass)

    # Storing kinematics
    s = deepcopy(obj.position)
    v = deepcopy(obj.velocity)
    a = zeros(2, n)

    # Compute acceleration from forces
    for i in 1:n
        # Gravity
        a[:, i] += [0., -settings.g]

        # Drag
        a[:, i] -= settings.drag * v[:, i]

        # Springs
        # TODO: should connections be memoized?
        for j in 1:n # Don't double-count!
            if obj.springs[i, j] != 0.
                Δs = s[:, j] - s[:, i]
                current_length = Distances.norm(Δs)
                direction = Δs / current_length

                a[:, i] += obj.springs[i,j] / obj.mass[i] *
                     (current_length - obj.rest_length[i,j]) *
                     direction # unit vector of force direction
            end
        end

    end

    # Applying acceleration
    for i in 1:n
        # Updating velocity from forces
        v[:, i] += a[:, i] * Δt

        # Continually attempt to apply velocity until result is in-bounds
        t_remaining = Δt
        while t_remaining > 0.

            # Movement
            proposed_position = s[:, i] + v[:, i] * Δt

            # Collision handling
            if !(bounds.x[1] <= proposed_position[1] <= bounds.x[2]) |
               !(bounds.y[1] <= proposed_position[2] <= bounds.y[2])

                # Determine whether horizontal or vertical collision is first
                # Catch 0 velocity cases
                if v[1, i] == 0.
                    x_or_y = :y

                    Δt_bottom = v[2, i] / (s[2, i] - bounds.y[1])
                    Δt_top    = v[2, i] / (s[2, i] - bounds.y[2])

                    Δt_c = max(Δt_bottom, Δt_top)
                elseif v[2, i] == 0.
                    x_or_y = :x

                    Δt_left  = v[1, i] / (s[1, i] - bounds.x[1])
                    Δt_right = v[1, i] / (s[1, i] - bounds.x[2])

                    Δt_c = max(Δt_left, Δt_right)
                else
                    Δt_left   = v[1, i] / (s[1, i] - bounds.x[1])
                    Δt_right  = v[1, i] / (s[1, i] - bounds.x[2])
                    Δt_bottom = v[2, i] / (s[2, i] - bounds.y[1])
                    Δt_top    = v[2, i] / (s[2, i] - bounds.y[2])

                    # Only one of top/bottom will be positive
                    # Select that one as the collision which occurs in the future
                    # Then select sooner of collisions that will occur
                    Δt_x = max(Δt_left, Δt_right)
                    Δt_y = max(Δt_top, Δt_bottom)

                    x_or_y = ifelse(Δt_x >= Δt_y, :x, :y)
                    Δt_c = min(Δt_x, Δt_y)
                end

                # Update proposal and time remaining in time step
                # Reverse movement to get to point where collision with wall occurs
                proposed_position -= v[:, i]*(Δt - Δt_c)
                t_remaining -= Δt_c

                # Change directions and lose energy
                if x_or_y == :x
                    v[1, i] *= -settings.elasticity
                else
                    v[2, i] *= -settings.elasticity
                end
            else
                s[:, i] = proposed_position
                t_remaining = 0
            end
        end
    end

    return (s, v)

end


#########################################
###Remove duplicates from a sorted vector
#########################################
function uniquesorted(in::AbstractVector)
  out = similar(in)
  c::Int = 0
  if length(in) > 0
    out[c+=1] = in[1]
    @inbounds for i=2:length(in)
      if in[i]!=in[i-1]
        out[c+=1]=in[i]
      end
    end
  end
  out[1:c]
end

############################################################
###Calculate a log-normal value for given location and scale
############################################################
@inline function randlognormal!{T<:AbstractFloat}(v::Vector{T},μ::T,σ::T)
    #tested different variations and the way this is implemented runs fastest
    #there is about 30% benefit in speed of using Float32 over Float64 for large vectors
    l = length(v)
    @simd for i=1:l
        @inbounds v[i] = T(randn())
    end
    @simd for i=1:l
        @inbounds v[i] = μ + σ*v[i]
    end
    @simd for i=1:l
        @inbounds v[i] = exp(v[i])
    end
    v
end