"""
 Tests for https://arxiv.org/abs/1303.5103 with test problems as in the paper.
 DRI1, RI1, RI3, RI5, RI6, RDI1WM, RDI2WM, RDI3WM, RDI4WM
"""

import Statistics # for mean values of trajectories
import LinearAlgebra # for the normn
using StochasticDiffEq
using Test
using Random
using DiffEqDevTools
#using DiffEqGPU

function prob_func(prob, i, repeat)
    remake(prob,seed=seeds[i])
end

"""
 Test Scalar SDEs (oop)
"""

@info "Scalar oop noise"

numtraj = Int(2e5) # in the paper they use 1e9
u₀ = 0.0
f(u,p,t) = 1//2*u+sqrt(u^2+1)
g(u,p,t) = sqrt(u^2+1)
dts = 1 .//2 .^(4:-1:1)
tspan = (0.0,2.0) # 2.0 in paper

h1(z) = z^3-6*z^2+8*z
#analytical_sol(t) = E(f(X(t))) = E(h1(arsinh(X(t))) = t^3-3*t^2+2*t
#analytical_sol(2) = 0 and analytical_sol(1)=0

seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

prob = SDEProblem(f,g,u₀,tspan)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h1(asinh(sol[end])),false),
        prob_func = prob_func
        )

sim = test_convergence(dts,ensemble_prob,DRI1(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("DRI1:", sim.𝒪est[:weak_final])

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RI1(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RI1:", sim.𝒪est[:weak_final])

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RI3(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RI3:", sim.𝒪est[:weak_final])

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RI5(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RI5:", sim.𝒪est[:weak_final])

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RI6(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RI6:", sim.𝒪est[:weak_final])

numtraj = Int(1e3)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RDI1WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-1) < 0.3
println("RDI1WM:", sim.𝒪est[:weak_final])

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RDI2WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RDI2WM:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,RDI3WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RDI3WM:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,RDI4WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3 # order is 2.568 with 3e5 trajectories
println("RDI4WM:", sim.𝒪est[:weak_final])

"""
 Test Scalar SDEs (iip)
"""

@info "Scalar iip noise"

u₀ = [0.0]
f1!(du,u,p,t) = @.(du = 1//2*u+sqrt(u^2 +1))
g1!(du,u,p,t) = @.(du = sqrt(u^2 +1))
dts = 1 .//2 .^(4:-1:1)
tspan = (0.0,2.0)

h1(z) = z^3-6*z^2+8*z

prob = SDEProblem(f1!,g1!,u₀,tspan)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h1(asinh(sol[end][1])),false),
        prob_func = prob_func
        )

numtraj = Int(2e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,DRI1(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("DRI1:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,DRI1NM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("DRI1NM:", sim.𝒪est[:weak_final])

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RI1(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RI1:", sim.𝒪est[:weak_final])

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RI3(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RI3:", sim.𝒪est[:weak_final])

numtraj = Int(5e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RI5(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RI5:", sim.𝒪est[:weak_final])

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RI6(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RI6:", sim.𝒪est[:weak_final])

numtraj = Int(1e3)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RDI1WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-1) < 0.3
println("RDI1WM:", sim.𝒪est[:weak_final])

numtraj = Int(7e5)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RDI2WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RDI2WM:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,RDI3WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RDI3WM:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,RDI4WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=0.0
)
@test -(sim.𝒪est[:weak_final]-2) < 0.3
println("RDI4WM:", sim.𝒪est[:weak_final])

"""
 Test Diagonal noise SDEs (iip), SIAM Journal on Numerical Analysis, 47 (2009), pp. 1713–1738
"""

@info "Diagonal noise"

u₀ = [0.1,0.1]
function f3!(du,u,p,t)
  du[1] = 3//2*u[1]
  du[2] = 3//2*u[2]
end
function g3!(du,u,p,t)
  du[1] = 1//10*u[1]
  du[2] = 1//10*u[2]
end
dts = 1 .//2 .^(3:-1:0)
tspan = (0.0,1.0)

h3(z) = z^2 # == 1//10**exp(3//2*t) if h3(z) = z and  == 1//100**exp(301//100*t) if h3(z) = z^2 )

prob = SDEProblem(f3!,g3!,u₀,tspan)
ensemble_prob = EnsembleProblem(prob;
        output_func = (sol,i) -> (h3(sol[end][1]),false),
        prob_func = prob_func
        )

numtraj = Int(5e4)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,DRI1(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test -(sim.𝒪est[:weak_final]-2) < 0.3 # order is 2.91
println("DRI1:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,DRI1NM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3 # order is 2.91
println("DRI1NM:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,RI1(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test -(sim.𝒪est[:weak_final]-2) < 0.3 # order is 3.05
println("RI1:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,RI3(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test -(sim.𝒪est[:weak_final]-2) < 0.3 # order is 2.77
println("RI3:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,RI5(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.3
println("RI5:", sim.𝒪est[:weak_final])

numtraj = Int(1e5)
seed = 70
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RI6(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.45 # order is 1.55
println("R6:", sim.𝒪est[:weak_final])

numtraj = Int(1e3)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RDI1WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test abs(sim.𝒪est[:weak_final]-1) < 0.45 # order is 1.44
println("RDI1WM:", sim.𝒪est[:weak_final])

numtraj = Int(1e5)
seed = 70
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RDI2WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test abs(sim.𝒪est[:weak_final]-2) < 0.45 # order is 1.55
println("RDI2WM:", sim.𝒪est[:weak_final])

numtraj = Int(5e4)
seed = 100
Random.seed!(seed)
seeds = rand(UInt, numtraj)

sim = test_convergence(dts,ensemble_prob,RDI3WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test -(sim.𝒪est[:weak_final]-2) < 0.3 # order is 2.84
println("RDI3WM:", sim.𝒪est[:weak_final])

sim = test_convergence(dts,ensemble_prob,RDI4WM(),
    save_everystep=false,trajectories=numtraj,save_start=false,adaptive=false,
    weak_timeseries_errors=false,weak_dense_errors=false,
    expected_value=1//100*exp(301//100)
)
@test -(sim.𝒪est[:weak_final]-2) < 0.3 # order is 2.91
println("RDI4WM:", sim.𝒪est[:weak_final])


abstract type AbstractExcitation end

λ₀(excitation::AbstractExcitation) = c₀ / frequency(excitation)
ω₀(excitation::AbstractExcitation) = 2π * frequency(excitation)
k₀(excitation::AbstractExcitation) = 2π / λ₀(excitation)

@with_kw struct Harmonic <: AbstractExcitation
    freq
    amp = 1.0
end
frequency(signal::Harmonic) = signal.freq
function (signal::Harmonic)(t)
    @unpack freq, amp = signal
    ω₀ = 2π * freq
    amp * exp(-im * (ω₀ * t))
end
@with_kw struct Gaussian <: AbstractExcitation
    freq
    df
    t0 = 2 * inv(df)
    amp = 1.0
end
frequency(signal::Gaussian) = signal.freq
function (signal::Gaussian)(t)
    @unpack freq, df, t0, amp = signal
    ω₀ = 2π * freq
    σₜ = 1 / (2π * df)
    amp * exp(-0.5((t - t0) * inv(σₜ))^2) * exp(-im * (ω₀ * t))
end


using CurveFit
using Polynomials
using Statistics

abstract type AbstractCalibr end
abstract type AbstractCalibr1d end

struct CalibrCurve{Fit} <: AbstractCalibr1d
    "Anemometer calibration output"
    E::Vector{Float64}
    "Calibration Velocity"
    U::Vector{Float64}
    "Polynomial fitting the calibration curve"
    fit::Fit
    "Reference temperature of the calibration"
    T0::Float64
end
Base.broadcastable(cal::CalibrCurve) = Ref(cal)

"""
    cal = calibr_curve(sensor, V, E, temp, N)

Create a calibration for a thermal anemometer

Parameters:
 * `sensor` an anemometer object (<:AbstractThermalAnemometer)
 * `V` a vector with calibration velocities
 * `E` a vector with calibration voltages
 * `temp` a vector with calibration temperatures or a single calibration temperature
 * `N` degree of polynomial fit

This method will correct all voltages to the mean calibration temperature. Then it will fit the nonzero calibration velocities using a polynomial fit of degree N.

The speeds between zero and the first nonzero velocity are then interpolated using King's law.

# Examples
```jldoctest
julia> U = [0.0, 2.0, 3.0, 4.0, 6.0, 8.0, 10.0, 12.0, 18.0];

julia> E = [0.979, 1.333, 1.385, 1.423, 1.48, 1.521, 1.551, 1.575, 1.624];

julia> temp = [20.02, 19.75, 19.79, 19.73, 19.67, 19.75, 19.85, 19.85, 20.12] .+ 273.15;

julia> Rw = Thermistor(5e3, 3500, 25.0+273.15);

julia> cta = CTASensor(Rw, 100.0, 20/120);

julia> cal = CalibrCurve(cta, U, E, temp, 5)
CalibrCurve{CTASensor{Thermistor}}(CTASensor{Thermistor}(Thermistor(5000.0, 3500.0, 298.15), 100.0, 0.16666666666666666, 1.0), Polynomial(2526.816258773263 - 3916.1946161935334*x - 1547.7951632370953*x^2 + 6109.915344234665*x^3 - 3944.5331575193045*x^4 + 814.5086726564189*x^5), 19.836666666666666, (0.9897386382393427, 0.5558472575990494, 0.5), 1.332625512625792)

julia> cal.(E, temp) - U  # Check the calibration
9-element Vector{Float64}:
  0.0
  0.00033274230130109217
 -0.002686962091977829
  0.005428070340446567
 -0.007450571094204328
  0.006889739326171096
 -0.0019912770775185606
 -0.0007380705280706223
  0.00021632881207622745

julia>

```
"""
function calibr_curve(sensor::CTA, V, E, temp, N=5; extrapolate=true) where{CTA<:Hotwire.AbstractThermalAnemometer}

    nz = V .> 0  # Points with velocity above zero
    Tm = mean(temp)

    # Correct output to reference temperature
    Ec = tempcorr.(Ref(sensor), E, temp, Tm)
    Uc = [u for u in V]

    V1 = V[nz]
    E1 = Ec[nz]

    # Fit the nonzero velocities to a polynomial of degree N
    fit = curve_fit(Polynomial, E1, V1, N)

    if extrapolate
        # Now if the speed is lower than te minimum calibration speed,
        # attempt to extrapolate using King's law: E^2 = A + B + U^n
        Ea = minimum(Ec[.!nz])
        Eb = minimum(E1)
        efit = ExtrapolateFit(fit, Ea, Eb)
        return CalibrCurve(Ec, Uc, efit, Tm)
    else
        return CalibrCurve(Ec, Uc, fit, Tm)
    end

end

struct ExtrapolateFit{Fit}
    "A curve fit object obtained during calibration"
    fit::Fit
    "Minimum output WITH velocity"
    E0::Float64
    "King's law A coefficient"
    A::Float64
    "King's law B coefficient"
    B::Float64
    "King's law n coefficient"
    n::Float64
end

"""
    `efit = ExtrapolateFit(fit, Ea, Eb)`

For very low speeds, the behavior of thermal anemometers change
as natural convection begins to dominate the heat transfer from the probe.
Most calibrations reach has a nonzero minimum velocity. This function
extrapolates the calibration curve to zero.

It uses the calibration point with lowest speed and its derivative.
It also uses the output at 0 velocity as a data point.

With these information, a King's law is fitted. Notice that the zero
velocity point is unstable since any air current whatsoever will enhance
heat transfer.

The following curve will be used to approximate the low speeds:

``
E² = A + B⋅Uⁿ
``

Arguments:

 - `fit` curve fit obtained from calibration
 - `Ea`  Anemometer output for 0 velocity
 - `Eb` Anemometer output for minimum calibration velocity

**Returns**: An `ExtrapolateFit` object

"""
function ExtrapolateFit(fit::Fit, Ea, Eb) where {Fit}

    Ua = 0.0
    Ub = fit(Eb)
    f = eps(Eb) * 1e7
    Eb1 = Eb + f
    Eb0 = Eb - f
    Ub1 = fit(Eb1)
    Ub0 = fit(Eb0)
    dE² = ( (Eb1 - Eb0) / (Ub1 - Ub0) ) * (Eb1 + Eb0)
    A = Ea*Ea
    n = Ub * dE² / (Eb*Eb - A)

    B = (Eb*Eb-A)/Ub^n

    return ExtrapolateFit(fit, Eb, A, B, n)
end
Base.broadcastable(fit::ExtrapolateFit) = Ref(cal)

"""
    `fit(E)`

Uses an `ExtrapolateFit` object as a function call.

If E > fit.E0, use the original fit. Otherwise use the extrapolation.
"""
function (fit::ExtrapolateFit)(E)
    if E < fit.E0
        E2 = E*E
        if E2 < fit.A
            return zero(E)
        else
            return ((E2 - fit.A) / fit.B)^(1/fit.n)
        end
    else
        return fit.fit(E)
    end
end

"""
    `BlendFit(fit₁, fit₂, E₁, E₂)`

Blend two curve fits together.
The approximation `fit₁` is used for ``E < E₁`` and
`fit₂` is used for ``E > E₂``. For intermediary values,
a linear blending is used

"""
struct BlendFit{Fit₁,Fit₂}
    "Fit for lower speeds"
    fit₁::Fit₁
    "Fit for higher speeds"
    fit₂::Fit₂
    "Lower threshold"
    E₁::Float64
    "Higher threshold"
    E₂::Float64
end
Base.broadcastable(fit::BlendFit) = Ref(cal)

"""
    `fit(E)`

Calculates the blended fit
"""
function (fit::BlendFit)(E)
    if E < fit.E₁
        return fit.fit₁(E)
    elseif E > fit.E₂
        return fit.fit₂(E)
    else
        u₁ =  fit.fit₁(E)
        u₂ = fit.fit₂(E)
        ΔE = fit.E₂ - fit.E₁
        return u₁ * (fit.E₂-E) / ΔE + u₂*(E - fit.E₁)/ΔE
    end
end

struct HWKingFit
    A::Float64
    B::Float64
    n::Float64
end
Base.broadcastable(fit::HWKingFit) = Ref(cal)

function HWKingFit(E::AbstractVector, U::AbstractVector, ϵ=1e-8, maxiter=200)

    A, B, n = king_fit(E, U, ϵ, maxiter)

    return HWKingFit(A, B, n)
end

function (fit::HWKingFit)(E)

    E² = E*E
    if E² < fit.A
        return zero(E)
    else
        return (  (E²-fit.A) / fit.B )^(1/fit.n)
    end
end

"""
    `U = cal(sensor, E [,temp])`

Apply the calibration to obtain wind speed.

Parameters:
 * `sensor` Structure representing the anemometer
 * `E` Output voltage
 * `temp` fluid temperature. If not provided assume calibration temperature
 * Returns velocity

# Examples
```jldoctest
julia> cal(w, 1.5)
6.925627034695708

julia> cal(w, 1.5, 15)
5.902143445069172

julia> cal(w, 1.5, 20)
6.9642311432613395

julia> cal(w, 1.5, 25)
8.309847269677324

```

"""
function (cal::CalibrCurve)(sensor, E, temp)
    Ec = tempcorr(sensor, E, temp, cal.T0)

    return cal.fit(Ec)
end

(cal::CalibrCurve)(sensor, E) = cal(sensor, E, cal.T0)


# Example 3.25, 3.35 and 3.38 of
# Blekherman, G., Parrilo, P. A., & Thomas, R. R. (Eds.).
# Semidefinite optimization and convex algebraic geometry SIAM 2013
@testset "Example 3.25 with $(factory.constructor)" for factory in sdp_factories
    @polyvar x y
    m = SOSModel(factory)
    @constraint m x^2-x*y^2+y^4+1 >= 0
    JuMP.optimize!(m)
    @test JuMP.primal_status(m) == MOI.FeasiblePoint
end
@testset "Example 3.35 with $(factory.constructor)" for factory in sdp_factories
    @polyvar x
    m = SOSModel(factory)
    @constraint m x^4+4x^3+6x^2+4x+5 in SOSCone() # equivalent to >= 0
    JuMP.optimize!(m)
    @test JuMP.primal_status(m) == MOI.FeasiblePoint
    # p(x) = (x^2+2x)^2 + 2(1+x)^2 + 3
end
@testset "Example 3.38 with $(factory.constructor)" for factory in sdp_factories
    @polyvar x y
    m = SOSModel(factory)
    @constraint m 2x^4 + 5y^4 - x^2*y^2 >= -2(x^3*y + x + 1)
    JuMP.optimize!(m)
    @test JuMP.primal_status(m) == MOI.FeasiblePoint
    # p(x) = (x^2+2x)^2 + 2(1+x)^2 + 3
end


#=  Here I try to simulate a realistic college-admissions market, where schools
    start with arbitrary cutoffs and adjust them over time in response to demand.
    At the same time, schools also try to improve their programs, so the equilibrium
    cutoff vector varies dynamically and we get an interesting plot. =#

using DeferredAcceptance
using Plots

m = 6

γ = rand(m)
γ ./= sum(γ)
capacities = 0.8 * ones(m) / m
cutoffs_container = zeros(m, 1)
avg_improvement = .95 .+ 0.1 .* rand(m)
sort!(avg_improvement, rev=true)

"""
Example of a custom demand function. Same as that in `../MNLuniform/` but the
qualities get randomly perturbed each iteration.
"""
function MNL_single_score_demand(cutoffs)
    global cutoffs_container, γ
    cutoffs_container = hcat(cutoffs_container, cutoffs)

    sort_order = sortperm(cutoffs)
    cutoffs[sort_order]

    demands = zeros(m)

    prob_of_th = diff([cutoffs[sort_order]; 1])

    for c in 1:m, d in c:m     # For each score threshold
        demands[sort_order[c]] += prob_of_th[d] *
                                  γ[sort_order[c]] / sum(γ[sort_order[1:d]])
    end

    # println("Qualities:  ", γ)
    γ .*= avg_improvement .+ .1 .* randn(m)
    γ ./= sum(γ)

    return demands
end

# Don't worry about the maxit exceeded warning; equilibrium is in the distant future.
nonatomictatonnement(MNL_single_score_demand, capacities, maxit=300, β = 0., verbose=false)

p = plot(cutoffs_container'[2:end, :],
         label = reshape(["p$i, avg. imp. $(round(avg_improvement[i], digits=3))" for i in 1:m], 1, :),
         ls = [:dashdot :dot :solid],
         c = [:crimson :dodgerblue :rebeccapurple :olivedrab :goldenrod :teal],
         title = "Tâtonnement process with $m schools, one score,\nequal capacities, MNL choice, random quality improvements",
         titlefontsize = 12,
         xlabel = "iteration",
         ylabel = "score cutoff",
         ylim = (0, 1),
         legend = :bottom)

# savefig(p, string("examples/tatonnementovertime/plot.pdf"))
# savefig(p, string("examples/tatonnementovertime/plot.png"))


module FiscomTools
include("test.jl")
include("incremental.jl")
include("bisection.jl")
include("ndiff.jl")
include("newton.jl")
include("gauss.jl")
include("intsimp.jl")
include("lucrout.jl")
include("trapz.jl")
include("diff_operators.jl")

export incremental, bisection_it, bisection_err, newtonR # For root approximation
export ndiff # Numerical differentiation
export redGauss, lgsolve, lucrout # Matrices and systems of equations
export trapz,intsimp # Integration techniques
export d1hat, d2hat # Differential operators of the first and second numerical derivative
end


### A Pluto.jl notebook ###
# v0.14.7

using Markdown
using InteractiveUtils

# ╔═╡ 7fbfbd62-5e18-11ec-36b7-9f00610a9206
begin
  using Ripserer
  using Plots
end

# ╔═╡ cffa6493-e4bd-4769-a0ce-8137fee1fbfe
function noisy_circle(n; r1=1, r2=1, noise=0.1)
    points = NTuple{2,Float64}[]
    for _ in 1:n
        θ = 2π * rand()
        point = (
            r1 * sin(θ) + noise * rand() - noise / 2,
            r2 * cos(θ) + noise * rand() - noise / 2,
        )
        push!(points, point)
    end
    return points
end

# ╔═╡ 153f893a-6c95-4794-b1f6-6014f64f86ea
anim = @animate for _ in 1:200
    points = noisy_circle(100; noise=0)
    result = ripserer(points)

    plt_pts = scatter(
        points;
        legend=false,
        aspect_ratio=1,
        xlim=(-2.2, 2.2),
        ylim=(-2.2, 2.2),
        title="Data",
    )
    plt_diag = plot(result; infinity=3)

    plot(plt_pts, plt_diag; size=(800, 400))
end

# ╔═╡ dcc743e1-5ef7-4df5-a1a6-0ec1bc17a95e

# ╔═╡ Cell order:
# ╠═7fbfbd62-5e18-11ec-36b7-9f00610a9206
# ╠═cffa6493-e4bd-4769-a0ce-8137fee1fbfe
# ╠═153f893a-6c95-4794-b1f6-6014f64f86ea
# ╠═dcc743e1-5ef7-4df5-a1a6-0ec1bc17a95e


include("typedefs.jl")

function calc_D_MN(providers::Vector{Node}, clients::Vector{Node}, avg_speed::Float64)::Matrix{Int}
    m = length(providers)
    n = length(clients)

    D_MN = zeros(Int, (m, n))

    for k in range(1, stop=m)
        for i in range(1, stop=n)
            p = providers[k]
            c = clients[i]
            d_ki = round(hypot(p.cx-c.cx, p.cy-c.cy) / avg_speed)
            if (d_ki == 0) && (k != i)
                d_ki += 1
            end
            D_MN[k, i] = d_ki
        end
    end

    return D_MN
end

function calc_D_NN(clients::Vector{Node}, avg_speed::Float64)::Matrix{Int}
    n = length(clients)

    D_NN = zeros((n, n))

    for i in range(1, stop=n)
        for j in range(1, stop=n)
            c_outer = clients[i]
            c_inner = clients[j]
            d_ij = round(hypot(c_outer.cx-c_inner.cx, c_outer.cy-c_inner.cy) / avg_speed)
            if (d_ij == 0) && (i != j)
                d_ij += 1
            end
            D_NN[i, j] = d_ij
        end
    end

    return D_NN
end


using LinearAlgebra
using jInv.Mesh
using Test

print("   test getInterpolationMatrix (RegularMesh) ... ")
domain = zeros(6); domain[2:2:end] = rand(3).+1; domain[1:2:end] = -rand(3).-1
n1    = [2; 3; 4]
n2    = 2*n1
for dim=2:3
	print(" dim=",dim)
	Mc    = getRegularMesh(domain[1:2*dim],n1[1:dim])
	Mf    = getRegularMesh(domain[1:2*dim],n2[1:dim])
	Pcf   = getInterpolationMatrix(Mc,Mf)
	Pfc   = getInterpolationMatrix(Mf,Mc)
	Xc    = getCellCenteredGrid(Mc)
	Xf    = getCellCenteredGrid(Mf)

	@test isapprox(sum(Pfc,dims=2),ones(size(Pfc,1)))
	@test norm(Xc-Pfc*Xf)/norm(Xc) < 1e-13
end
print(" passed\n")
print("   test getInterpolationMatrix (TensorMesh) ... ")
x0    = randn(3)
h1    = rand(4)
h2    = rand(6)
h3    = rand(8)
H1    = h1[1:2:end]+h1[2:2:end]
H2    = h2[1:2:end]+h2[2:2:end]
H3    = h3[1:2:end]+h3[2:2:end]

Mct    = getTensorMesh3D(H1,H2,H3,x0)
Mft    = getTensorMesh3D(h1,h2,h3,x0)
Pcft   = getInterpolationMatrix(Mct,Mft)
Pfct   = getInterpolationMatrix(Mft,Mct)
Xc    = getCellCenteredGrid(Mct)
Xf    = getCellCenteredGrid(Mft)

@test isapprox(sum(Pfct,dims=2),ones(size(Pfct,1)))
@test norm(Xc-Pfct*Xf)/norm(Xc) < 1e-13
print("passed\n")


struct LorentzVectorCyl{T <: AbstractFloat}
    pt::T
    eta::T
    phi::T
    mass::T
end

Base.show(io::IO, v::LorentzVectorCyl) = print(io, "$(typeof(v))(pt=$(v.pt), eta=$(v.eta), phi=$(v.phi), mass=$(v.mass))")

Base.broadcastable(v::LorentzVectorCyl) = Ref(v)

"""
    zero(LorentzVectorCyl)
    zero(LorentzVectorCyl{T})
Constructs a zero four-vector.
"""
Base.zero(::Type{LorentzVectorCyl{T}}) where T = LorentzVectorCyl{T}(zero(T), zero(T), zero(T), zero(T))
Base.zero(::Type{LorentzVectorCyl}) = zero(LorentzVectorCyl{Float64})

px(v::LorentzVectorCyl) = v.pt * cos(v.phi)
py(v::LorentzVectorCyl) = v.pt * sin(v.phi)
pz(v::LorentzVectorCyl) = v.pt * sinh(v.eta)
energy(v::LorentzVectorCyl) = sqrt(px(v)^2 + py(v)^2 + pz(v)^2 + v.mass^2)

function *(v::LorentzVectorCyl{T}, k::Real) where T
    LorentzVectorCyl{T}(v.pt*k, v.eta, v.phi, v.mass*k)
end

function +(v1::LorentzVectorCyl{T}, v2::LorentzVectorCyl{W}) where {T,W}
    m1, m2 = max(v1.mass, zero(v1.pt)), max(v2.mass, zero(v2.pt))

    px1, px2 = px(v1), px(v2)
    py1, py2 = py(v1), py(v2)
    pz1, pz2 = pz(v1), pz(v2)
    e1 = sqrt(px1^2 + py1^2 + pz1^2 + m1^2)
    e2 = sqrt(px2^2 + py2^2 + pz2^2 + m2^2)

    sumpx = px1+px2
    sumpy = py1+py2
    sumpz = pz1+pz2

    ptsq = sumpx^2 + sumpy^2
    pt = sqrt(ptsq)
    eta = asinh(sumpz/pt)
    phi = atan(sumpy, sumpx)
    mass = sqrt(max(muladd(m1, m1, m2^2) + 2*e1*e2 - 2*(muladd(px1, px2, py1*py2) + pz1*pz2), zero(v1.pt)))
    return LorentzVectorCyl(pt,eta,phi,mass)
end

function fast_mass(v1::LorentzVectorCyl, v2::LorentzVectorCyl)
    # Calculate mass directly. Same as (v1+v2).mass except
    # this skips the intermediate pt, eta, phi calculations.
    # ~4x faster than (v1+v2).mass
    pt1, pt2 = v1.pt, v2.pt
    eta1, eta2 = v1.eta, v2.eta
    phi1, phi2 = v1.phi, v2.phi
    m1, m2 = v1.mass, v2.mass

    # note, massless approximation is
    # mass = sqrt(max(2*pt1*pt2*(cosh(eta1-eta2) - cos(phi1-phi2)), zero(pt1)))

    # @fastmath sinh() is very slow in some cases
    sinheta1 = sinh(eta1)
    sinheta2 = sinh(eta2)
    tpt12 = 2*pt1*pt2
    return @fastmath sqrt(max(muladd(m1, m1, m2^2)
        + 2*sqrt((pt1^2*(1+sinheta1^2) + m1^2)*(pt2^2*(1+sinheta2^2) + m2^2))
        - tpt12*sinheta1*sinheta2
        - tpt12*cos(phi1-phi2), zero(pt1)))
end

# https://root.cern.ch/doc/v606/GenVector_2VectorUtil_8h_source.html#l00061
@inline function deltaphi(v1::LorentzVectorCyl, v2::LorentzVectorCyl)
    dphi = v2.phi - v1.phi
    dphi = dphi - 2*pi*(dphi > pi) + 2*pi*(dphi <= -pi)
    return dphi
end
const Δϕ = deltaphi

@inline function deltaeta(v1::LorentzVectorCyl, v2::LorentzVectorCyl)
    return v2.eta - v1.eta
end
const Δη = deltaeta

@inline function deltar2(v1::LorentzVectorCyl, v2::LorentzVectorCyl)
    dphi = deltaphi(v1,v2)
    deta = deltaeta(v1,v2)
    return muladd(dphi, dphi, deta^2)
end
@inline deltar(v1::LorentzVectorCyl, v2::LorentzVectorCyl) = sqrt(deltar2(v1, v2))
const ΔR = deltar

function tocartesian(v::LorentzVectorCyl)
    x = px(v)
    y = py(v)
    z = pz(v)
    t = energy(v)
    (t=t, x=x, y=y, z=z)
end

function fromcartesian(t, x, y, z)
    pt = sqrt(muladd(x, x, y^2))
    eta = asinh(z/pt)
    phi = atan(y, x)
    mass = sqrt(max(t^2 - x^2 - y^2 - z^2, 0))
    return LorentzVectorCyl(pt,eta,phi,mass)
end


#'---
#'author: Garrett Smith
#'title: Parameter recovery with a hierarchical model
#'---
#'
#' # Parameter recovery test using a hierarchical model
#'
#' The goal of this script is to test the recovery of parameters from a hierarchical model
#' with a `fpdistribution` likelihood.

#+ echo=false
using Distributions
using Turing
using Zygote
Turing.setadbackend(:zygote)
using Plots, StatsPlots
using Pkg
Pkg.activate("../../FirstPassageTools.jl/")
using FirstPassageTools

#' ## Generating the true data
#'
#' First, we set up the transition rate matrices for the first-passage time distribution we
#' want to fit.

#T = 4*[-1.0 0 0; 1 -1 1; 0 1 -2]
T = 4*[-1 1.0 0; 1 -2 1; 0 1 -2]
A = 4*[0 0 1.0]
p0 = [1.0, 0, 0]

#' Scaling the transition rate matrices by `τ = 2.5` should give mean first-passage times of
#' around 400ms. Generating and fitting the paramters (τ and the separate τᵢ) will be done
#' on the log scale and then exponentiated in order to keep the transition rates positive.

nparticipants = 50
true_tau = 2.5
true_sd = 0.2
true_tau_i = exp.(rand(Normal(0, true_sd), nparticipants));

#' The data will be saved in wide format: Each participant's data corresponds to a row, and
#' each column is a data point.

#+ echo=false
ndata = 10
data = zeros(nparticipants, ndata)
param = true_tau .* true_tau_i
for i = 1:nparticipants
    data[i,:] = rand(fpdistribution(param[i]*T, param[i]*A, p0), ndata)
end

#' ## Specifying the model
#'
#' Now, we can write the full model including the likelihood. Note that we're using a
#' non-centered parameterization. This is because pilot simulations suggested that
#' sampling was biased in the centered parameterization.

# Switch to param = exp(tau) + exp(tau_i). This will prevent really small params b/c multiplication.
@model function mod(y, Tmat=T, Amat=A, p0vec=p0)
    np, nd = size(y)
    # Priors
    # Using the non-centered parameterization for τ
    τ ~ Normal()
    τ̂ = 1 + 0.1*τ  # Corresponds to Normal(1, 0.1)
    sd ~ Exponential(0.25)
    τᵢ ~ filldist(Normal(), np)
    τ̂ᵢ = sd .* τᵢ  # Corresponds to MvNormal(0, sd)
    # Likelihood
    mult = exp.(τ̂ .+ τ̂ᵢ)
    y ~ filldist(arraydist([fpdistribution(mult[p]*Tmat, mult[p]*Amat, p0vec) for p in 1:np]), nd)
    return τ̂, τ̂ᵢ, mult
end

#+ echo=false
#@model function mod_centered(y)
#    np, nd = size(y)
#    # Priors
#    # Using the centered parameterization for τ
#    τ ~ Normal(1, 0.1)
#    sd ~ Exponential(0.25)
#    τᵢ ~ filldist(Normal(0, sd), np)
#    # Likelihood
#    mult = exp.(τ .+ τᵢ)
#    y ~ filldist(arraydist([fpdistribution(mult[p]*T, mult[p]*A, p0) for p in 1:np]), nd)
#end
#
#@model function mod_noncentered_tau(y)
#    np, nd = size(y)
#    # Priors
#    # Using the centered parameterization for τ
#    τ ~ Normal()
#    τ̂ = 1.0 + 0.1*τ
#    sd ~ Exponential(0.25)
#    τᵢ ~ filldist(Normal(0, sd), np)
#    # Likelihood
#    mult = exp.(τ̂ .+ τᵢ)
#    y ~ filldist(arraydist([fpdistribution(mult[p]*T, mult[p]*A, p0) for p in 1:np]), nd)
#end
#
#@model function mod_noncentered_tau_i(y)
#    np, nd = size(y)
#    # Priors
#    # Using the centered parameterization for τ
#    τ ~ Normal(1, 0.1)
#    sd ~ Exponential(0.25)
#    τᵢ ~ filldist(Normal(), np)
#    τ̂ᵢ = sd .* τᵢ
#    # Likelihood
#    mult = exp.(τ .+ τ̂ᵢ)
#    y ~ filldist(arraydist([fpdistribution(mult[p]*T, mult[p]*A, p0) for p in 1:np]), nd)
#end
#' ## Sampling
#'
#' Here, we'll use the NUTS sampler with a burnin of 100 samples and an acceptance rate of
#' 0.65 posterior. We'll use four chains of 1000 samples each. Make sure to execute this
#' script with `julia -t 4 HierarchicalParameterRecovery.jl`.

#posterior = sample(mod(data), NUTS(250, 0.65), MCMCThreads(), 1000, 4)
posterior = sample(mod(data), NUTS(250, 0.7), 1000);
#posterior_centered = sample(mod_centered(data), NUTS(100, 0.65), 1000);
#posterior_noncentered_tau = sample(mod_noncentered_tau(data), NUTS(100, 0.65), 1000);
#posterior_noncentered_tau_i = sample(mod_noncentered_tau_i(data), NUTS(100, 0.65), 1000);

#' ## Evaluating parameter recovery
#'
#' First, we summarize the chains:

posterior

#' Let's also look at the Gelman-Rubin statistic for the chains:
#gelmandiag(posterior)

# #' And the centered version:
#posterior_centered

# #' τ non-centered, τᵢ centered:
#posterior_noncentered_tau

# #' τ centered, τᵢ non-centered
#posterior_noncentered_tau_i

#' And plot histograms of the parameters on the millisecond scale:
#+ echo=false
gen = generated_quantities(mod(data), posterior);
gen = vcat(gen...);
gen_tau = [x[1] for x in gen];
gen_tau_i = transpose(hcat([x[2] for x in gen]...));
shrunken = transpose(hcat([x[3] for x in gen]...));

#+ echo=false; fig_cap="Posterior distribution of τ. Vertical line shows the true value"
histogram(exp.(gen_tau), xlabel="τ")
vline!([true_tau])

#+ echo=false; fig_cap="Posterior distribution of the SD. Vertical lines show the true value."
histogram(posterior[:sd][:], xlabel="Std. deviation of the τᵢ")
vline!([true_sd], label="True value")

#+ echo=false; fig_cap="Posterior distributions of the τᵢ. Vertical lines show the true value."
plot_vec = [];
for p = 1:nparticipants
    plt = histogram(exp.(gen_tau_i[:,p]), xticks=false, yticks=false)
    plt = vline!(plt, [true_tau_i[p]])
    push!(plot_vec, plt)
end
plot(plot_vec..., legend=false)

#+ echo=false; fig_cap="Posterior distributions of τ*τᵢ. Vertical lines show the true value."
plot_vec = [];
for p = 1:nparticipants
    plt = histogram(shrunken[:,p], xticks=false, yticks=false)
    plt = vline!(plt, [param[p]])
    push!(plot_vec, plt)
end
plot(plot_vec..., legend=false)

#' If the posterior contains the true values of the parameters, we can say the parameters
#' were recovered successfully.
#'


using LinearAlgebra ## "norm"
using StatsBase ## "sample"
using Statistics ## "mean"

## Some parts have been vectorized for better performance.
"""
**BFO** is a function that computes the **classical (non-adaptive) BFO (Bacterial Foraging Optimization)**

**Inputs**:
1. **J**: a function with domain R^n
2. **Range**: exploration range: Range^n

3. **S** = 10, number of bacteria
4. **Sr** = 4,  number of bacteria removed in reproductive step
5. **Nc** = 20, number of chemotactic steps
6. **Ns** = 5, number of swim steps
7. **Nre** = 5, number of reproductive steps
8. **Ned** = 2, elimination and dispersal steps
9. **Ped** = 0.3, probability of elimination
10. **Ci** = (mean(Range[:,2].-Range[:,1])/S), the run-length unit
**Output**: a dictionary that stores
1. the minimum value of J
2. the point achieving this minimu value
3. the path of each bacterium (for plotting illustration)
"""
function BFO(J, Range::Array{Float64,2}; S = 10::Int, Sr = 4::Int, Nc = 20::Int, Ns = 5::Int,
        Nre = 5::Int, Ned = 2::Int, Ped = 0.3::Float64, Ci = (mean(Range[:,2].-Range[:,1])/S)::Float64)

    n = size(Range,1) ## dimension of the input of J
    ## randomly generate S bacteria in Range
    B_loc = RandPts(Range,S)
    ## a dictionary recording the path of bacterium i
    Path_Dict = Dict{Int,Array{Float64,2}}(i=>[zeros(n,0) B_loc[:,i]] for i=1:S)
    for l = 1:Ned ## index of elimination-dispersal steps
        for k = 1:Nre ## index of reproductive steps
            for j = 1:Nc ## index of chemotactic steps
                ## Chemotactic Step; variable used: Range, S, B_loc, Ci, J
                println("Chemotactic step $j")
                @time Chemo = ChemotacticStep(J, Range::Array{Float64,2}, S::Int, Ns::Int, Ci::Float64, B_loc::Array{Float64,2})
                B_loc = Chemo[1]
                for i=1:S
                    Path_Dict[i] = [Path_Dict[i] Chemo[2][i]]
                end
            end ## end of chemotactic steps

            if k<Nre
                ## Reproductive Step; variable used: B_loc, Sr, S, J
                println("Reproduction step $k")
                @time Repro = ReproductionStep(J, B_loc::Array{Float64,2}, Sr::Int, S::Int, n::Int)
                B_loc = Repro[1]
                for i=1:S
                    Path_Dict[i] = [Path_Dict[i] Repro[2][i]]
                end
            end
        end
        ## Elimination-Dispersal Step; variable used: Ped, Range, B_loc
        println("Elimination-Dispersal step $l")
        if l<Ned
            @time EliDis = ElimDispStep(B_loc, Ped, Range, S, n)
            B_loc = EliDis[1]
            for i=1:S
                Path_Dict[i] = [Path_Dict[i] EliDis[2][i]]
            end
        end
    end
    B_best = Int(sortslices([[J(B_loc[:,i]) for i=1:S] collect(1:S)], dims = 1)[1,2]) ## best bacterium
    X_best = B_loc[:,B_best] ## best location for minimizing J
    J_best = J(X_best) ## best (minimum) J value
    return Dict("Minimum"=>J_best, "Minimum Point"=>X_best, "Path_Dict"=>Path_Dict)
end

"""
**ChemotacticStep** is a function performing the chemotactic step in BFO.
"""
function ChemotacticStep(J, Range::Array{Float64,2}, S::Int, Ns::Int, Ci::Float64, B_loc::Array{Float64,2})
    Path = Dict{Int, Array{Float64,2}}(i=>zeros(size(Range,1),2) for i=1:S)
    ToSwim = collect(1:S) ## monitor the swimming bacteria
    Tumble = RandUnit(Range,S)
    m = 0 ## index of swimming
    while (length(ToSwim)!=0)&(m<Ns) ## tumble/swim
        m = m + 1
        J_old = [J(B_loc[:,i]) for i in ToSwim]

        B_loc_new = copy(B_loc)
        B_loc_new[:,ToSwim] = B_loc[:,ToSwim].+Ci*Tumble[:,ToSwim]

        ## Mirror back the out-of-range bacteria
        ## Range_tail = Range[:,2]
        ## Range_head = Range[:,1]
        ## 1. mirror back bacteria tumbling out from below
        B_loc_new[:,ToSwim] = (abs.(B_loc_new[:,ToSwim].-Range[:,1])).+Range[:,1]
        ## 2. mirror back bacteria tumbling out from above
        B_loc_new[:,ToSwim] = Range[:,2].-(abs.(Range[:,2].-B_loc_new[:,ToSwim]))

        ## Evaluate J at new locations
        J_new = [J(B_loc_new[:,i]) for i in ToSwim]

        ## Find out improved bacteria and update the swimming ones
        ToSwim = ToSwim[findall(J_new.<J_old)]

        for i in ToSwim
            Path[i] = [Path[i] B_loc_new[:,i]]
        end
        B_loc[:,ToSwim] = B_loc_new[:,ToSwim]
    end ## end of tumble/swim
    return (B_loc,Path)
end

"""
**ReproductionStep** is a function performing the reproduction step in BFO.
"""
function ReproductionStep(J, B_loc::Array{Float64,2}, Sr::Int, S::Int, n::Int)
    Path = Dict{Int,Array{Float64,2}}(i=>zeros(n,2) for i=1:S)
    ## (!!!) I define the health of a bacterium as the J value of its current location
    Health = [J(B_loc[:,i]) for i=1:S]
    Health_sort = sortslices([Health collect(1:S)], dims = 1) ## sort the bacteria according to health
    B_survive = Array{Int,1}(Health_sort[1:S-Sr,2]) ## ## pick out the most healthy (S-Sr) bacteria
    B_die = setdiff(collect(1:S), B_survive)
    B_loc[:,B_die] = B_loc[:,B_survive[1:Sr]] ## reproduce the most healthy Sr bacteria
    for i in B_die
        Path[i] = [Path[i] B_loc[:,i]]
    end
    return (B_loc, Path)
end

"""
**ElimDispStep** is a function performing the elimination-dispersal step in BFO.
"""
function ElimDispStep(B_loc::Array{Float64,2}, Ped::Float64, Range::Array{Float64,2}, S::Int, n::Int)
    Path = Dict{Int,Array{Float64,2}}(i=>zeros(n,0) for i=1:S)
    Alive = findall([sample([false, true], aweights([Ped,1-Ped])) for i=1:S]) # true = alive, false = kill
    Killed = setdiff(collect(1:S),Alive) ## kiilled bacteria
    ## Rebuild the killed bacteria in random locations in Range
    B_loc[:,Killed] = RandPts(Range,length(Killed))
    ## update the path dictionary
    for i in Killed
        Path[i] = [Path[i] B_loc[:,i]]
    end
    return (B_loc, Path)
end

"""
**RandPts** is a function that randomly uniformly generates m points in dimension size(Range,2) inside Range.

Inputs:

m: number of random points to generate

Range: dim x 2 matrix; the d-th row is the range of the d-th coordinate

Output:

an dim x m matrix; each column is a desired random vector.

"""
function RandPts(Range::Array{Float64,2}, m::Int)
    ## dim = size(Range,1)
    ## randPts = rand(dim,m)
    ## Coeff = (Range[:,2].-Range[:,1])
    ## Intercept = Range[:,1]
    ## randPts = ((Range[:,2].-Range[:,1]).*rand(size(Range,1),m)).+Range[:,1]
    return ((Range[:,2].-Range[:,1]).*rand(size(Range,1),m)).+Range[:,1]
end

"""
**RandUnit** is a function generating m random unit vectors in dimension d. The vector generated is uniform with respect to a Range where random walk is to be applied.

**Inputs**:

1. Range: a d x m matrix; the kth row is the range for the kth variable.
2. m: number of vectors to generate

**Output**:

an d x m matrix whose columns are the desired vectors.

"""
function RandUnit(Range::Array{Float64,2}, m::Int)
    ## d = size(Range,1)
    ## Scales = Range[:,2].-Range[:,1]
    Rand = (rand(size(Range,1),m).-1/2).*(Range[:,2].-Range[:,1])
    for k = 1:m
        Rand[:,k]./=norm(Rand[:,k])
    end
    return Rand
end


include("SortSparse.jl")
include("KoLesky.jl")
include("Cholesky.jl")
using LinearAlgebra
using RecursiveArrayTools
using Distributions
using Distances
using Random
using Profile
using JLD
using MAT
using BlockArrays
Random.seed!(123)

NTrain = 10000
NTest = 5
N = NTest + NTrain
d =  2
y = randn(N)
x = randn(d, N)

setSplit = randperm(N)
trainInds = setSplit[1:NTrain]
testInds = setSplit[(NTrain + 1) : end]

xTrain = x[ :, trainInds ]
yTrain = y[ trainInds ]
xTest = x[:, testInds ]
yTest = y[ testInds ]

n = 10
h = 1 / (n - 1)

xTrain = Matrix(VectorOfArray([[i * h; j * h] for i in 1:n, j in 1:n][:]))
NTrain = size(xTrain,2)
yTrain = randn(NTrain)
N = NTrain + NTest

l = 1.0
σ = 1e-15
cov = r -> exp(-r/l) # + Float64(r <= 1e-10)

rhovals = [2.0, 3.0, 4.0, 5.0]

ρ = 3.0
λ = 1.3

# reordering points and forming skeletons
@time colptr, rowval, P, revP, distances = sortSparseRev(xTrain,ρ, 1)
xTrain = xTrain[:, P]
yTrain = yTrain[P]
x = hcat(xTest, xTrain)
skeletons = construct_skeletons(colptr, rowval, distances, λ)

# Creating
Θ = cov.(pairwise(Euclidean(), x, dims=2))
L_Θ = cholesky(Matrix(Θ))
y = L_Θ.L * randn(N)
Θ = Θ + σ * I

Θ = PseudoBlockArray(Θ, [NTest, NTrain], [NTest, NTrain])
y = PseudoBlockArray(y, [NTest, NTrain])
y[Block(2)] += √(σ) * randn(NTrain)

L = assembleL(skeletons, xTrain, cov)
U = sparse(triu(L * L') + I / max(σ, 1e-30))
icholU_high_level!(U)

μTrue = Θ[Block(1,2)] * (Θ[Block(2,2)] \ y[Block(2)])

# Direct application of Vecchia
# μ = Θ[Block(1,2)] * (L * (L' * y[Block(2)]))

# Version with Cholesky factor
μ = Θ[Block(1,2)] * (σ \ (  U \ (U' \ (L * (L' * y[Block(2)])))))

display(μTrue)
display(μ)
display(y[Block(1)])

@show norm(μTrue - μ) / norm(μTrue)


# Use DiffEqOperators to generate weights and calculate derivative orders
struct DifferentialDiscretizer{T,D1}
    approx_order::Int
    map::D1
    halfoffsetmap::Tuple{Dict,Dict}
    windmap::Tuple{Dict,Dict}
    interpmap::Dict{Num,DiffEqOperators.DerivativeOperator}
    orders::Dict{Num,Vector{Int}}
end

function DifferentialDiscretizer(pdesys, s, discretization)
    pdeeqs = pdesys.eqs isa Vector ? pdesys.eqs : [pdesys.eqs]
    bcs = pdesys.bcs isa Vector ? pdesys.bcs : [pdesys.bcs]
    approx_order = discretization.approx_order
    upwind_order = discretization.upwind_order
    d_orders(x) = reverse(sort(collect(union((differential_order(pde.rhs, x) for pde in pdeeqs)..., (differential_order(pde.lhs, x) for pde in pdeeqs)..., (differential_order(bc.rhs, x) for bc in bcs)..., (differential_order(bc.lhs, x) for bc in bcs)...))))

    # central_deriv_rules = [(Differential(s)^2)(u) => central_deriv(2,II,j,k) for (j,s) in enumerate(s.x̄), (k,u) in enumerate(s.ū)]
    differentialmap = Array{Pair{Num,DiffEqOperators.DerivativeOperator},1}()
    nonlinlap_inner = []
    nonlinlap_outer = []
    windpos = []
    windneg = []
    interp = []
    orders = []
    # Hardcoded to centered difference, generate weights for each differential

    for x in s.x̄
        orders_ = d_orders(x)
        push!(orders, x => orders_)
        _orders = Set(vcat(orders_, [1, 2]))

        if s.grid[x] isa StepRangeLen # Uniform grid case
            # TODO: Only generate weights for derivatives that are actually used and avoid redundant calculations
            rs = [(Differential(x)^d) => CompleteCenteredDifference(d, approx_order, s.dxs[x]) for d in _orders]

            windpos = vcat(windpos, [(Differential(x)^d) => CompleteUpwindDifference(d, upwind_order, s.dxs[x], 0) for d in orders_[isodd.(orders_)]])
            windneg = vcat(windneg, [(Differential(x)^d) => CompleteUpwindDifference(d, upwind_order, s.dxs[x], d + upwind_order - 1) for d in orders_[isodd.(orders_)]])

            nonlinlap_inner = vcat(nonlinlap_inner, [Differential(x)^d => CompleteHalfCenteredDifference(d, approx_order, s.dxs[x]) for d in _orders])
            nonlinlap_outer = push!(nonlinlap_outer, Differential(x) => CompleteHalfCenteredDifference(1, approx_order, s.dxs[x]))
            differentialmap = vcat(differentialmap, rs)
            # A 0th order derivative off the grid is an interpolation
            push!(interp, x => CompleteHalfCenteredDifference(0, approx_order, s.dxs[x]))

        elseif s.grid[x] isa AbstractVector # The nonuniform grid case
            rs = [(Differential(x)^d) => CompleteCenteredDifference(d, approx_order, s.grid[x]) for d in _orders]

            windpos = vcat(windpos, [(Differential(x)^d) => CompleteUpwindDifference(d, upwind_order, s.grid[x], 0) for d in orders_[isodd.(orders_)]])
            windneg = vcat(windneg, [(Differential(x)^d) => CompleteUpwindDifference(d, upwind_order, s.grid[x], d + upwind_order - 1) for d in orders_[isodd.(orders_)]])

            discx = s.grid[x]
            nonlinlap_inner = vcat(nonlinlap_inner, [Differential(x)^d => CompleteHalfCenteredDifference(d, approx_order, s.grid[x]) for d in _orders])
            nonlinlap_outer = push!(nonlinlap_outer, Differential(x) => CompleteHalfCenteredDifference(1, approx_order, [(discx[i+1] + discx[i]) / 2 for i in 1:length(discx)-1]))
            differentialmap = vcat(differentialmap, rs)
            # A 0th order derivative off the grid is an interpolation
            push!(interp, x => CompleteHalfCenteredDifference(0, approx_order, s.grid[x]))
        else
            @assert false "s.grid contains nonvectors"
        end
    end

    return DifferentialDiscretizer{eltype(orders),typeof(Dict(differentialmap))}(approx_order, Dict(differentialmap), (Dict(nonlinlap_inner), Dict(nonlinlap_outer)), (Dict(windpos), Dict(windneg)), Dict(interp), Dict(orders))
end

"""
Performs a centered difference in `x` centered at index `II` of `u`
ufunc is a function that returns the correct discretization indexed at Itap, it is designed this way to allow for central differences of arbitrary expressions which may be needed in some schemes
"""
function central_difference(D::DerivativeOperator{T,N,Wind,DX}, II, s, b, jx, u, ufunc) where {T,N,Wind,DX<:Number}
    j, x = jx
    ndims(u, s) == 0 && return Num(0)
    # unit index in direction of the derivative
    I1 = unitindex(ndims(u, s), j)
    # offset is important due to boundary proximity

    if (II[j] <= D.boundary_point_count) & (b isa Val{false})
        weights = D.low_boundary_coefs[II[j]]
        offset = 1 - II[j]
        Itap = [II + (i + offset) * I1 for i in 0:(D.boundary_stencil_length-1)]
    elseif (II[j] > (length(s, x) - D.boundary_point_count)) & (b isa Val{false})
        weights = D.high_boundary_coefs[length(s, x)-II[j]+1]
        offset = length(s, x) - II[j]
        Itap = [II + (i + offset) * I1 for i in (-D.boundary_stencil_length+1):1:0]
    else
        weights = D.stencil_coefs
        Itap = [wrapperiodic(II + i * I1, s, b, u, jx) for i in half_range(D.stencil_length)]
    end
    # Tap points of the stencil, this uses boundary_point_count as this is equal to half the stencil size, which is what we want.

    return dot(weights, ufunc(u, Itap, x))
end

function central_difference(D::DerivativeOperator{T,N,Wind,DX}, II, s, b, jx, u, ufunc) where {T,N,Wind,DX<:AbstractVector}
    j, x = jx
    @assert b isa Val{false} "Periodic boundary conditions are not yet supported for nonuniform dx dimensions, such as $x, please post an issue to https://github.com/SciML/MethodOfLines.jl if you need this functionality."
    ndims(u, s) == 0 && return Num(0)
    # unit index in direction of the derivative
    I1 = unitindex(ndims(u, s), j)
    # offset is important due to boundary proximity

    if (II[j] <= D.boundary_point_count) & (b isa Val{false})
        weights = D.low_boundary_coefs[II[j]]
        offset = 1 - II[j]
        Itap = [II + (i + offset) * I1 for i in 0:(D.boundary_stencil_length-1)]
    elseif (II[j] > (length(s, x) - D.boundary_point_count)) & (b isa Val{false})
        weights = D.high_boundary_coefs[length(s, x)-II[j]+1]
        offset = length(s, x) - II[j]
        Itap = [II + (i + offset) * I1 for i in (-D.boundary_stencil_length+1):1:0]
    else
        weights = D.stencil_coefs[II[j]-D.boundary_point_count]
        Itap = [wrapperiodic(II + i * I1, s, b, u, jx) for i in half_range(D.stencil_length)]
    end
    # Tap points of the stencil, this uses boundary_point_count as this is equal to half the stencil size, which is what we want.

    return dot(weights, ufunc(u, Itap, x))
end

@inline function _upwind_difference(D::DerivativeOperator{T,N,Wind,DX}, II, s, b, ispositive, u, jx) where {T,N,Wind,DX<:Number}
    j, x = jx
    I1 = unitindex(ndims(u, s), j)
    if !ispositive
        if (II[j] > (length(s, x) - D.boundary_point_count)) & (b isa Val{false})
            weights = D.high_boundary_coefs[length(s, x)-II[j]+1]
            offset = length(s, x) - II[j]
            Itap = [II + (i + offset) * I1 for i in (-D.boundary_stencil_length+1):0]
        else
            weights = D.stencil_coefs
            Itap = [wrapperiodic(II + i * I1, s, b, u, jx) for i in 0:D.stencil_length-1]
        end
    else
        if (II[j] <= D.offside) & (b isa Val{false})
            weights = D.low_boundary_coefs[II[j]]
            offset = 1 - II[j]
            Itap = [II + (i + offset) * I1 for i in 0:(D.boundary_stencil_length-1)]
        else
            weights = D.stencil_coefs
            Itap = [wrapperiodic(II + i * I1, s, b, u, jx) for i in -D.stencil_length+1:0]
        end
    end
    return weights, Itap
end

@inline function _upwind_difference(D::DerivativeOperator{T,N,Wind,DX}, II, s, b, ispositive, u, jx) where {T,N,Wind,DX<:AbstractVector}
    j, x = jx
    @assert b isa Val{false} "Periodic boundary conditions are not yet supported for nonuniform dx dimensions, such as $x, please post an issue to https://github.com/SciML/MethodOfLines.jl if you need this functionality."
    I1 = unitindex(ndims(u, s), j)
    if !ispositive
        @assert D.offside == 0

        if (II[j] > (length(s, x) - D.boundary_point_count))
            weights = D.high_boundary_coefs[length(s, x)-II[j]+1]
            offset = length(s, x) - II[j]
            Itap = [II + (i + offset) * I1 for i in (-D.boundary_stencil_length+1):0]
        else
            weights = D.stencil_coefs[II[j]]
            Itap = [II + i * I1 for i in 0:D.stencil_length-1]
        end
    else
        if (II[j] <= D.offside)
            weights = D.low_boundary_coefs[II[j]]
            offset = 1 - II[j]
            Itap = [II + (i + offset) * I1 for i in 0:(D.boundary_stencil_length-1)]
        else
            weights = D.stencil_coefs[II[j] - D.offside]
            Itap = [II + i * I1 for i in -D.stencil_length+1:0]
        end
    end
    return weights, Itap
end

function upwind_difference(d::Int, II::CartesianIndex, s::DiscreteSpace, b, derivweights, jx, u, ufunc, ispositive)
    j, x = jx
    ndims(u, s) == 0 && return Num(0)
    D = !ispositive ? derivweights.windmap[1][Differential(x)^d] : derivweights.windmap[2][Differential(x)^d]
    #@show D.stencil_coefs, D.stencil_length, D.boundary_stencil_length, D.boundary_point_count
    # unit index in direction of the derivative
    weights, Itap = _upwind_difference(D, II, s, b, ispositive, u, jx)
    return dot(weights, ufunc(u, Itap, x))
end

"""
Get the weights and stencil for the inner half offset centered difference for the nonlinear laplacian for a given index and differentiating variable.

Does not discretize so that the weights can be used in a replacement rule
TODO: consider refactoring this to harmonize with centered difference

Each index corresponds to the weights and index for the derivative at index i+1/2
"""
function get_half_offset_weights_and_stencil(D::DerivativeOperator{T,N,Wind,DX}, II::CartesianIndex, s::DiscreteSpace, b, u, jx, len=0) where {T,N,Wind,DX<:Number}
    j, x = jx
    len = len == 0 ? length(s, x) : len
    @assert II[j] != length(s, x)

    # unit index in direction of the derivative
    I1 = unitindex(ndims(u, s), j)
    # offset is important due to boundary proximity

    if (II[j] <= D.boundary_point_count) & (b isa Val{false})
        weights = D.low_boundary_coefs[II[j]]
        offset = 1 - II[j]
        Itap = [II + (i + offset) * I1 for i in 0:(D.boundary_stencil_length-1)]
    elseif (II[j] > (len - D.boundary_point_count)) & (b isa Val{false})
        weights = D.high_boundary_coefs[len-II[j]]
        offset = len - II[j]
        Itap = [II + (i + offset) * I1 for i in (-D.boundary_stencil_length+1):0]
    else
        weights = D.stencil_coefs
        Itap = [wrapperiodic(II + i * I1, s, b, u, jx) for i in (1-div(D.stencil_length, 2)):(div(D.stencil_length, 2))]
    end

    return (weights, Itap)
end

function get_half_offset_weights_and_stencil(D::DerivativeOperator{T,N,Wind,DX}, II::CartesianIndex, s::DiscreteSpace, b, u, jx, len=0) where {T,N,Wind,DX<:AbstractVector}
    j, x = jx
    @assert b isa Val{false} "Periodic boundary conditions are not yet supported for nonuniform dx dimensions, such as $x, please post an issue to https://github.com/SciML/MethodOfLines.jl if you need this functionality."
    len = len == 0 ? length(s, x) : len
    @assert II[j] != length(s, x)

    # unit index in direction of the derivative
    I1 = unitindex(ndims(u, s), j)
    # offset is important due to boundary proximity

    if (II[j] <= D.boundary_point_count)
        weights = D.low_boundary_coefs[II[j]]
        offset = 1 - II[j]
        Itap = [II + (i + offset) * I1 for i in 0:(D.boundary_stencil_length-1)]
    elseif (II[j] > (len - D.boundary_point_count))
        weights = D.high_boundary_coefs[len-II[j]]
        offset = len - II[j]
        Itap = [II + (i + offset) * I1 for i in (-D.boundary_stencil_length+1):0]
    else
        weights = D.stencil_coefs[II[j]-D.boundary_point_count]
        Itap = [II + i * I1 for i in (1-div(D.stencil_length, 2)):(div(D.stencil_length, 2))]
    end

    return (weights, Itap)
end

# i is the index of the offset, assuming that there is one precalculated set of weights for each offset required for a first order finite difference
function half_offset_centered_difference(D, II, s, b, jx, u, ufunc)
    ndims(u, s) == 0 && return Num(0)
    j, x = jx
    weights, Itap = get_half_offset_weights_and_stencil(D, II, s, b, u, jx)
    return dot(weights, ufunc(u, Itap, x))
end


module day07

using LightGraphs
using SimpleWeightedGraphs
using DataStructures

include("inputs.jl")

# Approach: problem is a search over a directed, weighted graph.

input = inputs.input_d7()
ex = inputs.ex1_d7()

function build_graph(input)
    """
    Builds weighted directed graph, g, noting which bags (src vertices) and can contain certain quantities (edge weights) of other bags (dest vertices). Since graph data type cannot store vertex labels, a dictionary is used to map bag names (strings) to vertex numbers."""

    g = SimpleWeightedDiGraph{Int, Int}()
    names_to_verts = Dict()

    function build_bagtag_dict(str)
        """ Builds dictionary of (bag_name=>vertex_number) pairs as a side effect (CODE SMELL).
        Returns vertex number of bag_name in names_to_verts dict as an integer.
        Creates new vertex if needed."""
        default = length(vertices(g))+1
        res = get!(names_to_verts, str, default)
        res == default && add_vertex!(g) # if node doesn't exist add new vert
        return res
    end
    for line in input
        thisbag, rest = split(line, " bags contain ")
        src_vert = build_bagtag_dict(thisbag)
        rest == "no other bags." && continue
        edges = map(split(rest, r" bags?(, |.)", keepempty=false)) do edge
            weight, innerbag = match(r"^(\d+) (.*)$", edge).captures
            weight = parse(Int, weight)
            dst_vert = build_bagtag_dict(innerbag)
            # add_edge!(g, src_vert, dst_vert, weight) || error("Couldn't add edge! ", join((thisbag, innerbag, dst_vert, length(g), src_vert, g), ", "))
            add_edge!(g, dst_vert, src_vert, weight) || error("Couldn't add edge! ", join((thisbag, innerbag, dst_vert, length(g), src_vert, g), ", "))
        end
    end
    verts_to_names = SortedDict([reverse(pair) for pair in pairs(names_to_verts)])
    g, names_to_verts, verts_to_names
end

# Part 1
function count_parents(g, start, parentset=Set{Int}())
    parents = neighbors(g, start)
    for parent in parents
        push!(parentset, parent)
        count_parents(g, parent, parentset)
    end
    return length(parentset)
end

function part1(input)
    g, ntv, vtn = build_graph(input)
    return count_parents(g, ntv["shiny gold"])
end

@assert part1(ex) == 4
@info "Day 7, Part 1 answer: $(part1(input))"

# Part 2
function count_child_bags(g, start)
    bagcount = 0
    children = inneighbors(g, start)
    for child in children
        # println("from parent: ", vtn[start], " to child: ", vtn[child], " with weight: ", g.weights[start, child])
        bagcount += g.weights[start, child] * (1 + count_child_bags(g, child))
    end
    return bagcount
end

function part2(input)
    g, ntv, vtn = build_graph(input)
    return count_child_bags(g, ntv["shiny gold"])
end

@assert part2(ex) == 32
@info "Day 7, Part 2 answer: $(part2(input))"

# Plots are always fun.
function plot_graph(g)
    gplot(SimpleDiGraph(g), nodelabel=values(rntn))
end

end


using Printf
using BenchmarkTools
function heat()
    N = 1001
    T0 = Matrix{Float64}(undef,N,N)
    T1 = Matrix{Float64}(undef,N,N)
    x  = Matrix{Float64}(undef,N,N)
    y  = Matrix{Float64}(undef,N,N)
    a = 0
    b = π;
    dx = (b-a)/(N-1)
    for j=1:N
        for i=1:N
            x[i,j] = (i-1)*dx
        end
    end
    for j=1:N
        for i=1:N
            y[i,j] = (j-1)*dx
        end
    end

    dt = dx
    tmax = (10-mod(10,dt))/dt
    k = 0.25*dx
    alpha = k*dt/dx^2
    for i=1:N
        T1[1,i] = 1.0
        T1[N,i] = 1.0
        T1[i,1] = 1.0
        T1[i,N] = 1.0
    end
    for i=1:N
        T0[1,i] = 1.0
        T0[N,i] = 1.0
        T0[i,1] = 1.0
        T0[i,N] = 1.0
    end
    t = 0*dt
    vel = open("Julia.dat", "w")
    write(vel, "title =\"ZoneTime_",string(t),"\""," \n")
    write(vel, "variables = \"x\", \"y\", \"T\""," \n")
    write(vel, "zone T=\"Zone_"*string(t)*"\" i=",@sprintf("%d",N)," j=",@sprintf("%d",N)," \n")
    for j=1:N
        for i=1:N
            write(vel, @sprintf("%1.9e",x[i,j])," ",@sprintf("%1.9e",y[i,j])," ",@sprintf("%1.9e",T0[i,j])," \n")
        end
    end

    for k=1:tmax
        for j=2:N-1
            for i=2:N-1
                T1[i,j] = T0[i,j]+alpha*((T0[i+1,j]-2*T0[i,j]+T0[i-1,j])+(T0[i,j+1]-2*T0[i,j]+T0[i,j-1]))
            end
        end
        T0 .= T1
        if mod(k,500)==0
            write(vel, "title =\"ZoneTime_",string(t),"\""," \n")
            write(vel, "variables = \"x\", \"y\", \"T\""," \n")
            write(vel, "zone T=\"Zone_"*string(t)*"\" i=",@sprintf("%d",N)," j=",@sprintf("%d",N)," \n")
            for j=1:N
                for i=1:N
                    write(vel, @sprintf("%1.9e",x[i,j])," ",@sprintf("%1.9e",y[i,j])," ",@sprintf("%1.9e",T0[i,j])," \n")
                end
            end
        end
    end
    t = tmax*dt
    write(vel, "title =\"ZoneTime_",string(t),"\""," \n")
    write(vel, "variables = \"x\", \"y\", \"T\""," \n")
    write(vel, "zone T=\"Zone_"*string(t)*"\" i=",@sprintf("%d",N)," j=",@sprintf("%d",N)," \n")
    for j=1:N
        for i=1:N
            write(vel, @sprintf("%1.9e",x[i,j])," ",@sprintf("%1.9e",y[i,j])," ",@sprintf("%1.9e",T0[i,j])," \n")
        end
    end
    close(vel)
end

function heat_Vec()
    N = 1001
    T0 = Matrix{Float64}(undef,N,N)
    T1 = Matrix{Float64}(undef,N,N)
    x  = Matrix{Float64}(undef,N,N)
    y  = Matrix{Float64}(undef,N,N)
    a = 0
    b = π;
    dx = (b-a)/(N-1)
    x[:,1] = (0:N-1)*dx
    x[:,:] .= x[:,1]
    y[1,:] = (0:N-1)*dx
    y[:,:] .= y[1,:]

    dt = dx
    tmax = (10-mod(10,dt))/dt
    k = 0.25*dx
    alpha = k*dt/dx^2
    T1[1,:] .= 1.0
    T1[N,:] .= 1.0
    T1[:,1] .= 1.0
    T1[:,N] .= 1.0

    T0[1,:] .= 1.0
    T0[N,:] .= 1.0
    T0[:,1] .= 1.0
    T0[:,N] .= 1.0
    t = 0*dt
    vel = open("Julia.dat", "w")
    write(vel, "title =\"ZoneTime_",string(t),"\""," \n")
    write(vel, "variables = \"x\", \"y\", \"T\""," \n")
    write(vel, "zone T=\"Zone_"*string(t)*"\" i=",@sprintf("%d",N)," j=",@sprintf("%d",N)," \n")
    for j=1:N
        for i=1:N
            write(vel, @sprintf("%1.9e",x[i,j])," ",@sprintf("%1.9e",y[i,j])," ",@sprintf("%1.9e",T0[i,j])," \n")
        end
    end

    for k=1:tmax
        T1[2:N-1,2:N-1] .= T0[2:N-1,2:N-1].+alpha.*((T0[3:N,2:N-1].-2 .*T0[2:N-1,2:N-1].+T0[1:N-2,2:N-1]).+
                                                    (T0[2:N-1,3:N].-2 .*T0[2:N-1,2:N-1].+T0[2:N-1,1:N-2]))
        T0 .= T1
        if mod(k,500)==0
            write(vel, "title =\"ZoneTime_",string(t),"\""," \n")
            write(vel, "variables = \"x\", \"y\", \"T\""," \n")
            write(vel, "zone T=\"Zone_"*string(t)*"\" i=",@sprintf("%d",N)," j=",@sprintf("%d",N)," \n")
            for j=1:N
                for i=1:N
                    write(vel, @sprintf("%1.9e",x[i,j])," ",@sprintf("%1.9e",y[i,j])," ",@sprintf("%1.9e",T0[i,j])," \n")
                end
            end
        end
    end
    t = tmax*dt
    write(vel, "title =\"ZoneTime_",string(t),"\""," \n")
    write(vel, "variables = \"x\", \"y\", \"T\""," \n")
    write(vel, "zone T=\"Zone_"*string(t)*"\" i=",@sprintf("%d",N)," j=",@sprintf("%d",N)," \n")
    for j=1:N
        for i=1:N
            write(vel, @sprintf("%1.9e",x[i,j])," ",@sprintf("%1.9e",y[i,j])," ",@sprintf("%1.9e",T0[i,j])," \n")
        end
    end
    close(vel)
end

function heat_wof()
    N = 1001
    T0 = Matrix{Float64}(undef,N,N)
    T1 = Matrix{Float64}(undef,N,N)
    x  = Matrix{Float64}(undef,N,N)
    y  = Matrix{Float64}(undef,N,N)
    a = 0
    b = π;
    dx = (b-a)/(N-1)
    for j=1:N
        for i=1:N
            x[i,j] = (i-1)*dx
        end
    end
    for j=1:N
        for i=1:N
            y[i,j] = (j-1)*dx
        end
    end

    dt = dx
    tmax = (10-mod(10,dt))/dt
    k = 0.25*dx
    alpha = k*dt/dx^2
    for i=1:N
        T1[1,i] = 1.0
        T1[N,i] = 1.0
        T1[i,1] = 1.0
        T1[i,N] = 1.0
    end
    for i=1:N
        T0[1,i] = 1.0
        T0[N,i] = 1.0
        T0[i,1] = 1.0
        T0[i,N] = 1.0
    end

    for k=1:tmax
        for j=2:N-1
            for i=2:N-1
                T1[i,j] = T0[i,j]+alpha*((T0[i+1,j]-2*T0[i,j]+T0[i-1,j])+(T0[i,j+1]-2*T0[i,j]+T0[i,j-1]))
            end
        end
        T0 .= T1
    end
end

function heat_Vec_wof()
    N = 1001
    T0 = Matrix{Float64}(undef,N,N)
    T1 = Matrix{Float64}(undef,N,N)
    x  = Matrix{Float64}(undef,N,N)
    y  = Matrix{Float64}(undef,N,N)
    a = 0
    b = π;
    dx = (b-a)/(N-1)
    x[:,1] = (0:N-1)*dx
    x[:,:] .= x[:,1]
    y[1,:] = (0:N-1)*dx
    y[:,:] .= y[1,:]

    dt = dx
    tmax = (10-mod(10,dt))/dt
    k = 0.25*dx
    alpha = k*dt/dx^2
    T1[1,:] .= 1.0
    T1[N,:] .= 1.0
    T1[:,1] .= 1.0
    T1[:,N] .= 1.0

    T0[1,:] .= 1.0
    T0[N,:] .= 1.0
    T0[:,1] .= 1.0
    T0[:,N] .= 1.0

    for k=1:tmax
        T1[2:N-1,2:N-1] .= T0[2:N-1,2:N-1].+alpha.*((T0[3:N,2:N-1].-2 .*T0[2:N-1,2:N-1].+T0[1:N-2,2:N-1]).+
                                                    (T0[2:N-1,3:N].-2 .*T0[2:N-1,2:N-1].+T0[2:N-1,1:N-2]))
        T0 .= T1
    end

end

results1 = @benchmark heat()
results2 = @benchmark heat_Vec()
results3 = @benchmark heat_wof()
results4 = @benchmark heat_Vec_wof()


export GAN,
jsgan,
klgan,
wgan,
rklgan,
lsgan,
sample,
predict,
dcgan_generator,
dcgan_discriminator

mutable struct GAN
    latent_dim::Int64
    batch_size::Int64
    dim::Int64
    dat::PyObject
    generator::Union{Missing, Function}
    discriminator::Union{Missing, Function}
    loss::Union{Missing, String, Function}
    d_vars::Union{Missing,Array{PyObject}}
    g_vars::Union{Missing,Array{PyObject}}
    d_loss::Union{Missing,PyObject}
    g_loss::Union{Missing,PyObject}
    is_training::Union{Missing,PyObject, Bool}
    update::Union{Missing,Array{PyObject}}
    fake_data::Union{Missing,PyObject}
    true_data::Union{Missing,PyObject}
    ganid::String
    noise::Union{Missing,PyObject}
    ids::Union{Missing,PyObject}
end

@doc raw"""
    GAN(dat::PyObject,
        generator::Function,
        discriminator::Function,
        loss::Union{String, Function, Missing}=missing;
        latent_dim::Union{Missing, Int64}=missing,
        batch_size::Union{Missing, Int64}=missing)

Creates a GAN instance.

- `dat` ``\in \mathbb{R}^{n\times d}`` is the training data for the GAN, where ``n`` is the number of training data, and ``d`` is the dimension per training data.
- `generator```:\mathbb{R}^{d'} \rightarrow \mathbb{R}^d`` is the generator function, ``d'`` is the hidden dimension.
- `discriminator```:\mathbb{R}^{d} \rightarrow \mathbb{R}`` is the discriminator function.
- `loss` is the loss function. See [`klgan`](@ref), [`rklgan`](@ref), [`wgan`](@ref), [`lsgan`](@ref) for examples.
- `latent_dim` (default=``d``, the same as output dimension) is the latent dimension.
- `batch_size` (default=32) is the batch size in training.
"""
function GAN(dat::Union{Array,PyObject}, generator::Function, discriminator::Function,
    loss::Union{String, Function, Missing}=missing; latent_dim::Union{Missing, Int64}=missing,
        batch_size::Union{Missing, Int64}=missing)
    dim = size(dat, 2)
    dat = convert_to_tensor(dat)
    if ismissing(latent_dim); latent_dim=dim; end
    if ismissing(batch_size); batch_size=32; end
    gan = GAN(latent_dim, batch_size, dim, dat, generator, discriminator, loss, missing, missing, missing,
        missing, placeholder(true), missing, missing, missing, randstring(), missing, missing)
    build!(gan)
    gan
end
GAN(dat::Array{T}) where T<:Real = GAN(constant(dat))

function Base.:show(io::IO, gan::GAN)
    print("GAN(id=$(gan.ganid), dat(shape=$(size(gan.dat)), type=$(get_dtype(gan.dat)), latent_dim=$(gan.latent_dim), batch_size=$(gan.batch_size), $(length(gan.d_vars)) d_vars, $(length(gan.g_vars)) g_vars, $(length(gan.update)) update ops, fake_data(shape=$(size(gan.fake_data))), true_data(shape=$(size(gan.true_data))))")
end
##################### GAN Library #####################
"""
    klgan(gan::GAN)

Computes the KL-divergence GAN loss function.
"""
function klgan(gan::GAN)
    P, Q = gan.true_data, gan.fake_data
    D_real = gan.discriminator(P, gan)
    D_fake = gan.discriminator(Q, gan)
    D_loss = -mean(log(D_real) + log(1-D_fake))
    G_loss = mean(log((1-D_fake)/D_fake))
    D_loss, G_loss
end

"""
    jsgan(gan::GAN)

Computes the vanilla GAN loss function.
"""
function jsgan(gan::GAN)
    P, Q = gan.true_data, gan.fake_data
    D_real = gan.discriminator(P, gan)
    D_fake = gan.discriminator(Q, gan)
    D_loss = -mean(log(D_real) + log(1-D_fake))
    G_loss = -mean(log(D_fake))
    D_loss, G_loss
end

"""
    wgan(gan::GAN)

Computes the Wasserstein GAN loss function.
"""
function wgan(gan::GAN)
    P, Q = gan.true_data, gan.fake_data
    D_real = gan.discriminator(P, gan)
    D_fake = gan.discriminator(Q, gan)
    D_loss = mean(D_fake)-mean(D_real)
    G_loss = -mean(D_fake)
    D_loss, G_loss
end

"""
    rklgan(gan::GAN)

Computes the reverse KL-divergence GAN loss function.
"""
function rklgan(gan::GAN)
    P, Q = gan.true_data, gan.fake_data
    D_real = gan.discriminator(P, gan)
    D_fake = gan.discriminator(Q, gan)
    G_loss = mean(log((1-D_fake)/D_fake))
    D_loss = -mean(log(D_fake)+log(1-D_real))
    D_loss, G_loss
end

"""
    lsgan(gan::GAN)

Computes the least square GAN loss function.
"""
function lsgan(gan::GAN)
    P, Q = gan.true_data, gan.fake_data
    D_real = gan.discriminator(P, gan)
    D_fake = gan.discriminator(Q, gan)
    D_loss = mean((D_real-1)^2+D_fake^2)
    G_loss = mean((D_fake-1)^2)
    D_loss, G_loss
end
#######################################################

"""
    build!(gan::GAN)

Builds the GAN instances. This function returns `gan` for convenience.
"""
function build!(gan::GAN)
    gan.noise = placeholder(get_dtype(gan.dat), shape=(gan.batch_size, gan.latent_dim))
    gan.ids = placeholder(Int32, shape=(gan.batch_size,))
    variable_scope("generator_$(gan.ganid)", initializer=random_uniform_initializer(0.0,1e-3)) do
        gan.fake_data = gan.generator(gan.noise, gan)
    end
    gan.true_data = tf.gather(gan.dat,gan.ids-1)
    if ismissing(gan.loss); gan.loss="jsgan"; end
    if isa(gan.loss, String)
        gan.loss = eval(Meta.parse(gan.loss))
    end
    variable_scope("discriminator_$(gan.ganid)", initializer=random_uniform_initializer(0.0,1e-3)) do
        gan.d_loss, gan.g_loss = gan.loss(gan)
    end
    gan.d_vars = Array{PyObject}(tf.get_collection(TRAINABLE_VARIABLES, scope="discriminator_$(gan.ganid)"))
    gan.g_vars = Array{PyObject}(tf.get_collection(TRAINABLE_VARIABLES, scope="generator_$(gan.ganid)"))
    gan.update = Array{PyCall.PyObject}(get_collection(UPDATE_OPS))
    gan
end

"""
    sample(gan::GAN, n::Int64)

Samples `n` instances from `gan`.
"""
function sample(gan::GAN, n::Int64)
    local out
    noise = normal(n, gan.latent_dim)
    gan.is_training = false
    variable_scope("generator_$(gan.ganid)", initializer=random_uniform_initializer(0.0,1e-3)) do
        out = gan.generator(noise, gan)
    end
    gan.is_training = placeholder(true)
    out
end

"""
    predict(gan::GAN, input::Union{PyObject, Array})

Predicts the GAN `gan` output given input `input`.
"""
function predict(gan::GAN, input::Union{PyObject, Array})
    local out
    flag = false
    if length(size(input))==1
        flag = true
        input = reshape(input, 1, length(input))
    end
    input = convert_to_tensor(input)
    variable_scope("generator_$(gan.ganid)", initializer=random_uniform_initializer(0.0,1e-3)) do
        out = gan.generator(input, gan)
    end
    if flag; out = squeeze(out); end
    out
end

# adapted from https://github.com/aymericdamien/TensorFlow-Examples/blob/master/examples/3_NeuralNetworks/dcgan.py
function dcgan_generator(x::PyObject, n::Int64=1)
    if length(size(x))!=2
        error("ADCME: input must have rank 2, rank $(length(size(x))) received")
    end
    variable_scope("generator", reuse=AUTO_REUSE) do
        # TensorFlow Layers automatically create variables and calculate their
        # shape, based on the input.
        x = tf.layers.dense(x, units=6n * 6n * 128)
        x = tf.nn.tanh(x)
        # Reshape to a 4-D array of images: (batch, height, width, channels)
        # New shape: (batch, 6, 6, 128)
        x = tf.reshape(x, shape=[-1, 6n, 6n, 128])
        # Deconvolution, image shape: (batch, 14, 14, 64)
        x = tf.layers.conv2d_transpose(x, 64, 4, strides=2)
        # Deconvolution, image shape: (batch, 28, 28, 1)
        x = tf.layers.conv2d_transpose(x, 1, 2, strides=2)
        # Apply sigmoid to clip values between 0 and 1
        x = tf.nn.sigmoid(x)
    end
    return squeeze(x)
end

function dcgan_discriminator(x)
    variable_scope("Discriminator", reuse=AUTO_REUSE) do
        # Typical convolutional neural network to classify images.
        x = tf.layers.conv2d(x, 64, 5)
        x = tf.nn.tanh(x)
        x = tf.layers.average_pooling2d(x, 2, 2)
        x = tf.layers.conv2d(x, 128, 5)
        x = tf.nn.tanh(x)
        x = tf.layers.average_pooling2d(x, 2, 2)
        x = tf.contrib.layers.flatten(x)
        x = tf.layers.dense(x, 1024)
        x = tf.nn.tanh(x)
        # Output 2 classes: Real and Fake images
        x = tf.layers.dense(x, 2)
    end
    return x
end


@testset verbose=true "Linesearch                  " begin

	n = 1

	f = x -> x[1]^2

	x = [-0.23]
	xnew = copy(x)
	d = [1.0]
	g = 2*x
	fval = f(x)
	cval = zeros(0)

	@testset "Armijo" begin
		flag, tot_iter1, tot_iter2, newf, f_diff, step_diff, α = LFPSQP.armijo!(xnew, x, n, d, g, f, fval, LFPSQP.Euclidean(), cval, nothing, LFPSQPParams(linesearch=LFPSQP.armijo), LFPSQP.ArmijoWork(1))

		@test flag == tot_iter1 == tot_iter2 == 0
		@test x[1] == -0.23							# ensure value not changed
		@test newf ≈ f(xnew)
		@test f_diff ≈ fval - newf
		@test step_diff ≈ α ≈ 0.25
	end

	@testset "Exact" begin
		flag, tot_iter1, tot_iter2, newf, f_diff, step_diff, α = LFPSQP.exact_linesearch!(xnew, x, n, d, f, fval, LFPSQP.Euclidean(), cval, nothing, LFPSQPParams(linesearch=LFPSQP.armijo), LFPSQP.ExactLinesearchWork(1))

		@test flag == tot_iter1 == tot_iter2 == 0
		@test x[1] == -0.23							# ensure value not changed
		@test newf ≈ f(xnew)
		@test f_diff ≈ fval - newf
		@test step_diff ≈ α && isapprox(α, 0.23, atol=1e-6)
	end
end


### A Pluto.jl notebook ###
# v0.14.5

using Markdown
using InteractiveUtils

# ╔═╡ 368b1310-b663-11eb-2d78-c79d87d7e01a
using Images, TestImages, ImageSegmentation, Random, ImageShow

# ╔═╡ abf9b1ed-308e-4e19-af87-81fa8f0cd028
function get_random_color(seed)
    Random.seed!(seed)
    rand(RGB{N0f8})
end

# ╔═╡ 2648125d-8574-4c56-8ee4-b086067949eb
begin
	img = Float32.(channelview(load(download("http://docs.opencv.org/3.1.0/water_coins.jpg")))[1, :, :]);
	Gray.(img)
end

# ╔═╡ c97329a0-6102-41d1-b794-abfd7339ddd2
begin
	img_thresh = img .> 0.8
	Gray.(img_thresh)
end

# ╔═╡ 6dc0d6f8-2f06-462b-9a8e-ab0d94aa4583
typeof(img_thresh)

# ╔═╡ af82e3dc-9502-4ff4-8ec7-614523fcae4b
begin
	img_dist = 1 .- distance_transform(feature_transform(img_thresh))
	Gray.(img_dist)
end

# ╔═╡ 4e96caa9-6976-488d-818e-713886d46c7e
begin
	markers = label_components(img_dist .< -15);
end

# ╔═╡ 4b047fc1-0496-4422-aa6b-0fa77854990d
segments = watershed(img_dist, markers)

# ╔═╡ d0a649b9-f046-44d4-a2cd-a9dc9a539fc3
RGB.(map(i->get_random_color(i), labels_map(segments)) .* (1 .-img_thresh))

# ╔═╡ Cell order:
# ╠═368b1310-b663-11eb-2d78-c79d87d7e01a
# ╠═abf9b1ed-308e-4e19-af87-81fa8f0cd028
# ╠═2648125d-8574-4c56-8ee4-b086067949eb
# ╠═c97329a0-6102-41d1-b794-abfd7339ddd2
# ╠═6dc0d6f8-2f06-462b-9a8e-ab0d94aa4583
# ╠═af82e3dc-9502-4ff4-8ec7-614523fcae4b
# ╠═4e96caa9-6976-488d-818e-713886d46c7e
# ╠═4b047fc1-0496-4422-aa6b-0fa77854990d
# ╠═d0a649b9-f046-44d4-a2cd-a9dc9a539fc3


"iterator type over simplicial multi-indices"
struct MultiIndexSet
    d::Int # dimension
    N::Int # degree
end

Base.length(a::MultiIndexSet) = prod(ntuple(i->(a.N+i),a.d)) ÷ factorial(a.d)

function Base.iterate(set::MultiIndexSet)
    α = ntuple(x->zero(Int),set.d)
    return α,α
end
function Base.iterate(set::MultiIndexSet,α::NTuple)
    if last(α) == set.N
        return nothing
    else
        αnext = increment(α,set,1)
        return αnext,αnext
    end
end

"function increment(I::NTuple,set::MultiIndexSet,i::Int):
    increments a multi-index of arbitrary dimension and degree.
    Intended to be used within the interface for Base.iterate"
function increment(I::NTuple,set::MultiIndexSet,i::Int)
    if length(I) != set.d
        error(DimensionMismatch("Length of I = $(length(I)) while set.d = $(set.d)"))
    end
    if last(I)==set.N
        return (I[1:end-1]...,last(I)+1) # set "incorrect" state to end iteration
    end
    if sum(I) < set.N
        return (I[1:i-1]...,I[i]+1,I[i+1:end]...)
    else
        Inext = (zero.(I[1:i])...,I[i+1:end]...)
        increment(Inext,set,i+1) # recursive
    end
end


using FEASTSolver
using LinearAlgebra
using Printf
using Random
using MatrixMarket
using NonlinearEigenproblems: nep_gallery, compute_Mder
using DelimitedFiles

function bf()
    N = diagm(-1 => [1,1,1,1,1,1,1])
    Mh0 = 1/6*(4*I + N + N')
    Mh1 = N - N'
    Mh2 = -1*(2*I - N - N')
    Mh3 = Mh1
    Mh4 = -Mh2
    c = [0.6 1.3; 1.3 0.1; 0.1 1.2; 1.0 1.0; 1.2 1.0]
    I8 = Matrix(I, 8, 8)
    M0 = c[1,1] * kron(I8, Mh0) + c[1,2] * kron(Mh0, I8)
    M1 = c[2,1] * kron(I8, Mh1) + c[2,2] * kron(Mh1, I8)
    M2 = c[3,1] * kron(I8, Mh2) + c[3,2] * kron(Mh2, I8)
    M3 = c[4,1] * kron(I8, Mh3) + c[4,2] * kron(Mh3, I8)
    M4 = c[5,1] * kron(I8, Mh4) + c[5,2] * kron(Mh4, I8)
    function butterfly(A0, A1, A2, A3, A4)
        f(z) =  z^4*A4 + z^3*A3 + z^2*A2 + z*A1 + A0
    end

    return butterfly(M0, M1, M2, M3, M4)
end

function deficient()
	# Random.seed!(1234)
	# A0 = complex.(rand(15,15))
	# A1 = complex.(rand(15,15))
	# A0[:,1] .= 0
	A0 = mmread("data/quadraticM0.mtx")
	A1 = mmread("data/quadraticM1.mtx")

	function T(z::ComplexF64)
		return (z + 0.2)*(z - 0.1).*A1 .+ A0
	end

	return T
end

function quadratic()
	A0 = mmread("data/system5A0.mtx")
	A1 = mmread("data/system5A1.mtx")
	A2 = mmread("data/system5A2.mtx")

	function T(z::ComplexF64)
		return (A2 .* (z*z)) .+ (A1 .* z) .+ A0
	end

	return T
end

function gun()
	nep_gun = nep_gallery("nlevp_native_gun")
	T(z) = compute_Mder(nep_gun, z)
	return T
end

function hadeler()
	nep_hadeler = nep_gallery("nlevp_native_hadeler", 100, 200)
	T(z) = compute_Mder(nep_hadeler, z)
	return T
end

function info(Λ, X, residuals, c, r)
    inside(x) = in_contour(x, c, r)
    in_eig = Λ[inside.(Λ)]
    in_res = residuals[inside.(Λ)]
    print("number inside : ")
    println(size(Λ[inside.(Λ)])[1])
    if sum(inside.(Λ)) > 0
        in_res_conv = in_res[in_res .<= 1e-3]
        in_eig_conv = in_eig[in_res .<= 1e-3]
        print("number inside converged : ")
        println(size(in_eig_conv)[1])
        print("max res inside: ")
        println(maximum(residuals[inside.(Λ)]))
        if size(in_res_conv, 1) > 0
            print("max res inside non spurious: ")
            println(maximum(in_res_conv))
        end
    end
    # display(in_eig_conv)
end

function res_table(feast_call, dim, max_iter, max_nodes, spurious, c, r, err_tol)
	all_res = zeros((max_iter, max_nodes))
	println(repeat("-", max_nodes*8*2))
	for j=1:max_nodes
	    print(2^(j+1))
	    print("\t\t")
	end
	println("\n"*repeat("-", max_nodes*8*2))
	for i=1:max_iter
	    for j=1:max_nodes
			if i>1 && minimum(all_res[1:i-1,j]) < err_tol
				# all_res[i,j] = all_res[i-1,j]
				all_res[i,j] = -Inf
				print("\t\t")
			else
		        inside(z) = in_contour(z, c, r)
		        e, v, res = feast_call(rand(ComplexF64, dim), 2^(j+1), i-1, err_tol)
		        in_eig = e[inside.(e)]
		        in_res = res[inside.(e)]
		        in_res_conv = in_res[in_res .<= spurious]
		        if size(in_res_conv, 1) == 0
		            all_res[i,j] = maximum(res[inside.(e)])
		        else
		            all_res[i,j] = maximum(in_res_conv)
		        end
		        @printf("%2.2e\t", all_res[i,j])
			end
	    end
	    println()
	end
	println(repeat("-", max_nodes*8*2))
	return all_res
end

err_tol_all = 10e-16
max_iter_all = 10

# for i=2:8
# 	println("\nBeyn N=$(2^i)")
# 	e, v, res = nlfeast!(bf(), rand(ComplexF64,64,30), 2^i, 0, c=complex(1.0,1.0), r=0.5)
# 	info(e, v, res, complex(1.0,1.0), 0.5)
# end

Random.seed!(1729)

# println("\nButterfly Problem")
# max_iter_butterfly = 25
# bf_feast_call(X, N, iter, err_tol) = nlfeast!(bf(), X, N, iter, c=complex(1.0,1.0), r=0.5, ϵ=err_tol)
# bf_res = res_table(bf_feast_call, (64,20), max_iter_butterfly, 7, 1e-0, complex(1.0,1.0), 0.5, err_tol_all)
# writedlm("test/data/bf.dat", zip(1:max_iter_butterfly, eachcol(bf_res)...))
#
#
# println("\nDeficient Quadratic Problem")
#
# deficient_feast_call(X, N, iter, err_tol) = nlfeast_moments!(deficient(), X, N, iter, c=complex(0.0,0.0), r=0.25, ϵ=err_tol, moments=2, store=true)
# def_res = res_table(deficient_feast_call, (15,4), max_iter_all, 7, 1e-5, complex(0.0,0.0), 0.25, err_tol_all)
# writedlm("test/data/deficient.dat", zip(1:max_iter_all, eachcol(def_res)...))
#
#
# println("\nLarger Quadratic Problem")
# quad_feast_call(X, N, iter, err_tol) = nlfeast_moments!(quadratic(), X, N, iter, c=complex(-1.55,0.0), r=0.05, ϵ=err_tol, moments=2, store=true)
# quad_res = res_table(quad_feast_call, (1000,70), max_iter_all, 6, 1e-6, complex(-1.55,0.0), 0.05, err_tol_all)
# writedlm("test/data/quad.dat", zip(1:max_iter_all, eachcol(quad_res)...))
#
#
# println("\nGun Cavity Problem")
# gun_feast_call(X, N, iter, err_tol) = nlfeast!(gun(), X, N, iter, c=complex(140000.0, 0.0), r=30000, ϵ=err_tol, store=false)
# gun_res = res_table(gun_feast_call, (9956,32), max_iter_all, 6, 1e-6, complex(140000.0, 0.0), 30000, err_tol_all)
# writedlm("test/data/gun.dat", zip(1:max_iter_all, eachcol(gun_res)...))

println("\nHadeler Problem")
hadeler_feast_call(X, N, iter, err_tol) = nlfeast!(hadeler(), X, N, iter, c=complex(-30.0, 0.0), r=10.0, ϵ=err_tol, store=true)
hadeler_res = res_table(hadeler_feast_call, (200,15), max_iter_all, 6, 1e-4, complex(-30.0, 0.0), 10.0, err_tol_all)
writedlm("test/data/hadeler.dat", zip(1:max_iter_all, eachcol(hadeler_res)...))


using PkgSRA
using Plots, Random, OrdinaryDiffEq, Statistics
using Distributions, Turing
using StatsPlots, DSP
pyplot()
Random.seed!(11)

#####
##### Define the controlled dynamical system
#####
# Load example problem
include("paper_settings.jl")
include(EXAMPLE_FOLDERNAME*"example_lorenz.jl")

# Define the forcing functions
num_ctr = 5;
    U_starts = rand(num_ctr,1) .* tspan[2]
    U_width = 0.05;
    my_U_func_time(t, u) = U_func_time(t, u, U_width, U_starts, F_dim=3)

#####
##### Produce data
#####
sol = solve_lorenz_system(my_U_func_time)
dat_clean = Array(sol)

# NEW: ADD NOISE
rng = MersenneTwister(128);
noise = 0.1
dat = dat_clean .+ noise.*randn(rng, eltype(dat_clean), size(dat_clean))

plot3d(dat[1, :], dat[2, :], dat[3, :])
    title!("Noisy Data")

# Derivatives
numerical_grad = numerical_derivative(dat, ts)
true_grad = zeros(size(dat))
for i in 1:size(dat,2)
    true_grad[:,i] = lorenz_system(true_grad[:,i], dat[:,i], p, [0])
end
# plot(numerical_grad[3,:], label="numerical")
#     plot!(true_grad[3,:], label="true uncontrolled")

# True control signal
U_true = zeros(size(dat))
for (i, t) in enumerate(ts)
    U_true[:,i] = my_U_func_time(t, dat[:,i])
end

# plot(U_true')

#####
##### Define the Bayesian modeling framework
#####

# Define the model
function lorenz_system(u, p, t)
    u = convert.(eltype(p),u)

    ρ, σ, β = p
    x, y, z = u
    du = [σ*(y-x); x*(ρ-z) - y; x*y - β*z]
    return du
end

@model lorenz_grad_residual(y) = begin
    # Lorenz parameters
    ρ ~ Normal(p[1], 1.0)
    σ ~ Normal(p[2], 0.1)
    β ~ Normal(p[3], 0.1)
    params = [ρ, σ, β]

    t = [0] # Not time dependent
    for i in 1:size(y,2)
        u = dat[:,i]
        du_params = lorenz_system(u, params, t)
        # TODO: Hard-coded noise
        y[:, i] ~ MvNormal(du_params, [10.0, 10.0, 10.0])
    end
end;

# Settings of the Hamiltonian Monte Carlo (HMC) sampler.
iterations = 1000
    ϵ = 0.02
    τ = 10
# Try to predict the GRADIENT from data
num_training_pts = 500
    train_ind = 201:num_training_pts+200
    y = numerical_grad[:,train_ind]
chain = sample(lorenz_grad_residual(y),
                HMC(iterations, ϵ, τ));
# plot3d(dat[1, train_ind], dat[2, train_ind], dat[3, train_ind])
#     title!("Training data")
# plot(chain)

# Generate test trajectories from the posterior
num_samples = 100
    param_samples = sample(chain, num_samples)
    save_ind = 3
    t = [0]
    num_test_pts = 500
    vars = [:ρ, :σ, :β]
    all_vals = zeros(num_samples, num_test_pts)
    plot_ind = 201:num_test_pts+200
for i in 1:num_samples
    these_params = [param_samples[v].value[i] for v in vars]
    for (i_save, i_dat) in enumerate(plot_ind)
        all_vals[i, i_save] = lorenz_system(dat[:,i_dat],
                                these_params, t)[save_ind]
    end
end

# Calculate the residuals
this_std = std(all_vals, dims=1)
    this_mean = mean(all_vals, dims=1)

# Align the signals
dat_grad = numerical_grad[save_ind, plot_ind]

# Function for processing the residual into a true guess
function process_residual(mean, std)
    low = mean .- std
    high = mean .+ std
    real_ind = .!((low .< 0.0) .& (high .> 0.0))

    ctr = zeros(size(mean))
    ctr[real_ind] = mean[real_ind]

    return ctr
end

#####
##### Plot everything
#####

this_plot_ind = 1:300

p1 = plot(vec(dat_grad[this_plot_ind]),
            linecolor=COLOR_DICT["true"], label="Data", lw=3);
    plot!(vec(this_mean[this_plot_ind]),
            linecolor=COLOR_DICT["intrinsic"], ribbon=vec(this_std),
            fillalpha=0.5, label="Models", lw=1);
    title!("Reconstruction");

residual = vec(dat_grad[this_plot_ind]) .- vec(this_mean[this_plot_ind])
p2 = plot(residual, lw=2, ribbon=vec(this_std), fillalpha=0.5,
            linecolor=COLOR_DICT["intrinsic"], legend=false);
    title!("Residual");

ctr_guess = process_residual(residual, this_std[this_plot_ind])
p3 = plot(ctr_guess, lw=2, label="Processed Residual",
            linecolor=COLOR_DICT["control_time"]);
    plot!(vec(U_true[3,plot_ind][this_plot_ind]), label="True control signal",
                linecolor=COLOR_DICT["control_true"]);
    title!("Learned signal");

## Start: Test plot
# residual_ratio = abs.(residual) ./ vec(this_std[this_plot_ind])
# plot(residual_ratio, label="Residual")
#     plot!(vec(U_true[3,plot_ind][this_plot_ind]), label="True control signal");
#     title!("??")
## End: Test plot

# All together
l = @layout [p1; p2; p3]
    p_all = plot(p1, p2, p3, layout=l)

fname = FIGURE_FOLDERNAME * "fig_bayesian_sra2.png";
savefig(p_all, fname)


using Gadfly, Compat

set_default_plot_size(6inch, 3inch)

f(x,y) = x*exp(-(x-(round(Int, x)))^2-y^2)
M = Array{Float64}(200, 200)
for (i, x) in enumerate(linspace(-8., 8., 200)), (j, y) in enumerate(linspace(-2., 2., 200))
    M[i, j] = f(x, y)
end

plot(z=M, Geom.contour, xmin=[-8.], xmax=[8], ymin=[-2.], ymax=[2])


# This file contains functions for the Radial Big Data Model, following Khezrimotlagh et al. (2019) algorithm (we use KZCT as acronym here)
# D. Khezrimotlagh, J. Zhu and W.D. Cook et al. / European Journal of Operational Research 274 (2019) 1047–1054

"""
    initialsubset(X,Y)

Creates the initial subset as in KZCT algorithm:
1. Determines p the size of the subset and create an empty vector or matrix to fill out
2. Determines the minimum value for each input m and the maximum value for each output s
3. Select a subset with one DMU per input, with DMU's input equal to the minimum value for this input m and one DMU per output
with DMU's output equal to the maximum value for this output s: m + s DMU's selected
4. To fill out the subset up to p DMUs, implement the following algorithm:
    For each unselected DMU_j,
    {u = 0,
        {For each i and k_i,
            If x_ij <= k_ith percentile of x^i then
                u = u + 1}
            {For each r and k'_r
                If y_rj >= k'th percentile of y^r then
                    u = u + 1}
    pre_score_j = u }

Where x^i and y^r represent the ith inputs and the rth outputs of all DMUS.
The indexes k_i and k'_r are changed from 0 to 100 and refer to the percentiles of x^i and y^r respectively.
Sorting DMUs in descending order by the assigned pre-scores, the remaining DMUs (to construct the subsample with size p)
are selected as those having the greatest pre-scores.

"""
function initialsubset(X::Union{Vector,Matrix},Y::Union{Vector,Matrix}, n::Int64, s::Int64, m::Int64)
    # Determines the size of the initial subset Dˢ
    p = convert(Int64, round(sqrt(n)))

    # Select DMU's with the minimum input or the maximum output
    dmus_selected = Vector{Int64}()
    for x in 1:m
        push!(dmus_selected, argmin(X[:,x]))
    end

    for y in 1:s
        push!(dmus_selected, argmax(Y[:,y]))
    end
    unique!(dmus_selected)

    initial_index = collect(1:n)
    dmus_unselected = initial_index[Not(dmus_selected)]

    # Algorithm implementation to fill out the subsample up to the size p

    # select unselected DMUs
    X_unselected = @view X[Not(dmus_selected), :]
    Y_unselected = @view Y[Not(dmus_selected), :]

    # create the pre-scores matrix (one column with the pre-score and the other to store the initial index)
    pre_scores = zeros(length(dmus_unselected))

    @inbounds for i in 1:size(X_unselected, 1)
        for x in 1:m
            xquant = quantile!(X[:,x], 0.01:0.01:1)
            for k in 1:100
                if X_unselected[i,x] <= xquant[k]
                    pre_scores[i] = pre_scores[i] + 1
                end
            end
        end
        for y in 1:s
            yquant = quantile!(Y[:,y], 0.01:0.01:1)
            for k in 1:100
                if Y_unselected[i,y] >= yquant[k]
                    pre_scores[i] = pre_scores[i] + 1
                end
            end
        end
    end

    # sort pre_scores to get the best DMUs in unselected DMUs
    score_perm = sortperm(pre_scores, rev = true)
    dmus_unselected = dmus_unselected[score_perm]
    pre_scores = pre_scores[score_perm, :]

    # add the number of DMUs needed to fill out the subsample
    z = p - length(dmus_selected)
    new_dmus = dmus_unselected[1:z]
    dmus_selected = vcat(dmus_selected, new_dmus)
    dmus_unselected = initial_index[Not(dmus_selected)]

    # Creates the subsets Dˢ and D_excluding_Dˢ
    return dmus_selected, dmus_unselected
end

"""
    bestpracticesfinder(Subset)

Identify best-practices in the subsample selected
"""
function bestpracticesfinder(scores::Vector, orient::Symbol, atol::Float64)
    index_bestpractices = Vector{Int64}()
    for i in 1:length(scores)
        if orient == :Input
            if scores[i] >= 1 - atol
                push!(index_bestpractices, i)
            end
        elseif orient == :Output
            if scores[i] <= 1 + atol
                push!(index_bestpractices, i)
            end
        end
    end
    return index_bestpractices
end

"""
    deabigdata(X, Y)

Compute the big data radial model using data envelopment analysis for inputs X and outputs Y.

# Optional Arguments
- `orient=:Input`: chooses the radially oriented input mode. For the radially oriented output model choose `:Output`.
- `rts=:CRS`: chooses constant returns to scale. For variable returns to scale choose `:VRS`.
- `atol=1e-6`: tolerance for DMU to be considered efficient.
- `names`: a vector of strings with the names of the decision making units.
"""
function deabigdata(X::Union{Matrix,Vector}, Y::Union{Matrix,Vector};
        orient::Symbol = :Input, rts::Symbol = :CRS, slack::Bool = true, atol::Float64 = 1e-6,
        names::Union{Vector{String},Nothing} = nothing,
        optimizer::Union{DEAOptimizer,Nothing} = nothing, progress::Bool = true)

    #=
    The algorithm is such as:
    1. Start
    2. D <- get a sample of DMUs
    3. D^S <- select a subsample of D
    4. B^S <- find best-practices in D^S
    5. E <- find exterior DMUs in D excluding D^S respect to the hull of B^S
    6. If E = {}, then F = B^S and all DMUs are altready evaluated (go to step 7)
    Otherwise
        6.1. F <- find best-practice DMUs in D^S Union E
        6.2. Evaluate DMUs in D excluding (D^S Union E) respoect to the F's hull
    7. End
    =#

    # Check parameters
    nx, m = size(X, 1), size(X, 2)
    ny, s = size(Y, 1), size(Y, 2)

    if nx != ny
        throw(DimensionMismatch("number of@rows in X and Y ($nx, $ny) are not equal"));
    end

    # Default optimizer
    if optimizer === nothing
        optimizer = DEAOptimizer(:LP)
    end

    n = nx
    disposX = :Strong
    disposY = :Strong

    # Select initial subset Dˢ using Khezrimotlagh et al. (2019) algorithm
    Dˢ, D_excluding_Dˢ = initialsubset(X, Y, n, s, m)

    # Find the best-practices Bˢ in Dˢ
    Dˢ_evaluation = dea(X[Dˢ, :], Y[Dˢ, :], orient = orient, rts = rts,slack = slack,
                    disposX = disposX, disposY = disposY, optimizer = optimizer, progress = progress)

    index_bestpractices = bestpracticesfinder(efficiency(Dˢ_evaluation), orient, atol)

    length(index_bestpractices) > 0 || throw(ErrorException("No efficient DMUs found in the initial subset. Consider increasing the tolerance `atol`"))

    Bˢ = Dˢ[index_bestpractices]

    # Find exterior DMUs in D excluding Dˢ respect to the hull of Bˢ
    D_excluding_Dˢ_evaluation = dea(X[D_excluding_Dˢ, :], Y[D_excluding_Dˢ, :], orient = orient,
                        rts = rts, slack = slack, Xref = X[Bˢ, :], Yref = Y[Bˢ, :],
                        disposX = disposX, disposY = disposY, optimizer = optimizer, progress = progress)

    index_exteriors = bestpracticesfinder(efficiency(D_excluding_Dˢ_evaluation), orient, atol)

    E = D_excluding_Dˢ[index_exteriors]

    # If E is not empty, then we have to find best practice DMUs in D^S Union E, otherwise, F = B_S and all DMUs are already evaluated
    if size(E,1) > 0
        # Create the subset Dˢ_union_E
        Dˢ_union_E = vcat(Dˢ, E)

        sort!(Dˢ_union_E, rev = false)

        Dˢ_union_E_evaluation = dea(X[Dˢ_union_E, :], Y[Dˢ_union_E, :], orient = orient, rts = rts,slack = slack,
                        disposX = disposX, disposY = disposY,optimizer = optimizer, progress = progress)

        # Find the index of best practices DMUs
        index_bestpractices = bestpracticesfinder(efficiency(Dˢ_union_E_evaluation), orient, atol)

        # Best practices DMUs are the efficient DMUs F of the sample D
        F = Dˢ_union_E[index_bestpractices]

        # # We then evaluate the rest of DMUs
        initial_index = collect(1:n)

        D_excluding_Dˢ_union_E = initial_index[Not(Dˢ_union_E)]

        D_excluding_Dˢ_union_E_evaluation = dea(X[D_excluding_Dˢ_union_E, :], Y[D_excluding_Dˢ_union_E, :], orient = orient,
                            rts = rts, slack = slack, Xref = X[F, :], Yref = Y[F, :],
                            disposX = disposX, disposY = disposY, optimizer = optimizer, progress = progress)

        subset1  = Dˢ_union_E
        results1 = Dˢ_union_E_evaluation
        subset2  = D_excluding_Dˢ_union_E
        results2 = D_excluding_Dˢ_union_E_evaluation
    else
        F = Bˢ

        subset1  = Dˢ
        results1 = Dˢ_evaluation
        subset2  = D_excluding_Dˢ
        results2 = D_excluding_Dˢ_evaluation
    end

    # Get results
    resultsperm = sortperm(vcat(subset1, subset2), rev = false)
    effi = vcat(efficiency(results1), efficiency(results2))
    effi = effi[resultsperm]

    if slack
        slackX = vcat(slacks(results1, :X), slacks(results2, :X))
        slackY = vcat(slacks(results1, :Y), slacks(results2, :Y))

        slackX = slackX[resultsperm, :]
        slackY = slackY[resultsperm, :]
    else
        slackX = Array{Float64}(undef, 0, 0)
        slackY = Array{Float64}(undef, 0, 0)
    end

    Xtarget = vcat(targets(results1, :X), targets(results2, :X))
    Ytarget = vcat(targets(results1, :Y), targets(results2, :Y))

    Xtarget = Xtarget[resultsperm, :]
    Ytarget = Ytarget[resultsperm, :]

    lambdas1 = peersmatrix(results1)
    lambdas2 = peersmatrix(results2)

    index_lambda_to_keep = findall(x -> x in F, subset1)
    lambdaeff = vcat(lambdas1[:,index_lambda_to_keep],lambdas2)
    lambdaindex = vcat(subset1, subset2)
    lambda_matrix = lambdaeff[resultsperm, :]

    # Create the sparse matrix for lambdas
    lambdaeff = spzeros(n, n)
    lambdaeff[:, convert.(Int, lambdaindex[index_lambda_to_keep])] = lambda_matrix

    # return results
    return RadialDEAModel(n, m, s, orient, rts, disposX, disposY, names, effi, slackX, slackY, lambdaeff, Xtarget, Ytarget)
end