function distance_weigth(a::NTuple{N,T}, b::NTuple{N,T}; order=4) where {N,T}
    return 1 / distance(a, b)^order
end

## CPU 2D

@inbounds function _gathering!(upper, lower, Fpi, p, x, y, dxi, order)

    # check that the particle is inside the grid
    isinside(p, x, y)

    # indices of lowermost-left corner of
    # the cell containing the particlex
    idx_x, idx_y = parent_cell(p, dxi)

    ω = (
        distance_weigth((x[idx_x], y[idx_y]), p; order=order),
        distance_weigth((x[idx_x + 1], y[idx_y]), p; order=order),
        distance_weigth((x[idx_x], y[idx_y + 1]), p; order=order),
        distance_weigth((x[idx_x + 1], y[idx_y + 1]), p; order=order),
    )

    nt = Threads.threadid()

    upper[nt][idx_x, idx_y] += ω[1] * Fpi
    upper[nt][idx_x + 1, idx_y] += ω[2] * Fpi
    upper[nt][idx_x, idx_y + 1] += ω[3] * Fpi
    upper[nt][idx_x + 1, idx_y + 1] += ω[4] * Fpi
    lower[nt][idx_x, idx_y] += ω[1]
    lower[nt][idx_x + 1, idx_y] += ω[2]
    lower[nt][idx_x, idx_y + 1] += ω[3]

    return lower[nt][idx_x + 1, idx_y + 1] += ω[4]
end

function gathering!(F::Array{T,2}, Fp::Vector{T}, xi, particle_coords; order=4) where {T}

    # unpack tuples
    px, py = particle_coords
    x, y = xi
    dxi = (x[2] - x[1], y[2] - y[1])

    # number of particles
    np = length(Fp)

    # TODO think about pre-allocating these 2 buffers
    upper = [zeros(size(F)) for _ in 1:Threads.nthreads()]
    lower = [zeros(size(F)) for _ in 1:Threads.nthreads()]

    # compute ∑ωᵢFᵢ and ∑ωᵢ
    Threads.@threads for i in 1:np
        _gathering!(upper, lower, Fp[i], (px[i], py[i]), x, y, dxi, order)
    end

    # compute Fᵢ=∑ωᵢFpᵢ/∑ωᵢ
    Threads.@threads for i in eachindex(F)
        @inbounds F[i] =
            sum(upper[nt][i] for nt in 1:Threads.nthreads()) /
            sum(lower[nt][i] for nt in 1:Threads.nthreads())
    end
end

## CPU 3D

@inbounds function _gathering!(upper, lower, Fpi, p, x, y, z, dxi, order)
    # check that the particle is inside the grid
    isinside(p, x, y, z)

    # indices of lowermost-left corner of
    # the cell containing the particle
    idx_x, idx_y, idx_z = parent_cell(p, dxi)

    ω = (
        distance_weigth((x[idx_x], y[idx_y], z[idx_z]), p; order=order),
        distance_weigth((x[idx_x + 1], y[idx_y], z[idx_z]), p; order=order),
        distance_weigth((x[idx_x], y[idx_y + 1], z[idx_z]), p; order=order),
        distance_weigth((x[idx_x + 1], y[idx_y + 1], z[idx_z]), p; order=order),
        distance_weigth((x[idx_x], y[idx_y], z[idx_z + 1]), p; order=order),
        distance_weigth((x[idx_x + 1], y[idx_y], z[idx_z + 1]), p; order=order),
        distance_weigth((x[idx_x], y[idx_y + 1], z[idx_z + 1]), p; order=order),
        distance_weigth((x[idx_x + 1], y[idx_y + 1], z[idx_z + 1]), p; order=order),
    )

    nt = Threads.threadid()

    upper[nt][idx_x, idx_y, idx_z] += ω[1] * Fpi
    upper[nt][idx_x + 1, idx_y, idx_z] += ω[2] * Fpi
    upper[nt][idx_x, idx_y + 1, idx_z] += ω[3] * Fpi
    upper[nt][idx_x + 1, idx_y + 1, idx_z] += ω[4] * Fpi
    upper[nt][idx_x, idx_y, idx_z + 1] += ω[5] * Fpi
    upper[nt][idx_x + 1, idx_y, idx_z + 1] += ω[6] * Fpi
    upper[nt][idx_x, idx_y + 1, idx_z + 1] += ω[7] * Fpi
    upper[nt][idx_x + 1, idx_y + 1, idx_z + 1] += ω[8] * Fpi
    lower[nt][idx_x, idx_y, idx_z] += ω[1]
    lower[nt][idx_x + 1, idx_y, idx_z] += ω[2]
    lower[nt][idx_x, idx_y + 1, idx_z] += ω[3]
    lower[nt][idx_x + 1, idx_y + 1, idx_z] += ω[4]
    lower[nt][idx_x, idx_y, idx_z + 1] += ω[5]
    lower[nt][idx_x + 1, idx_y, idx_z + 1] += ω[6]
    lower[nt][idx_x, idx_y + 1, idx_z + 1] += ω[7]

    return lower[nt][idx_x + 1, idx_y + 1, idx_z + 1] += ω[8]
end

function gathering!(F::Array{T,3}, Fp::Vector{T}, xi, particle_coords; order=4) where {T}
    # unpack tuples
    px, py, pz = particle_coords
    x, y, z = xi
    dxi = (x[2] - x[1], y[2] - y[1], z[2] - z[1])

    # number of particles
    np = length(Fp)

    # TODO think about pre-allocating these 2 buffers
    upper = [zeros(size(F)) for _ in 1:Threads.nthreads()]
    lower = [zeros(size(F)) for _ in 1:Threads.nthreads()]

    # compute ∑ωᵢFᵢ and ∑ωᵢ
    Threads.@threads for i in 1:np
        _gathering!(upper, lower, Fp[i], (px[i], py[i], pz[i]), x, y, z, dxi, order)
    end

    # compute Fᵢ=∑ωᵢFpᵢ/∑ωᵢ
    Threads.@threads for i in eachindex(F)
        @inbounds F[i] =
            sum(upper[nt][i] for nt in 1:Threads.nthreads()) /
            sum(lower[nt][i] for nt in 1:Threads.nthreads())
    end
end

## CUDA 2D

function _gather1!(
    upper::CuDeviceMatrix{T,1},
    lower::CuDeviceMatrix{T,1},
    Fpd::CuDeviceVector{T,1},
    xi,
    dxi,
    p;
    order=4,
) where {T}
    idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x

    # unpack tuples
    px, py = p
    x, y = xi

    @inbounds if idx ≤ length(px)
        # check that the particle is inside the grid
        isinside(px[idx], py[idx], x, y)

        p_idx = (px[idx], py[idx])

        # indices of lowermost-left corner of
        # the cell containing the particle
        idx_x, idx_y = parent_cell(p_idx, dxi)

        ω1::Float64 = distance_weigth((x[idx_x], y[idx_y]), p_idx; order=order)
        ω2::Float64 = distance_weigth((x[idx_x + 1], y[idx_y]), p_idx; order=order)
        ω3::Float64 = distance_weigth((x[idx_x], y[idx_y + 1]), p_idx; order=order)
        ω4::Float64 = distance_weigth((x[idx_x + 1], y[idx_y + 1]), p_idx; order=order)

        Fpi::Float64 = Fpd[idx]

        CUDA.@atomic upper[idx_x, idx_y] += ω1 * Fpi
        CUDA.@atomic upper[idx_x + 1, idx_y] += ω2 * Fpi
        CUDA.@atomic upper[idx_x, idx_y + 1] += ω3 * Fpi
        CUDA.@atomic upper[idx_x + 1, idx_y + 1] += ω4 * Fpi
        CUDA.@atomic lower[idx_x, idx_y] += ω1
        CUDA.@atomic lower[idx_x + 1, idx_y] += ω2
        CUDA.@atomic lower[idx_x, idx_y + 1] += ω3
        CUDA.@atomic lower[idx_x + 1, idx_y + 1] += ω4
    end

    return nothing
end

function _gather2!(Fd::CuDeviceArray{T,2}, upper, lower) where {T}
    idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    idy = (blockIdx().y - 1) * blockDim().y + threadIdx().y

    if (idx < size(Fd, 1)) && (idy < size(Fd, 2))
        @inbounds Fd[idx, idy] = upper[idx, idy] / lower[idx, idy]
    end

    return nothing
end

function gathering!(
    Fd::CuArray{T,2}, Fpd::CuArray{T,1}, xi, dxi, particle_coords; nt=512
) where {T}
    upper = CUDA.zeros(T, size(Fd))
    lower = CUDA.zeros(T, size(Fd))
    dxi = (x[2] - x[1], y[2] - y[1])

    # first kernel that computes ∑ωᵢFᵢ and ∑ωᵢ
    N = length(Fpd)
    numblocks = ceil(Int, N / nt)
    CUDA.@sync begin
        @cuda threads = nt blocks = numblocks _gather1!(
            upper, lower, Fpd, xi, dxi, particle_coords
        )
    end

    # seond and final kernel that computes Fᵢ=∑ωᵢFpᵢ/∑ωᵢ
    nx, ny = size(Fd)
    nblocksx = ceil(Int, nx / 32)
    nblocksy = ceil(Int, ny / 32)
    CUDA.@sync begin
        @cuda threads = (32, 32) blocks = (nblocksx, nblocksy, nblocksz) _gather2!(
            Fd, upper, lower
        )
    end
end

## CUDA 3D

function _gather1!(
    upper::CuDeviceArray{T,3},
    lower::CuDeviceArray{T,3},
    Fpd::CuDeviceVector{T,1},
    xi,
    dxi,
    p;
    order=4,
) where {T}
    idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x

    # unpack tuples
    px, py, pz = p
    x, y, z = xi

    @inbounds if idx ≤ length(px)
        # check that the particle is inside the grid
        isinside(px[idx], py[idx], pz[idx], x, y, z)

        p_idx = (px[idx], py[idx], pz[idx])

        # indices of lowermost-left corner of
        # the cell containing the particle
        idx_x, idx_y, idx_z = parent_cell(p_idx, dxi)

        ω1::Float64 = distance_weigth((x[idx_x], y[idx_y], z[idx_z]), p_idx; order=order)
        ω2::Float64 = distance_weigth(
            (x[idx_x + 1], y[idx_y], z[idx_z]), p_idx; order=order
        )
        ω3::Float64 = distance_weigth(
            (x[idx_x], y[idx_y + 1], z[idx_z]), p_idx; order=order
        )
        ω4::Float64 = distance_weigth(
            (x[idx_x + 1], y[idx_y + 1], z[idx_z]), p_idx; order=order
        )
        ω5::Float64 = distance_weigth(
            (x[idx_x], y[idx_y], z[idx_z + 1]), p_idx; order=order
        )
        ω6::Float64 = distance_weigth(
            (x[idx_x + 1], y[idx_y], z[idx_z + 1]), p_idx; order=order
        )
        ω7::Float64 = distance_weigth(
            (x[idx_x], y[idx_y + 1], z[idx_z + 1]), p_idx; order=order
        )
        ω8::Float64 = distance_weigth(
            (x[idx_x + 1], y[idx_y + 1], z[idx_z + 1]), p_idx; order=order
        )

        Fpi::Float64 = Fpd[idx]

        CUDA.@atomic upper[idx_x, idx_y, idx_z] += ω1 * Fpi
        CUDA.@atomic upper[idx_x + 1, idx_y, idx_z] += ω2 * Fpi
        CUDA.@atomic upper[idx_x, idx_y + 1, idx_z] += ω3 * Fpi
        CUDA.@atomic upper[idx_x + 1, idx_y + 1, idx_z] += ω4 * Fpi
        CUDA.@atomic upper[idx_x, idx_y, idx_z + 1] += ω5 * Fpi
        CUDA.@atomic upper[idx_x + 1, idx_y, idx_z + 1] += ω6 * Fpi
        CUDA.@atomic upper[idx_x, idx_y + 1, idx_z + 1] += ω7 * Fpi
        CUDA.@atomic upper[idx_x + 1, idx_y + 1, idx_z + 1] += ω8 * Fpi
        CUDA.@atomic lower[idx_x, idx_y, idx_z] += ω1
        CUDA.@atomic lower[idx_x + 1, idx_y, idx_z] += ω2
        CUDA.@atomic lower[idx_x, idx_y + 1, idx_z] += ω3
        CUDA.@atomic lower[idx_x + 1, idx_y + 1, idx_z] += ω4
        CUDA.@atomic lower[idx_x, idx_y, idx_z + 1] += ω5
        CUDA.@atomic lower[idx_x + 1, idx_y, idx_z + 1] += ω6
        CUDA.@atomic lower[idx_x, idx_y + 1, idx_z + 1] += ω7
        CUDA.@atomic lower[idx_x + 1, idx_y + 1, idx_z + 1] += ω8
    end

    return nothing
end

function _gather2!(
    Fd::CuDeviceArray{T,3}, upper::CuDeviceArray{T,3}, lower::CuDeviceArray{T,3}
) where {T}
    idx = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    idy = (blockIdx().y - 1) * blockDim().y + threadIdx().y
    idz = (blockIdx().z - 1) * blockDim().z + threadIdx().z

    if (idx ≤ size(Fd, 1)) && (idy ≤ size(Fd, 2)) && (idz ≤ size(Fd, 3))
        @inbounds Fd[idx, idy, idz] = upper[idx, idy, idz] / lower[idx, idy, idz]
    end

    return nothing
end

function gathering!(
    Fd::CuArray{T,3}, Fpd::CuArray{T,1}, xi, particle_coords; nt=512
) where {T}
    x, y, z = xi
    dxi = (x[2] - x[1], y[2] - y[1], z[2] - z[1])
    upper = CUDA.zeros(T, size(Fd))
    lower = CUDA.zeros(T, size(Fd))

    # first kernel that computes ∑ωᵢFᵢ and ∑ωᵢ
    N = length(Fpd)
    numblocks = ceil(Int, N / nt)
    CUDA.@sync begin
        @cuda threads = nt blocks = numblocks _gather1!(
            upper, lower, Fpd, xi, dxi, particle_coords
        )
    end

    # second and final kernel that computes Fᵢ=∑ωᵢFpᵢ/∑ωᵢ
    nx, ny, nz = size(Fd)
    nblocksx = ceil(Int, nx / 8)
    nblocksy = ceil(Int, ny / 8)
    nblocksz = ceil(Int, nz / 8)
    CUDA.@sync begin
        @cuda threads = (8, 8, 8) blocks = (nblocksx, nblocksy, nblocksz) _gather2!(
            Fd, upper, lower
        )
    end
end


#=
   Sudoku in Julia + ConstraintSolver

   This is based on the Sudoku example in
   https://wikunia.github.io/ConstraintSolver.jl/stable/tutorial/

   Model created by Hakan Kjellerstrand, hakank@gmail.com
   See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")
include("sudoku_problems.jl")

function sudoku(grid,print_solution=true,timeout=Inf,all_solutions=false)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    # Solver options:
    # https://wikunia.github.io/ConstraintSolver.jl/stable/options/
    # m = Model(optimizer_with_attributes(CS.Optimizer, "all_solutions"=>true,"logging"=>[]))
    m = Model(optimizer_with_attributes(CS.Optimizer,
                                        "all_solutions"=>all_solutions,
                                        "logging"=>[], # :true,

                                        "traverse_strategy"=>:BFS,
                                        # "traverse_strategy"=>:DFS,
                                        # "traverse_strategy"=>:DBFS,

                                        "branch_split"=>:Smallest,
                                        # "branch_split"=>:Biggest,
                                        # "branch_split"=>:InHalf,

                                        # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                        "branch_strategy" => :IMPS, # default
                                        # "branch_strategy" => :ABS, # Activity Based Search
                                        # "activity.decay" => 0.999, # default 0.999
                                        # "activity.max_probes" => 100, # default, 10
                                        # "activity.max_confidence_deviation" => 20, # default 20

                                        # "simplify"=>false,
                                        # "simplify"=>true,
                                        "time_limit"=>timeout,

                                        # "lp_optimizer" => cbc_optimizer,
                                        # "lp_optimizer" => glpk_optimizer,
                                        # "lp_optimizer" => ipopt_optimizer,

                                        ))

    n,_ = size(grid)  # 9
    ns = round(Int, sqrt(n)) # 3
    n2 = ns-1 # 2

    # define the variables
    @variable(m, 1 <= x[1:n,1:n] <= n, Int)

    # hints
    for r=1:n, c=1:n
        if grid[r,c] != 0
            @constraint(m, x[r,c] == grid[r,c])
        end
    end

    # Row and column constraints
    for rc = 1:n
        @constraint(m, x[rc,:] in CS.AllDifferent())
        @constraint(m, x[:,rc] in CS.AllDifferent())
    end

    # Alterntive using transpose (just checking)
    # xt = x' # Transposed version
    # for rc = 1:n
    #    @constraint(m, x[rc,:] in CS.AllDifferent())
    #    @constraint(m, xt[rc,:] in CS.AllDifferent())
    # end

    # Cell constraints
    for i in 1:ns:n, j in 1:ns:n
        @constraint(m, vec([x[i+k,j+l] for k in 0:n2, l in 0:n2]) in CS.AllDifferent())
    end

    # Solve the problem
    optimize!(m)

    status = JuMP.termination_status(m)
    println("status:$status")
    if status == MOI.OPTIMAL
        x = convert.(Integer,JuMP.value.(x))
        if print_solution
            # display(@show x)
            print_grid(x)
        end
        # num_sols = MOI.get(m, MOI.ResultCount())
        # println("num_sols:$num_sols")
        return x
    end
end

#
# Solve the problems in list a
# The problems are in sudoku_problems.jl
#
function check_sudoku(problems, a::Array,print_solutions=true, timeout=Inf)
    println("a:$a")
    for p in a
        println("p:$p")
        p2 = resize_matrix(problems[p])
        @time sudoku(p2,print_solutions,timeout)
    end
end

#
# Solve all problems of size n x n
# The problems are in sudoku_problems.jl
#
function check_sudoku(problems, n::Int,print_solutions=true,timeout=Inf)
    num_problems = 0
    total_time = 0
    for p in sort(collect(keys(problems)),by=k->Symbol(k))
        pp = problems[p]
        len = length(pp)
        if len == n
            num_problems += 1
            println("\n\nproblem $p")
            p2 = resize_matrix(problems[p])
            t = @timed @time sudoku(p2,print_solutions,timeout)
            total_time += t.time
        end
    end
    println("\nnum_problems:$num_problems total_time:$total_time")
end

#
# Solve Norvig's 95 Sudoku's from http://norvig.com/top95.txt
#
# 0.47s
function norvig_problems(print_solutions=false)
   file = "top95.txt"
   # Try to download the file if not exist
   if !isfile(file)
       println("$file does not exist. Trying to download it from http://norvig.com/top95.txt")
       tmp = download("http://norvig.com/top95.txt")
       mv(tmp,file)
       println("downloaded $file")
   end
   if !isfile(file)
       println("$file does not exist and download() didn't work. Please download from http://norvig.com/top95.txt")
       return
   end
   lines = readlines(file).|>line->replace(line,"."=>"0")
   c = 0
   for line in lines
       c += 1
       p = reshape(split(line,"").|>i->parse(Int,i),(9,9))
       @time sudoku(p,print_solutions)
   end
   println("Solve $c Sudoku instances.")

end

timeout = 600.0 # seconds. Only needed for the 25x25 instances
println("Timeout: $timeout")
print_solutions=false

#
# Run the 9x9 instances
#
@time check_sudoku(all_sudoku_problems, 9, print_solutions,timeout) # 40 problems, 0.33s

#
# Run the 16x16 instances
#
@time check_sudoku(all_sudoku_problems, 16,print_solutions,timeout) # 55 problems: 0.59s

#
# Run the 25x25 instances
# Note The 25x25 takes a lots of time so we have to use a timeout when running them.
#
#=
   With a timeout of 600s (10 min)
   - #18: Timeout -> 118s  --> 115s (v0.6.7)
   - #19: 146.5s -> 52.9 (v0.6.7)
   - #20: 515.9s -> 108s (v0.6.7)
   - #21-#29: Timeout
   - #30: 582s
   - #31-#33: Timeout
   - #89: 3.8s -> 1.6s
   - #90: 7.3s -> 8.4s (v0.6.7)
=#
# @time check_sudoku(all_sudoku_problems, 25, true,timeout) # test all instances
# @time check_sudoku(all_sudoku_problems, [19,89,90],true, timeout) # just test the simplest problems
# @time check_sudoku(all_sudoku_problems, [18,19,20,21,89,90],true, timeout) # just test the simplest problems
# @time check_sudoku(all_sudoku_problems, [30],true, timeout) # just test the simplest problems

#
# Run Norvig's 95 Sudoku problems: 0.47s
#
# @time norvig_problems(false)
#


using CircularArrays
using Test

v1 = CircularVector(rand(Int64, 5))

@test IndexStyle(CircularArray) == IndexCartesian()
@test IndexStyle(CircularVector) == IndexLinear()

@test size(v1, 1) == 5
@test typeof(v1) == CircularVector{Int64}
@test isa(v1, CircularVector)
@test isa(v1, AbstractVector{Int})
@test !isa(v1, AbstractVector{String})
@test v1[2] == v1[2 + length(v1)]
v1[2] = 0
v1[3] = 0
@test v1[2] == v1[3]
@test_throws MethodError v1[2] = "Hello"

v2 = CircularVector("abcde", 5)

@test prod(v2) == "abcde"^5

@test_throws MethodError push!(v1, 15)

b_arr = [2 4 6 8; 10 12 14 16; 18 20 22 24]
a1 = CircularArray(b_arr)
@test size(a1) == (3, 4)
@test a1[2, 3] == 14
a1[2, 3] = 17
@test a1[2, 3] == 17
@test !isa(a1, CircularVector)
@test !isa(a1, AbstractVector)
@test isa(a1, AbstractArray)

@test size(reshape(a1, (2, 2, 3))) == (2, 2, 3)

a2 = CircularArray(4, (2, 3))
@test isa(a2, CircularArray{Int, 2})


using Formulas
using Base.Test

## from Formulas
f = Formula(nothing, 0)
t = Terms(f)
@test t.intercept == false
@test t.response == false
@test t.terms == []
@test t.eterms == []

t = Terms(a ~ b)
@test t.terms == Term[:b]
@test t.eterms == [:a, :b]
@test t.intercept == true
@test t.response == true

## totally empty
t = Terms(Formula(nothing, 0))
@test t.response == false
@test t.intercept == false
@test t.terms == []
@test t.eterms == []

## empty RHS
t = Terms(y ~ 0)
@test t.intercept == false
@test t.terms == []
@test t.eterms == [:y]
t = Terms(y ~ -1)
@test t.intercept == false
@test t.terms == []

## intercept-only
t = Terms(y ~ 1)
@test t.response == true
@test t.intercept == true
@test t.terms == []
@test t.eterms == [:y]

## terms add
t = Terms(y ~ 1 + x1 + x2)
@test t.intercept == true
@test t.terms == Term[:x1, :x2]
@test t.eterms == [:y, :x1, :x2]

## implicit intercept behavior:
t = Terms(y ~ x1 + x2)
@test t.intercept == true
@test t.terms == Term[:x1, :x2]
@test t.eterms == [:y, :x1, :x2]

## no intercept
t = Terms(y ~ 0 + x1 + x2)
@test t.intercept == false
@test t.terms == Term[:x1, :x2]

@test t == Terms(y ~ -1 + x1 + x2) == Terms(y ~ x1 - 1 + x2) == Terms(y ~ x1 + x2 -1)

## can't subtract terms other than 1
@test_throws ErrorException Terms(y ~ x1 - x2)

t = Terms(y ~ x1 & x2)
@test t.terms == Term[:(x1 & x2)]
@test t.eterms == [:y, :x1, :x2]

## `*` expansion
t = Terms(y ~ x1 * x2)
@test t.terms == Term[:x1, :x2, :(x1 & x2)]
@test t.eterms == [:y, :x1, :x2]

## associative rule:
## +
t = Terms(y ~ x1 + x2 + x3)
@test t.terms == Term[:x1, :x2, :x3]

## &
t = Terms(y ~ x1 & x2 & x3)
@test t.terms == Term[:((&)(x1, x2, x3))]
@test t.eterms == [:y, :x1, :x2, :x3]

## distributive property of + and &
t = Terms(y ~ x1 & (x2 + x3))
@test t.terms == Term[:(x1&x2), :(x1&x3)]

t = Terms(y ~ (x1 + x2) & x3)
@test t.terms == Term[:(x1&x3), :(x2&x3)]

t = Terms(y ~ (x1 + x2) & (x3 + x4))
@test t.terms == Term[:(x1&x3), :(x1&x4), :(x2&x3), :(x2&x4)]

## three-way *
t = Terms(y ~ x1 * x2 * x3)
@test t.terms == Term[:x1, :x2, :x3,
                  :(x1&x2), :(x1&x3), :(x2&x3),
                  :((&)(x1, x2, x3))]
@test t.eterms == [:y, :x1, :x2, :x3]


T_memoize = Dict{Expr,Expr}()

@testset "@T" begin
    T = collect(1:100)

    @T(A = t)
    @test A == collect(1:100)

    @T(B = t)
    @test B == collect(1:100)

    @T begin D = t + T[t] end
    @test D == collect(1:100)*2

    @T(C = ( t==1  ? 51 : C[t-1]+1 ))
    @test C == collect(51:50+100)

    q = 10
    @T(D = C + B + 7 + q)
    @test D == @.(C + B + 7 + q)

    @T begin sum_assured_t0 =
        10000
    end
    @test sum_assured_t0 == 10000

    @T H = t==1 ? 1 : A[t-1] + B
    @test H == [t==1 ? 1 : A[t-1] + B[t] for t ∈ T]

    @test_throws CalculationSequenceError @T a = a[t]

    T = collect(1:600)
    @T begin
        age_t0 = 30
        duration_t0 = 6
        premium_t0 = 100
        sum_assured_t0 = 10000
        age = floor(Int, age_t0 + t/12)
        duration = duration_t0 + t
        qx = age / 1000
        lapse_rate_m = 1 / qx / 1000
        no_pols_som = if t==1 1 else no_pols_eom[t-1] end
        no_deaths = no_pols_som*qx/12*(1-0.5*lapse_rate_m)
        no_surrenders = no_pols_som*lapse_rate_m*(1-0.5*qx)
        no_pols_eom = if t==1 1 else no_pols_eom[t-1] end - no_deaths - no_surrenders
        premium_income = premium_t0*no_pols_eom/12
        claims_outgo = sum_assured_t0*no_deaths
        profit = premium_income-claims_outgo
        yield_curve = 0.1
        profit_discounted  = profit*(1+yield_curve)^(-t)
        liability = -profit_discounted+ if t<600 liability[t+1] else 0 end
    end

end

@testset "cluster_topology" begin
    formulas = OrderedDict(
        :A => :((t == 1 ? 1 : A[t - 1]- C) - D),
        :B => :(t == 1 ? 1 : A[t - 1]),
        :C => :(((B * E) / 12) * (1 - 0.5*F)),
        :D => :(B * F * (1 - 0.5*E))
      )
      @test cluster_topology(formulas) == (-1, [:B, :C, :D, :A])
end

"
For a given cluster of formulae, derive the order in which these formulae
should be computed within `for t ∈ T`, and return this order along with
how t is referenced: [-1, 0, +1, nothing]
"
function cluster_topology(formulas::OrderedDict)
    refs_dict = Dict{Symbol, Any}()
    Δ_dict = Dict{Symbol, Any}()
    all_Δ = Set()

    # Get all the t references and convert them to offsets Δ
    for (key, value) ∈ formulas
        refs = (Set(var=>ex for (var, ex) ∈ get_t_refs(value) if haskey(formulas, var)) ∪
                Set(var=>0 for var ∈ vars_without_t_index(value) if haskey(formulas, var)))

        Δ = Set(var=>   if t_linearcombinations(expr) == true
                            eval(assign_variable(expr, :t, 0))
                        else
                            0 # TODO: more inspection needed
                        end for (var, expr) ∈ refs)

        refs_dict[key] = refs
        Δ_dict[key] = Δ
        all_Δ = all_Δ ∪ [δ for (_, δ) ∈ Δ]
    end

    # Continue only if all the references are either backwards or forwards looking and not mixed
    rtype = all(all_Δ .<= 0) ? -1 : all(all_Δ .>= 0) ? +1 : nothing
    if rtype === nothing
        return rtype, keys(formulas)
    end

    edges = Vector{Tuple{Symbol, Symbol}}()
    for (key, Δ) in Δ_dict
        for (var, δ) ∈ Δ
            if δ==0
                push!(edges, (var, key))
            end
        end
    end

    # Get sequence... flatten... append lost keys
    sequence = RowSheets.traversalsequence!(RowSheets.DiGraph(edges))
    sequence = [var for cluster in sequence for var in cluster]
    sequence = [sequence; setdiff!(OrderedSet(keys(formulas)), sequence)...]

    return rtype, sequence
end

function transform_expr_T(expr; line::Union{LineNumberNode, Nothing}=nothing)
    empty!(T_memoize) # make empty to suppress memoization
    expr′ = striplines(expr)
    expr isa Expr && haskey(T_memoize, expr′) && return T_memoize[expr′]

    formulas = expr_to_assignment_dict(expr; line)
    ret = transform_expr_all(formulas)

    # Add the last entry as the return value
    (lastvar, _) = last(formulas)
    push!(ret.args, lastvar)

    T_memoize[expr′] = ret
    return ret
end

"
Transform all formulae/equations in a given variable library `:(B = t==1 ? 1 : A[t-1]; A = B[t]; C=A+100)`
into a coherent forloop `:(for t ∈ T A[t] = B[t]; B[t] = t==1 ? 1 : A[t-1] end; @.(C = A+100))`.

This will traverse the full graph dependencies
"
function transform_expr_all(formulas::OrderedDict{Symbol, CodePoint})
    # Convert the dictionary into a sequence
    symbol_links = Dict(key => Vector{Symbol}() for key ∈ keys(formulas))
    for (varⱼ, codepoint) ∈ formulas
        ex = codepoint.expr
        for varᵢ ∈ get_variables(ex)
            haskey(formulas, varᵢ) && push!(symbol_links[varᵢ], varⱼ)
        end
    end
    sequence = RowSheets.traversalsequence!(RowSheets.DiGraph(symbol_links))
    sequence_bias = OrderedDict(key=>i for (i,key) ∈ enumerate(keys(formulas)))

    # Keep order in cluster same as input order (to allow user intervension)
    expr′_list = []
    for cluster ∈ sequence
        formulas′ = OrderedDict(i=>formulas[i] for i ∈ sort(cluster, by=x->(sequence_bias[x])))
        push!(expr′_list, transform_expr_cluster(formulas′))
    end

    return Expr(:block, expr′_list...)
end

"
Transform a cluster of formulae/equations that is possibly cyclic `:(B = t==1 ? 1 : A[t-1]; A = B[t])`
into a coherent forloop `:(for t ∈ T A[t] = B[t]; B[t] = t==1 ? 1 : A[t-1] end)`.

This can be used in each traversal step when traversing the full formuale dependency graph.
"
function transform_expr_cluster(formulas::OrderedDict{Symbol, CodePoint})

    # Single definition might have some good shortcuts
    ret = nothing
    if length(formulas) == 1
        (var, codepoint) = first(formulas)
        e = codepoint.expr
        if !has_symbol(e, :t)
            # Do `$var = @.($expr)`
            ret = Expr(:block,
                        codepoint.line,
                        Expr(:(=), var, Expr(:macrocall, Symbol("@__dot__"), codepoint.line, e)))

        elseif !has_symbol(e, var)
            # Do  [$var[t] = $expr′[t] for t ∈ T]
            ret =  Expr(:block,
                        codepoint.line,
                        Expr(:(=), var, Expr(:comprehension, Expr(:generator, wrap_t_getter(e), Expr(:(=), :t, :T)))))
        end
    end
    if ret === nothing
        # Long way around
        rtype, seq = cluster_topology(OrderedDict(i=>j.expr for (i,j) ∈ formulas))

        definitions = Expr(:block)
        for var ∈ seq
            push!(definitions.args, formulas[var].line)
            push!(definitions.args, :($(var) = Vector(undef, length(T))))
        end

        loopover =  rtype > 0 ? :(reverse(T)) : (:T)

        assignments = Expr(:block)
        for var ∈ seq
            push!(assignments.args, formulas[var].line)
            push!(assignments.args, :($(var)[t] = $(wrap_t_getter(formulas[var].expr))))
        end

        @gensymx e
        ret = (quote
            $definitions
            try
                for t in $loopover
                    $assignments
                end
            catch $e
                $e isa UndefRefError && throw($CalculationSequenceError())
                rethrow($e)
            end
            nothing
        end)
    end

    return ret
end

macro T(expr)
    esc(transform_expr_T(expr; line=__source__))
end

macro no_op(x...)
end


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module SeqUtils

using Nemo

export ModuleSeqUtils
export SeqShow, SeqPrint, SeqName, SeqNum, Println, ZArray, Nemofmpz

"""
Nemo is a library designed, developed and maintained by William Hart with the help of others. Many functions in our project are based on Nemo.
"""
const ModuleSeqUtils = ""

"""
The basic integer type of our project is 'Nemo.fmpz'. Think of a 'Nemo.fmpz' as a 'BigInt'. 'fmpz' stands for 'fast multiple precision zahl' (zahl=integer). The mathematical symbol for the ring of integers is the blackboard (double-struck) Z, also written ZZ. In reference to this ZZ(n) defines the integer n as a fmpz.
"""
function Nemofmpz(n::Int) return ZZ(n) end

"""
Return an array of type `fmpz` of length ``n`` preset with ``0``.
"""
ZArray(len::Int) = zeros(ZZ, len)

"""
Return an array of type fmpz and of size ``n`` filled by the values of the function ``f`` in the range `offset:n`.
"""
function ZArray(n::Int, f::Function, offset=1)
    n ≤ 0 && return fmpz[]
    fmpz[f(k) for k in offset:offset+n-1] # !
end

"""
Return the size of the array ``A``.
"""
SeqSize(A) = Base.length(A)

"""
Return the range of a SeqArray.
"""
SeqRange(A) = 1:SeqSize(A)

"""
Print the array ``A`` in the format ``n ↦ A[n]`` for n in the given range.
"""
function SeqShow(A::Array, R::AbstractRange, offset=0)
    for k in R
        if isassigned(A, k)
            println(k + offset, " ↦ ", A[k])
        else
            println(k + offset, " ↦ ", "undef")
        end
    end
end

"""
Print the array ``A`` in the format ``n ↦ A[n]`` for n in the range first:last.
"""
SeqShow(A, first::Int, last::Int) = SeqShow(A, first:last)

"""
Print the array ``A`` in the format ``n ↦ A[n]``.
"""
SeqShow(A::Array, offset=0) = SeqShow(A, 1:length(A), offset - 1)

function print_without_type(io, v::AbstractVector)
    print(io, "[")
    for (i, el) in enumerate(v)
        i > 1 && print(io, ", ")
        print(io, el)
    end
    println(io, "]")
end

"""
Print the array without typeinfo.
"""
Println(v::AbstractVector) = print_without_type(IOContext(stdout), v)

"""
Print the array with or without typeinfo.
"""
function SeqPrint(v::AbstractVector, typeinfo = false)
    typeinfo ? println(v) : Println(v)
end

"""
Valid prefixes to the numerical part of the OEIS A-numbers.

* C => Channel
* F => Filter (all below n)
* G => Generating function
* I => Iteration
* L => List (array based)
* M => Matrix
* R => RealFunction
* S => Staircase (iteration)
* T => Triangle (iteration)
* TA => Triangle (triangular array)
* TL => Triangle (flat-list array)
* V => Value
* is => is a (predicate), boolean
"""
const ValidPrefixes = ['C', 'F', 'G', 'I', 'L', 'M', 'R', 'S', 'T', 'V']

"""
Return the name of a OEIS sequence given a similar named function as a string.
"""
function SeqName(fun)
    aname = string(fun)
    for X in ValidPrefixes
        aname = replace(aname, X => 'A')
    end

    if !occursin(r"^A[0-9]{6}$", aname)
        fullname = split(aname, ".")
        aname = String(fullname[end])
        if !occursin(r"^A[0-9]{6}$", aname)
            @error(aname, " is not a valid A-name!")
            return
        end
    end
    aname
end

"""
Return the A-number of a OEIS sequence given a similar named function as an integer.
"""
function SeqNum(seq)
    aname = SeqName(seq)
    parse(Int, aname[2:end])
end

#START-TEST-########################################################

using Test

function test()
    #@testset "SeqUtils" begin
    #end
end

function demo()
    A = fmpz[1, 2, 3, 4, 5]

    println()
    SeqShow(A)
    println()

    SeqShow(A, -4)
    println()

    A |> Println
end

function perf()
end

function main()
    test()
    demo()
    perf()
end

main()

end # module


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,jl:hydrogen
#     text_representation:
#       extension: .jl
#       format_name: hydrogen
#       format_version: '1.3'
#       jupytext_version: 1.11.2
#   kernelspec:
#     display_name: Julia 1.8.0-DEV
#     language: julia
#     name: julia-1.8
# ---

# %% [markdown]
# https://nbviewer.jupyter.org/gist/antimon2/dd23b1e62d5fcd08d91fd420d7ba996d/Mandelbrot.mt.cu.jl.ipynb

# %%
using BenchmarkTools
using Plots
gr(fmt=:png)
using CUDA

plotmandelbrot(m) = heatmap(m; c=reverse(cgrad(:jet1)),
    size=(300, 300), colorbar=false, ticks=false, frame=false)

function mandelbrot(c; maxiters=2^10, threshold_abs2=2^10)
    z = zero(c)
    for i in 1:maxiters
        z = z * z + c
        abs2(z) ≥ threshold_abs2 && return i
    end
    maxiters + 1
end

n = 2^8
x = range(-0.714689, -0.714679; length=n)
y = range( 0.299872,  0.299882; length=n)
c = complex.(x', y)

@show typeof(c)
@time m = mandelbrot.(c)
@time m = mandelbrot.(c)
@show typeof(m)
plotmandelbrot(m)

# %%
c_cuda = CuMatrix{ComplexF64}(c)

@show typeof(c_cuda)
@time m_cuda = collect(mandelbrot.(c_cuda))
@time m_cuda = collect(mandelbrot.(c_cuda))
@show typeof(mandelbrot.(c_cuda))
@show typeof(collect(mandelbrot.(c_cuda)))
plotmandelbrot(m_cuda)

# %%
x32 = range(-0.714689f0, -0.714679f0; length=n)
y32 = range( 0.299872f0,  0.299882f0; length=n)
c32 = complex.(x32', y32)

@show typeof(c32)
@time m32 = mandelbrot.(c32)
@time m32 = mandelbrot.(c32)
@show typeof(m32)
plotmandelbrot(m32)

# %%
c32_cuda = cu(c32)

@show typeof(c32_cuda)
@time m32_cuda = collect(mandelbrot.(c32_cuda))
@time m32_cuda = collect(mandelbrot.(c32_cuda))
@show typeof(mandelbrot.(c32_cuda))
@show typeof(collect(mandelbrot.(c32_cuda)))
plotmandelbrot(m32_cuda)

# %%
function mandelbrot_threads(c)
    m = similar(c, typeof(mandelbrot(c[end])))
    Threads.@threads for i in keys(c)
        m[i] = mandelbrot(c[i])
    end
    m
end

@show Threads.nthreads()
@time m_th = mandelbrot_threads(c)
@time m_th = mandelbrot_threads(c)
plotmandelbrot(m_th)

# %%
@benchmark mandelbrot.($c) # CPU Float64

# %%
@benchmark collect(mandelbrot.($c_cuda)) # GPU Float64 (GPU → GPU → CPU)

# %%
@benchmark collect(mandelbrot.(CuMatrix{ComplexF64}($c))) # GPU Float64 (CPU → GPU → GPU → CPU)

# %%
@benchmark mandelbrot.($c32) # CPU Float32

# %%
@benchmark collect(mandelbrot.($c32_cuda)) # GPU Float32 (GPU → GPU → CPU)

# %%
@benchmark collect(mandelbrot.(cu($c32))) # GPU Float32 (CPU → GPU → GPU → CPU)

# %%
@show Threads.nthreads()
@benchmark mandelbrot_threads($c) # CPU Float64 multi-threaded version with nthreads = 12

# %%
N = 2^10
X = range(-0.714689, -0.714679; length=N)
Y = range( 0.299872,  0.299882; length=N)
C = complex.(X', Y)
X32 = range(-0.714689f0, -0.714679f0; length=N)
Y32 = range( 0.299872f0,  0.299882f0; length=N)
C32 = complex.(X32', Y32)

C_cuda = CuMatrix{ComplexF64}(C)
C32_cuda = cu(C32);

# %%
@benchmark mandelbrot.($C) # CPU Float64

# %%
@benchmark collect(mandelbrot.($C_cuda)) # GPU Float64 (GPU → GPU → CPU)

# %%
@benchmark collect(mandelbrot.(CuMatrix{ComplexF64}($C))) # GPU Float64 (CPU → GPU → GPU → CPU)

# %%
@benchmark mandelbrot.($C32) # CPU Float32

# %%
@benchmark collect(mandelbrot.($C32_cuda)) # GPU Float32 (GPU → GPU → CPU)

# %%
@benchmark collect(mandelbrot.(cu($C32))) # GPU Float32 (CPU → GPU → GPU → CPU)

# %%


using StatsBase
using Test

wechsler = Float32[
     7  5  9  8
     8  8  5  6
    16 18 11  9
     8  3  7  9
     6  3 13  9
    11  8 10 10
    12  7  9  8
     8 11  9  3
    14 12 11  4
    13 13 13  6
    13  9  9  9
    13 10 15  7
    14 11 12  8
    15 11 11 10
    13 10 15  9
    10  5  8  6
    10  3  7  7
    17 13 13  7
    10  6 10  7
    10 10 15  8
    14  7 11  5
    16 11 12 11
    10  7 14  6
    10 10  9  6
    10  7 10 10
     7  6  5  9
    15 12 10  6
    17 15 15  8
    16 13 16  9
    13 10 17  8
    13 10 17 10
    19 12 16 10
    19 15 17 11
    13 10  7  8
    15 11 12  8
    16  9 11 11
    14 13 14  9
]

@test @inferred(partialcor(wechsler[:,1], wechsler[:,2], wechsler[:,3:4])) ≈ 0.7118787 rtol=1e-6

X = [ 2 1 0
      4 2 0
     15 3 1
     20 4 1]
@test @inferred(partialcor(view(X,:,1), view(X,:,2), view(X,:,3))) ≈ 0.919145 rtol=1e-6


using Cairo
using ComplexVisual
@ComplexVisual.import_huge

const CVG = try using ComplexVisualGtk; ComplexVisualGtk; catch nothing; end

const fontface = cv_fontface("sans-serif")

function create_scene(trafo,
        domain, codomain; cut_test=nothing, gap=80,
        axis_label_style=cv_black → fontface → cv_fontsize(20), padding=30) # {{{
    layout = CV_StateLayout(CV_2DLayout(), CV_CyclicValue(2))
    layout = cv_do_lr_layout(cv_add(layout, trafo, domain, codomain), gap)

    setup = cv_setup_cycle_state(CV_LRSetupChain(layout))
    setup = cv_setup_lr_painters(setup; cut_test,
        parallel_hlines_style=cv_op_source → cv_antialias_best →
            cv_linewidth(3) → cv_black,
        parallel_vlines_style=cv_op_source → cv_antialias_best →
            cv_linewidth(3) → cv_white,
        parallel_hlines=cv_parallel_lines(2.0+0.0im),
        parallel_vlines= cv_parallel_lines(0.0+2.0im))
    setup = cv_setup_lr_axis(setup; label_style=axis_label_style)
    setup = cv_setup_lr_border(setup)
    padding > 0 && cv_add_padding!(setup.layout, padding)
    setup = cv_setup_domain_codomain_scene(setup)
    cv_scene_lr_start(setup; state_start=2)
    return setup.layout
end

trafo = exp
domain   = CV_Math2DCanvas(-4.0 + 4.0im, 4.0 - 4.0im, 50)
codomain = CV_Math2DCanvas(-4.0 + 4.0im, 4.0 - 4.0im, 50)

scene = create_scene(trafo, domain, codomain)

cv_save_image(cv_get_can_layout(scene), "./Exp02.png")

if CVG !== nothing
    handler = CVG.cvg_visualize(scene)
    CVG.cvg_wait_for_destroy(handler.window)

    CVG.cvg_close(handler); CVG.cv_destroy(scene)
end

# vim:syn=julia:cc=79:fdm=marker:sw=4:ts=4:


using ForwardDiff: gradient!

export update!, optimize!, optimize

"""
    update!(x, x̄)

Updates the array `x` according to `x .-= x̄`. `x` is of type `AbstractArray`
"""
function update!(x::Vector{T}, x̄) where {T<:AbstractFloat}
    x .-= x̄
end

"""
    update!(opt, x, g)

Updates array `x` according to gradient `g`
with respect to optimizer `opt`.
"""
function update!(opt, x, g)
    x .-= apply!(opt, x, g)
end

"""
    optimize!(x, opt, f, iter)

Uses optimizer `opt` to update `x` to try and minimize `f(x)` in `iter` iterations. This is the mutating version, which changes the input vector.
"""
function optimize!(
    x::Vector{T}, f::Function, opt; iter::Int=10, chunk::Union{Int, Nothing}=2) where {T <: AbstractFloat}
    grad = similar(x)
    cfg = grad_cfg(f,x ; chunk=chunk)
    for i in 1:iter
        gradient!(grad, f, x, cfg)
        update!(opt, x, grad)
    end
end

function optimize(x::Vector{T}, f::Function, opt; iter::Int=10, chunk::Union{Int, Nothing}=2) where {T <: AbstractFloat}
   y =copy(x)
   optimize!(y, f, opt; iter=iter, chunk=chunk)
   return y
end

function optimize!(A::Matrix{T}, x::Vector{T}, f::Function, opt; iter::Int=10, chunk::Union{Int, Nothing}=2) where {T <: AbstractFloat}
    if size(A)!=(size(x)[1]+1, iter)
        throw(ArgumentError("size of the matrix is not $(size(x)[1]+1)x$(iter)"))
    end
    grad = similar(x)
    cfg = grad_cfg(f, x; chunk=chunk)
    A[:, 1] = x
    for i in 1:iter
        A[:, i+1] = @view A[:, i]
        gradient!(grad, f, A[:, i+1], cfg)
        update!(opt, A[:, i+1], grad)
    end
end


################################################################################
#VERSION 1
#linesearch is an intermediary between NewtonStopLS and a 1D solver.
function preparelinesearch(stp  :: NLPStopping,
                           x₀   :: AbstractVector,
                           d    :: AbstractVector,
                           f₀   :: AbstractFloat,
                           ∇f₀  :: AbstractVector,
                           τ₀   :: Real,
                           τ₁   :: Real)

    # initialize the LSModel
    g₀ = dot(∇f₀,d)
    ϕ = LSModel(stp.pb, x₀, d, τ₀, f₀, g₀ )

    # convert the Armijo and Wolfe criteria to an asymetric interval [α,β]
    α =  (τ₁-τ₀)*g₀
    β = -(τ₁+τ₀)*g₀

    # reuse the stopping
    ϕstp = LS_Stopping(ϕ, LSAtT(0.0, ht = f₀, gt = ∇f₀),
                       optimality_check = (p,s) -> optim_check_LS(p,s,α,β),
                       main_stp = stp,
                       max_iter = 40,
                       atol = 0.0, # to rely only on the Armijo-Wolfe conditions
                       rtol = 0.0, # otherwise, tolerance may prohibit convergence
                       unbounded_threshold = 1e100)

    # optimize in the interval [0.0,Inf]
    ϕstp = min_1D(ϕ, 0.0, Inf, α, β, stp = ϕstp)

    # unpack the results
    if ϕstp.meta.optimal
        t, ft, gt = ϕstp.current_state.x, ϕstp.current_state.ht, ϕstp.current_state.gt
    else
        t, ft, gt = 0.0, f₀, ∇f₀
        stp.meta.fail_sub_pb =  true
    end

    return t, ft, gt
end

function optim_check_LS(p, s :: LSAtT, α :: Float64, β :: Float64)
    return max(α - s.gt, s.gt - β, 0.0)
end

################################################################################
#VERSION 2
#Not an equivalent way would be with tol_check:
function preparelinesearch(stp  :: NLPStopping,
                           x₀   :: AbstractVector,
                           d    :: AbstractVector,
                           f₀   :: AbstractFloat,
                           ∇f₀  :: AbstractVector,
                           τ₀   :: Real,
                           τ₁   :: Real)

    # initialize the LSModel
    g₀ = dot(∇f₀,d)
    ϕ = LSModel(stp.pb, x₀, d, τ₀, f₀, g₀ )
    #Instead of redefining LSModel we can use the known ADNLPModel ? (not optimal though)
    #ϕ = ADNLPModel(t -> (obj(stp.pb, x₀ + t*d) - f₀ - τ₀ * t * g₀), [0.0], lvar = [0.0], uvar = [Inf])

    # convert the Armijo and Wolfe criteria to an asymetric interval [α,β]
    α =  (τ₁-τ₀)*g₀
    β = -(τ₁+τ₀)*g₀

    # reuse the stopping
    ϕstp = NLPStopping(ϕ, NLPAtX([0.0], fx = f₀, gx = ∇f₀),
                       optimality_check = unconstrained_check,
                       main_stp = stp,
                       max_iter = 40,
                       unbounded_threshold = 1e100,
                       tol_check = (a,b,c)     -> β,
                       tol_check_neg = (a,b,c) -> α)

    # optimize in the interval [0.0,Inf]
    ϕstp = min_1D(ϕ, 0.0, Inf, stp = ϕstp)

    # unpack the results
    if ϕstp.meta.optimal
        t, ft, gt = ϕstp.current_state.x, ϕstp.current_state.ht, ϕstp.current_state.gt
    else
        t, ft, gt = 0.0, f₀, ∇f₀
        stp.meta.fail_sub_pb =  true
    end

    return t, ft, gt
end


a = Set([1, 3, 4, 5]);
println(a)

# BitSet: Construct a sorted set of Ints generated by the given iterable object, or an empty set. Implemented as a bit string, and therefore designed for dense integer sets.
b = BitSet([1, 7, 4, 5]);
println(b)

# union
println(union([1, 2], [2, 4]))

# union!
# Construct the union of passed in sets and overwrite s with the result. Maintain order with arrays.
c = Set([1, 3, 4, 5]);
union!(c, 1:2:8);
println(c)

# intersect
# Construct the intersection of sets. Maintain order with arrays.
println(intersect([1, 2, 3], [3, 4, 5]))

# setdiff
println(setdiff([1,2,3], [3,4,5]))

# symdiff
# Construct the symmetric difference of elements in the passed in sets. When s is not an AbstractSet, the order is maintained. Note that in this case the multiplicity of elements matters.
println(symdiff([1,2,3], [3,4,5], [4,5,6]))
println([1,2,1], [2, 1, 2])

# issubset
println(issubset([1, 2], [1, 2, 3]))

# ⊈
println((1, 2) ⊈ (2, 3))


# plot kernek density estimation

module MarginalKde200
    using DataFrames, CSV, StatsPlots
    pyplot()

    function main(is_test=false)
        # input
        data_path = joinpath(split(@__FILE__, "src")[1], "data/sensor_data_200.txt")
        df = CSV.read(data_path, DataFrame,
                      header=["date", "time", "ir", "lidar"],
                      delim=' ')

        # plot marginal kde
        marginalkde(df.ir, df.lidar, c=:ice,
                    xlabel="ir", ylabel="liDAR")

        if is_test == false
            save_path = joinpath(split(@__FILE__, "src")[1], "src/prob_stats/multi_dim_gauss_dist/marginal_kde/marginal_kde_200.png")
            savefig(save_path)
        end
    end
end


#=
testODE:
- Julia version:
- Author: ymocquar
- Date: 2019-11-13
=#
include("../src/twoscales_pure_ab.jl")
using DifferentialEquations
using LinearAlgebra
using Plots
using Random

function twoscales_solve(par_u0::PrepareU0, order, t, nb)

    parGen = PrepareTwoScalePureAB(nb, t, order, par_u0)

    return twoscales_pure_ab(parGen, only_end = true)

end

function getmindif(tab::Vector{Vector{BigFloat}})
    nmmin = Inf
    ret = (1, 2)
    for i = 1:size(tab, 1)
        for j = (i+1):size(tab, 1)
            nm = norm(tab[i] - tab[j], Inf)
            if nm < nmmin
                nmmin = nm
                ret = i, j
            end
        end
    end
    return ret, min(nmmin, 1.1)
end

function fctMain(n_tau)

    #   u0 =[big"0.12345678", big"0.1209182736", big"0.1290582671", big"0.1239681094" ]
    seed = 12996
    Random.seed!(seed)
    u0 = rand(BigFloat, 4)
    println("seed = $seed")
    #   tab_eps = zeros(BigFloat,7)
    #   tab_eps= [big"0.5", big"0.2", big"0.1",big"0.05", big"0.02", big"0.01",big"0.005", big"0.002", big"0.001",big"0.0005", big"0.0002", big"0.0001"]
    tab_eps = [
        big"0.1",
        big"0.01",
        big"0.001",
        big"0.0001",
        big"0.00001",
        big"0.000001",
        big"0.0000001",
        big"0.00000001",
        big"0.000000001",
    ]
    # epsilon=big"0.1"
    # for i=1:7
    #     tab_eps[i] = epsilon
    #     epsilon /= 100
    # end
    nbmaxtest = 11
    t_max = big"1.0"
    A = [0 0 1 0; 0 0 0 0; -1 0 0 0; 0 0 0 0]
    tabtabsol = Vector{Vector{Vector{BigFloat}}}()
    for order = 2:14
        ordprep = order + 2
        y = ones(Float64, nbmaxtest, size(tab_eps, 1))
        y_big = ones(BigFloat, nbmaxtest, size(tab_eps, 1))
        x = zeros(Float64, nbmaxtest)
        ind = 1
        for epsilon in tab_eps
            #        fct = u -> [ 1.11u[2]^2+0.8247u[4]^3+0.12647u[3]^4,
            #       0.4356789141u[1]-0.87u[3]*u[4],
            #       1.9898985/(1+1.1237u[4]^3), 0.97u[1]*u[3]+0.8111u[2]^2 ]
            fct =
                (u, p, t) -> [
                    u[3]^3 + sin(t^2),
                    u[4] + u[1]^2 + 1 / (1 + exp(t)) - 0.35,
                    u[1] * u[2] * u[4] * (1 + t),
                    u[2] - u[2]^2 + 1 / (1 + u[4]^2) + u[1]^2 + cos(exp(t)),
                ]
            parphi = PreparePhi(epsilon, n_tau, A, fct)
            println("prepareU0 eps=$epsilon n_tau=$n_tau")
            nb = 10
            @time par_u0 = PrepareU0(parphi, ordprep, u0)
            @time pargen = PrepareTwoScalesPureAB(nb * 2^nbmaxtest, t_max, order, par_u0)
            @time solref = twoscales_pure_ab(pargen, only_end = true)
            if size(tabtabsol, 1) < ind
                push!(tabtabsol, Vector{Vector{BigFloat}}())
            end
            tabsol = tabtabsol[ind]
            push!(tabsol, solref)
            res_gen = Array{Array{BigFloat,1},1}(undef, nbmaxtest)
            indref = 1
            eps_v = convert(Float32, epsilon)
            println("epsilon = $eps_v solref=$solref")
            indc = 1
            labels = Array{String,2}(undef, 1, ind)
            while indc <= nbmaxtest
                @time pargen = PrepareTwoScalesPureAB(nb, t_max, order, par_u0)
                @time sol = twoscales_pure_ab(pargen, only_end = true)
                push!(tabsol, sol)
                res_gen[indc] = sol
                diff = solref - sol
                println("tabsol=$tabsol")
                (a, b), nm = getmindif(tabsol)
                if a != indref
                    println("New solref !!!! a=$a, b=$b nm=$nm")
                    indref = a
                    solref = tabsol[a]
                    for i = 1:indc
                        nm2 = min(norm(res_gen[i] - solref, Inf), 1.1)
                        y[i, ind] = nm2 == 0 ? nm : nm2
                        y_big[i, ind] = nm2 == 0 ? nm : nm2
                    end
                else
                    diff = solref - sol
                    y[indc, ind] = min(norm(diff, Inf), 1.1)
                    y_big[indc, ind] = min(norm(diff, Inf), 1.1)
                end
                println("solref=$solref")
                println("nb=$nb sol=$sol")
                x[indc] = 1.0 / nb
                println("epsilon=$epsilon result=$y")
                println("epsilon=$epsilon reslog2=$(log2.(y))")
                println("epsilon=$epsilon result=$y_big")
                println("epsilon=$epsilon reslog2=$(log2.(y_big))")
                nb *= 2
                indc += 1
            end
            if ind == size(tab_eps, 1)
                for i = 1:ind
                    labels[1, i] = " epsilon,order=$(convert(Float32,tab_eps[i])),$order "
                end
                gr()
                p = Plots.plot(
                    x,
                    view(y, :, 1:ind),
                    xlabel = "delta t",
                    xaxis = :log,
                    ylabel = "error",
                    yaxis = :log,
                    legend = :bottomright,
                    label = labels,
                    marker = 2,
                )

                prec_v = precision(BigFloat)
                Plots.savefig(
                    p,
                    "out/r20_$(prec_v)_$(eps_v)_$(order)_$(ordprep)_$(n_tau)_epsilon_fct.pdf",
                )
            end
            ind += 1
        end
    end
end

# testODESolverEps()

# for i=3:9
#     fctMain(2^i)
# end
# setprecision(512)
fctMain(32)


using Plots, StatsBase, Statistics, LaTeXStrings, JLD, ProgressMeter, LsqFit

function InitialNetwork(dim, P)
    Network = sample([-1,0], Weights([1-P, P]),(dim,dim))
    return Network
end

function FindNeighbors(i, j, Network, dim)
    Numbereds = 0
    colorList = []
    for neighbor in ([0,-1], [-1,0])
        if j + neighbor[2] >= 1 && i + neighbor[1] >= 1 && Network[([i, j] + neighbor)...] > 0
            push!(colorList,Network[([i, j] + neighbor)...])
            Numbereds += 1
        end
    end
    return colorList, Numbereds #colorList contain the node and its neighbors
end

function RecursionMapping(x, L)
    while L[x] != L[L[x]]
        x = L[x]
    end
    return L[x]
end

@inline function HKNetworkDynamic(dim, P)
    Network = InitialNetwork(dim, P)
    S = []
    L = []
    k = 1

    for j in 1:dim
        for i in 1:dim
            if Network[i,j] != 0
                continue
            end
            NColor, NNum = FindNeighbors(i,j,Network,dim)
            if NNum == 0
                Network[i,j] = k
                push!(L, k)
                push!(S, 1)
                k += 1
            elseif NNum == 1
                Network[i,j] = RecursionMapping(NColor[1],L)
                S[RecursionMapping(NColor[1],L)] += 1
            elseif NNum == 2 && RecursionMapping(NColor[1],L) == RecursionMapping(NColor[2],L)
                Network[i,j] = RecursionMapping(NColor[1],L)
                S[RecursionMapping(NColor[1],L)] += 1
            else
                S[RecursionMapping(NColor[1],L)] += 1 + S[RecursionMapping(NColor[2],L)]
                S[RecursionMapping(NColor[2],L)] = 0
                Network[i,j] = RecursionMapping(NColor[1],L)
                L[NColor[2]] = RecursionMapping(NColor[1],L)
            end
        end
    end
    return Network, S, L
end

function JoinColors(Network, L)
    dim = size(Network)[1]
    for i in 1:dim
        for j in 1:dim
            if Network[i,j]!= -1
                Network[i,j] = RecursionMapping(Network[i,j],L)
            end
        end
    end
    return Network
end

function FindCorrelationLength(Network, S)
    if length(Set(S)) <= 2
        return 0.0
    end
    if findfirst(x->x==max(S...),S) ∈ intersect(Network[1,:],Network[end,:])
        BiggestFinite = findfirst(x->x==sort(S, rev=true)[2],S)
    else
        BiggestFinite = findfirst(x->x==max(S...),S)
    end
    BFiniteCluster = Tuple.(findall(x->x==BiggestFinite,Network))
    BFiniteClusterSize = length(BFiniteCluster)
    MassCenter = reduce(.+, BFiniteCluster) ./ BFiniteClusterSize
    return √(sum(pos -> reduce(.+, (pos .- MassCenter).^2),Tuple.(BFiniteCluster))/BFiniteClusterSize)
end

function ReturnXi(dim, p)
    Network, S, L = HKNetworkDynamic(dim, p)
    Network = JoinColors(Network, L)
    return  FindCorrelationLength(Network, S)
end

dimlist = [10,15,20,25,30,35,40,45,50,55,60,65,70,75,100,125,150]
# runlist = [10000,5000,2000,1000,500,200,100,100,100]
PList = hcat(0.52:0.001:0.62)
runnum = 100
progress = Progress(length(PList)*runnum*length(dimlist); showspeed=true)
# progress = Progress(length(PList)*sum(runlist); showspeed=true)
for n in 1:length(dimlist)
    dim = dimlist[n]
    AvgData = []
    STDData = []
    TopAvg = []
    for p in PList
        totalruns = []
        for i in 1:runnum
            push!(totalruns, ReturnXi(dim, p))
            next!(progress)
            update!(progress)
        end
        push!(TopAvg,mean(sort(totalruns,rev = true)[1:Int(runnum/5)]))
        push!(AvgData,mean(totalruns))
        push!(STDData,std(totalruns))
    end
    save("../../Data/Q2/Q2-$dim-fc.jld", "TopAvg", TopAvg, "AvgData", AvgData, "STDData", STDData)
end

Data = zeros(length(dimlist),length(PList),3)

for n in 1:length(dimlist)
    print("\r$n")
    dim = dimlist[n]
    data = load("../../Data/Q2/Q2-$dim-fc.jld")
    Data[n,:,1] = data["TopAvg"]
    Data[n,:,2] = data["AvgData"]
    Data[n,:,3] = data["STDData"]
end

PmaxList = []
for n in 1:length(dimlist)
    push!(PmaxList, PList[findfirst(x->x==max(Data[n,:,2]...),Data[n,:,2])])
end

@. modelP∞(x, p) = abs(x - p[1])^(-1.2)
@. modelν(x, p) = abs(x - 0.5927)^(-p[1])

xdata = PmaxList
ydata = dimlist
InitP∞ = [0.59]
Initν = [1.2]

ν = curve_fit(modelν, xdata, ydata, Initν).param[1]
P∞ = curve_fit(modelP∞, xdata, ydata, InitP∞).param[1]

P∞range = hcat(0.52:0.001:0.596)
νrange = abs.(P∞range .- P∞).^(-ν)

scatter(PmaxList,dimlist, label=L"Data\ Points", title=L"Curve\ fitting\ Plot\ (\nu = %$(round(ν,digits= 2)) ,\ Pc_{\infty} = %$(round(P∞,digits= 2)))",legend = 150)
plot!(P∞range,νrange,c= :purple,lw = 2,label = L"L=|Pc_{(L)} - Pc_{\infty}|^{-\nu}")
plot!([0.6,0.6],[0.0,200.0], linestyle = :dash,c = :black,lw = 2,
    label = L"X=Pc_{\infty}= %$(round(P∞,digits= 2))")
savefig("../../Figs/Q2/Q2-CF-res.pdf")


export KooshballTrajectory, kooshballNodes, kooshballDensity, kooshBallTimes

"""
    KooshballTrajectory(numProfiles, numSamplingPerProfile
                  ; TE::Float64=0.0
                  , AQ::Float64=1.e-3
                  , kargs...)

returns a 3d kooshball trajectory.

# Arguments
* `numProfiles::Int64`            - number of profiles
* `numSamplingPerProfile::Int64`  - number of sampling points per profile
* (`TE::Float64=0.0`)             - echo time in s
* (`AQ::Float64=1.e-3`)           - readout duration in s (per profile)
"""
function KooshballTrajectory(::Type{T}, numProfiles, numSamplingPerProfile
                  ; TE=0.0
                  , AQ=1.e-3
                  , kargs...) where T
  nodes = kooshballNodes(T, numProfiles, numSamplingPerProfile; numSlices=numSlices)
  times = kooshballTimes(T, numProfiles, numSamplingPerProfile; TE=TE, AQ=AQ)
  return  Trajectory("Kooshball", nodes, times, TE, AQ, numProfiles, numSamplingPerProfile, 1, false, true)
end

function kooshballNodes(::Type{T}, numProfiles::Int64, numSamplingPerProfile::Int64; kargs...) where T
  nodes = zeros(T, 3, numSamplingPerProfile, numProfiles)
  pos = collect((0:numSamplingPerProfile-1)/numSamplingPerProfile .- 0.5)
  for l = 1:numProfiles
    for k = 1:numSamplingPerProfile
      nodes[1,k,l] = (-1)^l * pos[k]*xax(l,2*numProfiles)
      nodes[2,k,l] = (-1)^l * pos[k]*yax(l,2*numProfiles)
      nodes[3,k,l] = (-1)^l * pos[k]*zax(l,2*numProfiles)
    end
  end
  return reshape(nodes, 3, numSamplingPerProfile*numProfiles)
end

function kooshballTimes(::Type{T}, numProfiles::Int64, numSamplingPerProfile::Int64; TE=0.0, AQ=1.e-3, kargs...) where T
  times = zeros(T, numSamplingPerProfile, numProfiles)
  for l = 1:numProfiles
    for k = 1:numSamplingPerProfile
      times[k,l] = (TE + AQ*(k-1)/numSamplingPerProfile )
    end
  end
  return return vec(times)
end

function kooshballDensity(numSamplingPerProfile::Int64, numProfiles::Int64, numSlices::Int64)
  density = zeros(numSamplingPerProfile, numProfiles)
  pos = collect((0:numSamplingPerProfile-1)/numSamplingPerProfile .- 0.5)
  for l = 1:numProfiles
    for k = 1:numSamplingPerProfile
      density[k,l] = abs(pos[k])
    end
  end
  return vec(density)
end


# Some Boiler-plate code for preconditioning
#
# Meaning of P:
#    P ≈ ∇²E, so the preconditioned gradient is P^{-1} ∇E
#    P can be an arbitrary type but at a minimum it MUST provide:
#        ldiv!(x, P, b) ->  x = P \ b
#        dot(x, P, b)  ->  x' P b
#
#    If `dot` is not provided, then Optim.jl will try to define it via
#        dot(x, P, b) = dot(x, mul!(similar(x), P, y))
#
#    finally the preconditioner can be updated after each x-update using
#        precondprep!
#    but this is passed as an argument at the moment!
#

#####################################################
#  [0] Defaults and aliases for easier reading of the code
#      these can also be over-written if necessary.

if VERSION < v"1.4.0-DEV.92"
# an inner product w.r.t. a metric P (=preconditioner)
    dot(x, P, y) = dot(x, mul!(similar(x), P, y))
end

# default preconditioner update
precondprep!(P, x) = nothing

#####################################################
#  [1] Empty preconditioner = Identity
#

# out =  P^{-1} * A
ldiv!(out, ::Nothing, A) = copyto!(out, A)
# A' * P B
dot(A, ::Nothing, B) = dot(A, B)

#####################################################
#  [2] Diagonal preconditioner
#      P = Diag(d)
#      unfortunately, Base does not implement
#      ldiv!(a, P, b) or mul! for this type, so we do it by hand
#      TODO: maybe implement this in Base
ldiv!(out::Array, P::Diagonal, A::Array) = copyto!(out, A ./ P.diag)
dot(A::Array, P::Diagonal, B::Array) = dot(A, P.diag .* B)

#####################################################
#  [3] Inverse Diagonal preconditioner
#      here, P is stored by the entries of its inverse
#      TODO: maybe implement this in Base?
mutable struct InverseDiagonal
   diag
end
ldiv!(out::Array, P::InverseDiagonal, A::Array) = copyto!(out, A .* P.diag)
dot(A::Array, P::InverseDiagonal, B::Vector) = dot(A, B ./ P.diag)

#####################################################
#  [4] Matrix Preconditioner
#  the assumption here is that P is given by its inverse, which is typical
#     > ldiv! is about to be moved to Base, so we need a temporary hack
#     > mul! is already in Base, which defines `dot`
#  nothing to do!
ldiv!(x, P::AbstractMatrix, b) = copyto!(x, P \ b)


using Test, EllZeroTrendFiltering

zero_deviation(x) = 0
abstract type FunctionDeviation end

struct ZeroDeviation <: FunctionDeviation end
(f::ZeroDeviation)(x) = 0
deviationcost(f::ZeroDeviation) = 0

# Zig-zag function
#
struct TriangleDeviation <: FunctionDeviation
    y1::Float64
    y2::Float64
    x1::Float64
    x2::Float64
end
function (f::TriangleDeviation)(xinput)
    # Relative x in range 0 to 1
    x = (xinput-f.x1)/(f.x2-f.x1)
    # Relative height
    y = (x < 0.3 ? 0 :
        (x < 0.35 ? (x-0.3)*2 :
        (x < 0.45 ?  0.1 - (x-0.35)*2 :
        (x < 0.5  ? -0.1 + (x-0.45)*2 :
        (x < 0.55 ?  0.0 - (x-0.5 )*2 :
        (x < 0.65 ? -0.1 + (x-0.55)*2 :
        (x < 0.7  ?  0.1 - (x-0.65)*2 :
         0)))))))
    return y*(f.y2 - f.y1)
end
deviationcost(f::TriangleDeviation) = 8/3*(0.05^3*2^2*(f.y2 - f.y1)^2*(f.x2 - f.x1))

function evaluatefunction(Xsol,Ysol,deviations,x)
    firsti = findfirst(v -> v >= x, Xsol)
    i2 = (firsti == nothing) ? 2 : firsti
    i1 = i2-1
    linearval = Ysol[i1] + (x-Xsol[i1])*((Ysol[i2]-Ysol[i1])/(Xsol[i2]-Xsol[i1]))
    return  linearval + deviations[i1](x)
end

function testfunction(Xsol, Ysol, intermediate)
    # Alternate Zero and TriangleDeviations
    deviations = vcat([
                  vcat([[ZeroDeviation(), TriangleDeviation(Ysol[i+1],Ysol[i+2],Xsol[i+1],Xsol[i+2])] for i = 1:2:(length(Xsol)-3)]...),
                  [ZeroDeviation() for i = 1:mod(length(Xsol),2)+1] # One or two zeros at end
                ]...)
    #Total cost of deviations if  line goes through middle
    extra_cost = sum(deviationcost.(deviations))
    f = x -> evaluatefunction(Xsol,Ysol, deviations, x)
    t = [0.]
    # Add intermediate number of points on x grid, exactly hitting Xsol
    for i = 2:(segments+1)
        append!(t, range(Xsol[i-1], stop= Xsol[i], length=intermediate+2)[2:end])
    end
    return f, t, extra_cost
end

Random.seed!(12345)

function random_problem(segments)
    Msol = segments - 1
    intermediate = 19

    # X and Y points of correct solution
    Ysol = randn(segments+1)
    Xsol = [0,cumsum(abs.(randn(segments)).+ 1.0)...]

    f, t, extra_cost = testfunction(Xsol, Ysol, intermediate)
    Ysol, Xsol, f, t, extra_cost, intermediate
end

segments = 5
@testset "Exact solution segments = $segments, i=$i" for i = 1:10
    local t  # TODO, can remove in julia 1.0?
    Ysol, Xsol, f, t, extra_cost, intermediate = random_problem(segments)

    I, Y, v = fit_pwl_constrained(f, t, segments, 1e-12, lazy=true)
    # Test that we find exactly expected solution, should almost always be true for small mbr of segments
    @test all(I[segments] .== 1:(intermediate+1):(segments*(intermediate+1)+1))
    @test norm(Ysol.-Y[end]) < 1e-13
    @test norm(extra_cost - v[end]) < 1e-13
end

segments = 20
@testset "Exact solution segments = $segments, i=$i" for i = 1:10
    local t  # TODO, can remove in julia 1.0?
    Ysol, Xsol, f, t, extra_cost = random_problem(segments)

    I, Y, v = fit_pwl_constrained(f, t, segments, 1e-12, lazy=true)
    # Test that we find same or better solution
    @test extra_cost > v[end] - 1e-12
end


# # Comparison to Neumann analytic solution

# Before reading this tutorial,
# we recommend that you look over the coupled energy
# and water [tutorial](../Coupled/equilibrium_test.md)
# and the freezing front [tutorial](freezing_front.md).
# The former shows how to solve the heat equation for soil volumetric
# internal energy `ρe_int` simultaneously
# with Richards equation for volumetric liquid water fraction `ϑ_l`, assuming zero
# volumetric ice fraction `θ_i` for all time, everywhere in the domain[^a].
# The latter shows how to include freezing and thawing, and explains the freeze thaw model employed
# by CliMA Land. This tutorial compares a simulated temperature profile
# from a freezing front with an analytic solution.

# The analytic solution applies to a freezing front propagating under certain assumptions. It assumes
# that there is no water movement, which we mimic by setting `K_sat=0`. It also assumes a semi-infinite
# domain, which we approximate by making the domain larger than the extent to which the front propagates.
# The solution also assumes that all water freezes. Our model does not satisfy that assumption, as discussed
# [here](freezing_front.md#Discussion-and-Model-Explanation), but as we will see, the model still matches the
# analytic expectation well in the frozen region.

# As such, our set of equations is

# ``
# \frac{∂ ρe_{int}}{∂ t} =  ∇ ⋅ κ(θ_l, θ_i; ν, ...) ∇T
# ``

# ``
# \frac{ ∂ ϑ_l}{∂ t} = -\frac{F_T}{ρ_l}
# ``

# ``
# \frac{ ∂ θ_i}{∂ t} = \frac{F_T}{ρ_i}
# ``

# Here

# ``t`` is the time (s),

# ``z`` is the location in the vertical (m),

# ``ρe_{int}`` is the volumetric internal energy of the soil (J/m^3),

# ``T`` is the temperature of the soil (K),

# ``κ`` is the thermal conductivity (W/m/K),

# ``ϑ_l`` is the augmented volumetric liquid water fraction,

# ``θ_i`` is the volumetric ice fraction,

# ``ν, ...`` denotes parameters relating to soil type, such as porosity, and

# ``F_T`` is the freeze-thaw term.

# Our domain is effectively 1-d, with ``z ∈ [-3,0]``,
# and with the following boundary and initial conditions:

# `` T(t, z=0) = 263.15 K``

# ``- κ ∇T(t, z= -3) =  0 ẑ ``

# `` T(t = 0, z) = 275.15 K``

# ``- K ∇h(t, z = 0) = 0 ẑ ``

# `` -K ∇h(t, z = -3) = 0 ẑ``

# `` ϑ(t = 0, z) = 0.33 ``.

# # Import necessary modules
# External (non - CliMA) modules
using MPI
using OrderedCollections
using StaticArrays
using Statistics
using Test
using DelimitedFiles
using Plots

# CliMA Parameters
using CLIMAParameters
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet()
using CLIMAParameters.Planet: ρ_cloud_liq
using CLIMAParameters.Planet: ρ_cloud_ice
using CLIMAParameters.Planet: LH_f0

# ClimateMachine modules
using ClimateMachine
using ClimateMachine.Land
using ClimateMachine.Land.SoilWaterParameterizations
using ClimateMachine.Land.SoilHeatParameterizations
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.DGMethods
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.DGMethods: BalanceLaw, LocalGeometry
using ClimateMachine.MPIStateArrays
using ClimateMachine.GenericCallbacks
using ClimateMachine.SystemSolvers
using ClimateMachine.ODESolvers
using ClimateMachine.VariableTemplates
using ClimateMachine.SingleStackUtils
using ClimateMachine.BalanceLaws:
    BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state
using SpecialFunctions
using ArtifactWrappers

# # Preliminary set-up

# Get the parameter set, which holds constants used across CliMA models:
struct EarthParameterSet <: AbstractEarthParameterSet end
const param_set = EarthParameterSet();
# Initialize and pick a floating point precision:
ClimateMachine.init()
const FT = Float64;

# # Simulation specific parameters
N_poly = 1
nelem_vert = 40
zmax = FT(0)
zmin = FT(-3)
t0 = FT(0)
dt = FT(50)
timeend = FT(3600 * 24 * 20)
n_outputs = 540
every_x_simulation_time = ceil(Int, timeend / n_outputs)
Δ = abs(zmin - zmax) / FT(nelem_vert);

# # Soil properties
# All units are mks.
porosity = FT(0.535)
vg_α = 1.11
vg_n = 1.48
Ksat = 0.0
ν_ss_quartz = FT(0.2)
ν_ss_minerals = FT(0.6)
ν_ss_om = FT(0.2)
ν_ss_gravel = FT(0.0);
κ_quartz = FT(7.7)
κ_minerals = FT(2.5)
κ_om = FT(0.25)
κ_liq = FT(0.57)
κ_ice = FT(2.29)
ρp = FT(2700)
κ_solid = k_solid(ν_ss_om, ν_ss_quartz, κ_quartz, κ_minerals, κ_om)
κ_sat_frozen = ksat_frozen(κ_solid, porosity, κ_ice)
κ_sat_unfrozen = ksat_unfrozen(κ_solid, porosity, κ_liq)
ρc_ds = FT((1 - porosity) * 2.3e6)
soil_param_functions = SoilParamFunctions(
    FT;
    porosity = porosity,
    ν_ss_gravel = ν_ss_gravel,
    ν_ss_om = ν_ss_om,
    ν_ss_quartz = ν_ss_quartz,
    ρc_ds = ρc_ds,
    ρp = ρp,
    κ_solid = κ_solid,
    κ_sat_unfrozen = κ_sat_unfrozen,
    κ_sat_frozen = κ_sat_frozen,
    water = WaterParamFunctions(FT; Ksat = Ksat, S_s = 1e-3),
);

# # Build the model
# Initial and Boundary conditions. The default initial condition for
# `θ_i` is zero everywhere, so we don't modify that. Furthermore, since
# the equation for `θ_i` does not involve spatial derivatives, we don't need
# to supply boundary conditions for it.
zero_flux = (aux, t) -> eltype(aux)(0.0)
ϑ_l0 = (aux) -> eltype(aux)(0.33)
surface_state = (aux, t) -> eltype(aux)(273.15 - 10.0)
T_init = aux -> eltype(aux)(275.15)
bc = LandDomainBC(
    bottom_bc = LandComponentBC(
        soil_heat = Neumann(zero_flux),
        soil_water = Neumann(zero_flux),
    ),
    surface_bc = LandComponentBC(
        soil_heat = Dirichlet(surface_state),
        soil_water = Neumann(zero_flux),
    ),
);

# Create the [`SoilWaterModel`](@ref ClimateMachine.Land.SoilWaterModel),
# [`SoilHeatModel`](@ref ClimateMachine.Land.SoilHeatModel),
# and the [`SoilModel`](@ref ClimateMachine.Land.SoilModel) instances.
# Note that we are still specifying a hydraulics model, because the matric potential
# and hydraulic conductivity functions are still evaluated (though they don't affect
# the outcome). Setting `Ksat =0` is just a
# hack for turning off water flow.
soil_water_model = SoilWaterModel(
    FT;
    hydraulics = vanGenuchten(FT; α = vg_α, n = vg_n),
    initialϑ_l = ϑ_l0,
);

soil_heat_model = SoilHeatModel(FT; initialT = T_init);

m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model);
# Create the source term instance. Our phase change model requires
# knowledge of the vertical spacing, so we pass
# that information in via an attribute of the
# [`PhaseChange`](@ref ClimateMachine.Land.PhaseChange) structure.
freeze_thaw_source = PhaseChange{FT}(Δz = Δ);

# Sources are added as elements of a list of sources. Here we just add freezing
# and thawing.
sources = (freeze_thaw_source,);

# Next, we define the required `init_soil!` function, which takes the user
# specified functions of space for `T_init` and `ϑ_l0` and initializes the state
# variables of volumetric internal energy and augmented liquid fraction. This requires
# a conversion from `T` to `ρe_int`.
function init_soil!(land, state, aux, localgeo, time)
    myFT = eltype(state)
    ϑ_l = myFT(land.soil.water.initialϑ_l(aux))
    θ_i = myFT(land.soil.water.initialθ_i(aux))
    state.soil.water.ϑ_l = ϑ_l
    state.soil.water.θ_i = θ_i
    param_set = land.param_set

    θ_l = volumetric_liquid_fraction(ϑ_l, land.soil.param_functions.porosity)
    ρc_ds = land.soil.param_functions.ρc_ds
    ρc_s = volumetric_heat_capacity(θ_l, θ_i, ρc_ds, param_set)

    state.soil.heat.ρe_int = volumetric_internal_energy(
        θ_i,
        ρc_s,
        land.soil.heat.initialT(aux),
        param_set,
    )
end;

# Lastly, package it all up in the `LandModel`:
m = LandModel(
    param_set,
    m_soil;
    boundary_conditions = bc,
    source = sources,
    init_state_prognostic = init_soil!,
);

# # Set up and run the simulation

driver_config = ClimateMachine.SingleStackConfiguration(
    "LandModel",
    N_poly,
    nelem_vert,
    zmax,
    param_set,
    m;
    zmin = zmin,
    numerical_flux_first_order = CentralNumericalFluxFirstOrder(),
);

solver_config =
    ClimateMachine.SolverConfiguration(t0, timeend, driver_config, ode_dt = dt);

state_types = (Prognostic(), Auxiliary(), GradientFlux())
all_data = Dict[dict_of_nodal_states(solver_config, state_types; interp = true)]
time_data = FT[0]
callback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do
    dons = dict_of_nodal_states(solver_config, state_types; interp = true)
    push!(all_data, dons)
    push!(time_data, gettime(solver_config.solver))
    nothing
end;

ClimateMachine.invoke!(solver_config; user_callbacks = (callback,));
z = get_z(solver_config.dg.grid; rm_dupes = true);

# # Analytic Solution of Neumann
# All details here are taken from [DallAmico2011](@cite) (see also [CarslawJaeger](@cite)), and the reader is referred to that
# for further information on the solution. It takes the form of a function
# for T(z) on each side of the freezing front interface, which depends on
# the thermal properties in that region, and which is also parameterized by
# a parameter (ζ), which we show how to solve for below. In computing the
# thermal properties, we evaluate the conductivity and heat capacity
# assuming that all of the water is either in liquid or frozen form, with total
# mass proportional to ``θ_{l,0}ρ_l`` (as we have no water flow).

# Compute the thermal conductivity and heat capacity in the frozen region - subscript 1.
θ_l0 = FT(0.33)
kdry = k_dry(param_set, soil_param_functions)
ksat = saturated_thermal_conductivity(
    FT(0.0),
    θ_l0 * ρ_cloud_liq(param_set) / ρ_cloud_ice(param_set),
    κ_sat_unfrozen,
    κ_sat_frozen,
)
kersten = kersten_number(
    θ_l0 * ρ_cloud_liq(param_set) / ρ_cloud_ice(param_set),
    θ_l0 * ρ_cloud_liq(param_set) / ρ_cloud_ice(param_set) / porosity,
    soil_param_functions,
)
λ1 = thermal_conductivity(kdry, kersten, ksat)
c1 = volumetric_heat_capacity(
    FT(0.0),
    θ_l0 * ρ_cloud_liq(param_set) / ρ_cloud_ice(param_set),
    ρc_ds,
    param_set,
)
d1 = λ1 / c1;

# Compute the thermal conductivity and heat capacity in the region
# with liquid water - subscript 2.
ksat =
    saturated_thermal_conductivity(θ_l0, FT(0.0), κ_sat_unfrozen, κ_sat_frozen)
kersten = kersten_number(FT(0.0), θ_l0 / porosity, soil_param_functions)
λ2 = thermal_conductivity(kdry, kersten, ksat)
c2 = volumetric_heat_capacity(θ_l0, FT(0.0), ρc_ds, param_set)
d2 = λ2 / c2;

# Initial T and surface T, in Celsius
Ti = FT(2)
Ts = FT(-10.0);

# The solution requires the root of the implicit equation below
# (you'll need to install `Roots` via the package manager).
# ``` julia
# using Roots
# function implicit(ζ)
#    term1 = exp(-ζ^2) / ζ / erf(ζ)
#    term2 =
#        -λ2 * sqrt(d1) * (Ti - 0) /
#        (λ1 * sqrt(d2) * (0 - Ts) * ζ * erfc(ζ * sqrt(d1 / d2))) *
#        exp(-d1 / d2 * ζ^2)
#    term3 =
#        -LH_f0(param_set) * ρ_cloud_liq(param_set) * θ_l0 * sqrt(π) / c1 /
#        (0 - Ts)
#    return (term1 + term2 + term3)
# end
# find_zero(implicit, (0.25,0.27), Bisection())
# ```

# The root is
ζ = 0.26447353269809687;

# This function plots the analytic solution
# and the simulated result, at the output time index
# `k`.
function f(k; ζ = ζ)
    T = all_data[k]["soil.heat.T"][:] .- 273.15
    plot(
        T,
        z[:],
        xlim = [-10.5, 3],
        ylim = [zmin, zmax],
        xlabel = "T",
        label = "simulation",
    )
    t = time_data[k]
    zf = 2.0 * ζ * sqrt(d1 * t)
    myz = zmin:0.001:0
    spatially_varying =
        (erfc.(abs.(myz) ./ (zf / ζ / (d1 / d2)^0.5))) ./
        erfc(ζ * (d1 / d2)^0.5)
    mask = abs.(myz) .>= zf
    plot!(
        Ti .- (Ti - 0.0) .* spatially_varying[mask],
        myz[mask],
        label = "analytic",
        color = "green",
    )
    spatially_varying = ((erf.(abs.(myz) ./ (zf / ζ)))) ./ erf(ζ)
    mask = abs.(myz) .< zf
    plot!(
        Ts .+ (0.0 - Ts) .* spatially_varying[mask],
        myz[mask],
        label = "",
        color = "green",
    )
end

# Now we will plot this and compare to other methods for modeling phase change without water movement.
# These solutions were produced by modifying the Supplemental Program 5:3 from [Bonan19a](@cite).

# Excess heat solution:
eh_dataset = ArtifactWrapper(
    @__DIR__,
    isempty(get(ENV, "CI", "")),
    "eh",
    ArtifactFile[ArtifactFile(
        url = "https://caltech.box.com/shared/static/6xs1r98wk7u1b0xjhpkdvt80q4sh16un.csv",
        filename = "bonan_data.csv",
    ),],
);
eh_dataset_path = get_data_folder(eh_dataset);
eh_data = joinpath(eh_dataset_path, "bonan_data.csv")
ds = readdlm(eh_data, ',');
# Apparent heat capacity solution:
ahc_dataset = ArtifactWrapper(
    @__DIR__,
    isempty(get(ENV, "CI", "")),
    "ahc",
    ArtifactFile[ArtifactFile(
        url = "https://caltech.box.com/shared/static/d6xciskzl2djwi03xxfo8wu4obrptjmy.csv",
        filename = "bonan_data_ahc.csv",
    ),],
);
ahc_dataset_path = get_data_folder(ahc_dataset);
ahc_data = joinpath(ahc_dataset_path, "bonan_data_ahc.csv")
ds2 = readdlm(ahc_data, ',');

k = n_outputs
f(k)
plot!(ds[:, 2], ds[:, 1], label = "Excess Heat")
plot!(ds2[:, 2], ds2[:, 1], label = "Apparent heat capacity")
plot!(legend = :bottomleft)
savefig("analytic_comparison.png")
# ![](analytic_comparison.png)

# # References
# - [DallAmico2011](@cite)
# - [CarslawJaeger](@cite)
# - [Bonan19a](@cite)

# [^a]:
#       Note that `θ_i` is always treated as a prognostic variable
#       in the `SoilWaterModel`, but with
#       zero terms on the RHS unless freezing and thawing is turn on, as demonstrated in this
#       tutorial. That means that the user could, in principle, set the initial condition to be nonzero
#       (`θ_i(x, y, z ,t=0) = 0` is the default), which in turn would allow a nonzero `θ_i`
#       profile to affect things like thermal conductivity, etc,
#       in a consistent way. However, it would not be enforced that ``θ_l+θ_i \leq ν``, because there would
#       be no physics linking the liquid and water content to each other, and they are independent
#       variables in our model. We don't envision this being a common use case.


# Various correlation

#
# spearman correlation functions
#

# spearman correlation between two vectors
function cor_spearman(x::AbstractVector, y::AbstractVector)
    if any(isnan(x)) || any(isnan(y)) return NaN end
    return cor(tiedrank(x), tiedrank(y))
end

# spearman correlation over all pairs of columns of two matrices
function cor_spearman(X::AbstractMatrix, Y::AbstractMatrix)
    return cor(mapslices(tiedrank, X, 1), mapslices(tiedrank, Y, 1))
end
function cor_spearman(X::AbstractMatrix, y::AbstractVector)
    return cor(mapslices(tiedrank, X, 1), tiedrank(y))
end
function cor_spearman(x::AbstractVector, Y::AbstractMatrix)
    return cor(tiedrank(x), mapslices(tiedrank, Y, 1))
end

# spearman correlation over all pairs of columns of a matrix
function cor_spearman(X::AbstractMatrix)
    csp = cor(mapslices(tiedrank, X, 1))
    nanindex = vec(mapslices(any, isnan(X), 1))
    csp[nanindex, :] = NaN
    csp[:, nanindex] = NaN
    return csp
end

#
# Kendall's rank correlation
#

# Knigh JASA (1966)
function cor_kendall!{T<:Real,S<:Real}(x::AbstractVector{T}, y::AbstractVector{S})
    if any(isnan(x)) || any(isnan(y)) return NaN end
    n = length(x)
    if n != length(y) error("Vectors must have same length") end

    # Initial sorting
    pm = sortperm(y)
    x[:] = x[pm]
    y[:] = y[pm]
    pm[:] = sortperm(x)
    x[:] = x[pm]

    # Counting ties in x and y
    iT = 1
    nT = 0
    iU = 1
    nU = 0
    for i = 2:n
        if x[i] == x[i-1]
            iT += 1
        else
            nT += iT*(iT - 1)
            iT = 1
        end
        if y[i] == y[i-1]
            iU += 1
        else
            nU += iU*(iU - 1)
            iU = 1
        end
    end
    if iT > 1 nT += iT*(iT - 1) end
    nT = div(nT,2)
    if iU > 1 nU += iU*(iU - 1) end
    nU = div(nU,2)

    # Sort y after x
    y[:] = y[pm]

    # Calculate double ties
    iV = 1
    nV = 0
    jV = 1
    for i = 2:n
        if x[i] == x[i-1] && y[i] == y[i-1]
            iV += 1
        else
            nV += iV*(iV - 1)
            iV = 1
        end
    end
    if iV > 1 nV += iV*(iV - 1) end
    nV = div(nV,2)

    nD = div(n*(n - 1),2)
    return (nD - nT - nU + nV - 2swaps!(y))/sqrt((nD - nT)*(nD - nU))
end
cor_kendall(x::AbstractVector, y::AbstractVector) = cor_kendall!(copy(x), copy(y))
cor_kendall(x::AbstractVector, Y::AbstractMatrix) = [cor_kendall(x, Y[:,i]) for i in 1:size(Y, 2)]
cor_kendall(X::AbstractMatrix, y::AbstractVector) = [cor_kendall(X[:,i], y) for i in 1:size(X, 2)]
cor_kendall(X::AbstractMatrix, Y::AbstractMatrix) = [cor_kendall(X[:,i], Y[:,j]) for i in 1:size(X, 2), j in 1:size(Y, 2)]
function cor_kendall(X::AbstractMatrix)
    n = size(X, 2)
    C = eye(n)
    for j = 2:n
        for i = 1:j-1
            C[i,j] = cor_kendall!(X[:,i],X[:,j])
            C[j,i] = C[i,j]
        end
    end
    return C
end

# Auxilliary functions for Kendall's rank correlation
function swaps!(x::AbstractVector)
    n = length(x)
    if n == 1 return 0 end
    n2 = div(n, 2)
    xl = sub(x, 1:n2)
    xr = sub(x, n2+1:n)
    nsl = swaps!(xl)
    nsr = swaps!(xr)
    sort!(xl)
    sort!(xr)
    return nsl + nsr + mswaps(xl,xr)
end

function mswaps(x::AbstractVector, y::AbstractVector)
    i = 1
    j = 1
    nSwaps = 0
    n = length(x)
    while i <= n && j <= length(y)
        if y[j] < x[i]
            nSwaps += n - i + 1
            j += 1
        else
            i += 1
        end
    end
    return nSwaps
end

# autocorrelation for range
function autocor(x::AbstractVector, lags::Ranges)
    lx = length(x)
    if max(lags) > lx error("Autocorrelation distance must be less than sample size") end
    mx = mean(x)
    sxinv = 1/stdm(x, mx)
    xs = Array(typeof(sxinv), lx)
    for i = 1:lx
        xs[i] = (x[i] - mx)*sxinv
    end
    acf = Array(typeof(sxinv), length(lags))
    for i in 1:length(lags)
        acf[i] = dot(xs[1:end - lags[i]], xs[lags[i] + 1:end])/(lx - 1)
    end
    return acf
end

# autocorrelation at a specific lag
autocor(x::AbstractVector, lags::Real) = autocor(x, lags:lags)[1]

# autocorrelation at a default of zero to 10log10(length(v)) lags
autocor(v::AbstractVector) = autocor(v, 0:min(length(v) - 1, 10log10(length(v))))


# # Magnetic liquid droplet in magnetic field

# The behaviour of droplets under the action of magnetic fields of different configurations is an important issue in many domains, including microfluidics (Seeman et al. 2012), mechanics of tissues (Douezan et al. 2011; Frasca et al. 2014), studies of dynamic self-assembly (Timonen et al. 2013) and many others. After the successful synthesis of magnetic fluids in the late sixties by (Rosensweig 1985) an exciting story about droplets of magnetic fluid began. At first, the elongation of the droplets in an external field was observed and explained by Arhipenko et al. (1978). Later on, different droplet configurations were experimentally observed in the high-frequency rotating field, which was subject to our numerical study in [^1].

# Here I present the essential ingredients to reproduce the results of the paper with the new tools SurfaceTopology, LaplaceBIE and ElTopo (optional).

using LinearAlgebra
using GeometryTypes
using SurfaceTopology
using LaplaceBIE
using AbstractPlotting, GLMakie
## using ElTopo

# The calculation requires normal vector and vertex areas methods. Latter one gives the area of 1/3 of the vertex ring, so the sum is the area of the droplet.

function normals(vertices,topology)
    n = Point{3,Float64}[]
    for v in 1:length(vertices)
        s = Point(0,0,0)
        for (v1,v2) in EdgeRing(v,topology)
            s += cross(vertices[v1],vertices[v2])
        end
        normal = s ./ norm(s)
        push!(n,normal)
    end
    return n
end

function vertexareas(points,topology)
    vareas = zeros(Float64,length(points))
    for face in Faces(topology)
        v1,v2,v3 = face
        area = norm(cross(points[v3]-points[v1],points[v2]-points[v1])) /2
        vareas[v1] += area/3
        vareas[v2] += area/3
        vareas[v3] += area/3
    end
    return vareas
end

# These are tools to keep track if the calculation is sensible. Since we are working in an incompressible fluid, the volume needs to be constant, and the energy of the droplet decreases until equilibrium is reached.

function surfacevolume(points,topology)
    normal0 = [0,0,1]
    s = 0
    for face in Faces(topology)
        y1 = points[face[1]]
        y2 = points[face[2]]
        y3 = points[face[3]]

        normaly = cross(y2-y1,y3-y1)
        normaly /= norm(normaly)

        area = norm(cross(y2-y1,y3-y1))/2
        areaproj = dot(normaly,normal0)*area
        volume = dot(y1 + y2 + y3,normal0)/3*areaproj

        s += volume
    end
    return s
end

function energy(points,normals,faces,psi,mup,gammap,H0)
    vareas = vertexareas(points,faces)
    Area = sum(vareas)

    s = 0
    for xkey in 1:length(points)
        s += psi[xkey]*dot(H0,normals[xkey]) * vareas[xkey]
    end

    Es = gammap * Area
    Em = 1/8/pi * (1 - mup) * s

    return Es+Em
end

# For calculating the equilibrium of the droplet, we use a curvatureless algorithm for a vicious droplet developed by Zinchenko 1997. He found a way to calculate velocity generated by a surface tension without explicitly calculating the curvature, which makes it easier to implement and from some tests also more stable. The following method accepts surface defined by points, normals and faces (or topology) and surface force as well as surface tension γ and viscosity η (which only affects the scaling of time).

function stokesvelocity(points,normals,faces,forcen,etaP,gammap)
    vareas = vertexareas(points,faces)
    velocityn = zeros(Float64,length(points))

    for xkey in 1:length(points)

        x = points[xkey]
        nx = normals[xkey]
        fx = forcen[xkey]

        s = 0
        for ykey in 1:length(points)
            if ykey==xkey
                continue
            end

            y = points[ykey]
            ny = normals[ykey]
            fy = forcen[ykey]

            ### Need to check a missing 2
            s += vareas[ykey]*1 ./8/pi/etaP* dot(y-x,nx+ny)/norm(y-x)^3*(1-3*dot(y-x,nx)*dot(y-x,ny)/norm(y-x)^2) * gammap
            s += vareas[ykey]*1 ./8/pi/etaP* ( dot(nx,ny)/norm(x-y) + dot(nx,x -y)*dot(ny,x-y)/norm(x-y)^3 )*(fy - fx)
        end
        velocityn[xkey] = s
    end
    return velocityn
end

# Now having methods defined, we can include a sphere and proceed with calculation.

include("sphere.jl")
msh = unitsphere(2)
vertices, faces = msh.vertices, msh.faces

## Now let's do something fun. Visualize the process with Makie in real time.
x = Node(msh)
y = lift(x->x,x)

scene = Scene(show_axis=false)

wireframe!(scene,y,linewidth = 3f0)
mesh!(scene,y, color = :white, shading = false)

display(scene)

## Initial parameters
H0 = [4.,0.,0.]
etap = 1.
gammap = 1.
μ = 10.

t = 0.
Δt = 0.1
N = 100
volume0 = surfacevolume(vertices,faces)

record(scene, "mdrop.gif", 1:N) do i # for i in 1:N
    n = normals(vertices,faces)

    psi = surfacepotential(vertices,n,faces,μ,H0)
    P∇ψ = tangentderivatives(vertices,n,faces,psi)
    Hn = normalderivatives(vertices,n,faces,P∇ψ,μ,H0)

    E = energy(vertices,n,faces,psi,μ,gammap,H0)
    rV = surfacevolume(vertices,faces)/volume0
    @show E,rV

    Ht = [norm(j) for j in P∇ψ]

    # The force generated (M⋅∇)H, which includes a jump of magnetization at the surface and force coming from the whole bulk.
    tensorn = μ*(μ-1)/8/pi * Hn.^2 + (μ-1)/8/pi * Ht.^2

    vn = stokesvelocity(vertices,n,faces,tensorn,etap,gammap)

    vertices .+= n .* vn * Δt

    msh = HomogenousMesh(vertices,faces)

    ### ElTopo stabilization
    ## par = SurfTrack(allow_vertex_movement=true)
    ## msh = stabilize(msh,par)

    push!(x,msh)
    AbstractPlotting.force_update!()

    global vertices, faces = msh.vertices, msh.faces
    global t += Δt
end

# ![](mdrop.gif)

# [^1]: Erdmanis, J. & Kitenbergs, G. & Perzynski, R. & Cebers, A. (2017) Magnetic micro-droplet in rotating field: numerical simulation and comparison with experiment


# # Figures neither in the paper nor in the supplement
#

using Rhone2020  #hide
const R = Rhone2020  #hide

using PyPlot #hide

include("../read_measurements.jl"); #hide
include("../derived_quantities.jl"); #hide
use_progress_meter = false #hide
include("../size_evolution_models.jl"); #hide
include("../heat_transfer.jl"); #hide

# ## Moody plot
#
# The Moody plot shows the relation between the Reynolds number and the friction factor.
# Between AM15 on 9 August and AM13 on 21 August the friction factor does not change even though the Reynolds numbers are considerably different.
rcParams = PyPlot.PyDict(PyPlot.matplotlib."rcParams") #hide
rcParams["font.size"] = 9 # general #hide
figure(figsize=(9,5)) #hide
for date in keys(pick) #hide
    errorbar(mean.(mid_309_265[date][:Re][pick[date]]),mean.(mid_309_265[date][:f][pick[date]]),xerr=std.(mid_309_265[date][:Re][pick[date]]),yerr=std.(mid_309_265[date][:f][pick[date]]),fmt="_",label=join(["Aug-",date[1:2]])) #hide
    xlabel("Reynolds number") #hide
    ylabel("Friction coefficient") #hide
end #hide
legend() #hide
xscale("log") #hide
yscale("log") #hide
gcf() #hide

# ### Water offset-temperature as a histogram
#
# This plot shows the same data as Figure 5d in the paper but plotted as a histogram instead of an errorbar plot.
R.plot_tauw_hist(tau_w, tauw_measured) #hide


parse_board(b) = parse.(Int, (reduce(hcat, split.(b, " "; keepempty=false))))
function load(::Day{4}, path)
    lines = readlines(path)
    draws = parse.(Int, split(lines[1], ","))
    boards = [parse_board(lines[i:(i + 4)]) for i in 3:6:(length(lines) - 4)]
    return draws, boards
end

# Part 1
@inline is_win(check) = is_win_col(check) || is_win_row(check)
@inline is_win_col(check) = any(sum.(eachcol(check)) .== 5)
@inline is_win_row(check) = any(sum.(eachrow(check)) .== 5)

# Score board without multiplying by draw. Return 0 if not won.
@inline score(board, check) = is_win(check) ? sum(board[.!check]) : 0

# We don't really need any error-handling as the there are 100 draws,
# so the end of the game is guaranteed for any player
function solve1(::Day{4}, data) # winning score
    draws, boards = data
    checks = [zeros(Bool, 5, 5) for i in 1:length(boards)]
    @inbounds for draw in draws
        scores = map(zip(boards, checks)) do (b, c)
            c[findall(n -> n == draw, b)] .= true
            score(b, c)
        end
        if any(scores .> 0)
            return maximum(scores) * draw
        end
    end
end

# Part 2
# Luckily, we don't need to find the index of the board, just its score.
function solve2(::Day{4}, data) # last score
    draws, boards = data
    checks = [zeros(Bool, 5, 5) for i in 1:length(boards)]
    @inbounds for draw in draws
        scores = map(zip(boards, checks)) do (b, c)
            c[findall(n -> n == draw, b)] .= true
            score(b, c)
        end
        if all(scores .> 0)
            return minimum(scores) * draw
        else
            keep = findall(iszero, scores)
            boards = boards[keep]
            checks = checks[keep]
        end
    end
end

testresult1(::Day{4}) = 4512
testresult2(::Day{4}) = 1924


abstract type AbstractBC end
struct NeumannBC <: AbstractBC end
struct DirichletBC <: AbstractBC end

struct One end
struct MinusOne end
Base.getindex(::One, args...) = 1
Base.getindex(::MinusOne, args...) = -1
Base.:-(::One) = MinusOne()
Base.:-(::MinusOne) = One()

struct Derivative{O,F,D,S}
    ∂::O
    f::F
    domain::D
    scale::S
end
Base.getindex(d::Derivative, I::CartesianIndex{3}) = d.scale*d.∂(I,d.f,d.domain)
Base.:-(d::Derivative) = Derivative(d.∂, d.f, d.domain, -d.scale)
Base.:*(a, d::Derivative) = Derivative(d.∂, d.f, d.domain, a*d.scale)
Base.:*(d::Derivative, a) = *(a, d)

∂x²f(f,d::Domain) = Derivative(∂x²,f,d,1)
∂y²f(f,d::Domain) = Derivative(∂y²,f,d,1)
∂z²f(f,d::Domain) = Derivative(∂z²,f,d,1)
∂xyf(f,d::Domain) = Derivative(∂xy,f,d,1)
∂xzf(f,d::Domain) = Derivative(∂xz,f,d,1)
∂yzf(f,d::Domain) = Derivative(∂yz,f,d,1)

function new_sparse_matrix(d::Domain)
    N = prod(size(d))
    S = spzeros(N,N)
    l = LinearIndices(d)
    ldx = l[2,1,1] - l[1,1,1]
    ldy = l[1,2,1] - l[1,1,1]
    ldz = l[1,1,2] - l[1,1,1]
    for K = -1:1
        for J = -1:1
            for I = -1:1
                d = I*ldx + J*ldy + K*ldz
                n = N - abs(d)
                S += spdiagm(d => ones(n))
            end
        end
    end
    @. S = 0
    return S
end

function add_f∂a²!(S, f, da, Δa, Na, domain::Domain, Aa)
    Na = Na-1
    Δa² = Δa^2
    c = CartesianIndices(domain)
    l = LinearIndices(domain)
    for I = first(c):last(c)-Na*da
        S[l[I],l[I     ]] += -Aa/Δa² * f[I]
        S[l[I],l[I + da]] +=   1/Δa² * f[I]
    end
    for I = first(c)+da:last(c)-da
        S[l[I],l[I - da]] +=  1/Δa² * f[I]
        S[l[I],l[I     ]] += -2/Δa² * f[I]
        S[l[I],l[I + da]] +=  1/Δa² * f[I]
    end
    for I = first(c)+Na*da:last(c)
        S[l[I],l[I - da]] +=   1/Δa² * f[I]
        S[l[I],l[I     ]] += -Aa/Δa² * f[I]
    end
    return S
end
add_f∂x²!(S, f, domain::Domain) = add_f∂a²!(S, f, dx, domain.Δx, domain.nx, domain, 2)
add_f∂y²!(S, f, domain::Domain) = add_f∂a²!(S, f, dy, domain.Δy, domain.ny, domain, 2)
add_f∂z²!(S, f, domain::Domain) = add_f∂a²!(S, f, dz, domain.Δz, domain.nz, domain, 1)
add_∂x²!(S, domain::Domain) = add_f∂x²!(S, One(), domain)
add_∂y²!(S, domain::Domain) = add_f∂y²!(S, One(), domain)
add_∂z²!(S, domain::Domain) = add_f∂z²!(S, One(), domain)

function add_f∂ab!(S, f, da, db, Δa, Δb, Na, Nb, domain::Domain, Aa, Ab)
    Na = Na-1
    Nb = Nb-1
    c = CartesianIndices(domain)
    l = LinearIndices(domain)
    for I = first(c):last(c)-Na*da-Nb*db # a bottom, b bottom
        S[l[I],l[I + da + db]] +=      1/(4*Δa*Δb) * f[I]
        S[l[I],l[I          ]] +=  Aa*Ab/(4*Δa*Δb) * f[I]
        S[l[I],l[I + da     ]] += -   Ab/(4*Δa*Δb) * f[I]
        S[l[I],l[I      + db]] += -Aa   /(4*Δa*Δb) * f[I]
    end
    for I = first(c)+da:last(c)-da-Nb*db # a interior, b bottom
        S[l[I],l[I + da + db]] +=   1/(4*Δa*Δb) * f[I]
        S[l[I],l[I - da     ]] +=  Ab/(4*Δa*Δb) * f[I]
        S[l[I],l[I + da     ]] += -Ab/(4*Δa*Δb) * f[I]
        S[l[I],l[I - da + db]] += - 1/(4*Δa*Δb) * f[I]
    end
    for I = first(c)+Na*da:last(c)-Nb*db # a top, b bottom
        S[l[I],l[I      + db]] +=  Aa   /(4*Δa*Δb) * f[I]
        S[l[I],l[I - da     ]] +=     Ab/(4*Δa*Δb) * f[I]
        S[l[I],l[I          ]] += -Aa*Ab/(4*Δa*Δb) * f[I]
        S[l[I],l[I - da + db]] += -    1/(4*Δa*Δb) * f[I]
    end
    for I = first(c)+db:last(c)-Na*da-db # a bottom, b interior
        S[l[I],l[I + da + db]] +=   1/(4*Δa*Δb) * f[I]
        S[l[I],l[I      - db]] +=  Aa/(4*Δa*Δb) * f[I]
        S[l[I],l[I + da - db]] += - 1/(4*Δa*Δb) * f[I]
        S[l[I],l[I      + db]] += -Aa/(4*Δa*Δb) * f[I]
    end
    for I = first(c)+da+db:last(c)-da-db # a interior, b interior
        S[l[I],l[I + da + db]] +=  1/(4*Δa*Δb) * f[I]
        S[l[I],l[I - da - db]] +=  1/(4*Δa*Δb) * f[I]
        S[l[I],l[I + da - db]] += -1/(4*Δa*Δb) * f[I]
        S[l[I],l[I - da + db]] += -1/(4*Δa*Δb) * f[I]
    end
    for I = first(c)+Na*da+db:last(c)-db # a top, b interior
        S[l[I],l[I      + db]] +=  Aa/(4*Δa*Δb) * f[I]
        S[l[I],l[I - da - db]] +=   1/(4*Δa*Δb) * f[I]
        S[l[I],l[I      - db]] += -Aa/(4*Δa*Δb) * f[I]
        S[l[I],l[I - da + db]] += - 1/(4*Δa*Δb) * f[I]
    end
    for I = first(c)+Nb*db:last(c)-Na*da # a bottom, b top
        S[l[I],l[I + da     ]] +=     Ab/(4*Δa*Δb) * f[I]
        S[l[I],l[I      - db]] +=  Aa   /(4*Δa*Δb) * f[I]
        S[l[I],l[I + da - db]] += -    1/(4*Δa*Δb) * f[I]
        S[l[I],l[I          ]] += -Aa*Ab/(4*Δa*Δb) * f[I]
    end
    for I = first(c)+da+Nb*db:last(c)-da # a interior, b top
        S[l[I],l[I + da     ]] +=  Ab/(4*Δa*Δb) * f[I]
        S[l[I],l[I - da - db]] +=   1/(4*Δa*Δb) * f[I]
        S[l[I],l[I + da - db]] += - 1/(4*Δa*Δb) * f[I]
        S[l[I],l[I - da     ]] += -Ab/(4*Δa*Δb) * f[I]
    end
    for I = first(c)+Na*da+Nb*db:last(c) # a top, b top
        S[l[I],l[I          ]] +=  Aa*Ab/(4*Δa*Δb) * f[I]
        S[l[I],l[I - da - db]] +=      1/(4*Δa*Δb) * f[I]
        S[l[I],l[I      - db]] += -Aa   /(4*Δa*Δb) * f[I]
        S[l[I],l[I - da     ]] += -   Ab/(4*Δa*Δb) * f[I]
    end
    return S
end
add_f∂xy!(S, f, domain::Domain) = add_f∂ab!(S, f, dx, dy, domain.Δx, domain.Δy, domain.nx, domain.ny, domain, 0, 0)
add_f∂xz!(S, f, domain::Domain) = add_f∂ab!(S, f, dx, dz, domain.Δx, domain.Δz, domain.nx, domain.nz, domain, 0, 1)
add_f∂yz!(S, f, domain::Domain) = add_f∂ab!(S, f, dy, dz, domain.Δy, domain.Δz, domain.ny, domain.nz, domain, 0, 1)
add_∂xy!(S, domain::Domain) = add_f∂xy!(S, One(), domain)
add_∂xz!(S, domain::Domain) = add_f∂xz!(S, One(), domain)
add_∂yz!(S, domain::Domain) = add_f∂yz!(S, One(), domain)

function generate_sparse_linearized_L(ψ0, ϕ0, d::Domain)

    UL = new_sparse_matrix(d)
    add_∂z²!(UL, d)
    add_f∂z²!(UL, ∂x²f(ψ0,d), d)
    add_f∂z²!(UL, ∂y²f(ψ0,d), d)
    add_f∂xz!(UL, -∂xzf(ψ0,d), d)
    add_f∂yz!(UL, -∂yzf(ψ0,d), d)

    UR = new_sparse_matrix(d)
    add_f∂x²!(UR, ∂z²f(ϕ0,d), d)
    add_f∂y²!(UR, ∂z²f(ϕ0,d), d)
    add_f∂xz!(UR, -∂xzf(ϕ0,d), d)
    add_f∂yz!(UR, -∂yzf(ϕ0,d), d)

    BL = new_sparse_matrix(d)
    add_∂x²!(BL, d)
    add_∂y²!(BL, d)

    BR = new_sparse_matrix(d)
    add_f∂x²!(BR, -One(), d)
    add_f∂y²!(BR, -One(), d)
    add_f∂x²!(BR, -2*∂y²f(ψ0,d), d)
    add_f∂y²!(BR, -2*∂x²f(ψ0,d), d)
    add_f∂xy!(BR, 4*∂xyf(ψ0,d), d)

    S = hvcat((2,2), UL, UR, BL, BR)
    return S

end