module Spline

## spline construction and evaluation using "dierckx" FORTRAN library via "Dierckx" julia package
## http://www.netlib.org/dierckx/
## https://github.com/kbarbary/Dierckx.jl
using Dierckx
using QuadGK
##using Gadfly

export spline, hazard, cumhaz, eval_basis, plot_spline, pad_knots, const_spline

## FIXME: consistent use of Q vs k

## construct spline of order k+1 on interval (outer[1], outer[2])
## with given internal knots and weights theta
## knots = vector of interior knots
## w = vector of weights (one for each basis spline)
function spline(outer::NTuple{2,Float64}, knots::Vector{Float64}, w::Vector{Float64}, k::Int=3)
    ## number of basis splines
    m = length(knots) + k + 1

    @assert (outer[1] >= 0.0) & (outer[2] >= 0.0)

    if m != length(w)
        error("number of coefficients must match number of basis splines")
        return NaN
    end

    ## if !issorted(knots)
    ##   knots = sort(knots)
    ## end

    ## total number of knots
    n = m + k + 1

    ## pad knots by repeating outer knot (k+1)-times on each side
    t = Vector{Float64}(undef, n) ##Vector{Float64}(undef,n)  # All knots
    t[(k+2):(end-k-1)] = knots
    t[1:(k+1)] .= outer[1]
    t[(end-k):end] .= outer[2]

    Spline1D(t, w, k, 3, 0.0)
end

"""
    const_spline(val::Float64, outer::NTuple{2,Float64}; Q::Int=4)

Create spline of order `Q` with outer knots at `outer` and all weights equal to `val`. Resulting spline
is a constant function with value `val` everywhere.
"""
function const_spline(val::Float64, outer::NTuple{2,Float64}; Q::Int=4)
    spline(outer, [outer[1] + (outer[2]-outer[1])/2], val .+ zeros(Float64, Q+1), Q-1)
end

get_inner_knots(s::Dierckx.Spline1D) = s.t[(s.k+2):(end-s.k-1)]
get_outer_knots(s::Dierckx.Spline1D) = (s.t[1], s.t[end])

function pad_knots(knots::Vector{Float64}, outer::NTuple{2,Float64}, p::Int)
    pk = Vector{Float64}(undef, length(knots) + 2*(p+1))
    pk[(p+2):(end-p-1)] = knots
    pk[1:(p+1)] .= outer[1]
    pk[(end-p):end] .= outer[2]
    pk
end

hazard(x, s::Dierckx.Spline1D) = evaluate(s, x)

## evaluate cumulative hazard from 0.0 to x
function cumhaz(x::Vector{T}, s::Dierckx.Spline1D) where T <: Real
    res = zeros(eltype(x), length(x))
    for i in eachindex(x)
        if (s.t[1] <= 0.0) & (x[i] <= s.t[end])
            res[i] = integrate(s, 0.0, x[i])
        else
            res[i] = quadgk(t -> hazard(t,s), 0.0, x[i])[1]
        end
    end
    res
end

cumhaz(x::T, s::Dierckx.Spline1D) where T <: Real = cumhaz([x], s)[1]

function plot_spline(s::Dierckx.Spline1D)
    a = s.t[1]
    b = s.t[end] + 10.0

    x = Vector(a:((b-a)/1000):b)
    y = evaluate(s, x)

    plot(x=x, y=y, Geom.line)
end

"""
    eval_basis(s::Dierckx.Spline1D, x::Vector{Float64}, y::Vector{Float64})

Evaluate all basis functions of 1D spline `s` at points `x`, and all integrals of spline functions at `y`.
"""
function eval_basis(s::Dierckx.Spline1D, x::Vector{T}, y::Vector{T}=x) where T <: Real
    eval_basis(x, y, get_inner_knots(s), get_outer_knots(s), s.k+1)
end

function eval_basis(x::Vector{T}, y::Vector{T}, knots::Vector{T}, outer=NTuple{2,T}, Q::Int=4) where T <: Real
    K = length(knots) + Q
    A = Matrix{T}(undef, length(x), K)
    B = Matrix{T}(undef, length(y), K)

    for k in 1:K
        w = zeros(T, K)
        w[k] = 1.0
        s = spline(outer, knots, w, Q-1)
        A[:,k] = hazard(x, s)
        B[:,k] = cumhaz(y, s)
    end

    A, B
end

end

function test_spline()
    outer = (0.0, 1.0)
    knots = [0.2, 0.4, 0.6, 0.8]

    x = rand(100)
    theta = rand(Normal(0, 1), 8)
    A = Array{Float64,2}(undef, 100, 8)

    for k in 1:8
        w = zeros(Float64, 8)
        w[k] = 1.0
        s = spline(outer, knots, w, 3)
        A[:,k] = hazard(x, s)
    end

    a = A*exp.(theta)
    b = hazard(x, spline(outer, knots, exp.(theta), 3))

    display(hcat(a, b))
    isapprox(a, b)
end

function test_splineA(M=10, n=100)
    outer = (0.0, 1.0)
    knots = [0.2, 0.4, 0.6, 0.8]

    x = rand(n)
    h = zeros(Float64, n)
    ch = zeros(Float64, n)
    A = Array{Float64,2}(undef, n, 8)
    B = Array{Float64,2}(undef, n, 8)

    for k in 1:8
        w = zeros(Float64, 8)
        w[k] = 1.0
        s = spline(outer, knots, w, 3)
        A[:,k] = hazard(x, s)
        B[:,k] = cumhaz(x, s)
    end

    for m in 1:M
        theta = rand(Normal(0, 1), 8)
        et = exp.(theta)
        h += A*et
        ch += B*et
    end

    h ./ M, ch ./ M
end

function test_splineB(M=10, n=100)
    outer = (0.0, 1.0)
    knots = [0.2, 0.4, 0.6, 0.8]

    x = rand(n)
    h = zeros(Float64, n)
    ch = zeros(Float64, n)

    for m in 1:M
        theta = rand(Normal(0, 1), 8)
        s = spline(outer, knots, exp.(theta), 3)
        h += hazard(x, s)
        ch += cumhaz(x, s)
    end

    h ./ M, ch ./ M
end


"""
    sinewave1D!(height, h₀, n, ϵ)

Creates a sine wave like height field with `n` full waves, average height `h₀` and displacement magnitude `ϵ`.
"""
function sinewave1D!(height, h₀, n::Int, ϵ, dims::Int)
    # Some helping dummy array
    hdummy = ones(size(height))
    # Used to perform the accumulate operation
    height .= 1.0
    # This we get n nice sine waves
    hdummy .= 2π .* n .* accumulate(+, height, dims=dims) ./ (size(height,dims)-1)
    # Still need to use the correct undulation ϵ and initial height
    height .= h₀ .* (1.0 .+ ϵ .* (map(sin, hdummy) .- 1/size(height,dims)))
    return nothing
end

"""
    randinterface!(height, h₀, ϵ)

Creates a random height field with average height `h₀` and displacement magnitude `ϵ`.
"""
function randinterface!(height, h₀, ϵ)
    # Some helping dummy array
    hdummy = ones(size(height))
    # Used to perform the accumulate operation
    height .= 1.0
    # This we get n nice sine waves
    randn!(hdummy)
    # Still need to use the correct undulation ϵ and initial height
    height .= h₀ .* (1.0 .+ ϵ .* hdummy)
    return nothing
end

"""
    singledroplet(T, size(θ,1), size(θ,2), radius, θ, center, device=false)

Generates a fluid configuration of a single droplet in the shape of spherical cap with contact angle `θ`, sphere radius `radius` and centered at `center`.

# Arguments

- `height::Array{undef, 2}`: numerical formate, either `Float64` or `Float32`
- `radius::AbstractFloat`: radius of the underlying sphere from which the spherical cap is cut off
- `θ::AbstractFloat`: contact angle in multiples of `π`
- `center::Tuple{Int,Int}`: x and y coordinates of the center of the droplet

# Examples

```jldoctest
julia> using Swalbe, Test

julia> rad = 50; θ = 1/3;

julia> height = Swalbe.singledroplet(ones(100,100), rad, θ, (50,50));

julia> @test maximum(height) == rad * (1 - cospi(θ)) # Simple geometry
Test Passed
  Expression: maximum(height) == rad * (1 - cospi(θ))
   Evaluated: 25.0 == 25.0

julia> argmax(height) # Which is constistent with the center!
CartesianIndex(50, 50)

```

# References

See also:
"""
function singledroplet(height, radius, θ, center)
    lx, ly = size(height)
    # area = 2π * radius^2 * (1- cospi(θ))
    @inbounds for i in 1:lx
        for j in 1:ly
            circ = sqrt((i-center[1])^2 + (j-center[2])^2)
            if circ <= radius
                height[i,j] = (cos(asin(circ/radius)) - cospi(θ)) * radius
            else
                height[i,j] = 0.05
            end
        end
    end
    @inbounds for i in 1:lx
        for j in 1:ly
            if height[i,j] < 0
                height[i,j] = 0.05
            end
        end
    end
    return height
end
# For a one dimensional droplet
function singledroplet(height::Vector, radius, θ, center)
    L = size(height)[1]
    # area = 2π * radius^2 * (1- cospi(θ))
    @inbounds for i in 1:L
        circ = sqrt((i-center)^2)
        if circ <= radius
            height[i] = (cos(asin(circ/radius)) - cospi(θ)) * radius
        else
            height[i] = 0.05
        end
    end
    @inbounds for i in 1:L
        if height[i] < 0
            height[i] = 0.05
        end
    end
    return height
end

"""
    two_droplets(sys)

Generates a fluid configuration of a two droplets in the shape of spherical cap with contact angles `θ₁`, `θ₂`, sphere radius `r₁`, `r₂` and centers at `center`.

This is work in progress, therefore so far it is only available for the lower dimension model.

# Arguments

- `r₁::Float64`: radius of the first sphere from which the cap is cut of
- `r₂::Float64`: radius of the second sphere from which the cap is cut of
- `θ₁::Float64`: contact angle of the first cap in multiples of `π`
- `θ₂::Float64`: contact angle of the second cap in multiples of `π`
- `center::Tuple{Int,Int}`: x coordinates of the centers of the droplets

# Examples

```jldoctest
julia> using Swalbe, Test

julia> rad = 45; θ = 1/4; sys = Swalbe.SysConst_1D(L=200);

julia> height = Swalbe.two_droplets(sys, r₁=rad, r₂=rad, θ₁=θ, θ₂=θ);

julia> @test maximum(height) ≈ rad * (1 - cospi(θ)) atol=0.01 # Simple geometry
Test Passed
  Expression: ≈(maximum(height), rad * (1 - cospi(θ)), atol = 0.01)
   Evaluated: ≈(13.178960261768571, 13.180194846605358; atol = 0.01)

```

# References

See also:
"""
function two_droplets(sys::SysConst_1D; r₁=230, r₂=230, θ₁=1/9, θ₂=1/9, center=(sys.L/3,2*sys.L/3))
    dum = zeros(sys.L)
    dum2 = zeros(sys.L)
    height = zeros(sys.L)
    # area = 2π * radius^2 * (1- cospi(θ))
    @inbounds for i in 1:sys.L
        circ = sqrt((i-center[1])^2)
        if circ <= r₁
            dum[i] = (cos(asin(circ/r₁)) - cospi(θ₁)) * r₁
        else
            dum[i] = 0.05
        end
    end

    @inbounds for i in 1:sys.L
        circ2 = sqrt((i-center[2])^2)
        if circ2 <= r₂
            dum2[i] = (cos(asin(circ2/r₂)) - cospi(θ₂)) * r₂
        else
            dum2[i] = 0.05
        end
    end

    @inbounds for i in 1:sys.L
        if dum[i] < 0
            dum[i] = 0.05
        end
        if dum2[i] < 0
            dum2[i] = 0.05
        end
    end

    height .= dum .+ dum2 .- 0.05
    @inbounds for i in 1:sys.L
        if height[i] < 0
            height[i] = 0.05
        end
    end
    return height
end

"""
    restart_from_height(data)

Restarts a simulation from already generated height data.

# Arguments

- `data::file`: Some file with computed data
- `kind::String`: File format, `.bson` or `.jld2` are valid options
- `timestep::Int`: The time step the height is read from file
- `size::Tuple{Int,Int}`: x and y limits of the computational domain

```julia
julia> using Swalbe, FileIO, Test

julia> h1 = rand(10,10); h2 = rand(10,10);

julia> f = Dict() # For storage
Dict{Any, Any}()

julia> f["h_1"] = vec(h1); f["h_2"] = vec(h2);

julia> save("file.jld2", f)

julia> h = Swalbe.restart_from_height("file.jld2", timestep=1, size=(10,10));

julia> @test all(h .== h1)
Test Passed

julia> rm("file.jld2")

```
"""
function restart_from_height(data; kind="jld2", timestep=0, size=(512,512))
    # Allocate the height
    height = zeros(size)
    # Use the correct file format
    if kind == "jld2"
        # Pipe the data to a dataframe
        df = load(data) |> DataFrame
        if timestep == 0
            height .= reshape(df[:,end], size[1], size[2])
        else
            height .= reshape(df[!, Symbol("h_$(timestep)")], size[1], size[2])
        end
    elseif kind == "bson"
        df = DataFrame(BSON.load(data))
        # Now choose only the timestep that is needed
        if timestep == 0
            height .= reshape(df[:,end], size[1], size[2])
        else
            height .= reshape(df[!, Symbol("h_$(timestep)")], size[1], size[2])
        end
    end

    return height

end

# ---------------------------------------------------------------- #
# ---------------------- Substrate patterns ---------------------- #
# ---------------------------------------------------------------- #
"""
    boxpattern(θ, θ₀; center=(size(θ,1)÷2, size(θ,2)÷2), δₐ=1/36, side=20)

Defines a quadratic box around a `center` with side length `side`.

# Arguments

- `θ::Array{undef,2}`: Contact angle map
- `θ₀`: Default contact angle `θ`
- `center::Tuple{Int, Int}`: center of the box pattern, default value is `(size(θ,1)÷2, size(θ,2)÷2)`
- `δₐ::Float64`: Contact angle contrast with the substrate, default is set to `1/36` ≈ 5 degrees difference
- `side::Int`: Length of the sides, default is set to `20`

# Examples

```jldoctest
julia> using Swalbe

julia> θ₀ = 1/9;

julia> pattern, polygon = Swalbe.boxpattern(ones(100,100), θ₀);

julia> polygon # Some cool thing we use to create the posize(θ,2)gones, a LazySet
LazySets.VPolygon{Float64, Vector{Float64}}([[40.0, 40.0], [60.0, 40.0], [60.0, 60.0], [40.0, 60.0]])

julia> pattern[50,50] == pattern[1,1] # In the center there is a different contact angle!
false

```

# References

Really not much to say here, check out [LazySets.jl](https://github.com/JuliaReach/LazySets.jl).

"""
function boxpattern(θ, θ₀; center=(size(θ,1)÷2, size(θ,2)÷2), δₐ=1/36, side=20)
    # Vertex vectors
    xcoords = [center[1]-side/2 center[1]-side/2 center[1]+side/2 center[1]+side/2]
    ycoords = [center[2]-side/2 center[2]+side/2 center[2]-side/2 center[2]+side/2]
    # Check if the vertices are on the grid
    for i in 1:length(xcoords)
        if xcoords[i] < 1 || xcoords[i] > size(θ,1)
            xcoords[i] = 1 # If not push them to the origin
        elseif ycoords[i] < 1 || ycoords[i] > size(θ,2)
            ycoords[i] = 1
        end
    end
    # Vertex coordinates
    vertex₁ = [xcoords[1], ycoords[1]]
    vertex₂ = [xcoords[2], ycoords[2]]
    vertex₃ = [xcoords[3], ycoords[3]]
    vertex₄ = [xcoords[4], ycoords[4]]
    # Build a posize(θ,2)gon
    P = VPolygon([vertex₁, vertex₂, vertex₃, vertex₄])
    for i in 1:size(θ,1), j in 1:size(θ,2)
        if [Float64(i), Float64(j)] ∈ P
            θ[j,i] = θ₀ + δₐ
        else
            θ[j,i] = θ₀
        end
    end
    return θ, P
end

"""
    ellipsepattern(θ, θ₀; center=(size(θ,1)÷2, size(θ,2)÷2), δₐ=1/36, a=10, b=5)

Creates an ellipse shaped contact angle defect with contact angle mismatch `δₐ`.

# Arguments

- `θ::Array{undef,2}`: Contact angle map
- `θ₀`: Default contact angle `θ`
- `center::Tuple{Int, Int}`: center of the created pattern, default values is `center = (size(θ,1)÷2, size(θ,2)÷2)`
- `δₐ::Float64`: contact angle mismatch between patch and rest of substrate, default is `δₐ = 1/36` or 5 degrees
- `a::Int`:: semimajor half ax of the ellipse, default value is `a=10`
- `b::Int`:: semiminor half ax of the ellipse, default value is `b=5`

# Examples

```jldoctest
julia> using Swalbe, Test

julia> θ₀ = 1/9;

julia> θ, P = Swalbe.ellipsepattern(ones(100,100), θ₀); # per default the center is in the middle!

julia> @test θ[1,1] == θ₀
Test Passed
  Expression: θ[1, 1] == θ₀
   Evaluated: 0.1111111111111111 == 0.1111111111111111

julia> @test θ[50,50] == θ₀ + 1/36 # The default increment, is about 5 degrees.
Test Passed
  Expression: θ[50, 50] == θ₀ + 1 / 36
   Evaluated: 0.1388888888888889 == 0.1388888888888889
```

# References
Nothing interesting here.

"""
function ellipsepattern(θ, θ₀; center=(size(θ,1)÷2, size(θ,2)÷2), δₐ=1/36, a=10, b=5)
    # Don't know why LazySets needs this, can not just parse center as argument of Ellipsoid constructor
    mid = ones(2)
    mid[1] = center[1]
    mid[2] = center[2]
    # The eigenvalues of the shape matrix, Q in the LazySets docs, needs to be spezified not the axis length
    shape₁ = a^(2)
    shape₂ = b^(2)
    P = Ellipsoid(mid, [shape₁ 0.0; 0.0 shape₂])
    for i=1:size(θ,1), j=1:size(θ,2)
        if [Float64(i), Float64(j)] ∈ P
            θ[j,i] = θ₀ + δₐ
        else
            θ[j,i] = θ₀
        end
    end
    return θ, P
end

"""
    trianglepattern(θ, θ₀; center=(size(θ,1)÷2, size(θ,2)÷2), δₐ=1/36, side=60)

Generates an equilateral triangle centered around `center` with contact angle contrast `δₐ` and side length `side`.

# Arguments
- `θ::Array{undef,2}`: Contact angle map
- `θ₀`: Default contact angle `θ`
- `center::Tuple(Int, Int)`: position of the center of the triangle
- `δₐ::Float64`: contact angle contrast with the rest of the substrate
- `side::Int`: length of the sides of the equilateral triangle

# Examples
```jldoctest
julia> using Swalbe

julia> θ, P = Swalbe.trianglepattern(ones(50,50), 1/9, side=20) # Returns a polygon and the contact angle field
([0.1111111111111111 0.1111111111111111 … 0.1111111111111111 0.1111111111111111; 0.1111111111111111 0.1111111111111111 … 0.1111111111111111 0.1111111111111111; … ; 0.1111111111111111 0.1111111111111111 … 0.1111111111111111 0.1111111111111111; 0.1111111111111111 0.1111111111111111 … 0.1111111111111111 0.1111111111111111], LazySets.VPolygon{Float64, Vector{Float64}}([[15.0, 19.226497308103742], [35.0, 19.226497308103742], [25.0, 36.547005383792516]]))

julia> P
LazySets.VPolygon{Float64, Vector{Float64}}([[15.0, 19.226497308103742], [35.0, 19.226497308103742], [25.0, 36.547005383792516]])

julia> θ[25,25]
0.1388888888888889
```

# References

See also:
"""
function trianglepattern(θ, θ₀; center=(size(θ,1)÷2, size(θ,2)÷2), δₐ=1/36, side=60)
    height = sqrt(3)/2*side
    # Vertices vectors
    xcoords = [center[1]-side/2 center[1]+side/2 center[1]]
    ycoords = [center[2]-height/3 center[2]-height/3 center[2]+2*height/3]
    # Check if the vertices are on the grid
    for i in 1:length(xcoords)
        if xcoords[i] < 1 || xcoords[i] > size(θ,1)
            xcoords[i] = 1 # If not push them to the origin
        elseif ycoords[i] < 1 || ycoords[i] > size(θ,2)
            ycoords[i] = 1
        end
    end
    vertex₁ = [xcoords[1], ycoords[1]]
    vertex₂ = [xcoords[2], ycoords[2]]
    vertex₃ = [xcoords[3], ycoords[3]]
    P = VPolygon([vertex₁, vertex₂, vertex₃])
    for i in 1:size(θ,1), j in 1:size(θ,2)
        if [Float64(i), Float64(j)] ∈ P
            θ[j,i] = θ₀ + δₐ
        else
            θ[j,i] = θ₀
        end
    end
    return θ, P
end


#=
# Import a simple triangle or quad mesh from a .msh file.
=#

### If the Finch package has already been added, use this line #########
using Finch # Note: to add the package, first do: ]add "https://github.com/paralab/Finch.git"

### If not, use these four lines (working from the examples directory) ###
# if !@isdefined(Finch)
#     include("../Finch.jl");
#     using .Finch
# end
##########################################################################

init_finch("unstruct2dtest");

useLog("unstruct2dlog", level=3)

domain(2, grid=UNSTRUCTURED)
functionSpace(order=2)

# This rectangle covers [0, 0.1]x[0, 0.3]
# Uncomment the desired mesh.
mesh("utriangle.msh")  # Using triangles
#mesh("uquad.msh")     # Using quads

u = variable("u")
testSymbol("v")

boundary(u, 1, DIRICHLET, 0)

coefficient("f", "(-10-(x+1)*200*pi*pi)*sin(10*pi*x)*sin(10*pi*y) + 10*pi*cos(10*pi*x)*sin(10*pi*y)")
coefficient("k", "x+1")
coefficient("C", 10)
weakForm(u, "k*dot(grad(u), grad(v)) + C*u*v+ f*v")

solve(u);

finalize_finch();

# exact solution is sin(10*pi*x)*sin(10*pi*y)
# check error
maxerr = 0;
exact(x,y) = sin(10*pi*x)*sin(10*pi*y);

for i=1:size(Finch.grid_data.allnodes,2)
    x = Finch.grid_data.allnodes[1,i];
    y = Finch.grid_data.allnodes[2,i];
    err = abs(u.values[i] - exact(x,y));
    global maxerr;
    maxerr = max(err,maxerr);
end
println("max error = "*string(maxerr));


using ResumableFunctions
using SimJulia
using Distributions
using DataFrames
using Random
using StatsPlots
using BenchmarkTools

function increment!(a::Array{Int64})
    push!(a,a[length(a)]+1)
end

function decrement!(a::Array{Int64})
    push!(a,a[length(a)]-1)
end

function carryover!(a::Array{Int64})
    push!(a,a[length(a)])
end;

mutable struct SIRPerson
    id::Int64 # numeric ID
    status::Symbol # :S,I,R
end;

mutable struct SIRModel
    sim::Simulation
    β::Float64
    c::Float64
    γ::Float64
    ta::Array{Float64}
    Sa::Array{Int64}
    Ia::Array{Int64}
    Ra::Array{Int64}
    allIndividuals::Array{SIRPerson}
end;

function infection_update!(sim::Simulation,m::SIRModel)
    push!(m.ta,now(sim))
    decrement!(m.Sa)
    increment!(m.Ia)
    carryover!(m.Ra)
end;

function recovery_update!(sim::Simulation,m::SIRModel)
    push!(m.ta,now(sim))
    carryover!(m.Sa)
    decrement!(m.Ia)
    increment!(m.Ra)
end;

@resumable function live(sim::Simulation, individual::SIRPerson, m::SIRModel)
  while individual.status==:S
      # Wait until next contact
      @yield timeout(sim,rand(Distributions.Exponential(1/m.c)))
      # Choose random alter
      alter=individual
      while alter==individual
          N=length(m.allIndividuals)
          index=rand(Distributions.DiscreteUniform(1,N))
          alter=m.allIndividuals[index]
      end
      # If alter is infected
      if alter.status==:I
          infect = rand(Distributions.Uniform(0,1))
          if infect < m.β
              individual.status=:I
              infection_update!(sim,m)
          end
      end
  end
  if individual.status==:I
      # Wait until recovery
      @yield timeout(sim,rand(Distributions.Exponential(1/m.γ)))
      individual.status=:R
      recovery_update!(sim,m)
  end
end;

function MakeSIRModel(u0,p)
    (S,I,R) = u0
    N = S+I+R
    (β,c,γ) = p
    sim = Simulation()
    allIndividuals=Array{SIRPerson,1}(undef,N)
    for i in 1:S
        p=SIRPerson(i,:S)
        allIndividuals[i]=p
    end
    for i in (S+1):(S+I)
        p=SIRPerson(i,:I)
        allIndividuals[i]=p
    end
    for i  in (S+I+1):N
        p=SIRPerson(i,:R)
        allIndividuals[i]=p
    end
    ta=Array{Float64,1}(undef,0)
    push!(ta,0.0)
    Sa=Array{Int64,1}(undef,0)
    push!(Sa,S)
    Ia=Array{Int64,1}(undef,0)
    push!(Ia,I)
    Ra=Array{Int64,1}(undef,0)
    push!(Ra,R)
    SIRModel(sim,β,c,γ,ta,Sa,Ia,Ra,allIndividuals)
end;

function activate(m::SIRModel)
     [@process live(m.sim,individual,m) for individual in m.allIndividuals]
end;

function sir_run(m::SIRModel,tf::Float64)
    SimJulia.run(m.sim,tf)
end;

function out(m::SIRModel)
    result = DataFrame()
    result[!,:t] = m.ta
    result[!,:S] = m.Sa
    result[!,:I] = m.Ia
    result[!,:R] = m.Ra
    result
end;

tmax = 40.0;

u0 = [990,10,0];

p = [0.05,10.0,0.25];

Random.seed!(1234);

des_model = MakeSIRModel(u0,p)
activate(des_model)
sir_run(des_model,tmax)

data_des=out(des_model);

@df data_des plot(:t,
    [:S :I :R],
    labels = ["S" "I" "R"],
    xlab="Time",
    ylab="Number")

@resumable function live_delay(sim::Simulation, individual::SIRPerson, m::SIRModel)
  while individual.status==:S
      # Wait until next contact
      @yield timeout(sim,rand(Distributions.Exponential(1/m.c)))
      # Choose random alter
      alter=individual
      while alter==individual
          N=length(m.allIndividuals)
          index=rand(Distributions.DiscreteUniform(1,N))
          alter=m.allIndividuals[index]
      end
      # If alter is infected
      if alter.status==:I
          infect = rand(Distributions.Uniform(0,1))
          if infect < m.β
              individual.status=:I
              infection_update!(sim,m)
          end
      end
  end
  # The following is modified to capture a fixed infection delay
  if individual.status==:I
      # Wait until recovery, here fixed
      @yield timeout(sim,1/m.γ)
      individual.status=:R
      recovery_update!(sim,m)
  end
end;

function activate_delay(m::SIRModel)
     [@process live_delay(m.sim,individual,m) for individual in m.allIndividuals]
end;

des_delay_model = MakeSIRModel(u0,p)
activate_delay(des_delay_model)
sir_run(des_delay_model,tmax)

data_delay_des=out(des_delay_model);

@df data_delay_des plot(:t,
    [:S :I :R],
    labels = ["S" "I" "R"],
    xlab="Time",
    ylab="Number")

@benchmark begin
    des_model = MakeSIRModel(u0,p)
    activate(des_model)
    sir_run(des_model,tmax)
end


# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.3'
#       jupytext_version: 0.8.6
#   kernelspec:
#     display_name: Julia 1.0.3
#     language: julia
#     name: julia-1.0
# ---

module DiscreteTimeSIR
using RandomNumbers
using DataFrames

@inline @fastmath function randbn(n,p,rng)
    q = 1.0 - p
    s = p/q
    a = (n+1)*s
    r = exp(n*log(q))
    x = 0
    u = rand(rng)
    while true
        if (u < r)
            return x
        end
        u -= r
        x += 1
        r *= (a/x)-s
    end
end

@inline @fastmath function sir(u, prams, rng)

    (S, I, R, Y) = u
    (β, γ, ι, N, δt) = prams
    λ = β * (I + ι) / N
    ifrac = 1.0 - exp(-λ * δt)
    rfrac = 1.0 - exp(-γ * δt)
    infection = randbn(S, ifrac, rng)
    recovery = randbn(I, rfrac, rng)
    return (S - infection, I + infection - recovery, R + recovery, Y + infection)
end

function simulate(r)
    prams = (0.1, 0.05, 0.01, 1000.0, 0.1)
    tf = 200
    t = 0:0.1:tf
    tl = length(t)
    S = zeros(tl)
    I = zeros(tl)
    R = zeros(tl)
    Y = zeros(tl)
    u0 = (999, 1, 0, 0)
    (S[1],I[1],R[1],Y[1]) = u0
    u = u0
    for i in 2:tl
        u = sir(u, prams, r)
        (S[i],I[i],R[i],Y[i]) = u
    end
    return DataFrame(Time=t,S=S,I=I,R=R,Y=Y)
end

seed = 42
r = Xorshifts.Xorshift128Plus(seed);

sir_out = simulate(r);

if(size(sir_out)[1] >= 6)
    first(sir_out,6)
else
    first(sir_out,size(sir_out)[1])
end

using Plots
using StatPlots

@df sir_out plot(:Time, [:S :I :R], colour = [:red :green :blue], xlabel="Time",ylabel="Number")

end


using Printf

abstract type SubspaceCache{T} end
abstract type HermitianSubspaceCache{T} <: SubspaceCache{T} end

include("stegr_cache.jl")

get_subspace_cache(Ks::KrylovSubspace{B,T,U}) where {B,T,U<:Complex} =
    error("Subspace exponential caches not yet available for non-Hermitian matrices.")
get_subspace_cache(Ks::KrylovSubspace{B,T,U}) where {B,T,U<:Real} =
    StegrCache(Ks)

"""
    expv!(w, t, A, b, Ks, cache)

Alternative interface for calculating the action of `exp(t*A)` on the
vector `b`, storing the result in `w`. The Krylov iteration is
terminated when an error estimate for the matrix exponential in the
generated subspace is below the requested tolerance. `Ks` is a
`KrylovSubspace` and `typeof(cache)<:HermitianSubspaceCache`, the
exact type decides which algorithm is used to compute the subspace
exponential.
"""
function expv!(w::AbstractVector{T}, t::Number, A, b::AbstractVector{T},
               Ks::KrylovSubspace{B, T, B}, cache::HSC;
               atol::B=1.0e-8, rtol::B=1.0e-4,
               m=min(Ks.maxiter, size(A,1)),
               verbose::Bool=false) where {B, T <: Number, HSC <: HermitianSubspaceCache}
    if m > Ks.maxiter
        resize!(Ks, m)
    else
        Ks.m = m # might change if error estimate is below requested tolerance
    end

    V, H = getV(Ks), getH(Ks)
    Ks.beta = norm(b)
    @. V[:, 1] = b / Ks.beta

    ε = atol + rtol * Ks.beta
    verbose && @printf("Initial norm: β₀ %e, stopping threshold: %e\n", Ks.beta, ε)

    α = @diagview(H)
    β = @diagview(H,-1)
    n = size(V, 1)

    for j ∈ 1:m
        lanczos_step!(j, A, V, α, β)
        expT!(@view(α[1:j]), @view(β[1:j]), t, cache)

        # This is practical error estimate Er₂ from
        #
        #   Saad, Y. (1992). Analysis of some Krylov subspace
        #   approximations. SIAM Journal on Numerical Analysis.
        σ = β[j]*Ks.beta*abs(cache.v[j])
        verbose && @printf("iter %d, α[%d] %e, β[%d] %e, σ %e\n",j, j, α[j], j, β[j], σ)
        if σ < ε
            Ks.m = j
            break
        end
    end
    verbose && println("Krylov subspace size: ", Ks.m)

    lmul!(Ks.beta, mul!(w, @view(Ks.V[:,1:Ks.m]), @view(cache.v[1:Ks.m])))
end


module DuffyQuadratures

using Gridap
using FastGaussQuadrature: gaussjacobi
using FastGaussQuadrature: gausslegendre
using StaticArrays

export DuffyQuadrature
import Gridap: coordinates
import Gridap: weights

struct DuffyQuadrature{D} <: Quadrature{D}
  x::Vector{Point{D,Float64}}
  w::Vector{Float64}
end

function DuffyQuadrature{D}(order::Integer) where D
  x,w = _duffy_quad_data(order,D)
  DuffyQuadrature{D}(x,w)
end

coordinates(q::DuffyQuadrature) = q.x

weights(q::DuffyQuadrature) = q.w

# Helpers

function _duffy_quad_data(order::Integer,D::Int)

  beta = 0
  dim_to_quad_1d = [
    _gauss_jacobi_in_0_to_1(order,(D-1)-(d-1),beta) for d in 1:(D-1) ]

  quad_1d = _gauss_legendre_in_0_to_1(order)
  push!(dim_to_quad_1d,quad_1d)

  x_pos = 1
  w_pos = 2
  dim_to_xs_1d = [quad_1d[x_pos] for quad_1d in dim_to_quad_1d]
  dim_to_ws_1d = [quad_1d[w_pos] for quad_1d in dim_to_quad_1d]

  a = 0.5
  for d in (D-1):-1:1
    ws_1d = dim_to_ws_1d[d]
    ws_1d[:] *= a
    a *= 0.5
  end

  x,w = _tensor_product(dim_to_xs_1d,dim_to_ws_1d)

  (_duffy_map.(x),w)

end

"""
Duffy map from the n-cube in [0,1]^d to the n-simplex in [0,1]^d
"""
function _duffy_map(q::Point{D,T}) where {D,T}
  m = zero(MVector{D,T})
  m[1] = q[1]
  a = one(T)
  for i in 2:D
    a *= (1-q[i-1])
    m[i] = a*q[i]
  end
  Point{D,T}(m)
end

_duffy_map(q::Point{1,T}) where T = q

function _gauss_jacobi_in_0_to_1(order,alpha,beta)
  n = _npoints_from_order(order)
  x,w = gaussjacobi(n,alpha,beta)
  _map_to(0,1,x,w)
end

function _gauss_legendre_in_0_to_1(order)
  n = _npoints_from_order(order)
  x,w = gausslegendre(n)
  _map_to(0,1,x,w)
end

"""
Transforms a 1-D quadrature from `[-1,1]` to `[a,b]`, with `a<b`.
"""
function _map_to(a,b,points,weights)
  points_ab = 0.5*(b-a)*points .+ 0.5*(a+b)
  weights_ab = 0.5*(b-a)*weights
  (points_ab, weights_ab)
end

function _npoints_from_order(order)
  ceil(Int, (order + 1.0) / 2.0 )
end

function _tensor_product(
  dim_to_xs_1d::Vector{Vector{T}},
  dim_to_ws_1d::Vector{Vector{W}}) where {T,W}

  D = length(dim_to_ws_1d)
  @assert D == length(dim_to_xs_1d)
  dim_to_n = [length(ws_1d) for ws_1d in dim_to_ws_1d]
  n = prod(dim_to_n)
  xs = zeros(Point{D,T},n)
  ws = zeros(W,n)
  cis = CartesianIndices(tuple(dim_to_n...))
  m = zero(MVector{D,T})
  _tensor_product!(xs,ws,dim_to_xs_1d,dim_to_ws_1d,cis,m)
  (xs,ws)
end

function _tensor_product!(
  xs,ws,dim_to_xs_1d,dim_to_ws_1d,cis::CartesianIndices{D},m) where D
  k = 1
  for ci in cis
    w = 1.0
    for d in 1:D
      xs_1d = dim_to_xs_1d[d]
      ws_1d = dim_to_ws_1d[d]
      i = ci[d]
      xi = xs_1d[i]
      wi = ws_1d[i]
      w *= wi
      m[d] = xi
    end
    xs[k] = m
    ws[k] = w
    k += 1
  end
end

end # module


type Adaboost
    n_clf::Int64
    clf::Matrix
end

function Adaboost(;
                  n_clf::Int64 = 10
                  )
    clf = zeros(4, n_clf)
    return Adaboost(n_clf, clf)
end

function train!(model::Adaboost, X::Matrix, y::Vector)
    n_sample, n_feature = size(X)
    w = ones(n_sample) / n_sample
    threshold = 0
    polarity = 0
    feature_index = 0
    alpha = 0
    for i = 1:model.n_clf
        err_max = 1e10
        for feature_ind = 1:n_feature
            for threshold_ind = 1:n_sample
                polarity_ = 1
                err = 0
                threshold_ = X[threshold_ind, feature_ind]

                for sample_ind = 1:n_sample
                    pred = 1
                    x = X[sample_ind, feature_ind]
                    if x < threshold_
                        pred = -1
                    end
                    err += w[sample_ind] * (y[sample_ind] != pred)
                end
                err = err/sum(w)
                if err > 0.5
                    err = 1 - err
                    polarity_ = -1
                end

                if err < err_max
                    err_max = err
                    threshold = threshold_
                    polarity = polarity_
                    feature_index = feature_ind
                end
            end
        end

        alpha = 1/2 * log((1.000001-err_max)/(err_max+0.000001))

        for j = 1:n_sample
            pred = 1
            x = X[j, feature_index]
            if polarity * x < polarity * threshold
                pred = -1
            end
            w[j] = w[j] * exp(-alpha * y[j] * pred)
        end
        model.clf[:, i] = [feature_index, threshold, polarity, alpha]
    end
end

function predict(model::Adaboost,
                 x::Matrix)
    n = size(x,1)
    res = zeros(n)
    for i = 1:n
        res[i] = predict(model, x[i,:])
    end
    return res
end

function predict(model::Adaboost,
                 x::Vector)
    s = 0
    for i = 1:model.n_clf
        pred = 1
        feature_index = trunc(Int64,model.clf[1,i])
        threshold = model.clf[2,i]
        polarity = model.clf[3,i]
        alpha = model.clf[4,i]
        x_temp = x[feature_index]
        if polarity * x_temp < polarity * threshold
            pred = -1
        end
        s += alpha * pred
    end

    return sign(s)

end

function test_Adaboost()
    X_train, X_test, y_train, y_test = make_cla(n_features = 8, n_samples = 1000)

    #Adaboost
    model = Adaboost()
    train!(model,X_train, y_train)
    predictions = predict(model,X_test)
    println("The number of week classifiers ", 10)
    println("classification accuracy: ", accuracy(y_test, predictions))

    #PCA

    pca_model = PCA()
    train!(pca_model, X_test)
    plot_in_2d(pca_model, X_test, predictions, "Adaboost")

end


export  EmptyQuad,
        Quad

struct Quad{T, V<:AbstractVector{<:T}} <: AbstractQuad{T}
    name::String
    Nquad::Int # number of qudrature points
    nodes::V
    weights::V

    function Quad(name::String, N::Int, nodes, weights)
        N <= 0 && throw(DomainError(N,"number of qudrature points has to be positive"))
        !(length(nodes) == length(weights)) && throw(InconsistencyError("inconsistent numbers of nodes and weights"))
        new{promote_type(eltype(nodes), eltype(weights)), promote_type(typeof(nodes), typeof(weights))}(lowercase(name), N, nodes, weights)
    end
end

struct EmptyQuad{T} <: AbstractQuad{T}
    EmptyQuad() = new{Float64}()
end

# general constructor
function Quad(N::Int, α::AbstractVector{<:Real}, β::AbstractVector{<:Real})
    !(length(α) == length(β)) && throw(InconsistencyError("inconsistent numbers of recurrence coefficients"))
    !(N <= length(α) - 1) && throw(DomainError(N,"requested number of quadrature points $N cannot be provided with $(length(α)) recurrence coefficients"))
    nodes, weights = gauss(N,α,β)
    Quad("golubwelsch", N, nodes, weights)
end

# quadrature rules for orthoPolys
Quad(N::Int, op::AbstractOrthoPoly) = Quad(N, op.α, op.β)
Quad(op::AbstractOrthoPoly) = Quad(op.deg, op)

#####################################################
#####################################################
#####################################################
# the constructor below is probably not relevant
#####################################################
#####################################################
#####################################################
# function Quad(N::Int, weight::Function, α::Vector{<:Real}, β::Vector{<:Real}, supp::Tuple{<:Real,<:Real}, symm::Bool, d::Dict=Dict())
#     m = Measure("fun_"*String(nameof(weight)),weight,supp,symm,d)
#     Quad(N,α,β,m)
# end
#####################################################
#####################################################
#####################################################

# all-purpose constructor (last resort!)
function Quad(N::Int, w::Function, dom::Tuple{<:Real,<:Real}; quadrature::Function=clenshaw_curtis)
    N <= 0 && throw(DomainError(N, "number of quadrature points has to be positive"))
    nodes, weights = quadgp(w, dom[1], dom[2], N; quadrature=quadrature)
    Quad("quadgp", N, nodes, weights)
end

function Quad(N::Int, measure::AbstractMeasure; quadrature::Function=clenshaw_curtis)
    typeof(measure) != Measure && @warn "For measures of type $(typeof(measure)) the quadrature rule should be based on the recurrence coefficients."
    Quad(N, measure.w, (measure.dom[1], measure.dom[2]); quadrature=quadrature)
end


using Zygote
using FiniteDifferences
using KernelFunctions
using Distances
using Test

x = rand(5)
y = rand(5)

@testset "Testing Zygote adjoints" begin
gzeucl = first(Zygote.gradient(xy->evaluate(Euclidean(),xy[1],xy[2]),[x,y]))
gzsqeucl =  first(Zygote.gradient(xy->evaluate(SqEuclidean(),xy[1],xy[2]),[x,y]))
gzdotprod = first(Zygote.gradient(xy->evaluate(KernelFunctions.DotProduct(),xy[1],xy[2]),[x,y]))

FDM = central_fdm(5,1)

gfeucl = collect(first(FiniteDifferences.grad(FDM,xy->evaluate(Euclidean(),xy[1],xy[2]),(x,y))))
gfsqeucl = collect(first(FiniteDifferences.grad(FDM,xy->evaluate(SqEuclidean(),xy[1],xy[2]),(x,y))))
gfdotprod =collect(first(FiniteDifferences.grad(FDM,xy->evaluate(KernelFunctions.DotProduct(),xy[1],xy[2]),(x,y))))

@test all(gzeucl .≈ gfeucl)
@test all(gzsqeucl .≈ gfsqeucl)
@test all(gzdotprod .≈ gfdotprod)

end


using Printf

struct KPCSupply{T<:Real}
    item::String
    weight::T
    value::T
    uvalue::T
end
function KPCSupply(item::AbstractString, weight::Real, value::Real)
    w, v = promote(weight, value)
    KPCSupply(item, w, v, v / w)
end

Base.show(io::IO, s::KPCSupply) = print(io, s.item, @sprintf " (%.2f kg, %.2f €, %.2f €/kg)" s.weight s.value s.uvalue)
Base.isless(a::KPCSupply, b::KPCSupply) = a.uvalue < b.uvalue

function solve(store::Vector{KPCSupply{T}}, capacity::Real) where T<:Real
    sack = similar(store, 0) # vector like store, but of length 0
    kweight = zero(T)
    for s in sort(store, rev = true)
        if kweight + s.weight ≤ capacity
            kweight += s.weight
            push!(sack, s)
        else
            w = capacity - kweight
            v = w * s.uvalue
            push!(sack, KPCSupply(s.item, w, v, s.value))
            break
        end
    end
    return sack
end


let
    function f_1(x::Vector)
        (x[1] - 5.0)^4
    end

    function g!_1(x::Vector, storage::Vector)
        storage[1] = 4.0 * (x[1] - 5.0)^3
    end

    function h!_1(x::Vector, storage::Matrix)
        storage[1, 1] = 12.0 * (x[1] - 5.0)^2
    end

    d = TwiceDifferentiableFunction(f_1, g!_1, h!_1)

    # Need to specify autodiff!
    @test_throws ErrorException Optim.optimize(DifferentiableFunction(f_1, g!_1), [0.0], Newton())
    Optim.optimize(DifferentiableFunction(f_1, g!_1), [0.0], Newton(), Optim.Options(autodiff = true))

    results = Optim.optimize(d, [0.0], Newton())
    @test_throws ErrorException Optim.x_trace(results)
    @assert Optim.g_converged(results)
    @assert norm(Optim.minimizer(results) - [5.0]) < 0.01

    eta = 0.9

    function f_2(x::Vector)
      (1.0 / 2.0) * (x[1]^2 + eta * x[2]^2)
    end

    function g!_2(x::Vector, storage::Vector)
      storage[1] = x[1]
      storage[2] = eta * x[2]
    end

    function h!_2(x::Vector, storage::Matrix)
      storage[1, 1] = 1.0
      storage[1, 2] = 0.0
      storage[2, 1] = 0.0
      storage[2, 2] = eta
    end

    d = TwiceDifferentiableFunction(f_2, g!_2, h!_2)
    results = Optim.optimize(d, [127.0, 921.0], Newton())
    @test_throws ErrorException Optim.x_trace(results)
    @assert Optim.g_converged(results)
    @assert norm(Optim.minimizer(results) - [0.0, 0.0]) < 0.01

    # Test Optim.newton for all twice differentiable functions in Optim.UnconstrainedProblems.examples
    for (name, prob) in Optim.UnconstrainedProblems.examples
    	if prob.istwicedifferentiable
    		ddf = TwiceDifferentiableFunction(prob.f, prob.g!,prob.h!)
    		res = Optim.optimize(ddf, prob.initial_x, Newton())
    		@assert norm(Optim.minimizer(res) - prob.solutions) < 1e-2
    	end
    end

    let
        prob=Optim.UnconstrainedProblems.examples["Himmelblau"]
        ddf = TwiceDifferentiableFunction(prob.f, prob.g!, prob.h!)
        res = optimize(ddf, [0., 0.], Newton())
        @assert norm(Optim.minimizer(res) - prob.solutions) < 1e-10
    end

    for (name, prob) in Optim.UnconstrainedProblems.examples
    	if prob.istwicedifferentiable
    		ddf = DifferentiableFunction(prob.f, prob.g!)
    		res = Optim.optimize(ddf, prob.initial_x, Newton(), Optim.Options(autodiff = true))
    		@assert norm(Optim.minimizer(res) - prob.solutions) < 1e-2
    		res = Optim.optimize(ddf.f, prob.initial_x, Newton(), Optim.Options(autodiff = true))
    		@assert norm(Optim.minimizer(res) - prob.solutions) < 1e-2
            res = Optim.optimize(ddf.f, ddf.g!, prob.initial_x, Newton(), Optim.Options(autodiff = true))
    		@assert norm(Optim.minimizer(res) - prob.solutions) < 1e-2
    	end
    end
end


"""
    NSVerify(optimizer, m::Float64)

NSVerify finds counter examples using mixed integer linear programming.

# Problem requirement
1. Network: any depth, ReLU activation
2. Input: hyperrectangle or hpolytope
3. Output: PolytopeComplement

# Return
`CounterExampleResult`

# Method
MILP encoding (using `m`). No presolve.
Default `optimizer` is `GLPKSolverMIP()`. Default `m` is `1000.0` (should be large enough to avoid approximation error).

# Property
Sound and complete.

# Reference
[A. Lomuscio and L. Maganti,
"An Approach to Reachability Analysis for Feed-Forward Relu Neural Networks,"
*ArXiv Preprint ArXiv:1706.07351*, 2017.](https://arxiv.org/abs/1706.07351)
"""
@with_kw struct NSVerify <: Solver
    optimizer = GLPK.Optimizer
    m = nothing
end

function solve(solver::NSVerify, problem::Problem)
    # Set up and solve the problem
    model = Model(solver)
    z = init_vars(model, problem.network, :z, with_input=true)
    δ = init_vars(model, problem.network, :δ, binary=true)
    # Set M automatically if not already set.
    if isnothing(solver.m)
        model[:M] = set_automatic_M(problem)
    else
        @warn "M should be chosen carefully. An M which is too small will cause
        the problem to be solved incorrectly. Not setting the `m` keyword, or setting
        it equal to `nothing` will cause a safe value to be calculated automatically.
        E.g. NSVerify()." maxlog = 1
        model[:M] = solver.m
    end

    add_set_constraint!(model, problem.input, first(z))
    add_complementary_set_constraint!(model, problem.output, last(z))
    encode_network!(model, problem.network, MixedIntegerLP())
    feasibility_problem!(model)
    optimize!(model)

    if termination_status(model) == OPTIMAL
        return CounterExampleResult(:violated, value(first(z)))
    elseif termination_status(model) == INFEASIBLE
        return CounterExampleResult(:holds)
    else
        return CounterExampleResult(:unknown)
    end
end

function set_automatic_M(problem)
    # Compute the largest pre-activation bound absolute value.
    # M must be larger than any value a variable in the problem
    # can take.
    bounds = get_bounds(problem, false)
    M = maximum(abs, Iterators.flatten(abs.(hr.center) + hr.radius for hr in bounds))
end


using Test
using MechGluecode
using MechanicalUnits
using MechanicalUnits: Unitfu.promote_to_derived
@import_expand(~m, s, N)
using DifferentialEquations
import MechGluecode: norm # should work without

@testset "Initial checks ArrayPartition" begin
    r0 = [1131.340, -2282.343, 6672.423]∙km
    v0 = [-5.64305, 4.30333, 2.42879]∙km/s
    Δt = 86400.0*365∙s
    μ = 398600.4418∙km³/s²
    rv0 = ArrayPartition(r0,v0)

    function fo(dy, y, μ, t)
        r = norm(y.x[1])
        dy.x[1] .= y.x[2]
        dy.x[2] .= -μ .* y.x[1] / r^3
    end
    prob = ODEProblem(fo, rv0, (0.0s, 0.035Δt), μ)
    @time sol = solve(prob, Tsit5(), dt = 0.1s);
    # Plot all variables vs time
    plot(sol)
    plot(sol, vars = (1))
    plot(sol, vars = (2))
    plot(sol, vars = (3))
    plot(sol, vars = (4))
    plot(sol, vars = (5))
    plot(sol, vars = (6))

    # Plot in phase space
    plot(sol, vars=(1,4))
    plot(sol, vars=(2,5))
    plot(sol, vars=(3,6))
    plot(sol, vars=(2,5))
    # Three distances vs time (0)
    plot(sol, vars=[(0,1), (0,2),(0,3)])
    # Three velocities vs time (0)
    plot(sol, vars=[(0,4), (0,5),(0,6)])
end
#=
@testset "Plot ODE system single pendulum" begin
    # Second order non-linear, undamped pendulum:
    # mLθ'' + mg∙sin(θ)
    # where m is mass, unit kg
    #       L is length, unit m
    #       θ is angle to vertical, strictly no unit
    #       g is gravity's acceleration
    # Acceleration
    # θ'' + k * sin(θ) = 0
    #    where k = g / L , unit(k) = s⁻²
    # Reduced to first order non-linear system by defining:
    # θ = y[1]
    # θ' = y[2]
    # It follows that
    # θ'' = y'[2]
    # and
    # y'[1] = y[2]
    # y'[2] = -k * sin(y[1])
    # To solve numerically, we define the derivative.
    # For efficiency, it's evaluated in place, i.e. dy changes at every call.
    function f!(dy, y, k, t)
        dy[1] = y[2]
        dy[2] = -k ∙ sin(y[1])
    end
    # Fixed parameter
    k = upreferred(g / 1.0m)
    # Initial conditions
    θ0 = 0.1
    θ0_vel = 0.0s⁻¹
    θ0_acc = -k *sin(θ0)

    y0 = [θ0, θ0_vel]
    y0 = ArrayPartition([θ0_vel, θ0_acc])
    tspan = (0.0, 100.0)s
    prob = ODEProblem(f!, y0, tspan, k)
    sol = solve(prob)
    # Plot all variables vs time
    plot(sol)
    # Plot in phase space
    plot(solve(prob), vars=(1,2))
end
=#


using ConservationLawsParticles
using SpecialFunctions, RecursiveArrayTools, DifferentialEquations, Plots

using Random
Random.seed!(42)

# model
V(t, x) = -x^3
W(t, x) = 0.0 # x^2/2 - abs(x)
W′(t, x) = 0.0 # x - sign(x)
mob(ρ) = 1.0 # max((ρ-1)^2 * (2ρ+1), 0)
smodel = DiffusiveSampledModel((V,), ((W′,),), (mob,), (Diffusion(1),))
imodel = DiffusiveIntegratedModel((V,), ((W,),), (mob,), (Diffusion(1),))

# initial condition
n = 200
x0 = ArrayPartition(vcat(range(-1, 1; length=n+1)))
x0 = ArrayPartition(sort(randn(n+1)))

refsol(t, x) = sqrt(2)/gamma(1/4) * exp(-x^4/4)

# time span
tspan = (0., 2.0)

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

abstol = reltol = 2e-8

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

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

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


import LessOLS: normal_sampler, ols, show, equation

# create sample with known properties
ns = normal_sampler(a=0, b=20, β_0=5, β=[1, 2.5], sd_e=0.5)
sample1 = ns(100000)

# attempt a linear model
lm1 = ols(sample1, intercept=false)
show(lm1)

# check the coefficients are about right (for a model with zero intercept)
function almost_equals(a, b, threshold)
    diff = map(abs, (a-b))
    return all(diff .< threshold)
end
@assert almost_equals(lm1.beta, [1.2143, 2.713], 0.01)


function op_K( pw::PWGrid, psi::Array{ComplexF64,2} )
    out = zeros(ComplexF64,size(psi))
    Ncol = size(psi,2)
    Ω  = pw.Ω
    G2 = pw.G2
    Npoints = pw.Npoints
    for is = 1:Ncol
        for ip = 1:Npoints
            out[ip,is] = psi[ip,is]*G2[ip]
        end
    end
    return 0.5*out # two minus signs
end


# This file includes an example file for HInterp1D interpolation.

using FractalTools
using Makie

# Generate data
xi, dx, xf = 0., 0.1, 10.
x = collect(xi : dx : xf)
y = sin.(x)
z = cos.(x)
pts = collect.(zip(x, y, z))
freevars = [0.01 * rand() * ones(2,2) for i in 1 : length(x) - 1]
interp = interpolate(pts, HInterp1D(freevars))

# Calculate interpolant
xt = collect(xi : 0.1dx : xf)
_yt = interp.(xt)
yt = getindex.(_yt, 1)

# Plot interpolation
fig, ax, plt = lines(xt, yt, color=yt, linewidth=5)
scatter!(ax, x, y, markersize=5, color=:black)
fig


using Base.Test
using QuantLib

cal = QuantLib.Time.TargetCalendar()
settlementDate = Date(2002, 2, 19)
todays_date = Date(2002, 2, 15)
set_eval_date!(settings, todays_date)

swaptionMats = [Dates.Year(1), Dates.Year(2), Dates.Year(3), Dates.Year(4), Dates.Year(5)]
swaptionVols = [0.1490, 0.1340, 0.1228, 0.1189, 0.1148, 0.1290, 0.1201, 0.1146, 0.1108,
                0.1040, 0.1149, 0.1112, 0.1070, 0.1010, 0.0957, 0.1047, 0.1021, 0.0980, 0.0951,
                0.1270, 0.1000, 0.0950, 0.0900, 0.1230, 0.1160]
swaptionLengths = [Dates.Year(1), Dates.Year(2), Dates.Year(3), Dates.Year(4), Dates.Year(5)]

flat_rate = Quote(0.04875825)
rhTermStructure = FlatForwardTermStructure(settlementDate, cal, flat_rate, QuantLib.Time.Actual365())

fixedLegFrequency = QuantLib.Time.Annual()
fixedLegConvention = QuantLib.Time.Unadjusted()
floatingLegConvention = QuantLib.Time.ModifiedFollowing()
fixedLegDayCounter = QuantLib.Time.EuroThirty360()
floatingLegFrequency = QuantLib.Time.Semiannual()

swapType = Payer()
dummyFixedRate = 0.03
indexSixMonths = euribor_index(QuantLib.Time.TenorPeriod(Dates.Month(6)), rhTermStructure)

startDate = QuantLib.Time.advance(Dates.Year(1), cal, settlementDate, floatingLegConvention)
maturity = QuantLib.Time.advance(Dates.Year(5), cal, startDate, floatingLegConvention)

fixedSchedule = QuantLib.Time.Schedule(startDate, maturity, QuantLib.Time.TenorPeriod(fixedLegFrequency),
                fixedLegConvention, fixedLegConvention, QuantLib.Time.DateGenerationForwards(), false, cal)
floatSchedule = QuantLib.Time.Schedule(startDate, maturity, QuantLib.Time.TenorPeriod(floatingLegFrequency),
                floatingLegConvention, floatingLegConvention, QuantLib.Time.DateGenerationForwards(), false, cal)

swap = VanillaSwap(swapType, 1000.0, fixedSchedule, dummyFixedRate, fixedLegDayCounter,
      indexSixMonths, 0.0, floatSchedule, indexSixMonths.dc, DiscountingSwapEngine(rhTermStructure))

fixedATMRate = fair_rate(swap)

atmSwap = VanillaSwap(swapType, 1000.0, fixedSchedule, fixedATMRate, fixedLegDayCounter, indexSixMonths,
          0.0, floatSchedule, indexSixMonths.dc, DiscountingSwapEngine(rhTermStructure))

modelG2 = G2(rhTermStructure)

numRows = 5
numCols = 5

times = zeros(0)
swaptions = Vector{SwaptionHelper}(numRows)

for i = 1:numRows
  j = numCols - (i - 1)
  k = (i - 1) * numCols + j

  sh = SwaptionHelper(swaptionMats[i], swaptionLengths[j], Quote(swaptionVols[k]), indexSixMonths,
        indexSixMonths.tenor, indexSixMonths.dc, indexSixMonths.dc, rhTermStructure,
        G2SwaptionEngine(modelG2, 6.0, 16))

  times = add_times_to!(sh, times)
  swaptions[i] = sh
end

tg = QuantLib.Time.TimeGrid(times, 30)

om = QuantLib.Math.LevenbergMarquardt()
calibrate!(modelG2, swaptions, om, QuantLib.Math.EndCriteria(400, 100, 1.0e-8, 1.0e-8, 1.0e-8))

for i=1:numRows
  j = numCols - (i - 1)
  k = (i - 1) * numCols + j

  npv = model_value!(swaptions[i])
  implied = implied_volatility!(swaptions[i], npv, 1e-4, 1000, 0.05, 0.50)
  diff = implied - swaptionVols[k]
end

swapLeg = swap.legs[1] # Fixed Leg

bermudanDates = Vector{Date}(length(swapLeg.coupons))
for i=1:length(swapLeg.coupons)
bermudanDates[i]  = accrual_start_date(swapLeg.coupons[i])
end

bermudanExercise = BermudanExercise(bermudanDates)

bermudanSwaption = Swaption(atmSwap, bermudanExercise)

bermudanSwaption = update_pricing_engine(bermudanSwaption, TreeSwaptionEngine(modelG2, 50))

@test npv(bermudanSwaption) == 14.11010300956427


using LinearAlgebra
using StatsBase

@noinline function compute_cordf(df, columns)
	result = DataFrame(col=columns)
	function compute_cor(name_1, name_2)
		col_1 = df[!, name_1]
		col_2 = df[!, name_2]
		StatsBase.cor(col_1, col_2)
	end
	transformations = map(columns) do name_1
		:col => ByRow(name_2 -> compute_cor(name_1, name_2)) => Symbol(name_1)
	end
	transform(result, transformations...)
end

function find_multicolinearity(cordf, ycol, threshold::Number)
	to_be_removed = String[]
	sizehint!(to_be_removed, size(cordf, 1))
	considers = setdiff(cordf.col, ["SalePrice"])
	name_index = String.(cordf.col)
	df = select(cordf, Not("col"))
	for n1 in considers, n2 in considers
		if n1 in to_be_removed || n2 in to_be_removed || n1 == n2
			continue
		end
		row_ind = findfirst(n1 .== name_index)
		col_ind = findfirst(n2 .== name_index)
		if abs(df[row_ind, col_ind]) > threshold
			append!(to_be_removed, [abs(df[row_ind, ycol]) > abs(df[col_ind, ycol]) ? n2 : n1])
		end
	end
	to_be_removed
end

function find_multicolinearity(df, columns, ycol; threshold_step=0.01, threshold_max=1.0, threshold_min=0.5)
	t_range = reverse(collect(range(start=threshold_min, stop=threshold_max, step=threshold_step)))
	df = select(df, columns)
	cordf = compute_cordf(df, columns)
	for t in t_range
		removed = find_multicolinearity(cordf, ycol, t)
		df2 = select(df, Not(removed))
		if rank(Matrix(df2)) == length(names(df2))
			return removed, t
		end
	end
	String[], threshold_min
end


using Test, Random, LinearAlgebra, LuxurySparse

using Yao
using YaoArrayRegister: invorder
using YaoExtensions
using FFTW

@testset "QFT" begin
    num_bit = 5
    fftblock = QFTCircuit(num_bit)
    ifftblock = fftblock'
    reg = rand_state(num_bit)
    rv = copy(statevec(reg))

    @test Matrix(mat(chain(3, QFTCircuit(3) |> adjoint, QFTCircuit(3)))) ≈ IMatrix(1<<3)

    # test ifft
    reg1 = apply!(copy(reg), ifftblock)

    # permute lines (Manually)
    kv = fft(statevec(reg))/sqrt(length(rv))
    @test statevec(reg1) ≈ invorder(kv)

    # test fft
    reg2 = apply!(invorder!(copy(reg)), fftblock)
    kv = ifft(rv) * sqrt(length(rv))
    @test statevec(reg2) ≈ kv
end

@testset "QFTBlock" begin
    num_bit = 5
    qft = QFTCircuit(num_bit)
    iqft = adjoint(qft)
    qftblock = QFTBlock{num_bit}()
    iqftblock = QFTBlock{num_bit}() |> adjoint
    @test openbox(qftblock) == qft
    @test openbox(iqftblock) == iqft
    reg = rand_state(num_bit)

    @test Matrix(mat(chain(3, QFTBlock{3}() |> adjoint, QFTBlock{3}()))) ≈ IMatrix(1<<3)

    # permute lines (Manually)
    @test apply!(copy(reg), iqft) ≈ apply!(copy(reg), QFTBlock{num_bit}() |> adjoint)

    # test fft
    @test apply!(copy(reg), qft) ≈ apply!(copy(reg), qftblock)

    # regression test for nactive
    @test apply!(focus!(copy(reg), 1:3), QFTBlock{3}()) |> isnormalized
end


# calculate pairwise distance between every point in a vector

using CUDAdrv, CUDAnative
const CuArray = CUDAnative.CuHostArray  # real applications: use CuArrays.jl

function haversine_cpu(lat1::Float32, lon1::Float32, lat2::Float32, lon2::Float32, radius::Float32)
    c1 = cospi(lat1 / 180.0f0)
    c2 = cospi(lat2 / 180.0f0)
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    d1 = sinpi(dlat / 360.0f0)
    d2 = sinpi(dlon / 360.0f0)
    t = d2 * d2 * c1 * c2
    a = d1 * d1 + t
    c = 2.0f0 * asin(min(1.0f0, sqrt(a)))
    return radius * c
end

function pairwise_dist_cpu(lat::Vector{Float32}, lon::Vector{Float32})
    # allocate
    n = length(lat)
    rowresult = Array{Float32}(undef, n, n)

    # brute force fill in each cell
    for i in 1:n, j in 1:n
        @inbounds rowresult[i, j] = haversine_cpu(lat[i], lon[i], lat[j], lon[j] , 6372.8f0)
    end

    return rowresult
end

# from https://devblogs.nvidia.com/parallelforall/fast-great-circle-distance-calculation-cuda-c/
function haversine_gpu(lat1::Float32, lon1::Float32, lat2::Float32, lon2::Float32, radius::Float32)
    # XXX: need to prefix math intrinsics with CUDAnative
    c1 = CUDAnative.cospi(lat1 / 180.0f0)
    c2 = CUDAnative.cospi(lat2 / 180.0f0)
    dlat = lat2 - lat1
    dlon = lon2 - lon1
    d1 = CUDAnative.sinpi(dlat / 360.0f0)
    d2 = CUDAnative.sinpi(dlon / 360.0f0)
    t = d2 * d2 * c1 * c2
    a = d1 * d1 + t
    c = 2.0f0 * CUDAnative.asin(CUDAnative.min(1.0f0, CUDAnative.sqrt(a)))
    return radius * c
end

# pairwise distance calculation kernel
function pairwise_dist_kernel(lat::CuDeviceVector{Float32}, lon::CuDeviceVector{Float32},
                              rowresult::CuDeviceMatrix{Float32}, n)
    i = (blockIdx().x-1) * blockDim().x + threadIdx().x
    j = (blockIdx().y-1) * blockDim().y + threadIdx().y

    if i <= n && j <= n
        # store to shared memory
        shmem = @cuDynamicSharedMem(Float32, 2*blockDim().x + 2*blockDim().y)
        if threadIdx().y == 1
            shmem[threadIdx().x] = lat[i]
            shmem[blockDim().x + threadIdx().x] = lon[i]
        end
        if threadIdx().x == 1
            shmem[2*blockDim().x + threadIdx().y] = lat[j]
            shmem[2*blockDim().x + blockDim().y + threadIdx().y] = lon[j]
        end
        sync_threads()

        # load from shared memory
        lat_i = shmem[threadIdx().x]
        lon_i = shmem[blockDim().x + threadIdx().x]
        lat_j = shmem[2*blockDim().x + threadIdx().y]
        lon_j = shmem[2*blockDim().x + blockDim().y + threadIdx().y]

        @inbounds rowresult[i, j] = haversine_gpu(lat_i, lon_i, lat_j, lon_j, 6372.8f0)
    end

    return
end

function pairwise_dist_gpu(lat::Vector{Float32}, lon::Vector{Float32})
    # upload
    lat_gpu = CuArray(lat)
    lon_gpu = CuArray(lon)

    # allocate
    n = length(lat)
    rowresult_gpu = CuArray(zeros(Float32, n, n))

    # calculate launch configuration
    function get_config(kernel)
        # calculate a 2D block size from the suggested 1D configuration
        # NOTE: we want our launch configuration to be as square as possible,
        #       because that minimizes shared memory usage
        function get_threads(threads)
            threads_x = floor(Int, sqrt(threads))
            threads_y = threads ÷ threads_x
            return (threads_x, threads_y)
        end

        # calculate the amount of dynamic shared memory for a 2D block size
        get_shmem(threads) = 2 * sum(threads) * sizeof(Float32)

        fun = kernel.fun
        config = launch_configuration(fun, shmem=threads->get_shmem(get_threads(threads)))

        # convert to 2D block size and figure out appropriate grid size
        threads = get_threads(config.threads)
        blocks = ceil.(Int, n ./ threads)
        shmem = get_shmem(threads)

        return (threads=threads, blocks=blocks, shmem=shmem)
    end

    @cuda config=get_config pairwise_dist_kernel(lat_gpu, lon_gpu, rowresult_gpu, n)

    return Array(rowresult_gpu)
end

using Test

# generate reasonable data
function main(n = 1000)
    lat = rand(Float32, n) .* 45
    lon = rand(Float32, n) .* -120

    @test pairwise_dist_cpu(lat, lon) ≈ pairwise_dist_gpu(lat, lon) rtol=1e-2
end
main()


using Random, SparseArrays

@testset "SparseMatrixCSC with matrices" begin
    rng = MersenneTwister(123)
    A = sprand(rng, 10, 5, 0.2)
    C = zeros(10, 2)
    B = randn(5, 2)

    cols = 1:5
    colsmul!(C, A, B, cols)
    @test C ≈ A[:, cols]*B[cols, :]

    cols = 2:4
    colsmul!(C, A, B, cols)
    @test C ≈ A[:, cols]*B[cols, :]

    cols = [1, 5]
    colsmul!(C, A, B, cols)
    @test C ≈ A[:, cols]*B[cols, :]

    α, β = rand(), rand()
    Cc = copy(C)
    colsmul!(C, A, B, cols, α, β)
    @test C ≈ α.*A[:, cols]*B[cols, :] .+ β.*Cc
end

@testset "Transposed SparseMatrixCSC with matrices" begin
    rng = MersenneTwister(123)
    A = sprand(rng, 10, 5, 0.2)
    C = zeros(5, 2)
    B = randn(10, 2)

    cols = 1:5
    tcolsmul!(C, A, B, cols)
    @test C[cols, :] ≈ A[:, cols]'*B

    cols = 2:4
    tcolsmul!(C, A, B, cols)
    @test C[cols, :] ≈ A[:, cols]'*B

    cols = [1, 5]
    tcolsmul!(C, A, B, cols)
    @test C[cols, :] ≈ A[:, cols]'*B

    α, β = rand(), rand()
    Cc = copy(C)
    tcolsmul!(C, A, B, cols, α, β)
    @test C[cols, :] ≈ α.*A[:, cols]'*B .+ β.*Cc[cols, :]
end

@testset "SparseMatrixCSC with vectors" begin
    rng = MersenneTwister(123)
    A = sprand(rng, 10, 5, 0.2)
    c = zeros(10)
    b = randn(5)

    cols = 1:5
    colsmul!(c, A, b, cols)
    @test c ≈ A[:, cols]*b[cols]

    cols = 2:4
    colsmul!(c, A, b, cols)
    @test c ≈ A[:, cols]*b[cols]

    cols = [1, 5]
    colsmul!(c, A, b, cols)
    @test c ≈ A[:, cols]*b[cols]

    α, β = rand(), rand()
    cc = copy(c)
    colsmul!(c, A, b, cols, α, β)
    @test c ≈ α.*A[:, cols]*b[cols] .+ β.*cc
end

@testset "Matrix with matrices" begin
    A = randn(10, 5)
    C = zeros(10, 2)
    B = randn(5, 2)

    cols = 1:5
    colsmul!(C, A, B, cols)
    @test C ≈ A[:, cols]*B[cols, :]

    cols = 2:4
    colsmul!(C, A, B, cols)
    @test C ≈ A[:, cols]*B[cols, :]

    cols = [1, 5]
    colsmul!(C, A, B, cols)
    @test C ≈ A[:, cols]*B[cols, :]

    α, β = rand(), rand()
    Cc = copy(C)
    colsmul!(C, A, B, cols, α, β)
    @test C ≈ α.*A[:, cols]*B[cols, :] .+ β.*Cc
end

@testset "Transposed Matrix with matrices" begin
    A = randn(10, 5)
    C = zeros(5, 2)
    B = randn(10, 2)

    cols = 1:5
    tcolsmul!(C, A, B, cols)
    @test C[cols, :] ≈ A[:, cols]'*B

    cols = 2:4
    tcolsmul!(C, A, B, cols)
    @test C[cols, :] ≈ A[:, cols]'*B

    cols = [1, 5]
    tcolsmul!(C, A, B, cols)
    @test C[cols, :] ≈ A[:, cols]'*B

    α, β = rand(), rand()
    Cc = copy(C)
    tcolsmul!(C, A, B, cols, α, β)
    @test C[cols, :] ≈ α.*A[:, cols]'*B .+ β.*Cc[cols, :]
end


abstract type EasingType end

struct Easing{T <: EasingType}
    easing::T
    n::Int
    yoyo::Bool
    prewait::Float64
    postwait::Float64
end

struct FuncEasing <: EasingType
    f::Function
    args::Tuple

    function FuncEasing(f, args)
        f0 = f(0, args...)
        f1 = f(1, args...)
        if !(f0 == 0 && f1 == 1)
            error("The given easing function is not 0 and 1 at t = 0 and t = 1, but $f0 and $f1")
        end
        new(f, args)
    end
end

FuncEasing(f::Function) = FuncEasing(f, ())

"""
Mixes two easings with a constant ratio r, so that final = e1 * r + e2 * (1 - r)
"""
struct MixedEasing{S <: Easing, T <: Easing} <: EasingType
    e1::S
    e2::T
    mix::Float64

    function MixedEasing(e1::S, e2::T, mix::Float64=0.5) where {S <: Easing, T <: Easing}
        if !(0 <= mix <= 1)
            error("Mix has to be between 0 and 1, but is $mix")
        end
        new{S, T}(e1, e2, mix)
    end
end

"""
Multiplies two easings with each other, so final = e1 * e2
"""
struct MultipliedEasing{S <: Easing, T <: Easing} <: EasingType
    e1::S
    e2::T
end

"""
Mixes two easings with a factor decided with a third easing, so that
final = (1 - easing()) * e1 + easing() * e2
"""
struct EasedEasing{S <: Easing, T <: Easing, U <: Easing} <: EasingType
    e1::S
    e2::T
    easing::U

    function EasedEasing(e1, e2, easing=Easing(easing=LinearEasing()))
        new{typeof(e1), typeof(e2), typeof(easing)}(e1, e2, easing)
    end
end

Easing(easing = FuncEasing(f_linear); n=1, yoyo=false, prewait=0.0, postwait=0.0) = Easing(easing, n, yoyo, prewait, postwait)

function opposite(f)
    function opfunc(fraction, args...)
        1 - f(1 - fraction, args...)
    end
end

funcease(f, args...; kwargs...) = Easing(FuncEasing(f, args); kwargs...)

f_noease(fraction) = fraction == 1 ? 1 : 0
noease(;kwargs...) = funcease(f_noease; kwargs...)

f_stepease(fraction, stepfraction) = fraction < stepfraction ? 0 : 1
stepease(stepfraction; kwargs...) = funcease(f_stepease, stepfraction; kwargs...)

f_sineio(fraction) = sin(pi * fraction - 0.5pi) * 0.5 + 0.5
sineio(;kwargs...) = funcease(f_sineio; kwargs...)

f_saccadic(fraction, power) = -(sin((-fraction + 1) ^ power * pi - pi/2) * 0.5 + 0.5) + 1
saccadic(power; kwargs...) = funcease(f_saccadic, power; kwargs...)

f_linear(fraction) = fraction
linear(;kwargs...) = funcease(f_linear; kwargs...)

f_polyin(fraction, power) = fraction ^ power
polyin(power; kwargs...) = funcease(f_polyin, power; kwargs...)

polyout(power; kwargs...) = funcease(opposite(f_polyin), power; kwargs...)

f_polyio(fraction, power) = (
    fraction < 0.5 ?
    f_polyin(2 * fraction, power) / 2 :
    opposite(f_polyin)(2 * (fraction - 0.5), power) / 2 + 0.5)
polyio(power; kwargs...) = funcease(f_polyio, power; kwargs...)

f_expin(fraction, exponent) = ((exponent ^ fraction) - 1) / (exponent - 1)
expin(exponent; kwargs...) = funcease(f_expin, exponent; kwargs...)

expout(exponent; kwargs...) = funcease(opposite(f_expin), exponent; kwargs...)

mixed(e1, e2; kwargs...) = Easing(MixedEasing(e1, e2); kwargs...)
multiplied(e1, e2; kwargs...) = Easing(MultipliedEasing(e1, e2); kwargs...)
eased(e1, e2, with_e; kwargs...) = Easing(EasedEasing(e1, e2, with_e); kwargs...)

Base.Broadcast.broadcastable(e::Easing) = Ref(e)


# # Postprocessing
#
# ![](heat_square_fluxes.png)

# ## Introduction
#
# After running a simulation, we usually want to visualize the results in different ways.
# The `L2Projector` and the `PointEvalHandler` build a pipeline for doing so. With the `L2Projector`,
# integration point quantities can be projected to the nodes. The `PointEvalHandler` enables evaluation of
# the finite element approximated function in any coordinate in the domain. Thus with the combination of both functionalities,
# both nodal quantities and integration point quantities can be evaluated in any coordinate, allowing for example
# cut-planes through 3D structures or cut-lines through 2D-structures.
#
# This example continues from the Heat equation example, where the temperature field was
# determined on a square domain. In this example, we first compute the heat flux in each
# integration point (based on the solved temperature field) and then we do an L2-projection
# of the fluxes to the nodes of the mesh. By doing this, we can more easily visualize
# integration points quantities. Finally, we visualize the temperature field and the heat fluxes along a cut-line.
#
# The L2-projection is defined as follows: Find projection ``q(\boldsymbol{x}) \in L_2(\Omega)`` such that
# ```math
# \int v q \ \mathrm{d}\Omega = \int v d \ \mathrm{d}\Omega \quad \forall v \in L_2(\Omega),
# ```
# where ``d`` is the quadrature data to project. Since the flux is a vector the projection function
# will be solved with multiple right hand sides, e.g. with ``d = q_x`` and ``d = q_y`` for this 2D problem.
#
# Ferrite has functionality for doing much of this automatically, as displayed in the code below.
# In particular [`L2Projector`](@ref) for assembling the left hand side, and
# [`project`](@ref) for assembling the right hand sides and solving for the projection.

# ## Implementation
#
# Start by simply running the Heat equation example to solve the problem
include("heat_equation.jl");

# Next we define a function that computes the heat flux for each integration point in the domain.
# Fourier's law is adopted, where the conductivity tensor is assumed to be isotropic with unit
# conductivity ``\lambda = 1 ⇒ q = - \nabla u``, where ``u`` is the temperature.
function compute_heat_fluxes(cellvalues::CellScalarValues{dim,T}, dh::DofHandler, a) where {dim,T}

    n = getnbasefunctions(cellvalues)
    cell_dofs = zeros(Int, n)
    nqp = getnquadpoints(cellvalues)

    ## Allocate storage for the fluxes to store
    q = [Vec{2,T}[] for _ in 1:getncells(dh.grid)]

    for (cell_num, cell) in enumerate(CellIterator(dh))
        q_cell = q[cell_num]
        celldofs!(cell_dofs, dh, cell_num)
        aᵉ = a[cell_dofs]
        reinit!(cellvalues, cell)

        for q_point in 1:nqp
            q_qp = - function_gradient(cellvalues, q_point, aᵉ)
            push!(q_cell, q_qp)
        end
    end
    return q
end
#md nothing # hide

# Now call the function to get all the fluxes.
q_gp = compute_heat_fluxes(cellvalues, dh, u);

# Next, create an [`L2Projector`](@ref) using the same interpolation as was used to approximate the
# temperature field. On instantiation, the projector assembles the coefficient matrix `M` and
# computes the Cholesky factorization of it. By doing so, the projector can be reused without
# having to invert `M` every time.
projector = L2Projector(ip, grid);

# Project the integration point values to the nodal values
q_projected = project(projector, q_gp, qr; project_to_nodes=false); # TODO: this should be default.

# ## Exporting to VTK
# To visualize the heat flux, we export the projected field `q_projected`
# to a VTK-file, which can be viewed in e.g. [ParaView](https://www.paraview.org/).
vtk_grid("heat_equation_flux", grid) do vtk
    vtk_point_data(vtk, projector, q_projected, "q")
end;

# ## Point Evaluation
# ![](heat_square_pointevaluation.png)

# Consider a cut-line through the domain, like the black line in the figure above.
# We will evaluate the temperature and the heat flux distribution along a horizontal line.
points = [Vec((x, 0.75)) for x in range(-1.0, 1.0, length=101)];

# First, we need to generate a `PointEvalHandler`. This will find and store the cells
# containing the input points.
ph = PointEvalHandler(grid, points);

# After the L2-Projection, the heat fluxes `q_projected` are stored in the DoF-ordering
# determined by the projector's internal DoFHandler, so to evalute the flux `q` at our points
# we give the `PointEvalHandler`, the `L2Projector` and the values `q_projected`.
q_points = get_point_values(ph, projector, q_projected);

# We can also extract the field values, here the temperature, right away from the result
# vector of the simulation, that is stored in `u`. These values are stored in the order of
# our initial DofHandler so the input is not the `PointEvalHandler`, the original `DofHandler`,
# the dof-vector `u`, and (optionally for single-field problems) the name of the field.
# from the `L2Projection`, the values are stored in the order of the degrees of freedom.
u_points = Ferrite.get_point_values(ph, dh, u, :u);

# Now, we can plot the temperature and flux values with the help of any plotting library, e.g. Plots.jl.
# To do this, we need to import the package:
import Plots

# Firstly, we are going to plot the temperature values along the given line.
Plots.plot(getindex.(points,1), u_points, xlabel="x (coordinate)", ylabel="u (temperature)", label=nothing)

# Secondly, the horizontal heat flux (i.e. the first component of the heat flux vector) is plotted.
Plots.plot(getindex.(points,1), getindex.(q_points,1), xlabel="x (coordinate)", ylabel="q_x (flux in x-direction)", label=nothing)

#md # ## [Plain Program](@id postprocessing-plain-program)
#md #
#md # Below follows a version of the program without any comments.
#md # The file is also available here: [postprocessing.jl](postprocessing.jl)
#md #
#md # ```julia
#md # @__CODE__
#md # ```


export WrightSampler

#=
This document implements a WrightSampler mutable struct type.
This object is used to generate Gibbs-sampler output on Wright equilibria
Also included are methods of the current state to generate stochastic gradients
And fisher information estimates
 =#

# This lists the neccesary parameters
# for sampling from the Wright-equilibrium distribution. The partition function is the most costly to compute and so it cahched
# θ_mut, θ_wt, σ are forward and backward diversity parameters respectively. At the moment, there is only code for  σ>0.
mutable struct WrightSampler
    θ_wt::Float64
    θ_mut::Float64
    σ::Float64
    k::Int64   # auxiliary variable for gibbs sampling
    x::Float64 # The freqeuncy of wt
end

function WrightSampler(θ_wt, θ_mut, σ)
#    z = gamma(θ_wt) * HypergeometricFunctions.pFq([θ_wt],[θ_mut + θ_wt],σ) /gamma(θ_wt +θ_mut)
    ws = WrightSampler(θ_wt, θ_mut, σ, 0, .5)
    for ii in 1:10
    incrementk!(ws)
    incrementx!(ws)
    end
    return ws
end

# This sampler is robuset for small-exponent gamma distributions
# Code was copied from somewhere unknown (citation needed)
function randlogGamma(a;scale = 1)
    if a > .2
        return log.(rand(Gamma(a)))
    else
        L = 1 / a - 1
        w = a / MathConstants.e / (1 - a)
        ww = 1 / (1 + w)
        η = z -> ( z >= 0 ? exp(-z) : w * L * exp(L * z))
        h = z -> exp(-z - exp(-z / a))
        function rh(a)
            while true
                U = rand()
                z = (U <= ww) ? -log(U / ww) : log(rand()) / L
                h(z) / η(z) > rand() && return(z)
            end
        end
        return log(scale) - rh(a) / a
    end
end

function sigmoid(x) # numerically stable sigmoid
    if x >= 0
        1 / ( 1 + exp(-x) )
    else
        exp(x) / ( 1 + exp(x) )
    end
end

function incrementk!(ws::WrightSampler)
    if ws.σ >= 0
        ws.k = rand(Poisson(abs(ws.σ * ws.x)))
    else
        ws.k = rand(Poisson(abs(ws.σ * (1-ws.x))))
    end
end

function incrementx!(ws::WrightSampler)
    if ws.σ >= 0
        a = randlogGamma(ws.θ_wt + ws.k;scale = 1)
        b = randlogGamma(ws.θ_mut;scale = 1)
    else
        a = randlogGamma(ws.θ_wt;scale = 1)
        b = randlogGamma(ws.θ_mut + ws.k;scale = 1)
    end
    ws.x = sigmoid(a-b)
end

function sample!(ws::WrightSampler)
    for ii in 1:5
    incrementk!(ws)
    incrementx!(ws)
    end
    return ws.x
end

function rejection_sample_k(θ_wt, θ_mut,σ)
    k = 0
    while true
        k = rand(Poisson(σ))
        k == 0 ? break : nothing
        p = exp( loggamma(θ_wt + θ_mut) + loggamma(k + θ_wt) - loggamma(k + θ_wt + θ_mut) -loggamma(θ_wt))
        rand(Bernoulli(p)) ? break : nothing
    end
    return k
end


"""
θ-method ODE solver
"""
struct ThetaMethod <: ODESolver
  nls::NonlinearSolver
  dt::Float64
  θ::Float64
  function ThetaMethod(nls,dt,θ)
    if θ > 0.0
      return new(nls,dt,θ)
    else
      return ForwardEuler(nls,dt)
    end
  end
end

BackwardEuler(nls,dt) = ThetaMethod(nls,dt,1.0)
MidPoint(nls,dt) = ThetaMethod(nls,dt,0.5)

function solve_step!(uf::AbstractVector,
                     solver::ThetaMethod,
                     op::ODEOperator,
                     u0::AbstractVector,
                     t0::Real,
                     cache) # -> (uF,tF)

  dt = solver.dt
  solver.θ == 0.0 ? dtθ = dt : dtθ = dt*solver.θ
  tθ = t0+dtθ

  if cache === nothing
    ode_cache = allocate_cache(op)
    vθ = similar(u0)
    nl_cache = nothing
  else
    ode_cache, vθ, nl_cache = cache
  end

  ode_cache = update_cache!(ode_cache,op,tθ)

  nlop = ThetaMethodNonlinearOperator(op,tθ,dtθ,u0,ode_cache,vθ)

  nl_cache = solve!(uf,solver.nls,nlop,nl_cache)

  if 0.0 < solver.θ < 1.0
    uf = uf*(1.0/solver.θ)-u0*((1-solver.θ)/solver.θ)
  end

  cache = (ode_cache, vθ, nl_cache)

  tf = t0+dt
  return (uf,tf,cache)

end

"""
Nonlinear operator that represents the θ-method nonlinear operator at a
given time step, i.e., A(t,u_n+θ,(u_n+θ-u_n)/dt)
"""
struct ThetaMethodNonlinearOperator <: NonlinearOperator
  odeop::ODEOperator
  tθ::Float64
  dtθ::Float64
  u0::AbstractVector
  ode_cache
  vθ::AbstractVector
end

function residual!(b::AbstractVector,op::ThetaMethodNonlinearOperator,x::AbstractVector)
  uθ = x
  vθ = op.vθ
  vθ = (x-op.u0)/op.dtθ
  residual!(b,op.odeop,op.tθ,(uθ,vθ),op.ode_cache)
end

function jacobian!(A::AbstractMatrix,op::ThetaMethodNonlinearOperator,x::AbstractVector)
  uF = x
  vθ = op.vθ
  vθ = (x-op.u0)/op.dtθ
  z = zero(eltype(A))
  fillstored!(A,z)
  jacobians!(A,op.odeop,op.tθ,(uF,vθ),(1.0,1/op.dtθ),op.ode_cache)
end

function allocate_residual(op::ThetaMethodNonlinearOperator,x::AbstractVector)
  allocate_residual(op.odeop,x,op.ode_cache)
end

function allocate_jacobian(op::ThetaMethodNonlinearOperator,x::AbstractVector)
  allocate_jacobian(op.odeop,x,op.ode_cache)
end

function zero_initial_guess(op::ThetaMethodNonlinearOperator)
  x0 = similar(op.u0)
  fill!(x0,zero(eltype(x0)))
  x0
end


using Crux
using Test
using POMDPModels

## Spaces
mdp = SimpleGridWorld()

s1 = DiscreteSpace(5)
@test s1 isa AbstractSpace
@test s1 isa DiscreteSpace
@test s1.N == 5
@test type(s1) == Bool
@test Crux.dim(s1) == (5,)
@test s1.vals == [1,2,3,4,5]
@test tovec(4, s1) == Bool[0,0,0,1,0]

s2 = DiscreteSpace(5, [:a, :b, :c, :d, :e])
@test s2 isa AbstractSpace
@test s2 isa DiscreteSpace
@test s2.N == 5
@test type(s2) == Bool
@test Crux.dim(s2) == (5,)
@test s2.vals == [:a, :b, :c, :d, :e]
@test tovec(:d, s2) == Bool[0,0,0,1,0]

s3 = ContinuousSpace((3,4), Float64)
@test s3 isa AbstractSpace
@test s3 isa ContinuousSpace
@test s3.dims == (3,4)
@test type(s3) == Float64
@test Crux.dim(s3) == (3,4)
@test tovec(zeros(3,4), s3) == zeros(3, 4)

s3 = ContinuousSpace((3,4), Float64, 1f0, 2f0)
@test s3 isa AbstractSpace
@test s3 isa ContinuousSpace
@test s3.dims == (3,4)
@test type(s3) == Float64
@test Crux.dim(s3) == (3,4)
@test tovec(zeros(3,4),s3) == -0.5*ones(3, 4)

s4 = state_space(mdp)
@test s4 isa ContinuousSpace
@test Crux.dim(s4) == (2,)


@testset "Utilities" begin
  a, b, c = P2(0,0), P2(1,0), P2(0,1)
  @test signarea(a, b, c) == T(0.5)
  a, b, c = P2(0,0), P2(0,1), P2(1,0)
  @test signarea(a, b, c) == T(-0.5)

  p1, p2, p3 = P2(0,0), P2(1,1), P2(0.25,0.5)
  s = Segment(P2(0.5,0.0), P2(0.0,1.0))
  @test sideof(p1, s) == :LEFT
  @test sideof(p2, s) == :RIGHT
  @test sideof(p3, s) == :ON
end


using HilbertSpaces
using Base.Test

S = AngularSpace

i = bra(S)(1, 0)
j = ket(S)(0, 0)
C = op(S)(1, 0)

@testset "Init" begin
    @test_throws ErrorException bra(S)(1, 2)
    @test_throws ErrorException ket(S)(1, 2)
    @test_throws ErrorException op(S)(1, 2)
end

@testset "Comparisons" begin
    @test i == HilbertSpaces.AngKet(1, 0)'
    @test j == HilbertSpaces.AngBra(0, 0)'
    @test C == HilbertSpaces.SphericalTensor(1, 0)
end

@testset "Spaces" begin
    @test space(i) == AngularSpace
    @test space(j) == AngularSpace
    @test space(C) == AngularSpace
    @test op(i) == HilbertSpaces.SphericalTensor
    @test op(j) == HilbertSpaces.SphericalTensor
    @test op(HilbertSpaces.AngularSpace) == HilbertSpaces.SphericalTensor
    @test identity(op(i)) == op(S)(0, 0)
    @test bra(HilbertSpaces.AngularSpace) == HilbertSpaces.AngBra
    @test ket(HilbertSpaces.AngularSpace) == HilbertSpaces.AngKet
    @test scalar(i) == HilbertSpaces.AngScalar
    @test scalar(j) == HilbertSpaces.AngScalar
    @test scalar(HilbertSpaces.AngularSpace) == HilbertSpaces.AngScalar
end

@testset "Products" begin
    @test typeof(i|j) == scalar(S)
    @test typeof((i|C)|j) == scalar(S)
    @test typeof(i|(C|j)) == scalar(S)
    @test float(i|j) == 0
    @test float(i|i') ≈ 1
    @test float(j'|j) ≈ 1

    @test_throws ErrorException C|i
    @test_throws ErrorException j|C
end

@testset "Scalars" begin
    v = (i|C|j)
    @test float(v) ≈ 1/√3
end

@testset "Strings" begin
    @test string(bra(S)(0,0)) == "〈s|"
    @test string(bra(S)(1,0)) == "〈p|"
    @test string(bra(S)(1,1)) == "〈p⁺|"
    @test string(bra(S)(1,-1)) == "〈p⁻|"
    @test string(bra(S)(3,-2)) == "〈f;-2|"
    @test string(bra(S)(3//2, -1//2)) == "〈³⁄₂;⁻¹⁄₂|"
    @test string(ket(S)(1,0)) == "|p〉"

    @test string(op(S)(2,1)) == "C₁⁽²⁾"

    @test string(C|j) == "C₀⁽¹⁾|s〉"
    @test string(i|C) == "〈p|C₀⁽¹⁾"
    @test string(i|C|j) == @sprintf("〈p|C₀⁽¹⁾|s〉 ≈ %0.5g", float(i|C|j))
    @test string(i|j) == "〈p|s〉 ≈ 0"
    @test string(i|i') == "〈p|p〉 ≈ 1"
end


using Test
using DataWrangler: normality

@testset "normality" begin

    x = 1:100
    @test normality(x) ≈ 6.002400480071892

    x = randn(100)
    @test normality(x) < 6

end


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

import Base:*
"""
  `*(A::WFST,B::WFST)`

Returns `C`, the concatenation/product of `A` and `B`. If `A` converts the sequence `a_i` to `a_o` with weight `wa` and `B` converts the sequence `b_i` to `b_o` with weight `wb` then C` converts the sequence `[a_i,b_i]` to `[a_o,b_o]` with weight `wa*wb`.

```julia
julia> sym = ["a","b","c","d","α","β","γ","δ"];

julia> A = WFST(sym);

julia> add_arc!(A,1=>2,"a"=>"α");

julia> add_arc!(A,2=>3,"b"=>"β");

julia> initial!(A,1); final!(A,3,5)
WFST #states: 3, #arcs: 2, #isym: 8, #osym: 8
|1/0.0f0|
a:α/0.0f0 → (2)
(2)
b:β/0.0f0 → (3)
((3/5.0f0))

julia> B = WFST(sym);

julia> add_arc!(B,1=>2,"c"=>"γ");

julia> add_arc!(B,2=>3,"d"=>"δ");

julia> initial!(B,1); final!(B,3,2)
WFST #states: 3, #arcs: 2, #isym: 8, #osym: 8
|1/0.0f0|
c:γ/0.0f0 → (2)
(2)
d:δ/0.0f0 → (3)
((3/2.0f0))

julia> C = A*B
WFST #states: 6, #arcs: 5, #isym: 8, #osym: 8
|1/0.0f0|
a:α/0.0f0 → (2)
(2)
b:β/0.0f0 → (3)
(3)
ϵ:ϵ/5.0f0 → (4)
(4)
c:γ/0.0f0 → (5)
(5)
d:δ/0.0f0 → (6)
((6/2.0f0))

julia> A(["a","b"])
(["α", "β"], 5.0f0)

julia> B(["c","d"])
(["γ", "δ"], 2.0f0)

julia> C(["a","b","c","d"])
(["α", "β", "γ", "δ"], 7.0f0)

```
"""
function *(A::WFST,B::WFST)
  isym = get_isym(A)
  osym = get_osym(A)
  W = typeofweight(A)
  if isym != get_isym(B)
    throw(ErrorException("WFSTs have different input symbol tables"))
  end
  if osym != get_osym(B)
    throw(ErrorException("WFSTs have different output symbol tables"))
  end
  if W != typeofweight(B)
    throw(ErrorException("WFSTs have different type of weigths $W and $(typeofweight(B))"))
  end
  C = deepcopy(A)
  offset = length(C)
  add_states!(C,length(B))
  for (q,arcs) in enumerate(get_states(B))
    for arc in arcs
      ilab, olab = get_ilabel(arc), get_olabel(arc)
      w, n = get_weight(arc), get_nextstate(arc)
      push!(C[q+offset],Arc(ilab,olab,w,n+offset))
    end
  end
  final_A = get_final(A; single=false)
  initial_B = get_initial(B; single=false)
  for (qf,wf) in final_A
    for (qi,wi) in initial_B
      push!(C[qf],Arc(0,0,wf*wi,qi+offset))
    end
    rm_final!(C,qf)
  end
  for (q,w) in get_final(B; single=false)
    final!(C,q+offset,w)
  end
  return C
end


function is_orthogonal(S, tol=1e-6)
    n = LinearAlgebra.checksquare(S)
    for i in 1:n
        for j in 1:n
            if i != j
                s = LinearAlgebra.dot(S[:, i], S[:, j])
                if abs(s) > tol
                    return false
                end
            end
        end
    end
    return true
end

function _projcoef(A, i, v, w = A[i, :])
    return LinearAlgebra.dot(v, w) / LinearAlgebra.dot(w, w)
end
function _proj(A, i, v)
    w = A[i, :]
    return w * _projcoef(A, i, v, w)
end

function orthogonalize(A)
    display(A)
    for i in 2:size(A, 1)
        row = A[i, :]
        for j in 1:(i - 1)
            row = row - _proj(A, j, row)
        end
        A[i, :] = row
    end
    return A
end

struct _OrthogonalMatrix end
struct _RowEchelonMatrix end

function __linsolve(A::Matrix{T}, b::Vector{T}, ::_OrthogonalMatrix) where {T}
    return map(1:size(A, 1)) do i
        return _projcoef(A, i, b)
    end
end

function __linsolve(A::Matrix{T}, b::Vector{T}, ::_RowEchelonMatrix) where {T}
    j = 0
    return map(1:size(A, 1)) do i
        while j < size(A, 2)
            j += 1
            if isone(A[i, j]) && all(k -> k == i || iszero(A[k, j]), 1:size(A, 1))
                return b[j]
            end
        end
        error("Not in row_echelon_form, cannot find for `$i`th entry.")
    end
end

function _linsolve(A::Matrix{T}, b::Vector{T}, form) where {T}
    x = __linsolve(A, b, form)
    #@assert transpose(A) * x == b
    return x
end


ENV["MPLBACKEND"]="agg"

using Test, ExtendableGrids, GridVisualize
import PyPlot

@testset "Basic" begin
    function test_geomspace()
        X1=geomspace(2.0,3.0,0.2,0.2)
        X2=collect(2:0.2:3)
        @assert X1 ≈  X2
        X1=geomspace(2.0,3.0,0.1,0.2)
        X2=geomspace(2.0,3.0,0.2,0.1)

        X1=geomspace(2.0,3.0,1.0e-5,0.2)
        X2=geomspace(2.0,3.0,0.2,1.0e-5)
        (X1[2]-X1[1])  ≈ (X2[end]-X2[end-1]) &&
            (X2[2]-X2[1])  ≈ (X1[end]-X1[end-1])
    end
    @test test_geomspace()

    function test_prepare_edges()
        ## Compared with pdelib; have more of these
        nodes=Matrix([
            0.0 1;
            1 0;
            0 1;
            1 1]')

        cells=Matrix([
            1 2 3;
            2 3 4
        ]')
        cellmat=[1,1]
        bfaces=Matrix([ 1 2;
                        1 3;
                        2 4]')
        bfacemat=[1,1]

        grid=simplexgrid(nodes,cells,cellmat,bfaces,bfacemat)

        grid[CellEdges]==[3 5; 2 4; 1 3] &&
        grid[EdgeNodes]==[2 3 3 4 4; 1 1 2 2 3] &&
            grid[EdgeCells]==[1 1 1 2 2; 0 0 2 0 0]

    end
    @test test_prepare_edges()

    @test let
        g=simplexgrid(0:10)
        bfc=g[BFaceCells]
        bfc[1,1]==1 && bfc[1,2]==10
    end

    @test let
        g=simplexgrid(0:10)
        bfn=g[BFaceNormals]
        bfn==[-1.0 1.0]
    end

    @test let
        g=simplexgrid(0:2, 0:2)
        bfn=g[BFaceNormals]
        bfn==[0.0  -1.0;
              -1.0   0.0;
              0.0  -1.0;
              1.0   0.0;
              0.0   1.0;
              -1.0   0.0;
              1.0   0.0;
              0.0   1.0]'
    end

    @test let
        g=simplexgrid(0:2, 0:2)
        bfc=g[BFaceCells]
        tryfix(bfc)==[1  2  3  3  6  6  7  8;]
    end

end

@testset "Read/Write" begin
    function testrw(grid)
        ftmp=tempname()
        write(ftmp,grid,format="sg")
        grid1=simplexgrid(ftmp,format="sg")
        seemingly_equal(grid1,grid)
    end
    X=collect(0:0.05:1)
    @test testrw(simplexgrid(X))
    @test testrw(simplexgrid(X,X))
    @test testrw(simplexgrid(X,X,X))
end

function testgrid(grid,testdata)
    (num_nodes(grid),num_cells(grid), num_bfaces(grid))==testdata
end

examples1d=joinpath(@__DIR__,"..","examples","examples1d.jl")
include(examples1d)
examples2d=joinpath(@__DIR__,"..","examples","examples2d.jl")
include(examples2d)
examples3d=joinpath(@__DIR__,"..","examples","examples3d.jl")
include(examples3d)
plotting=joinpath(@__DIR__,"..","examples","plotting.jl")
include(plotting)

@testset "1D" begin
    @test testgrid(interval_from_vector(),(21,20,2))
    @test testgrid(interval_localref(),   (27,26,2))
    @test testgrid(interval_multiregion(),(21,20,3))
    @test testgrid(interval_subgrid(),(51,50,2))
end

@testset "2D" begin
    @test testgrid(rectangle(),(441,800,80))
    @test testgrid(rectangle_localref(),(729, 1352, 104))
    @test testgrid(rectangle_multiregion(),(441,800,100))
    @test testgrid(rectangle_subgrid(),(360, 600, 120))
end

@testset "3D" begin
    @test testgrid(quadrilateral(),(330,1200,440))
end

if !Sys.isapple()
    @testset "plotting examples" begin
        include("../docs/makeplots.jl")
        picdir=mktempdir()

        @test makeplot("interval_from_vector",picdir)
        @test makeplot("interval_localref",picdir)
        @test makeplot("interval_multiregion",picdir)
        @test makeplot("interval_subgrid",picdir)
        @test makeplot("rectangle",picdir)
        @test makeplot("rectangle_localref",picdir)
        @test makeplot("rectangle_multiregion",picdir)
        @test makeplot("rectangle_subgrid",picdir)
        @test makeplot("quadrilateral",picdir)
    end
end


# open file and store content into a string
input = open("inputs//day1.txt", "r") do io
    read(io, String)
end

# split the string into a vector and parse it as integers
data = parse.(Int, split(input))

##### part 1 #####

# bit-sum of all the "true" when checking for positive in the derivation of data
result1 = sum(0 .< diff(data))

println("part1: ", result1)

##### part 2 #####

# generate the new data (gliding window of 3 terms sum)
data3 = zeros(Int, length(data) - 2)
for i = 1:length(data)-2
    data3[i] = sum(data[i:i+2])
end

# bit-sum of all the "true" when checking for positive in the derivation of data
result2 = sum(0 .< diff(data3))

println("part2: ", result2)


using Revise
	using BifurcationKit, LinearAlgebra, Plots, SparseArrays, Setfield, Parameters
	const BK = BifurcationKit

f1(u, v) = u * u * v
norminf = x -> norm(x, Inf)

function plotsol(x; kwargs...)
	N = div(length(x), 2)
	plot!(x[1:N], label="u"; kwargs...)
	plot!(x[N+1:2N], label="v"; kwargs...)
end

function Fbru!(f, x, p)
	@unpack α, β, D1, D2, l = p
	n = div(length(x), 2)
	h = 1.0 / n; h2 = h*h
	c1 = D1 / l^2 / h2
	c2 = D2 / l^2 / h2

	u = @view x[1:n]
	v = @view x[n+1:2n]

	# Dirichlet boundary conditions
	f[1]   = c1 * (α	  - 2u[1] + u[2] ) + α - (β + 1) * u[1] + f1(u[1], v[1])
	f[end] = c2 * (v[n-1] - 2v[n] + β / α)			 + β * u[n] - f1(u[n], v[n])

	f[n]   = c1 * (u[n-1] - 2u[n] +  α   ) + α - (β + 1) * u[n] + f1(u[n], v[n])
	f[n+1] = c2 * (β / α  - 2v[1] + v[2])			 + β * u[1] - f1(u[1], v[1])

	for i=2:n-1
		  f[i] = c1 * (u[i-1] - 2u[i] + u[i+1]) + α - (β + 1) * u[i] + f1(u[i], v[i])
		f[n+i] = c2 * (v[i-1] - 2v[i] + v[i+1])			  + β * u[i] - f1(u[i], v[i])
	end
	return f
end

Fbru(x, p) = Fbru!(similar(x), x, p)

function Jbru_sp(x, p)
	@unpack α, β, D1, D2, l = p
	# compute the Jacobian using a sparse representation
	n = div(length(x), 2)
	h = 1.0 / n; h2 = h*h

	c1 = D1 / p.l^2 / h2
	c2 = D2 / p.l^2 / h2

	u = @view x[1:n]
	v = @view x[n+1:2n]

	diag   = zeros(eltype(x), 2n)
	diagp1 = zeros(eltype(x), 2n-1)
	diagm1 = zeros(eltype(x), 2n-1)

	diagpn = zeros(eltype(x), n)
	diagmn = zeros(eltype(x), n)

	@. diagmn = β - 2 * u * v
	@. diagm1[1:n-1] = c1
	@. diagm1[n+1:end] = c2

	@. diag[1:n]    = -2c1 - (β + 1) + 2 * u * v
	@. diag[n+1:2n] = -2c2 - u * u

	@. diagp1[1:n-1]   = c1
	@. diagp1[n+1:end] = c2

	@. diagpn = u * u
	J = spdiagm(0 => diag, 1 => diagp1, -1 => diagm1, n => diagpn, -n => diagmn)
	return J
end

function finalise_solution(z, tau, step, contResult)
	n = div(length(z.u), 2)
	printstyled(color=:red, "--> Solution constant = ", norm(diff(z.u[1:n])), " - ", norm(diff(z.u[n+1:2n])), "\n")
	return true
end

n = 500
####################################################################################################
# test for the Jacobian expression
# using ForwardDiff
# sol0 = rand(2n)
# J0 = ForwardDiff.jacobian(x-> Fbru(x, par_bru), sol0) |> sparse
# J1 = Jbru_sp(sol0, par_bru)
# J0 - J1
####################################################################################################
# different parameters to define the Brusselator model and guess for the stationary solution
par_bru = (α = 2., β = 5.45, D1 = 0.008, D2 = 0.004, l = 0.3)
	sol0 = vcat(par_bru.α * ones(n), par_bru.β/par_bru.α * ones(n))

# # parameters for an isola of stationary solutions
# par_bru = (α = 2., β = 4.6, D1 = 0.0016, D2 = 0.008, l = 0.061)
# 	xspace = LinRange(0, par_bru.l, n)
# 	sol0 = vcat(		par_bru.α .+ 2 .* sin.(pi*xspace/par_bru.l),
# 			par_bru.β/par_bru.α .- 0.5 .* sin.(pi*xspace/par_bru.l))

# eigls = EigArpack(1.1, :LM)
# 	opt_newton = BK.NewtonPar(eigsolver = eigls)
# 	out, hist, flag = @time BK.newton(
# 		x ->    Fbru(x, par_bru),
# 		x -> Jbru_sp(x, par_bru),
# 		sol0, opt_newton, normN = norminf)
#
# 	plot();plotsol(out);plotsol(sol0, label = "sol0",line=:dash)
####################################################################################################
eigls = EigArpack(1.1, :LM)
opts_br_eq = ContinuationPar(dsmin = 0.03, dsmax = 0.05, ds = 0.03, pMax = 1.9, detectBifurcation = 3, nev = 21, plotEveryStep = 50, newtonOptions = NewtonPar(eigsolver = eigls, tol = 1e-9), maxSteps = 1060, nInversion = 6, tolBisectionEigenvalue = 1e-20, maxBisectionSteps = 30)

	br, = @time continuation(
		Fbru, Jbru_sp, sol0, par_bru, (@lens _.l),
		opts_br_eq, verbosity = 0,
		plot = true,
		plotSolution = (x, p; kwargs...) -> (plotsol(x; label="", kwargs... )),
		printSolution = (x, p) -> x[div(n,2)], normC = norminf)
####################################################################################################
jet  = BK.get3Jet(Fbru, Jbru_sp)
hopfpt = computeNormalForm(jet..., br, 1; verbose = true)
#################################################################################################### Continuation of the Hopf Point using Jacobian expression
ind_hopf = 1
	# hopfpt = BK.HopfPoint(br, ind_hopf)
	optnew = opts_br_eq.newtonOptions
	hopfpoint, _, flag = @time newton(
		Fbru, Jbru_sp,
		br, ind_hopf;
		d2F = d2Fbru,
		options = (@set optnew.verbose=true), normN = norminf)
	flag && printstyled(color=:red, "--> We found a Hopf Point at l = ", hopfpoint.p[1], ", ω = ", hopfpoint.p[2], ", from l = ", br.specialpoint[ind_hopf].param, "\n")

if 1==0
	br_hopf, u1_hopf = @time continuation(
		Fbru, Jbru_sp,
		br, ind_hopf, (@lens _.β),
		ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= 0.01, pMax = 6.5, pMin = 0.0, newtonOptions = optnew); verbosity = 2, normC = norminf)

	plot(br_hopf, label="")
end

hopfpoint, hist, flag = @time newton(
	Fbru, Jbru_sp,
	br, ind_hopf;
	options = (@set optnew.verbose = true),
	d2F = d2Fbru, normN = norminf)

if 1==1
	br_hopf, u1_hopf = @time continuation(
		Fbru, Jbru_sp,
		br, ind_hopf, (@lens _.β),
		ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds= 0.01, pMax = 10.5, pMin = 5.1, detectBifurcation = 0, newtonOptions = optnew); plot = true,
		updateMinAugEveryStep = 1,
		d2F = d2Fbru,
		verbosity = 2, normC = norminf, bothside = true)
end
plot(br_hopf)
####################################################################################################Continuation of Periodic Orbit
# number of time slices
M = 51
l_hopf, Th, orbitguess2, hopfpt, vec_hopf = guessFromHopf(br, ind_hopf, opts_br_eq.newtonOptions.eigsolver, M, 2*2.7; phase = 0.25)
#
# orbitguess_f2 = orbitguess2[1]; for ii=2:M; global orbitguess_f2 = hcat(orbitguess_f2, orbitguess2[ii]);end
orbitguess_f2 = reduce(vcat, orbitguess2)
orbitguess_f = vcat(vec(orbitguess_f2), Th) |> vec

poTrap = PeriodicOrbitTrapProblem(Fbru, Jbru_sp, real.(vec_hopf), hopfpt.u, M, 2n)

poTrap(orbitguess_f,  @set par_bru.l = l_hopf + 0.01) |> plot
BK.plotPeriodicPOTrap(orbitguess_f, n, M; ratio = 2)

ls0 = GMRESIterativeSolvers(N = 2n, reltol = 1e-9)#, Pl = lu(I + par_cgl.Δ))

poTrapMF = PeriodicOrbitTrapProblem(
			Fbru, (x, p) -> (dx -> d1Fbru(x, p, dx)),
			real.(vec_hopf),
			hopfpt.u,
			M, ls0)

deflationOp = DeflationOperator(2.0, (x,y) -> dot(x[1:end-1], y[1:end-1]),1.0, [zero(orbitguess_f)])
# deflationOp = DeflationOperator(2.0, (x,y) -> dot(x[1:end-1], y[1:end-1]),1.0, [outpo_f])
####################################################################################################
opt_po = NewtonPar(tol = 1e-10, verbose = true, maxIter = 14)
# opt_po = NewtonPar(tol = 1e-10, verbose = true, maxIter = 14, linsolver = BK.LSFromBLS())
outpo_f, _, flag = @time newton(poTrap,
			orbitguess_f, (@set par_bru.l = l_hopf + 0.01),
			opt_po;
			# deflationOp,
			linearPO = :BorderedSparseInplace, #  3.2 seconds (756.76 k allocations: 2.363 GiB, 6.67% gc time)
			# linearPO = :BorderedLU,
			# linearPO = :FullSparseInplace, # 3.634019 seconds (775.86 k allocations: 2.452 GiB, 6.40% gc time)
			normN = norminf,
			callback = (x, f, J, res, iteration, itl, options; kwargs...) -> (println("--> amplitude = ", BK.amplitude(x, n, M; ratio = 2));true)
			)
	flag && printstyled(color=:red, "--> T = ", outpo_f[end], ", amplitude = ", BK.amplitude(outpo_f, n, M; ratio = 2),"\n")
	BK.plotPeriodicPOTrap(outpo_f, n, M; ratio = 2)

opt_po = @set opt_po.eigsolver = EigKrylovKit(tol = 1e-5, x₀ = rand(2n), verbose = 2, dim = 40)
opt_po = @set opt_po.eigsolver = DefaultEig()
# opt_po = @set opt_po.eigsolver = EigArpack(; tol = 1e-5, v0 = rand(2n))
opts_po_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.1, ds= 0.01, pMax = 3.0, maxSteps = 20, newtonOptions = opt_po, saveSolEveryStep = 2, plotEveryStep = 5, nev = 11, precisionStability = 1e-6, nInversion = 4, detectBifurcation = 0, dsminBisection = 1e-6, maxBisectionSteps = 15)
	br_po, = @time continuation(poTrap,
			outpo_f, (@set par_bru.l = l_hopf + 0.01), (@lens _.l),
			opts_po_cont;
			# linearPO = :FullLU,
			updateSectionEveryStep = 1,
			# linearPO = :FullSparseInplace,
			linearPO = :BorderedSparseInplace,
			# tangentAlgo = BorderedPred(),
			verbosity = 3,	plot = true,
			# callbackN = (x, f, J, res, iteration, options; kwargs...) -> (println("--> amplitude = ", BK.amplitude(x, n, M));true),
			# finaliseSolution = (z, tau, step, contResult; k...) ->
				# (Base.display(contResult.eig[end].eigenvals) ;true),
			plotSolution = (x, p;kwargs...) -> heatmap!(reshape(x[1:end-1], 2*n, M)'; ylabel="time", color=:viridis, kwargs...),
			# printSolution = (x, p;kwargs...) -> BK.amplitude(x, n, M; ratio = 2),
			normC = norminf)

####################################################################################################
using IncompleteLU
Jpo = @time poTrap(Val(:JacFullSparse), orbitguess_f, @set par_bru.l = l_hopf + 0.01)

@time poTrap(Val(:JacFullSparseInplace), Jpo, orbitguess_f, @set par_bru.l = l_hopf + 0.01)
_indx = BK.getBlocks(Jpo, 2n, M)
@time poTrap(Val(:JacFullSparseInplace), Jpo, orbitguess_f, (@set par_bru.l = l_hopf + 0.01), _indx)

@time lu(Jpo)
Jpo = @time poTrap(Val(:JacCyclicSparse), orbitguess_f, @set par_bru.l = l_hopf + 0.01)
@time lu(Jpo)

Precilu = @time ilu(Jpo, τ = 0.01)
Precilu = @time lu(Jpo)
ls = GMRESIterativeSolvers(verbose = false, reltol = 1e-5, N = size(Jpo,1), restart = 30, maxiter = 150, log=true, Pl = Precilu)
	@time ls(Jpo, rand(ls.N))

opt_po = NewtonPar(tol = 1e-10, verbose = true, maxIter = 20)
	outpo_f, = @time newton(poTrap,
			orbitguess_f, (@set par_bru.l = l_hopf + 0.01),
			(@set opt_po.linsolver = ls); linearPO = :BorderedMatrixFree,
			normN = norminf,
			# callback = (x, f, J, res, iteration, options; kwargs...) -> (println("--> amplitude = ", amplitude(x), " T = ", x[end], ", T0 = ",orbitguess_f[end]);true)
			)
	printstyled(color=:red, "--> T = ", outpo_f[end], ", amplitude = ", BK.amplitude(outpo_f, n, M; ratio = 2),"\n")
	BK.plotPeriodicPOTrap(outpo_f, n, M; ratio = 2)

opts_po_cont = ContinuationPar(dsmin = 0.0001, dsmax = 0.05, ds= 0.01, pMax = 2.2, maxSteps = 3000, newtonOptions = (@set opt_po.linsolver = ls))
	br_pok, = @time continuation(poTrap,
			outpo_f, (@set par_bru.l = l_hopf + 0.01), (@lens _.l),
			opts_po_cont; linearPO = :BorderedMatrixFree,
			verbosity = 2,
			plot = true,
			# callbackN = (x, f, J, res, iteration, options; kwargs...) -> (println("--> amplitude = ", BK.amplitude(x, n, M));true),
			# plotSolution = (x, p;kwargs...) -> heatmap!(reshape(x[1:end-1], 2*n, M)'; ylabel="time", color=:viridis, kwargs...)
			normC = norminf)
####################################################################################################
# automatic branch switching from Hopf point
opt_po = NewtonPar(tol = 1e-10, verbose = true, maxIter = 15)
opts_po_cont = ContinuationPar(dsmin = 0.001, dsmax = 0.04, ds = 0.01, pMax = 2.2, maxSteps = 200, newtonOptions = opt_po, saveSolEveryStep = 2,
	plotEveryStep = 1, nev = 11, precisionStability = 1e-6,
	detectBifurcation = 0, dsminBisection = 1e-6, maxBisectionSteps = 15, tolBisectionEigenvalue = 3., nInversion = 4)

M = 51
probFD = PeriodicOrbitTrapProblem(M = M)
br_po, _ = continuation(
	# arguments for branch switching
	jet..., br, 1,
	# arguments for continuation
	opts_po_cont, probFD;
	δp = 0.01,
	verbosity = 3,	plot = true, linearPO = :FullSparseInplace,
	finaliseSolution = (z, tau, step, contResult; k...) ->
		(Base.display(contResult.eig[end].eigenvals) ;true),
	plotSolution = (x, p; kwargs...) -> heatmap!(getTrajectory(p.prob, x, par_bru).u'; ylabel="time", color=:viridis, kwargs...),
	normC = norminf)

# branches = [br_po]
push!(branches, br_po)
plot(branches..., legend = :bottomright, plotfold = false)

####################################################################################################
# semi-automatic branch switching from bifurcation BP-PO
br_po2, _ = BK.continuationPOTrapBPFromPO(
	# arguments for branch switching
	br_po, 1,
	# arguments for continuation
	opts_po_cont;
	δp = 0.01,
	verbosity = 3,	plot = true, linearPO = :FullSparseInplace,
	finaliseSolution = (z, tau, step, contResult; k...) ->
		(Base.display(contResult.eig[end].eigenvals) ;true),
	plotSolution = (x, p; kwargs...) -> (plot!(br_po,legend = :bottomright, subplot=1)),
	normC = norminf)
push!(branches, br_po2)

plot();for _b in branches; plot!(_b; label = ""); end; title!("")

plot(branches[2])

br_po.functional