# ------------------------------------------------------------------
# Includes para os outros arquivos do resto da turma
#include("funcoes_relax.jl")
include("feasible_solution.jl")

# Código original do Raphael
mutable struct node
  level::Int
  model::JuMP.Model
end

## Checks if solution is binary
function isBinary(model::JuMP.Model, binaryIndices::Vector{Int64})
  isZero = model.colVal[binaryIndices] .== 0
  isOne  = model.colVal[binaryIndices] .== 1
  if sum(isZero + isOne) == length(binaryIndices)
    return true
  end
  return false
end

## Checks if model is max: if it is, converts to min
function convertSense!(m::JuMP.Model)
  if m.objSense == :Max
    m.objSense = :Min
    m.obj = -m.obj
  else
    m.objSense = :Max
    m.obj = -m.obj
  end
end

## Strong branching -- if amount of branches == Inf --> full strong branching
function strong(currentNode::node, binaryIndices::Vector{Int64}, amountOfBranches::Int64,
                pscMatrix::Matrix{Float64})

  isZero = currentNode.model.colVal[binaryIndices] .== 0
  isOne  = currentNode.model.colVal[binaryIndices] .== 1
  indFracIndices = find(isZero .+ isOne .== 0)
  fracIndices = binaryIndices[indFracIndices]

  n = min(amountOfBranches, length(fracIndices))

  indCandidatesToBranch = randperm(length(fracIndices))[1:n]
  candidatesToBranch = fracIndices[indCandidatesToBranch]
  bounds = Array{Float64}(2,n)
  for i = 1:n
    leftModel = deepcopy(currentNode.model)
    leftModel.colUpper[candidatesToBranch[i]] = 0
    leftModel.colLower[candidatesToBranch[i]] = 0

    rightModel = deepcopy(currentNode.model)
    rightModel.colUpper[candidatesToBranch[i]] = 1
    rightModel.colLower[candidatesToBranch[i]] = 1

    solve(leftModel)
    bounds[1,i] = leftModel.objVal
    solve(rightModel)
    bounds[2,i] = rightModel.objVal
  end

  indBestFrac = ceil(Int,indmax(bounds)/2)
  indToBranch = candidatesToBranch[indBestFrac]

  leftModel = deepcopy(currentNode.model)
  leftModel.colUpper[indToBranch] = 0
  leftModel.colLower[indToBranch] = 0

  rightModel = deepcopy(currentNode.model)
  rightModel.colUpper[indToBranch] = 1
  rightModel.colLower[indToBranch] = 1

  leftDelta = abs(currentNode.model.objVal - leftModel.objVal)/currentNode.model.colVal[indToBranch]
  rightDelta = abs(currentNode.model.objVal - rightModel.objVal)/(1-currentNode.model.colVal[indToBranch])

  pscMatrix[indBestFrac, 2]+=1
  pscMatrix[indBestFrac, 1] = (pscMatrix[indBestFrac,1] + ((leftDelta + rightDelta)/2))/pscMatrix[indBestFrac,2]

  leftChild = node(currentNode.level+1, leftModel)
  rightChild = node(currentNode.level+1, rightModel)

  return leftChild, rightChild, pscMatrix
end

## Most fractional branching
function fractional(currentNode::node, binaryIndices::Vector{Int64})

  distance = abs.(currentNode.model.colVal[binaryIndices] - 0.5)
  indFrac = indmin(distance)
  indToSet = binaryIndices[indFrac]

  leftModel = deepcopy(currentNode.model)
  leftModel.colUpper[indToSet] = 0
  leftModel.colLower[indToSet] = 0

  rightModel = deepcopy(currentNode.model)
  rightModel.colUpper[indToSet] = 1
  rightModel.colLower[indToSet] = 1

  leftChild = node(currentNode.level+1, leftModel)
  rightChild = node(currentNode.level+1, rightModel)

  return leftChild, rightChild
end

## Pseudo-cost branching
function pseudocost(currentNode::node, binaryIndices::Vector{Int64}, pscMatrix::Matrix{Float64})

    isZero = currentNode.model.colVal[binaryIndices] .== 0
    isOne  = currentNode.model.colVal[binaryIndices] .== 1
    indFracIndices = find(isZero .+ isOne .== 0)
    fracIndices = binaryIndices[indFracIndices]

    indMaxPSC = indmax(pscMatrix[indFracIndices,1].*(1-currentNode.model.colVal[indFracIndices]))
    indToBranch = fracIndices[indMaxPSC]

    leftModel = deepcopy(currentNode.model)
    leftModel.colUpper[indToBranch] = 0
    leftModel.colLower[indToBranch] = 0

    rightModel = deepcopy(currentNode.model)
    rightModel.colUpper[indToBranch] = 1
    rightModel.colLower[indToBranch] = 1

    leftDelta = abs(currentNode.model.objVal - leftModel.objVal)/currentNode.model.colVal[indToBranch]
    rightDelta = abs(currentNode.model.objVal - rightModel.objVal)/(1-currentNode.model.colVal[indToBranch])

    pscMatrix[indMaxPSC, 2]+=1
    pscMatrix[indMaxPSC, 1] = (pscMatrix[indMaxPSC,1] + ((leftDelta + rightDelta)/2))/pscMatrix[indMaxPSC, 2]

    leftChild = node(currentNode.level+1, leftModel)
    rightChild = node(currentNode.level+1, rightModel)

    return leftChild, rightChild, pscMatrix
end

## Receives node and creates two children by setting a variable to 0 and 1 respectively
function branch(currentNode::node, binaryIndices::Vector{Int64}, method::Symbol;
                amountOfBranches::Int64 = 10^10, pscMatrix::Matrix{Float64} = ones(1,1))

  if method == :fractional
    leftChild, rightChild = fractional(currentNode, binaryIndices)
  elseif method == :strong
    leftChild, rightChild, pscMatrix = strong(currentNode, binaryIndices, amountOfBranches, pscMatrix)
  elseif method == :pseudocost
    leftChild, rightChild, pscMatrix = pseudocost(currentNode, binaryIndices, pscMatrix)
  elseif method == :hybrid
    if currentNode.level <= 3
      leftChild, rightChild, pscMatrix = strong(currentNode, binaryIndices, 10^10, pscMatrix)
    else
      leftChild, rightChild, pscMatrix = pseudocost(currentNode, binaryIndices, pscMatrix)
    end
  else
    println("Error on branching method defition")
  end

  return leftChild, rightChild, pscMatrix
end

function obtainBoundList(nodeList)
  boundList = Array{Float64}(length(nodeList))
  for i = 1 : length(nodeList)
    boundList[i] = nodeList[i].model.objVal
  end
  return boundList
end

## Receives a mixed binary linear JuMP model
function solveMIP(m::JuMP.Model; branchMethod = :strong, traverseMethod = :mixed, boolgrasptsp::Bool = false)

  tic()

  # Check if model is max; if it is, converts to min
  flagConverted = 0
  if m.objSense == :Max
    convertSense!(m)
    flagConverted = 1
  end

  # Best bounds: start as Inf
  bestBound = -1e200
  bestVal = 1e200

  # Create vector of indices of the binary variables
  binaryIndices = find(m.colCat .== :Bin)
  binarySolutions = 0

  flagOpt = 0 # flag that indicates if a viable solution has been found

  # Solve using initial heuristic (GRASP)
  if boolgrasptsp
    m_aux = deepcopy(m)
    status = feasible_solution.grasp(m_aux,true,200,0.03,10000000)
    if status == :SubOptimal && m_aux.objVal < bestVal
      bestVal = m_aux.objVal
      m.colVal = m_aux.colVal
      binarySolutions += 1
      flagOpt+=1
    end
  end

  # Solve linear relaxation
  m.colCat[:] = :Cont
  status = solve(m)
  nodes = Vector{node}(0)
  if status == :Optimal && isBinary(m, binaryIndices)
    # Solution of the relaxed problem is binary: optimal solution
    bestBound = m.objVal
    bestVal = m.objVal
    binarySolutions = 1
  else
    push!(nodes, node(0, m)) # Add root to branch and bound tree
    lastNodeLevel = 0
  end

  iter = 1 # number of visited nodes
  branched = false # flag that indicates if a branch has occurred in this iteration
  tol = 0.01 # tolerance (%)

  if traverseMethod == :depth
    traverse = -1
  else
    traverse = 1
  end

  # Initializing pseudo-cost matrix
  pscMatrix = [ones(length(binaryIndices)) zeros(length(binaryIndices))]

  time0 = time_ns()
  timeLimit = 1800

  while !isempty(nodes) && abs((bestVal - bestBound)/bestVal) > tol && (time_ns()-time0)/1e9 < timeLimit

    # Change traverse method every 10 iterations for better bound discovery
    if traverseMethod == :mixed && iter%10 == 0
      traverse = (-1)*traverse
    end

    branched = false
    # Check node lower bound. If greater than current best UB, prune by limit
    if nodes[1].model.objVal <= bestVal
      status = solve(nodes[1].model)
      if status == :Optimal
        boundList = obtainBoundList(nodes)
        bestBound = minimum(boundList)
        if isBinary(nodes[1].model, binaryIndices)
          # Relaxed solution is binary: optimal solution -- don't branch
          if nodes[1].model.objVal < bestVal
            bestVal = nodes[1].model.objVal
            m.colVal = copy(nodes[1].model.colVal)
            flagOpt = 1
            binarySolutions+=1
          end
        elseif nodes[1].model.objVal <= bestVal
          # Relaxed solution is not binary and should not be pruned by limit -- branch
          leftChild, rightChild, pscMatrix = branch(nodes[1], binaryIndices, branchMethod; pscMatrix = pscMatrix)
          branched = true
        end
      end
    end

    # Check if branched -- if true, add children to the list
    if branched == true
      if traverse == 1 # breadth -- insert children at the end of the list
        push!(nodes, leftChild, rightChild)
        lastNodeLevel = nodes[1].level
        deleteat!(nodes, 1)
      else # depth -- insert children at the beginning of the list
        lastNodeLevel = nodes[1].level
        deleteat!(nodes, 1)
        unshift!(nodes, leftChild, rightChild)
      end
    else
      deleteat!(nodes, 1)
    end
    iter+=1

    if iter == 1 || iter%10 == 0
      println("UB: $bestVal")
      println("LB: $bestBound")
    end
  end

  if flagConverted == 1
    convertSense!(m)
    m.objVal = -bestVal
    m.objBound = -bestBound
  else
    m.objVal = bestVal
    m.objBound = bestBound
  end

  # Return binary variables to original state
  m.colCat[binaryIndices] = :Bin

  # Outputs
  m.ext[:status] = status
  if flagOpt == 1
    m.ext[:status] = :Optimal
  end
  m.ext[:nodes] = iter
  m.ext[:solutions] = binarySolutions
  t = toc()
  m.ext[:time] = t

  # Retorna status :userlimit caso tenha extrapolado o tempo
  if m.ext[:time] >= timeLimit
    m.ext[:status] = :userlimit
  end

  return m.ext[:status]
end


using OffsetArrays

setprecision(2000)
T = BigFloat

function rs_find_λ(T, n)
    # https://oeis.org/A087617
    λ = OffsetArray(zeros(T, n + 1), 0:n)
    λ[0] = 1
    eulers = rs_find_eulers(T, n)
    for i in 0:n-1
        c = 0
        for k in 0:i
            c += 2^(4k + 1) * eulers[k + 1] * λ[i - k]
        end
        λ[i + 1] = c / (i + 1)
    end

    λ
end

function rs_find_eulers(T, n)
    # https://oeis.org/A000364
    a = OffsetArray(zeros(T, n + 1), 0:n)
    a[0] = 1
    for i in 1:n
        for j in 0:i-1
            a[i] -= (-1)^(i + j) * a[j] * T(binomial(BigInt(2 * i), BigInt(2 * j)))
        end
    end
    a
end

function rs_find_d(T, m)
    d = OffsetArray(zeros(T, m + 1, m + 1), 0:m, 0:m)
    λ = rs_find_λ(T, m)
    for i in 0:m÷4
        d[i * 4, i * 3] = λ[i]
    end
    for n in 0:m-1
        for k in 0:(3n+2)÷4  # k < 3(n+1)/4  <=> k <= (3n+2)/4
            d[n + 1, k] = (3n + 1 - 4k) * (3n + 2 - 4k) * d[n, k] + (k >= 1 ? d[n, k - 1] : 0)
        end
    end

    d
end

function get_last_nonzero(ps)
    n = length(ps) - 1
    for i in n:-1:0
        if ps[i] != 0
            return i
        end
    end
    0
end

function ps_mul(a, b)
    n = length(a) - 1
    c = OffsetArray(zeros(typeof(a[0]), n+1), 0:n)
    max_i = get_last_nonzero(a)
    max_j = get_last_nonzero(b)
    for i in 0:min(n, max_i)
        for j in 0:min(n-i, max_j)
            c[i + j] += a[i] * b[j]
        end
    end
    c
end

function ps_div(a, b)
    n = length(a) - 1
    c = OffsetArray(zeros(typeof(a[0]), n+1), 0:n)
    for i in 0:n
        v = a[i]
        for j in 0:i-1
            v = v - b[i - j] * c[j]
        end
        c[i] = v / b[0]
    end
    c
end

function ps_compose(coeffs, ps)
    x = ps[0]
    n = length(ps) - 1
    a = copy(ps)
    a[0] = 0

    b = OffsetArray(zeros(typeof(a[0]), n+1), 0:n)
    for i in n:-1:0
        b = ps_mul(b, a)
        b[0] += coeffs[i]
    end
    b
end

function ps_apply(derivatives, ps)
    x = ps[0]
    n = length(ps) - 1
    a = copy(ps)
    a[0] = 0

    b = OffsetArray(zeros(typeof(a[0]), n+1), 0:n)
    factorial = typeof(a[0])(1.0)
    for i in 0:n
        derivatives[i] /= factorial
        factorial *= i + 1
    end
    for i in n:-1:0
        b = ps_mul(b, a)
        b[0] += derivatives[i]
    end
    b
end

function ps_cos(ps)
    n = length(ps) - 1
    x = ps[0]

    cos_x = cos(x)
    sin_x = sin(x)

    derivatives = [-sin_x, -cos_x, sin_x, cos_x]
    fn = OffsetArray(zeros(typeof(x), n+1), 0:n)
    for i in 0:n
        fn[i] = derivatives[mod1(i, 4)]
    end
    fn

    ps_apply(fn, ps)
end

function rs_find_ps(z, n)
    @assert n >= 2

    T = typeof(z)
    numer = OffsetArray(zeros(T, n + 1), 0:n)
    n_pi = T(π)
    numer[0] = z^2 * n_pi / 2 + 3 * n_pi / 8
    numer[1] = n_pi * z
    numer[2] = n_pi / 2
    numer = ps_cos(numer)

    denom = OffsetArray(zeros(T, n + 1), 0:n)
    denom[0] = n_pi * z
    denom[1] = n_pi
    denom = ps_cos(denom)

    ps = ps_div(numer, denom)
#     println(ps_mul(denom, ps) - numer)
    ps
end

function diff_ps(ps)
    n = length(ps) - 1
    ps = copy(ps)
    for i in 1:n
        ps[i - 1] = ps[i] * i
    end
    ps[n] = 0
    ps
end

n = 10
m = 1000

d = rs_find_d(T, n)
ps = rs_find_ps(T(0), m)
derivatives = OffsetArray(zeros(T, 3n+1, m + 1), 0:3n, 0:m)
derivatives[0, :] = ps
for i in 1:3n
    derivatives[i, :] = diff_ps(derivatives[i-1, :])
end

C = OffsetArray(zeros(T, n+1, m + 1), 0:n, 0:m)
for j in 0:n
    for k in 0:(3j ÷ 4)
        C[j, :] = C[j, :] + derivatives[3j - 4k, :] / factorial(T(3j - 4k)) * d[j, k] / T(π)^(2j - 2k)
    end
    C[j, :] /= T(2)^(2j)
end

fac = OffsetArray(zeros(T, 1001), 0:1000)
for i in 0:1000
    fac[i] = factorial(T(i))
end

function binom(n::Int, m::Int)
    fac[n] / fac[m] / fac[n - m]
end

function transform_to_chebyshev(ps) # transform a power series to Chebyshev's
    γ = OffsetArray(zeros(T, m+1), 0:m)

    # even case
    if ps[0] == 0
        ps = OffsetArray([ps[1:m]; 0], 0:m)
    end
    for k in 0:m÷2
        for l in k:m÷2
            γ[2k] += binom(2l, l-k) * ps[2l] / T(2)^(2l - 1)
        end
    end
    γ
end

γ = OffsetArray(zeros(T, n+1, m+1), 0:n, 0:m)
for j in 0:n
    γ[j, :] = transform_to_chebyshev(C[j, :])
end

function repr(x::BigFloat)
    hi = Float64(x)
    lo = Float64(x - hi)
    "f64x2::new($hi, $lo)"
    # "$hi"
end

output_m = 50
println("[")
for i in 0:n
    println("  [")
    for j in 0:2:output_m
        println("    $(repr(j == 0 ? γ[i, 0] / 2 : γ[i, j])),")
    end
    println("  ],")
end
println("]")

# println("[")
# for i in 0:n
#     println("  [")
#     for j in 0:2:output_m
#         println("    $(repr(j == 0 ? γ[i, 0] / 2 : γ[i, j])),")
#     end
#     println("  ], ")
# end
# println("]")


using Revise
using ADSeismic
using ADCME
using JLD2
ADCME.options.customop.verbose = false
mpi_init()
M = N = Int(sqrt(mpi_size()))
n = 10÷M

param = MPIElasticPropagatorParams(NX=10, NY=10, n=n, NSTEP=10, DELTAT=1e-4,
    DELTAX=1.0, DELTAY=1.0,
    USE_PML_XMIN=true, USE_PML_XMAX = true, USE_PML_YMIN = true, USE_PML_YMAX = true)
compute_PML_Params!(param)
vp = 3300.
vs = 3300. / 1.732
rho = 2800.
λ, μ, ρ = compute_lame_parameters(param, vp, vs, rho)
source = Ricker(param, 50.0, 200.0)

# close("all")
# plot(run(sess, source))
# savefig("test.png")
srci = [div(param.NX,5)]
srcj = [div(param.NY,2)]
srctype = [0]
srcv = reshape(source, :, 1)
src = MPIElasticSource(param, srci, srcj, srctype, srcv)

propagator = MPIElasticPropagatorSolver(param, src, ρ, λ, μ)

sess = Session(); init(sess)
@info "starting..."
Vx, Vy, Sxx, Syy, Sxy = run(sess, [propagator.vx, propagator.vy, propagator.sigmaxx, propagator.sigmayy, propagator.sigmaxy])

# Vx = run(sess, propagator.vx)

# Vx[2,:,:]
@info "saving..."
@save "data/dat$(mpi_rank())-$(mpi_size()).jld2" Vx Vy Sxx Syy Sxy param

if mpi_size()>1
    mpi_finalize()
end
# if mpi_rank()==0 && mpi_size()==1
#     # close("all")
#     # pcolormesh(rcvv, cmap="gray")
#     # colorbar()
#     # xlabel("Location")
#     # ylabel("Time")
#     # savefig("receiver.png")

#     close("all")
#     visualize_wavefield(Vx, param)

# end


"""
    AbstractOperator

An abstract type for linear operators on Banach spaces.
"""
abstract type AbstractOperator end

"""
    KernelOperator{K,S} <: AbstractOperator

A type for representing kernel integral operators over Hilbert spaces.
"""
struct KernelOperator{K,S} <: AbstractOperator
    kernel::K
    solver::S
end

KernelOperator(k) = KernelOperator(k,HCubatureJL())

function eval_(ko::KernelOperator, f; solver=ko.solver)
    k = ko.kernel
    fun = function(u,p)
        k(p,u) * f(u)
    end
    return function(x)
        prob = QuadratureProblem{false}(fun,0,1,x)
        solve(prob, solver).u
    end
end

function (ko::KernelOperator)(x)
    eval_(ko, x)
end

Base.:*(ko::KernelOperator, x) = ko(x)

function ChainRulesCore.rrule(ko::KernelOperator, x)
    y = ko(x)
    function ko_pullback(δy)
        δko = NoTangent()
        δx = t -> kernel_transpose_integral(ko.kernel, δy, t, ko.solver)
        return δko, δx
    end
    return y, ko_pullback
end

"""
    FourierKernel{T<:Real}

An object representing a Fourier kernel function ``k(x,y) = sin (ω  π xy + θ)``
"""
struct FourierKernel{T<:Real}
    ω::T
    θ::T
end

FourierKernel(; a = 0.0, b = 1.0) = FourierKernel((b-a), a)

(ff::FourierKernel)(x,y) = sin(ff.ω*π*x*y + ff.θ)

"""
    PolynomialKernel{M<:AbstractMatrix}

A type representing a polynomial kernel function.
"""
struct PolynomialKernel{M<:AbstractMatrix}
    m::M
end

function (pk::PolynomialKernel)(x,y)
    (m,n) = size(pk.m)
    X = [x^(k-1) for k in 1:m]
    Y = [y^(k-1) for k in 1:n]
    X * pk.m * Y
end


#==========================================================
This routine initializes an optimization problem to recover
a CNOT gate on a single qudit with 4 energy levels (and 2
guard states). The drift Hamiltonian in the rotating frame
is
              H_0 = - 0.5*ξ_a(a^†a^†aa),
where a is the annihilation operator for the qudit.
Here the control Hamiltonian includes the usual symmetric
and anti-symmetric terms
     H_{sym} = p(t)(a + a^†),    H_{asym} = q(t)(a - a^†)
which come from the rotating frame approximation and hence
we refer to these as "coupled" controls.
The problem parameters are:
                ω_a =  2π × 4.10595   Grad/s,
                ξ_a =  2π × 2(0.1099) Grad/s.
We use Bsplines with carrier waves with frequencies
0, ξ_a, 2ξ_a Grad/s.
==========================================================#
using LinearAlgebra
using Plots
pyplot()
using FFTW
using DelimitedFiles
using Printf
using Ipopt
using Random

Base.show(io::IO, f::Float64) = @printf(io, "%20.13e", f)

using Juqbox # quantum control module

Nosc = 1 # Number of coupled sub-systems = oscillators
N = 4 # Number of essential energy levels

Nguard = 2 # Number of extra guard levels
Ntot = N + Nguard # Total number of energy levels

T = 100.0 # Duration of gate.

# frequencies (in GHz, will be multiplied by 2*pi to get angular frequencies in the Hamiltonian matrix)
fa = 4.10336
xa = 0.2198
rot_freq = [fa] # Used to calculate the lab frame ctrl function

# setup drift Hamiltonian
number = Diagonal(collect(0:Ntot-1))

H0 = -0.5*(2*pi)*xa* (number*number - number) # xa is in GHz

# lowering matrix
amat = Bidiagonal(zeros(Ntot),sqrt.(collect(1:Ntot-1)),:U) # standard lowering matrix
adag = Array(transpose(amat));
Hsym_ops=[Array(amat + adag)]
Hanti_ops=[Array(amat - adag)]
H0 = Array(H0)

# Estimate time step
maxctrl = 0.001*2*pi * 8.5 #  9, 10.5, 12, 15 MHz

# K1 =  H0 + maxctrl.*( amat +  amat') + 1im* maxctrl.*(amat - amat')
# lamb = eigvals(Array(K1))
# maxeig = maximum(abs.(lamb))

# Pmin = 40
# samplerate1 = maxeig*Pmin/(2*pi)
# nsteps = ceil(Int64, T*samplerate1)

nsteps = calculate_timestep(T, H0, Hsym_ops, Hanti_ops, [maxctrl])
println("# time steps: ", nsteps)

use_bcarrier = true # Use carrier waves in the control pulses?

if use_bcarrier
    Nfreq = 3 # number of carrier frequencies
else
    Nfreq = 1
end

Ncoupled = length(Hsym_ops) # Here, Ncoupled = 1
om = zeros(Ncoupled,Nfreq)

if use_bcarrier
    om[1:Ncoupled,2] .= -2.0*pi *xa       # Note negative sign
    om[1:Ncoupled,3] .= -2.0*pi* 2.0*xa
end
println("Carrier frequencies [GHz]: ", om[1,:]./(2*pi))

maxamp = zeros(Nfreq)

if Nfreq >= 3
    const_fact = 0.45
    maxamp[1] = maxctrl*const_fact
    maxamp[2:Nfreq] .= maxctrl*(1.0-const_fact)/(Nfreq-1) # split the remainder equally
else
    # same threshold for all frequencies
    maxamp .= maxctrl/Nfreq
end

maxpar = maximum(maxamp)

Ident = Matrix{Float64}(I, Ntot, Ntot)

# CNOT trarget
utarget = Ident[1:Ntot,1:N]
utarget[:,3] = Ident[:,4]
utarget[:,4] = Ident[:,3]

omega1 = Juqbox.setup_rotmatrices([N], [Nguard], [fa])

# Compute Ra*utarget
rot1 = Diagonal(exp.(im*omega1*T))

# target in the rotating frame
vtarget = rot1*utarget

# Initial conditions
U0 = Ident[1:Ntot,1:N]

params = Juqbox.objparams([N], [Nguard], T, nsteps, Uinit=U0, Utarget=vtarget, Cfreq=om, Rfreq=rot_freq,
                          Hconst=H0, Hsym_ops=Hsym_ops, Hanti_ops=Hanti_ops)

# initial parameter guess

# D1 smaller than 5 does not work
D1 = 10 # Number of B-spline coefficients per segment
nCoeff = 2*Ncoupled*Nfreq*D1 # factor '2' is for sin/cos

# Random.seed!(12456)

startFromScratch = true # false
startFile="cnot-pcof-opt.dat"

# initial parameter guess
if startFromScratch
    pcof0 = maxpar*0.01 * rand(nCoeff)
    println("*** Starting from random pcof with amplitude ", maxpar*0.01)
else
    # use if you want to read the initial coefficients from file
    pcof0 = vec(readdlm(startFile))
    println("*** Starting from B-spline coefficients in file: ", startFile)
    nCoeff = length(pcof0)
    D1 = div(nCoeff, 2*Ncoupled*Nfreq)  # number of B-spline coeff per control function
end

# min and max coefficient values
useBarrier = true
minCoeff, maxCoeff = Juqbox.assign_thresholds_freq(maxamp, Ncoupled, Nfreq, D1)

samplerate = 32 # only used for plotting
casename = "cnot1" # base file name (used in optimize-once.jl)

maxIter = 200 # 0  # optional argument
lbfgsMax = 250 # optional argument
ipTol = 1e-5   # optional argument
acceptTol = ipTol # 1e-4 # acceptable tolerance
acceptIter = 15

println("*** Settings ***")
println("System Hamiltonian coefficients [GHz]: (fa, xa) =  ", fa, xa)
println("Total number of states, Ntot = ", Ntot, " Total number of guard states, Nguard = ", Nguard)
if use_bcarrier
  println("Using B-spline basis functions with carrier wave, # freq = ", Nfreq)
else
  println("Using regular B-spline basis functions")
end
println("Number of coefficients per spline = ", D1, " Total number of parameters = ", nCoeff)
for q=1:Nfreq
    println("Carrier frequency: ", om[q]/(2*pi), " GHz, max parameter amplitude = ", 1000*maxamp[q]/(2*pi), " MHz")
end
println("Tikhonov coefficients: tik0 = ", params.tik0)

# Estimate number of terms in Neumann series for time stepping (Default 3)
tol = eps(1.0); # machine precision
Juqbox.estimate_Neumann!(tol, params, [maxpar])

wa = Juqbox.Working_Arrays(params,nCoeff)
prob = Juqbox.setup_ipopt_problem(params, wa, nCoeff, minCoeff, maxCoeff, maxIter, lbfgsMax, startFromScratch)

println("Initial coefficient vector stored in 'pcof0'")


using LightGraphs
using LightGraphs.SimpleGraphs
using Plots
using ColorSchemes

include("metaplex.jl")
include("metapopulation.jl")
include("contact_process.jl")

N = 100
M = 10

g = complete_graph(N*M)
g = erdos_renyi(N*M,0.5)
h = erdos_renyi(M, 0.5)
h = complete_graph(M)
h = path_graph(M)

β = 0.3
D = 0.1

tmax = 7.0
nmax = 50000

# Metapopulation
mpp = Metapopulation(h, fill(1, M), β, D)
s0 = fill(N, M)
i0 = fill(0, M)
i0[1] = 10
s0[1] = N-10

t_mpp, s_mpp, i_mpp = metapopulation_gillespie(mpp, s0, i0, nmax=nmax, tmax=tmax)

plot(t_mpp, i_mpp,
    label="",
    line=:steppre,
    xlabel="Time",
    ylabel="# infected"
    )

plot(t_mpp, sum(i_mpp, dims=2)/(N*M),
    label="",
    line=:steppre,
    xlabel="Time",
    ylabel="Fraction of infected individuals",
    title ="Metapopulation (single simulation)"
    )

t_mpp_mc, s_mpp_mc, i_mpp_mc, sd_s_mpp_mc, sd_i_mpp_mc = metapopulation_montecarlo(mpp,s0,i0,tmax=tmax, nmax=nmax, nbins=400)

t_mpp_mf, s_mpp_mf, i_mpp_mf, mpp_mf = metapopulation_ode(mpp, s0, i0, tmax=tmax, saveat=t_mpp_mc)

plot(t_mpp_mc, i_mpp_mc,
    label="",
    linestyle=:dash,
    xlabel="Time",
    ylabel="#infected individuals"
)
    plot!(t_mpp_mf, i_mpp_mf,
    label=""
)

plot(t_mpp_mc, sum(i_mpp_mc, dims=2)/(N*M),
    label="Average (Metapopulation)",
    legend=:bottomright,
    xlabel="Time",
    ylabel="Fraction of infected individuals",
    linestyle=:dash
    )
    plot!(t_mpp_mf, sum(i_mpp_mf, dims=2)/(N*M),
    label="Mean Field (Metapopulation)",
    )
plot!(t_mpp_mc, sum(i_mpp_mc+sd_i_mpp_mc, dims=2)/(N*M), label="", color=:lightblue)
    plot!(t_mpp_mc, sum(i_mpp_mc-sd_i_mpp_mc, dims=2)/(N*M), label="", color=:lightblue)

plot(t_mpp_mc, sum(i_mpp_mf-i_mpp_mc, dims=2)/(N*M))

# Metaplex
mpx = Metaplex(g,h)
Xi0 = falses(N*M)
Xi0[1:10] .= true
Xμ0 = Xμ0 = [div(i-1,N)+1 for i in 1:N*M]

t_mpx, pcs = metaplex_gillespie(mpx, Xi0, Xμ0, β, [D, D], nmax=nmax, tmax=tmax)

plot(t_mpx, pcs[:,2,:], label="", line=:steppre, title="Metaplex (Gillespie)")

t_mpx_mc, s_mpx_mc, i_mpx_mc, sd_s_mpx_mc, sd_i_mpx_mc = metaplex_montecarlo(mpx, Xi0, Xμ0, β, [D, D], tmax=tmax, nmax=nmax)

t_mpx_mf, s_mpx_mf, i_mpx_mf, sμ_mpx_mf, iμ_mpx_mf, mpx_mf = metaplex_ode(mpx, Xi0, Xμ0, β, [D, D], tmax=tmax, saveat=t_mpx_mc)

plot(t_mpx_mc, i_mpx_mc,
    title="Metaplex (Mean vs MF)",
    label="",
    xlabel="Time",
    ylabel="#infected individuals",
    linestyle=:dash
)
    plot!(t_mpx_mf, iμ_mpx_mf, label="")

plot(t_mpx_mc, sum(i_mpx_mc, dims=2)/(N*M),
    label="Average (Metaplex)",
    xlabel="Time",
    ylabel="Fraction of infected individuals")
    plot!(t_mpx_mf, sum(iμ_mpx_mf, dims=2)/(N*M),
    label="Mean Field (Metaplex)")
plot!(t_mpx_mc, sum(i_mpx_mc+sd_i_mpx_mc, dims=2)/(N*M), label="", color=:lightblue)
    plot!(t_mpx_mf, sum(i_mpx_mc-sd_i_mpx_mc, dims=2)/(N*M), label="", color=:lightblue)

t_cp_mc, X_cp_mc, sd_cp_mc = contact_process_montecarlo(g, Xi0, β, tmax=tmax, nmax=nmax)

t_cp_mf, X_cp_mf = contact_process_ode(g, Xi0, β, tmax=tmax, saveat=t_cp_mc)

# Compare the two
colors = ColorSchemes.tab20;
plot(t_mpp_mc, sum(i_mpp_mc, dims=2)/(N*M),
    label="Average (Metapopulation)",
    legend=:bottomright,
    color = colors[1],
    xlabel="Time",
    ylabel="Fraction of infected individuals",
    linestyle=:dash
    )
    plot!(t_mpp_mf, sum(i_mpp_mf, dims=2)/(N*M),
    label="Mean Field (Metapopulation)",
    color=colors[1]
    )
    plot!(t_mpx_mc, sum(i_mpx_mc, dims=2)/(N*M),
    label="Average (Metaplex)",
    linestyle=:dashdot,
    color=colors[3]
    )
    plot!(t_mpx_mf, sum(iμ_mpx_mf, dims=2)/(N*M),
    label="Mean Field (Metaplex)",
    color=colors[3]
    )
plot!(t_mpp_mc, sum(i_mpp_mc+sd_i_mpp_mc, dims=2)/(N*M), label="", color=colors[2])
    plot!(t_mpp_mc, sum(i_mpp_mc-sd_i_mpp_mc, dims=2)/(N*M), label="", color=colors[2])
    plot!(t_mpx_mc, sum(i_mpx_mc+sd_i_mpx_mc, dims=2)/(N*M), label="", color=colors[4])
    plot!(t_mpx_mf, sum(i_mpx_mc-sd_i_mpx_mc, dims=2)/(N*M), label="", color=colors[4])

plot!(t_cp_mc, sum(X_cp_mc, dims=2)/(N*M),
    label="Average (Contact Process)",
    linestyle=:dash
)
plot!(t_cp_mf, sum(X_cp_mf, dims=2)/(N*M),
    label="Mean Field (Contact Process)"
)


# This file was generated, do not modify it. # hide
using PlotlyJS
function random_line()
    a=0.1

    us = 0:0.05π:2π
    ts = -1:0.05:1
    traces=typeof(scatter3d())[]
    marker=attr(
        color="#1f77b4",
        size=12,
        symbol="circle",
        line=attr(color="rgb(0,0,0)", width=0),
    )
    for t in ts
        x=@. (t^2+a)*cos(us)
        y=@. (t^2+a)*sin(us)
        z=t*ones(length(x))
        push!(
            traces,
            scatter3d(
                x=x,y=y,z=z,
                mode="lines",
                marker=marker,
                showlegend=false,
            )
        )
    end

    for u in us
        x = @. (ts^2+a)*cos(u)
        y = @. (ts^2+a)*sin(u)
        z = ts
        push!(
            traces,
            scatter3d(
                x=x,y=y,z=z,
                mode="lines",
                marker=marker,
                showlegend=false
            )
        )
    end

    layout = Layout(
        autosize=true,
        width=500,
        height=500,
        margin=attr(l=0, r=0, b=150, t=0),
    )
    plot(traces,layout)
end

p = random_line()
fdplotly(json(p)) # hide


# # clip-01.jl

# Load Julia packages (libraries) needed  for the snippets in chapter 0

using StatisticalRethinking

# ### snippet 3.1

PrPV = 0.95
PrPM = 0.01
PrV = 0.001
PrP = PrPV*PrV + PrPM*(1-PrV)
PrVP = PrPV*PrV / PrP

# End of `03/clip-01.jl`


abstract type AbstractMMAApprox <: AbstractFunction end

"""
struct MMAApprox{T} <: AbstractMMAApprox
    parent::AbstractFunction # exact function
    xk::AbstractVector # current approximation point
    fk::T # f(xk)
    ∇fk::AbstractVecOrMat # ∇f(xk)
    σ::AbstractVector # move limit
    l::AbstractVector # xk .- σ
    u::AbstractVector # xk .+ σ
    ρ::T # lift - only used in MMA 2002
    p::AbstractVecOrMat # p[i, j] = σ[j]^2 max(0, ∇fk[j]) + ρ[i] σ[j] / 4
    q::AbstractVecOrMat # q[i, j] = σ[i]^2 max(0, -∇fk[j]) + ρ[i] σ[j] / 4
    r::T # r[i] = fk[i] + sum_j (p[i, j] + q[i, j]) / σ[j]
    out::T # cached output of the approximation
end

An instance of `MMAApprox` stores information about the method of moving asymptotes (MMA) approximation of the function `parent`. The notation used here is from the original [1987 MMA paper](https://onlinelibrary.wiley.com/doi/abs/10.1002/nme.1620240207) and the [2002 paper](https://epubs.siam.org/doi/abs/10.1137/S1052623499362822). `xk` is the point around which the approximation is built. `fk` is the value of approximated function at `xk`. And `∇fk` is the gradient of the approximated function at `xk` if it is a scalar-valued function and the Jacobian if it is a vector-valued function. `l` and `u` are the lower and upper asymptotes. `σ` is the move limit such that `l = xk .- σ` and `u = xk .+ σ`. `r`, `p`, `q` and `ρ` are explained in the paper. When an instance of `MMAApprox` is called, its output will be cached in the field `out`.
"""
@params struct MMAApprox{T} <: AbstractMMAApprox
    parent::AbstractFunction # exact function
    xk::AbstractVector # current approximation point
    fk::T # f(xk)
    ∇fk::AbstractVecOrMat # ∇f(xk)
    σ::AbstractVector # move limit
    l::AbstractVector # xk .- σ
    u::AbstractVector # xk .+ σ
    ρ::T # lift - only used in MMA 2002
    p::AbstractVecOrMat # p[i, j] = σ[j]^2 max(0, ∇fk[j]) + ρ[i] σ[j] / 4
    q::AbstractVecOrMat # q[i, j] = σ[i]^2 max(0, -∇fk[j]) + ρ[i] σ[j] / 4
    r::T # r[i] = fk[i] + sum_j (p[i, j] + q[i, j]) / σ[j]
    out::T # cached output of the approximation
end

"""
    MMAApprox(f::AbstractFunction, x::AbstractVector; σ = fill(10.0, length(x)), ρ = 0.0)

Constructs an MMA approximation of the function `f` around the point `x` using a move limit `σ` and `ρ`. See [this paper](https://epubs.siam.org/doi/abs/10.1137/S1052623499362822) for an explanation of the notation.
"""
function MMAApprox(f::AbstractFunction, x::AbstractVector; kwargs...)
    if getdim(f) == 1
        val, grad = value_gradient(f, x)
        return MMAApprox(f, x, val, copy(grad'); kwargs...)
    else
        val, jac = value_jacobian(f, x)
        return MMAApprox(f, x, val, jac; kwargs...)
    end
end
function MMAApprox(
    parent::AbstractFunction,
    x::AbstractVector,
    f::Real,
    ∇f::AbstractVector;
    σ = fill(10.0, length(x)),
    ρ = 0.0,
)
    l = x .- σ
    u = x .+ σ
    ρ = Ref(ρ)

    # All matrices are stored as Adjoint{<:Real, <:AbstractMatrix} for cache efficiency
    p = similar(x)
    q = similar(x)

    r = Ref(zero(f))
    out = Ref(zero(f))

    approx = MMAApprox(parent, copy(x), Ref(f), copy(∇f), σ, l, u, ρ, p, q, r, out)
    updateapprox!(approx, x, f, ∇f)
    return approx
end
function MMAApprox(
    parent::AbstractFunction,
    x::AbstractVector,
    f::Union{Real, AbstractVector},
    ∇f::AbstractMatrix;
    σ = fill(10.0, length(x)),
    ρ = f isa Real ? Ref(zero(f)) : zeros(length(f)),
)
    l = x .- σ
    u = x .+ σ

    # All matrices are stored as Adjoint{<:Real, <:AbstractMatrix} for cache efficiency
    p = similar(∇f, length(x), length(f))'
    q = similar(∇f, length(x), length(f))'

    r = f isa Real ? Ref(f) : similar(f)
    out = f isa Real ? Ref(f) : similar(f)
    fk = f isa Real ? Ref(f) : similar(f)

    approx = MMAApprox(parent, copy(x), fk, copy(∇f')', σ, l, u, ρ, p, q, r, out)
    updateapprox!(approx, x, f, ∇f)
    return approx
end

"""
    (f::MMAApprox)(x::AbstractVector)

Returns an approximation of the function `f.parent` at point `x`. See [`MMAApprox`](@ref) for an explanation.
"""
function (f::MMAApprox{<:Base.RefValue{<:Real}})(x::AbstractVector)
    @unpack p, q, l, u, r, out = f
    @assert all(l[j] < x[j] < u[j] for j in 1:length(x))
    _out = r[]
    for j in 1:length(x)
        _out += p[j] / (u[j] - x[j]) + q[j] / (x[j] - l[j])
    end
    out[] = ForwardDiff.value(_out)
    return _out
end
function (f::MMAApprox{<:AbstractVector})(x::AbstractVector)
    @unpack p, q, l, u, r, out = f
    @assert all(l[j] < x[j] < u[j] for j in 1:length(x))
    if eltype(x) <: ForwardDiff.Dual
        dual_out = map(1:getdim(f)) do i
            output = r[i]
            for j in 1:length(x)
                output += p[i, j] / (u[j] - x[j]) + q[i, j] / (x[j] - l[j])
            end
            return output
        end
        out .= ForwardDiff.value.(dual_out)
        return dual_out
    else
        for i in 1:getdim(f)
            out[i] = r[i]
            for j in 1:length(x)
                out[i] += p[i, j] / (u[j] - x[j]) + q[i, j] / (x[j] - l[j])
            end
        end
        return copy(out)
    end
end

function ChainRulesCore.rrule(f::MMAApprox{<:Base.RefValue{<:Real}}, x::AbstractVector)
    @unpack p, q, l, u, r = f
    @assert all(f.l[j] < x[j] < f.u[j] for j in 1:length(x))
    out = f(x)
    adjoint = Δ -> begin
        if p isa AbstractMatrix
            (nothing, (p' ./ (u .- x) .^ 2 .- q' ./ (x .- l) .^ 2) .* Δ)
        else
            (nothing, (p ./ (u .- x) .^ 2 .- q ./ (x .- l) .^ 2) .* Δ)
        end
    end
    f.out[] = ForwardDiff.value(out)
    return out, adjoint
end
function ChainRulesCore.rrule(f::MMAApprox{<:AbstractVector}, x::AbstractVector)
    @unpack p, q, l, u, r = f
    @assert all(f.l[j] < x[j] < f.u[j] for j in 1:length(x))
    out = f(x)
    adjoint = Δ -> begin
        (nothing, (p ./ (u' .- x') .^ 2 .- q ./ (x' .- l') .^ 2)' * Δ)
    end
    f.out .= ForwardDiff.value.(out)
    return out, adjoint
end

getmmaapprox(approx::MMAApprox) = approx

getasymptotes(approx::MMAApprox) = approx.l, approx.u

getapproxfg(approx::MMAApprox) = approx.out

getσ(approx::MMAApprox{<:Base.RefValue{<:Real}}) = approx.σ[]
getσ(approx::MMAApprox{<:AbstractVector}) = approx.σ

function setσ!(approx::MMAApprox{<:Base.RefValue{<:Real}}, σ::Real)
    approx.σ[] = σ
    return approx
end
function setσ!(approx::MMAApprox{<:AbstractVector}, σ::AbstractVector)
    approx.σ .= σ
    return approx
end

getρ(approx::MMAApprox{<:Base.RefValue{<:Real}}) = approx.ρ[]
getρ(approx::MMAApprox{<:AbstractVector}) = approx.ρ

function setρ!(approx::MMAApprox{<:Base.RefValue{<:Real}}, ρ::Real)
    approx.ρ[] = ρ
    return approx
end
function setρ!(approx::MMAApprox{<:AbstractVector}, ρ::AbstractVector)
    approx.ρ .= ρ
    return approx
end

getfk(approx::MMAApprox{<:Base.RefValue{<:Real}}) = approx.fk[]
getfk(approx::MMAApprox{<:AbstractVector}) = approx.fk

function setfk!(approx::MMAApprox{<:Base.RefValue{<:Real}}, f::Real)
    approx.fk[] = f
    return approx
end
function setfk!(approx::MMAApprox{<:AbstractVector}, f::AbstractVector)
    approx.fk .= f
    return approx
end

getxk(approx::MMAApprox) = approx.xk

function setxk!(approx::MMAApprox, x::AbstractVector)
    approx.xk .= x
    return approx
end

get∇fk(approx::MMAApprox) = approx.∇fk

function set∇fk!(approx::MMAApprox, ∇f::AbstractVecOrMat)
    approx.∇fk .= ∇f
    return approx
end

getparent(approx::MMAApprox) = approx.parent
getdim(approx::MMAApprox) = getdim(getparent(approx))

function updateapprox!(
    approx::MMAApprox{<:Base.RefValue{<:Real}},
    x::AbstractVector,
)
    val, grad = value_gradient(getparent(approx), x)
    return updateapprox!(approx, x, val, grad)
end
function updateapprox!(
    approx::MMAApprox{<:AbstractVector},
    x::AbstractVector,
)
    val, jac = value_jacobian(getparent(approx), x)
    return updateapprox!(approx, x, val, jac)
end
function updateapprox!(
    approx::MMAApprox{<:Base.RefValue{<:Real}},
    x::AbstractVector,
    f::Real,
    ∇f::Union{AbstractVector, Adjoint{<:Any, <:AbstractVector}},
)
    setxk!(approx, x)
    setfk!(approx, f)
    set∇fk!(approx, ∇f)
    return updateapprox!(approx)
end
function updateapprox!(
    approx::MMAApprox{<:AbstractVector},
    x::AbstractVector,
    f::AbstractVector,
    ∇f::AbstractMatrix,
)
    setxk!(approx, x)
    setfk!(approx, f)
    set∇fk!(approx, ∇f)
    return updateapprox!(approx)
end
function updateapprox!(approx::MMAApprox{<:Base.RefValue{<:Real}})
    @unpack xk, fk, ∇fk, l, u, σ, ρ, r, p, q, out = approx
    T = eltype(∇fk)
    l .= xk .- σ
    u .= xk .+ σ
    for j in 1:length(p)
        p[j] = σ[j] * σ[j] * max(0, ∇fk[j]) + ρ[] * σ[j] / 4
        q[j] = σ[j] * σ[j] * max(0, -∇fk[j]) + ρ[] * σ[j] / 4
    end
    r[] = fk[]
    for j in 1:length(p)
        r[] -= (p[j] + q[j]) / σ[j]
    end
    return approx
end
function updateapprox!(approx::MMAApprox{<:AbstractVector})
    @unpack xk, fk, ∇fk, l, u, σ, ρ, r, p, q, out = approx
    T = eltype(∇fk)
    l .= xk .- σ
    u .= xk .+ σ
    # All matrices are stored as Adjoint{<:Real, <:AbstractMatrix} for cache efficiency
    for i in 1:size(p, 1)
        for j in 1:size(p, 2)
            p[i, j] = σ[j] * σ[j] * max(0, ∇fk[i, j]) + ρ[i] * σ[j] / 4
            q[i, j] = σ[j] * σ[j] * max(0, -∇fk[i, j]) + ρ[i] * σ[j] / 4
        end
    end
    for i in 1:size(p, 1)
        r[i] = fk[i]
        for j in 1:size(p, 2)
            r[i] -= (p[i, j] + q[i, j]) / σ[j]
        end
    end
    return approx
end


using GLMakie
function peaks(; n = 49)
    x = LinRange(-3, 3, n)
    y = LinRange(-3, 3, n)
    a = 3 * (1 .- x').^2 .* exp.(-(x'.^2) .- (y .+ 1).^2)
    b = 10 * (x' / 5 .- x'.^3 .- y.^5) .* exp.(-x'.^2 .- y.^2)
    c = 1 / 3 * exp.(-(x' .+ 1).^2 .- y.^2)
    return (x, y, a .- b .- c)
end
x, y, z = peaks(; n=30)
set_theme!(theme_dark())
fig =  Figure(resolution = (1200,800))
axs = [Axis3(fig[1,i]; aspect = :data) for i in 1:3]
contour3d!(axs[1], x, y, z; levels = 20, transparency = true)
contour!(axs[1], x, y, z; levels = 20, transformation = (:xy, minimum(z)),
    transparency = true)
lines!(axs[2], cat(x, NaN, dims=1), y, vcat(z, fill(NaN,30)'),
    color = maximum(z, dims=1)[1,:], transparency = false)
lines!(axs[2], cat(x, NaN, dims=1), y,
  	vcat(z, fill(NaN,30)')*0 .+ minimum(z), color = maximum(z, dims=1)[1,:],
  	transparency = true)
wireframe!(axs[3], x, y, z; color = :grey90, transparency = true)
wireframe!(axs[3], x, y, z*0 .+ minimum(z); color = :grey90,
	transparency = true)
[hidedecorations!(ax; grid = false) for ax in axs]
save("topoLines.png", fig)
fig


# Tests for Revenue DEA Models
@testset "RevenueDEAModel" begin

    ## Test Revenue DEA Model with Cooper et al. (2007)
    # Test agains book results
    X = [3 2; 1 3; 4 6]
    Y = [3; 5; 6]
    P = [6; 6; 6]

    dearevenuecooper = dearevenue(X, Y, P, rts = :CRS)

    @test typeof(dearevenuecooper) == RevenueDEAModel

    @test nobs(dearevenuecooper) == 3
    @test ninputs(dearevenuecooper) == 2
    @test noutputs(dearevenuecooper) == 1
    @test ismonetary(dearevenuecooper) == false

    @test efficiency(dearevenuecooper, :Economic)   ≈ [0.9; 1; 0.6] atol = 1e-3
    @test efficiency(dearevenuecooper, :Technical)  ≈ [0.9; 1; 0.6  ] atol = 1e-3
    @test efficiency(dearevenuecooper, :Allocative) ≈ [1; 1; 1] atol = 1e-3

    ## Test Revenue DEA Model with Zofío and Prieto (2006) data.
    # Test agains results in R
    X = [5 3; 2 4; 4 2; 4 8; 7 9]
    Y = [7 4; 10 8; 8 10; 5 4; 3 6]
    P = [3 2; 3 2; 3 2; 3 2; 3 2.0]

    # Revnue CRS
    dearevenuecrs = dearevenue(X, Y, P, rts = :CRS)
    @test efficiency(dearevenuecrs) ≈ [0.4915254237;
                                    1.000;
                                    1.000;
                                    0.250;
                                    0.1735537190]
    @test efficiency(dearevenuecrs, :Technical) ≈ [0.6364;
                                1.000;
                                1.000;
                                0.250;
                                0.2609]  atol = 1e-3
    @test efficiency(dearevenuecrs, :Allocative) ≈ [0.7723970944;
                                1.000;
                                1.000;
                                1.000;
                                0.6652892562]

    @test efficiency(dearevenue(targets(dearevenuecrs, :X), targets(dearevenuecrs, :Y), P, rts = :CRS)) ≈ ones(5)

    # Revenue VRS
    dearevenuevrs = dearevenue(X, Y, P, rts = :VRS)
    @test efficiency(dearevenuevrs) ≈ [0.6444444444 ;
                                    1.000;
                                    1.000;
                                    0.500;
                                    0.4565217391]
    @test efficiency(dearevenuevrs, :Technical) ≈ [0.7777777778 ;
                                1.000;
                                1.000;
                                0.500 ;
                                0.6000000000]
    @test efficiency(dearevenuevrs, :Allocative) ≈ [0.8285714286;
                                1.000;
                                1.000;
                                1.000;
                                0.7608695652]

    @test efficiency(dearevenue(targets(dearevenuevrs, :X), targets(dearevenuevrs, :Y), P, rts = :VRS)) ≈ ones(5)

    # Check defaults
    @test efficiency(dearevenue(X, Y, P)) == efficiency(dearevenuevrs)
    @test efficiency(dearevenuevrs, :Economic) == efficiency(dearevenuevrs)

    # Print
    show(IOBuffer(), dearevenuecooper)

    # Test errors
    @test_throws DimensionMismatch dearevenue([1; 2 ; 3], [4 ; 5], [1; 1; 1]) #  Different number of observations
    @test_throws DimensionMismatch dearevenue([1; 2; 3], [4; 5; 6], [1; 2; 3; 4]) # Different number of observation in prices
    @test_throws DimensionMismatch dearevenue([1; 2; 3], [4 4; 5 5; 6 6], [4 4 4; 5 5 5; 6 6 6]) # Different number of output prices and outputs
    @test_throws ArgumentError dearevenue([1; 2; 3], [4; 5; 6], [1; 2; 3], rts = :Error) # Invalid returns to scale
    @test_throws ArgumentError dearevenue([1; 2; 3], [4; 5; 6], [1; 2; 3], dispos = :Error) # Invalid disposability
    @test_throws ArgumentError normfactor(dearevenue(X, Y, P)) # ERROR: RevenueDEAModel has no normalization factor

    # ------------------
    # Weak Disposability Tests
    # ------------------

    X = [1; 2; 3; 2; 4]
    Y = [2; 3; 4; 1; 3]
    P = [1; 1; 1; 1; 1]

    dearevenueStrong = dearevenue(X, Y, P, dispos = :Strong)
    @test efficiency(dearevenueStrong, :Economic) ≈ [1.0; 1.0; 1.0; 0.3333333333333333; 0.75]
    @test efficiency(dearevenueStrong, :Technical) ≈ [1.0; 1.0; 1.0; 0.3333333333333333; 0.75]
    @test efficiency(dearevenueStrong, :Allocative) ≈ [1.0; 1.0; 1.0; 1.0; 1.0]

    dearevenueWeak = dearevenue(X, Y, P, dispos = :Weak)
    @test efficiency(dearevenueWeak, :Economic) ≈ [1.0; 1.0; 1.0; 0.3333333333333333; 1.0]
    @test efficiency(dearevenueWeak, :Technical) ≈ [1.0; 1.0; 1.0; 0.3333333333333333; 1.0]
    @test efficiency(dearevenueWeak, :Allocative) ≈ [1.0; 1.0; 1.0; 1.0; 1.0]

    # ------------------
    # Test Vector and Matrix inputs and outputs
    # ------------------
    # Tests against results in R

    # Inputs is Matrix, Outputs is Vector
    X = [2 2; 1 4; 4 1; 4 3; 5 5; 6 1; 2 5; 1.6	8]
    Y = [1; 1; 1; 1; 1; 1; 1; 1]
    P = [1; 1; 1; 1; 1; 1; 1; 1]

    @test efficiency(dearevenue(X, Y, P)) ≈ [1; 1; 1; 1; 1; 1; 1; 1]

    # Inputs is Vector, Output is Matrix
    X = [1; 1; 1; 1; 1; 1; 1; 1]
    Y = [7 7; 4 8; 8 4; 3 5; 3 3; 8 2; 6 4; 1.5 5]
    P = [1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1]

    @test efficiency(dearevenue(X, Y, P)) ≈ [1; 0.8571428571; 0.8571428571; 0.5714285714; 0.4285714286; 0.7142857143; 0.7142857143; 0.4642857143]

    # Inputs is Vector, Output is Vector
    X = [2; 4; 8; 12; 6; 14; 14; 9.412]
    Y = [1; 5; 8; 9; 3; 7; 9; 2.353]
    P = [1; 1; 1; 1; 1; 1; 1; 1]

    @test efficiency(dearevenue(X, Y, P)) ≈ [1; 1; 1; 1; 0.4615384615; 0.7777777778; 1; 0.2816951993]

end


function verif(n::Int64,f)
	if n < 0
        return true
    elseif f(n)
        return verif(n-1,f)
    else
        return false
    end
end

function f(x) true end
function g(x) false end
function h(x) x <= 43 end

println(verif(314151692, f))
println(verif(1, f))
println(verif(0, f))
println(verif(-42, f))
println(verif(314151692, g))
println(verif(1, g))
println(verif(0, g))
println(verif(-42, g))
println(verif(314151692, h))
println(verif(1, h))
println(verif(0, h))
println(verif(-42, h))
println(verif(42, h))
println(verif(43, h))


# durbin_watson.jl
# Durbin-Watson test for autocorrelation
#
# Copyright (C) 2017   Benjamin Born
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
# LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
# WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

export DurbinWatsonTest

struct DurbinWatsonTest <: HypothesisTest
    xmat::Array{Float64}  # regressor matrix
    n::Int                # number of observations
    DW::Float64           # test statistic
    p_compute::Symbol     # determines how p-values are computed
end

"""
    DurbinWatsonTest(X::AbstractArray, e::AbstractVector; p_compute::Symbol = :ndep)

Compute the Durbin-Watson test for serial correlation in the residuals of a regression model.

`X` is the matrix of regressors from the original regression model and `e` the vector of
residuals. Note that the Durbin-Watson test is not valid if `X` includes a lagged dependent
variable. The test statistic is computed as
```math
DW = \\frac{\\sum_{t=2}^n (e_t - e_{t-1})^2}{\\sum_{t=1}^n e_t^2}
```
where `n` is the number of observations.

By default, the choice of approach to compute p-values depends on the sample size
(`p_compute=:ndep`). For small samples (n<100), Pan's algorithm (Farebrother, 1980) is
employed. For larger samples, a normal approximation is used (Durbin and Watson, 1950). To
always use Pan's algorithm, set `p_compute=:exact`. `p_compute=:approx` will always use the
normal approximation.

Default is a two-sided p-value for the alternative hypothesis of positive or negative
serial correlation. One-sided p-values can be requested by calling
`pvalue(x::DurbinWatsonTest; tail=)` with the options `:left` (negative serial correlation)
and `:right` (positive serial correlation).

# References

  * J. Durbin and G. S. Watson, 1951, "Testing for Serial Correlation in Least Squares
    Regression: II", Biometrika, Vol. 38, No. 1/2, pp. 159-177,
    [http://www.jstor.org/stable/2332325](http://www.jstor.org/stable/2332325).
  * J. Durbin and G. S. Watson, 1950, "Testing for Serial Correlation in Least Squares
    Regression: I", Biometrika, Vol. 37, No. 3/4, pp. 409-428,
    [http://www.jstor.org/stable/2332391](http://www.jstor.org/stable/2332391).
  * R. W. Farebrother, 1980, "Algorithm AS 153: Pan's Procedure for the Tail Probabilities
    of the Durbin-Watson Statistic", Journal of the Royal Statistical Society, Series C
    (Applied Statistics), Vol. 29, No. 2, pp. 224-227,
    [http://www.jstor.org/stable/2986316](http://www.jstor.org/stable/2986316).

# External links

  * [Durbin-Watson test on Wikipedia](https://en.wikipedia.org/wiki/Durbin–Watson_statistic)
"""
function DurbinWatsonTest(xmat::AbstractArray{T}, e::AbstractArray{T};
    p_compute::Symbol = :ndep) where T<:Real

    n = length(e)
    DW = sum(diff(e) .^2) / sum(e .^2)

    DurbinWatsonTest(xmat, n, DW, p_compute)
end

testname(::DurbinWatsonTest) = "Durbin-Watson autocorrelation test"
population_param_of_interest(x::DurbinWatsonTest) =
    ("sample autocorrelation parameter", "0", 1 - x.DW / 2)
default_tail(test::DurbinWatsonTest) = :both

function show_params(io::IO, x::DurbinWatsonTest, ident)
    println(io, ident, "number of observations:     ", x.n)
    println(io, ident, "DW statistic:               ", x.DW)
end

"""
    pan_algorithm(a::AbstractArray, x::Float64, m::Int, n::Int)

Compute exact p-values for the Durbin-Watson statistic using Pan's algorithm (Farebrother,
1980).

`a` is the vector of non-zero Eigenvalues of ``(I-(X(X'X)^{-1}X'))A`` (see Durbin and
Watson, 1971, p. 2), `x` is the value of the Durbin-Watson statistic, `m` the number of
elements in `a`, and `n` the number of approximation terms (see Farebrother, 1980, eq. 5).

# References

  * J. Durbin and G. S. Watson, 1971, "Testing for Serial Correlation in Least Squares
  Regression: III", Biometrika, Vol. 58, No. 1, pp. 1-19,
  [http://www.jstor.org/stable/2334313](http://www.jstor.org/stable/2334313).
  * R. W. Farebrother, 1980, "Algorithm AS 153: Pan's Procedure for the Tail Probabilities
  of the Durbin-Watson Statistic", Journal of the Royal Statistical Society, Series C
  (Applied Statistics), Vol. 29, No. 2, pp. 224-227,
  [http://www.jstor.org/stable/2986316](http://www.jstor.org/stable/2986316).

"""
function pan_algorithm(a::AbstractArray, x::Float64, m::Int, n::Int)

    ν = findfirst(ai -> ai >= x, a)
    if ν == 0 || ν === nothing
        return 1.0
    elseif ν == 1
        return 0.0
    else
        k = 1
        ν = ν - 1
        h = m - ν

        if ν <= h
            d  = 2; h  = ν; k  = - k; j1 = 0; j2 = 2; j3 = 3; j4 = 1
        else
            d  = - 2; ν  = ν + 2; j1 = m - 2; j2 = m - 1; j3 = m + 1; j4 = m
        end

        pin = pi / (2n)
        sum = (k + 1) / 2
        sgn = k / n
        n2  = 2n - 1

        # first integral
        for  f1 = h - 2 * floor(Int,h/2) : -1 : 0
            for f2 = j2:d:ν
                sum1 = a[j4]
                if f2 == 0
                    prod = x
                else
                    prod = a[f2]
                end
                u = 0.5 * (sum1 + prod)
                v = 0.5 * (sum1 - prod)
                sum1 = 0.0
                for  i = n2:-2:1
                    y = u - v * cos(i * pin)
                    num = y - x
                    prod = 1.0
                    for k in [(1:j1)' (j3:m)']
                        prod *= num / (y - a[k])
                    end
                    sum1 += sqrt(abs(prod))
                end
                sgn = -sgn
                sum += sgn * sum1
                j1 += d
                j3 += d
                j4 += d
            end

            # second integral
            if d == 2
                j3 = j2
            else
                j1 = j2
            end
            j2 = 0
            ν  = 0
        end
        return sum
    end

end

function pvalue(x::DurbinWatsonTest; tail=:both)

    exact_problem_flag = 0
    if (x.p_compute == :ndep && x.n <= 100) || x.p_compute == :exact
        # p-vales based on Pan's algorithm (see Farebrother, 1980)

        # the following setup is, e.g, described in Durbin and Watson (1971)
        A = diagm(-1 => -ones(x.n - 1), 0 => 2 * ones(x.n), 1 => -ones(x.n - 1))
        A[1, 1] = 1
        A[x.n, x.n] = 1
        EV_temp = sort(real(eigvals(A - x.xmat / (x.xmat' * x.xmat) * (x.xmat'*A))))
        EV = EV_temp[EV_temp .> 1e-10]
        p_temp = pan_algorithm(EV, x.DW, length(EV), 15)

        if p_temp < 0.0 || p_temp > 1.0
            # println(p_temp)
            @warn("Exact p-values outside [0,1]. Approximate p-values reported instead.")
            exact_problem_flag = 1
        end
    end

    if exact_problem_flag == 1 || (x.p_compute == :ndep && x.n > 100
        ) || x.p_compute == :approx
        # p-values based on normal approximation (see Durbin and Watson, 1950)

        # the following derivations follow Durbin and Watson (1951, p. 164)
        X = x.xmat
        k = size(X, 2)
        inv_XX = inv(X' * X)

        AX = zeros(x.n, k)
        AX[[1, x.n], :] = X[[1, x.n], :] - X[[2, x.n - 1], :]
        for i = 2:(x.n - 1)
            AX[i, :] = - X[i - 1, :] + 2 * X[i, :] - X[i + 1, :]
        end

        temp_mat = X' * AX * inv_XX
        P = 2 * (x.n - 1) - tr(temp_mat) # first term: tr(A)
        Q = 2 * (3 * x.n - 4) - 2 * tr(AX' * AX * inv_XX) + tr(temp_mat^2)
        # first term: tr(A^2)

        dw_mean = P / (x.n - k)
        dw_var = 2 / ((x.n - k) * (x.n - k + 2)) * (Q - P * dw_mean)

        p_temp = cdf(Normal(dw_mean, sqrt(dw_var)), x.DW)
    end

    if tail == :both
        2 * min(p_temp, 1 - p_temp)
    elseif tail == :right
        p_temp
    elseif tail == :left
        1 - p_temp
    else
        throw(ArgumentError("tail=$(tail) is invalid"))
    end

end


export
    LineSegment1D

"""
A line segment in one dimension with a ≤ b
"""
immutable LineSegment1D
    a::Float64
    b::Float64
    LineSegment1D(a::Real, b::Real) = new(convert(Float64, min(a,b)), convert(Float64, max(a,b)))
end

Base.:+(seg::LineSegment1D, v::Real) = LineSegment1D(seg.a + v, seg.b + v)
Base.:-(seg::LineSegment1D, v::Real) = LineSegment1D(seg.a - v, seg.b - v)

"""
The distance between the line segment and the point P
"""
function get_distance(seg::LineSegment1D, P::Real)::Float64
    if P < seg.a
        return seg.a - P
    elseif P > seg.b
        return P - seg.b
    else
        return 0.0
    end
end

Base.contains(P::LineSegment1D, v::Float64) = P.a ≤ v ≤ P.b
Base.contains(P::LineSegment1D, Q::LineSegment1D) = Q.a ≥ P.a && Q.b ≤ P.b

intersects(P::LineSegment1D, Q::LineSegment1D) = P.b ≥ Q.a && Q.b ≥ P.a


include("testFunctions.jl")
include("polyFunctions.jl")
include("../src/algorithm.jl")

@testset "countZeros counts zeros" begin
    @test_throws ErrorException countZeros(-2 - 5im, 2 - 3im, poly10, 0.01, 0.3)
    @test countZeros(-2 + 2im, 2 - 2im, poly1, 0.1, 0.5) == 2
    @test countZeros(-2 + 2im, 2 - 2im, poly2, 0.1, 0.5) == 5
    @test countZeros(-2 + 2im, 2 - 2im, poly3, 0.01, 0.5) == 2
    @test countZeros(-2 + 2im, 3 - 2im, poly4, 0.01, 0.5) == 2
    @test countZeros(-2 + 2im, 2 - 2im, poly5, 0.01, 0.5) == 2
    @test countZeros(-2 + 5im, 2 - 2im, poly6, 0.01, 0.6) == 1
    @test countZeros(-2 + 2im, 2 - 2im, poly7, 0.01, 0.5) == 2
    @test countZeros(-2 + 2im, 2 - 2im, poly8, 0.01, 0.8) == 3
    @test countZeros(-4 + 5im, 4 - 2im, poly9, 0.01, 0.3) == 3
    @test countZeros(-2 + 4im, 2 - 3im, poly10, 0.01, 0.1) == 2
    @test countZeros(-2 + 4im, 5 - 3im, poly10, 0.01, 0.1) == 3
    @test countZeros(-2 + 5im, 3 - 3im, poly11, 0.01, 0.1) == 3
end;

# Randomised testing

pass = randomised_countzeros(100)
@testset "Randomised testing for countZeros" begin @test pass == 0 end


using IntUtils

# Why can't I write the following on one line?
@enum DIR Right Down Left Up

function next_grid_elt(dir::DIR, r, c)
    if dir == Right
        return (r, c + 1)
    elseif dir == Down
        return (r + 1, c)
    elseif dir == Left
        return (r, c - 1)
    elseif dir == Up
        return (r - 1, c)
    else
        return (0, 0)
    end
end

function next_dir(dir)
    return DIR((Int(dir) + 1) % 4)
end

function lookup_index{T<:Integer}(dim::T, row::T, col::T)::T
    dir = Right
    runlen = 1
    r = div(dim, 2) + 1
    c = div(dim, 2) + 1
    value = 1
    if r == row && c == col
        return value
    end

    runlen = 1
    max_value = dim * dim

    while value < max_value
        for step in 1:2*runlen
            (r, c) = next_grid_elt(dir, r, c)
            value += 1
            if r == row && c == col
                return value
            end
            if step == runlen
                dir = next_dir(dir)
            end
        end
        runlen += 1
        dir = next_dir(dir)
    end
    error(@sprintf("dim=%d, row=%d, col=%d, r=%d, c=%d, value=%d\n",
                   dim, row, col, r, c, value))
end

function main()
    dim = 1001
    ul_diag = 0
    ll_diag = 0
    center = div(dim, 2) + 1
    for d in 1:dim
        ul_diag += lookup_index(dim, d, d)
        ll_diag += lookup_index(dim, dim - d + 1, d)
    end
    v = ul_diag + ll_diag - lookup_index(dim, center , center)
    @printf("array sum over diagonals for dim = %d is %d\n", dim, v)
end

main()


module SpecialFuncs

using PyCall
using SymPy

import Base: gamma, polygamma, beta,
             airyai, airybi,
             besseli, besselj, besselk, bessely

for meth in (
             :jacobi,
             :gegenbauer,
             :chebyshevt, :chebyshevu,
             :legendre, :assoc_legendre,
             :hermite,
             :laguerre, :assoc_laguerre,
             :Ynm, :Ynm_c,
             :hankel1, :hankel2,
             :jn, :yn
             )
    meth_name = string(meth)
    @eval begin
        @doc """
            `$($meth_name)`: a SymPy function.
                The SymPy documentation can be found through: http://docs.sympy.org/latest/search.html?q=$($meth_name)
                """ ->
        ($meth)(xs...;kwargs...) = sympy_meth($meth_name, xs...; kwargs...)
    end
    eval(Expr(:export, meth))
end

## these have (parameter, x) signature. Use symbolic x to call sympy version, othewise
## should dispatch to julia version.
for fn in (:besselj, :bessely, :besseli, :besselk)
    meth = string(fn)
#    eval(Expr(:import, :Base, fn))
    @eval ($fn)(nu::Number, x::Sym; kwargs...) = sympy_meth($meth, nu, x; kwargs...)
    @eval ($fn)(nu::Number, a::Array{Sym}) = map(x ->$fn(nu, x), a)
end

for meth in (:fresnels, :fresnelc, :Ei, :Si, :Ci,
             :airyai, :airybi
           )
    meth_name = string(meth)

    @eval begin
        @doc """
`$($meth_name)`: a SymPy function.
The SymPy documentation can be found through: http://docs.sympy.org/latest/search.html?q=$($meth_name)
""" ->
        ($meth)(x::Sym;kwargs...) = sympy_meth($meth_name, x; kwargs...)
    end
    @eval ($meth)(a::Array{Sym}) = map($meth, a)
    eval(Expr(:export, meth))
end

## :gamma, :beta... # need import
beta(x::Sym, y::Sym) = sympy_meth(:beta, x, y)
gamma(x::Sym) = sympy_meth(:gamma, x)
polygamma(nu::Sym, x::Sym) = sympy_meth(:polygamma, nu, x)

## Hyper and friends don't really have symbolic use...
"""

Evaluates the generalized hypergeometric function:
[hyper](http://docs.sympy.org/dev/modules/mpmath/functions/hypergeometric.html#hyper)

"""
hyper{T<:Number, S<:Number}(as::Vector{T}, bs::Vector{S}, z::Number) = sympy_meth(:hyper, as, bs, z)
export hyper

"""

[docs](http://docs.sympy.org/dev/modules/mpmath/functions/hypergeometric.html#meijer-g-function)

"""
function meijerg{T<:Number, S<:Number}(a1s::Vector{T}, a2s::Vector{T}, b1s::Vector{S}, b2s::Vector{S}, z::Number, r=1;kwargs...)
    sympy_meth(:meijerg, a1s, a2s, b1s, b2s, z;  kwargs...)
end
export meijerg

end


using SLFN

function gridmake(x, y, z)
    nx, ny, nz = length(x), length(y), length(z)
    hcat(repeat(x, outer=[ny*nz]),
         repeat(y, outer=[nz], inner=[nx]),
         repeat(z, inner=[nx*ny]))
end

function main(new_data=:same, n=40)
    println("Using $(new_data) data on each online training epoch")

    # create some data
    nx, ny, nz = 10, 8, 9
    x = linspace(0, 1, nx)
    y = linspace(0, 1, ny)
    z = linspace(0.1, 1, nz)

    # create some arbitrary functions
    f1(x) = sin(2pi/nx * (x-1))
    f2(x, y) = f1(x) .* cos(4pi/ny * (y-1))
    f3(x, y, z) = f2(x, y) .* log(z)

    # construct grids on the data
    X = gridmake(x, y, z)
    Y = f3(X[:, 1], X[:, 2], X[:, 3])

    # in the next line we train a single hidden layer neural network with tanh
    # neurons and random hidden parameters. The hidden neurons are assigned
    # random parameters such that on the initial training data, the output
    # matrix of the hidden layer is well conditioned.
    elm = ROSELM(X, Y, s=60, neuron_type=Linear(Tanh()))

    # see how well it generalizes
    x2 = linspace(0, 1, nx+5)
    y2 = linspace(0, 1, ny+5)
    z2 = linspace(0.1, 1, nz+5)

    X2 = gridmake(x2, y2, z2)
    Y2 = f3(X2[:, 1], X2[:, 2], X2[:, 3])

    function myprint(elm, X, Y, it)
        @printf("%7i | %8.3f | %8.3f\n",  it, maxabs(elm(X2) - Y2), norm(elm.v))
    end

    @printf("%7s | %8s | %8s\n", "init", "max err", "norm(v)")

    myprint(elm, X2, Y2, 0)

    for i in 1:n
	      for _ in 1:5
            if new_data == :same
	              SLFN.fit!(elm, X, Y)
            elseif new_data == :rand
                # generate random data on the same grid
                _X = gridmake(rand(nx), rand(ny), 0.9*rand(nz)+0.1)
                _Y = f3(_X[:, 1], _X[:, 2], _X[:, 3])
                SLFN.fit!(elm, _X, _Y)
            end
	      end
        myprint(elm, X2, Y2, i)
    end

end

main(:same)

println("\n\n")
main(:rand)


julia> print(j(41,3))
Prisoner killing order: 2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35, 38, 0, 4, 9, 13, 18, 22, 27, 31, 36, 40, 6, 12, 19, 25, 33, 39, 7, 16, 28, 37, 10, 24, 1, 21, 3, 34, 15.
Survivor: 30


"""
    L(x, y, k) = -log(mean(k(xi, y) for xi in x))
"""
L(x, y, k) = -log(mean(k(xi, y) for xi in x))
"""
    KernelLoss(; loss = L, K = 50,
                 kerneltype = ModifiedGaussian,
                 gamma = .1,
                 heuristic = MedianHeuristic(),
                 prior = nothing,
                 distance = euclidean,
                 kernel = Kernel(type = kerneltype,
                                 heuristic = heuristic,
                                 gamma = gamma,
                                 distance = distance))

Constructs an structure, used to compare `K` samples generated from a model
to some data with a `kernel` and loss function `loss`. The loss function has the
signature `L(list_of_samples, data, kernel)`; by default it is the negative log
of the average distance between samples and data
(see [LikelihoodfreeInference.L](@ref)).
The `KernelLoss` can be used as a loss in a [`PointEstimator`](@ref).

The returned structure is callable.
# Example
```julia
model(x) = randn() .+ x
data = [2.]
k = KernelLoss()
k([1.], model, data)
```

`K` can be an integer or a function that returns integers, to implement a schedule.
# Example
```julia
using Statistics
model(x) = randn() .+ x
data = [2.]
k = KernelLoss(K = (n) -> n > 20 ? 10^3 : 10)
std(k([1.], model, data, 1) for _ in 1:100)   # large std, because K = 10
std(k([1.], model, data, 21) for _ in 1:100)  # small std, because K = 10^3
```

This `KernelLoss` is inspired by
[Bertl, J., Ewing, G., Kosiol, C., and Futschik, A., (2017)
Approximate maximum likelihood estimation for population genetic inference,
Statistical Applications in Genetics and Molecular Biology. 16:5-6
](http://dx.doi.org/10.1515/sagmb-2017-0016)
"""
struct KernelLoss{P,Tl,TK,Tk}
    loss::Tl
    K::TK
    kernel::Tk
    prior::P
end
function KernelLoss(; loss = L, K = 50,
                      kerneltype = ModifiedGaussian,
                      gamma = .1,
                      heuristic = MedianHeuristic(),
                      prior = nothing,
                      distance = euclidean,
                      kernel = Kernel(type = kerneltype,
                                      heuristic = heuristic,
                                      gamma = gamma,
                                      distance = distance))
    KernelLoss(loss, K, kernel, prior)
end
"""
    MMDLoss(; single_observation = false, kwargs...)

Returns a [`KernelLoss`](@ref) structure with the maximum mean discrepancy (MMD)
loss; to be used as a `loss` in a [`PointEstimator`](@ref). `kwargs` are passed
to the [`KernelLoss`](@ref) constructor.
If `single_observation = true` the data is expected to be a single vector of
numbers; otherwise a vector of vectors.

This `MMDLoss` is inspired by
[Briol, F.-X., Barp, A., Duncan, A. B. and Girolami, M., (2019),
Statistical Inference for Generative Models with Maximum Mean Discrepancy,
arXiv:1906.05944]
(https://arxiv.org/abs/1906.05944)
"""
function MMDLoss(; single_observation = false, kwargs...)
    loss = single_observation ? (S, data, k) -> mmd(S, [data], k) :
                                (S, data, k) -> mmd(vcat(S...), data, k)
    KernelLoss(; loss = loss, kwargs...)
end
computeK(K::Number, ::Any) = K
computeK(K::Function, n) = K(n)
computeK(K::Base.RefValue, ::Any) = K[]
function (k::KernelLoss)(θ, model, data, n = 1)
    S = [model(θ) for _ in 1:computeK(k.K, n)]
    n == 1 && update!(k.kernel, S, data)
    k.loss(S, data, k.kernel) - logpdf(k.prior, θ)
end
mutable struct QDLoss{D,EPS,Tk,P}
    distance::D
    ϵ::EPS
    K::Tk
    prior::P
    mins::Vector{Float64}
end
"""
    QDLoss(; K = 50, epsilon = 2/K, prior = nothing, distance = euclidean)

Constructs an structure, used to compare `K` samples generated from a model
to some data by computing the first epsilon-quantile of the `distance` between
samples and data.
The `QDLoss` can be used as a loss in a [`PointEstimator`](@ref).

The returned structure is callable.
# Example
```julia
model(x) = randn() .+ x
data = [2.]
q = QDLoss()
q([1.], model, data)
```

`K` can be an integer or a function that returns integers, to implement a schedule.
Similarly, `epsilon` can be a number or a function that returns numbers.
# Example
```julia
using Statistics
model(x) = randn() .+ x
data = [2.]
q = QDLoss(K = (n) -> n > 20 ? 10^3 : 10)
std(q([1.], model, data, 1) for _ in 1:100)   # large std, because K = 10
std(q([1.], model, data, 21) for _ in 1:100)  # small std, because K = 10^3
```
"""
function QDLoss(; K = 50, epsilon = 2/computeK(K, 1),
                  prior = nothing, distance = euclidean)
    QDLoss(distance, epsilon, K, prior, Float64[])
end
function lquantile(k::QDLoss{typeof(euclidean)}, S,
                   data::AbstractVector{<:AbstractVector}, n)
    length(k.mins) == 0 && (k.mins = zeros(length(data)))
    ds = [[k.distance(x[i], data[i]) for x in S] for i in eachindex(data)]
    sum(lquantile(ds[i], computeK(k.ϵ, n), Ref(k.mins, i)) for i in eachindex(data))
end
function lquantile(k, S, data, n)
    length(k.mins) == 0 && (k.mins = [0.])
    d = [k.distance(x, data) for x in S]
    lquantile(d, computeK(k.ϵ, n), Ref(k.mins, 1))
end
function lquantile(d, ϵ, min::Base.RefArray)
    q = quantile(d, ϵ)
    q == 0 && return -log(mean(d .== 0)/ϵ) + min[]
    lq = log(q)
    lq < min[] && (min[] = lq)
    lq
end
function (k::QDLoss)(θ, model, data, n = 1)
    S = [model(θ) for _ in 1:computeK(k.K, n)]
    lquantile(k, S, data, n) - 1/length(θ) * logpdf(k.prior, θ)
end
function objective(model, data, loss;
                   callback = (x, y) -> nothing)
    cb = hasmethod(callback, Tuple{Vector{<:Number}, Number}) ?
         callback : (x, y) -> callback()
    n = 0
    θ -> begin
        n += 1
        result = loss(θ, model, data, n)
        cb(θ, result)
        result
    end
end
struct EnergyLoss{TK}
    K::TK
end
"""
    EnergyLoss(K = 50)

Return an energy loss minization structure, that can be used as a `loss` in
[`PointEstimator`](@ref).

See e.g. section 6.1 of
[Székely, G. J. and Rizzo, M. L., (2013)
Energy statistics: A class of statistics based on distances,
Journal of Statistical Planning and Inference, 8:1249-1272]
(http://dx.doi.org/10.1016/j.jspi.2013.03.018).
"""
function EnergyLoss(; prior = nothing, K = 50)
    prior !== nothing && @warn "Priors are ignored in `EnergyLoss`."
    EnergyLoss(K)
end
function (e::EnergyLoss)(θ, model, data::T, n = 1) where T
    S = [model(θ) for _ in 1:computeK(e.K, n)]
    if T <: AbstractVector{<:AbstractVector}
        energydistance(vcat(S...), data)
    elseif T <: AbstractVector{<:Number}
        energydistance(S, [data])
    end
end
struct KLLoss{TK}
    K::TK
end
"""
    KLLoss(K = 50)

Return a KL-loss minization structure, that can be used as a `loss` in
[`PointEstimator`](@ref).

[Jiang, B., Wu, T.-Y. and Wong, W. H., (2018)
Approximate Bayesian Computation with Kullback-Leibler Divergence as Data Discrepancy,
Proceedings of Machine Learning Research,
(http://proceedings.mlr.press/v84/jiang18a.html).
"""
function KLLoss(; prior = nothing, K = 50)
    prior !== nothing && @warn "Priors are ignored in `KLLoss`."
    KLLoss(K)
end
function (e::KLLoss)(θ, model, data::T, n = 1) where T
    S = [model(θ) for _ in 1:computeK(e.K, n)]
    if T <: AbstractVector{<:AbstractVector}
        kldistance(vcat(S...), data)
    elseif T <: AbstractVector{<:Number}
        kldistance(S, data)
    end
end
function (e::KLLoss)(θ, model, data::AbstractVector{<:Number}, n = 1)
    S = [model(θ) for _ in 1:computeK(e.K, n)]
    energydistance(S, [data])
end
"""
    PointEstimator(optimizer, loss)
"""
struct PointEstimator{O,L} <: AbstractPointABC
    optimizer::O
    loss::L
end
"""
    PointEstimator(; prior = nothing, lower = nothing, upper = nothing,
                     optimizer = nlopt(prior),
                     losstype = QDLoss, kwargs...)

Creates a `PointEstimator` with given `prior`, `optimizer` and
`loss = losstype(; prior = prior, kwargs...)`. If the prior is `nothing` a
uniform prior is assumed. In this case the optimizer has to be initialized with
bounds, , e.g. `nlopt(lower, upper)` where `lower` and `upper` are arrays.

# Example
```julia
model(x) = randn() .+ x
data = [2.]
p = PointEstimator(lower = [-5], upper = [5], losstype = QDLoss, K = 10^3)
res = run!(p, model, data, maxfevals = 10^5)
```
"""
function PointEstimator(; prior = nothing, lower = nothing, upper = nothing,
                        optimizer = prior === nothing ? (lower === nothing || upper ===nothing) ? error("Please set `prior` or `lower` and `upper`") : nlopt(lower, upper) :
                                                          nlopt(prior),
                          losstype = QDLoss, kwargs...)
    PointEstimator(optimizer, losstype(; prior = prior, kwargs...))
end

function run!(p::PointEstimator, model, data;
              maxfevals = Inf, callback = () -> nothing, verbose = true)
    _optimize(p.optimizer, objective(model, data, p.loss, callback = callback),
              maxfevals = max(div(maxfevals, computeK(p.loss.K, 1)), 1),
              verbose = verbose)
end


using Test
using HMatrices
using LinearAlgebra
using Random
using StaticArrays

using HMatrices: RkMatrix

include(joinpath(HMatrices.PROJECT_ROOT,"test","testutils.jl"))

Random.seed!(1)

m = 1000
T = Float64
X    = points_on_sphere(m)
Y    = X
struct ExponentialKernel <: AbstractMatrix{Float64}
    X::Vector{SVector{3,Float64}}
    Y::Vector{SVector{3,Float64}}
end
function Base.getindex(K::ExponentialKernel,i::Int,j::Int)
    x,y = K.X[i], K.Y[j]
    exp(-norm(x-y))
end
Base.size(K::ExponentialKernel) = length(K.X), length(K.Y)
K = ExponentialKernel(X,X)

splitter  = CardinalitySplitter(nmax=50)
Xclt      = ClusterTree(X,splitter)
Yclt      = ClusterTree(Y,splitter)
H = assemble_hmat(K,Xclt,Yclt;threads=false,distributed=false)
H_full      = Matrix(H;global_index=false)

@testset "ldiv!" begin
    B = rand(m,2)
    R = RkMatrix(rand(m,3),rand(m,3))
    R_full = Matrix(R)

    ## 3.1
    exact  = ldiv!(UnitLowerTriangular(H_full),copy(B))
    approx = ldiv!(UnitLowerTriangular(H),copy(B))
    @test exact ≈ approx

    exact  = ldiv!(UpperTriangular(H_full),copy(B))
    approx = ldiv!(UpperTriangular(H),copy(B))
    @test exact ≈ approx

    ## 3.2
    exact  = ldiv!(UnitLowerTriangular(H_full),copy(R_full))
    approx = ldiv!(UnitLowerTriangular(H),copy(R))
    @test exact ≈ Matrix(approx)

    ## 3.3
    compressor = PartialACA(;atol=1e-8)
    exact      = ldiv!(UnitLowerTriangular(H_full),copy(H_full))
    approx     = ldiv!(UnitLowerTriangular(H),deepcopy(H),compressor)
    @test exact ≈ Matrix(approx;global_index=false)
end

@testset "rdiv!" begin
    B = rand(2,m)
    R = RkMatrix(rand(m,3),rand(m,3))
    R_full = Matrix(R)
    # 3.1
    exact  = rdiv!(copy(B),UpperTriangular(H_full))
    approx = rdiv!(copy(B),UpperTriangular(H))
    @test exact ≈ approx

    ## 3.2
    exact  = rdiv!(copy(R_full),UpperTriangular(H_full))
    approx = rdiv!(copy(R),UpperTriangular(H))
    @test exact ≈ approx

    ## 3.3
    compressor = PartialACA(;atol=1e-10)
    exact  = rdiv!(deepcopy(H_full),UpperTriangular(H_full))
    approx = rdiv!(deepcopy(H),UpperTriangular(H),compressor)
    @test exact ≈ Matrix(approx;global_index=false)
end


# addprocs(2);

using Revise #used during development to detect changes in module
using PubMedMiner

#Settings
const mh = "Epilepsy"
const concepts = ("Disease or Syndrome", "Mental or Behavioral Dysfunction", "Neoplastic Process");

overwrite = false
@time save_semantic_occurrences(mh, concepts...; overwrite = overwrite)

using FreqTables

occurrence_df = get_semantic_occurrences_df(mh, concepts...);

#temporary fix - FreqTables breaking with Unions of Missings
#see https://github.com/nalimilan/FreqTables.jl/issues/23

occurrence_df[:descriptor] = convert(Vector{String}, occurrence_df[:descriptor])
occurrence_df[:pmid] = convert(Vector{Int}, occurrence_df[:pmid])

@time mesh_frequencies = freqtable(occurrence_df, :pmid, :descriptor);

using PlotlyJS
using NamedArrays

# Visualize frequency
topn = 50
mesh_counts = vec(sum(mesh_frequencies, 1))
count_perm = sortperm(mesh_counts, rev=true)
mesh_names = collect(keys(mesh_frequencies.dicts[2]))

#traces
#most frequent is epilepsy - remove from plot for better scaling
freq_trace = PlotlyJS.bar(; x = mesh_names[count_perm[2:topn]], y= mesh_counts[count_perm[2:topn]], marker_color="orange")

data = [freq_trace]
layout = Layout(;title="$(topn)-Most Frequent MeSH ",
                 showlegend=false,
                 margin= Dict(:t=> 70, :r=> 0, :l=> 50, :b=>200),
                 xaxis_tickangle = 90,)
plot(data, layout)

using BCBIStats.COOccur
using StatsBase

#co-occurrance matrix - only for topp MeSH
# min_frequency = 5 -- alternatively compute topn based on min-frequency
top_occ = mesh_frequencies.array[:, count_perm[2:topn]]
top_mesh_labels = mesh_names[count_perm[2:topn]]
top_occ_sp = sparse(top_occ)
top_coo_sp = top_occ_sp' * top_occ_sp

#Point Mutual Information
pmi_sp = BCBIStats.COOccur.pmi_mat(top_coo_sp)
#chi2
top_chi2= BCBIStats.COOccur.chi2_mat(top_occ, min_freq=0);
#correlation
corrcoef = BCBIStats.COOccur.corrcoef(top_occ);

function plot_stat_mat(stat_mat, labels)
    stat_trace = heatmap(x=labels, y=labels, z=full(stat_mat- spdiagm(diag(stat_mat))))

    data = [stat_trace]
    layout = Layout(;
                     showlegend=false,
                     height = 900, width=900,
                     margin= Dict(:t=> 300, :r=> 0, :l=> 200, :b=>0),
                     xaxis_tickangle = 90, xaxis_autotick=false, yaxis_autotick=false,
                     yaxis_autorange = "reversed",
                     xaxis_side = "top",
                     xaxis_ticks = "", yaxis_ticks = "")
    plot(data, layout)
end

plot_stat_mat(corrcoef, top_mesh_labels)

plot_stat_mat(pmi_sp, top_mesh_labels)

using PlotlyJSFactory

p = create_chord_plot(top_coo_sp, labels = top_mesh_labels)
relayout!(p, title="Co-occurrances between top 50 MeSH terms")
JupyterPlot(p)

using ARules
using DataFrames

# Remove serch MeSH (set column to 0). Rules of lenght 2 with search term correspond to histogram,
# since it is in every transaction
mh_occ = convert(BitArray{2}, mesh_frequencies.array)

mh_col = mesh_frequencies.dicts[2][mh]
mh_occ[:, mh_col] = zeros(size(mh_occ,1))

mh_lkup = convert(DataStructures.OrderedDict{String,Int16}, mesh_frequencies.dicts[2])
@time mh_rules = apriori(mh_occ, supp = 0.001, conf = 0.1, maxlen = 9)

#Pretty print of rules
mh_lkup = Dict(zip(values(mesh_frequencies.dicts[2]), keys(mesh_frequencies.dicts[2])))
rules_dt= ARules.rules_to_dataframe(mh_rules, mh_lkup, join_str = " | ");

println(head(rules_dt))
println("Found ", size(rules_dt, 1), " rules")

supp_int = round(Int, 0.001 * size(mh_occ, 1))
@time root = frequent_item_tree(mh_occ, supp_int, 9);

supp_lkup = gen_support_dict(root, size(mh_occ, 1))
item_lkup = mesh_frequencies.dicts[2]
item_lkup_t = Dict(zip(values(item_lkup), keys(item_lkup)))
freq = ARules.suppdict_to_dataframe(supp_lkup, item_lkup_t);

println(head(freq))
println("Found ", size(freq, 1), " frequent itemsets")

function fill_sankey_data!(node, sources, targets, vals)
    if length(node.item_ids) >1
        push!(sources, node.item_ids[end-1]-1)
        push!(targets, node.item_ids[end]-1)
        push!(vals, node.supp)
    end
    if has_children(node)
        for nd in node.children
            fill_sankey_data!(nd,  sources, targets, vals)
        end
    end
end

sources = []
targets = []
vals = []
fill_sankey_data!(root, sources, targets, vals)

# size(sources)
topn_links = 50
freq_vals_perm = sortperm(vals, rev=true)
s = sources[freq_vals_perm[1:topn_links]]
t = targets[freq_vals_perm[1:topn_links]]
v = vals[freq_vals_perm[1:topn_links]]
l = mesh_names

println("Found, ", length(sources), "links, showing ", topn_links)

pad = 1e-7
trace=sankey(orientation="h",
             node = attr(domain=attr(x=[0,1], y=[0,1]), pad=pad, thickness=pad, line = attr(color="black", width= 0.5),
                         label=l),
             link = attr(source=s, target=t, value = v))

layout = Layout(width=900, height=1100)

plot([trace], layout)


#####
##### unit tests for building blocks
#####

using TrustRegionMethods: ResidualModel, MinimizationModel, cauchy_point, dogleg_boundary,
    ges_kernel, solve_model, unconstrained_optimum

"Return a closure that evaluates to the objective function of a model."
function model_objective(model::ResidualModel)
    @unpack r, J = model
    function(p)
        0.5 * sum(abs2, r .+ J * p)
    end
end

"Basic sanity check for solver results (splat into tail arguments)."
function is_consistent_solver_results(Δ, p, p_norm, on_boundary)
    all(isfinite, p) &&
        (0 ≤ p_norm ≤ Δ) &&
        (norm(p, 2) ≈ p_norm) &&
        ((p_norm ≈ Δ) == on_boundary)
end

@testset "Nonlinear model constructor sanity checks" begin
    r = ones(2)
    J = ones(2, 2)
    @test_throws ArgumentError ResidualModel(fill(Inf, 2), J)
    @test_throws ArgumentError ResidualModel(fill(NaN, 2), J)
    @test_throws ArgumentError ResidualModel(r, fill(-Inf, 2, 2))
    @test_throws ArgumentError ResidualModel(r, fill(NaN, 2, 2))
    @test_throws ArgumentError ResidualModel(r, ones(2, 3))
    @test_throws ArgumentError ResidualModel(r, ones(3, 3))
end

@testset "Cauchy point" begin
    for _ in 1:100
        # random parameters
        n = rand(2:10)
        r = randn(n)
        J = randn(n, n)
        Δ = abs(randn())
        model = ResidualModel(r, J)
        m_obj = model_objective(model)

        # brute force calculation
        pS = - Δ .* normalize(model.J' * model.r, 2)
        opt = Optim.optimize(τ -> m_obj(pS .* τ), 0, 1, Optim.Brent())

        # calculate and compare
        pC, _, _= pC_results = @inferred cauchy_point(Δ, model)
        @test is_consistent_solver_results(Δ, pC_results...)
        @test pC ≈ (Optim.minimizer(opt) .* pS) atol = eps()^0.25 * n

        # estimated decrease invariant
        g_norm = norm(J' * r, 2)
        @test m_obj(zeros(n)) - m_obj(pC) ≥ 0.5 * g_norm * min(Δ, g_norm / norm(J' * J, 2))
    end
end

@testset "dogleg quadratic" begin
    for _ in 1:100
        n = rand(2:10)
        pC = randn(n)
        D = randn(n)
        while dot(D, pC) ≤ 0
            D = randn(n)
        end
        τ = rand()
        Δ = norm(pC .+ τ .* D, 2)
        @test @inferred(dogleg_boundary(Δ, D, pC)) ≈ τ
    end
end

@testset "eigensolver type stability" begin
    # we test this separately because it makes it easier to debug type inference failures
    n = 10
    model = ResidualModel(rand(n), rand(n, n))
    λ, gap, y1, y2 = @inferred ges_kernel(1.0, MinimizationModel(model), I)
    @test λ isa Float64
    @test gap::Float64 ≥ 0
    @test y1 isa Vector{Float64} && y2 isa Vector{Float64}
end

@testset "trust region solver tests" begin
    for _ in 1:100

        # random problem
        n = rand(2:10)
        model = ResidualModel(rand(n), rand(n, n))
        Δ = abs(randn())
        m_obj = model_objective(model)

        # unconstrained optimum and Cauchy point
        pU, pU_norm = @inferred unconstrained_optimum(model)
        @test pU_norm ≥ 0
        pC, _, _ = pC_results = cauchy_point(Δ, model)
        @test is_consistent_solver_results(Δ, pC_results...)

        # Dogleg
        pD, _, _, = pD_results = @inferred solve_model(Dogleg(), Δ, model)
        @test is_consistent_solver_results(Δ, pD_results...)
        @test m_obj(pD) ≤ m_obj(pC) # improve on Cauchy point

        # GES
        pG, _, _ = pG_results = @inferred solve_model(GeneralizedEigenSolver(), Δ, model)
        @test is_consistent_solver_results(Δ, pG_results...)
        @test m_obj(pG) ≤ m_obj(pD) # improve on dogleg
    end
end

@testset "singularities" begin
    # just some basic sanity check to see if the solver can deal with these
    singular1 = ResidualModel(ones(2), ones(2, 2))
    Δ = 1.0
    @test is_consistent_solver_results(Δ, cauchy_point(1.0, singular1)...)
    @test is_consistent_solver_results(Δ, solve_model(Dogleg(), 1.0, singular1)...)
    # FIXME: solver below could do better, just that hard case is not implemented.
    # check for that when it is.
    @test is_consistent_solver_results(Δ, solve_model(GeneralizedEigenSolver(), 1.0,
                                                      singular1)...)
end

@testset "ForwardDiff wrapper" begin
    @testset "different input types" begin
        r = 0:3
        f(x) = x .+ r
        ff = ForwardDiff_wrapper(f, 4)
        J = Diagonal(ones(4))
        for x ∈ (randn(4), zeros(Int, 4), range(1//5, 3//7; length = 4))
            @test ff(x) == (residual = Float64.(x .+ r), Jacobian = J)
        end
    end

    @testset "handle infeasible" begin
        ff = ForwardDiff_wrapper(x -> all(x .> 0) ? x : x .+ NaN, 4)
        @test !all(isfinite, ff(-ones(4)).residual)
        @test ff(ones(4)) == (residual = ones(4), Jacobian = Diagonal(ones(4)))
    end
end

@testset "printing" begin       # just test that printing is defined
    ff = ForwardDiff_wrapper(x -> Diagonal(ones(2)) * x, 2)
    @test repr(ff) isa AbstractString
    res = trust_region_solver(ff, ones(2))
    @test repr(res) isa AbstractString
    @test repr(TrustRegionResult(1, [1.0], nothing, 1.0, false, 99)) isa AbstractString
end


#-----------------------------------------------------------------------------------------------
# 1st run - packages installation

# ERROR: LoadError: ArgumentError: Package DataFrames not found in current path:
# - Run `import Pkg; Pkg.add("DataFrames")` to install the DataFrames package.
# import Pkg; Pkg.add("DataFrames")
# import Pkg;
# Pkg.add("CSV")
# Pkg.add("StatsBase")

#-----------------------------------------------------------------------------------------------
using DataFrames, CSV, Statistics, StatsBase
data = CSV.read("../../../externals/Core.Math.Samples/data/Pejcic_318.csv", copycols = true)

# @show data.ATV
# println()

for col in eachcol(data,true);
    if  (
        eltype(col[2]) <: Real)
        println(mean(col[2])
        );
    end;
end;

# println(names(data),"\n")

# println("Sample Mean: ", mean(data[:ATT]))
# println("Harmonic <= Geometric <= Arithmetic ",(harmmean(data), geomean(data), mean(data)))
# println("Sample Variance: ",var(data))
# println("Sample Standard Deviation: ",std(data))
# println("Minimum: ", minimum(data))
# println("Maximum: ", maximum(data))
# println("Median: ", median(data))
# println("95th percentile: ", percentile(data, 95))
# println("0.95 quantile: ", quantile(data, 0.95))
# println("Interquartile range: ", iqr(data),"\n")


function total_population(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    return(length(y_true))
end

function true_positive(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    return(sum((y_true .== y_pred) & (y_true .== 1)))
end

function true_negative(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    return(sum((y_true .== y_pred) & (y_true .== 0)))
end

function false_positive(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    return(sum(!(y_true .== y_pred) & (y_true .== 1)))
end

function false_negative(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    return(sum(!(y_true .== y_pred) & (y_true .== 0)))
end

function condition_positive(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    return(sum(y_true .== 1))
end

function condition_negative(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    return(sum(y_true .== 0))
end

function predicted_condition_positive(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    return(sum(y_true .== 1))
end

function predicted_condition_negative(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    return(sum(y_pred .== 0))
end

function accuracy_score(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    tp = true_positive(y_true, y_pred)
    tn = true_negative(y_true, y_pred)
    tot_pop = total_population(y_true, y_pred)
    return((tp + tn) / tot_pop)
end

accuracy = accuracy_score

function prevalence(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    cp = condition_positive(y_true, y_pred)
    tot_pop = total_population(y_true, y_pred)
    return(cp / tot_pop)
end

function positive_predictive_value(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    tp = true_positive(y_true, y_pred)
    pcp = predicted_condition_positive(y_true, y_pred)
    return(tp / pcp)
end

precision = positive_predictive_value

function false_discovery_rate(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    fp = false_positive(y_true, y_pred)
    pcp = predicted_condition_positive(y_true, y_pred)
    return(fp / pcp)
end

function negative_predictive_value(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    tn = true_negative(y_true, y_pred)
    pcn = predicted_condition_negative(y_true, y_pred)
    return(tn / pcn)
end

function false_omission_rate(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    fn = false_negative(y_true, y_pred)
    pcn = predicted_condition_negative(y_true, y_pred)
    return(fn / pcn)
end

function true_positive_rate(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    tp = true_positive(y_true, y_pred)
    cp = condition_positive(y_true, y_pred)
    return(tp / cp)
end

sensitivity = true_positive_rate
recall = true_positive_rate

function false_positive_rate(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    fp = false_positive(y_true, y_pred)
    cn = condition_negative(y_true, y_pred)
    return(fp / cn)
end

function false_negative_rate(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    fn = false_negative(y_true, y_pred)
    cp = condition_positive(y_true, y_pred)
    return(fn / cp)
end

function true_negative_rate(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    tn = true_negative(y_true, y_pred)
    cn = condition_negative(y_true, y_pred)
    return(tn / cn)
end

specificity = true_negative_rate

function positive_likelihood_ratio(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    tpr = true_positive_rate(y_true, y_pred)
    fpr = false_positive_rate(y_true, y_pred)
    return(tpr / fpr)
end

function negative_likelihood_ratio(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    fnr = false_negative_rate(y_true, y_pred)
    tnr = true_negative_rate(y_true, y_pred)
    return(fnr / tnr)
end

function diagnostic_odds_ratio(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    plr = positive_likelihood_ratio(y_true, y_pred)
    nlr = negative_likelihood_ratio(y_true, y_pred)
    return(plr / nlr)
end

function f1_score(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    tp = true_positive(y_true, y_pred)
    fp = false_positive(y_true, y_pred)
    fn = false_negative(y_true, y_pred)
    return((2 * tp) / (2 * tp + fp + fn))
end

function matthews_corrcoef(y_true::Vector, y_pred::Vector)
    check_args(y_true, y_pred)
    tp = true_positive(y_true, y_pred)
    tn = true_negative(y_true, y_pred)
    fp = false_positive(y_true, y_pred)
    fn = false_negative(y_true, y_pred)
    numerator = (tp * tn) - (fp * fn)
    denominator = (tp + fp) * (tp + fn) * (tn + fp) * (tn + fn)
    return(numerator / (denominator ^ 0.5))
end


module SignedDistanceFields

using Statistics

# This package calculates the signed distance field for 2d images
# using an approach due to Saito and Toriwaki (1994). The
# generalization to 3d is straightforward but not implemented.
#
# A recent comparative survey of algorithms for the Euclidean distance
# transform (upon which the SDF is based) found Saito's algorithm to be
# a good tradeoff between simplicity and performance. Those interested
# in the fastest approach should look at Meijster's algorithm, which
# is an optimization of the basic approach implemented below.
#
# Link to survey:
# http://www.agencia.fapesp.br/arquivos/survey-final-fabbri-ACMCSurvFeb2008.pdf

export edf, sdf

function xsweep!(img, out, y, xitr)
    dist = -1
    for x in xitr
        val = img[y, x]
        dist < 0 && !val && continue
        dist = val ? 0 : dist + 1
        out[y, x] = min(out[y, x], dist^2)
    end
end

function edf_sq(img)
    ndims(img) == 2 || throw(ArgumentError("Image must be two-dimensional"))

    # An upper bound for the distance between two pixels
    maxval = prod(size(img))^2
    ncols, nrows = size(img)

    # Calculate the row-wise distance transform for each row
    # in two passes, taking the minimum of the distance-from-
    # left and distance-from-right.
    rowdf_sq = fill(maxval, size(img))
    for y in 1:ncols
        xsweep!(img, rowdf_sq, y, 1:nrows)
        xsweep!(img, rowdf_sq, y, reverse(1:nrows))
    end

    # Use the row-wise information to compute the full distance transform
    df_sq = fill(maxval, size(img))
    for x in 1:nrows
        for y in 1:ncols
            for yp in 1:ncols
                ydist_sq = (y - yp)^2
                ydist_sq > df_sq[y, x] && break
                df_sq[y, x] = min(df_sq[y, x], rowdf_sq[yp, x] + ydist_sq)
            end
        end
    end
    df_sq
end

edf(img) = sqrt.(edf_sq(img))
edf(img, xsize, ysize=xsize) = downsample(edf(img), xsize, ysize)

sdf(img) = sqrt.(edf_sq(img)) .- sqrt.(edf_sq((!).(img)))
sdf(img, xsize, ysize=xsize) = downsample(sdf(img), xsize, ysize)

function downsample(img, xsize, ysize=xsize)
    yscale, yrem = divrem(size(img, 1), ysize)
    xscale, xrem = divrem(size(img, 2), xsize)

    # Make sure we're downsampling by integer amounts
    @assert xrem == yrem == 0 "Downsampling by noninteger amounts is not supported"

    out = Matrix{Float64}(undef, ysize, xsize)
    for y in 1:ysize
        for x in 1:xsize
            yinds = (1 + (y-1) * yscale):(y * yscale)
            xinds = (1 + (x-1) * xscale):(x * xscale)
            out[y, x] = mean(view(img, yinds, xinds))
        end
    end
    out
end

end # module


include("sample.jl")

using PyPlot
using Distances
###
### Box Constrained Minimization
###
using Optim, NLSolversBase #hide
import NLSolversBase: clear! #hide

fun(x) =  (1.0 - x[1])^2 + 100.0 * (x[2] - x[1]^2)^2

function fun_grad!(g, x)
g[1] = -2.0 * (1.0 - x[1]) - 400.0 * (x[2] - x[1]^2) * x[1]
g[2] = 200.0 * (x[2] - x[1]^2)
end

function fun_hess!(h, x)
h[1, 1] = 2.0 - 400.0 * x[2] + 1200.0 * x[1]^2
h[1, 2] = -400.0 * x[1]
h[2, 1] = -400.0 * x[1]
h[2, 2] = 200.0
end;

x0 = [0.0, 0.0]
df = TwiceDifferentiable(fun, fun_grad!, fun_hess!, x0)
lx = [-0.5, -0.5]; ux = [0.5, 0.5]
dfc = TwiceDifferentiableConstraints(lx, ux)
res = optimize(df, dfc, x0, IPNewton())
if false
ux = fill(Inf, 2)
dfc = TwiceDifferentiableConstraints(lx, ux)
clear!(df)
res = optimize(df, dfc, x0, IPNewton())
lx = fill(-Inf, 2); ux = fill(Inf, 2)
dfc = TwiceDifferentiableConstraints(lx, ux)
clear!(df)
res = optimize(df, dfc, x0, IPNewton())
lx = Float64[]; ux = Float64[]
dfc = TwiceDifferentiableConstraints(lx, ux)
clear!(df)
res = optimize(df, dfc, x0, IPNewton())
con_c!(c, x) = (c[1] = x[1]^2 + x[2]^2; c)
function con_jacobian!(J, x)
    J[1,1] = 2*x[1]
    J[1,2] = 2*x[2]
    J
end
function con_h!(h, x, λ)
    h[1,1] += λ[1]*2
    h[2,2] += λ[1]*2
end;
lx = Float64[]; ux = Float64[]
lc = [-Inf]; uc = [0.5^2]
dfc = TwiceDifferentiableConstraints(con_c!, con_jacobian!, con_h!,
                                     lx, ux, lc, uc)
x0 = [0.2, 0.2]
res = optimize(df, dfc, x0, IPNewton())

lc = [0.1^2]
dfc = TwiceDifferentiableConstraints(con_c!, con_jacobian!, con_h!,
                                     lx, ux, lc, uc)
res1 = optimize(df, dfc, x0, IPNewton())

function con2_c!(c, x)
    c[1] = x[1]^2 + x[2]^2     ## First constraint
    c[2] = x[2]*sin(x[1])-x[1] ## Second constraint
    c
end

function con2_jacobian!(J, x)
    # First constraint
    J[1,1] = 2*x[1]
    J[1,2] = 2*x[2]
    # Second constraint
    J[2,1] = x[2]*cos(x[1])-1.0
    J[2,2] = sin(x[1])
    J
end

function con2_h!(h, x, λ)
    # First constraint
    h[1,1] += λ[1]*2
    h[2,2] += λ[1]*2
    # Second constraint
    h[1,1] += λ[2]*x[2]*-sin(x[1])
    h[1,2] += λ[2]*cos(x[1])
    # Symmetrize h
    h[2,1]  = h[1,2]
    h
end;

x0 = [0.25, 0.25]
lc = [-Inf, 0.0]; uc = [0.5^2, 0.0]
dfc = TwiceDifferentiableConstraints(con2_c!, con2_jacobian!, con2_h!,
                                     lx, ux, lc, uc)
res2 = optimize(df, dfc, x0, IPNewton())

if true
    f = x -> (x[1]^2 + x[2] -11)^2 + (x[1]+x[2]^2-7)^2
    figure(1)
    x = collect(Float16, range(-5,length=200,stop=5));
    y = collect(Float16, range(-5,length=200, stop=5));
    pairs = collect(Iterators.product(x, y));
    z = f.(pairs);
    #z = hcat(x, y); z = mapslices(fun, z, dims = 2)
    surf(x, y, z);
    (res, loc) = sample(f, [-5, -5], [5, 5]; num = 1000)
    scatter3D(loc[:, 1], loc[:, 2], res, color = "red")
    #contour(x, y, z)
    #scatter3D([0], [0], [100], color = "red")
end
end


module MyModule
function complexfunc(dims::Int, iters::Int)::Float64
    x::Array{Float64, 2} = ones(Float64, (dims, dims))
    for i in 1:iters
        x .+= 1.0
    end

    z::Float64 = 0.0
    for j in 1:dims
        z += x[j, j]
    end

    z
end
end