"""
    module ExactlyRemoveSingularities - exactly remove singularities from A*v = 0 with A an Integer matrix.

Use this module in the following way:

   using ExactlyRemoveSingularities
   result = removeSingularities(A,ix,iy)    # result = (iya, eqr, ix1, ix2, eqx, A1, A2)
   printRemoveSingularities(A,result,names)

where A is a dense or sparse Int matrix A[:,:] describing a linear system

    A*v = 0

with

  - v[ix]: Variables used in the derivative operator, der(v[ix]), so these are the potential states.

  - v[iy]: Variables that must be solved from ```A*v = 0```, because these variables do not appear in the remaining part of the model equations (e.g. absolute potentials). If variables v[iy] cannot be uniquely solved from ```A*v = 0```, then in argument iya the variables are listed that must be set arbitrarily so that v[iy] can be uniquely solved from the equations. For example, if v[iy] contains the absolute potentials and no ground object is present, then v[iya] are the absolute potentials that must be explicitly set, e.g., by introducing an equation v[iya]=0, in order that the equations have a unique solution. Additionally, a warning message should be printed that v[iya] was set to 0.)

The function return arguments are all of type Int and have the following meaning:

  - iya: variables v[iya] can be arbitrarily set (iya is a sub-set of iy)

  - eqr: equations ```A[eqr,:]*v = 0``` are redundant and should be removed (they do not change the result)

  - ix1: x[ix1] are the dependent states (ix1 is a sub-set of ix)

  - ix2: i[ix2] are the independent states (ix2 is a sub-set of ix)

  - eqx, A1, A2: equations ```A[eqx,:]*v = 0``` should be replaced by the equations ```A1*x1 + A2*x2 = 0```, where A1 is square, upper triangular and has full rank with A1[i,i] != 0. In other words, x1 can be computed explicitely from these equations given x2.

The calling program should do the following:

   - Remove equations ```A[eqr,:]*v = 0```, because these are redundant equations.
   - Replace equations ```A[eqx,:]*v = 0``` by ```A1*x1 + A2*x2 = 0```, in order that index reduction can be performed by a structural algorithm (e.g. Pantelides algorithm).
   - Introduce equations v[iya] = 0 and print a warning message that these variables have been arbitrarily set to zero. This is needed, in order that the equations have a unique solution.
"""
module ExactlyRemoveSingularities

export removeSingularities, upperTrapezoidal
export printobj, printRemoveSingularities

"""
    (Afac, rk, p1, p2, Bfac) = upperTrapezoidal(A,B)
    (Afac, rk, p1, p2)       = upperTrapezoidal(A)

A and B are **Int** matrices with dimensions A[n1,n2], B[n1,:] or B[n1], so B can be a matrix or a vector. Furthermore A and B can be dense or sparse matrices (SparseMatrixCSC). If B is not present, it is interpreted
as zero matrix. Typically A has more columns as rows (n2 > n1). The function transforms the equation

    A*X = B

to

    Afac*X[p2,:] = Bfac

with

- Afac an **upper trapezoidal** Int matrix with Afac[i,i] != 0 for i<=rk; Afac[rk+1:end,:] = 0

- Bfac an Int matrix or vector (not present if B is not present)

- rk the rank of A and of Afac

- p1 the row permutation vector of A

- p2 the column permutation vector of A

The transformation is performed by a fraction-free Gauss elimination with
exact (integer) computation. It is then possible to solve for X by splitting
X[p2,:] = [X1;X2] with X1[rk,rk], providing any values for X2 and computing X1 from

    Afac[1:rk,1:rk]*X1 = Bfac[1:rk,:] - Afac[1:rk,rk+1:end]*X2

Since Afac[1:rk,1:rk] is **upper triangular** with non-zeros on the diagonal,
this is just a forward recursion. Note, the (full) pivoting is performed
in such a way that it is tried to utilize the **upper** part of X in X1.
Therefore, variables that must be solved for should be in the upper part of X.

The permutation vector p1 can be used to determine the equations in A and B that
can be removed without changing the result (provided Bfac[rk+1:end,:] = 0):

    p1red = p1[rg+1:end]
    A[p1red,:]*X = B[p1red,:] are redundant equations that can be removed

This implementation is based on:

  - Bareiss E.H. (1968): Sylvester’s identity and multistep integer-preserving Gaussian elimination. Mathematics of Computation, Vol. 22, No. 103, pp. 565 – 578. Download: http://www.ams.org/journals/mcom/1968-22-103/S0025-5718-1968-0226829-0/

  - Turner P.R. (1995): A simplified fraction-free Integer elimination algorithm. Naval Air Warfare Center. Report NAWCADPAX--96-196-TR. Download: http://www.dtic.mil/cgi-bin/GetTRDoc?AD=ADA313755

  - Dureisseix D. (2012): Generalized fraction-free LU factorization for singular systems with kernel extraction. Linear Algebra and its Applications, Volume 436, Issue 1, pp. 27-40. Download: http://www.sciencedirect.com/science/article/pii/S0024379511004617#
"""
function upperTrapezoidal(A::Matrix{Int}, B::Matrix{Int})
   assert(size(A,1) == size(B,1))
   (n1,n2) = size(A)
   nb   = size(B,2)
   p1   = collect(1:n1)
   p2   = collect(1:n2)
   Afac = copy(A)
   Bfac = copy(B)
   rk   = n1
   p1k  = 0
   p2k  = 0
   Aik  = 0
   temp = 0
   tempRow    = fill(0,1,nb)
   pivot      = 1
   oldPivot   = 1
   pivotFound = false

   for k = 1:n1
      # Search for a pivot in Afac[k:n1,min(k,n2):n2]
      if k <= n2
         if Afac[k,k] == 0
            pivotFound = false
            for k2 = k:n2
               for k1 = k:n1
                  if Afac[k1,k2] != 0
                     pivotFound = true
                     p1k = k1
                     p2k = k2
                     break
                  end
               end
               pivotFound && break
            end

            if pivotFound
               # Row interchange
               if p1k != k
                  for i = p2k:n2
                     temp = Afac[k,i]
                     Afac[k  ,i] = Afac[p1k,i]
                     Afac[p1k,i] = temp
                  end
                  tempRow = Bfac[k,:]
                  Bfac[k,:]   = Bfac[p1k,:]
                  Bfac[p1k,:] = tempRow

                  temp    = p1[k]
                  p1[k]   = p1[p1k]
                  p1[p1k] = temp
               end

               # Column interchange
               if p2k != k
                  for i = 1:n1
                     temp = Afac[i,k]
                     Afac[i,k  ] = Afac[i,p2k]
                     Afac[i,p2k] = temp
                  end
                  temp    = p2[k]
                  p2[k]   = p2[p2k]
                  p2[p2k] = temp
               end
            else
               # no pivot found: submatrix Afac[k:n1,:] has only zeros
               rk = k-1
               return (Afac, rk, p1, p2, Bfac)
            end
         end
      else
         # submatrix Afac[k:n1,:] has only zeros
         rk = k-1
         return (Afac, rk, p1, p2, Bfac)
      end

      # Afac[k,k] != 0
      pivot = Afac[k,k]
      for i = k+1:n1
         Aik  = Afac[i,k]
         Bfac[i,:] = div.(pivot*Bfac[i,:] - Aik*Bfac[k,:], oldPivot)  # guaranteed no remainder
         for j = k+1:n2
            Afac[i,j] = div(pivot*Afac[i,j] - Aik*Afac[k,j], oldPivot)  # guaranteed no remainder
         end
         Afac[i,k] = 0
      end
      oldPivot = pivot
   end
   return (Afac, rk, p1, p2, Bfac)
end

toVector(result) = (result[1], result[2], result[3], result[4], vec(result[5]))
upperTrapezoidal(A::Matrix{Int}, b::Vector{Int}) = toVector(upperTrapezoidal(A, reshape(b,size(b,1),1)))
upperTrapezoidal(A::Matrix{Int})                 =  upperTrapezoidal(A, fill(0,size(A,1),0))

function upperTrapezoidal(A::SparseMatrixCSC{Int,Int}, B::SparseMatrixCSC{Int,Int})
   assert(size(A,1) == size(B,1))
   (n1,n2) = size(A)
   nb   = size(B,2)
   p1   = collect(1:n1)
   p2   = collect(1:n2)
   Afac = copy(A)
   Bfac = copy(B)

   rk   = n1
   p1k  = 0
   p2k  = 0
   Aik  = 0
   temp = 0
   pivot      = 1
   oldPivot   = 1
   pivotFound = false

   for k = 1:n1
      # Search for a pivot in Afac[k:n1,min(k,n2):n2] columnwise
      if k <= n2
         if Afac[k,k] == 0
            Arows = rowvals(Afac)
            pivotFound = false
            Avals = nonzeros(Afac)
            for k2 = k:n2
                for kk1 in nzrange(Afac, k2)
                   k1  = Arows[kk1]
                   val = Avals[kk1]
                    if k1 >= k && val != 0
                     pivotFound = true
                     p1k = k1
                     p2k = k2
                     break
                  end
               end
               pivotFound && break
            end

            if pivotFound
               # Row interchange
               if p1k != k
                  ii = p2k:n2
                  tempAr       = Afac[k,ii]
                  Afac[k  ,ii] = Afac[p1k,ii]
                  Afac[p1k,ii] = tempAr

                  if nb > 0
                     tempRow = Bfac[k,:]
                     Bfac[k,:]   = Bfac[p1k,:]
                     Bfac[p1k,:] = tempRow
                  end

                  temp    = p1[k]
                  p1[k]   = p1[p1k]
                  p1[p1k] = temp
               end

               # Column interchange
               if p2k != k
                  tempAc      = Afac[:,k]
                  Afac[:,k  ] = Afac[:,p2k]
                  Afac[:,p2k] = tempAc

                  temp    = p2[k]
                  p2[k]   = p2[p2k]
                  p2[p2k] = temp
               end
            else
               # no pivot found: submatrix Afac[k:n1,:] has only zeros
               rk = k-1
               return (Afac, rk, p1, p2, Bfac)
            end
         end
      else
         # no pivot found: submatrix Afac[k:n1,:] has only zeros
         rk = k-1
         return (Afac, rk, p1, p2, Bfac)
      end

      # Subtract row k from row i = k+1:n1
      j = k+1:n2
      for ii in nzrange(Afac, k)    # [i,k], i = k:n1
         Arows = rowvals(Afac)
         Avals = nonzeros(Afac)
         i = Arows[ii]
         if i == k
            # Afac[k,k] != 0
            pivot = Avals[ii]
         elseif i > k
            Aik = Avals[ii]
            Avals[ii] = 0
            if nb > 0
               Bfac[i,:] = div.(pivot*Bfac[i,:] - Aik*Bfac[k,:], oldPivot)  # guaranteed no remainder
            end
            Afac[i,j] = div.(pivot*Afac[i,j] - Aik*Afac[k,j], oldPivot)  # guaranteed no remainder
         end
      end
      oldPivot = pivot
   end
   return (Afac, rk, p1, p2, Bfac)
end

upperTrapezoidal(A::SparseMatrixCSC{Int,Int}) =  upperTrapezoidal(A, sparse(fill(0,size(A,1),0)))

"""-----------------------------------------------------------------------
    (iya, eqr, ix1, ix2, eqx, A1, A2) = removeSingularities(A, ix, iy)

A is a dense or sparse Int matrix A[:,:] describing a linear system

    A*v = 0

where

  - v[ix]: Variables used in the derivative operator, der(v[ix]), so these are the potential states.

  - v[iy]: Variables that must be solved from ```A*v = 0```, because these variables do not appear in the remaining part of the model equations (e.g. absolute potentials). If variables v[iy] cannot be uniquely solved from ```A*v = 0```, then in argument iya the variables are listed that must be set arbitrarily so that v[iy] can be uniquely solved from the equations. For example, if v[iy] contains the absolute potentials and no ground object is present, then v[iya] are the absolute potentials that must be explicitly set, e.g., by introducing an equation v[iya]=0, in order that the equations have a unique solution. Additionally, a warning message should be printed that v[iya] was set to 0.)

The function return arguments are all of type Int and have the following meaning:

  - iya: variables v[iya] can be arbitrarily set (iya is a sub-set of iy)

  - eqr: equations ```A[eqr,:]*v = 0``` are redundant and should be removed (they do not change the result)

  - ix1: x[ix1] are the dependent states (ix1 is a sub-set of ix)

  - ix2: i[ix2] are the independent states (ix2 is a sub-set of ix)

  - eqx, A1, A2: equations ```A[eqx,:]*v = 0``` should be replaced by the equations ```A1*x1 + A2*x2 = 0```, where A1 is square, upper triangular and has full rank with A1[i,i] != 0. In other words, x1 can be computed explicitely from these equations given x2.

The calling program should do the following:

   - Remove equations ```A[eqr,:]*v = 0```, because these are redundant equations.
   - Replace equations ```A[eqx,:]*v = 0``` by ```A1*x1 + A2*x2 = 0```, in order that index reduction can be performed by a structural algorithm (e.g. Pantelides algorithm).
   - Introduce equations v[iya] = 0 and print a warning message that these variables have been arbitrarily set to zero. This is needed, in order that the equations have a unique solution.
"""
function removeSingularities(A::T, ix::Vector{Int}, iy::Vector{Int}) where T <: AbstractMatrix{Int}
  # Split equation system A*v = 0 into V*[vy,vr] = X*(-x), where
  #   x  = v[ix]
  #   vy = v[iy]
  #   vr = the remaining variables
  (ne,nv) = size(A)   # number of equations (ne) and variables (nv)
  nx = length(ix)     # number of potential states
  set_iy = Set(iy)
  ivr = collect( setdiff( setdiff(Set(1:nv), Set(ix)), set_iy ) )  # indices of vr
  ivall = [iy;ivr]
  V = [ A[:,iy] A[:,ivr] ]
  X = A[:,ix];

  # Transform V to upper trapezoidal form
  #    Vfac*v[pv2] = Xfac*(-x)
  (Vfac, rk, pv1, pv2, Xfac) = upperTrapezoidal(V, X)

  # iya are the independent variables that are in set_iy
  ivind = ivall[ pv2[rk+1:end] ]
  iya   = collect( intersect(Set(ivind), set_iy) )

  if nx == 0
     eqr = pv1[rk+1:ne]
     eqx = fill(0,0)
     ix1 = fill(0,0)
     ix2 = fill(0,0)
     A1  = fill(0,0,0)
     A2  = fill(0,0,0)

  else
     # Transform state constraints to upper trapezoidal form
     #    Xc*x = 0     // Xc = Xfac[rk+1:end,:]
     #
     #    [Xcfac1 XcFac2;
     #       0       0  ] * [x1,x2] = 0
     Xc = Xfac[rk+1:end,:]
     (Xcfac, rk2, px1, px2) = upperTrapezoidal( Xc )
     eqx2 = px1[1:rk2]
     eqr2 = px1[rk2+1:end]

     eqx  = pv1[rk+eqx2]
     eqr  = pv1[rk+eqr2]
     ix1  = ix[ px2[1:rk2] ]
     ix2  = ix[ px2[rk2+1:end] ]
     A1   = Xcfac[1:rk2, 1:rk2]
     A2   = Xcfac[1:rk2, rk2+1:end]
  end

  return (iya, eqr, ix1, ix2, eqx, A1, A2)
end

"""---------------------------------------------------------------
    str = displayAsString(x)

Return the result of the Julia built-in command display(x) as string
instead of displaying it on the terminal.
"""
displayAsString(x) = sprint((io,x)->display(TextDisplay(io),x),x)

"""---------------------------------------------------------------
    str = displayAsString2(x)

Return the result of the Julia built-in command display(x) as string
instead of displaying it on the terminal, remove the preceding type
information and shift the whole output to the right
"""
function displayAsString2(obj)
   str  = displayAsString(obj)
   str2 = str[searchindex(str,"\n")+1:end]
   str3 = "   " * replace(str2,"\n","\n   ")
end

"""---------------------------------------------------------------
    printobj(name,obj)

Pretty print an instance obj as "<name> = <display(obj)>".
"""
function printobj(name,obj)
   str = displayAsString2(obj)
   println("\n$name =\n$str")
end

function printEquation(coeff, names::Vector{String})
   assert(length(coeff) == length(names))
   print("   ")
   first = true
   for i=1:length(names)
      if coeff[i] != 0
         if first
            first = false
            if coeff[i] < 0
               print("-")
            end
         else
            if coeff[i] < 0
               print(" - ")
            elseif coeff[i] > 0
               print(" + ")
            end
         end

         c = abs(coeff[i])
         if c != 1
            print(c,"*")
         end
         print(names[i])
      end
   end
   println(" = 0")
end

"""
    printRemoveSingularities(A, r, names)

Prints the result of r=removeSingularities(..), where

  - A: matrix of removeSingularities(..)

  - r: the tuple returned by removeSingularities.

  - names: string vector containing the names of the variables v from A*v = 0.
"""
function printRemoveSingularities(A, r, names)
   (iya, eqr, ix1, ix2, eqx, A1, A2) = r

   if length(iya) > 0
      iya_names = displayAsString2(names[iya])
      println("\n  Variables that can be arbitrarily set:\n", iya_names)
   end

   if length(eqr) > 0
      println("\n  Equations that can be removed, since redundant: ", eqr)
      for i=1:length(eqr)
         printEquation(A[eqr[i],:],names)
      end
   end

   if length(ix1) > 0
      ix1_names = displayAsString2(names[ix1])
      ix2_names = displayAsString2(names[ix2])
      println("\n  Dependent state variables x1:\n"  ,ix1_names)
      println("\n  Independent state variables x2:\n",ix2_names)
      println("\n  Equations: ", eqx)
      for i=1:length(eqx)
         printEquation(A[eqx[i],:],names)
      end
      println("\n  that shall be replaced by A1*x1 + A2*x2 = 0: ")
      printobj("   A1", A1)
      printobj("   A2", A2)
      print("\n")
      for i=1:size(A1,1)
         printEquation([A1 A2][i,:],[names[ix1];names[ix2]])
      end
   end
end

end


using Revise
using Test
using NNFEM
using PyCall
using PyPlot
using JLD2
using ADCME
using LinearAlgebra
reset_default_graph()
include("nnutil.jl")

testtype = "NeuralNetwork1D"

prop = Dict("name"=> testtype, "rho"=> 0.1, "E"=> 200.0, "B"=> 10.0,
            "sigmaY"=>0.300, "K"=>1/9*200, "A0"=> 1.0, "eta"=> 10.0, "nn"=>post_nn)

include("NNTrussPull_Domain.jl")

# domain = Domain(nodes, elements, ndofs, EBC, g, FBC, fext)
@load "Data/domain1.jld2" domain
state = zeros(domain.neqs)
∂u = zeros(domain.neqs)
globdat = GlobalData(state,zeros(domain.neqs), zeros(domain.neqs),∂u, domain.neqs, gt, ft)
assembleMassMatrix!(globdat, domain)

nntype = "ae_scaled"
# nntype = "linear"
n_data = [1,2,4,5]
losses = Array{PyObject}(undef, length(n_data))
for (i, ni) in enumerate(n_data)
    state_history, fext_history = read_data("$(@__DIR__)/Data/$i.dat")
    losses[i] = DynamicMatLawLoss(domain, globdat, state_history, fext_history, nn,Δt)
end
loss = sum(losses)

sess = Session(); init(sess)
@show run(sess, loss)
BFGS!(sess, loss, 2000)
ADCME.save(sess, "Data/trained_nn_fem.mat")


function divergence_kernel(
    u, v, w,
    dxi, dyi, dzi,
    is, ie, js, je, ks, ke)

    divmax = 0
    @fast3d begin
        div = @fd (u, v, w) gradx(u) + grady(v) + gradz(w)
        divmax = max(divmax, abs(div))
    end

    return divmax
end

function cfl_kernel(
    u, v, w,
    dxi, dyi, dzi, dt,
    is, ie, js, je, ks, ke)

    cflmax = 0
    @fast3d begin
        cfl = @fd (u, v, w) abs(interpx(u))*dxi + abs(interpy(v))*dyi + abs(interpz(w))*dzi[k]
        cflmax = max(cflmax, cfl)
    end

    return cflmax*dt
end

function calc_divergence(f::Fields, g::Grid)
    div = divergence_kernel(
        f.u, f.v, f.w,
        g.dxi, g.dyi, g.dzi,
        g.is, g.ie, g.js, g.je, g.ks, g.ke)
end

function calc_cfl(f::Fields, g::Grid, t::Timeloop)
    cfl = cfl_kernel(
        f.u, f.v, f.w,
        g.dxi, g.dyi, g.dzi, t.dt,
        g.is, g.ie, g.js, g.je, g.ks, g.ke)
end


using StatsBase: var
function solver_default(X,data::AbstractArray{T,2};stderror=false) where {T<:Union{Missing, <:Number}}
    # likely much larger matrix, using lsqr
    minfo = []
    beta = Array{Float64}(undef,size(data,1),size(X,2))

    @showprogress .1 for ch in 1:size(data,1)
        ix = .!ismissing.(data[ch,:])
        beta[ch,:],h = lsmr(X[ix,:],data[ch,ix],log=true)
        push!(minfo,h)
    end

    if stderror
        stderror = calculate_stderror(X,data,beta)
        minfo = modelinfo(modelinfo,stderror)
    end
    return beta, minfo
end
struct modelinfo
    history
    stderror
end

function calculate_stderror(Xdc,data::Matrix{T},beta) where {T<:Union{Missing, <:Number}}

    # remove missings
    ix = any(.!ismissing.(data),dims=1)[1,:]
    if length(ix)!=size(data,2)
        @warn("Limitation: Missing data are calculated over all channels for standard error")
    end
    data = data[:,ix]
    Xdc = Xdc[ix,:]

    # Hat matrix only once
    hat_prime = inv(Matrix( Xdc'*Xdc))
    # Calculate residual variance
    @warn("Autocorrelation was NOT taken into account. Therefore SE are UNRELIABLE. Use at your own discretion")

    se = Array{Float64}(undef,size(data,1),size(Xdc,2))
    for ch = 1:size(data,1)
        residualVar = var(data[ch,:] .- Xdc*beta[ch,:]);
        @assert(!isnan(residualVar),"residual Variance was NaN")
        hat = hat_prime .* residualVar
        #se = sqrt(diag(cfg.contrast(:,:)*hat*cfg.contrast(:,:)'));
        se[ch,:] = sqrt.(diag(hat))
    end
    return se
end
function calculate_stderror(X,  data::AbstractArray{T,3},beta) where {T<:Union{Missing, <:Number}}
#function calculate_stderror(Xdc,data::AbstractArray{T,2},beta) where {T<:Union{Missing, <:Number}}

    # Hat matrix
    hat_prime = inv(Matrix( X'*X))
    # Calculate residual variance
    @warn("Autocorrelation was NOT taken into account. Therefore SE are UNRELIABLE. Use at your own discretion")

    se = Array{Float64}(undef,size(data,1),size(data,2),size(X,2))
    for ch = 1:size(data,1)
        for t = 1:size(data,2)
            ix = .!ismissing.(data[ch,t,:])
            residualVar = var(data[ch,t,ix] .- X[ix,:]*beta[ch,t,:]);
            @assert(!isnan(residualVar),"residual Variance was NaN")
            hat = hat_prime .* residualVar
            #se = sqrt(diag(cfg.contrast(:,:)*hat*cfg.contrast(:,:)'));
            se[ch,t,:] = sqrt.(diag(hat))
        end
    end
    return se
end
function solver_default(X,data::AbstractArray{T,3};stderror=false) where {T<:Union{Missing, <:Number}}
       beta = Array{Union{Missing,Number}}(undef,size(data,1),size(data,2),size(X,2))
       @showprogress .1 for ch in 1:size(data,1)
           for t in 1:size(data,2)
               @debug("$(ndims(data,)),$t,$ch")
               ix = .!ismissing.(data[ch,t,:])
               beta[ch,t,:] = X[ix,:] \ data[ch,t,ix]
           end
       end

       minfo = [undef] # no history implemented (yet?)
       if stderror
            stderror = calculate_stderror(X,data,beta)
            minfo = modelinfo(minfo,stderror)
        end
       return beta, minfo
end

solver_b2b(X,data,cross_val_reps) = solver_b2b(X,data,cross_val_reps = cross_val_reps)
function solver_b2b(X,data::AbstractArray{T,3};cross_val_reps = 10) where {T<:Union{Missing, <:Number}}

    X,data = dropMissingEpochs(X,data)

    E = zeros(size(data,2),size(X,2),size(X,2))
    W = Array{Float64}(undef,size(data,2),size(X,2),size(data,1))

    prog = Progress(size(data,2)*cross_val_reps,.1)
    for t in 1:size(data,2)

        for m in 1:cross_val_reps
            k_ix = collect(Kfold(size(data,3),2))
            Y1 = data[:,t,k_ix[1]]
            Y2 = data[:,t,k_ix[2]]
            X1 = X[k_ix[1],:]
            X2 = X[k_ix[2],:]

            G = (Y1' \ X1)
            H = X2 \ (Y2'*G)

            E[t,:,:] = E[t,:,:]+Diagonal(H[diagind(H)])
            ProgressMeter.next!(prog; showvalues = [(:time,t), (:cross_val_rep,m)])
        end
        E[t,:,:] = E[t,:,:] ./ cross_val_reps
        W[t,:,:] = (X*E[t,:,:])' / data[:,t,:]

    end

    # extract diagonal
    beta = mapslices(diag,E,dims=[2,3])
    # reshape to conform to ch x time x pred
    beta = permutedims(beta,[3 1 2])
    modelinfo = Dict("W"=>W,"E"=>E,"cross_val_reps"=>cross_val_reps) # no history implemented (yet?)
    return beta, modelinfo
end


@doc raw"""
    SpecialLinear{n,𝔽} <:
        AbstractGroupManifold{𝔽,MultiplicationOperation,DefaultEmbeddingType}

The special linear group ``\mathrm{SL}(n,𝔽)`` that is, the group of all invertible matrices
with unit determinant in ``𝔽^{n×n}``.

The Lie algebra ``𝔰𝔩(n, 𝔽) = T_e \mathrm{SL}(n,𝔽)`` is the set of all matrices in
``𝔽^{n×n}`` with trace of zero. By default, tangent vectors ``X_p ∈ T_p \mathrm{SL}(n,𝔽)``
for ``p ∈ \mathrm{SL}(n,𝔽)`` are represented with their corresponding Lie algebra vector
``X_e = p^{-1}X_p ∈ 𝔰𝔩(n, 𝔽)``.

The default metric is the same left-``\mathrm{GL}(n)``-right-``\mathrm{O}(n)``-invariant
metric used for [`GeneralLinear(n, 𝔽)`](@ref). The resulting geodesic on
``\mathrm{GL}(n,𝔽)`` emanating from an element of ``\mathrm{SL}(n,𝔽)`` in the direction of
an element of ``𝔰𝔩(n, 𝔽)`` is a closed subgroup of ``\mathrm{SL}(n,𝔽)``. As a result, most
metric functions forward to `GeneralLinear`.
"""
struct SpecialLinear{n,𝔽} <:
       AbstractGroupManifold{𝔽,MultiplicationOperation,TransparentIsometricEmbedding} end

SpecialLinear(n, 𝔽::AbstractNumbers=ℝ) = SpecialLinear{n,𝔽}()

function allocation_promotion_function(::SpecialLinear{n,ℂ}, f, args::Tuple) where {n}
    return complex
end

function check_point(G::SpecialLinear{n,𝔽}, p; kwargs...) where {n,𝔽}
    mpv = check_point(Euclidean(n, n; field=𝔽), p; kwargs...)
    mpv === nothing || return mpv
    detp = det(p)
    if !isapprox(detp, 1; kwargs...)
        return DomainError(
            detp,
            "The matrix $(p) does not lie on $(G), since it does not have a unit " *
            "determinant.",
        )
    end
    return nothing
end
check_point(::GT, ::Identity{GT}; kwargs...) where {GT<:SpecialLinear} = nothing
function check_point(G::SpecialLinear, e::Identity; kwargs...)
    return DomainError(e, "The identity element $(e) does not belong to $(G).")
end

function check_vector(G::SpecialLinear, p, X; kwargs...)
    mpv = check_vector(decorated_manifold(G), p, X; kwargs...)
    mpv === nothing || return mpv
    trX = tr(inverse_translate_diff(G, p, p, X, LeftAction()))
    if !isapprox(trX, 0; kwargs...)
        return DomainError(
            trX,
            "The matrix $(X) does not lie in the tangent space of $(G) at $(p), since " *
            "its Lie algebra representation is not traceless.",
        )
    end
    return nothing
end

decorated_manifold(::SpecialLinear{n,𝔽}) where {n,𝔽} = GeneralLinear(n, 𝔽)

default_metric_dispatch(::SpecialLinear, ::EuclideanMetric) = Val(true)
default_metric_dispatch(::SpecialLinear, ::LeftInvariantMetric{EuclideanMetric}) = Val(true)

inverse_translate_diff(::SpecialLinear, p, q, X, ::LeftAction) = X
inverse_translate_diff(::SpecialLinear, p, q, X, ::RightAction) = p * X / p

function inverse_translate_diff!(G::SpecialLinear, Y, p, q, X, conv::ActionDirection)
    return copyto!(Y, inverse_translate_diff(G, p, q, X, conv))
end

function manifold_dimension(G::SpecialLinear)
    return manifold_dimension(decorated_manifold(G)) - real_dimension(number_system(G))
end

@doc raw"""
    project(G::SpecialLinear, p)

Project ``p ∈ \mathrm{GL}(n, 𝔽)`` to the [`SpecialLinear`](@ref) group
``G=\mathrm{SL}(n, 𝔽)``.

Given the singular value decomposition of ``p``, written ``p = U S V^\mathrm{H}``, the
formula for the projection is

````math
\operatorname{proj}_{\mathrm{SL}(n, 𝔽)}(p) = U S D V^\mathrm{H},
````
where

````math
D_{ij} = δ_{ij} \begin{cases}
    1            & \text{ if } i ≠ n \\
    \det(p)^{-1} & \text{ if } i = n
\end{cases}.
````
"""
project(::SpecialLinear, p)

function project!(::SpecialLinear{n}, q, p) where {n}
    detp = det(p)
    isapprox(detp, 1) && return copyto!(q, p)
    F = svd(p)
    q .= F.U .* F.S'
    q[:, n] ./= detp
    mul!_safe(q, q, F.Vt)
    return q
end

@doc raw"""
    project(G::SpecialLinear, p, X)

Orthogonally project ``X ∈ 𝔽^{n × n}`` onto the tangent space of ``p`` to the
[`SpecialLinear`](@ref) ``G = \mathrm{SL}(n, 𝔽)``. The formula reads
````math
\operatorname{proj}_{p}
    = (\mathrm{d}L_p)_e ∘ \operatorname{proj}_{𝔰𝔩(n, 𝔽)} ∘ (\mathrm{d}L_p^{-1})_p
    \colon X ↦ X - \frac{\operatorname{tr}(X)}{n} I,
````
where the last expression uses the tangent space representation as the Lie algebra.
"""
project(::SpecialLinear, p, X)

function project!(G::SpecialLinear{n}, Y, p, X) where {n}
    inverse_translate_diff!(G, Y, p, p, X, LeftAction())
    Y[diagind(n, n)] .-= tr(Y) / n
    translate_diff!(G, Y, p, p, Y, LeftAction())
    return Y
end

function decorator_transparent_dispatch(::typeof(project), ::SpecialLinear, args...)
    return Val(:parent)
end

Base.show(io::IO, ::SpecialLinear{n,𝔽}) where {n,𝔽} = print(io, "SpecialLinear($n, $𝔽)")

translate_diff(::SpecialLinear, p, q, X, ::LeftAction) = X
translate_diff(::SpecialLinear, p, q, X, ::RightAction) = p \ X * p

function translate_diff!(G::SpecialLinear, Y, p, q, X, conv::ActionDirection)
    return copyto!(Y, translate_diff(G, p, q, X, conv))
end


L = 20
labels = rand(Bool, L);
scores = add_noise.(labels, 0.6)

curve = DET(labels, scores)
B = BinaryMetrics(labels, scores)
curve = DET(B)
println(curve)

OP = op(B)
@test op(curve) == OP
@test fpr(curve) == fpr(B,OP)
@test tpr(curve) == tpr(B,OP)


using Plots, DifferentialEquations
dt = 0.001
σnoise = 0.007
τAMPA = 2e-3
RNG = MersenneTwister(10)
RNG2 = MersenneTwister(10)

function model1(du, u, p, t, W)
    x, I = u
    du[1] = 0.0
    du[2] = (W[1]*σnoise*√τAMPA - I) / τAMPA
end

prob = RODEProblem(model1, [0.0, 0.0], (0.0, 1.0), noise=WienerProcess(0.0, 0.0, rng=RNG, reseed=false))
sol = solve(prob, saveat=dt, dt=dt, reltol=1e-15, abstol=1e-15)
display(plot(sol))

function model2(du, u, p, t, W)
    du[1] = 2*u[1]*sin(W[1])
end

W = OrnsteinUhlenbeckProcess(τAMPA, 0.0, σnoise*√τAMPA, 0.0, [0.0], rng=RNG2, reseed=false)
prob2 = RODEProblem(model2, [0.0], (0.0, 1.0), noise=W)
sol2 = solve(prob2, saveat=dt, dt=dt, reltol=1e-15, abstol=1e-15)
display(plot(sol2.W.t, reduce(vcat, sol2.W.u)))


export policy_evaluation!, policy_improvement!, policy_iteration!

using Distributions:probs

"""
    policy_evaluation!(;V, π, model, γ, θ)

# Keyword arguments

- `V`, an [`AbstractApproximator`](@ref).
- `π`, an `AbstractPolicy`.
- `model`, a distribution based environment model(given a state and action
  pair, return all possible reward, next state, termination info and
  corresponding probability).
- `γ::Float64`, discount rate.
- `θ::Float64`, threshold to stop evaluation.
"""
function policy_evaluation!(
    ;
    V::AbstractApproximator,
    π::AbstractPolicy,
    model::AbstractEnvironmentModel,
    γ::Float64 = 0.9,
    θ::Float64 = 1e-4,
)
    states, actions = state_space(model), action_space(model)
    while true
        Δ = 0.0
        for s in states
            v = sum(
                prob(π, s, a) * sum(
                    p * (r + (1-t) * γ * V(s′))
                    for ((r, t, s′), p) in model(s, a)
                )
                for a in actions
            )
            δ = V(s) - v
            update!(V, s => δ)
            Δ = max(Δ, abs(δ))
        end
        Δ < θ && break
    end
    V
end

function policy_improvement!(
    ;
    V::AbstractApproximator,
    π::AbstractPolicy,
    model::AbstractEnvironmentModel,
    γ::Float64 = 0.9,
)
    states, actions = state_space(model), action_space(model)
    is_policy_stable = true
    for s in states
        old_a = π(s)
        best_action_inds = find_all_max(
            [
                sum(
                    p * (r + (1-t)* γ * V(s′))
                    for ((r, t, s′), p) in model(s, a,)
                ) for a in actions
            ]
        )[2]
        new_a = rand(best_action_inds)  # break tie
        if new_a != old_a
            update!(π, s => new_a)
            is_policy_stable = false
        end
    end
    is_policy_stable
end

function policy_iteration!(
    ;
    V::AbstractApproximator,
    π::AbstractPolicy,
    model::AbstractEnvironmentModel,
    γ::Float64 = 0.9,
    θ::Float64 = 1e-4,
    max_iter = typemax(Int),
)
    for i = 1:max_iter
        policy_evaluation!(; V = V, π = π, model = model, γ = γ, θ = θ)
        policy_improvement!(; V = V, π = π, model = model, γ = γ) && return i
    end
    return max_iter
end

function value_iteration!(
    ;
    V::AbstractApproximator,
    model::AbstractEnvironmentModel,
    γ::Float64 = 0.9,
    θ::Float64 = 1e-4,
    max_iter = typemax(Int),
)
    states, actions = state_space(model), action_space(model)
    for i = 1:max_iter
        Δ = 0.0
        for s in states
            v = maximum(sum(p * (r + (1-t) * γ * V(s′)) for ((r, t, s′), p) in model(s, a)) for a in actions)
            δ = V(s) - v
            update!(V, s => δ)
            Δ = max(Δ, abs(δ))
        end
        Δ < θ && break
    end
end


using NPZ
using TomoForward
using Images
using Plots
using XfromProjections
using StaticArrays
using PyCall
using Logging
include("../../phantoms/simple_phantoms.jl")
include("../tv_primaldual_flow.jl")
include("../optical_flow.jl")
include("../../../src/analytic/gridrec.jl")

cwd = @__DIR__
cd(cwd)

like(x::T, y) where T = convert(T, y)

proj = npzread("data/projections.npy")
theta = npzread("data/theta.npy")

nangles, imageZ, detcount = size(proj)

H, W = 100,100

function radon_operator(H,W, detcount, angles)
    proj_geom = ProjGeom(1.0, detcount, angles)
    A = fp_op_parallel2d_line(proj_geom, H, W)
    return A
end

pr = Float64.(reshape(sum(proj,dims=2), (nangles, detcount)))

window_step = 2
nframes = Int(nangles/window_step)

frames = zeros(H,W,nframes)
As = map(t -> radon_operator(H,W,detcount,theta[t:t+1]),1:window_step:nangles)
bs = reshape(pr,(detcount*window_step,nframes))
niter=500
u0s = zeros(H,W,nframes)

w_tv = 0.3
w_flow  = 0.05

@info "Reconstructing using joint motion estimation and reconstruction"
c=10.0
us_flow = recon2d_tv_primaldual_flow(As, bs, u0s, 20, niter, w_tv, w_flow, c)

@info "Reconstruction using tv regularization frame by frame"
us_tv = zeros(H,W,nframes)
for t = 1:nframes
    A = As[t]
    p = bs[:,t]
    u0 = u0s[:,:,t]
    us_tv[:,:,t] .= recon2d_tv_primaldual!(us_tv[:,:,t], A, p, niter, w_tv, c)
end

@info "Preparing results in human readable format"
anim = @animate for t=1:nframes
    l = @layout [a b c]
    p1 = plot(Gray.(frames[:,:,t]), aspect_ratio=:equal, framestyle=:none, title="Ground truth")
    p2 = plot(Gray.(us_flow[:,:,t]), aspect_ratio=:equal, framestyle=:none, title="Flow")
    p3 = plot(Gray.(us_tv[:,:,t]), aspect_ratio=:equal, framestyle=:none, title="TV")
    plot(p1, p2, p3, layout = l)
end

cwd = @__DIR__
path = normpath(joinpath(@__DIR__, "result"))
!isdir(path) && mkdir(path)
cd(path)
gif(anim, "reconstruction_flow.gif", fps = 1)
cd(cwd)


# Guizar-Sicairos2004JotOSoAA - Computation of quasi-discrete Hankel transforms of integer order for propagating optical wave fields

#=
Calculates an approximation to
\[f_2(\nu) = 2\pi\int\limits_0^\infty dr\; r f_1(r)J_p(2\pi \nu r)\]

as \[\vec{F}_2=\mathsf{T}\vec{F}_1\]

where \[\vec{F}_2(m)=\left\{\frac{f_2\left(\frac{\alpha_{pm}}{2\pi R}\right)V}{|J_{p+1}(\alpha_{pm})|}\right\}\],
\[\vec{F}_1(n)=\left\{\frac{f_1\left(\frac{\alpha_{pn}}{2\pi V}\right)R}{|J_{p+1}(\alpha_{pn})|}\right\}\],
and \[\mathsf{T}_{mn}=\frac{2J_p(\alpha_{pn}\alpha_{pm}/S)}{|J_{p+1}(\alpha_{pn})||J_{p+1}(\alpha_{pm})|S}\].

\(R\) is maximum radius, \(V\) the maximum frequency and \(S=2\pi RV\).

\(f_1\) is assumed to be sampled on these gridpoints.
=#

export Guizar, guizar

using GSL

type Guizar
    p::Int
    R
    N::Int
    r
    ν
    V
    Jp1
    T
end

function Guizar(p::Int, R, N::Int)
    α = sf_bessel_zero_Jnu(p, 1:N)
    αNp1 = sf_bessel_zero_Jnu(p, N+1)
    r = α*R/αNp1
    ν = α/(2π*R)
    V = αNp1/(2π*R)
    S = 2π*R*V
    Jp1 = abs(besselj(p+1,α))
    T = 2besselj(p,(α*α')/S)./(S*Jp1*Jp1')
    Guizar(p,R,N,r,ν,V,Jp1,T)
end

function (gt::Guizar)(v::AbstractArray)
    F1 = (v./gt.Jp1)*gt.R
    F2 = gt.T*F1
    gt.ν,F2.*gt.Jp1/gt.V
end

function (gt::Guizar)(f1::Function)
    gt(f1(gt.r))
end

function guizar(f1, p::Int, R, N::Int)
    Guizar(p,R,N)(f1)
end


using ModelingToolkit, DiffEqBase, LinearAlgebra

# Define some variables
@parameters t x y
@variables u(..)

ConstrainedEquation([x ~ 0, y < 1 / 2], u(t, x, y) ~ x + y^2)


using ForneyLab
using LinearAlgebra
import ForneyLab: SoftFactor, @ensureVariables, generateId, addNode!, associate!,
                  averageEnergy, differentialEntropy, Interface, Variable, slug, ProbabilityDistribution,
				  unsafeLogMean, unsafeMean, unsafeCov, unsafePrecision, unsafeMeanCov,
                  ultimatePartner, region, Region, isClamped, currentInferenceAlgorithm
export LatentAutoregressiveX, LARX

"""
Description:

    A Latent Autoregressive model with eXogenous input (LARX).

    Consider states x ∈ R, with y = [x_k, x_k-1, ... x_k-M+1] and
    z = [x_k-1, ... x_k-M]. The first element of y (x_k) is produced by an inner
    product between autoregression coefficients θ and M previous states (i.e. z),
    as well as a coefficient η times an input variable u. The remainder of y is
    produced by the first M-1 elements of z.

    The node function is a Gaussian with mean-precision parameterization:

    f(y, θ, z, η, u, γ) = 𝒩(y | A(θ)z + B(η)u, V(γ)),

    where A(θ) is a state transition matrix dependent on unknown coefficients θ.
    The matrix is split into a matrix

        S = |0 .. 0
              I   0|,

    which shifts the elements of the state vector downwards and drops the
    oldest state x_k-M, and a vector s = [1 .. 0]', which can be used to
    construct a matrix with the autoregression coefficients as top row;

        sθ' = | θ_1 .. θ_M
                 ⋮       ⋮
                 0  ..  0 |

    B(η)u is a scaled linear additive control and V(γ) is a covariance
    matrix based on process precision γ.

Interfaces:

    1. y (state vector at time k; i.e. [x_k, x_k-1, ... x_k-M+1])
    2. θ (autoregression coefficients; i.e. [θ_1, ... θ_M])
    3. z (state vector at time k-1; i.e. [x_k-1, ... x_k-M])
    4. η (control coefficient)
    5. u (control / input)
    6. γ (precision)

Construction:

    LatentAutoregressiveX(y, θ, z, η, u, γ, id=:some_id)
"""

mutable struct LatentAutoregressiveX <: SoftFactor
    id::Symbol
    interfaces::Vector{Interface}
    i::Dict{Symbol,Interface}

    function LatentAutoregressiveX(y, θ, z, η, u, γ; id=generateId(LatentAutoregressiveX))
        @ensureVariables(y, z, θ, η, u, γ)
        self = new(id, Array{Interface}(undef, 6), Dict{Symbol,Interface}())
        addNode!(currentGraph(), self)
        self.i[:y] = self.interfaces[1] = associate!(Interface(self), y)
        self.i[:z] = self.interfaces[2] = associate!(Interface(self), z)
        self.i[:θ] = self.interfaces[3] = associate!(Interface(self), θ)
        self.i[:η] = self.interfaces[4] = associate!(Interface(self), η)
        self.i[:u] = self.interfaces[5] = associate!(Interface(self), u)
        self.i[:γ] = self.interfaces[6] = associate!(Interface(self), γ)
        return self
    end
end

slug(::Type{LatentAutoregressiveX}) = "LARX"

function averageEnergy(::Type{LatentAutoregressiveX},
                       marg_y::ProbabilityDistribution{Multivariate},
                       marg_z::ProbabilityDistribution{Multivariate},
                       marg_θ::ProbabilityDistribution{Multivariate},
                       marg_η::ProbabilityDistribution{Univariate},
                       marg_u::ProbabilityDistribution{Univariate},
                       marg_γ::ProbabilityDistribution{Univariate})

    # Expectations of marginal beliefs
    my, Vy = unsafeMeanCov(marg_y)
    mz, Vz = unsafeMeanCov(marg_z)
    mθ, Vθ = unsafeMeanCov(marg_θ)
    mη, vη = unsafeMeanCov(marg_η)
    mu, vu = unsafeMeanCov(marg_u)
    mγ = unsafeMean(marg_γ)

    error("Not implemented yet")

    # Compute
    Az = mθ*mz
    Eg2 = Eg*Eg' + Jx'*Vz*Jx + Jθ'*Vθ*Jθ

    # Expand square and pre-compute terms
    sq1 = my[1]^2 + Vy[1,1]
	sq2 = my[1]*(Eg + mη*mu)
	sq3 = Eg2 + 2*Eg*mη*mu + (mη^2 + vη)*(mu + vu)

	# Compute average energy
	AE = 1/2*log(2*π) -1/2*unsafeLogMean(marg_γ) +1/2*mγ*(sq1 -2*sq2 + sq3)

    # Correction
    AE += differentialEntropy(marg_y)
    marg_y1 = ProbabilityDistribution(Univariate, GaussianMeanVariance, m=my[1], v=Vy[1,1])
    AE -= differentialEntropy(marg_y1)

    return AE
end


module BayesianIntegral

using Distributions
using LinearAlgebra
using Random
using Optim

# Basic operation of Kriging model
include("kriging_structs.jl")
export KrigingModel, GaussianKernelHyperparameters
include("making_covariances.jl")
export gaussian_kernel, marginal_gaussian_kernel, K_matrix
export marginal_likelihood_gaussian_derivatives, log_likelihood
export correlation_vector_of_a_point
include("evaluate_kriging_model.jl")
export evaluate, predicted_error
# Integration of Kriging Model
include("integrate_given_hyperparameters.jl")
export integrate
# Calibration of Kriging Model
include("calibrate_weights.jl")
export calibrate_by_ML_with_SGD, RProp_params, calibrate_by_ML_with_Rprop

# Searching for Optima.
include("Expected_improvement_optimisation.jl")
export expected_improvement, get_next_query_point_through_expected_improvement, get_predicted_minimum

end


module Camera

export SimpleCamera, FovCamera, getray

using CthulhuVision.Math
using CthulhuVision.Random
using CthulhuVision.Light

struct SimpleCamera
    origin     :: Vector3
    lowerleft  :: Vector3
    horizontal :: Vector3
    vertical   :: Vector3
end

@inline function getray(cam::SimpleCamera, u::Float32, v::Float32, rng::UniformRNG) :: Ray
    direction = cam.lowerleft + u*cam.horizontal + v*cam.vertical - cam.origin
    Ray(cam.origin, direction)
end

@inline unithost(v::Vector3) = 1.0f0 / sqrt(v.x * v.x + v.y * v.y + v.z * v.z) * v

@inline function randominunitdisk(rng::UniformRNG) :: Vector3
    p = Vector3(1.0f0, 1.0f0, 0.0f0)

    while dot(p, p) >= 1.0f0
        p = 2.0f0 * Vector3(next(rng), next(rng), 0.0f0) - Vector3(1.0f0, 1.0f0, 0.0f0)
    end

    p
end

struct FovCamera
    origin     :: Vector3
    lowerleft  :: Vector3
    horizontal :: Vector3
    vertical   :: Vector3
    u          :: Vector3
    v          :: Vector3
    w          :: Vector3
    lensradius :: Float32

    function FovCamera(lookfrom::Vector3, lookat::Vector3, vup::Vector3, vfov::Float32, aspect::Float32, aperture::Float32, focusdist::Float32)
        lensradius = aperture / 2.0f0
        theta = vfov * pi / 180.0f0
        halfheight = tan(theta / 2.0f0)
        halfwidth = aspect * halfheight

        origin = lookfrom
        w = unithost(lookfrom - lookat)
        u = unithost(cross(vup, w))
        v = cross(w, u)

        lowerleft  = origin - halfwidth * focusdist * u - halfheight * focusdist * v - focusdist * w
        horizontal = 2.0f0 * halfwidth * focusdist * u
        vertical   = 2.0f0 * halfheight * focusdist * v

        new(origin, lowerleft, horizontal, vertical, u, v, w, lensradius)
    end
end

@inline function getray(cam::FovCamera, s::Float32, t::Float32, rng::UniformRNG) :: Ray
    rd = cam.lensradius * randominunitdisk(rng)
    offset = rd.x * cam.u + rd.y * cam.v
    direction = cam.lowerleft + s*cam.horizontal + t*cam.vertical - cam.origin - offset
    Ray(cam.origin + offset, direction)
end

end


module model
using Distributions
include("data_repeatrate.jl")
include("tree_number_estimate.jl")
include("kmean.jl")
include("sigma_estimate.jl")

function fit(df)
    x1 = 30
    x2 = 38
    y1 = 65
    y2 = 75

    one_tree,df_close = data_repeatrate.Onetree(df,x1,x2,y1,y2)
    #println(one_tree)

    fruit_number = data_repeatrate.Areafruit(one_tree)
    #println(fruit_number)

    df_c = df[df.Close.>0,:]

    tree_number = round(Int, tree_number_estimate.number_estimate(df_c, fruit_number))
    println(tree_number)
    df_k = [df_c.X df_c.Y]

    #df_k = DataFrames(X = df_c.X, Y = df_c.Y)
    #println(size(df_k))
    centers = kmean.K_mean(df_k', tree_number, df_c.Close .^(1/2))
    # println("tree centers")
    # println(centers)
    clustInd = kmean.k_mean_predict(df_k', centers)
    # println(clustInd)
    df_centers = centers[:, clustInd]
    loc = [df.X df.Y]
    esigma_x, esigma_y = sigma_estimate.sigma_est(loc, df_centers, df.Close')
    esigma = [esigma_x 0; 0 esigma_y]
    return fruit_number, centers, esigma
end

function predict(loc::AbstractMatrix{<:Real},
    fruit_number::Float64,
    centers::AbstractMatrix{<:Real},
    esigma::AbstractMatrix{<:Real})
    d, n = size(loc)
    d2, k = size(centers)
    (d == 2) || throw(ArgumentError(
    "dims of data $d should be 2"))
    (d == d2) || throw(ArgumentError(
    "dims of data $d and center $d2 should be equal"))
    pred_fruit_arr = Array{Float64, 2}(undef, n, k)

    for j in 1:k
        center_t = centers[:, j]
        distriNormal = MvNormal(center_t, esigma)
        for i in 1:n
            loc_t = loc[:, i]
            pred_fruit_arr[i, j] = fruit_number * pdf(distriNormal, loc_t)
        end
    end
    pred_fruit = sum(pred_fruit_arr, dims=2)[:,1]
    return pred_fruit

end # function

function mse(pred_fruit, gt_fruit)
    len1 = length(pred_fruit)
    len2 = length(gt_fruit)
    (len1 == len2) || throw(ArgumentError(
    "length of pred_fruit $len1 and gt_fruit $len2 should be equal"))
    return mean((pred_fruit - gt_fruit).^ 2)
end # function

function accuarcy_soft(pred_fruit, gt_fruit)
    len1 = length(pred_fruit)
    len2 = length(gt_fruit)
    (len1 == len2) || throw(ArgumentError(
    "length of pred_fruit $len1 and gt_fruit $len2 should be equal"))
    correct_soft = sum(abs.(pred_fruit - gt_fruit) .< 3)
    return correct_soft/len1
end

end  # module model


module HRP

using Statistics, LinearAlgebra, Clustering

export hrp, ivp

dist(x, y) = norm(collect(skipmissing(x .- y)), 2)

function covar(x, y)
    id = findall(z -> !ismissing(z), x .- y)
    return cov(x[id], y[id])
end

function ret2dist(r)
    n = size(r, 2)
    Σ = fill(0.0, n, n)
    D = fill(0.0, n, n)
    for i = 1:n
        for j = i:n
            D[i, j] = dist(r[:, i], r[:, j])
            D[j, i] = D[i, j]
            Σ[i, j] = covar(r[:, i], r[:, j])
            Σ[j, i] = Σ[i, j]
        end
    end
    return Σ, D
end

function ivp(C)
    iv = 1.0 ./ diag(C)
    return iv ./ sum(iv)
end

function clustervar(C, cluster)
    V = C[cluster, cluster]
    u = ivp(V)
    return dot(u' * V, u)
end

function bisect(items)
    tmp = Vector{Int}[]
    for item in items
        n1 = length(item)
        if n1 > 1
            n2 = floor(Int, n1 / 2)
            push!(tmp, item[1:n2])
            push!(tmp, item[n2+1:end])
        end
    end
    return tmp
end

function update!(w, items, Σ)
    for k in 1:2:length(items)
        c1, c2 = items[k], items[k+1]
        v1, v2 = clustervar(Σ, c1), clustervar(Σ, c2)
        α = v1 / (v1 + v2)
        w[c1] .= w[c1] .* α
        w[c2] .= w[c2] .* (1.0 - α)
    end
end

function hrp(x; input_type="ret")
    Σ = fill(0.0, size(x, 2), size(x, 2))
    D = similar(Σ)
    if input_type == "ret"
        Σ, D = ret2dist(x)
    elseif input_type == "cov"
        Σ, D = x, cov2dist(x)
    end
    h = hclust(D).order
    n = length(h)
    w = fill(1.0, n)
    items = [h]
    while !isempty(items)
        items = bisect(items)
        update!(w, items, Σ)
    end
    return ivp(Σ), w
end

#= function hrp(Σ)
    D = sqrt.(2.0*(1.0.-cov2corr(Σ)))
    id = 1:size(D,2)
    w = fill(1.0, size(D,2))
    while length(id) > 1
        id, k = split(D, id)
        v1 = clustervar(Σ, id)
        v2 = Σ[k,k]
        α = ivp([v1,v2])[1]
        w[id] .= w[id] * α
        w[k] = w[k] * (1. - α)
    end
    return w
end =#

function split(D::Matrix{Float64}, id::Union{UnitRange{Int64},Vector{Int64}})
    if length(id) == 1
        return nothing, id
    end
    tmp = []
    for i in id
        push!(tmp, minimum([D[i, k] for k in setdiff(id, i)]))
    end
    k = id[argmax(tmp)]
    id = setdiff(id, k)
    return id, k
end

function hcluster_order(D)
    n = size(D, 2)
    id = 1:n
    out = []
    for i = 1:n
        id, k = split(D, id)
        push!(out, k)
    end
    return reverse(collect(Iterators.flatten(out)))
end

function cov2corr(C)
    σinv = diagm(1.0 ./ sqrt.(diag(C)))
    return σinv * C * σinv
end

function cov2dist(C)
    R = cov2corr(C)
    for i ∈ 1:size(R, 2)
        R[i, i] = 1.0
    end
    D = sqrt.(0.5 * (1.0 .- R))
    for i ∈ 1:size(D, 2)
        for j ∈ i+1:size(D, 2)
            D[i, j] = D[j, i]
        end
    end
    return D
end

end


using Random: randn

@testset "self-coherence" begin
    N, M = 1024, 128
    t = 1:N
    X = zeros(Float64, N, M)
    for i in 1:M
        X[:,i] = sin.(2pi*t/50) + randn(N)
    end
    wt = ContinuousWaveletTransform(CauchyWavelet(10), [50])
    C = coherence(wt, X; nsurrogate=1)
    @test sum(abs.(C)[:,1])/N ≈ 1.0 atol=1e-1
end

@testset "cross-coherence" begin
    N, M = 1024, 128
    t = 1:N
    X = zeros(Float64, N, M)
    Y = zeros(Float64, N, M)
    for i in 1:M
        X[:,i] = sin.(2pi*t/50) + randn(N)
        Y[:,i] = cos.(2pi*t/50) + randn(N)
    end
    wt = ContinuousWaveletTransform(CauchyWavelet(10), [50])
    C = coherence(wt, X, Y, nsurrogate=1)
    @test sum(abs.(C)[:,1])/N ≈ 1.0 atol=1e-1
end


∂ˣ = PointWiseSpectralOperator(im * kˣ)

ϕ = ScalarField(ψ, grid)

∂ˣϕ = ∂ˣ(ϕ)

∂ˣψ = ∂ˣ(ψ)


function pressure(model::EoSModel, v, T,  z=SA[1.])
    return -∂f∂v(model,v,T,z)
end

function vt_entropy(model::EoSModel, v, T,   z=SA[1.])
    return -∂f∂t(model,v,T,z)
end

function vt_chemical_potential(model::EoSModel, v, T, z= SA[1.])
    fun(x) = eos(model, v, T,x)
    return ForwardDiff.gradient(fun,z)
end

function vt_internal_energy(model::EoSModel, v, T,  z=SA[1.])
    _df,f =  ∂f(model,v,T,z)
    dv,dt = _df
    return f  - dt*T
end

function vt_enthalpy(model::EoSModel, v, T,  z=SA[1.])
    _df,f =  ∂f(model,v,T,z)
    dv,dt = _df
    return f  - dv*v - dt*T
end

function vt_gibbs_free_energy(model::EoSModel, v, T,  z=SA[1.])
    _df,f =  ∂f(model,v,T,z)
    dv,dt = _df
    return f  - dv*v
end

function vt_helmholtz_free_energy(model::EoSModel, v, T,  z=SA[1.])
    return eos(model, v, T,z)
end

function vt_isochoric_heat_capacity(model::EoSModel, v, T,  z=SA[1.])
    fun(x)  = eos(model, v, x,z)
    df(x)   = ForwardDiff.derivative(fun,x)
    d2f(x)  = ForwardDiff.derivative(df,x)
    return -T*d2f(T)
end

function vt_isobaric_heat_capacity(model::EoSModel, v, T,  z=SA[1.])
    d2f = f_hess(model,v,T,z)
    return T*(d2f[1,2]^2/d2f[1]-d2f[2,2])
end

function vt_isothermal_compressibility(model::EoSModel, v, T,  z=SA[1.])
    p0,dpdv = p∂p∂v(model,v,T,z)
    return -1/v*dpdv^-1
end

function vt_isentropic_compressibility(model::EoSModel, v, T,  z=SA[1.])
    d2f = f_hess(model,v,T,z)
    return 1/v*(d2f[1]-d2f[1,2]^2/d2f[2,2])^-1
end

function vt_speed_of_sound(model::EoSModel, v, T,  z=SA[1.])
    Mr      = molecular_weight(model,z)
    d2f = f_hess(model,v,T,z)
    return v*sqrt((d2f[1]-d2f[1,2]^2/d2f[2,2])/Mr)
end

function vt_isobaric_expansivity(model::EoSModel, v, T,  z=SA[1.])
    d2f = f_hess(model,v,T,z)
    return d2f[1,2]/(v*d2f[1])
end

function vt_joule_thomson_coefficient(model::EoSModel, v, T,  z=SA[1.])

    d2f = f_hess(model,v,T,z)
    return -(d2f[1,2]-d2f[1]*((T*d2f[2,2]+v*d2f[1,2])/(T*d2f[1,2]+v*d2f[1])))^-1
end

function second_virial_coefficient(model::EoSModel, T,  z=SA[1.])
    TT = promote_type(eltype(z),typeof(T))
    V = 1/sqrt(eps(TT))
    _∂2f = ∂2f(model,V,T,z)
    hessf,gradf,f = _∂2f
    _p,dpdv = p∂p∂v(model,V,T,z)
    return -V^2/(R̄*T)*(_p+V*dpdv)
end

function vt_compressibility_factor(model::EoSModel, v, T,  z=SA[1.])
    p = pressure(model,v,T,z)
    return p*v/(R̄*T)
end

"""
    pip(model::EoSModel,v,T,z=[1.0])

Phase identification parameter `Π`. as described in _1_. If `Π > 1`, then the phase is clasified as a liquid or a liquid-like vapor, being a vapor or vapor-like liquid otherwise.

This identification parameter fails at temperatures and pressures well above the critical point.

Calculated as:
```
Π = v*((∂²p/∂v∂T)/(∂p/∂T) - (∂²p/∂v²)/(∂p/∂v))
```

1.  G. Venkatarathnama, L.R. Oellrich, Identification of the phase of a fluid using partial derivatives of pressure, volume,and temperature without reference to saturation properties: Applications in phase equilibria calculations, Fluid Phase Equilibria 301 (2011) 225–233

"""
function pip(model::EoSModel,v,T,z=SA[1.0])
    _∂2p = ∂2p(model,v,T,z)
    hess_p, grad_p, _ = _∂2p
    Π = v*(hess_p[1,2]/grad_p[2]  - hess_p[1,1]/grad_p[1])
end

export second_virial_coefficient


## Packages and settings.
include("../src/StencilBuilder.jl")

using BenchmarkTools
using LoopVectorization
using .StencilBuilder

float_type = Float32

## CPU dynamics kernel.
macro make_kernel(do_advec::Bool, do_diff::Bool, do_buoy::Bool)
    ex_rhs_l = []
    if do_advec
        push!(ex_rhs_l, :( -gradx(interpz(u) * interpx(w)) ))
        push!(ex_rhs_l, :( -grady(interpz(v) * interpy(w)) ))
        push!(ex_rhs_l, :( -gradz(interpz(w) * interpz(w)) ))
    end
    if do_diff
        push!(ex_rhs_l, :( visc * (gradx(gradx(w))) ))
        push!(ex_rhs_l, :( visc * (grady(grady(w))) ))
        push!(ex_rhs_l, :( visc * (gradz(gradz(w))) ))
    end
    if do_buoy
        push!(ex_rhs_l, :( alpha*interpz(s) ))
    end

    if length(ex_rhs_l) > 0
        ex_rhs = Expr(:call, :+, ex_rhs_l...)

        ex = quote
            function dynamics_w_kernel!(
                    wt, u, v, w, s,
                    visc, alpha,
                    dxi, dyi, dzi, dzhi,
                    is, ie, js, je, ks, ke)

                @fast3d begin
                    @fd (wt, u, v, w, s) wt += $ex_rhs
                end
            end
        end
    else
        ex = quote
            function dynamics_w_kernel!(
                    wt, u, v, w, s,
                    visc, alpha,
                    dxi, dyi, dzi, dzhi,
                    is, ie, js, je, ks, ke)
            end
        end
    end

    return esc(ex)
end

## User input
itot = 384; jtot = 384; ktot = 384;
igc = 1; jgc = 1; kgc = 1;

do_advec = true; do_diff = true; do_buoy = true

icells = itot + 2igc; jcells = jtot + 2jgc; kcells = ktot + 2kgc
is = igc+1; ie = igc+itot; js = jgc+1; je = jgc+jtot; ks = kgc+1; ke = kgc+ktot

wt = rand(float_type, icells, jcells, kcells)
u = rand(float_type, icells, jcells, kcells)
v = rand(float_type, icells, jcells, kcells)
w = rand(float_type, icells, jcells, kcells)
s = rand(float_type, icells, jcells, kcells)
dxi = rand(float_type)
dyi = rand(float_type)
dzi = rand(float_type, kcells)
dzhi = rand(float_type, kcells)
visc = convert(float_type, 1.)
alpha = convert(float_type, 9.81/300)

## Run CPU kernel
@eval @make_kernel($do_advec, $do_diff, $do_buoy)

@btime dynamics_w_kernel!(
    $wt, $u, $v, $w, $s,
    $visc, $alpha,
    $dxi, $dyi, $dzi, $dzhi,
    $is, $ie, $js, $je, $ks, $ke)


using Revise
using AdFem
using PyCall
using LinearAlgebra
using PyPlot
using SparseArrays
using MAT
using ADCMEKit
np = pyimport("numpy")
using PyPlot
using SpecialFunctions

# simulation parameter setup
n = 20
NT = 400
ρ = 0.1 # design variable in α-schemes
m = 5n
h = 1/n
Δt = 10000. /NT

# coordinates
xo = zeros((m+1)*(n+1))
yo = zeros((m+1)*(n+1))
for i = 1:m+1
  for j = 1:n+1
    idx = (j-1)*(m+1)+i
    xo[idx] = (i-1)*h
    yo[idx] = (j-1)*h
  end
end

# Dirichlet boundary condition on three sides, and Neumann boundary condition (traction-free) on the top
bdnode = bcnode("left  | lower | right", m, n, h)

# viscoelasticity η and shear modulus μ

ηf = (x,y)->begin
  if y<=0.25
    return 10000.
  else
    return 1.
  end
end

η = constant(eval_f_on_gauss_pts(ηf, m, n, h))
μ = 0.001 * constant(ones(4m*n))

# linear elasticity matrix
coef = 2μ*η/(η + μ*Δt)
mapH = c->begin
  c * diagm(0=>ones(2))
end
H = map(mapH, coef)

Δt = Δt * ones(NT)

# generalized mass matrix and stiffness matrix
M = 100. * constant(compute_fem_mass_matrix1(m, n, h))
K = compute_fem_stiffness_matrix1(H, m, n, h)
C = spzero((m+1)*(n+1))

C = 0.1 * M + 0.1 * K

# fixed displacement
db = zeros((m+1)*(n+1))
for j = 1:n+1
  idx = (j-1)*(m+1)+1
  if j<=div(n, 4)
    db[idx] = 1.
  else
    db[idx] = (1-(j-div(n, 4))/(3div(n, 4)))
  end
  # y = (j-1)*h
  # db[idx] = y * (1-y) * 0.0
end

# cast the problem into homogeneous Dirichlet problem
idof = ones(Bool, (m+1)*(n+1))
idof[bdnode] .= false
idof = findall(idof)

d0 = db
v0 = zeros((m+1)*(n+1))
# v0[bcnode("left", m, n, h)] .= 1.0

a0 = vector(idof, M[idof, idof]\((-K*db)[idof]), (m+1)*(n+1))

function solver(A, rhs)
  A, _ = fem_impose_Dirichlet_boundary_condition1(A, bdnode, m, n, h)
  rhs = scatter_update(rhs, bdnode, zeros(length(bdnode)))
  A\rhs
end

function antiplane_visco_αscheme(M::Union{SparseTensor, SparseMatrixCSC},
  K::Union{SparseTensor, SparseMatrixCSC},
  d0::Union{Array{Float64, 1}, PyObject},
  v0::Union{Array{Float64, 1}, PyObject},
  a0::Union{Array{Float64, 1}, PyObject},
  Δt::Array{Float64};
  ρ::Float64 = 1.0)
  nt = length(Δt)
  αm = (2ρ-1)/(ρ+1)
  αf = ρ/(1+ρ)
  γ = 1/2-αm+αf
  β = 0.25*(1-αm+αf)^2
  d = length(d0)

  M = isa(M, SparseMatrixCSC) ? constant(M) : M
  K = isa(K, SparseMatrixCSC) ? constant(K) : K
  d0, v0, a0, Δt = convert_to_tensor([d0, v0, a0, Δt], [Float64, Float64, Float64, Float64])

  A = (1-αm)*M + (1-αf)*K*β*Δt[1]^2
  A, _ = fem_impose_Dirichlet_boundary_condition1(A, bdnode, m, n, h)

  function equ(dc, vc, ac, dt, εc, σc, i)
    dn = dc + dt*vc + dt^2/2*(1-2β)*ac
    vn = vc + dt*((1-γ)*ac)

    df = (1-αf)*dn + αf*dc
    vf = (1-αf)*vn + αf*vc
    am = αm*ac

    σ_dt = (1-αf)*dt
    Σ = 2* repeat(μ.*η/(η+μ*σ_dt),1,2) .*εc - repeat(η/(η+μ*σ_dt), 1, 2) .* σc
    Force = compute_strain_energy_term1(Σ, m, n, h)

    rhs = - (M*am + C*vf + K*df) + Force

    rhs = scatter_update(rhs, bdnode, zeros(length(bdnode)))
    A\rhs

  end

  function condition(i, tas...)
    return i<=nt
  end
  function body(i, tas...)
    dc_arr, vc_arr, ac_arr, σc_arr = tas
    dc = read(dc_arr, i)
    vc = read(vc_arr, i)
    ac = read(ac_arr, i)
    σc = read(σc_arr, i)
    εc = eval_strain_on_gauss_pts1(dc, m, n, h)

    y = equ(dc, vc, ac, Δt[i],εc, σc, i)
    dn = dc + Δt[i]*vc + Δt[i]^2/2*((1-2β)*ac+2β*y)
    vn = vc + Δt[i]*((1-γ)*ac+γ*y)

    εn = eval_strain_on_gauss_pts1(dn, m, n, h)
    σn = 2* repeat(μ.*η/(η+μ*Δt[i]),1,2) .* (εn - εc) + repeat(η/(η+μ*Δt[i]), 1, 2) .* σc

    i+1, write(dc_arr, i+1, dn), write(vc_arr, i+1, vn), write(ac_arr, i+1, y), write(σc_arr, i+1, σn)
  end

  dM = TensorArray(nt+1); vM = TensorArray(nt+1); aM = TensorArray(nt+1); σM = TensorArray(nt+1);
  dM = write(dM, 1, d0)
  vM = write(vM, 1, v0)
  aM = write(aM, 1, a0)

  ε0 = eval_strain_on_gauss_pts1(d0, m, n, h)
  σ0 = batch_matmul(H, ε0)
  σM = write(σM, 1, σ0)

  i = constant(1, dtype=Int32)
  _, d, v, a = while_loop(condition, body, [i,dM, vM, aM, σM])
  set_shape(stack(d), (nt+1, length(a0))), set_shape(stack(v), (nt+1, length(a0))), set_shape(stack(a), (nt+1, length(a0)))
end

# d, v, a = αscheme(M, C, K, zeros(NT, (m+1)*(n+1)), d0, v0, a0, Δt; solve = solver)

d, v, a = antiplane_visco_αscheme(M, K, d0, v0, a0, Δt, ρ=ρ)

sess = Session(); init(sess)
d_, v_, a_ = run(sess, [d, v, a])

# close("all")
# for (k,tid) in enumerate(LinRange{Int64}(1, NT+1, 5))
#   t = (tid-1)*Δt[1]
#   plot((d_[tid, :]+db)[(1:m+1)],"C$k-", label="$t", markersize=3)
# end
# legend()

# close("all")
# for tid in LinRange{Int64}(1, NT+1, 5)
#   plot((a_[tid, :])[1:m+1], label="$tid")
# end
# legend()

# plot(d_[end,1:m+1])
# plot(d_[div(NT+1,2),1:m+1])
# plot(d_[div(NT+1,3),1:m+1])

# pcolormesh(reshape(d_[end,:], m+1, n+1)')
figure()
pl, = plot([], [], "o-", markersize = 3)
t = title("0")
xi = (0:m)*h
xlim(-h, (m+1)*h)
xlabel("Distance")
ylim(-0.1, 1.1)
ylabel("Displacement")
tight_layout()
function update(i)
  pl.set_data(xi[:], d_[i,1:m+1])
  t.set_text("time = $(i*Δt[1])")
end
p = animate(update, [1:1:20;25:5:NT+1])
saveanim(p, "displacement.gif")

figure()
pl, = plot([], [], "o-", markersize = 3)
t = title("time = 0")
xi = (0:m)*h
xlim(-h, (m+1)*h)
xlabel("Distance")
ylim(-0.0001, 0.0005)
ylabel("Velocity")
tight_layout()
function update(i)
  pl.set_data(xi[:], v_[i,1:m+1])
  t.set_text("time = $(i*Δt[1])")
end
p = animate(update, [1:1:20;25:5:NT+1])
saveanim(p, "velocity.gif")

figure()
pl, = plot([], [], "o-", markersize = 3)
t = title("time = 0")
xi = (0:m)*h
xlim(-h, (m+1)*h)
xlabel("Distance")
ylim(-0.001, 0.001)
ylabel("Strain Rate")
tight_layout()
function update(i)
  pl.set_data(xi[1:end-1], (v_[i,2:m+1]-v_[i,1:m])/h)
  t.set_text("time = $(i*Δt[1])")
end
p = animate(update, [1:1:20;25:5:NT+1])
saveanim(p, "strain_rate.gif")


# Find a way to run these example in isolation
using Finch

init_finch("poisson1d")

useLog("poisson1dlog")

domain(1)

functionSpace(order=3)
mesh(LINEMESH, elsperdim=20)

u = variable("u")
testSymbol("v")
boundary(u, 1, DIRICHLET, 0)

coefficient("f", "-100*pi*pi*sin(10*pi*x)*sin(pi*x) -
             pi*pi*sin(10*pi*x)*sin(pi*x) + 20*pi*pi*cos(10*pi*x)*cos(pi*x)")
weakForm(u, "-grad(u)*grad(v) - f*v")

import Finch: grid_data, mesh_data, refel
dofs_per_node = 1

N1 = size(grid_data.allnodes,2);
Nn = dofs_per_node * N1;
Np = refel.Np;
nel = mesh_data.nel;

rhsvec = zeros(Nn);
lhsmatI = zeros(Int, nel*dofs_per_node*Np*dofs_per_node*Np);
lhsmatJ = zeros(Int, nel*dofs_per_node*Np*dofs_per_node*Np);
lhsmatV = zeros(nel*dofs_per_node*Np*dofs_per_node*Np);
allocated_vecs = [rhsvec, lhsmatI, lhsmatJ, lhsmatV];

Finch.CGSolver.assemble(u, Finch.bilinears[u.index], Finch.linears[u.index], allocated_vecs)

print("END")

#SUITE["poisson1d-assemble"] = @benchmarkable Finch.CGSolver.assemble($u,
                         #$Finch.bilinears[u.index], $Finch.linears[u.index])


########################################################################################################
## Details
# This file will create the structure for the UQ inference problems.
# This package is designed to be used with a family of julia packages,
#           - Potential.jl (*)
#               - Creates potential types, has abstract 'Potential' type (*), let p <: 'Potential'
#               - Assume mutable struct Potential{T<:AbstractFloat}
#                            x       # Potential trainable parameters
#                            p       # Potential nontrainable parameters
#                        end
#               - Assume mutable struct SNAP{T<:AbstractFloat}
#                            A
#                            b
#                            β
#                        end
#           - PotentialLearning.jl
#               - Fits data-driven potential
#               - We assume that potential learning updates parameters p.x
#           - PotentialUQ.jl (this package)
#               - Defines distribution over potential parameters.
#               - Initially must provide distribution over p.x (assume user provides distribution, in form of log likelihood)
#               - In case of SNAP, potential is predetermined.
#               - Allow users to modify distribution with priors by adding to log-likelihood.
#
#########################################################################################################
module PotentialUQ

using Base: NamedTuple
using LinearAlgebra: tr, logdet, transpose, Symmetric, diag, I, pinv
using Statistics: cov, std, mean
using SciMLBase, GalacticOptim, Optim
using TransformVariables
import Turing
using Distributions

include("distribution_types.jl")
include("distributions.jl")
include("utils.jl")
include("map.jl")
include("sampling.jl")
end


using Combinatorics
function henry_q(Pk,l,n,kmin,kmax; r)
    function I_m(m, smin, smax, kmin, kmax, pmin, pmax)
        kturn = (kmin^smin * kmax^(-smax) * abs(pmax/pmin))^(1/(smin-smax))
        lo = kmin^(-smin) * pmin * (kturn^(m+smin) - kmin^(m+smin))/(m+smin)
        hi = kmax^(-smax) * pmax * (kmax^(m+smax) - kmin^(m+smax))/(m+smax)
        return lo + hi
    end
    pmin = Pk(kmin)
    pmax = Pk(kmax)

    kvals = exp10.(range(log10(kmin), stop=log10(kmax), length=1024))
    if Pk(kvals[1]) > 0
        smin = (log(Pk(kvals[20])) - log(Pk(kvals[1]))) / (log(kvals[20]) - log(kvals[1]))
    else
        smin = (log(abs(Pk(kvals[20]))) - log(abs(Pk(kvals[1])))) / (log(kvals[20]) - log(kvals[1]))
    end
    if Pk(kvals[end]) > 0
        smax = (log(Pk(kvals[end])) - log(Pk(kvals[end-19]))) / (log(kvals[end]) - log(kvals[end-19]))
    else
        smax = (log(abs(Pk(kvals[end]))) - log(abs(Pk(kvals[end-19])))) / (log(kvals[end]) - log(kvals[end-19]))
    end

    Ia = I_m(3+l+n, smin, smax, kmin, kmax, pmin, pmax)
    Ib = I_m(1+n, smin, smax, kmin, kmax, pmin, pmax)
    G = log(kmax/kmin)
    qnu = 1 - l/2 + (1/2G)*log((r^(l+2)/Int64(doublefactorial(2l+1)))*(Ia/Ib))
    #@show (r^(l+1)/Int64(doublefactorial(2l+1)))*(Ia/Ib)
    qmin = max(smax+4-1+n,-l+.5)
    qmax = 3+smin+n #min(3+n1+n,2)

    if qnu > qmax
        q = qmax
    elseif qnu < qmin
        q = qmin
    else
        q= qnu
    end

    return q#+.5
end

function best_q(Pk,ℓ,n, kmin, kmax)
    kvals = exp10.(range(log10(kmin), stop=log10(kmax), length=1024))
    #if Pk(kvals[1]) > 0
    #    smin = (log(Pk(kvals[20])) - log(Pk(kvals[1]))) / (log(kvals[20]) - log(kvals[1]))
    #else
        smin = (log(abs(Pk(kvals[20]))) - log(abs(Pk(kvals[1])))) / (log(kvals[20]) - log(kvals[1]))
    #end
    #if Pk(kvals[end]) > 0
    #    smax = (log(Pk(kvals[end])) - log(Pk(kvals[end-19]))) / (log(kvals[end]) - log(kvals[end-19]))
    #else
        smax = (log(abs(Pk(kvals[end]))) - log(abs(Pk(kvals[end-19])))) / (log(kvals[end]) - log(kvals[end-19]))
    #end

    #@show smin
    #@show smax
    n1 = smin
    n2 = smax+4
    qmin = max(n2-1+n,-ℓ+.5)
    qmax = 3+n1+n#min(3+n1+n,2)
    qbest = n - (smin+smax)/2

    #@show qbest
    #@show qmin
    #@show qmax
    if qbest > qmax
        q = qmax
    elseif qbest < qmin
        q = qmin
    else
        q=qbest
    end

    if sign(Pk(kvals[1])) != sign(Pk(kvals[end]))
        #println("Changing Sign correction")
        if Pk(kvals[1]) > 0
            q -= .5
        else
            q += .5
        end
    end
    #if (smin < -1e-10) && (smax < 0)
    if (smin < -1e-4) && (smax < 0)
        #println("Montonic Decreasing correction")
        if smin > -.3
            #println("Almost Flat correction")
            q -= 1.1
        else
            #q -= .2
            q -= .4
        end
    end

    if smax < -6.5
        #println("Rapid Decrease correction")
        q -= 1
    #elseif (smax-smin) < -5.5
    #    println("Large Second Derivative correction")
    #    q += .4
    end

    if (smin > 3.5) && (smax < -3.5)
        #println("Both Slopes Large correction")
        q += 3
    end

    if (smin < .5) && (smax < -3.5) && (smin > 0)
        #println("Flat to Quick Fall correction")
        q -= 1.2
    end

    #@show q
    return q
end


# This file is to get the Cl and Cd of Blade Element

#the excel interpolation method
function cla(α,ma)
  # NACA0012
  x = α/π*180
  if ma<0.1
    cla = 2e-17*x^6 + 1e-10*x^5 - 8e-13*x^4 - 4e-06*x*3 + 1e-08*x^2 + 0.0351*x + 1e-06
  elseif ma<0.2
    cla = 9e-18*x^6 + 1e-10*x^5 - 5e-13*x^4 - 4e-06*x^3 + 6e-09*x^2 + 0.0351*x - 3e-06
  elseif ma<0.3
    cla = 2e-17*x^6 + 1e-10*x^5 - 8e-13*x^4 - 4e-06*x^3 + 1e-08*x^2 + 0.0351*x + 1e-06
  elseif ma<0.4
    cla = -1e-15*x^6 + 1e-10*x^5 + 6e-11*x^4 - 4e-06*x^3 - 8e-07*x^2 + 0.0349*x + 0.0025
  elseif ma<0.5
    cla = -3e-17*x^6 + 1e-10*x^5 + 1e-12*x^4 - 4e-06*x^3 - 3e-09*x^2 + 0.0349*x - 0.0001
  elseif ma<0.6
    cla = -6e-17*x^6 + 1e-10*x^5 + 2e-12*x^4 - 4e-06*x^3 - 9e-09*x^2 + 0.0347*x - 0.0002
  elseif ma<0.7
    cla = -6e-17*x^6 + 1e-10*x^5 + 2e-12*x^4 - 4e-06*x^3 - 9e-09*x^2 + 0.0345*x - 0.0002
  elseif ma<0.8
    cla = -6e-17*x^6 + 1e-10*x^5 + 2e-12*x^4 - 4e-06*x^3 - 9e-09*x^2 + 0.0345*x - 0.0002
  elseif ma<0.9
    cla = 2e-16*x^6 + 9e-11*x^5 - 7e-12*x^4 - 4e-06*x^3 + 4e-08*x^2 + 0.0313*x + 0.0004
  elseif ma<1.0
    cla = 1e-16*x^6 + 9e-11*x^5 - 5e-12*x^4 - 4e-06*x^3 - 4e-09*x^2 + 0.03*x + 0.0007
  else
    cla = 1e-16*x^6 + 9e-11*x^5 - 5e-12*x^4 - 4e-06*x^3 - 7e-09*x^2 + 0.03*x + 0.0007
  end
  return cla
end

#the simple interpolation method
function liftcoffi(α,Ma)
  if abs(α)>(20.0*π/180.0)
    Cl = 0.0
  else
    Cl = 5.73*α+0.085
  end
  return Cl
end

function dragcoffi(α,Ma)
  if abs(α)<(3.0*π/180.0)
    Cd = 0.0076
  elseif abs(α)<(5.0*π/180.0)
    Cd = 0.008
  elseif abs(α)<(6.0*π/180.0)
    Cd = 0.0086
  elseif abs(α)<(7.0*π/180.0)
    Cd = 0.0093
  elseif abs(α)<(8.0*π/180.0)
    Cd = 0.0103
  elseif abs(α)<(9.0*π/180.0)
    Cd = 0.012
  elseif abs(α)<(10.0*π/180.0)
    Cd = 0.0146
  elseif abs(α)<(11.0*π/180.0)
    Cd = 0.0177
  elseif abs(α)<(12.0*π/180.0)
    Cd = 0.0215
  elseif abs(α)<(13.0*π/180.0)
    Cd = 0.0266
  elseif abs(α)<(14.0*π/180.0)
    Cd = 0.055
  elseif abs(α)<(15.0*π/180.0)
    Cd = 0.125
  elseif abs(α)<(16.0*π/180.0)
    Cd = 0.21
  else
    Cd = 0.5
  end

  Cd = Cd*1.05  #The real blade correction
  return Cd
end

function clcdget(α_aero,vall_r)
  Ma = Array{Float64}(NR,Nbe)
  for k in 1:NR #take Mach Number into consideration
    for i in 1:Nbe
      Ma[k,i] = norm(vall_r[k,i])/v_sound
    end
  end
  clift = Array{Float64}(NR,Nbe)
  cdrag = Array{Float64}(NR,Nbe)
  for k in 1:NR #solve the local clcd
    for i in 1:Nbe
      clift[k,i] = liftcoffi(α_aero[k,i],Ma[k,i])
      cdrag[k,i] = dragcoffi(α_aero[k,i],Ma[k,i])
    end
  end
  return clift,cdrag,Ma
end


# -----------------------------------------------------------
#
# An example of interpolating a given function in 1D and 2D
#
# -----------------------------------------------------------

using GalerkinSparseGrids
using LinearAlgebra

D = 1
# Random function to interpolate:
truesoln = x -> 1.152 * norm(x)^2 * sin(dot([1.214], x))^3
k_max = 5
n_max = 5

println("Interpolation in ", D, "D.")
println("Going to k_max = ", k_max, ", n_max = ", n_max, ":")

for k_used in 1:k_max
    for n_used in 1:n_max
        # By default, sparse grids are used. Change scheme to "full"
        # to compare the different accuraces and costs of the
        # interpolations
        dict = coeffs_DG(D, k_used, n_used, truesoln)
        err = mcerr(x->reconstruct_DG(dict, [x...]), truesoln, D)
        println("(k = ", k_used, ", n = ", n_used, ") : err = ", err)
    end
end

D = 2
# Random function to interpolate:
truesoln = x -> 1.152 * norm(x)^2 * sin(dot([1.214, -0.5], x))^3
k_max = 5
n_max = 5

println("Interpolation in ", D, "D.")
println("Going to k_max = ", k_max, ", n_max = ", n_max, ":")

for k_used in 1:k_max
    for n_used in 1:n_max
        # By default, sparse grids are used. Change scheme to "full"
        # to compare the different accuraces and costs of the
        # interpolations
        dict = coeffs_DG(D, k_used, n_used, truesoln)
        err = mcerr(x->reconstruct_DG(dict, [x...]), truesoln, D)
        println("(k = ", k_used, ", n = ", n_used, ") : err = ", err)
    end
end


using Ipopt
using MathProgBase
using SpecialFunctions

mutable struct GenTDist <: ContinuousUnivariateDistribution
    mu::Float64
    sigma::Float64
    t::TDist
end

GenTDist( mu::Float64, sigma::Float64, nu::Float64 ) = GenTDist( mu, sigma, TDist( nu ) )

Base.rand( t::GenTDist ) = t.sigma .* rand( t.t ) .+ t.mu

Base.rand( t::GenTDist, n::Int ) = t.sigma .* rand( t.t, n ) .+ t.mu

Distributions.pdf( t::GenTDist, x::Float64 ) = pdf( t.t, (x - t.mu)/t.sigma)/t.sigma

Distributions.cdf( t::GenTDist, x::Float64 ) = cdf( t.t, (x - t.mu)/t.sigma)

Distributions.mean( t::GenTDist ) = t.mu

Distributions.std( t::GenTDist ) = t.sigma * sqrt(t.t.ν/(t.t.ν-2))

randomparameters( ::Type{GenTDist} ) = [randn(), rand( Exponential() ), 1 + rand(Exponential())]

randomparameters( ::Type{GenTDist}, n::Int ) =
    [randn( 1, n ); rand( Exponential(), 1, n ); 1 .+ rand( Exponential(), 1, n )]

mutable struct GenTDistOptimizer <: MathProgBase.AbstractNLPEvaluator
    y::Vector{Float64}
    w::Vector{Float64}
    debug::Int
end

GenTDistOptimizer( y::Vector, w::Vector ) = GenTDistOptimizer( y, w, 0 )

function MathProgBase.initialize( ::GenTDistOptimizer, requested_features::Vector{Symbol} )
    unimplemented = setdiff( requested_features, [:Grad] )
    if !isempty( unimplemented )
        error( "The following features aren't implemented: " * join( string.(unimplemented), ", " ) )
    end
end

MathProgBase.features_available( ::GenTDistOptimizer ) = [:Grad]

function MathProgBase.eval_f( t::GenTDistOptimizer, x; debug=0 )
    if t.debug > 0
        println( "eval_f called with $x" )
    end
    (mu, sigma, nu) = x

    sigma <= 0.0 && return -Inf

    y = t.y
    w = t.w
    n = sum(w)
    normalysq = ((y .- mu)/sigma).^2
    constant = n*(lgamma((nu+1)/2) - log(nu*pi)/2 - lgamma(nu/2) - log(sigma))
    return constant - (nu+1)/2*sum(w .* log.(1 .+ normalysq/nu))
end

function MathProgBase.eval_grad_f( t::GenTDistOptimizer, g, x )
    (mu, sigma, nu) = x
    y = t.y
    w = t.w
    n = sum(w)
    normaly = (y .- mu)/sigma
    normalysq = normaly.^2
    g[1] = (nu+1)*sum(w .* (y .- mu)./(nu*sigma^2 .+ (y .- mu).^2))
    g[2] = -n/sigma + (nu+1)*sum(w .* normalysq ./ (sigma*(nu .+ normalysq)))
    g[3] = n*(digamma((nu+1)/2)/2 - 1/(2*nu) - digamma(nu/2)/2)
    g[3] += -sum(w .* log.(1 .+ normalysq/nu))/2 + (nu+1)/(2*nu)*sum(w .* normalysq ./ (nu .+ normalysq))
    if t.debug > 0
        println( "grad_f called with $x, returning $g" )
    end
end

function fit_mle!(
    ::Type{GenTDist},
    parameters::AbstractVector{Out},
    x::Vector{Out},
    w::Vector{Calc},
    scratch::Dict{Symbol,Any};
    max_iter::Int = 3000,
    print_level::Int = 0,
) where {Calc,Out}
    t = GenTDistOptimizer( x, w )

    solver = IpoptSolver(print_level=print_level, max_iter=max_iter)

    model = MathProgBase.NonlinearModel(solver)
    MathProgBase.loadproblem!(model, 3, 0, [-Inf; inequalityconstraintvector(GenTDist)], fill(Inf,3), Float64[], Float64[], :Max, t)
    MathProgBase.setwarmstart!( model, parameters )
    MathProgBase.optimize!(model)

    status = MathProgBase.status(model)
    if status != :Optimal
        @warn( "Non-optimal optimization result: $status" )
    end
    parameters[:] = MathProgBase.getsolution(model)
end


# Provides dynamic risk measures for SDEs

"""
    dynamicRM(problem::SDEProblem,
                eval = x-> sum(x),
                RM=1.0,
                dt = 0.01,
                MCs = 200,
                iter = 100,
                u0::Flux.Chain=Flux.Chain() )
implements the unique dynamic risk measure for the SDE ``problem``associated to the nonlinear term
```math
s_\\rho(t,x) \\lVert \\sigma(t,x) \\nabla u(t,x) \\rVert.
```
For more details on dynamic risk measures for SDEs, its relationship to BSDEs and parabolic PDEs as well as discrete time dynamic risk measures
see the Section [BSDE]().
"""
function dynamicRM(problem::SDEProblem,eval = x-> sum(x), RM=1.0, dt = 0.01, MCs = 200, iter = 100,
                    u0::Flux.Chain=Flux.Chain() )

    time_steps = div(problem.tspan[2]-problem.tspan[1],dt, RoundNearest)

    f(X,u,σᵀ∇u,p,t) = RM*abs.(σᵀ∇u)[1]                 # nonlinear riskMeasure part

    BSDE = TerminalPDEProblem(eval, f, problem.f, problem.g, problem.u0, problem.tspan);

    d = length(problem.u0)
    hls  = 10 + d #hide layer size
    opt = Flux.ADAM(0.001)

    if (u0 == Flux.Chain())
        u0 = Flux.Chain(Dense(d,hls,relu),
                Dense(hls,hls,relu),
                Dense(hls,hls,relu),
                Dense(hls,1))
    end

    σᵀ∇u = [Flux.Chain(Dense(d,hls,relu),
        Dense(hls,hls,relu),
        Dense(hls,d)) for i in 1:time_steps]

    pdealg = NNPDEHan(u0,σᵀ∇u;opt=Flux.ADAM(0.1))
    ans = solve(BSDE, pdealg, verbose=true, maxiters=iter, trajectories=MCs, dt=dt)

    return(ans)
end

# Provide a jump diffusion version of dynamicRM


using ModelingToolkit, Test
using ModelingToolkit: value
using SymbolicUtils: <ₑ
@parameters α β δ
expr = (((1 / β - 1) + δ) / α) ^ (1 / (α - 1))
ref = sort([β, δ, α], lt = <ₑ)
sol = sort(Num.(ModelingToolkit.get_variables(expr)), lt = <ₑ)
@test all(x->x isa Num, sol[i] == ref[i] for i in 1:3)
@test all(simplify∘value, sol[i] == ref[i] for i in 1:3)

@parameters γ
s = α => γ
expr = (((1 / β - 1) + δ) / α) ^ (1 / (α - 1))
sol = ModelingToolkit.substitute(expr, s)
new = (((1 / β - 1) + δ) / γ) ^ (1 / (γ - 1))
@test iszero(sol - new)

# Continuous
using ModelingToolkit: isdifferential, isdifference, vars, difference_vars, collect_difference_variables, collect_differential_variables, collect_ivs
@variables t u(t) y(t)
D = Differential(t)
eq = D(y) ~ u
v = vars(eq)
@test v == Set([D(y), u])

ov = collect_differential_variables(eq)
@test ov == Set(Any[y])

aov = ModelingToolkit.collect_applied_operators(eq, Differential)
@test aov == Set(Any[D(y)])

ts = collect_ivs([eq])
@test ts == Set([t])

# Discrete
z = Difference(t; dt=1, update=true)
eq = z(y) ~ u
v = difference_vars(eq)
@test v == Set([z(y), u])

ov = collect_difference_variables(eq)
@test ov == Set(Any[y])

aov = ModelingToolkit.collect_applied_operators(eq, Difference)
@test aov == Set(Any[z(y)])

ts = collect_ivs([eq], Difference)
@test ts == Set([t])

@variables t
function UnitDelay(dt; name)
    @variables u(t)=0.0 [input=true] y(t)=0.0 [output=true]
    z = Difference(t; dt=dt, update=true)
    eqs = [
        z(y) ~ u
    ]
    DiscreteSystem(eqs, t, name=name)
end

dt = 0.1
@named int = UnitDelay(dt)
collect_difference_variables(int) == Set(Any[@nonamespace(int.y)])


todisplay = VERSION <= VersionNumber("1.2") ? true : false

using CAGD
using SparseArrays
using LinearAlgebraicRepresentation
Lar = LinearAlgebraicRepresentation

if todisplay
	using ViewerGL
	GL = ViewerGL
	include("../examples/3d/views.jl")
end

V = [
    0.0 +1  0  0  0
    0.0  0  0 +1 -1
    0.0  0 +1  0  0
];

W = [
     0.0  0.0  0.0
    +1.0  0.0  0.0
     0.0  0.0 +1.0
     0.0 +1.0  0.0
     0.0 -1.0  0.0
];

EV = SparseArrays.sparse(Int8.([
    -1 +1  0  0  0
    -1  0 +1  0  0
    -1  0  0 +1  0
    -1  0  0  0 +1
     0 -1 +1  0  0
     0 -1  0 +1  0
     0 -1  0  0 +1
     0  0 -1 +1  0
     0  0 -1  0 +1
]));

FE = SparseArrays.sparse(Int8.([
    +1 -1  0  0 +1   0  0  0  0
    +1  0 -1  0  0  +1  0  0  0
    +1  0  0 -1  0   0 +1  0  0
     0 +1 -1  0  0   0  0 +1  0
     0 +1  0 -1  0   0  0  0 +1
     0  0  0  0 +1  -1  0 +1  0
     0  0  0  0 +1   0 -1  0 +1
]));

CF = SparseArrays.sparse(Int8.([
    -1 +1  0 -1  0 +1  0
    +1  0 -1  0 +1  0 -1
     0 -1 +1 +1 -1 -1 +1
]));

model = CAGD.Model(V)
CAGD.addModelCells!(model, 1, EV)
CAGD.addModelCells!(model, 2, FE)
CAGD.addModelCells!(model, 3, CF)

atol = 1e-6;
if todisplay  displayModel(model)  end

split_model = CAGD.facesplitting(model, atol = atol)

if todisplay  displayModel(split_model)  end

congr_model = CAGD.mergeModelVertices(split_model, err=atol, signed_merge=true)
#split_model = CAGD.facesplitting(congr_model, atol = atol)
#congr_model = CAGD.mergeModelVertices(split_model, err=atol, signed_merge=true)

if todisplay  displayModel(congr_model)  end

gift_model = deepcopy(congr_model)
FCcycles, bicon_comps = CAGD.tgw(congr_model, 3, atol = atol)
CAGD.addModelCells!(gift_model, 3, convert(Lar.ChainOp, FC'))

if todisplay  viewExplode(gift_model, complete = false)  end

# Since first face is inner, then retrieving the external is needed

externals = [bc[CAGD.getExternalBoundary(congr_model, 3, FCcycles, bc)] for bc in bicon_comps]
internals = [setdiff(bicon_comps[i], externals[i]) for i = 1 : length(externals)]
assembled_model = deepcopy(congr_model)

FC = FCcycles[:, internals...]
JK = ones(Int8, length(externals))
FC = [FCcycles * SparseArrays.sparse(externals, JK, JK, size(FCcycles, 2), 1) FC]
CAGD.addModelCells!(assembled_model, 3, convert(Lar.ChainOp, FC'))

if todisplay  viewExplode(assembled_model)  end


@testset "Transforms" begin
  for c in [Composition(1,2,3),
            Composition(0,2,3),
            Composition(1,2,0)]
    @test alrinv(alr(c)) == c
    @test clrinv(clr(c)) == c
    @test ilrinv(ilr(c)) == c
  end
end


@testset "Basics" begin

    @testset "constructors" begin
        A = round.(Int, max.(0.0, 5 * rand(6, 10) .- 1))
        f = Polynomial(rand(10), A)
        @test typeof(f) <: Polynomial{Float64}

        @test_throws AssertionError Polynomial(rand(9), A)

        @polyvar x y
        f2 = Polynomial(2x^2+4y^2+3x*y+1)
        @test exponents(f2) == [0 2 1 0; 0 0 1 2]
        @test nvariables(f2) == 2
        @test coefficients(f2) == [1, 2, 3, 4]
        @test coefficienttype(f2) == Int64
        @test permutation(f2) == [4, 1, 2, 3]
        f2_2 = Polynomial(2x^2+4y^2+3x*y+1)
        @test f2 == f2_2

        @polyvar x y a b
        f = Polynomial(2x+b*y+a, parameters=[a, b])
        @test f isa Polynomial
        @test nvariables(f) == 2
        @test parameters(f) == [:a, :b]

        @test_throws ErrorException Polynomial(2x+b*y+a, variables=[x], parameters=[a, b])
    end

    @testset "evaluation" begin
        @polyvar x y
        f2 = 2x^2+4y^2+3x*y^4+1
        g = Polynomial(f2)
        w = rand(2)

        @test abs(SP.evaluate(g, w) - f2(x => w[1], y => w[2])) < 1e-13

        f2_x = MP.differentiate(2x^2+4y^2+3x*y^4+1, x)
        f2_y = MP.differentiate(2x^2+4y^2+3x*y^4+1, y)
        @test norm(SP.gradient(g, w) - [f2_x(x => w[1], y => w[2]), f2_y(x => w[1], y => w[2])]) < 1e-13

        u = zeros(2)
        SP.gradient!(u, g, w)
        @test u ≈ SP.gradient(g, w)

        @test all(SP.evaluate_and_gradient(g, w) .≈ (g(w), u))

        @test all(SP.evaluate_and_gradient!(u, g, w) .≈ g(w))
        @test u ≈ SP.gradient(g, w)
    end

    @testset "Evaluation with parameters" begin
        @polyvar x y z a b
        f = 2x^2+4a*y^2+(a^2+b^4+a+b+1)*3x*y^4+1+a*z+b-2x*b^6
        g = Polynomial(f, parameters=[a, b])

        w = [5, 3, -4]
        p = [-2, 7]
        f_wp = f([x, y, z] => w, [a, b] => p)
        ∇_wp = map(fi -> fi([x, y, z] => w, [a, b] => p), MP.differentiate(f, [x, y, z]))
        ∇p_wp = map(fi -> fi([x, y, z] => w, [a, b] => p), MP.differentiate(f, [a, b]))

        @test g(w, p) == f_wp
        @test gradient(g, w,p) == ∇_wp
        @test evaluate_and_gradient(g, w, p) == (f_wp, ∇_wp)
        @test differentiate_parameters(g, w, p) == ∇p_wp

        u = zeros(Int, 3)
        gradient!(u, g, w, p)
        @test u == ∇_wp
        u .= 0
        @test evaluate_and_gradient!(u, g, w, p) == f_wp
        @test u == ∇_wp

        v = zeros(Int, 2)
        differentiate_parameters!(v, g, w, p)
        @test v == ∇p_wp

        f = Polynomial(a + b + x + y, parameters=[a, b])
        @test differentiate_parameters(f, [0, 0], [1, 1]) == [1, 1]

        f = Polynomial(x*y + x + y, parameters=[a, b])
        @test differentiate_parameters(f, [0, 0], [1, 1]) == [0, 0]
    end

    @testset "show" begin
        @polyvar x[0:9]
        @test string(Polynomial(x[1])) == "x₀"
        @test string(Polynomial(sum((-1)^i * x[i]^i for i=1:length(x)))) ==
            "-x₀ + x₁² - x₂³ + x₃⁴ - x₄⁵ + x₅⁶ - x₆⁷ + x₇⁸ - x₈⁹ + x₉¹⁰"
        @test string(Polynomial(2im*x[1] - x[2])) == "(0 + 2im)x₀ - x₁"
        @test sprint(show, Polynomial((2+0im)*x[1] - x[2])) == "2x₀ - x₁"
        @test_nowarn string(typeof(Polynomial(x[1]^2+x[4]^2)))

        @polyvar x a b
        @test sprint(show, Polynomial(x+a^2+b, parameters=[a, b])) == "a² + b + x"
        @test sprint(show, Polynomial(zero(x))) == "0"
    end

    @testset "helpers" begin
        x = rand()
        z = rand(Complex{Float64})
        for k = 4:15
            @test abs(SP.pow(x, k) - x^k) < 1e-13
            @test abs(SP.pow(z, k) - z^k) < 1e-13
        end
    end

    @testset "system constructor" begin
        @polyvar x y
        f1 = x^2+y^2
        f2 = 2x^2+4y^2+3x*y^4+1
        g1 = Polynomial(f1)
        g2 = Polynomial(f2)
        @test_deprecated SP.system(g1, g2)
        @test PolynomialSystem(g1, g2) isa PolynomialSystem{2, 2}
        @test PolynomialSystem(g1, g2, g2) isa PolynomialSystem{3, 2}
        @test PolynomialSystem([f1, f2, y, x]) isa PolynomialSystem{4, 2}
        @test length(SP.PolynomialSystem([f1, f2, y, x])) == 4

        @test_nowarn sprint(show, PolynomialSystem([f1, f2, y, x]))
    end

    @testset "System evaluation and Jacobian" begin
        @polyvar x y
        f1 = x^2+y^2
        f2 = 2x^2+4y^2+3x*y^4+1
        g1 = Polynomial(f1)
        g2 = Polynomial(f2)

        G = PolynomialSystem(g1, g2)

        w = rand(2)
        @test [evaluate(g1, w), evaluate(g2, w)] ≈ evaluate(G, w)
        @test [g1(w), g2(w)] ≈ evaluate(G, w)

        w = SVector{2}(w)
        @test evaluate(G, w) isa SVector{2}
        @test [evaluate(g1, w), evaluate(g2, w)] ≈ evaluate(G, w)

        @test jacobian(G, [2, 3]) == [4 6; 251 672]
        U = zeros(Int, 2, 2)
        @test jacobian!(U, G, [2, 3]) == [4 6; 251 672] == U
        u = zeros(Int, 2)
        U .= 0
        evaluate_and_jacobian!(u, U, G, [2, 3])
        @test U == jacobian(G, [2, 3])
        @test u == G([2, 3])

        @polyvar x y a b

        F = PolynomialSystem(x^2+x*y+a+a*x, x+y+b, parameters=[a, b])

        @test F([0,0], [1, 2]) == [1, 2]
        @test jacobian(F, [0, 0], [3, 2]) == [3 0; 1 1]

        @test parameters(F) == [:a, :b]
        @test nparameters(F) == 2
        @test variables(F) == [:x, :y]
        @test nvariables(F) == 2
        @test F isa PolynomialSystem{2, 2, 2}

        @polyvar x y z a b
        f = 2x^2+4a*y^2+(a^2+b^4+a+b+1)*3x*y^4+1+a*z+b-2x*b^6
        g = (x^2+z)*(x+b*y+a)

        F = PolynomialSystem(f, g, parameters=[a, b])
        differentiate_parameters(F, [2, 3, -3], [4, -2]) == [4407 -14297; 1 3]
        differentiate_parameters(F, [2, 3, -3], @SVector([4, -2])) == [4407 -14297; 1 3]
        @test differentiate_parameters(F, [2, 3, -3], @SVector([4, -2])) isa SMatrix{2, 2, Int, 4}
        U = zeros(Int, 2, 2)
        differentiate_parameters!(U, F, [2, 3, -3], [4, -2])
        @test U == [4407 -14297; 1 3]
    end

    @testset "Hessian" begin
        @polyvar x y z

        w = [2, 3, 5]

        f = x^4*y+x*y*z+y+x+z
        g = Polynomial(f)
        grad_g = Polynomial.(MP.differentiate(f, [x, y, z]))
        hess_w = vcat(map(grad_g) do g_i
            gradient(g_i, w)'
        end...)

        @test hessian(g, w) == hess_w
        @test gradient_and_hessian(g, w) == (gradient(g, w), hess_w)

        U = zeros(Int, 3, 3)
        @test hessian!(U, g, w) == hess_w

        p = [x^4*y+x*y*z+y+x+z, (x^3+y^2)*z]
        P = PolynomialSystem(p)

        u = zeros(2, 3)
        U = zeros(2, 3, 3)

        hessian!(U, P, w)
        @test U[1,:,:] == hess_w
        jacobian_and_hessian!(u, U, P, w)
        @test u[1,:] == gradient(g, w)
        @test U[1,:,:] == hess_w
    end

    @testset "foreach" begin
        @polyvar x y
        g = [Polynomial(x^2+y^2), Polynomial(2x^2+4y^2+3x*y^4+1)]
        G = PolynomialSystem(g...)
        i = 1
        foreach(G) do gi
            @test exponents(gi) == exponents(g[i])
            i += 1
        end
    end

    @testset "scale coefficients" begin
        @polyvar x y
        f = Polynomial(x^2+y^2)
        scale_coefficients!(f, 2)
        @test coefficients(f) == [2, 2]

        g1 = Polynomial(x^2+y^2)
        g2 = Polynomial(2x^2+4y^2+3x*y^4+1)
        G = PolynomialSystem(g1, g2)
        w = rand(2)
        x1 = evaluate(G, w)
        scale_coefficients!(G, [-2, 3])
        x2 = evaluate(G, w)
        @test x2 ≈ (-2, 3) .* x1
    end
end


using Test, Random, LSHFunctions

include(joinpath("..", "utils.jl"))

#==================
Tests
==================#
@testset "MinHash tests" begin
    Random.seed!(RANDOM_SEED)

    @testset "Can construct a MinHash hash function" begin
        symbol_collections = [
            collect(1:100),
            Set(1:30),
            ["a", "a", "b", "c"]
        ]
        n_hashes = [1, 8, 4]

        for (symbols, nh) in zip(symbol_collections, n_hashes)
            hashfn = MinHash(nh; symbols=symbols)

            @test isa(hashfn, MinHash{eltype(symbols)})
            @test LSHFunctions.n_hashes(hashfn) == nh
            @test hashtype(hashfn) == UInt32
        end
    end

    @testset "MinHash range is 1:length(symbols)" begin
        symbols = collect(1:10)
        hashfn = MinHash(1000; symbols = symbols)

        # Verify that all hashes computed by hashfn are in the appropriate range
        has_correct_output_types = true
        has_correct_output_ranges = true

        for ii = 1:100
            if !(has_correct_output_types && has_correct_output_ranges)
                break
            end

            n = mod(rand(Int32), length(symbols)) + 1
            x = shuffle(symbols)[1:n]
            hashes = hashfn(x)

            has_correct_output_types  &= isa(hashes, Vector{UInt32})
            has_correct_output_ranges &= all(1 .≤ hashes .≤ 10 - (n - 1))
        end

        @test has_correct_output_types
        @test has_correct_output_ranges
    end

    @testset "Hashes are correctly computed" begin
        symbols = collect(50:100)
        hashfn = MinHash(10; symbols = symbols)

        dataset = shuffle(symbols)[1:10]
        hashes = hashfn(dataset)

        @test eltype(hashes) == hashtype(hashfn)

        @test let hashes_match = true
            for (hash, mapping) in zip(hashes, hashfn.mappings)
                if !hashes_match
                    break
                end

                # Compute MinHash manually
                expected_hash = minimum(mapping[x] for x in dataset)
                hashes_match &= (hash == expected_hash)
            end
            hashes_match
        end
    end

    @testset "Collision probabilities correlated with Jaccard similarity" begin
        # Create three different datasets, such that the first and second datasets
        # have high Jaccard similarity while the first and third have low Jaccard
        # similarity. Test whether we in fact get higher collision rates for the
        # two sets with high similarity than the two sets with low similarity.
        symbols = collect(1:200)
        n_hashes = 100
        hashfn = MinHash(n_hashes; symbols=symbols)

        shuffle!(symbols)
        dataset_1 = symbols[1:100]
        dataset_2 = [symbols[1:75]; symbols[101:125]]
        dataset_3 = [symbols[1:25]; symbols[101:175]]

        hashes_1 = hashfn(dataset_1)
        hashes_2 = hashfn(dataset_2)
        hashes_3 = hashfn(dataset_3)

        @test sum(hashes_1 .== hashes_2) > sum(hashes_1 .== hashes_3)
    end

    @testset "MinHash observed collision frequencies match probabilities" begin
        # In theory, the probability of collision for two datasets should be
        # roughly equal to the Jaccard similarity between those datasets
        symbols = collect(1:200)
        n_hashes = 10_000
        hashfn = MinHash(n_hashes; symbols=symbols)

        shuffle!(symbols)
        dataset_1 = Set(symbols[1:100])
        dataset_2 = Set([symbols[1:75]; symbols[101:125]])
        dataset_3 = Set([symbols[1:25]; symbols[101:175]])

        hashes_1 = hashfn(dataset_1)
        hashes_2 = hashfn(dataset_2)
        hashes_3 = hashfn(dataset_3)

        sim_12 = similarity(hashfn)(Set(dataset_1), Set(dataset_2))
        sim_13 = similarity(hashfn)(Set(dataset_1), Set(dataset_3))
        prob_12 = collision_probability(hashfn, sim_12; n_hashes=1)
        prob_13 = collision_probability(hashfn, sim_13; n_hashes=1)

        mean(x) = sum(x) / length(x)

        @test abs(mean(hashes_1 .== hashes_2) - prob_12) ≤ 0.01
        @test abs(mean(hashes_1 .== hashes_3) - prob_13) ≤ 0.01
    end

    @testset "Can omit symbol set to lazily update hash functions" begin
        # If the users want, they can decide not to pre-specify the symbol
        # set. In that case, the hash functions should be lazily updated.
        hashfn = MinHash(100)

        @test isa(hashfn, MinHash{Any, UInt64})

        symbols = [collect(1:10); "a"; "b"; "c"]
        dataset = shuffle(symbols)[1:3]
        hashes = hashfn(dataset)

        @test isa(hashes, Vector{UInt64})

        # Should still be able to hash new elements, updating the hash
        # functions further.
        new_dataset = ["d"; "e"; "f"]
        new_hashes = hashfn(new_dataset)

        @test sum(new_hashes .== hashes) == 0
    end

    @testset "Achieve collision probabilities with lazily updated hash functions" begin
        symbols = collect(1:1_000)
        dataset_1 = Set(shuffle(symbols)[1:500])
        dataset_2 = Set(shuffle(symbols)[1:500])
        similarity = jaccard(dataset_1, dataset_2)

        n_hashes = 10_000
        hashfn = MinHash(n_hashes; dtype=Int64)
        hashes_1 = hashfn(dataset_1)
        hashes_2 = hashfn(dataset_2)

        collision_proportion = sum(hashes_1 .== hashes_2) / n_hashes
        @test abs(collision_proportion - similarity) ≤ 0.01
    end
end


using Curvelet
using Base.Test

N           = 512
sample_data = randn(N,N)
plan        = Curvelet.planCurveletTransform(N, Curvelet.C1)
tmp         = Curvelet.curveletTransform(sample_data, plan)
sample_hat  = Curvelet.inverseCurveletTransform(tmp, plan)
@test norm(sample_hat.-sample_hat) < 0.01


module ChemicalReactionNetworks

export
    Reaction,
    ReactionNetwork,
    SpeciesComplex,
    ==,
    hash,
    isredundant,
    stoichiometric_matrix,
    mass_action,
    mass_action!,
    reaction_current,
    jacobian,
    from_stoichiometric_matrix,
    normal_form,
    conservation_laws,
    cycles,
    n_conservation_laws,
    n_cycles,
    rand_concentrations,
    chemostatted!,
    relabel!,
    complexes,
    complex_decomposition,
    deficiency,
    complex_concentration,
    equilibrium_state,
    equilibrium_state_space,
    cycle_affinities,
    specificity,
    complex_basis_reactions,
    free_energies,
    reaction_rates,
    petri_net

include("crn.jl")
include("algebra.jl")
include("sampling.jl")
include("chemostat.jl")
include("complexes.jl")
include("thermodynamics.jl")
include("petri.jl")

"""
A Julia package for studying and simulating Chemical
Reaction Networks, of the kind that would be encountered in
the study of complex metabolic networks, self-replication,
and so on.

API Overview:
-`Reaction(...)` creates a reaction.
-`ReactionNetwork(...)` creates a network of reactions.
-`mass_action(reactions, concentration)` returns change in concentration over time.
-`chemostatted!(reactions, chemostatted_species, chemostatted_concentrations)` chemostats certain species and modifies the network.
-`equilibrium_state(reaction_net)` returns detailed-balance equilibrium state.
-`conservation_laws(reaction_net)` returns conservation laws of net
-`cycles(reaction_net)` returns cycles of net.
-`deficiency(reaction_net)` returns network deficiency.
-`complex_decomposition(reaction_net)` decomposes a CRN in terms of a set of complexes and incidence matrix.
"""

end


using HomotopyContinuation.DoubleDouble
const DD = HomotopyContinuation.DoubleDouble

@testset "DoubleF64" begin
    @test DoubleF64(Float32(2.0)) == DoubleF64(2.0)
    @test DoubleF64(Float16(2.0)) == DoubleF64(2.0)
    @test DoubleF64(2) == DoubleF64(2.0)
    @test DoubleF64(big(π)) isa DoubleF64
    @test DoubleF64(π) == DoubleF64(big(π))

    @test DD.hi(DoubleF64(2.0)) == 2.0
    @test DD.lo(DoubleF64(2.0)) == 0.0

    @test isbitstype(DoubleF64)
    @test isbits(DoubleF64(3))

    @test zero(DoubleF64(2.0)) == zero(DoubleF64)

    @test convert(Integer, DoubleF64(2.0)) isa Int64
    @test 32 ≤ length(string(DoubleF64(rand()))) ≤ 40

    @test convert(DoubleF64, 2.0) isa DoubleF64
    @test convert(DoubleF64, Float32(2.0)) isa DoubleF64
    @test convert(DoubleF64, Float16(2.0)) isa DoubleF64
    @test convert(DoubleF64, 2) isa DoubleF64
    @test convert(DoubleF64, π) isa DoubleF64

    @test promote_type(DoubleF64, Float64) == DoubleF64
    @test promote_type(DoubleF64, Float32) == DoubleF64
    @test promote_type(DoubleF64, Float16) == DoubleF64

    @test DD.wide_sub(rand(), rand()) isa DoubleF64
    @test DD.wide_add(rand(), rand()) isa DoubleF64
    @test DD.wide_div(rand(), rand()) isa DoubleF64
    @test DD.wide_mul(rand(), rand()) isa DoubleF64

    x = DoubleF64(2.321)
    @test x isa DoubleF64
    @test x / 3.2 isa DoubleF64

    @test x + 2 isa DoubleF64
    @test 2 + x isa DoubleF64

    @test im * DoubleF64(rand()) isa Complex{DoubleF64}
    @test im * DoubleF64(rand()) isa ComplexDF64
    @test x < 2 * x
    @test x ≤ x
    @test convert(Float64, x) ≤ x
    @test x ≤ 2 * convert(Float64, x)
    @test DoubleF64(2.0) == 2.0
    @test 2.0 == DoubleF64(2.0)

    @test isnan(DoubleF64(NaN))
    @test isinf(DoubleF64(Inf))

    @test floor(DoubleF64(3.2)) == 3.0
    @test floor(Int, DoubleF64(3.2)) == 3
    @test ceil(DoubleF64(3.2)) == 4.0
    @test ceil(Int, DoubleF64(3.2)) == 4
    @test trunc(DoubleF64(3.2)) == 3.0
    @test trunc(Int, DoubleF64(3.2)) == 3
    @test isinteger(DoubleF64(3.2)) == false
    @test isinteger(DoubleF64(3.0)) == true

    @test DoubleF64(0.0)^0.14 == 1.0

    for k = 1:10
        x = DoubleF64(rand()) * 20 - 10
        y = DoubleF64(rand()) * 20 - 10
        @test x * y ≈ big(x) * big(y) atol = 1e-29
        @test x + y ≈ big(x) + big(y) atol = 1e-30
        @test x - y ≈ big(x) - big(y) atol = 1e-30
        @test x / y ≈ big(x) / big(y) atol = 1e-26
        u = rand()
        @test x / u ≈ big(x) / big(u) atol = 1e-26
        @test x^5 ≈ BigFloat(x)^5 atol = 1e-25

        @test exp(x) ≈ exp(big(x)) atol = 1e-26
        @test log(abs(x)) ≈ log(big(abs(x))) atol = 1e-29
        @test sin(x) ≈ sin(big(x)) atol = 1e-30
        @test cos(x) ≈ cos(big(x)) atol = 1e-30
        @test atan(x) ≈ atan(big(x)) atol = 1e-30
        @test atan(y, x) ≈ atan(big(y), big(x)) atol = 1e-30
        @test tan(x) ≈ tan(big(x)) atol = 1e-21
        sin_x, cos_x = DD.sincos(x)
        @test sin_x ≈ sin(big(x)) atol = 1e-30
        @test cos_x ≈ cos(big(x)) atol = 1e-30

        z = DoubleF64(rand() * 2 - 1)
        @test asin(z) ≈ asin(big(z)) atol = 1e-30
        @test acos(z) ≈ acos(big(z)) atol = 1e-30

        n, p, d = Base.decompose(x)
        @test x ≈ BigFloat(n) * BigFloat(2)^p / BigFloat(d)
    end

    @test eps(DoubleF64(1.0)) == 4.93038065763132e-32 # 2^-104
    @test eps(DoubleF64) == 4.93038065763132e-32 # 2^-104
    @test eps(DoubleF64) == 4.93038065763132e-32 # 2^-104
end


using CurveProximityQueries
using ConvexBodyProximityQueries
using IntervalArithmetic
using Plots
using StaticArrays

pgfplots()

include("EllipseBounds.jl")
using .EllipseBounds

curve = Bernstein([[-0.6, 0.0], [-0.6, 0.0], [0.1, -0.4], [0.4, 1.6], [-0.2, -1.6], [-0.1, 0.4], [0.8, 0.0], [0.8, 0.0]])
ell = Ellipse(curve, 0.2..0.8)

curve2 = Bernstein([[0.6, 0.6], [0.4, 1.0], [0.2, 0.2], [0.0, 0.6]])
ell2 = Ellipse(curve2, 0.2..0.9)

pts = closest_points(ell, ell2, @SVector(ones(2)))

plot(ell,
    linecolor=RGBA(102/255,130/225,223/255,0.5),
    fillcolor=RGBA(102/255,130/225,223/255,0.5),
    )
plot!(ell2,
     linecolor=RGBA(102/255,130/225,223/255,0.5),
     fillcolor=RGBA(102/255,130/225,223/255,0.5),
     )
plot!(curve, 0.0..0.2,
      linewidth=2.0,
      linestyle=:dash,
      linecolor=Gray(0.5),
     )
plot!(curve, 0.2..0.8,
      linewidth=2.0,
      linecolor=:black,
     )
plot!(curve, 0.8..1.0,
      linewidth=2.0,
      linestyle=:dash,
      linecolor=Gray(0.5),
     )
scatter!(curve(0.2), markersize=4, markercolor=:black)
scatter!(curve(0.8), markersize=4, markercolor=:black)

plot!(curve2, 0.0..0.2,
      linewidth=2.0,
      linestyle=:dash,
      linecolor=Gray(0.5),
     )
plot!(curve2, 0.2..0.9,
      linewidth=2.0,
      linecolor=:black,
     )
plot!(curve2, 0.9..1.0,
      linewidth=2.0,
      linestyle=:dash,
      linecolor=Gray(0.5),
     )
scatter!(curve2(0.2), markersize=4, markercolor=:black)
scatter!(curve2(0.9), markersize=4, markercolor=:black)

plot!(pts,
      linestyle=:dash,
      markersize=8.0,
      linewidth=2.0,
     )
savefig("out/fig4b.pdf")