function metro!(S::Config, M::Model)
    for _ in eachindex(S.config) #length(M)
        mcstep!(S, M)
    end
end

function mcstep!(S::Config, M::Model)
    proposed_site = rand(S.sites)
    dE = energy_singleflip(S, M, proposed_site) # SLOW 1
    if dE < 0 || rand() < exp(-dE)
        S.config[proposed_site] *= -1
    end
end

function energy_singleflip(S::Config, M::Model, site::CartesianIndex{D}) where D
    E = 0.0
    T = length(M.J[:X])
    @views for nn in S.neighbours_table[:, site]
        if 1 ≤ nn[1] ≤ M.dims[1] &&  # OBNC
           1 ≤ nn[2] ≤ M.dims[2] &&
           1 ≤ nn[3] ≤ M.dims[3]

            if nn[D] == site[D] # horizontal couplings
                J = M.J[:ZZ][nn[D]]
            elseif nn[D] == mymod1(site[D] - 1, T)
                J = M.J[:X][nn[D]]
            elseif nn[D] == mymod1(site[D] + 1, T)
                J = M.J[:X][site[D]]
            end
            E += ifelse(S.config[nn] == S.config[site], J, -J)  # kinda slow
        end
    end
    # Exteral field
    # E += M.J[:Z][site[D]] * S.config[site]
    return 2.0 * E
end


using FractionalDiffEq, Plots

function classical!(du, u, p, t)
    A, B = 1, 3
    du[1]=A-(B+1)*u[1]+u[1]^2*u[2]
    du[2]=B*u[1]-u[1]^2*u[2]
end
α = [0.8, 0.7]
u0 = [1.2, 2.8]
h=0.01
prob = FODESystem(classical!, α, u0)
result = solve(prob, h, 10, NonLinearAlg())
tspan = collect(0:h:10)
plot(result[:, 1], result[:, 2])


function modes_x(y,byp,H,Np,B,m,omega)

# Computes the longitudinal component of the magnetic
# perturbation using a solution for the meridional
# component
#
# input
#  y     - grid points in meridional coordinate
#  byp   - meridional magnetic perturbation
#  H     - thickness of layer (km)
#  Np    - dimensionless buoyancy frequency  (N/Omega)
#  B     - radial magnetic field  (T)
#  m     - angular order
#  omega0 - initial estimate of frequency

# a few constants   (must be consistent with ModelParameters)
Omega = 0.7292e-4;
R = 3.48e6;
N = Np * Omega;
k = pi/(1000.0*H);

# compute model parameters
C,M = ModelParameters(H,Np,B,m,omega);

# evaluate derivative matrices
d1,d2 = derivatives(y);

#   right-hand side of Eq. (17) for bxp
n = length(y);
y2 = y.*y;
yc2 = -(y2 .- 1.0);
a1 = -(m^2*ones(n) + M * yc2);

#   left-hand side
a2 = -((C)im * y .* byp);
a3 = (m * d1*byp)im;

bxp = (a2 + a3)./a1;

return bxp

end


"""
Performs a simple single-threaded Simulated Annealing optimization.

  Args:
    f (Function): Objective function.
    x (Array): Initial design variable values.
    t (Function): A function representing the annealing schedule (temperature).
    k_max (Integer): Maximum iterations.

  Returns:
    Tuple(Array, Float): A tuple representing the minimizers and minimum for
      the given objective function.
"""
function simulated_annealing(f, x, t, k_max)
    y = f(x)
    x_best, y_best = x, y
    for k in 1 : k_max
        x′ = x + randn(length(x))
        y′ = f(x′)
        Δy = y′ - y
        if Δy ≤ 0 || rand() < exp(-Δy/t(k))
            x, y = x′, y′
        end
        if y′ < y_best
            x_best, y_best = x′, y′
        end
    end
    return x_best, y_best
end

"""
Performs a simple multithreaded Simulated Annealing optimization.

  Args:
    f (Function): Objective function.
    x (Array): Initial design variable values.
    t (Function): A function representing the annealing schedule (temperature).
    k_max (Integer): Maximum iterations.

  Returns:
    Tuple(Array, Float): A tuple representing the minimizers and minimum for
      the given objective function.
"""
function simulated_annealing_p(f, x, t, k_max)
    y = f(x)
    x_best, y_best = x, y
    # create a lock to use in reading/writing current best solution
    lk = ReentrantLock()

    # split max iterations evenly over threads
    Threads.@threads for k in 1 : k_max
        x′ = x + randn(length(x))
        y′ = f(x′)
        Δy = y′ - y
        if Δy ≤ 0 || rand() < exp(-Δy/t(k))
            x, y = x′, y′
        end

        # acquire lock to potentially update current best while avoiding race conditions
        lock(lk) do
            if y′ < y_best
                x_best, y_best = x′, y′
            end
        end
    end
    return x_best, y_best
end


using ADCME
using PyCall
using LinearAlgebra
using PyPlot
using Random
using AdFem
using SparseArrays
Random.seed!(233)

function fast_solve(rhs)
    fast_solve_ = load_op_and_grad("./build/libFastSolve","fast_solve")
    rhs = convert_to_tensor(rhs, dtype = Float64)
    fast_solve_(rhs)
end

m = 10
n = 5
h = 0.1

K_ = compute_fem_stiffness_matrix1(diagm(0=>ones(2)), m, n, h)
ind = collect(1:m+1)
K_[ind, :] .= 0.0
K_[ind, ind] = spdiagm(0=>ones(length(ind)))
rhs = Float64.(collect((1:(m+1)*(n+1))))
x = K_\rhs
ccall((:factorize_matrix, "./build/libSaver"), Cvoid, (Cint, Cint, Cdouble, Ref{Cint}, Cint),
            Int32(m), Int32(n), h, Int32.(ind .- 1) , Int32(length(ind)))

tf_x = fast_solve(constant(rhs))
sess = Session(); init(sess)
@show norm(run(sess, tf_x)-x)
# uncomment it for testing gradients
# error()

# TODO: change your test parameter to `m`
#       in the case of `multiple=true`, you also need to specify which component you are testings
# gradient check -- v
function scalar_function(m)
    return sum(fast_solve(m)^2)
end

# TODO: change `m_` and `v_` to appropriate values
m_ = constant(rand((m+1)*(n+1)))
v_ = rand((m+1)*(n+1))
y_ = scalar_function(m_)
dy_ = gradients(y_, m_)
ms_ = Array{Any}(undef, 5)
ys_ = Array{Any}(undef, 5)
s_ = Array{Any}(undef, 5)
w_ = Array{Any}(undef, 5)
gs_ =  @. 1 / 10^(1:5)

for i = 1:5
    g_ = gs_[i]
    ms_[i] = m_ + g_*v_
    ys_[i] = scalar_function(ms_[i])
    s_[i] = ys_[i] - y_
    w_[i] = s_[i] - g_*sum(v_.*dy_)
end

sess = Session(); init(sess)
sval_ = run(sess, s_)
wval_ = run(sess, w_)
close("all")
loglog(gs_, abs.(sval_), "*-", label="finite difference")
loglog(gs_, abs.(wval_), "+-", label="automatic differentiation")
loglog(gs_, gs_.^2 * 0.5*abs(wval_[1])/gs_[1]^2, "--",label="\$\\mathcal{O}(\\gamma^2)\$")
loglog(gs_, gs_ * 0.5*abs(sval_[1])/gs_[1], "--",label="\$\\mathcal{O}(\\gamma)\$")

plt.gca().invert_xaxis()
legend()
xlabel("\$\\gamma\$")
ylabel("Error")


"""

    iterateCSteps(setting, subsetindices, h)

Perform a concentration step for a given subset of a regression setting.

# Arguments
- `setting::RegressionSetting`: RegressionSetting object with a formula and dataset.
- `subsetindices::Array{Int, 1}`: Indicies of observations in the initial subset.
- `h::Int`: A constant at least half of the number of observations.

# Notes
    This function is a helper for the lts function. A concentration step starts with a
    initial subset. The size of the subset is enlarged to h, a constant at least half of the
    observations. Please refer to the citation given below.

# References
Rousseeuw, Peter J., and Katrien Van Driessen. "An algorithm for positive-breakdown
regression based on concentration steps." Data Analysis.
Springer, Berlin, Heidelberg, 2000. 335-346.
"""
function iterateCSteps(setting::RegressionSetting, subsetindices::Array{Int,1}, h::Int)
    X, y = @extractRegressionSetting setting
    return iterateCSteps(X, y, subsetindices, h)
end

function iterateCSteps(X::Array{Float64,2}, y::Array{Float64,1}, subsetindices::Array{Int,1}, h::Int)
    starterset = copy(subsetindices)
    oldobjective = Inf
    objective = Inf
    iter = 0
    maxiter = 10000
    n, p = size(X)
    while iter < maxiter
        try
            Xsub = X[subsetindices, :]
            ysub = y[subsetindices]
            olsreg = ols(Xsub, ysub)
            betas = coef(olsreg)
            res = [y[i] - sum(X[i,:] .* betas) for i in 1:n]
            sortedresindices = sortperm(abs.(res))
            subsetindices = sortedresindices[1:h]
            objective = sum(sort(res.^2.0)[1:h])
            if oldobjective == objective
                break
            end
            oldobjective = objective
            iter += 1
        catch er
            @warn er
            return (objective, subsetindices)
        end
    end
    if iter >= maxiter
        @warn "in c-step stage of LTS, a h-subset is not converged for starting indices " starterset
    end
    return (objective, subsetindices)
end

function iterateCSteps(setting::RegressionSetting, initialBetas::Array{Float64,1}, h::Int)
    X = designMatrix(setting)
    y = responseVector(setting)
    return iterateCSteps(X, y, initialBetas, h)
end

function iterateCSteps(X::Array{Float64,2}, y::Array{Float64,1}, initialBetas::Array{Float64,1}, h::Int)
    n, p = size(X)
    res = [y[i] - sum(X[i,:] .* initialBetas) for i in 1:n]
    sortedresindices = sortperm(abs.(res))
    subsetindices = sortedresindices[1:p]
    return iterateCSteps(X, y, subsetindices, h)
end

"""

    lts(setting; iters, crit)

Perform the Fast-LTS (Least Trimmed Squares) algorithm for a given regression setting.

# Arguments
- `setting::RegressionSetting`: RegressionSetting object with a formula and dataset.
- `iters::Int`: Number of iterations.
- `crit::Float64`: Critical value.

# Description
The algorithm searches for estimations of regression parameters which minimize the sum of first h
ordered squared residuals where h is Int(floor((n + p + 1.0) / 2.0)). Specifically, our implementation,
uses the algorithm Fast-LTS in which concentration steps are used for enlarging a basic
subset to subset of clean observation of size h.

# Output
- `["betas"]`: Estimated regression coefficients
- `["S"]`: Standard error of regression
- `["hsubset"]`: Best subset of clean observation of size h.
- `["outliers"]`: Array of indices of outliers
- `["scaled.residuals"]`: Array of scaled residuals
- `["objective"]`: LTS objective value.

# Examples
```julia-repl
julia> reg = createRegressionSetting(@formula(calls ~ year), phones);
julia> lts(reg)
Dict{Any,Any} with 6 entries:
  "betas"            => [-56.5219, 1.16488]
  "S"                => 1.10918
  "hsubset"          => [11, 10, 5, 6, 23, 12, 13, 9, 24, 7, 3, 4, 8]
  "outliers"         => [14, 15, 16, 17, 18, 19, 20, 21]
  "scaled.residuals" => [2.41447, 1.63472, 0.584504, 0.61617, 0.197052, -0.222066, -0.551027, -0.970146, -0.397538, -0.185558  …  …
  "objective"        => 3.43133
```

# References
Rousseeuw, Peter J., and Katrien Van Driessen. "An algorithm for positive-breakdown
regression based on concentration steps." Data Analysis.
Springer, Berlin, Heidelberg, 2000. 335-346.
"""
function lts(setting::RegressionSetting; iters=nothing, crit=2.5)
    X = designMatrix(setting)
    y = responseVector(setting)
    return lts(X, y, iters=iters, crit=crit)
end

function lts(X::Array{Float64,2}, y::Array{Float64,1}; iters=nothing, crit=2.5)
    n, p = size(X)
    h = Int(floor((n + p + 1.0) / 2.0))
    if iters === nothing
        iters = minimum([500 * p, 3000])
    end
    allindices = collect(1:n)
    bestobjective = Inf
    besthsubset = []
    for iter in 1:iters
        subsetindices = sample(allindices, p, replace=false)
        objective, hsubsetindices = iterateCSteps(X, y, subsetindices, h)
        if objective < bestobjective
            bestobjective = objective
            besthsubset = hsubsetindices
        end
    end
    ltsreg = ols(X[besthsubset, :], y[besthsubset])
    ltsbetas = coef(ltsreg)
    ltsres = [y[i] - sum(X[i,:] .* ltsbetas) for i in 1:n]
    ltsS = sqrt(sum((ltsres.^2.0)[1:h]) / (h - p))
    ltsresmean = mean(ltsres[besthsubset])
    ltsScaledRes = (ltsres .- ltsresmean) / ltsS
    outlierindices = filter(i -> abs(ltsScaledRes[i]) > crit, 1:n)
    result = Dict()
    result["objective"] = bestobjective
    result["hsubset"] = besthsubset
    result["betas"] = ltsbetas
    result["S"] = ltsS
    result["outliers"] = outlierindices
    result["scaled.residuals"] = ltsScaledRes
    return result
end


using OrdinaryDiffEq, ODE, ODEInterfaceDiffEq, LSODA, Sundials, DiffEqDevTools, Plots; gr()

## Define the ThreeBody Problem
const threebody_μ = parse(Float64,"0.012277471")
const threebody_μ′ = 1 - threebody_μ

f = (du,u,p,t) -> begin
  @inbounds begin
  # 1 = y₁
  # 2 = y₂
  # 3 = y₁'
  # 4 = y₂'
  D₁ = ((u[1]+threebody_μ)^2 + u[2]^2)^(3/2)
  D₂ = ((u[1]-threebody_μ′)^2 + u[2]^2)^(3/2)
  du[1] = u[3]
  du[2] = u[4]
  du[3] = u[1] + 2u[4] - threebody_μ′*(u[1]+threebody_μ)/D₁ - threebody_μ*(u[1]-threebody_μ′)/D₂
  du[4] = u[2] - 2u[3] - threebody_μ′*u[2]/D₁ - threebody_μ*u[2]/D₂
  end
end

t₀ = 0.0; T = parse(Float64,"17.0652165601579625588917206249")
tspan = (t₀,2T)

prob = ODEProblem(f,[0.994, 0.0, 0.0, parse(Float64,"-2.00158510637908252240537862224")],tspan)

test_sol = TestSolution(T,[prob.u0])
abstols = 1.0 ./ 10.0 .^ (3:13); reltols = 1.0 ./ 10.0 .^ (0:10);

sol = solve(prob,Vern9(),abstol=1e-14,reltol=1e-14)
@show sol[1] - sol[end]
@show sol[end] - prob.u0;

apr = appxtrue(sol,test_sol)
@show sol[end]
@show apr.u[end]
@show apr.errors

setups = [Dict(:alg=>DP5())
          #Dict(:alg=>ode45()) #fails
          Dict(:alg=>BS5())
          Dict(:alg=>Tsit5())
          Dict(:alg=>dopri5())];
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,save_everystep=false,numruns=100)
plot(wp)

setups = [Dict(:alg=>DP5(),:dense=>false)
          #Dict(:alg=>ode45()) # Fails
          Dict(:alg=>BS5(),:dense=>false)
          Dict(:alg=>Tsit5(),:dense=>false)
          Dict(:alg=>dopri5())];
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,numruns=100)
plot(wp)

setups = [Dict(:alg=>DP5())
          #Dict(:alg=>ode45()) #fails
          Dict(:alg=>BS5())
          Dict(:alg=>Tsit5())
          Dict(:alg=>dopri5())];
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,numruns=100)
plot(wp)

setups = [Dict(:alg=>DP5())
          Dict(:alg=>Vern6())
          Dict(:alg=>Vern7())
          Dict(:alg=>TanYam7())
          Dict(:alg=>Vern8())
          Dict(:alg=>DP8())
          Dict(:alg=>dop853())
          Dict(:alg=>Vern9())];
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,save_everystep=false,numruns=100)
plot(wp)

setups = [Dict(:alg=>DP5())
          Dict(:alg=>Vern6())
          Dict(:alg=>Vern7())
          Dict(:alg=>TanYam7())
          Dict(:alg=>Vern8())
          Dict(:alg=>DP8())
          Dict(:alg=>dop853())
          Dict(:alg=>Vern9())];
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,dense=false,numruns=100,verbose=false)
plot(wp)

setups = [Dict(:alg=>DP5())
          Dict(:alg=>Vern6())
          Dict(:alg=>Vern7())
          Dict(:alg=>TanYam7())
          Dict(:alg=>Vern8())
          Dict(:alg=>DP8())
          Dict(:alg=>dop853())
          Dict(:alg=>Vern9())];
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,numruns=100)
plot(wp)

#setups = [Dict(:alg=>ode78())
#          Dict(:alg=>VCABM())
#          Dict(:alg=>CVODE_Adams())];
#wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,dense=false,numruns=100)

setups = [Dict(:alg=>DP5())
          Dict(:alg=>lsoda())
          Dict(:alg=>Vern8())
          Dict(:alg=>ddeabm())
          Dict(:alg=>odex())
          Dict(:alg=>ARKODE(Sundials.Explicit(),order=6))
    ];
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,save_everystep=false,numruns=100)
plot(wp)

abstols = 1.0 ./ 10.0 .^ (3:13); reltols = 1.0 ./ 10.0 .^ (0:10);
setups = [Dict(:alg=>Tsit5())
          Dict(:alg=>Vern9())
          Dict(:alg=>AitkenNeville(min_order=1, max_order=9, init_order=4, threading=true))
          Dict(:alg=>ExtrapolationMidpointDeuflhard(min_order=1, max_order=9, init_order=4, threading=true))
          Dict(:alg=>ExtrapolationMidpointHairerWanner(min_order=2, max_order=11, init_order=4, threading=true))]
solnames = ["Tsit5","Vern9","AitkenNeville","Midpoint Deuflhard","Midpoint Hairer Wanner"]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,names=solnames,
                      save_everystep=false,verbose=false,numruns=100)
plot(wp)

setups = [Dict(:alg=>ExtrapolationMidpointDeuflhard(min_order=1, max_order=9, init_order=9, threading=false))
          Dict(:alg=>ExtrapolationMidpointHairerWanner(min_order=2, max_order=11, init_order=4, threading=false))
          Dict(:alg=>ExtrapolationMidpointHairerWanner(min_order=2, max_order=11, init_order=4, threading=true))
          Dict(:alg=>ExtrapolationMidpointHairerWanner(min_order=2, max_order=11, init_order=4, sequence = :romberg, threading=true))
          Dict(:alg=>ExtrapolationMidpointHairerWanner(min_order=2, max_order=11, init_order=4, sequence = :bulirsch, threading=true))]
solnames = ["Deuflhard","No threads","standard","Romberg","Bulirsch"]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,names=solnames,
                      save_everystep=false,verbose=false,numruns=100)
plot(wp)

setups = [Dict(:alg=>ExtrapolationMidpointHairerWanner(min_order=2, max_order=11, init_order=10, threading=true))
          Dict(:alg=>ExtrapolationMidpointHairerWanner(min_order=2, max_order=11, init_order=4, threading=true))
          Dict(:alg=>ExtrapolationMidpointHairerWanner(min_order=5, max_order=11, init_order=10, threading=true))
          Dict(:alg=>ExtrapolationMidpointHairerWanner(min_order=2, max_order=15, init_order=10, threading=true))
          Dict(:alg=>ExtrapolationMidpointHairerWanner(min_order=5, max_order=7, init_order=6, threading=true))]
solnames = ["1","2","3","4","5"]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,names=solnames,
                      save_everystep=false,verbose=false,numruns=100)
plot(wp)

using SciMLBenchmarks
SciMLBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


using ConservationLawsParticles
using SpecialFunctions, RecursiveArrayTools, DifferentialEquations, Plots
using Random
Random.seed!(42)

# model
V(t, x) = 0.0 # 2 + 0.5sin(2x) + 0.3sin(2sqrt(2)*x) + 0.2cos(2sqrt(7)*x)
W(t, x) = 0.0 # x^2/2 - abs(x)
W′(t, x) = 0.0 # x - sign(x)
mob(ρ) = 0.0 # max((ρ-1)^2 * (2ρ+1), 0)
smodel = DiffusiveSampledModel((V,), ((W′,),), (mob,), (Diffusion(1),))
imodel = DiffusiveIntegratedModel((V,), ((W,),), (mob,), (Diffusion(1),))

# initial condition
x0 = ArrayPartition(sort(randn(50)))

boxsol(t, x; a, b) = (erf((x-a)/2sqrt(t)) - erf((x-b)/2sqrt(t))) / (2*(b-a))
refsol(t, x) = sum(boxsol(t, x; a=x0[i-1], b=x0[i])
    for i in eachindex(x0)[2:end]) / (length(x0)-1)

# time span
tspan = (0., 0.1)

# ODE system for the particles
sprob = ODEProblem(velocities_diff!, x0, tspan, smodel)
# iprob = ODEProblem(velocities_diff!, x0, tspan, imodel)

abstol = reltol = 2e-8

@time ssol = solve(sprob, BS5(); abstol=abstol, reltol=reltol);
# @time isol = solve(iprob, BS5(); abstol=abstol, reltol=reltol);

# plot the particle trajectories
plot(legend=false, title="Trajectories", xlabel="time", ylabel="position")
plot!(ssol, vars=eachindex(x0)[1:5:end]; color=:blue)
savefig("diffusion5.png")

# plot an animation of the density
anim = @animate for t in range(sqrt.(tspan)...; length=101).^2
    plot(title="Density", xlabel="position", ylabel="density",
        xlims=(-5,3), ylims=(0,.8))
    plot!(x -> refsol(t, x); color=:blue, label="analytic")
    densityplot!(ssol(t); color=:red, label="particles")
end
gif(anim, "diffusion5.gif")

# plot an animation of the density
tspan = (0., .001)
anim = @animate for t in range(sqrt.(tspan)...; length=101).^2
    plot(title="Density", xlabel="position", ylabel="density",
        xlims=(-4,2), ylims=(0,1), legend=:topleft)
    plot!(x -> refsol(t, x); color=:blue, label="analytic")
    densityplot!(ssol(t); color=:red, label="particles")
end
gif(anim, "diffusion5.1.gif")


using LinearAlgebra
using ParallelStencil
using ParallelStencil.FiniteDifferences2D
using CUDA

include("part2_utils.jl")
include("krylov.jl")

""" Iterative methods that can be used as coarse solvers in Multigrid. """
@enum CoarseSolver_t begin
    jacobi
    conjugate_gradient
end

""" Input parameters for Multigrid. """
mutable struct MGOpt
    coarse_solve_size :: Int
    coarse_solver :: CoarseSolver_t
    execution_policy :: ExecutionPolicy_t

    MGOpt() = new(5, jacobi, parallel_shmem)
end

""" Preallocates all buffers/arrays required for all levels of Multigrid. """
function preallocate_buffers(nx, ny) :: Union{Dict{Symbol, Dict{Int, AbstractArray}}, Nothing}
     # prealloc all the buffers
     λx = (nx > ny) ? (nx-1) ÷ (ny-1) : 1
     λy = (ny > nx) ? (ny-1) ÷ (nx-1) : 1
     k_max = Int(log2(min(nx, ny)-1))
     prealloc_dict = Dict(
         :res_f   => Dict( k => @zeros(λx*2^k    +1, λy*2^k    +1) for k in 1:k_max ),
         :corr_f  => Dict( k => @zeros(λx*2^k    +1, λy*2^k    +1) for k in 1:k_max ),
         :corr_c  => Dict( k => @zeros(λx*2^(k-1)+1, λy*2^(k-1)+1) for k in 1:k_max ),
         :res_c   => Dict( k => @zeros(λx*2^(k-1)+1, λy*2^(k-1)+1) for k in 1:k_max ),
    )

    return prealloc_dict
end

""" Solves (∇^2 - c) u = f using matrix-free V-cycle Multigrid. """
function MGsolve_2DPoisson!(u::AbstractArray{Float64}, f::AbstractArray{Float64}, h::Float64, c::Float64, tol::Float64, niters::Int, apply_BCs::Bool; opt=MGOpt(), verbose=false, prealloc_dict::Union{Dict{Symbol, Dict{Int, AbstractArray}}, Nothing} = nothing)
    # set nx, ny
    nx, ny = size(u)

    @assert opt.coarse_solve_size <= min(nx, ny)
    @assert isinteger(log2(opt.coarse_solve_size - 1))

    # prealloc all the buffers
    if prealloc_dict == nothing
        prealloc_dict = preallocate_buffers(nx, ny)
    end

    f_rms = sqrt(sum(f.^2)/(nx*ny))

    tolf = tol * f_rms

    r_rms = 0.0
    for iter = 1:niters
        # apply Dirichlet and Neumann BCs for temperature T
        if (apply_BCs)
            apply_boundary_conditions!(u)
        end
        # execute V-cycle iteration
        r_rms = Vcycle_2DPoisson!(u, f, h, c, tol, opt.coarse_solve_size, opt.coarse_solver, opt.execution_policy, apply_BCs, prealloc_dict=prealloc_dict)
        @synchronize()

        if verbose
            println("$(iter) $(r_rms / f_rms)")
        end
        if (r_rms < tolf)
            if verbose
                println("V-cycle multigrid converged in $(iter) iterations.")
            end
            break
        end
    end

    if (r_rms > tolf)
        @warn "V-cycle multigrid failed to converge within" niters "iterations."
    end

    return r_rms

end

""" Runs a single V-cycle iteration of Multigrid.

Note: nx, ny must be = (2^k)+1, where k ∈ N.

"""
function Vcycle_2DPoisson!(u_f::AbstractArray{Float64}, rhs::AbstractArray{Float64}, h::Float64, c::Float64, tol::Float64, coarse_solve_size::Int, coarse_solver::CoarseSolver_t, execution_policy::ExecutionPolicy_t, apply_BCs::Bool; prealloc_dict::Union{Dict{Symbol, Dict{Int, AbstractArray}}, Nothing} = nothing)

    nx, ny = size(u_f)

    if ( nx-1 != 2*((nx-1)÷2) || ny-1 != 2*((ny-1)÷2) )
        error("ERROR:not a power of 2")
    end

    # coarse grid size
    nxc = 1+(nx-1)÷2
    nyc = 1+(ny-1)÷2

    k_curr = Int(log2(min(nx, ny)-1))

    if prealloc_dict == nothing
        prealloc_dict = preallocate_buffers(nx, ny)
    end

    res_f   = prealloc_dict[:res_f][k_curr]
    corr_f  = prealloc_dict[:corr_f][k_curr]
    corr_c  = prealloc_dict[:corr_c][k_curr]
    res_c   = prealloc_dict[:res_c][k_curr]

    res_f   .= 0.0
    corr_f  .= 0.0
    corr_c  .= 0.0
    res_c   .= 0.0

    res_rms = 0.0

    if (min(nx, ny) > coarse_solve_size)   # not the coarsest level

        #---------- take 2 iterations on the fine grid--------------
        res_rms = iteration_2DPoisson!(u_f, rhs, h, c, res_f, execution_policy)
        res_rms = iteration_2DPoisson!(u_f, rhs, h, c, res_f, execution_policy)

        #--------- restrict the residual to the coarse grid --------
        residual_2DPoisson_wrapper!(u_f, rhs, h, c, res_f, execution_policy)
        restrict_wrapper!(res_f, res_c, apply_BCs, execution_policy)

        #---------- solve for the coarse grid correction -----------
        corr_c .= 0.
        res_rms = Vcycle_2DPoisson!(corr_c, res_c, h*2, c, tol, coarse_solve_size, coarse_solver, execution_policy, apply_BCs) # *RECURSIVE CALL*

        #---- prolongate (interpolate) the correction to the fine grid
        prolongate_wrapper!(corr_c, corr_f, apply_BCs, execution_policy)

        #---------- correct the fine-grid solution -----------------
        u_f .= u_f - corr_f

        #---------- two more smoothing iterations on the fine grid---
        res_rms = iteration_2DPoisson!(u_f, rhs, h, c, res_f, execution_policy)
        res_rms = iteration_2DPoisson!(u_f, rhs, h, c, res_f, execution_policy)

    else

        if coarse_solver == jacobi
            #----- coarsest level (ny=5): iterate to get 'exact' solution
            coarse_solve_iters = 20*coarse_solve_size
            tol_rhs = tol * sqrt(sum(rhs.^2)/(nx*ny))
            for i = 1:coarse_solve_iters  # heuristic
                res_rms = iteration_2DPoisson!(u_f, rhs, h, c, res_f, execution_policy)
                if res_rms < tol_rhs
                    break
                end
            end
            if res_rms > tol_rhs
                @debug "Coarse solve of V-cycle multigrid did not converge to " tol " within " coarse_solve_iters "iterations. This is usually fine as long as it converges in a fixed amount of Vcycles."
            end
        elseif coarse_solver == conjugate_gradient
            coarse_solve_iters = 20*coarse_solve_size
            res_rms = cg!(u_f, rhs, h, h, c, tol, coarse_solve_iters, execution_policy=execution_policy)
        else
            error()
        end

    end

    return res_rms   # returns the rms. residual
end

""" Computes the residual res = (∇^2 - c) u - f, without using shared memory. """
@parallel_indices (ix, iy) function residual_2DPoisson!(u::AbstractArray{Float64}, f::AbstractArray{Float64}, h::Float64, c::Float64, res::AbstractArray{Float64})
    nx, ny = size(u)
    C = (4.0 + c * h^2)
    _h2 = 1 / h^2
    if (1 < ix < nx && 1 < iy < ny)
        res[ix, iy] = (
                       (  u[ix+1, iy]
                        + u[ix-1, iy]
                        + u[ix, iy+1]
                        + u[ix, iy-1]
                        - C * u[ix, iy] ) * _h2
                       - f[ix, iy]
                      )
    end
    return nothing
end

""" Computes the residual res = (∇^2 - c) u - f, using shared memory. """
@parallel_indices (ix, iy) function residual_2DPoisson_shmem!(u_::AbstractArray{Float64}, f::AbstractArray{Float64}, h::Float64, c::Float64, res::AbstractArray{Float64})
    nx, ny = size(u_)
    C = (4.0 + c * h^2)
    _h2 = 1 / h^2

    tx = @threadIdx().x + 1
    ty = @threadIdx().y + 1
    u = @sharedMem(eltype(u_), (@blockDim().x+2, @blockDim().y+2))
    if (1 <= ix <= nx && 1 <= iy <= ny)
        u[tx, ty] = u_[ix, iy]
    end
    @sync_threads()
    if (1 < ix < nx && 1 < iy < ny)
        u[tx, ty] = u_[ix, iy]
        if (@threadIdx().x == 1)             u[tx-1, ty] = u_[ix-1, iy] end
        if (@threadIdx().x == @blockDim().x) u[tx+1, ty] = u_[ix+1, iy] end
        if (@threadIdx().y == 1)             u[tx, ty-1] = u_[ix, iy-1] end
        if (@threadIdx().y == @blockDim().y) u[tx, ty+1] = u_[ix, iy+1] end
        @sync_threads()
        res[ix, iy] = (
                       (  u[tx+1, ty]
                        + u[tx-1, ty]
                        + u[tx, ty+1]
                        + u[tx, ty-1]
                        - C * u[tx, ty] ) * _h2
                       - f[ix, iy]
                      )
    end
    return nothing
end

""" Wrapper function computing residual res = (∇^2 - c) u - f, for various execution policies. """
function residual_2DPoisson_wrapper!(u_f, rhs, h, c, res_f, execution_policy)
    nx, ny = size(u_f)

    if execution_policy == parallel_shmem
        threads = (32, 8)
        blocks = ((nx, ny) .- 1) .÷ (threads) .+ 1
        shmem = prod(threads.+2) * sizeof(Float64)
        @parallel blocks threads shmem=shmem residual_2DPoisson_shmem!(u_f, rhs, h, c, res_f)
    elseif execution_policy == parallel
        @parallel residual_2DPoisson!(u_f, rhs, h, c, res_f)
    elseif execution_policy == serial
        error() # Not implemented yet
    else
        error()
    end
end

""" Performs a single Jacobi iteration on field u.

Note: alpha = 4.0 / 5.0 is the optimal damping factor for smoothing using Jacobi iterations.

"""
function iteration_2DPoisson!(u, f, h, c, res, execution_policy; alpha = 4.0 / 5.0)
    nx, ny = size(u)

    # Do one Jacobi iteration
    residual_2DPoisson_wrapper!(u, f, h, c, res, execution_policy)

    # update r_rms
    r_rms = sqrt(sum(res.^2) / (nx * ny))

    # update u
    u .+= alpha * (h^2 / (4.0 + c * h^2)) .* res

    return r_rms
end

""" Performs a single Gauss-Seidel iteration on field u.

Note: alpha = 1.0 is the optimal damping factor for smoothing using Gauss-Seidel iterations.

Gauss-Seidel is highly serial so in our current Multigrid implementation we replaced it by Jacobi.

We only use this function for debugging.

"""
@views function iteration_2DPoisson_gs!(u, f, h, c; alpha = 1.0)

    nx, ny = size(u)

    # Gauss-Seidel iteration
    r_rms = 0.0
    for j = 2:ny-1
        for i = 2:nx-1

            # compute residual
            r = ( u[i+1, j]
                 + u[i-1, j]
                 + u[i, j+1]
                 + u[i, j-1]
                 - (4.0 + c * h^2) * u[i, j] ) / h^2 - f[i, j]

            # update u
            u[i, j] = u[i, j] + alpha * (h^2 / (4.0 + c * h^2)) * r

            # update r_rms
            r_rms = r_rms + r^2

        end
    end

    r_rms = sqrt(r_rms / (nx * ny))

    return r_rms
end

""" Restricts a field from a fine grid to a coarse grid of half the resolution.

This is a serial version, mainly used for debugging.

"""
@views function  restrict_serial!(fine, coarse, apply_BCs)
    nx = size(fine, 1)
    ny = size(fine, 2)

    # initalize coarse to zero (+ Dirichlet(0) BCs)
    coarse .= 0.0

    # apply restriction stencil
    jc = 2
    for j = 3:2: ny-2
        ic = 2
        for i = 3:2: nx-2
            coarse[ic, jc] = fine[i, j]
            ic = ic + 1
        end
        jc = jc + 1
    end

    # apply Neumann BCs for temperature T
    if apply_BCs
        apply_boundary_conditions_neumann!(coarse)
    end

end

""" Restricts a field from a fine grid to a coarse grid of half the resolution. """
@parallel_indices (ix, iy) function restrict!(fine, coarse)
    nx, ny = size(fine)
    if (ix % 2 == 1 && iy % 2 == 1) && (3 <= ix <= nx-2 && 3 <= iy <= ny-2)
        ix_c, iy_c = ((ix-1)÷2)+1, ((iy-1)÷2)+1
        coarse[ix_c, iy_c] = fine[ix, iy]
    end
    return nothing
end

""" Restrict wrapper.

Calls the proper restrict function given an execution policy and applies boundary conditions on the coarser level.

"""
@views function restrict_wrapper!(fine, coarse, apply_BCs, execution_policy)
    # initalize coarse to zero (+ Dirichlet(0) BCs)
    coarse[:] .= 0.0

    if execution_policy in [parallel, parallel_shmem]
        @parallel restrict!(fine, coarse)
    else
        restrict_serial!(fine, coarse)
    end

    # apply Neumann BCs for temperature T
    if apply_BCs
        apply_boundary_conditions_neumann!(coarse)
    end
end

""" Prolongates a field from a coarse grid to a fine grid of double the resolution.

This is a serial version, mainly used for debugging.

"""
@views function prolongate_serial!(coarse, fine, apply_BCs)
    a2, a4 = 1.0 / 2.0, 1.0 / 4.0

    nx = size(fine, 1)
    ny = size(fine, 2)

    # initialize fine to zero (+ Dirichlet(0) BCs)
    fine .= 0.0

    for j = 1:ny
        for i = 1: nx
            ic, jc = ((i-1)÷2)+1, ((j-1)÷2)+1
            if (i % 2 == 1 && j % 2 == 1) && (3 <= i <= nx-2 && 3 <= j <= ny-2)
                fine[i, j] = coarse[ic, jc]
                fine[i+1, j] = fine[i+1, j] + a2 * coarse[ic, jc]
                fine[i-1, j] = fine[i-1, j] + a2 * coarse[ic, jc]
                fine[i, j+1] = fine[i, j+1] + a2 * coarse[ic, jc]
                fine[i, j-1] = fine[i, j-1] + a2 * coarse[ic, jc]
                fine[i+1, j+1] = fine[i+1, j+1] + a4 * coarse[ic, jc]
                fine[i+1, j-1] = fine[i+1, j-1] + a4 * coarse[ic, jc]
                fine[i-1, j+1] = fine[i-1, j+1] + a4 * coarse[ic, jc]
                fine[i-1, j-1] = fine[i-1, j-1] + a4 * coarse[ic, jc]
            end
        end
    end

    # apply Neumann BCs for temperature T
    if apply_BCs
        apply_boundary_conditions_neumann!(fine)
    end

end

""" Prolongates a field from a coarse grid to a fine grid of double the resolution.

Uses atomic operation for avoiding issues when run on the GPU.

"""
@parallel_indices (ix, iy) function prolongate_with_atomic!(coarse, fine)
    nx, ny = size(fine)
    a2 = 1.0 / 2.0
    a4 = 1.0 / 4.0
    if (ix % 2 == 1 && iy % 2 == 1) && (3 <= ix <= nx-2 && 3 <= iy <= ny-2)
        ix_c, iy_c = ((ix-1)÷2)+1, ((iy-1)÷2)+1
        @atomic fine[ix, iy] = fine[ix, iy] + coarse[ix_c, iy_c]
        @atomic fine[ix+1, iy] = fine[ix+1, iy] + a2 * coarse[ix_c, iy_c]
        @atomic fine[ix-1, iy] = fine[ix-1, iy] + a2 * coarse[ix_c, iy_c]
        @atomic fine[ix, iy+1] = fine[ix, iy+1] + a2 * coarse[ix_c, iy_c]
        @atomic fine[ix, iy-1] = fine[ix, iy-1] + a2 * coarse[ix_c, iy_c]
        @atomic fine[ix+1, iy+1] = fine[ix+1, iy+1] + a4 * coarse[ix_c, iy_c]
        @atomic fine[ix+1, iy-1] = fine[ix+1, iy-1] + a4 * coarse[ix_c, iy_c]
        @atomic fine[ix-1, iy+1] = fine[ix-1, iy+1] + a4 * coarse[ix_c, iy_c]
        @atomic fine[ix-1, iy-1] = fine[ix-1, iy-1] + a4 * coarse[ix_c, iy_c]
    end
    return nothing
end

""" Prolongates a field from a coarse grid to a fine grid of double the resolution.

This version does not use atomics. We only used for execution on the CPU.

"""
@parallel_indices (ix, iy) function prolongate!(coarse, fine)
    nx, ny = size(fine)
    a2 = 1.0 / 2.0
    a4 = 1.0 / 4.0
    if (ix % 2 == 1 && iy % 2 == 1) && (3 <= ix <= nx-2 && 3 <= iy <= ny-2)
        ix_c, iy_c = ((ix-1)÷2)+1, ((iy-1)÷2)+1
        fine[ix, iy] = fine[ix, iy] + coarse[ix_c, iy_c]
        fine[ix+1, iy] = fine[ix+1, iy] + a2 * coarse[ix_c, iy_c]
        fine[ix-1, iy] = fine[ix-1, iy] + a2 * coarse[ix_c, iy_c]
        fine[ix, iy+1] = fine[ix, iy+1] + a2 * coarse[ix_c, iy_c]
        fine[ix, iy-1] = fine[ix, iy-1] + a2 * coarse[ix_c, iy_c]
        fine[ix+1, iy+1] = fine[ix+1, iy+1] + a4 * coarse[ix_c, iy_c]
        fine[ix+1, iy-1] = fine[ix+1, iy-1] + a4 * coarse[ix_c, iy_c]
        fine[ix-1, iy+1] = fine[ix-1, iy+1] + a4 * coarse[ix_c, iy_c]
        fine[ix-1, iy-1] = fine[ix-1, iy-1] + a4 * coarse[ix_c, iy_c]
    end
    return nothing
end

""" Prolongate wrapper.

Calls the proper prolongate function given an execution policy and applies boundary conditions on the finer level.

"""
@views function prolongate_wrapper!(coarse, fine, apply_BCs, execution_policy)
    # initialize fine to zero (+ Dirichlet(0) BCs)
    fine[:] .= 0.0

    # We need the @atomic macro to avoid race-conditions on the GPU,
    # but this is not supported for CPU :(
    if execution_policy in [parallel, parallel_shmem]
        if USE_GPU
            @parallel prolongate_with_atomic!(coarse, fine)
        else
            @parallel prolongate!(coarse, fine)
        end
    else
        prolongate_serial!(coarse, fine)
    end

    # apply Neumann BCs for temperature T
    if apply_BCs
        apply_boundary_conditions_neumann!(fine)
    end

end


using LinearAlgebra
using Test

include("test_solve_eye_plus_at_kron_b.jl")


# Julia1_3.1.jl  Conditional Control Structures
#
# Boolean values are
true, false # are Boolean values (type Bool)
#   (one byte values that are 1 or 0).
# Flow control - Conditionals if-block
# if-elseif-else-end can be nested to any level
x=2; y=1
if x<y
    println("x is less than y")
elseif x>y
    println("x is greater than y")
else
    println("x is equal to y")
end

# you can have a chain of elseifs for
# handling different cases

# watch out for x=1//3; y=1/3, won't be equal
# however rationalize(1/3)==1//3
# 1 ≈ 0.999999999 where ≈ means approximately equal
# 3.825==153/40 this is false
# 3.825≈153/40  this is true

# one liner conditional borrowed from C
# cond? true-expr:false-expr
x=-2
y = x<0 ? -x : x  # gets absolute value of x
# same as y=abs(x)

# other one liner conditional executions
# if <cond>; <expression>; end can be reduced to:
#  where && read as "and then"
x=-3
x<0 && (x=-x) # if x<0 and then x set to -x
# if !<cond>; <expression>; end   can be reduced to:
#  where || read as "or else"
x=-4
!(x<0) || (x=-x) # if not x<0 bypass else x set to -x
# in this case both perform x=abs(x)

# the above operations use a side effect of eval
# when used in this fashion
#  It can obsfuscate code

# Now as clarity intended.
# logical operator && is "and"
# logical operator || is non-exclusice "or"
# a && b is true, if and only if, both a and b are true.
# a || b is true, if either a or b or both are true.
 3<5 && 7>6  # is true
 3>5 || 2<5  # is true, note had to continue to eval
 3>5 && 7>6  # 3>5 false so needn't eval rest

# numeric(and alphanumeric)comparison operators
# < less than, > greater than, == equal to
# <= or ≤  less than or equal,     ≤ is Char(0x2264)
# >= or ≥  greater than or equal,  ≥ is Char(0x2265)
# != or ≠  not equal,              ≠ is Char(0x2260)
# != or ≠  same as !(x==y) which is often preferred
# many other operators for checking types, etc.

# operators can be chained:
1<2≤2<3==3>2≥1==1<3≠5
# results in true

# logical Operators
# && = AND, || = OR,
# xor(x,y) XOR true if x or y true but not both!
xor(true,true)

x=1; y=2
x<5 && y>2
# false
0.0≤x≤1.0 && 1.0≤y≤2.0
# true


const MP = MultivariatePolynomials
@testset "Division" begin
    @testset "GCD and LCM" begin
        Mod.@polyvar x y z
        @test gcd(x*y, y) == y
        @test lcm(x*y, y) == x*y
        @test gcd(x*y^2*z, y*z^3) == y*z
        @test lcm(x*y^2*z, y*z^3) == x*y^2*z^3
        @test gcd(x^2*y^7*z^3, x^4*y^5*z^2) == x^2*y^5*z^2
        @test lcm(x^2*y^7*z^3, x^4*y^5*z^2) == x^4*y^7*z^3
    end
    # Taken from
    # Ideals, Varieties, and Algorithms
    # Cox, Little and O'Shea, Fourth edition
    # They have been adapted to the grlex ordering
    @testset "Divides" begin
        Mod.@polyvar x y z
        @test divides(leadingmonomial(x+y), x) # leadingmonomial(x+y) will be x^1*y^0 -> tricky test !
        @test !divides(leadingmonomial(x^2+y), x)
        @test divides(x*y, x^2*y)
        @test divides(x*y, x*y^2)
        @test divides(y*z, x*y*z)
        @test !divides(y*z, x*z)
        @test !divides(x^2*y, x*y^2)
    end
    @testset "Leading function" begin
        Mod.@polyvar x y z
        @test @inferred(monic(2x^2 + 4y + 2)) == x^2 + 2y + 1
        # See page 60
        p = 4x*y^2*z + 4z^2 + 7x^2*z^2 - 5x^3
        @test leadingcoefficient(p) == 7
        @test leadingmonomial(p) == x^2*z^2
        @test leadingterm(p) == 7x^2*z^2
    end
    @testset "Division examples" begin
        Mod.@polyvar x y
        @test iszero(@inferred MP.proddiff(2x*y, 3y^2*x))
        @test (@inferred div(x*y^2 + 1, x*y + 1)) == y
        @test (@inferred rem(x*y^2 + 1, x*y + 1)) == -y + 1
        @test (@inferred div(x*y^2 + x, y)) == x*y
        @test (@inferred rem(x*y^2 + x, y)) == x
        @test (@inferred rem(x^4 + x^3 + (1+1e-10)*x^2 + 1, x^2 + x + 1)) == 1
        @test (@inferred rem(x^4 + x^3 + (1+1e-10)*x^2 + 1, x^2 + x + 1; ztol=1e-11)) ≈ -((1+1e-10)-1)x + 1
    end
    @testset "Division by multiple polynomials examples" begin
        function testdiv(p, ps)
            q, r = @inferred divrem(p, ps)
            @test p == dot(q, ps) + r
            q, r
        end
        Mod.@polyvar x y
        # Example 1
        q, r = testdiv(x*y^2 + 1, [x*y + 1, y + 1])
        @test q == [y, -1]
        @test r == 2
        # Example 2
        q, r = testdiv(x^2*y + x*y^2 + y^2, [x*y - 1, y^2 - 1])
        @test q == [x + y, 1]
        @test r == x + y + 1
        # Example 4
        q, r = testdiv(x^2*y + x*y^2 + y^2, [y^2 - 1, x*y - 1])
        @test q == [x + 1, x]
        @test r == 2x + 1
        # Example 5
        q, r = testdiv(x*y^2 - x, [x*y - 1, y^2 - 1])
        @test q == [y, 0]
        @test r == -x + y
        q, r = testdiv(x*y^2 - x, [y^2 - 1, x*y - 1])
        @test q == [x, 0]
        @test r == 0

        @test (@inferred rem(x^2*y + (1+1e-10)*x*y + 1, [x^2 + x, y + 1])) == 1
        @test (@inferred rem(x^2*y + (1+1e-10)*x*y + 1, [x^2 + x, y + 1]; ztol=1e-11)) == -((1+1e-10)-1)x + 1
    end
    @testset "gcd" begin
        Mod.@polyvar x
        @test -2(x + 1) == @inferred gcd(x^2 - 1, x^2 + 2x + 1)
        @test  2(x + 1) == @inferred gcd(x^2 + 2x + 1, x^2 - 1)
        @test    x + 1  == @inferred gcd(x^2 - 1, x + 1)
        @test    x + 1  == @inferred gcd(x + 1, x^2 - 1)
        @test    x + 1  == @inferred gcd(x - x, x + 1)
        @test    x + 1  == @inferred gcd(x + 1, x - x)
        @test        0  == @inferred gcd(x - x, x^2 - x^2)
        @test        0  == @inferred gcd(x^2 - x^2, x - x)
    end
    @testset "lcm" begin
        Mod.@polyvar x
        @test -(x-1) * (x+1)^2 / 2 == @inferred lcm(x^2 - 1, x^2 + 2x + 1)
        @test  (x-1) * (x+1)^2 / 2 == @inferred lcm(x^2 + 2x + 1, x^2 - 1)
        @test x^2 - 1 == @inferred lcm(x^2 - 1, x - 1)
        @test x^2 - 1 == @inferred lcm(x - 1, x^2 - 1)
        @test       0 == @inferred lcm(x - x, x + 1)
        @test       0 == @inferred lcm(x + 1, x - x)
        @test       0 == @inferred lcm(x - x, x^2 - x^2)
        @test       0 == @inferred lcm(x^2 - x^2, x - x)
    end
end


using OrdinaryDiffEq, ParameterizedFunctions, ODE, ODEInterfaceDiffEq,
      LSODA, Sundials, DiffEqDevTools

f = @ode_def FitzhughNagumo begin
  dv = v - v^3/3 -w + l
  dw = τinv*(v +  a - b*w)
end a b τinv l

p = [0.7,0.8,1/12.5,0.5]
prob = ODEProblem(f,[1.0;1.0],(0.0,10.0),p)

abstols = 1.0 ./ 10.0 .^ (6:13)
reltols = 1.0 ./ 10.0 .^ (3:10);

sol = solve(prob,Vern7(),abstol=1/10^14,reltol=1/10^14)
test_sol = TestSolution(sol)
using Plots; gr()

plot(sol)

setups = [Dict(:alg=>DP5())
          #Dict(:alg=>ode45()) #fails
          Dict(:alg=>dopri5())
          Dict(:alg=>BS5())
          Dict(:alg=>Tsit5())
          Dict(:alg=>Vern6())
]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,save_everystep=false,numruns=100,maxiters=1000)
plot(wp)

setups = [Dict(:alg=>DP5())
          #Dict(:alg=>ode45()) # fails
          Dict(:alg=>BS5())
          Dict(:alg=>Tsit5())
          Dict(:alg=>Vern6())
]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,numruns=100,maxiters=10000,error_estimate=:L2,dense_errors=true)
plot(wp)

setups = [Dict(:alg=>DP8())
          #Dict(:alg=>ode78()) # fails
          Dict(:alg=>Vern7())
          Dict(:alg=>Vern8())
          Dict(:alg=>dop853())
          Dict(:alg=>Vern6())
]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,save_everystep=false,numruns=100,maxiters=1000)
plot(wp)

setups = [Dict(:alg=>DP8())
          Dict(:alg=>Vern7())
          Dict(:alg=>CVODE_Adams())
          Dict(:alg=>ARKODE(Sundials.Explicit(),order=6))
          #Dict(:alg=>lsoda())
          Dict(:alg=>odex())
          Dict(:alg=>ddeabm())
]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,save_everystep=false,numruns=100,maxiters=1000)
plot(wp)

setups = [Dict(:alg=>DP8())
          #Dict(:alg=>ode78()) # fails
          Dict(:alg=>Vern7())
          Dict(:alg=>Vern8())
          Dict(:alg=>Vern6())
]
wp = WorkPrecisionSet(prob,abstols,reltols,setups;appxsol=test_sol,numruns=100,maxiters=1000,error_estimate=:L2,dense_errors=true)
plot(wp)

using DiffEqBenchmarks
DiffEqBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


"""
ContinuousTimeMarkovChain{T}

Type for storing a continuous time Markov chain.

Can be constructed directly by supplying the infinitesimal generator infinitesimal_generator, or by supplying an ItoDiffusionProcess to be approximated.

See also https://en.wikipedia.org/wiki/Continuous-time_Markov_chain#Stationary_distribution

"""
struct ContinuousTimeMarkovChain{T}

	infinitesimal_generator::AbstractMatrix{<:Real}
	states::AbstractVector{T}

	# Standard Constructor with Consistency Checks
	function ContinuousTimeMarkovChain(infinitesimal_generator::AbstractMatrix{<:Real},
									   states::AbstractVector{T}) where T

		# Check infinitesimal generator is a square matrix
		if ~ size(infinitesimal_generator, 1) == size(infinitesimal_generator, 2)
			error("infinitesimal generator must be a squadre matrix.")
		end

		# Check infinitesimal generator rows sum to 0
		# Numerically, this means that the distance between sum(row_i) and 0.0
		# can at most be the machine precision around the value of the diagonal
		# element. On top of that we add 2 orders of magnitude.
		row_sums = abs.(sum(infinitesimal_generator, dims=2))
		T_i = typeof(infinitesimal_generator[1,1])
		if T_i <: Integer
			if any(row_sums .> zero(T_i))
				error("Rows of the infinitesimal generator must sum to 0.")
			end
		else
			check = false
			for (idx, row_sum) in enumerate(row_sums)
				check = ~isapprox(row_sum, zero(T_i), atol=1e2 .* eps(infinitesimal_generator[idx,idx]))
				check ? error("Rows of the infinitesimal generator must sum to 0.") : continue
			end
		end

		# Check sign restriction
		if any(diag(infinitesimal_generator, 0) .> 0)
			error("Diagonal elements of the infinitesimal generator must be negative or zero.")
		end

		# Check infinitesimal generator and supplied state vector have compatible sizes
		if size(infinitesimal_generator, 1) != length(states)
			error("Size mismatch between infinitesimal generator infinitesimal_generator and the state vector states.")
		end

		# Allocate object
		new{T}(infinitesimal_generator, states)
	end

end

# Constructor which omits State Vector, assume state space is a range vector.
function ContinuousTimeMarkovChain(infinitesimal_generator::AbstractMatrix{<:Real})

	# If State Vector not provided, then use 1:size(infinitesimal_generator,1) as placeholder
	states = Vector{Int64}(1:size(infinitesimal_generator, 1))

	# Allocate object
	return ContinuousTimeMarkovChain(infinitesimal_generator, states)
end

# Constructor for single-dimensional Diffusion Process
function ContinuousTimeMarkovChain(dp::ItoDiffusionProcess, grid::Vector{<:Real})

	# Grab length of State Vector
	n = length(grid)

	# Preallocate intensity vectors
	diag = Vector{Float64}(undef, n)
	fwd = similar(diag[1:end-1])
	bwd = similar(fwd)
	dS = Vector{Float64}(undef, n-1)

	# Unequally spaced grids are supported
	dS = diff(grid)
	dS_dwn = vcat(dS[1], dS)
	dS_up = vcat(dS, dS[end])
	dS_avg = (dS_up + dS_dwn) ./ 2
	# dS_avg[1] = dS_avg[1] ./2
	# dS_avg[end] = dS_avg[end] ./2

	# Compute drift and diffusion terms
	drift = dp.μ.(grid)
	diffusion = dp.σ.(grid) .^ 2 ./ 2

	# Compute intensity up and down using Upwind Scheme
	bwd = - min.(drift, 0) ./ dS_dwn + diffusion ./ (dS_avg .* dS_dwn)
	fwd = max.(drift, 0) ./ dS_up + diffusion ./ (dS_avg .* dS_up)

	# By construction, outflow intensity on diagonal s.t. rows sum to 0
	bwd[1], fwd[end] = 0, 0
	diag = - fwd - bwd

	# Populate Infinitesimal Generator as a Sparse Matrix
	infinitesimal_generator = spdiagm(-1 => bwd[2:end], 0 => diag, 1 => fwd[1:end-1])

	# Allocate Object
	return ContinuousTimeMarkovChain(infinitesimal_generator, grid)
end

# Overloading the show operator to get pretty-printing behavior.
function Base.show(io::IO, ::MIME"text/plain", mc::ContinuousTimeMarkovChain)
	print(io, "Continuous Time Markov Chain object.\n")
	bigchain = length(mc.states) > 10
	if bigchain
		print(io, "Stored objects are too large. Showing only first few elements.\n")
		print(io, "\nInfinitesimal generator:\n")
		display(Matrix(mc.infinitesimal_generator[1:10, 1:10]))
		print(io, "\nState space:\n")
		display(mc.states[1:10])
	else
		print(io, "\nInfinitesimal generator:\n")
		display(Matrix(mc.infinitesimal_generator))
		print(io, "\nState space:\n")
		display(mc.states)
	end
end

"""
	stationary_distribution(mc::ContinuousTimeMarkovChain, tol::T)

Computes the stationary distribution of the continuous time Markov chain mc.

By convention, the stationary distribution is treated as a PMF.

See also https://en.wikipedia.org/wiki/Continuous-time_Markov_chain#Stationary_distribution
"""
function stationary_distribution(mc::ContinuousTimeMarkovChain{T},
								 step_size::Real=1e8,
								 maxit::Integer=20,
								 tol::Real=1e-8) where T

	# Delta parameter for implicit method
	Δ = step_size

	# Initial guess is uniform over all states
	g0 = ones(Float64, size(mc.states)) ./ length(mc.states)
	g = similar(g0)

	# Compute RHS matrix in KFE
	B = I - Δ .* transpose(mc.infinitesimal_generator)

	#
	it = 0
	supnorm = 1.0e6
	while supnorm > tol && it < maxit
		g = B \ g0
		supnorm = maximum(abs.(g .- g0))
		g0 = g
		it += 1
	end

	if it == maxit
		error("Algorithm for finding the stationary distribution did not converge.")
	end

	g ./= sum(g)

	if isa(Real, T)
		return DiscreteNonParametric(mc.states, g)
	else
		return Categorical(g)
	end
end

"""
	random_sample(mc::ContinuousTimeMarkovChain, Is::Integer=1, draws::Real=1000)

Given a starting index Is of type Integer, it returns a sequence of draws of length N.

The output Tuple contains a cumulative elapsed time T and the pseudorandom sequence of indices.

"""
function random_sample(mc::ContinuousTimeMarkovChain, Is::Integer=1, N::Real=1000)

	spy = pattern(mc)
	degenerate = size(spy,1) == 1

	traj = Vector{Int64}(undef, N + 1)
	traj[1] = Is

	T = Vector{Float64}(undef, N + 1)
	T[1] = 0

	if ~ degenerate
		for draw = 1:N

			Δ, traj[draw+1] = conditional_draw(mc, traj[draw], spy)

			T[draw+1] = T[draw] + Δ

		end
	else

		traj[2:end] = Is
		T[2:end] = Inf

	end

	return T, traj
end

function conditional_draw(mc::ContinuousTimeMarkovChain, orig_idx::Integer, pattern::AbstractMatrix{Tv}) where Tv<:Bool

	ran = 1:length(mc.states)

	cols = ran[view(pattern, orig_idx, :)]

	Δ, idx = findmin(rand.(Exponential.(1 ./ view(mc.infinitesimal_generator, orig_idx, cols))))

	return Δ, cols[idx]

end

# Better not to have to use this one!!!
function conditional_draw(mc::ContinuousTimeMarkovChain, orig_idx::Integer)

	return conditional_draw(mc, orig_idx, pattern(mc))

end

function pattern(mc::ContinuousTimeMarkovChain)

	return mc.infinitesimal_generator .> 0.0

end


# ------------------------------------------------------------------------------------------
# # Linear models, loss functions, gradients, SGD
# (c) Deniz Yuret, 2019
# * Objectives: Define, train and visualize a simple model; understand gradients and SGD;
# learn to use the GPU.
# * Prerequisites: [Callable
# objects](https://docs.julialang.org/en/v1/manual/methods/#Function-like-objects-1),
# [Generator expressions](https://docs.julialang.org/en/v1/manual/arrays/#Generator-
# Expressions-1), [MNIST](20.mnist.ipynb), [Iterators](25.iterators.ipynb)
# * New functions:
# [mnistdata](https://github.com/denizyuret/Knet.jl/blob/master/data/mnist.jl),
# [accuracy](http://denizyuret.github.io/Knet.jl/latest/reference/#Knet.accuracy),
# [zeroone](http://denizyuret.github.io/Knet.jl/latest/reference/#Knet.zeroone),
# [nll](http://denizyuret.github.io/Knet.jl/latest/reference/#Knet.nll),
# [Param, @diff, value, params,
# grad](http://denizyuret.github.io/Knet.jl/latest/reference/#AutoGrad),
# [sgd](http://denizyuret.github.io/Knet.jl/latest/reference/#Knet.sgd),
# [progress,
# progress!](http://denizyuret.github.io/Knet.jl/latest/reference/#Knet.progress),
# [gpu](http://denizyuret.github.io/Knet.jl/latest/reference/#Knet.gpu),
# [KnetArray](http://denizyuret.github.io/Knet.jl/latest/reference/#Knet.KnetArray),
# [load](http://denizyuret.github.io/Knet.jl/latest/reference/#Knet.load),
# [save](http://denizyuret.github.io/Knet.jl/latest/reference/#Knet.save)
#
#
# <img src="https://www.oreilly.com/library/view/tensorflow-for-
# deep/9781491980446/assets/tfdl_0401.png" alt="A linear model" width=300/> ([image
# source](https://www.oreilly.com/library/view/tensorflow-for-deep/9781491980446/ch04.html))
#
# In Knet, a machine learning model is defined using plain Julia code. A typical model
# consists of a **prediction** and a **loss** function. The prediction function takes some
# input, returns the prediction of the model for that input. The loss function measures how
# bad the prediction is with respect to some desired output. We train a model by adjusting
# its parameters to reduce the loss.
#
# In this section we will implement a simple linear model to classify MNIST digits. The
# prediction function will return 10 scores for each of the possible labels 0..9 as a linear
# combination of the pixel values. The loss function will convert these scores to normalized
# probabilities and return the average -log probability of the correct answers. Minimizing
# this loss should maximize the scores assigned to correct answers by the model. We will
# make use of the loss gradient with respect to each parameter, which tells us the direction
# of the greatest loss increase. We will improve the model by moving the parameters in the
# opposite direction (using a GPU if available). We will visualize the model weights and
# performance over time. The final accuracy of about 92% is close to the limit of what we
# can achieve with this type of model. To improve further we must look beyond linear models.
# ------------------------------------------------------------------------------------------

# Set display width, load packages, import symbols
ENV["COLUMNS"]=72
using Pkg; for p in ("Knet","AutoGrad","Plots","Images","ImageMagick"); haskey(Pkg.installed(),p) || Pkg.add(p); end
using Statistics: mean
using Base.Iterators: flatten
import Random # seed!
using Knet: Knet, AutoGrad, dir, Data, Param, @diff, value, params, grad, progress, progress!, gpu, KnetArray, load, save
# The following are defined for instruction even though they are provided in Knet
# using Knet: accuracy, zeroone, nll, sgd

# Load data (mnistdata basically replicates mnist.ipynb)
include(Knet.dir("data","mnist.jl"))
dtrn,dtst = mnistdata(xsize=(784,:),xtype=Array)
println.(summary.((dtrn,dtst)));

# ------------------------------------------------------------------------------------------
# ## Model definition
# ------------------------------------------------------------------------------------------

# In Julia we define a new datatype using `struct`:
struct Linear; w; b; end

# The new struct comes with a default constructor:
model = Linear(0.01 * randn(10,784), zeros(10))

# We can define other constructors with different inputs:
Linear(i::Int,o::Int,scale=0.01) = Linear(scale * randn(o,i), zeros(o))

# This one allows instances to be defined using input and output sizes:
model = Linear(784,10)

# ------------------------------------------------------------------------------------------
# ## Prediction
# ------------------------------------------------------------------------------------------

# We turn Linear instances into callable objects for prediction:
(m::Linear)(x) = m.w * x .+ m.b

x,y = first(dtst) # The first minibatch from the test set
summary.((x,y))

Int.(y)' # correct answers are given as an array of integers (remember we use 10 for 0)

ypred = model(x)  # Predictions on the first minibatch: a 10x100 score matrix

# We can calculate the accuracy of our model for the first minibatch
accuracy(model,x,y) = mean(y' .== map(i->i[1], findmax(Array(model(x)),dims=1)[2]))
accuracy(model,x,y)

# We can calculate the accuracy of our model for the whole test set
accuracy(model,data) = mean(accuracy(model,x,y) for (x,y) in data)
accuracy(model,dtst)

# ZeroOne loss (or error) is defined as 1 - accuracy
zeroone(x...) = 1 - accuracy(x...)
zeroone(model,dtst)

# ------------------------------------------------------------------------------------------
# ## Loss function
# ------------------------------------------------------------------------------------------

# For classification we use negative log likelihood loss (aka cross entropy, softmax loss, NLL)
# This is the average -log probability assigned to correct answers by the model
function nll(scores, y)
    expscores = exp.(scores)
    probabilities = expscores ./ sum(expscores, dims=1)
    answerprobs = (probabilities[y[i],i] for i in 1:length(y))
    mean(-log.(answerprobs))
end

# model(x) gives predictions, let model(x,y) give the loss
(m::Linear)(x, y) = nll(m(x), y)
model(x,y)

# We can also use the Knet nll implementation for efficiency
(m::Linear)(x, y) = Knet.nll(m(x), y)
model(x,y)

# If the input is a dataset compute average loss:
(m::Linear)(data::Data) = mean(m(x,y) for (x,y) in data)

# Here is per-instance average negative log likelihood for the whole test set
model(dtst)

# ------------------------------------------------------------------------------------------
# **Bonus question:** What is special about the loss value 2.3?
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Calculating the gradient using AutoGrad
# ------------------------------------------------------------------------------------------

@doc AutoGrad

# Redefine the constructor to use Param's so we can compute gradients
Linear(i::Int,o::Int,scale=0.01) =
    Linear(Param(scale * randn(o,i)), Param(zeros(o)))

# Set random seed for replicability
Random.seed!(9);

# Use a larger scale to get a large initial loss
model = Linear(784,10,1.0)

# We can still do predictions and calculate loss:
model(x,y)

# And we can do the same loss calculation also computing gradients:
J = @diff model(x,y)

# To get the actual loss value from J:
value(J)

# params(J) returns an iterator of Params J depends on (i.e. model.b, model.w):
params(J) |> collect

# To get the gradient of a parameter from J:
∇w = grad(J,model.w)

# Note that each gradient has the same size and shape as the corresponding parameter:
@show ∇b = grad(J,model.b);

# ------------------------------------------------------------------------------------------
# ## Checking the gradient using numerical approximation
#
# What does ∇b represent?
#
# ∇b[10] = 0.79 means if I increase b[10] by ϵ, loss will increase by 0.79ϵ
# ------------------------------------------------------------------------------------------

# Loss for the first minibatch with the original parameters
@show value(model.b)
model(x,y)

# To numerically check the gradient let's increase the last entry of b by +0.1.
model.b[10] = 0.1

# We see that the loss moves by ≈ +0.79*0.1 as expected.
@show value(model.b)
model(x,y)

# Reset the change.
model.b[10] = 0

# ------------------------------------------------------------------------------------------
# ## Checking the gradient using manual implementation
# ------------------------------------------------------------------------------------------

# Without AutoGrad we would have to define the gradients manually:
function nllgrad(model,x,y)
    scores = model(x)
    expscores = exp.(scores)
    probabilities = expscores ./ sum(expscores, dims=1)
    for i in 1:length(y); probabilities[y[i],i] -= 1; end
    dJds = probabilities / length(y)
    dJdw = dJds * x'
    dJdb = vec(sum(dJds,dims=2))
    dJdw,dJdb
end;

∇w2,∇b2 = nllgrad(model,x,y)

∇w2 ≈ ∇w

∇b2 ≈ ∇b

# ------------------------------------------------------------------------------------------
# ## Training with Stochastic Gradient Descent (SGD)
# ------------------------------------------------------------------------------------------

# Here is a single SGD update:
function sgdupdate!(func, args; lr=0.1)
    fval = @diff func(args...)
    for param in params(fval)
        ∇param = grad(fval, param)
        param .-= lr * ∇param
    end
    return value(fval)
end

# We define SGD for a dataset as an iterator so that:
# 1. We can monitor and report the training loss
# 2. We can take snapshots of the model during training
# 3. We can pause/terminate training when necessary
sgd(func, data; lr=0.1) =
    (sgdupdate!(func, args; lr=lr) for args in data)

# Let's train a model for 10 epochs to compare training speed on cpu vs gpu.
# progress!(itr) displays a progress bar when wrapped around an iterator like this:
# 2.94e-01  100.00%┣████████████████████┫ 6000/6000 [00:10/00:10, 592.96/s] 2.31->0.28
model = Linear(784,10)
@show model(dtst)
progress!(sgd(model, repeat(dtrn,10)))
@show model(dtst);

# ------------------------------------------------------------------------------------------
# ## Using the GPU
# ------------------------------------------------------------------------------------------

# The training would go a lot faster on a GPU:
# 2.94e-01  100.00%┣███████████████████┫ 6000/6000 [00:02/00:02, 2653.45/s]  2.31->0.28
# To work on a GPU, all we have to do is convert Arrays to KnetArrays:
if gpu() >= 0  # gpu() returns a device id >= 0 if there is a GPU, -1 otherwise
    atype = KnetArray{Float32}  # KnetArrays are stored and operated in the GPU
    dtrn,dtst = mnistdata(xsize=(784,:),xtype=atype)
    Linear(i::Int,o::Int,scale=0.01) =
        Linear(Param(atype(scale * randn(o,i))),
               Param(atype(zeros(o))))

    model = Linear(784,10)
    @show model(dtst)
    progress!(sgd(model,repeat(dtrn,10)))
    @show model(dtst)
end;

# ------------------------------------------------------------------------------------------
#
# ## Recording progress
# ------------------------------------------------------------------------------------------

function trainresults(file, model)
    if (print("Train from scratch? (~77s) "); readline()[1]=='y')
        # This will take every nth element of an iterator:
        takeevery(n,itr) = (x for (i,x) in enumerate(itr) if i % n == 1)
        # We will use it to snapshot model and results every epoch (i.e. 600 iterations)
        lin = ((deepcopy(model),model(dtrn),model(dtst),zeroone(model,dtrn),zeroone(model,dtst))
        # (progress displays a bar like progress! but returns an iterator, progress! returns nothing)
               for x in takeevery(length(dtrn), progress(sgd(model,repeat(dtrn,100)))))
        # Save it as a 5x100 array
        lin = reshape(collect(flatten(lin)),(5,:))
        # Knet.save and Knet.load can be used to store models in files
        Knet.save(file,"results",lin)
    else
        isfile(file) || download("http://people.csail.mit.edu/deniz/models/tutorial/$file", file)
        lin = Knet.load(file,"results")
    end
    return lin
end

# 2.43e-01  100.00%┣████████████████▉┫ 60000/60000 [00:44/00:44, 1349.13/s]
lin = trainresults("lin113.jld2",Linear(784,10));

# ------------------------------------------------------------------------------------------
# ## Linear model shows underfitting
# ------------------------------------------------------------------------------------------

using Plots; default(fmt = :png)

# Demonstrates underfitting: training loss not close to 0
# Also slight overfitting: test loss higher than train
plot([lin[2,:], lin[3,:]],ylim=(.0,.4),labels=[:trnloss :tstloss],xlabel="Epochs",ylabel="Loss")

# this is the error plot, we get to about 7.5% test error, i.e. 92.5% accuracy
plot([lin[4,:], lin[5,:]],ylim=(.0,.12),labels=[:trnerr :tsterr],xlabel="Epochs",ylabel="Error")

# ------------------------------------------------------------------------------------------
# ## Visualizing the learned weights
# ------------------------------------------------------------------------------------------

# Let us visualize the evolution of the weight matrix as images below
# Each row is turned into a 28x28 image with positive weights light and negative weights dark gray
using Images, ImageMagick
for t in 10 .^ range(0,stop=log10(size(lin,2)),length=20) #logspace(0,2,20)
    i = ceil(Int,t)
    f = lin[1,i]
    w1 = reshape(Array(value(f.w))', (28,28,1,10))
    w2 = clamp.(w1.+0.5,0,1)
    IJulia.clear_output(true)
    display(hcat([mnistview(w2,i) for i=1:10]...))
    display("Epoch $(i-1)")
    sleep(1) # (0.96^i)
end


using PtFEM

ProjDir = @__DIR__

l = 1.0       # Total length [m]
N = 5         # Number of nodes
els = N - 1   # Number of finite elements
nod = 2       # Number of nodes per finite elements
nodof = 1     # Degrees of freedom for each node
np_types = 1  # Number of proerty types
EA = 1.0e5    # Strain stiffness
nip = 1       # Number of integration points

struct_el = :Rod
fin_el = :Line

data = Dict(
  # Rod(nels, np_types, nip, finite_element(nod, nodof))
  :struc_el => Rod(4, 1, 1, Line(2, 1)),
  :properties => [2.0e3; 1.0e3],
  :etype => [2, 2, 1, 1],
  :x_coords => collect(range(0, stop=1, length=5)),
  :support => [
    (1, [0])
    ],
  :fixed_freedoms => [
    (5, 1, 0.05)
    ]
)

data |> display
println()

@time fem, dis_df, fm_df = p41(data)
println()

display(fem.actions)
println()
display(fem.displacements)
println()

display(fm_df)
println()
display(dis_df)

using Plots
gr(size=(400,500))

x = 0.0:l/els:l
u = convert(Array, dis_df[:, :x_translation])

p = Vector{Plots.Plot{Plots.GRBackend}}(undef, 2)
titles = ["PtFEM Ex42.1 u(x)", "PtFEM Ex42.1 N(x)"]

p[1]=plot(ylim=(0.0, 1.0), xlim=(0.0, 70.0),
  yflip=true, xflip=false, xlab="Normal force [N]",
  ylab="x [m]", title=titles[2]
)
vals = fm_df[:, :normal_force_2]
for i in 1:els
    plot!(p[1],
      [vals[i], vals[i]],
      [(i-1)*l/els, i*l/els], color=:blue,
      fill=true, fillalpha=0.1, leg=false
    )
    delta = abs(((i-1)*l/els) - (i*l/els)) / 20.0
    y1 = collect(((i-1)*l/els):delta:(i*l/els))
    for j in 1:length(y1)
      plot!(p[1],
        [vals[i], 0.0],
        [y1[j], y1[j]], color=:blue, alpha=0.5
      )
    end
end

p[2] = plot(u, x, xlim=(0.0, 0.05), yflip=true,
  xlab="Displacement [m]", ylab="x [m]",
  fill=true, fillalpha=0.1, leg=false, title=titles[1])

plot(p..., layout=(1, 2))
savefig(ProjDir*"/Ex41.2.png")


#this includes methods for aproximation for saturation pressures and
#temperatures
"""
    Antoine(a,b,c,from,to,base)
    Antoine(a,b,c;from =u"K",to = u"Pa",base=ℯ)

## Antoine Vapor Pressure

Calculates vapor pressure of a chemical using the Antoine equation. Parameters `a`, `b`, and `c` are chemical-dependent. Parameters can be found in numerous sources; however units of the coefficients used vary. Originally proposed by Antoine (1888)

The vapor pressure is given by:

    log(base,p) = a + b/(c+t)

by default, `t` in in kelvin, `p` in Pa and the base is natural (ℯ)

the coefficients are normally defined for a specific combination of pressure and temperature units, you can specify those using the `from` (temperature) and `to` (pressure) arguments, passing a corresponding Unitful unit.

## References

1. Poling, Bruce E. The Properties of Gases and Liquids. 5th edition. New York: McGraw-Hill Professional, 2000.
2. Antoine, C. 1888. Tensions des Vapeurs: Nouvelle Relation Entre les Tensions et les Tempé. Compt.Rend. 107:681-684.
3. Yaws, Carl L. The Yaws Handbook of Vapor Pressure: Antoine Coefficients. 1 edition. Houston, Tex: Gulf Publishing Company, 2007.
"""
struct Antoine{V,B,F,T} <: SaturationModel
    a::V
    b::V
    c::V
    base::B
    from::F
    to::T
end

Antoine(a,b,c;from =u"K",to = u"Pa",base=ℯ) = Antoine(a,b,c,base,from,to)

function pressure_impl(mt::SingleSatT,model::T,t) where T<: Antoine
    _t = convert_unit(u"K",model.from,t)
    _a = model.a
    _b = model.b
    _c = model.c
    _base = model.base
    p = _base^(_a+_b/(_c+_t))
    return convert_unit(model.to,u"Pa",p)
end

function temperature_impl(mt::SingleSatP,model::T,p) where T<:Antoine
    _p = convert_unit(u"Pa",model.to,p)
    _a = model.a
    _b = model.b
    _c = model.c
    _base = model.base
    _t = _b/(log(_base,_p)-_a) -_c
    return convert_unit(model.from,u"K",_t)
end
"""
    LeeKesler(;tc,pc,ω)
    LeeKesler(tc,pc,ω)

## Lee Kesler Vapor Pressure

Calculates vapor pressure of a fluid at arbitrary temperatures using a
corresponding states relationship by [1]; requires a chemical's critical temperature and
acentric factor.

The vapor pressure is given by:

    ln(Pᵣ) = f⁰ + ωf¹
    f⁰ = 5.92714 - 6.09648/Tᵣ - 1.28862*ln(Tᵣ) + 0.169347*Tᵣ^6
    f¹ = 15.2518 - 15.6875/Tᵣ - 13.4721*ln(Tᵣ) + 0.43577*Tᵣ^6

## References
1. Lee, Byung Ik, and Michael G. Kesler. "A Generalized Thermodynamic Correlation Based on Three-Parameter Corresponding States." AIChE Journal 21, no. 3 (1975): 510-527. doi:10.1002/aic.690210313.
2. Reid, Robert C..; Prausnitz, John M.;; Poling, Bruce E. The Properties of Gases and Liquids. McGraw-Hill Companies, 1987.

"""
struct LeeKesler{T} <: SaturationModel
    tc::T
    pc::T
    ω::T
end

function LeeKesler(;tc,pc,ω)
    _tc = ThermoState.normalize_units(tc)
    _pc = ThermoState.normalize_units(pc)
    return LeeKesler(_tc,_pc,ω)
end

function LeeKesler(model::ThermoModel)
    _tc = temperature(model,CriticalPoint(),u"K")
    _pc = pressure(model,CriticalPoint(),u"Pa")
    _ω = acentric_factor(model)
    return LeeKesler(_tc,_pc,_ω)
end

function temperature(model::LeeKesler,st::CriticalPoint,unit=u"K")
    return convert_unit(u"K",unit,model.tc)
end

function pressure(model::LeeKesler,st::CriticalPoint,unit=u"Pa")
    return convert_unit(u"Pa",unit,model.pc)
end

function acentric_factor(model::LeeKesler)
    return model.ω
end

function pressure_impl(mt::SingleSatT,model::LeeKesler,t)
    Pc = model.pc
    Tᵣ = t/model.tc
    logTᵣ = log(Tᵣ)
    ω = model.ω
    Tᵣ6 = Tᵣ*Tᵣ
    Tᵣ6 *= Tᵣ6*Tᵣ6
    f⁰ = 5.92714 - 6.09648/Tᵣ - 1.28862*logTᵣ+ 0.169347*Tᵣ6
    f¹ = 15.2518 - 15.6875/Tᵣ - 13.4721*logTᵣ + 0.43577*Tᵣ6
    return Pc*exp(muladd(f¹,ω,f⁰))
end

initial_temperature(model::LeeKesler,p) = critical_sat_interpolation(model,p)

"""
    Wagner(;tc,pc,a,b,c,d)
    Wagner(tc,pc,a,b,c,d)

# Wagner vapor pressure

Calculates vapor pressure using the Wagner equation (2.5, 5 form).

Requires critical temperature and pressure as well as four coefficients specific to each chemical.

The vapor pressure is given by:

    ln(Pᵣ)=  (aτ  + bτ^1.5 + cτ^2.5+ dτ^5)/Tᵣ

    τ = 1 - Tᵣ

## References

1. Wagner, W. “New Vapour Pressure Measurements for Argon and Nitrogen and a New Method for Establishing Rational Vapour Pressure Equations.” Cryogenics 13, no. 8 (August 1973): 470-82. doi:10.1016/0011-2275(73)90003-9
2. Poling, Bruce E. The Properties of Gases and Liquids. 5th edition. New York: McGraw-Hill Professional, 2000.

"""
struct Wagner{T} <: SaturationModel
    tc::T
    pc::T
    a::T
    b::T
    c::T
    d::T
end

function Wagner(;tc,pc,a,b,c,d)
    _tc = normalize_units(tc)
    _pc = normalize_units(pc)
    return Wagner(_tc,_pc,a,b,c,d)
end

initial_temperature(model::Wagner,p) = critical_sat_interpolation(model,p)

function temperature(model::Wagner,st::CriticalPoint,unit=u"K")
    return convert_unit(u"K",unit,model.tc)
end

function pressure(model::Wagner,st::CriticalPoint,unit=u"Pa")
    return convert_unit(u"Pa",unit,model.pc)
end

function pressure_impl(mt::SingleSatT,model::Wagner,t)
    Tc = model.tc
    Pc = model.pc
    Tr = t/Tc
    τ = one(Tr) - Tr
    τhalf = sqrt(tau)
    bτ = τ*τhalf
    cτ = bτ*τ
    dτ = cτ*cτ
    return Pc*exp((a*τ + b*bτ+ c*cτ + d*dτ)/Tr)
end

"""
WagnerOriginal(;tc,pc,a,b,c,d)
WagnerOriginal(tc,pc,a,b,c,d)

# Wagner (original) vapor pressure

Calculates vapor pressure using the Wagner equation (3, 6 form).

Requires critical temperature and pressure as well as four coefficients specific to each chemical.

The vapor pressure is given by:

    ln(Pᵣ)=  (aτ  + bτ^1.5 + cτ^3 dτ^6)/Tᵣ

    τ = 1 - Tᵣ

## References

1. Poling, Bruce E. The Properties of Gases and Liquids. 5th edition. New York: McGraw-Hill Professional, 2000.
2. McGarry, Jack. “Correlation and Prediction of the Vapor Pressures of Pure Liquids over Large Pressure Ranges.” Industrial & Engineering Chemistry Process Design and Development 22, no. 2 (April 1, 1983): 313-22. doi:10.1021/i200021a023.
"""
struct WagnerOriginal{T} <: SaturationModel
    tc::T
    pc::T
    a::T
    b::T
    c::T
    d::T
end

function WagnerOriginal(;tc,pc,a,b,c,d)
    _tc = normalize_units(tc)
    _pc = normalize_units(pc)
    return WagnerOriginal(_tc,_pc,a,b,c,d)
end

function temperature(model::WagnerOriginal,st::CriticalPoint,unit=u"K")
    return convert_unit(u"K",unit,model.tc)
end

function pressure(model::WagnerOriginal,st::CriticalPoint,unit=u"Pa")
    return convert_unit(u"Pa",unit,model.pc)
end

function pressure_impl(mt::SingleSatT,model::WagnerOriginal,t)
    Tc = model.tc
    Pc = model.pc
    Tr = t/Tc
    τ = one(Tr) - Tr
    τhalf = sqrt(tau)
    bτ = τ*τhalf
    cτ = bτ*bτ
    dτ = cτ*cτ
    return Pc*exp((a*τ + b*bτ+ c*cτ + d*dτ)/Tr)
end

initial_temperature(model::WagnerOriginal,p) = critical_sat_interpolation(model,p)

"""
    TRCAntoine(tc,t0,A,B,C,n,E,F)
    TRCAntoine(;tc,t0,B,A,C,n,E,F)

# TRC extended Antoine model for vapor pressure

Calculates vapor pressure of a chemical using the TRC Extended Antoine equation. Parameters are chemical dependent, and said to be from the Thermodynamics Research Center (TRC) at Texas A&M. Coefficients for various chemicals can be found in [1].

The vapor pressure is given by:

    log10(Psat) = A - B/(t + C) + 0.43429*x^n + Ex^8 + Fx^12

    x = max((t-t₀-273.15)/tc, 0)

## References

1. Poling, Bruce E. The Properties of Gases and Liquids. 5th edition. New York: McGraw-Hill Professional, 2000.
"""

struct TRCAntoine{T} <: SaturationModel
    tc::T
    t0::T
    A::T
    B::T
    C::T
    n::T
    E::T
    F::T
end

TRCAntoine(;tc,t0,B,A,C,n,E,F) = TRCAntoine(tc,t0,A,B,C,n,E,F)

function temperature(model::TRCAntoine,st::CriticalPoint,unit=u"K")
    return convert_unit(u"K",unit,model.tc)
end

function pressure_impl(mt::SingleSatT,model::TRCAntoine,t)
    _x = (t - to - 273.15)/Tc
    x = max(_x,zero(_x))
    A = model.A
    B = model.B
    C = model.C
    n = model.n
    E = model.E
    F = model.F
    x4 = x^4
    return exp10(A - B/(t+C) + 0.43429*x^n + x4*x4*(E + F*x4))
end

function initial_temperature(model::TRCAntoine,p)
    A = model.A
    B = model.B
    C = model.C
    return B/(log10(p)-A) -C
end

"""
    AmbroseWalton(;tc,pc,ω)
    AmbroseWalton(tc,pc,ω)

# Ambrose-Walton relation for vapor pressure

Calculates vapor pressure of a fluid at arbitrary temperatures using a corresponding states relationship by [1]; requires a chemical’s critical temperature and acentric factor.

The vapor pressure is given by:

    ln(Pᵣ) = f⁰ + f¹ω + f²ω^2
    f⁰ = -5.97616τ+ 1.29874τ^1.5 - 0.60394τ^2.5 - 1.06841τ^5
    f¹ = -5.03365τ + 1.11505τ^1.5 - 5.41217τ^2.5 - 7.46628τ^5
    f² = -0.64771τ + 2.41539τ^1.5 - 4.26979τ^2.5 + 3.25259τ^5
    τ = 1 - Tᵣ

## References

1. Ambrose, D., and J. Walton. “Vapour Pressures up to Their Critical Temperatures of Normal Alkanes and 1-Alkanols.” Pure and Applied Chemistry 61, no. 8 (1989): 1395-1403. doi:10.1351/pac198961081395.
2. Poling, Bruce E. The Properties of Gases and Liquids. 5th edition. New York: McGraw-Hill Professional, 2000.
"""

struct AmbroseWalton{T} <: SaturationModel
    tc::T
    pc::T
    ω::T
end

function AmbroseWalton(;tc,pc,ω)
    _tc = ThermoState.normalize_units(tc)
    _pc = ThermoState.normalize_units(pc)
    return AmbroseWalton(_tc,_pc,ω)
end

function AmbroseWalton(model::ThermoModel)
    _tc = temperature(model,CriticalPoint(),u"K")
    _pc = pressure(model,CriticalPoint(),u"Pa")
    _ω = acentric_factor(model)
    return AmbroseWalton(_tc,_pc,_ω)
end

function temperature(model::AmbroseWalton,st::CriticalPoint,unit=u"K")
    return convert_unit(u"K",unit,model.tc)
end

function pressure(model::AmbroseWalton,st::CriticalPoint,unit=u"Pa")
    return convert_unit(u"Pa",unit,model.pc)
end

function acentric_factor(model::AmbroseWalton)
    return model.ω
end
function pressure_impl(mt::SingleSatT,model::AmbroseWalton,_t)
    t = normalize_units(_t)
    Pc = model.pc
    τ = 1 - t/model.tc
    ω = model.ω
    τ15 = sqrt(τ)*τ
    τ25 = τ15*τ
    τ5 = τ25*τ25
    f⁰ = -5.97616τ+ 1.29874τ15 - 0.60394τ25 - 1.06841τ5
    f¹ = -5.03365τ + 1.11505τ15 - 5.41217τ25 - 7.46628τ5
    f² = -0.64771τ + 2.41539τ15 - 4.26979τ25 + 3.25259τ5

    return Pc*exp(muladd(muladd(f²,ω,f¹),ω,f⁰))
end

initial_temperature(model::AmbroseWalton,p) = critical_sat_interpolation(model,p)

"""
    Sanjari(;tc,pc,ω)
    Sanjari(tc,pc,ω)

# Sanjari relation for vapor pressure

Calculates vapor pressure of a fluid at arbitrary temperatures using a corresponding states relationship by [1]; requires a chemical’s critical temperature and acentric factor. Although developed for refrigerants, this model should have some general predictive ability.

The vapor pressure is given by:

    ln(Pᵣ) = f⁰ + f¹ω + f²ω^2
    f⁰ = 6.83377 + -5.76051/Tᵣ + 0.90654*log(Tᵣ) + -1.16906*Tᵣ^1.9
    f¹ = 5.32034 + -28.1460/Tᵣ + -58.0352*log(Tᵣ) + 23.57466*Tᵣ^1.9
    f² = 18.19967 + 16.33839/Tᵣ + 65.6995*log(Tᵣ) + -35.9739*Tᵣ^1.9

## References

(1, 2) Sanjari, Ehsan, Mehrdad Honarmand, Hamidreza Badihi, and Ali Ghaheri. “An Accurate Generalized Model for Predict Vapor Pressure of Refrigerants.” International Journal of Refrigeration 36, no. 4 (June 2013): 1327-32. doi:10.1016/j.ijrefrig.2013.01.007.
"""

struct Sanjari{T} <: SaturationModel
    tc::T
    pc::T
    ω::T
end

function Sanjari(;tc,pc,ω)
    _tc = ThermoState.normalize_units(tc)
    _pc = ThermoState.normalize_units(pc)
    return Sanjari(_tc,_pc,ω)
end

function Sanjari(model::ThermoModel)
    _tc = temperature(model,CriticalPoint(),u"K")
    _pc = pressure(model,CriticalPoint(),u"Pa")
    _ω = acentric_factor(model)
    return Sanjari(_tc,_pc,_ω)
end

function temperature(model::Sanjari,st::CriticalPoint,unit=u"K")
    return convert_unit(u"K",unit,model.tc)
end

function pressure(model::Sanjari,st::CriticalPoint,unit=u"Pa")
    return convert_unit(u"Pa",unit,model.pc)
end

function acentric_factor(model::Sanjari)
    return model.ω
end

function pressure_impl(mt::SingleSatT,model::Sanjari,_t)
    t = normalize_units(_t)
    ω = model.ω
    Tr = t/model.tc
    Tr_inv = 1.0/Tr
    log_Tr = log(Tr)
    Tr_19 = Tr^1.9
    f⁰ = 6.83377 + -5.76051*Tr_inv + 0.90654*log_Tr + -1.16906*Tr_19
    f¹ = 5.32034 + -28.1460*Tr_inv + -58.0352*log_Tr + 23.57466*Tr_19
    f² = 18.19967 + 16.33839*Tr_inv + 65.6995*log_Tr + -35.9739*Tr_19

    return model.pc*exp(muladd(muladd(f²,ω,f¹),ω,f⁰))
end

initial_temperature(model::Sanjari,p) = critical_sat_interpolation(model,p)

"""
    Edalat(;tc,pc,ω)
    Edalat(tc,pc,ω)

# Edalat relation for vapor pressure

Calculates vapor pressure of a fluid at arbitrary temperatures using a CSP relationship by [1]. Requires a chemical’s critical temperature, pressure, and acentric factor. Claimed to have a higher accuracy than the Lee-Kesler relationship.

The vapor pressure is given by:

    ln(Pᵣ) = f⁰ + f¹ω + f²ω^2
    f⁰ = 6.83377 + -5.76051/Tᵣ + 0.90654*log(Tᵣ) + -1.16906*Tᵣ^1.9
    f¹ = 5.32034 + -28.1460/Tᵣ + -58.0352*log(Tᵣ) + 23.57466*Tᵣ^1.9
    f² = 18.19967 + 16.33839/Tᵣ + 65.6995*log(Tᵣ) + -35.9739*Tᵣ^1.9

## References

(1, 2) Edalat, M., R. B. Bozar-Jomehri, and G. A. Mansoori. “Generalized Equation Predicts Vapor Pressure of Hydrocarbons.” Oil and Gas Journal; 91:5 (February 1, 1993).
"""

struct Edalat{T} <: SaturationModel
    tc::T
    pc::T
    ω::T
end

function Edalat(;tc,pc,ω)
    _tc = ThermoState.normalize_units(tc)
    _pc = ThermoState.normalize_units(pc)
    return Edalat(_tc,_pc,ω)
end

function Edalat(model::ThermoModel)
    _tc = temperature(model,CriticalPoint(),u"K")
    _pc = pressure(model,CriticalPoint(),u"Pa")
    _ω = acentric_factor(model)
    return Edalat(_tc,_pc,_ω)
end

function temperature(model::Edalat,st::CriticalPoint,unit=u"K")
    return convert_unit(u"K",unit,model.tc)
end

function pressure(model::Edalat,st::CriticalPoint,unit=u"Pa")
    return convert_unit(u"Pa",unit,model.pc)
end

function acentric_factor(model::Edalat)
    return model.ω
end

function pressure_impl(mt::SingleSatT,model::Edalat,_t)
    t = normalize_units(_t)
    Pc = model.pc
    τ = 1 - t/model.tc
    ω = model.ω
    τ15 = sqrt(τ)*τ
    τ25 = τ15*τ
    τ5 = τ25*τ25
    tau = 1. - T/Tc
    a = -6.1559 - 4.0855ω
    c = -0.8747 - 7.8874ω
    d = 1/(-0.4893 - 0.9912ω + 3.1551ω^2)
    b = 1.5737 - 1.0540ω - 4.4365E-3*d
    τ15 = sqrt(τ)*τ
    τ3 = τ15*τ15
    lnPr = (a*τ + b*τ15 + c*τ3 + d*τ3*τ3)*model.tc/t
    return exp(lnPr)*Pc
end

initial_temperature(model::Edalat,p) =  critical_sat_interpolation(model,p)

"""
    WaterSat()
    Antoine(a,b,c;from =u"K",to = u"Pa",base=ℯ)

## Water saturation acording to IAPWS-IF97

Calculates vapor pressure of water acording IAPWS Industrial Formulation

The implementation is from https://github.com/braamvandyk/SteamTables.jl
"""
struct WaterSat <: SaturationModel end

const watersat_data = (;n = [0.116_705_214_527_67E4,
    -0.724_213_167_032_06E6,
    -0.170_738_469_400_92E2,
    0.120_208_247_024_70E5,
    -0.323_255_503_223_33E7,
    0.149_151_086_135_30E2,
    -0.482_326_573_615_91E4,
    0.405_113_405_420_57E6,
    -0.238_555_575_678_49,
    0.650_175_348_447_98E3]
)

function pressure_impl(mt::SingleSatT,::WaterSat,t)
    n = watersat_data.n
    Θ = t + n[9]/(t - n[10])
    A =      Θ^2 + n[1]*Θ + n[2]
    B = n[3]*Θ^2 + n[4]*Θ + n[5]
    C = n[6]*Θ^2 + n[7]*Θ + n[8]
    P = (2C / (-B + √(B^2 - 4*A*C)))^4
    return convert_unit(u"MPa",u"Pa",P)
end

function temperature_impl(mt::SingleSatP,::WaterSat,p)
    n = watersat_data.n
    P = convert_unit(u"Pa",u"MPa",p)
    β = P^0.25
    E =      β^2 + n[3]*β + n[6]
    F = n[1]*β^2 + n[4]*β + n[7]
    G = n[2]*β^2 + n[5]*β + n[8]
    D = 2G / (-F - √(F^2 - 4*E*G))
    T = (n[10]+D-√((n[10]+D)^2 - 4(n[9]+n[10]*D)))/2.0
    return T
end


builtin_joltage( adapters ) = maximum(adapters) + 3

function load( filename )
    open(filename) do f
        return parse.(Int, readlines(f))
    end
end

function part1( puzzle )
    adapt = vcat(0, puzzle..., builtin_joltage( puzzle ) )
    sort!(adapt)
    differences = []
    adapter = 0
    for x in adapt
         possibles = 0 .< adapt .- adapter .<= 3
         if any(possibles)
             new_adapter = adapt[findfirst(possibles)]
             push!(differences, new_adapter - adapter )
             adapter = new_adapter
         else
             return count(==(1), differences), count(==(3), differences)
         end
    end
end

function part2( puzzle ) # brute force is too exepensive

    function count_path( current_adapter, adapters)

        possibles = 0 .< adapters .- current_adapter .<= 3

        if any(possibles)
            for adapter in adapters[possibles]
                count_path( adapter, adapters)
            end
        else
            count += 1
        end
    end

    adapters = vcat(0, puzzle...)
    sort!(adapters)
    count = 0
    count_path( 0, adapters)

    return count

end

# I found the solution with the help of this reference
# https://www.geeksforgeeks.org/count-possible-paths-source-destination-exactly-k-edges/

function count_arrangements(puzzle)

    adapt = vcat(0, puzzle..., builtin_joltage(puzzle))

    n = length(adapt)

    graph = zeros(Bool,n, n)
    for i in 1:n, j in 1:n
        graph[i,j] = (0 < adapt[j] - adapt[i] <= 3)
    end

    count =  zeros(Int,n,n,n)

    e = 1
    while e <= n && adapt[e] < builtin_joltage(adapt)

        for i in 1:n, j in 1:n

            if (e == 1 && i == j)
                count[i,j,e] = 1
            end
            if (e == 2 && graph[i,j])
                count[i,j,e] = 1
            end

            if e > 2
                for k in 1:n
                    if graph[i,k]
                        count[i,j,e] += count[k,j,e-1]
                    end
                end
            end
        end
        e = e+1

    end

    return sum(count[1,n,:])

end

puzzle = load("input10.txt")
println(prod(part1(puzzle)))
puzzle = load("input10.txt")
println(count_arrangements(puzzle))


module ShaleDrillingLikelihood_OverallLikelihood_Test

# using Revise
# using Juno
# using Profile
# using ProfileView
# using PProf

const DOBTIME=false

using ShaleDrillingLikelihood
using Test
using StatsFuns
using Calculus
using Optim
using Random
using InteractiveUtils
using LinearAlgebra
using BenchmarkTools

using Base: OneTo

using ShaleDrillingLikelihood: AbstractDataSet,
    DataRoyalty, DataProduce, DataDrill,
    TestDrillModel,
    simloglik!,
    SimulationDraws,
    idx_drill_ρ, idx_royalty_ρ,
    idx_drill_ψ, idx_produce_ψ,
    _nparm,
    thetas

println("testing overall likelihood")

@testset "Testing Joint Likelihoods" begin
    num_i = 1_200
    L = 3
    M = 1_000
    u = randn(num_i)
    v = randn(num_i)
    psi2 = u

    # set up coefs
    θρ = 0.5
    αψ = 0.6

    # set up drilling coefs
    model_drill = TestDrillModel()
    θ_drill   = [αψ, 2.0, -2.0, -0.75, θρ]
    θ_royalty = [θρ, 1.0,    -1.0, 1.0, 1.0,    -0.6, 0.6]  # dψdρ, ψ, β, κ
    θ_produce = vcat(αψ, rand(3), 0.3, 0.4)

    θ = vcat(θ_drill, θ_royalty[2:end], θ_produce[2:end])
    @test idx_produce_ψ isa Function
    @test idx_drill_ψ isa Function
    coef_links = [(idx_produce_ψ, idx_drill_ψ,),]

    # make data
    data_drill_opt = (num_zt=200, minmaxleases=1:2, nper_initial=10:20, nper_development=0:10, tstart=1:50,)
    data_drill   = DataDrill(u, v, TestDrillModel(), θ_drill; data_drill_opt...)
    data_produce = DataProduce(u, 10, 10:20, θ_produce)
    data_royalty = DataRoyalty(u,v,θ_royalty,L)
    data_full = DataSetofSets(data_drill, data_royalty, data_produce, coef_links)
    sim = SimulationDraws(M, data_drill)

    @testset "Joint likelihood of Royalty + Pdxn" begin

        println("testing joint likelihood of royalty + pdxn")
        data = DataSetofSets(EmptyDataSet(), data_royalty, data_produce)
        theta = vcat(θ_royalty, θ_produce)
        grad = similar(theta)
        nparm = _nparm(data)
        hess = Matrix{eltype(theta)}(undef, nparm, nparm)
        tmpgrads = Matrix{eltype(theta)}(undef, nparm, num_i)

        simloglik!(grad, hess, tmpgrads, data, theta, sim, false)
        simloglik!(grad, hess, tmpgrads, data, theta, sim, true)

        fd = Calculus.gradient(xx -> simloglik!(grad, hess, tmpgrads, data, xx, sim, false), theta, :central)

        fill!(grad,0)
        fill!(hess,0)
        simloglik!(grad, hess, tmpgrads, data, theta, sim, true)
        @test !all(grad.==0)
        @test isapprox(fd, grad; rtol=2e-5)

        if DOBTIME
            print("")
            @show @benchmark simloglik!($grad, $hess, $tmpgrads, $data, $theta, $sim, false)
            @show @benchmark simloglik!($grad, $hess, $tmpgrads, $data, $theta, $sim, true)
            print("")
        end

        # @code_warntype simloglik!(grad, hess, tmpgrads, data, theta, sim, false)
        # println("\n\n\n----------------------------\n\n\n")
        # @code_warntype simloglik!(grad, hess, tmpgrads, data, theta, sim, true)

        # Profile.clear()
        # @profile simloglik!(grad, hess, tmpgrads, data, theta, sim, false)
        # # Juno.profiletree()
        # # Juno.profiler()
        # Profile.print(format=:flat)
        # # ProfileView.view()
        # pprof()

        # Profile.clear()
        # @profile simloglik!(grad, hess, tmpgrads, data, theta, sim, true)
        # ProfileView.view()
    end

    @testset "Joint likelihood of Drill + Roy + Pdxn" begin

        println("testing joint likelihood of FULL data!")
        data = data_full
        theta = θ
        grad = similar(theta)
        nparm = _nparm(data)
        hess = zeros(nparm, nparm)
        tmpgrads = zeros(nparm, num_i)

        println("simloglik, no grad")
        simloglik!(grad, hess, tmpgrads, data, theta, sim, false)
        println("simloglik, with grad")
        simloglik!(grad, hess, tmpgrads, data, theta, sim, true)

        println("simloglik finite diff")
        fd = Calculus.gradient(xx -> simloglik!(grad, hess, tmpgrads, data, xx, sim, false), theta, :central)

        fill!(grad,0)
        fill!(hess,0)
        println("simloglik gradient")
        simloglik!(grad, hess, tmpgrads, data, theta, sim, true)
        @test !all(grad.==0)
        @test isapprox(fd, grad; rtol=2e-5)
        println("done")

        if DOBTIME
            print("")
            @show @benchmark simloglik!($grad, $hess, $tmpgrads, $data, $theta, $sim, false)
            @show @benchmark simloglik!($grad, $hess, $tmpgrads, $data, $theta, $sim, true)
            print("")
        end
        # @code_warntype simloglik!(grad, hess, tmpgrads, data, theta, sim, false)
        # println("\n\n\n----------------------------\n\n\n")
        # @code_warntype simloglik!(grad, hess, tmpgrads, data, theta, sim, true)

        # Profile.clear()
        # @profile simloglik!(grad, hess, tmpgrads, data, theta, sim, false)
        # # Juno.profiletree()
        # # Juno.profiler()
        # Profile.print(format=:flat)
        # # ProfileView.view()
        # pprof()

        # Profile.clear()
        # @profile simloglik!(grad, hess, tmpgrads, data, theta, sim, true)
        # ProfileView.view()
    end

end

end # module


function anisodiff(im,niter,kappa,lambda,option)
    diff=zeros(Float64,size(im))
    anisodiff!(im,niter,kappa,lambda,option,diff)
    diff
end

function anisodiff!(im::Array{Float64,2}, niter::Int64, kappa::Float64, lambda::Float64, option::Int64,dd=zeros(Float64,size(dd)))
#=
% Arguments:
%         im     - input image
%         niter  - number of iterations.
%         kappa  - conduction coefficient 20-100 ?
%         lambda - max value of .25 for stability
%         option - 1 Perona Malik diffusion equation No 1
%                  2 Perona Malik diffusion equation No 2
%
% Returns:
%         diff   - diffused image.
%
% kappa controls conduction as a function of gradient.  If kappa is low
% small intensity gradients are able to block conduction and hence diffusion
% across step edges.  A large value reduces the influence of intensity
% gradients on conduction.
%
% lambda controls speed of diffusion (you usually want it at a maximum of
% 0.25)
%
% Diffusion equation 1 favours high contrast edges over low contrast ones.
% Diffusion equation 2 favours wide regions over smaller ones.

% Reference:
% P. Perona and J. Malik.
% Scale-space and edge detection using ansotropic diffusion.
% IEEE Transactions on Pattern Analysis and Machine Intelligence,
% 12(7):629-639, July 1990.
%
% Peter Kovesi
% www.peterkovesi.com/matlabfns/
%
% June 2000  original version.
% March 2002 corrected diffusion eqn No 2.
=#

    (rows,cols) = size(dd);
    #rows=480
    #cols=640

    for i = 1:niter
        for j=1:length(im)
            dd[j]=im[j]
        end

        for k=1:cols
            for j=1:rows

                if j==1
                    deltaN = -1 * dd[j,k]
                else
                    @inbounds deltaN = dd[j-1,k] - dd[j,k]
                end

                if j==rows
                    deltaS = -1 * dd[j,k]
                else
                    @inbounds deltaS = dd[j+1,k] - dd[j,k]
                end

                if k==1
                    deltaW = -1 * dd[j,k]
                else
                    @inbounds deltaW = dd[j,k-1] - dd[j,k];
                end

                if k==cols
                    deltaE = -1 * dd[j,k]
                else
                    @inbounds deltaE = dd[j,k+1] - dd[j,k];
                end

                #if option == 1
                    cN = exp(-(deltaN/kappa)*(deltaN/kappa))
                    cS = exp(-(deltaS/kappa)*(deltaS/kappa))
                    cE = exp(-(deltaE/kappa)*(deltaE/kappa))
                    cW = exp(-(deltaW/kappa)*(deltaW/kappa))
                #else
                    #cN = 1/(1 + (deltaN/kappa)*(deltaN/kappa))
                    #cS = 1/(1 + (deltaS/kappa)*(deltaS/kappa))
                    #cE = 1/(1 + (deltaE/kappa)*(deltaE/kappa))
                    #cW = 1/(1 + (deltaW/kappa)*(deltaW/kappa))
                #end

                @inbounds im[j,k] = dd[j,k] + lambda*(cN*deltaN + cS*deltaS + cE*deltaE + cW*deltaW);
            end
        end
    end

    nothing
end

#The clahe method from Images is not optomized for iterations.
#Should make a new method that preallocates a temp array to be reused.
function local_contrast_enhance!(img,out_img)

    for i=1:length(img)
        out_img[i] = img[i]/255
    end

    img2 = Images.clahe(out_img,256,xblocks=30,yblocks=30,clip=15)

    for i=1:length(img2)
        out_img[i] = img2[i]*255
    end

    nothing
end

function local_contrast_enhance(img)

    img2 = img ./ 255

    img2 = Images.clahe(img2,256,xblocks=30,yblocks=30,clip=15) .* 255

    img2
end

function subtract_background(han::Tracker_Handles)

    mydiff = han.current_frame2 .- mean_image(han)
    new_diff = (mydiff - minimum(mydiff))
    new_diff = new_diff ./ maximum(new_diff)

    han.current_frame = round.(UInt8,new_diff .* 255)

    nothing
end

# https://web.stanford.edu/class/cs448f/lectures/2.1/Sharpening.pdf
#=
Filter type, 1 = median, 2 = gaussian
=#
function sharpen_image(img::AbstractArray{T,2},sigma=21,rep = 5,filter_type=1) where T
    if iseven(sigma)
        sigma = sigma + 1
    end
    kern = KernelFactors.gaussian((sigma,sigma))
    img_out = zeros(T,size(img))
    img_temp = zeros(Float64,size(img))
    img_median = zeros(UInt8,size(img))
    if filter_type == 1
        for i=1:length(img)
            img_median[i] = round(UInt8,img[i])
        end
    end

    for i=1:length(img)
        img_temp[i] = img[i]
    end

    for i=1:rep
        if filter_type == 1
            img_temp = 1.5 .* img_temp .- 0.5 .* median_filter(img_median,(sigma,sigma))
        else
            img_temp = 1.5 .* img_temp .- 0.5 .* imfilter(img_temp,kern)
        end
        adjust_contrast(img_temp,0,255)
        if filter_type == 1
            for i=1:length(img)
                img_median[i] = round(UInt8,img_temp[i])
            end
        end
    end
    if eltype(img) <: Integer
        for i=1:length(img)
            img_out[i] = round(T,img_temp[i])
        end
    else
        img_out[:] = img_temp
    end
    img_out
end

#=
Median Filter adapted from Julia code by Tejus Gupta (c) 2018
https://gist.github.com/tejus-gupta/2f29e101e18a58dcf2477991151b3a42

Based on algorithm in 'A Fast Two-Dimensional Median Filtering Algorithm' by Huang, Yang and Tang.
=#
function median_filter(img::Array{UInt8, 2}, window_size::Tuple{Int64,Int64})

    function update_median(median_val, n_pixels_below_median, hist, half_pixels)
        if n_pixels_below_median <= half_pixels
            for val in median_val:2^8
                if n_pixels_below_median + hist[val] > half_pixels
                    median_val = val
                    break
                else
                    n_pixels_below_median += hist[val]
                end
            end
        elseif n_pixels_below_median > half_pixels
            for val in median_val-1:-1:1
                n_pixels_below_median -= hist[val]
                if n_pixels_below_median <= half_pixels
                    median_val = val
                    break
                end
            end
        end
        return median_val, n_pixels_below_median
    end

    result = copy(img)

    if iseven(window_size[1]) || iseven(window_size[2])
        error("window width and height must be odd.")
    end

    half_window_size = window_size.>>1
    img_size = size(img)
    half_pixels = (window_size[1]*window_size[2])>>1

    for j in 1 + half_window_size[2]: img_size[2] - half_window_size[2]
        hist = zeros(Int, 2^8)
        for I in CartesianIndices((1:window_size[1], j-half_window_size[2]:j+half_window_size[2]))
            hist[img[I]+1] += 1
        end

        median_val, n_pixels_below_median = update_median(1, 0, hist, half_pixels)
        result[1+half_window_size[2], j] = UInt8(median_val-1)

        for i in 2 + half_window_size[1]: img_size[1] - half_window_size[1]
            for I in CartesianIndices((i-half_window_size[1]-1, j-half_window_size[2]:j+half_window_size[2]))
                hist[img[I]+1] -= 1
                if img[I] + 1 < median_val
                    n_pixels_below_median -= 1
                end
            end

            for I in CartesianIndices((i+half_window_size[1], j-half_window_size[2]:j+half_window_size[2]))
                hist[img[I]+1] += 1
                if img[I] + 1 < median_val
                    n_pixels_below_median += 1
                end
            end

            median_val, n_pixels_below_median = update_median(median_val, n_pixels_below_median, hist, half_pixels)
            result[i, j] = UInt8(median_val-1)
        end
    end
    result
end

function adjust_contrast(img::AbstractArray{T,2},min_c::Real,max_c::Real) where T

    out_max = 255
    out_min = 0

    out_min_rel = min_c / out_max
    out_max_rel = max_c / out_max
    out_range_f = (max_c - min_c) / (out_max - out_min)

    for i=1:length(img)

        inten = img[i] / 255 #normalize to float of intensity

        inten = inten - out_min_rel # remove bottom
        if inten < 0.0
            inten = 0.0
        end

        inten = inten / out_range_f

        if inten > 1.0
            inten = 1.0
        end

        if eltype(img) <: Integer
            img[i] = round(T,inten * out_max)
        else
            img[i] = inten * out_max
        end
    end

    nothing
end